| """The unified GRM model: encode, reason, select, verify, answer. |
|
|
| Pipeline (PLAN sections 5 and 48):: |
|
|
| graph --> GraphEncoder --------+ |
| >-- ReasoningCore --> PathSelection --> Verifier --> ReasoningResult |
| query --> QueryEncoder --> PAT + |
|
|
| Two behaviours are worth calling out because they were the substantive defects in |
| the previous version. |
|
|
| ``reason`` now *selects*. It used to return every path the verifier accepted, |
| which made the answer set "everything reachable within max_hops" and put exact |
| match out of reach on any task with distractors. Selection runs through |
| :class:`~graph.heads.PathSelectionHead` with a learned threshold, so the model |
| commits to a subset and is measured on it. |
|
|
| ``forward_example`` drives PAT and the reasoning core on every training step. |
| PAT used to be absent from the forward pass entirely (the old |
| ``GRMForwardOutput.pat`` was hard-coded to ``None``), so its only gradient came |
| from a side call in the trainer. |
| """ |
| |
| from __future__ import annotations |
| |
| import itertools |
| from collections.abc import Mapping, Sequence |
| from dataclasses import dataclass, replace |
| from itertools import islice |
| |
| import torch |
| from torch import Tensor, nn |
| |
| from graph.config import GraphConfig |
| from graph.constraints import ConstraintState |
| from graph.graph import GraphStore |
| from graph.graph_encoder import GraphEncoder, GraphEncoderOutput |
| from graph.heads import ( |
| ConfidenceHead, |
| ContradictionHead, |
| CountHead, |
| EntityAnswerHead, |
| NextHopHead, |
| PathSelectionHead, |
| RelationAnswerHead, |
| SelectionOutput, |
| SupportHead, |
| VerificationStack, |
| ) |
| from graph.pat import PathAwareTransformer, PATOutput |
| from graph.query_encoder import QueryEncoder, QueryEncoding |
| from graph.reasoner import BeamSearchReasoner, SearchOutput |
| from graph.reasoning_core import ReasoningCore |
| from graph.text_features import build_featuriser |
| from graph.types import ( |
| EntityAnswer, |
| GraphPath, |
| GraphQuery, |
| GraphTensors, |
| ReasoningResult, |
| ReasoningStatus, |
| ) |
| from graph.verification import ( |
| classify_evidence, |
| detect_conflicting_paths, |
| verify_path, |
| ) |
| |
| STATUS_INDEX: Mapping[str, int] = { |
| "supported": 0, |
| "unsupported": 1, |
| "contradicted": 2, |
| "insufficient_evidence": 3, |
| "error": 3, |
| } |
| INDEX_STATUS: tuple[ReasoningStatus, ...] = ( |
| ReasoningStatus.SUPPORTED, |
| ReasoningStatus.UNSUPPORTED, |
| ReasoningStatus.CONTRADICTED, |
| ReasoningStatus.INSUFFICIENT_EVIDENCE, |
| ) |
| |
| |
| @dataclass(frozen=True) |
| class EncodedGraph: |
| """Graph encoding plus the id-keyed views the search and PAT consume.""" |
| |
| output: GraphEncoderOutput |
| node_states: dict[str, Tensor] |
| relation_states: dict[str, Tensor] |
| memory: Tensor |
| memory_mask: Tensor |
| |
| |
| @dataclass(frozen=True) |
| class GRMForwardOutput: |
| """Every tensor the multi-task loss needs from one example.""" |
| |
| graph: GraphEncoderOutput |
| query: QueryEncoding |
| pat: PATOutput | None |
| entity_logits: Tensor |
| relation_logits: Tensor |
| support_logit: Tensor |
| contradiction_logits: Tensor |
| count_logits: Tensor |
| confidence_logit: Tensor |
| selection_logits: Tensor |
| selection_labels: Tensor |
| #: The threshold the model would apply, so the set-level objective can train it. |
| selection_threshold: Tensor | None |
| next_hop_logits: Tensor |
| next_hop_labels: Tensor |
| pat_positive: Tensor | None |
| pat_negative: Tensor | None |
| node_states: Tensor |
| entity_order: tuple[str, ...] |
| |
| |
| class GraphModel(nn.Module): |
| """Graph-native reasoning model with an explicit query-conditioned core.""" |
| |
| def __init__(self, config: GraphConfig) -> None: |
| super().__init__() |
| self.config = config |
| checkpointing = config.training.gradient_checkpointing |
| # One featuriser is shared by the query encoder and the graph tensorisation, so |
| # a relation string gets the same vector wherever the model meets it. It holds |
| # no trainable parameters, so it is deliberately not a submodule: registering a |
| # frozen pretrained encoder would put it in the checkpoint and in the parameter |
| # budget, neither of which it belongs in. |
| self.featuriser = ( |
| build_featuriser( |
| config.text_features, |
| config.text_feature_dim, |
| # The graph model builds its featuriser for training as well as |
| # evaluation, so it asks for an accelerator when there is one. |
| device="auto", |
| model_name=config.text_feature_model, |
| ) |
| if config.text_features |
| else None |
| ) |
| if self.featuriser is not None and self.featuriser.dimension != ( |
| config.graph_input_dim |
| ): |
| raise ValueError( |
| f"text_features produces {self.featuriser.dimension}-dimensional " |
| f"vectors but graph_input_dim is {config.graph_input_dim}; set them " |
| "equal so the encoder receives the features it was configured for" |
| ) |
| self.query_encoder = QueryEncoder(config, self.featuriser) |
| self.graph_encoder = GraphEncoder(config) |
| self.pat = PathAwareTransformer( |
| config.pat, |
| gradient_checkpointing=checkpointing, |
| logit_cap=config.attention_logit_cap, |
| qk_norm=config.attention_qk_norm, |
| ) |
| self.reasoning_core = ReasoningCore( |
| config.hidden_dim, |
| config.num_reasoning_layers, |
| config.num_reasoning_heads, |
| dropout=config.dropout, |
| max_hops=config.max_hops, |
| ffn_multiplier=config.reasoning_ffn_multiplier, |
| gradient_checkpointing=checkpointing, |
| logit_cap=config.attention_logit_cap, |
| qk_norm=config.attention_qk_norm, |
| ) |
| self.next_hop_head = NextHopHead(config.hidden_dim, config.dropout) |
| self.path_selection_head = PathSelectionHead(config.hidden_dim, config.dropout) |
| self.entity_answer_head = EntityAnswerHead(config.hidden_dim) |
| self.relation_answer_head = RelationAnswerHead( |
| config.hidden_dim, config.num_relations |
| ) |
| self.verification_stack = VerificationStack( |
| config.hidden_dim, |
| depth=config.verifier_layers, |
| width_multiplier=config.verifier_width_multiplier, |
| dropout=config.dropout, |
| num_signals=config.verification_signals, |
| ) |
| self.support_head = SupportHead(config.hidden_dim) |
| self.contradiction_head = ContradictionHead(config.hidden_dim) |
| self.count_head = CountHead(config.hidden_dim, config.max_count) |
| self.confidence_head = ConfidenceHead(config.hidden_dim) |
| self.reasoner = BeamSearchReasoner( |
| self.next_hop_head, |
| config.beam_width, |
| config.candidate_width, |
| config.max_hops, |
| config.relation_top_k, |
| ) |
|
|
| # -- encoding ----------------------------------------------------------- |
|
|
| def graph_to_example(self, graph: GraphStore): |
| """Tensorise ``graph`` at the widths this model's encoder expects. |
|
|
| Single point of truth for the node/edge feature widths so a config with |
| ``graph_input_dim != edge_input_dim`` cannot produce a shape mismatch. |
| """ |
| return graph.to_example( |
| feature_dim=self.config.graph_input_dim, |
| edge_feature_dim=self.config.edge_input_dim, |
| featuriser=self.featuriser, |
| ) |
|
|
| def encode_graph( |
| self, |
| graph_tensors: GraphTensors, |
| query_state: Tensor, |
| node_ids: Sequence[str], |
| relation_ids: Sequence[str], |
| *, |
| start_ids: Sequence[str], |
| use_types: bool = True, |
| ) -> EncodedGraph: |
| """Encode a graph and expose id-keyed node/relation states. |
|
|
| ``start_ids`` are the query's start entities, marked in the encoder's |
| input when the model was built with ``start_marking``. |
| """ |
| if not use_types: |
| # The untyped ablation must blank the identity the encoder actually |
| # reads, not only the local bookkeeping index; zeroing just the latter |
| # would leave the type embedding fully informed and report the |
| # unablated model. |
| graph_tensors = replace( |
| graph_tensors, |
| node_type_ids=torch.zeros_like(graph_tensors.node_type_ids), |
| node_type_hashes=torch.zeros_like(graph_tensors.node_type_identity), |
| ) |
| encoded = self.graph_encoder( |
| graph_tensors, |
| query_state, |
| start_mask=_start_mask(node_ids, start_ids, query_state.device), |
| ) |
| node_states = { |
| node_id: encoded.node_states[index] |
| for index, node_id in enumerate(node_ids) |
| } |
| relation_bank = self._pool_relation_states( |
| encoded.relation_states, graph_tensors.edge_type_ids, len(relation_ids) |
| ) |
| relation_states = { |
| relation_id: relation_bank[index] |
| for index, relation_id in enumerate(relation_ids) |
| } |
| memory, memory_mask = self.reasoning_core.build_memory( |
| encoded.node_states, relation_bank |
| ) |
| return EncodedGraph(encoded, node_states, relation_states, memory, memory_mask) |
|
|
| @staticmethod |
| def _pool_relation_states( |
| edge_states: Tensor, edge_type_ids: Tensor, num_relations: int |
| ) -> Tensor: |
| """Mean-pool per-edge states into one state per relation type. |
| |
| Vectorised with ``index_add_`` rather than accumulated in a Python loop. |
| The loop built one autograd node per edge, so a graph with a few hundred |
| edges produced a backward chain hundreds of nodes deep and the backward |
| pass cost roughly twenty times the forward pass. Scatter-add is a single |
| node and the ratio returns to normal. |
| """ |
| dim = edge_states.shape[-1] |
| if not num_relations or edge_states.shape[0] == 0: |
| return edge_states.new_zeros((max(num_relations, 0), dim)) |
| safe_ids = edge_type_ids.clamp(0, num_relations - 1) |
| sums = edge_states.new_zeros((num_relations, dim)) |
| sums = sums.index_add(0, safe_ids, edge_states) |
| counts = edge_states.new_zeros((num_relations,)) |
| counts = counts.index_add( |
| 0, safe_ids, torch.ones_like(safe_ids, dtype=edge_states.dtype) |
| ) |
| return sums / counts.clamp_min(1.0).unsqueeze(-1) |
|
|
| # -- training ----------------------------------------------------------- |
|
|
| def forward_example( |
| self, |
| graph: GraphStore, |
| query: GraphQuery, |
| *, |
| supporting_paths: Sequence[GraphPath], |
| competing_paths: Sequence[GraphPath], |
| answers: Sequence[str], |
| status: ReasoningStatus, |
| feature_dim: int | None = None, |
| ) -> GRMForwardOutput: |
| """Compute every training signal for one example in a single pass.""" |
| example = self.graph_to_example(graph) |
| device = next(self.parameters()).device |
| graph_tensors = example.graph.to(device) |
| query_encoding = self.query_encoder.encode(query) |
| query_state = query_encoding.pooled[0] |
| encoded = self.encode_graph( |
| graph_tensors, |
| query_state, |
| example.node_external_ids, |
| example.relation_external_ids, |
| start_ids=_query_starts(query), |
| ) |
| constraints = ConstraintState.from_query(query) |
|
|
| positives = tuple(supporting_paths)[: self.config.max_supervised_paths] |
| if self.config.on_policy_selection: |
| positives, negatives = self._on_policy_candidates( |
| graph, query, encoded, query_state, positives, competing_paths |
| ) |
| else: |
| negatives = self._selection_negatives( |
| graph, query, positives, competing_paths |
| ) |
| labelled = (*positives, *negatives) |
|
|
| pat_output: PATOutput | None = None |
| pat_positive: Tensor | None = None |
| pat_negative: Tensor | None = None |
| if labelled: |
| pat_output = self.pat.encode_paths( |
| labelled, |
| encoded.node_states, |
| encoded.relation_states, |
| query_state, |
| constraints, |
| ) |
| path_states = pat_output.path_states |
| if positives: |
| pat_positive = pat_output.compatibility[: len(positives)] |
| if negatives: |
| pat_negative = pat_output.compatibility[len(positives) :] |
| core = self.reasoning_core( |
| query_state, |
| encoded.memory, |
| encoded.memory_mask, |
| path_states.unsqueeze(0), |
| hop=1, |
| ) |
| selection = self.path_selection_head( |
| core.candidate_states[0], path_states, query_state |
| ) |
| selection_logits = selection.logits |
| selection_threshold = selection.threshold |
| selection_labels = torch.cat( |
| ( |
| selection_logits.new_ones(len(positives)), |
| selection_logits.new_zeros(len(negatives)), |
| ) |
| ) |
| reasoning_state = core.state[0] |
| evidence_state = path_states.mean(0) |
| else: |
| zeros = query_state.new_zeros((0,)) |
| selection_logits, selection_labels = zeros, zeros |
| selection_threshold = None |
| core = self.reasoning_core( |
| query_state, |
| encoded.memory, |
| encoded.memory_mask, |
| query_state.new_zeros((1, 1, query_state.shape[-1])), |
| hop=1, |
| ) |
| reasoning_state = core.state[0] |
| evidence_state = torch.zeros_like(reasoning_state) |
|
|
| entity_logits = self.entity_answer_head( |
| reasoning_state, encoded.output.node_states |
| ) |
| relation_context = ( |
| torch.stack(list(encoded.relation_states.values())).mean(0) |
| if encoded.relation_states |
| else torch.zeros_like(reasoning_state) |
| ) |
| relation_logits = self.relation_answer_head( |
| reasoning_state.unsqueeze(0), relation_context.unsqueeze(0) |
| ) |
| verified_state = self.verification_stack( |
| reasoning_state.unsqueeze(0), |
| evidence_state.unsqueeze(0), |
| query_state.unsqueeze(0), |
| self._evidence_signals( |
| pat_output.compatibility if pat_output is not None else None, |
| len(labelled), |
| query_state.device, |
| ), |
| ) |
| support_logit = self.support_head(verified_state) |
| contradiction_logits = self.contradiction_head(verified_state) |
| count_logits = self.count_head(verified_state) |
| selected_fraction = ( |
| (selection_logits > selection_threshold).float().mean() |
| if selection_logits.numel() and selection_threshold is not None |
| else torch.zeros((), device=device) |
| ) |
| confidence_logit = self.confidence_head( |
| verified_state, |
| support_logit.detach().reshape(1), |
| selected_fraction.reshape(1), |
| ) |
| next_hop_logits, next_hop_labels = self._next_hop_supervision( |
| graph, |
| query, |
| positives, |
| encoded, |
| core.candidate_states[0] if labelled else None, |
| ) |
| return GRMForwardOutput( |
| graph=encoded.output, |
| query=query_encoding, |
| pat=pat_output, |
| entity_logits=entity_logits, |
| relation_logits=relation_logits, |
| support_logit=support_logit, |
| contradiction_logits=contradiction_logits, |
| count_logits=count_logits, |
| confidence_logit=confidence_logit, |
| selection_logits=selection_logits, |
| selection_labels=selection_labels, |
| selection_threshold=selection_threshold, |
| next_hop_logits=next_hop_logits, |
| next_hop_labels=next_hop_labels, |
| pat_positive=pat_positive, |
| pat_negative=pat_negative, |
| node_states=encoded.output.node_states, |
| entity_order=example.node_external_ids, |
| ) |
|
|
| def forward_batch( |
| self, |
| graphs: Sequence[GraphStore], |
| queries: Sequence[GraphQuery], |
| supporting_paths: Sequence[Sequence[GraphPath]], |
| competing_paths: Sequence[Sequence[GraphPath]], |
| ) -> list[GRMForwardOutput]: |
| """:meth:`forward_example` for a batch, with each network run once. |
| |
| Per-example training launched every layer once per example: graphs have |
| different shapes, and the model's layers are small, so a step was bound by |
| kernel launches and Python overhead with the GPU about 15% busy. Here the |
| graphs are encoded as one disjoint union, the queries, PAT paths and |
| reasoning tokens are padded and masked, and the heads run over flattened |
| rows. Each example's outputs equal :meth:`forward_example`'s up to |
| floating-point reordering (and dropout draws); the equivalence is tested. |
| """ |
| if self.config.on_policy_selection: |
| # The on-policy search is sequential per example; nothing to batch. |
| raise ValueError("forward_batch does not support on_policy_selection") |
| count = len(graphs) |
| if not count: |
| return [] |
| device = next(self.parameters()).device |
| hidden = self.config.hidden_dim |
|
|
| # -- tensorise (text features encoded in one call) ---------------------- |
| if self.featuriser is not None: |
| self.featuriser.prefetch_identifiers( |
| [ |
| text |
| for graph in graphs |
| for text in ( |
| *( |
| f"{node_id} {graph.nodes[node_id]}" |
| for node_id in graph.nodes |
| ), |
| *( |
| f"{edge.relation_id} {edge.relation_name or ''}" |
| for edge in graph.edges |
| ), |
| ) |
| ] |
| ) |
| examples = [self.graph_to_example(graph) for graph in graphs] |
| tensors = [example.graph for example in examples] |
| node_counts = [item.num_nodes for item in tensors] |
| edge_counts = [item.num_edges for item in tensors] |
| relation_counts = [len(example.relation_external_ids) for example in examples] |
| node_starts = [0, *itertools.accumulate(node_counts)] |
| edge_starts = [0, *itertools.accumulate(edge_counts)] |
| relation_starts = [0, *itertools.accumulate(relation_counts)] |
| union = GraphTensors( |
| node_ids=torch.arange(node_starts[-1], dtype=torch.long), |
| node_type_ids=torch.cat([item.node_type_ids for item in tensors]), |
| node_features=torch.cat([item.node_features for item in tensors]), |
| edge_index=torch.cat( |
| [ |
| item.edge_index + start |
| for item, start in zip(tensors, node_starts, strict=False) |
| ], |
| dim=1, |
| ), |
| edge_type_ids=torch.cat( |
| [ |
| item.edge_type_ids + start |
| for item, start in zip(tensors, relation_starts, strict=False) |
| ] |
| ), |
| edge_features=torch.cat([item.edge_features for item in tensors]), |
| graph_id=torch.repeat_interleave( |
| torch.arange(count), torch.tensor(node_counts) |
| ), |
| node_type_hashes=torch.cat([item.node_type_identity for item in tensors]), |
| edge_type_hashes=torch.cat([item.edge_type_identity for item in tensors]), |
| ).to(device) |
| edge_owner = torch.repeat_interleave( |
| torch.arange(count, device=device), |
| torch.tensor(edge_counts, device=device), |
| ) |
| node_owner = union.graph_id |
| relation_owner = torch.repeat_interleave( |
| torch.arange(count, device=device), |
| torch.tensor(relation_counts, device=device), |
| ) |
|
|
| # -- query and graph encoding ------------------------------------------ |
| query_encodings = self.query_encoder.encode_batch(queries) |
| query_states = torch.cat([encoding.pooled for encoding in query_encodings]) |
| encoded = self.graph_encoder( |
| union, |
| query_states, |
| edge_query=query_states[edge_owner], |
| start_mask=torch.cat( |
| [ |
| _start_mask(ex.node_external_ids, _query_starts(query), device) |
| for ex, query in zip(examples, queries, strict=True) |
| ] |
| ), |
| node_query=query_states[node_owner], |
| ) |
| # ``edge_type_ids`` were offset per graph, so pooling over the union |
| # yields every graph's relation bank, stacked. |
| relation_bank = self._pool_relation_states( |
| encoded.relation_states, union.edge_type_ids, relation_starts[-1] |
| ) |
| node_rows = [ |
| { |
| node_id: node_starts[i] + j |
| for j, node_id in enumerate(ex.node_external_ids) |
| } |
| for i, ex in enumerate(examples) |
| ] |
| relation_rows = [ |
| { |
| relation_id: relation_starts[i] + j |
| for j, relation_id in enumerate(ex.relation_external_ids) |
| } |
| for i, ex in enumerate(examples) |
| ] |
| memory, memory_mask = self._batch_memory( |
| encoded.node_states, |
| relation_bank, |
| node_counts, |
| relation_counts, |
| node_starts, |
| relation_starts, |
| ) |
|
|
| # -- candidate paths through PAT --------------------------------------- |
| positives_per: list[tuple[GraphPath, ...]] = [] |
| negatives_per: list[tuple[GraphPath, ...]] = [] |
| for graph, query, supporting, competing in zip( |
| graphs, queries, supporting_paths, competing_paths, strict=True |
| ): |
| positives = tuple(supporting)[: self.config.max_supervised_paths] |
| positives_per.append(positives) |
| negatives_per.append( |
| self._selection_negatives(graph, query, positives, competing) |
| ) |
| labelled_per = [ |
| (*positives, *negatives) |
| for positives, negatives in zip(positives_per, negatives_per, strict=True) |
| ] |
| path_counts = [len(paths) for paths in labelled_per] |
| owners = [i for i, paths in enumerate(labelled_per) for _ in paths] |
| path_starts = [0, *itertools.accumulate(path_counts)] |
| # Example and within-example position of every labelled path. |
| path_owner = torch.tensor(owners, dtype=torch.long, device=device) |
| slots = torch.tensor( |
| [j for n in path_counts for j in range(n)], dtype=torch.long, device=device |
| ) |
| pat_output: PATOutput | None = None |
| if owners: |
| pat_input = self.pat.build_batch_input( |
| [path for paths in labelled_per for path in paths], |
| owners, |
| node_rows, |
| relation_rows, |
| encoded.node_states, |
| relation_bank, |
| query_states, |
| [ConstraintState.from_query(query) for query in queries], |
| ) |
| pat_output = self.pat(pat_input, query_states[path_owner]) |
|
|
| # -- reasoning core over each example's candidates --------------------- |
| width = max(1, *path_counts) |
| candidates = query_states.new_zeros((count, width, hidden)) |
| candidate_mask = torch.zeros((count, width), dtype=torch.long, device=device) |
| # An example without candidates is given one zero candidate, exactly as |
| # forward_example does. |
| candidate_mask[:, 0] = 1 |
| if pat_output is not None: |
| candidates = candidates.index_put( |
| (path_owner, slots), pat_output.path_states.to(candidates.dtype) |
| ) |
| candidate_mask[path_owner, slots] = 1 |
| core = self.reasoning_core( |
| query_states, |
| memory, |
| memory_mask, |
| candidates, |
| hop=1, |
| candidate_mask=candidate_mask, |
| ) |
| reasoning_states = core.state |
|
|
| # -- heads ------------------------------------------------------------- |
| path_totals = torch.tensor(path_counts, device=device).clamp_min(1) |
| selection_logits_all: Tensor | None = None |
| thresholds: Tensor | None = None |
| evidence = torch.zeros_like(reasoning_states) |
| if pat_output is not None: |
| path_states = pat_output.path_states |
| candidate_rows = core.candidate_states[path_owner, slots] |
| head = self.path_selection_head |
| selection_logits_all = head.score( |
| torch.cat( |
| (candidate_rows, path_states, query_states[path_owner]), dim=-1 |
| ) |
| ).squeeze(-1) |
| evidence = ( |
| path_states.new_zeros((count, path_states.shape[-1])).index_add( |
| 0, path_owner, path_states |
| ) |
| / path_totals.unsqueeze(-1).to(path_states.dtype) |
| ).to(reasoning_states.dtype) |
| thresholds = head.threshold + head.threshold_offset( |
| torch.cat((query_states, evidence.to(query_states.dtype)), dim=-1) |
| ).squeeze(-1) |
|
|
| projected = self.entity_answer_head.project(reasoning_states) |
| entity_logits_all = self.entity_answer_head.scale * ( |
| encoded.node_states * projected[node_owner] |
| ).sum(-1) |
| relation_context = relation_bank.new_zeros((count, hidden)).index_add( |
| 0, relation_owner, relation_bank |
| ) / torch.tensor(relation_counts, device=device).clamp_min(1).unsqueeze(-1).to( |
| relation_bank.dtype |
| ) |
| relation_logits = self.relation_answer_head( |
| reasoning_states, relation_context.to(reasoning_states.dtype) |
| ) |
| signals = self._batch_evidence_signals(pat_output, owners, path_counts, device) |
| verified = self.verification_stack( |
| reasoning_states, evidence, query_states, signals |
| ) |
| support_logits = self.support_head(verified) |
| contradiction_logits = self.contradiction_head(verified) |
| count_logits = self.count_head(verified) |
| if selection_logits_all is not None and thresholds is not None: |
| kept = (selection_logits_all > thresholds[path_owner]).to(verified.dtype) |
| selected_fraction = verified.new_zeros((count,)).index_add( |
| 0, path_owner, kept |
| ) / path_totals.to(verified.dtype) |
| else: |
| selected_fraction = verified.new_zeros((count,)) |
| confidence_logits = self.confidence_head.score( |
| torch.cat( |
| ( |
| verified, |
| support_logits.detach().unsqueeze(-1), |
| selected_fraction.unsqueeze(-1), |
| ), |
| dim=-1, |
| ) |
| ).squeeze(-1) |
| next_hop = self._batch_next_hop_supervision( |
| graphs, |
| positives_per, |
| path_counts, |
| node_rows, |
| relation_rows, |
| encoded.node_states, |
| relation_bank, |
| core.candidate_states, |
| ) |
|
|
| # -- per-example views ------------------------------------------------- |
| outputs: list[GRMForwardOutput] = [] |
| for i in range(count): |
| nodes = slice(node_starts[i], node_starts[i + 1]) |
| edges = slice(edge_starts[i], edge_starts[i + 1]) |
| paths = slice(path_starts[i], path_starts[i + 1]) |
| n_pos = len(positives_per[i]) |
| n_paths = path_counts[i] |
| if n_paths and selection_logits_all is not None and thresholds is not None: |
| selection_logits = selection_logits_all[paths] |
| selection_labels = torch.cat( |
| ( |
| selection_logits.new_ones(n_pos), |
| selection_logits.new_zeros(n_paths - n_pos), |
| ) |
| ) |
| threshold: Tensor | None = thresholds[i] |
| else: |
| selection_logits = query_states.new_zeros((0,)) |
| selection_labels = selection_logits |
| threshold = None |
| compatibility = ( |
| pat_output.compatibility[paths] |
| if pat_output is not None and n_paths |
| else None |
| ) |
| outputs.append( |
| GRMForwardOutput( |
| graph=GraphEncoderOutput( |
| encoded.node_states[nodes], |
| encoded.relation_states[edges], |
| encoded.edge_states[edges], |
| None, |
| ), |
| query=query_encodings[i], |
| pat=_slice_pat(pat_output, paths) if n_paths else None, |
| entity_logits=entity_logits_all[nodes], |
| relation_logits=relation_logits[i : i + 1], |
| support_logit=support_logits[i : i + 1], |
| contradiction_logits=contradiction_logits[i : i + 1], |
| count_logits=count_logits[i : i + 1], |
| confidence_logit=confidence_logits[i : i + 1], |
| selection_logits=selection_logits, |
| selection_labels=selection_labels, |
| selection_threshold=threshold, |
| next_hop_logits=next_hop[i][0], |
| next_hop_labels=next_hop[i][1], |
| pat_positive=compatibility[:n_pos] |
| if compatibility is not None and n_pos |
| else None, |
| pat_negative=compatibility[n_pos:] |
| if compatibility is not None and n_paths > n_pos |
| else None, |
| node_states=encoded.node_states[nodes], |
| entity_order=examples[i].node_external_ids, |
| ) |
| ) |
| return outputs |
|
|
| def _batch_memory( |
| self, |
| node_states: Tensor, |
| relation_bank: Tensor, |
| node_counts: Sequence[int], |
| relation_counts: Sequence[int], |
| node_starts: Sequence[int], |
| relation_starts: Sequence[int], |
| ) -> tuple[Tensor, Tensor]: |
| """Padded ``[B, M, D]`` graph memory: each example's nodes, then relations. |
|
|
| Matches :meth:`ReasoningCore.build_memory` row for row; the padding is |
| masked out of the cross-attention. |
| """ |
| core = self.reasoning_core |
| device = node_states.device |
| count = len(node_counts) |
| lengths = [n + r for n, r in zip(node_counts, relation_counts, strict=True)] |
| width = max(lengths) |
| roles = core.memory_role_embedding(torch.arange(2, device=device)) |
| nodes = core.memory_projection(node_states) + roles[0] |
| destination = [ |
| i * width + j for i, n in enumerate(node_counts) for j in range(n) |
| ] |
| rows = [nodes] |
| if relation_bank.shape[0]: |
| rows.append(core.memory_projection(relation_bank) + roles[1]) |
| destination += [ |
| i * width + node_counts[i] + j |
| for i, r in enumerate(relation_counts) |
| for j in range(r) |
| ] |
| values = torch.cat(rows) |
| memory = values.new_zeros((count * width, values.shape[-1])).index_copy( |
| 0, torch.tensor(destination, dtype=torch.long, device=device), values |
| ) |
| mask = torch.zeros((count, width), dtype=torch.long) |
| for i, length in enumerate(lengths): |
| mask[i, :length] = 1 |
| return memory.view(count, width, -1), mask.to(device) |
|
|
| def _batch_evidence_signals( |
| self, |
| pat_output: PATOutput | None, |
| owners: Sequence[int], |
| path_counts: Sequence[int], |
| device: torch.device, |
| ) -> Tensor | None: |
| """:meth:`_evidence_signals` for every example in a batch.""" |
| if not self.config.verification_signals: |
| return None |
| count = len(path_counts) |
| maximum = torch.zeros((count,), device=device) |
| mean = torch.zeros((count,), device=device) |
| if pat_output is not None: |
| owner = torch.tensor(owners, dtype=torch.long, device=device) |
| compatibility = pat_output.compatibility.reshape(-1).float() |
| maximum = maximum.scatter_reduce( |
| 0, owner, compatibility, reduce="amax", include_self=False |
| ) |
| mean = ( |
| mean.index_add(0, owner, compatibility) |
| / torch.tensor(path_counts, device=device).clamp_min(1).float() |
| ) |
| sizes = torch.log1p(torch.tensor(path_counts, device=device).float()) / 10.0 |
| signals = torch.stack((maximum, mean, sizes), dim=-1) |
| return signals[:, : self.config.verification_signals] |
|
|
| def _batch_next_hop_supervision( |
| self, |
| graphs: Sequence[GraphStore], |
| positives_per: Sequence[Sequence[GraphPath]], |
| path_counts: Sequence[int], |
| node_rows: Sequence[Mapping[str, int]], |
| relation_rows: Sequence[Mapping[str, int]], |
| node_states: Tensor, |
| relation_bank: Tensor, |
| candidate_states: Tensor, |
| ) -> list[tuple[Tensor, Tensor]]: |
| """:meth:`_next_hop_supervision` for a batch, scored in one head call.""" |
| device = node_states.device |
| empty = torch.zeros((0,), device=device) |
| width = candidate_states.shape[1] |
| # Row 0 of each lookup table is the zero vector a missing id falls back to. |
| state_index: list[int] = [] |
| source_index: list[int] = [] |
| relation_index: list[int] = [] |
| target_index: list[int] = [] |
| labels: list[float] = [] |
| spans: list[tuple[int, int] | None] = [] |
| for i, (graph, positives) in enumerate(zip(graphs, positives_per, strict=True)): |
| if not positives or not positives[0].steps or not path_counts[i]: |
| spans.append(None) |
| continue |
| gold = positives[0].steps[0] |
| siblings = graph.outgoing(gold.source_id) |
| if len(siblings) < 2: |
| spans.append(None) |
| continue |
| start = len(labels) |
| nodes, relations = node_rows[i], relation_rows[i] |
| for edge in siblings[: self.config.candidate_width]: |
| state_index.append(i * width) |
| source = nodes.get(edge.source_id) |
| relation = relations.get(edge.relation_id) |
| target = nodes.get(edge.target_id) |
| source_index.append(0 if source is None else source + 1) |
| relation_index.append(0 if relation is None else relation + 1) |
| target_index.append(0 if target is None else target + 1) |
| labels.append( |
| float( |
| edge.relation_id == gold.relation_id |
| and edge.target_id == gold.target_id |
| ) |
| ) |
| spans.append((start, len(labels))) |
| if not labels: |
| return [(empty, empty) for _ in graphs] |
|
|
| def table(values: Tensor) -> Tensor: |
| return torch.cat((values.new_zeros((1, values.shape[-1])), values)) |
|
|
| def rows(values: Tensor, index: list[int]) -> Tensor: |
| return values[torch.tensor(index, dtype=torch.long, device=device)] |
|
|
| node_table = table(node_states) |
| logits = self.next_hop_head( |
| rows(candidate_states.reshape(-1, candidate_states.shape[-1]), state_index), |
| rows(node_table, source_index), |
| rows(table(relation_bank), relation_index), |
| rows(node_table, target_index), |
| ) |
| label_tensor = torch.tensor(labels, device=device) |
| return [ |
| (empty, empty) |
| if span is None |
| else (logits[span[0] : span[1]], label_tensor[span[0] : span[1]]) |
| for span in spans |
| ] |
|
|
| def _selection_negatives( |
| self, |
| graph: GraphStore, |
| query: GraphQuery, |
| positives: Sequence[GraphPath], |
| competing_paths: Sequence[GraphPath], |
| ) -> tuple[GraphPath, ...]: |
| """Negatives drawn from the paths the search will actually surface. |
| |
| Training on the generator's curated ``competing_paths`` alone leaves a |
| train/inference mismatch: those are same-length routes from the same start, |
| while the beam returns paths of *every* admissible length. Worse, an |
| example with no competing paths has an all-positive label vector, so both |
| the per-path and set-level objectives are minimised by keeping everything — |
| and the measured result was a model that selected all 14 candidates when |
| the gold set averaged 2. |
| |
| Curated competitors come first, because they are the hardest negatives; |
| the remainder is filled from the reachable-path distribution the beam sees. |
| """ |
| budget = max( |
| self.config.max_supervised_paths, |
| len(positives) * self.config.selection_negative_ratio, |
| ) |
| gold = {_path_key(path) for path in positives} |
| chosen: list[GraphPath] = [] |
| seen: set[tuple[tuple[str, str, str], ...]] = set() |
| for path in competing_paths: |
| key = _path_key(path) |
| if key in gold or key in seen: |
| continue |
| seen.add(key) |
| chosen.append(path) |
| if len(chosen) >= budget: |
| return tuple(chosen) |
| starts = query.start_entity_ids or query.start_entities |
| horizon = min(query.max_hops, self.config.max_hops) |
| for path in islice( |
| graph.iter_paths(starts, horizon), self.config.reachable_negative_pool |
| ): |
| if path.hop_count < query.min_hops: |
| continue |
| key = _path_key(path) |
| if key in gold or key in seen: |
| continue |
| seen.add(key) |
| chosen.append(path) |
| if len(chosen) >= budget: |
| break |
| return tuple(chosen) |
|
|
| def _on_policy_candidates( |
| self, |
| graph: GraphStore, |
| query: GraphQuery, |
| encoded: EncodedGraph, |
| query_state: Tensor, |
| positives: Sequence[GraphPath], |
| competing_paths: Sequence[GraphPath], |
| ) -> tuple[tuple[GraphPath, ...], tuple[GraphPath, ...]]: |
| """Label the candidates the search itself produces (on-policy selection). |
| |
| Training on curated negatives leaves the selection head facing a different |
| distribution at inference than the one it was fitted on. The beam picks |
| candidates with ``next_hop_head``, which is trained to rank the gold *first* |
| hop highly, so its output is dominated by paths that start correctly and |
| diverge later -- precisely the negatives an enumeration-order sample misses. |
| Off-policy training therefore produced a threshold that kept ten candidates |
| where the gold set held two, no matter how the base rate was balanced. |
| |
| The search runs without gradients (path enumeration needs none); the chosen |
| paths are then re-encoded with gradients by the caller. Gold paths the beam |
| missed are appended, because a batch with no positive carries no signal for |
| what to keep. |
| """ |
| budget = self.config.max_supervised_paths |
| gold_keys = {_path_key(path) for path in positives} |
| with torch.no_grad(): |
| search = self.reasoner.search( |
| graph, |
| query, |
| query_state, |
| encoded.node_states, |
| encoded.relation_states, |
| self.pat, |
| self.reasoning_core, |
| encoded.memory, |
| encoded.memory_mask, |
| use_pat=self.config.pat.enabled, |
| ) |
| found_positive: list[GraphPath] = [] |
| found_negative: list[GraphPath] = [] |
| for path in search.paths: |
| if _path_key(path) in gold_keys: |
| found_positive.append(path) |
| else: |
| found_negative.append(path) |
| # Keep every gold path so the objective always has something to keep. |
| missed = [ |
| path |
| for path in positives |
| if _path_key(path) not in {_path_key(p) for p in found_positive} |
| ] |
| chosen_positive = (*found_positive[:budget], *missed[:budget])[:budget] |
| if not chosen_positive: |
| chosen_positive = tuple(positives[:budget]) |
| limit = max(budget, len(chosen_positive) * self.config.selection_negative_ratio) |
| chosen_negative = tuple(found_negative[:limit]) |
| if not chosen_negative: |
| chosen_negative = self._selection_negatives( |
| graph, query, chosen_positive, competing_paths |
| ) |
| return tuple(chosen_positive), chosen_negative |
|
|
| def _next_hop_supervision( |
| self, |
| graph: GraphStore, |
| query: GraphQuery, |
| positives: Sequence[GraphPath], |
| encoded: EncodedGraph, |
| candidate_states: Tensor | None, |
| ) -> tuple[Tensor, Tensor]: |
| """Label the first transition of the gold path against its siblings. |
| |
| This is the signal that teaches the beam which relation to take out of |
| the start entity -- the decision the leaked composition used to make. |
| """ |
| device = encoded.output.node_states.device |
| if not positives or not positives[0].steps or candidate_states is None: |
| empty = torch.zeros((0,), device=device) |
| return empty, empty |
| gold = positives[0].steps[0] |
| start = gold.source_id |
| siblings = graph.outgoing(start) |
| if len(siblings) < 2: |
| empty = torch.zeros((0,), device=device) |
| return empty, empty |
| zero = torch.zeros_like(encoded.output.node_states[0]) |
| state = candidate_states[0] |
| rows, labels = [], [] |
| for edge in siblings[: self.config.candidate_width]: |
| rows.append( |
| ( |
| state, |
| encoded.node_states.get(edge.source_id, zero), |
| encoded.relation_states.get(edge.relation_id, zero), |
| encoded.node_states.get(edge.target_id, zero), |
| ) |
| ) |
| labels.append( |
| float( |
| edge.relation_id == gold.relation_id |
| and edge.target_id == gold.target_id |
| ) |
| ) |
| logits = self.next_hop_head( |
| torch.stack([row[0] for row in rows]), |
| torch.stack([row[1] for row in rows]), |
| torch.stack([row[2] for row in rows]), |
| torch.stack([row[3] for row in rows]), |
| ) |
| return logits, torch.tensor(labels, device=device) |
|
|
| # -- inference ---------------------------------------------------------- |
|
|
| @torch.no_grad() |
| def reason( |
| self, |
| graph: GraphStore, |
| query: GraphQuery, |
| *, |
| use_pat: bool | None = None, |
| verify: bool | None = None, |
| beam_width: int | None = None, |
| use_types: bool = True, |
| query_conditioned: bool = True, |
| select: bool = True, |
| abstain: bool = True, |
| ) -> ReasoningResult: |
| """Answer ``query`` against ``graph`` and return a verified typed result. |
| |
| ``abstain=False`` disables the learned support gate, leaving the deterministic |
| verifier as the only thing that can empty the answer set. It exists because the |
| gate is distribution-sensitive: its evidence signals come from PAT |
| query-compatibility, which saturates low on unseen relation vocabulary, and on |
| real typed KGs it fired on every query and drove every metric to 0.000. |
| """ |
| was_training = self.training |
| self.eval() |
| try: |
| return self._reason( |
| graph, |
| query, |
| use_pat=self.config.pat.enabled if use_pat is None else use_pat, |
| verify_enabled=query.verify if verify is None else verify, |
| beam_width=beam_width, |
| use_types=use_types, |
| query_conditioned=query_conditioned, |
| select=select, |
| abstain=abstain, |
| ) |
| finally: |
| if was_training: |
| self.train() |
|
|
| def _reason( |
| self, |
| graph: GraphStore, |
| query: GraphQuery, |
| *, |
| use_pat: bool, |
| verify_enabled: bool, |
| beam_width: int | None, |
| use_types: bool, |
| query_conditioned: bool, |
| select: bool, |
| abstain: bool = True, |
| ) -> ReasoningResult: |
| example = self.graph_to_example(graph) |
| device = next(self.parameters()).device |
| graph_tensors = example.graph.to(device) |
| query_encoding = self.query_encoder.encode(query) |
| query_state = query_encoding.pooled[0] |
| if not query_conditioned: |
| # The ablation must remove query conditioning from *every* path the |
| # query takes -- graph encoding, the reasoning core and path |
| # selection. Zeroing it in only one of them leaves the other two |
| # conditioned and the ablation reports the unablated model. |
| query_state = torch.zeros_like(query_state) |
| encoded = self.encode_graph( |
| graph_tensors, |
| query_state, |
| example.node_external_ids, |
| example.relation_external_ids, |
| use_types=use_types, |
| start_ids=_query_starts(query), |
| ) |
| search = self.reasoner.search( |
| graph, |
| query, |
| query_state, |
| encoded.node_states, |
| encoded.relation_states, |
| self.pat, |
| self.reasoning_core, |
| encoded.memory, |
| encoded.memory_mask, |
| use_pat=use_pat, |
| beam_width=beam_width, |
| query_conditioned=query_conditioned, |
| ) |
|
|
| candidates = search.paths |
| selection_scores: list[float] = [] |
| kept: tuple[GraphPath, ...] = candidates |
| predicted_count: int | None = None |
| if candidates and search.candidate_states is not None: |
| selection = self.path_selection_head( |
| search.candidate_states, search.path_states, query_state |
| ) |
| selection_scores = selection.logits.detach().cpu().tolist() |
| scored = tuple( |
| replace( |
| path, |
| score=float(score), |
| query_compatibility=_compat(search.pat_output, index), |
| ) |
| for index, (path, score) in enumerate( |
| zip(candidates, selection_scores, strict=True) |
| ) |
| ) |
| if select: |
| # Abstention is decided first and independently of the selection |
| # rule. Count-guided selection always returns at least one path, |
| # so gating inside it would make INSUFFICIENT_EVIDENCE unreachable |
| # exactly as the old unconditional fallback did. Whether the graph |
| # supports *any* answer is the verifier's question (PLAN 19/20); |
| # how many paths to commit to is the selection rule's. |
| if abstain and not self._supports_any_answer(search, query_state): |
| kept = () |
| else: |
| predicted_count = self._predicted_answer_count(search, query_state) |
| kept = self._select_paths(scored, selection, predicted_count) |
| if not kept: |
| kept = (max(scored, key=lambda item: item.score),) |
| else: |
| kept = scored |
|
|
| verified = ( |
| tuple(verify_path(path, graph, query) for path in kept) |
| if verify_enabled |
| else kept |
| ) |
| supporting = tuple( |
| path for path in verified if path.verified or not verify_enabled |
| ) |
| conflicts = tuple( |
| path |
| for group in detect_conflicting_paths(supporting, graph) |
| for path in group |
| ) |
| status = self._status( |
| supporting, |
| conflicts, |
| graph, |
| query, |
| verify_enabled, |
| search, |
| encoded, |
| query_state, |
| ) |
| answers_by_id: dict[str, float] = {} |
| for path in supporting: |
| if path.endpoint_id is not None: |
| answers_by_id[path.endpoint_id] = max( |
| answers_by_id.get(path.endpoint_id, float("-inf")), path.score |
| ) |
| ranked = sorted(answers_by_id.items(), key=lambda item: item[1], reverse=True) |
| answers = tuple( |
| EntityAnswer(entity_id, float(score)) |
| for entity_id, score in ranked[: self.config.answer_top_k] |
| ) |
| relation_answers = tuple( |
| dict.fromkeys( |
| path.steps[-1].relation_id for path in supporting if path.steps |
| ) |
| ) |
| boolean_answer, aggregate_value = self._mode_answers( |
| query, status, answers_by_id, supporting |
| ) |
| confidence = self._confidence(supporting, status) |
| if ( |
| query.target_entity_ids |
| and status is ReasoningStatus.SUPPORTED |
| and not set(query.target_entity_ids) & set(answers_by_id) |
| ): |
| status = ReasoningStatus.UNSUPPORTED |
| boolean_answer = False if query.answer_mode == "boolean" else boolean_answer |
| return ReasoningResult( |
| status=status, |
| answers=answers |
| if status is ReasoningStatus.SUPPORTED and query.answer_mode == "entity" |
| else (), |
| supporting_paths=supporting if query.return_paths else (), |
| conflicting_paths=conflicts if query.return_paths else (), |
| confidence=confidence, |
| calibrated=False, |
| model_version=self.config.model_version, |
| graph_version=graph.graph_version, |
| reasoning_steps=max((path.hop_count for path in supporting), default=0), |
| metadata={ |
| "pat_enabled": use_pat, |
| "verifier_enabled": verify_enabled, |
| "selection_enabled": select, |
| "query_conditioned": query_conditioned, |
| "typed_graph": use_types, |
| "beam_width": beam_width or self.config.beam_width, |
| "candidates_considered": len(candidates), |
| "candidates_selected": len(kept), |
| "answers_committed": len( |
| {path.endpoint_id for path in kept if path.endpoint_id} |
| ), |
| "predicted_answer_count": predicted_count, |
| # ``max_count`` bounds what the count head can express, and real KG |
| # queries can have more true answers than that. At the cap the |
| # prediction is a floor, not an estimate, and count-guided selection |
| # truncates the commit. Surfacing it distinguishes "the model |
| # thinks there are 32 answers" from "it cannot count past 32". |
| "count_head_saturated": predicted_count is not None |
| and predicted_count >= self.config.max_count, |
| "reasoning_state_count": len(search.reasoning_states), |
| "pat_path_scores": [path.query_compatibility for path in supporting], |
| }, |
| relation_answers=relation_answers |
| if query.answer_mode == "relation" |
| else (), |
| boolean_answer=boolean_answer, |
| aggregate_value=aggregate_value, |
| ) |
|
|
| def _select_paths( |
| self, |
| scored: Sequence[GraphPath], |
| selection: SelectionOutput, |
| predicted: int | None, |
| ) -> tuple[GraphPath, ...]: |
| """Choose the committed subset of candidate paths. |
| |
| Two rules are available. A learned threshold is the obvious one, but it is |
| fitted on the training candidate distribution and does not transfer: it kept |
| roughly nine paths where the gold set held two, which is what capped exact |
| match and made a wider beam score worse than greedy. |
| |
| Count-guided selection instead asks the count head how many answers the |
| query has and takes that many highest-scoring paths. Cardinality then comes |
| from a head trained directly on ``len(answers)`` rather than from a decision |
| boundary that has to be calibrated against a distribution it never saw. |
| """ |
| threshold = float(selection.threshold.item()) |
| if not self.config.count_guided_selection or predicted is None: |
| return tuple(path for path in scored if path.score > threshold) |
| ranked = sorted(scored, key=lambda item: item.score, reverse=True) |
| # Paths are the evidence, answers are the endpoints: keep ranked paths |
| # until the requested number of distinct endpoints is covered, so several |
| # routes to one answer are not charged against the budget. A path with no |
| # endpoint names no answer, so it is evidence for nothing and is dropped |
| # rather than kept free of charge against that budget. |
| kept: list[GraphPath] = [] |
| endpoints: set[str] = set() |
| for path in ranked: |
| endpoint = path.endpoint_id |
| if endpoint is None: |
| continue |
| if endpoint not in endpoints: |
| if len(endpoints) >= predicted: |
| break |
| endpoints.add(endpoint) |
| kept.append(path) |
| return tuple(kept) |
|
|
| def _predicted_answer_count( |
| self, search: object, query_state: Tensor |
| ) -> int | None: |
| """How many answers the count head expects, clamped to a sane range.""" |
| pooled = getattr(search, "pooled_state", None) |
| path_states = getattr(search, "path_states", None) |
| if pooled is None: |
| return None |
| evidence = ( |
| path_states.mean(0) |
| if path_states is not None and path_states.numel() |
| else torch.zeros_like(pooled) |
| ) |
| verified = self.verification_stack( |
| pooled.unsqueeze(0), |
| evidence.unsqueeze(0), |
| query_state.unsqueeze(0), |
| self._evidence_signals( |
| getattr(getattr(search, "pat_output", None), "compatibility", None), |
| len(getattr(search, "paths", ()) or ()), |
| query_state.device, |
| ), |
| ) |
| logits = self.count_head(verified).reshape(-1) |
| return max(1, int(logits.argmax().item())) |
|
|
| def _evidence_signals( |
| self, compatibility: Tensor | None, count: int, device: torch.device |
| ) -> Tensor | None: |
| """Scalar summaries of how well *any* candidate matches the query. |
| |
| Computed the same way from training candidates and from beam candidates, which |
| is the whole point: the pooled evidence mean the trunk otherwise sees comes |
| from different distributions in the two regimes, and the support head learned |
| a rule that inverted at inference. PAT query-compatibility does not have that |
| problem -- it is the same function of a path and a query wherever it is called, |
| and it is already trained to discriminate (0.959). |
| |
| The maximum is the informative one: "is there any candidate compatible with |
| this query at all" is exactly the abstention question. |
| """ |
| if not self.config.verification_signals: |
| return None |
| if compatibility is None or compatibility.numel() == 0: |
| values = [ |
| torch.zeros((), device=device), |
| torch.zeros((), device=device), |
| ] |
| else: |
| flat = compatibility.reshape(-1).float() |
| values = [flat.max(), flat.mean()] |
| # Candidate count, compressed: "the beam found nothing" is itself evidence. |
| values.append(torch.log1p(torch.tensor(float(count), device=device)) / 10.0) |
| signals = torch.stack(values)[: self.config.verification_signals] |
| return signals.reshape(1, -1) |
|
|
| def _supports_any_answer(self, search: SearchOutput, query_state: Tensor) -> bool: |
| """Neural verification of whether the graph supports any answer. |
| |
| Consumes the pooled reasoning state, the pooled candidate evidence and |
| the query, exactly as the training-time support objective does, so the |
| head is used at inference in the same shape it was trained in. |
| """ |
| pooled = getattr(search, "pooled_state", None) |
| path_states = getattr(search, "path_states", None) |
| if pooled is None: |
| return True |
| evidence = ( |
| path_states.mean(0) |
| if path_states is not None and path_states.numel() |
| else torch.zeros_like(pooled) |
| ) |
| verified = self.verification_stack( |
| pooled.unsqueeze(0), |
| evidence.unsqueeze(0), |
| query_state.unsqueeze(0), |
| self._evidence_signals( |
| getattr(getattr(search, "pat_output", None), "compatibility", None), |
| len(getattr(search, "paths", ()) or ()), |
| query_state.device, |
| ), |
| ) |
| probability = torch.sigmoid(self.support_head(verified)).reshape(-1)[0] |
| return bool(probability.item() > self.config.abstention_threshold) |
|
|
| def _status( |
| self, |
| supporting: Sequence[GraphPath], |
| conflicts: Sequence[GraphPath], |
| graph: GraphStore, |
| query: GraphQuery, |
| verify_enabled: bool, |
| search: object, |
| encoded: EncodedGraph, |
| query_state: Tensor, |
| ) -> ReasoningStatus: |
| """Combine the neural status classifier with deterministic evidence. |
| |
| The classifier may propose ``CONTRADICTED``; it may never promote an |
| answer the verifier rejected, which is why the deterministic branch wins |
| whenever there is no verified support (PLAN section 20). |
| """ |
| if not supporting: |
| if query.target_entity_ids or query.answer_mode == "boolean": |
| return ReasoningStatus.UNSUPPORTED |
| return ReasoningStatus.INSUFFICIENT_EVIDENCE |
| if conflicts: |
| return ReasoningStatus.CONTRADICTED |
| pooled = getattr(search, "pooled_state", None) |
| if pooled is not None: |
| evidence = getattr(search, "path_states", None) |
| evidence_state = ( |
| evidence.mean(0) |
| if evidence is not None and evidence.numel() |
| else torch.zeros_like(pooled) |
| ) |
| verified_state = self.verification_stack( |
| pooled.unsqueeze(0), |
| evidence_state.unsqueeze(0), |
| query_state.unsqueeze(0), |
| self._evidence_signals( |
| getattr(getattr(search, "pat_output", None), "compatibility", None), |
| len(getattr(search, "paths", ()) or ()), |
| query_state.device, |
| ), |
| ) |
| logits = self.contradiction_head(verified_state) |
| predicted = INDEX_STATUS[int(logits.argmax(dim=-1).item())] |
| if predicted is ReasoningStatus.CONTRADICTED: |
| return ReasoningStatus.CONTRADICTED |
| if verify_enabled: |
| return classify_evidence( |
| candidate_paths=tuple(supporting), graph=graph, query=query |
| ) |
| return ReasoningStatus.SUPPORTED |
|
|
| def _mode_answers( |
| self, |
| query: GraphQuery, |
| status: ReasoningStatus, |
| answers_by_id: Mapping[str, float], |
| supporting: Sequence[GraphPath], |
| ) -> tuple[bool | None, int | float | None]: |
| boolean_answer: bool | None = None |
| aggregate_value: int | float | None = None |
| if query.answer_mode == "boolean": |
| if status is ReasoningStatus.SUPPORTED: |
| boolean_answer = not query.target_entity_ids or bool( |
| set(query.target_entity_ids) & set(answers_by_id) |
| ) |
| elif status in {ReasoningStatus.UNSUPPORTED, ReasoningStatus.CONTRADICTED}: |
| boolean_answer = False |
| elif query.answer_mode in {"count", "aggregate"}: |
| aggregate_value = len(answers_by_id) |
| return boolean_answer, aggregate_value |
|
|
| def _confidence( |
| self, supporting: Sequence[GraphPath], status: ReasoningStatus |
| ) -> float: |
| if not supporting: |
| return float(torch.sigmoid(torch.tensor(-4.0)).item()) |
| best = max(path.score for path in supporting) |
| return float(torch.sigmoid(torch.tensor(best)).item()) |
|
|
| # -- accounting --------------------------------------------------------- |
|
|
| def parameter_buckets(self) -> Mapping[str, int]: |
| """Partition parameters into the five PLAN section 6 budget categories. |
| |
| The buckets are disjoint and sum exactly to the model total. Ownership |
| follows the section 6 wording: graph embeddings belong to the graph |
| encoder and PAT's own embeddings belong to PAT, while the separate |
| ``embeddings`` budget covers the query vocabulary and positional tables. |
| The previous implementation counted query embeddings inside |
| ``reasoning_transformer`` *and* again under ``embeddings``, so the |
| reported buckets summed to 108% of the model. |
| """ |
| buckets = dict.fromkeys( |
| ("graph_encoder", "pat", "reasoning_transformer", "heads", "embeddings"), 0 |
| ) |
| query_embedding_prefixes = ( |
| "query_encoder.token_embedding", |
| "query_encoder.position_embedding", |
| "query_encoder.structure_embedding", |
| ) |
| for name, parameter in self.named_parameters(): |
| count = parameter.numel() |
| if name.startswith(query_embedding_prefixes): |
| buckets["embeddings"] += count |
| elif name.startswith("graph_encoder"): |
| buckets["graph_encoder"] += count |
| elif name.startswith("pat"): |
| buckets["pat"] += count |
| elif name.startswith(("query_encoder", "reasoning_core")): |
| buckets["reasoning_transformer"] += count |
| else: |
| buckets["heads"] += count |
| total = sum(parameter.numel() for parameter in self.parameters()) |
| if sum(buckets.values()) != total: |
| raise RuntimeError( |
| f"parameter buckets sum to {sum(buckets.values()):,} but the model has {total:,}" |
| ) |
| return buckets |
|
|
|
|
| def _path_key(path: GraphPath) -> tuple[tuple[str, str, str], ...]: |
| return tuple( |
| (step.source_id, step.relation_id, step.target_id) for step in path.steps |
| ) |
|
|
|
|
| def _compat(pat_output: PATOutput | None, index: int) -> float | None: |
| if pat_output is None: |
| return None |
| values = pat_output.compatibility |
| if index >= values.shape[0]: |
| return None |
| return float(values[index].item()) |
|
|
|
|
| def _slice_pat(output: PATOutput | None, rows: slice) -> PATOutput | None: |
| """One example's paths out of a batched PAT pass.""" |
| if output is None: |
| return None |
| return PATOutput( |
| *(getattr(output, field)[rows] for field in PATOutput.__dataclass_fields__) |
| ) |
|
|
|
|
| def _query_starts(query: GraphQuery) -> tuple[str, ...]: |
| """The node ids a query starts from, as the search resolves them.""" |
| return tuple(query.start_entity_ids or query.start_entities) |
|
|
|
|
| def _start_mask( |
| node_ids: Sequence[str], start_ids: Sequence[str], device: torch.device |
| ) -> Tensor: |
| """Boolean mask over ``node_ids`` flagging the query's start entities.""" |
| starts = set(start_ids) |
| return torch.tensor( |
| [node_id in starts for node_id in node_ids], dtype=torch.bool, device=device |
| ) |
|
|
