"""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 )