| """Configuration objects for reproducible GRM model construction.""" | |
| from __future__ import annotations | |
| from dataclasses import asdict, dataclass, field | |
| from typing import Any, Literal | |
| class PATConfig: | |
| enabled: bool = True | |
| hidden_dim: int = 128 | |
| num_layers: int = 2 | |
| num_heads: int = 4 | |
| max_path_length: int = 8 | |
| num_token_types: int = 10 | |
| dropout: float = 0.1 | |
| relation_aware_attention: bool = True | |
| query_conditioning: bool = True | |
| type_vocab_size: int = 1024 | |
| attribute_vocab_size: int = 2048 | |
| def __post_init__(self) -> None: | |
| if self.hidden_dim <= 0 or self.num_layers <= 0 or self.num_heads <= 0: | |
| raise ValueError("PAT dimensions and layer counts must be positive") | |
| if ( | |
| self.max_path_length <= 0 | |
| or self.num_token_types <= 0 | |
| or self.type_vocab_size <= 0 | |
| or self.attribute_vocab_size <= 0 | |
| ): | |
| raise ValueError("PAT path length and token vocabulary must be positive") | |
| if not 0.0 <= self.dropout < 1.0: | |
| raise ValueError("PAT dropout must be in [0, 1)") | |
| class LossConfig: | |
| graph: float = 0.1 | |
| link: float = 0.2 | |
| next_hop: float = 1.0 | |
| relation: float = 0.5 | |
| path: float = 1.0 | |
| pat: float = 1.0 | |
| answer: float = 1.0 | |
| verification: float = 0.5 | |
| contradiction: float = 0.5 | |
| calibration: float = 0.1 | |
| count: float = 0.3 | |
| #: Set-level selection objective. Weighted above the per-path term because | |
| #: committing to the right number of paths is what exact match measures. | |
| selection_set: float = 1.5 | |
| #: Listwise ranking objective over the candidate set. ``count_guided_selection`` | |
| #: commits the top-k candidates by score, so the ranking -- not the threshold -- | |
| #: is what inference consumes; optimising only the thresholded objective left a | |
| #: train/inference mismatch of the same kind the on-policy experiment targeted. | |
| selection_rank: float = 1.0 | |
| #: Weight on the *unsupported* class in the verification objective. About 85% of | |
| #: curriculum examples are answerable, and with a per-example batch the unweighted | |
| #: objective is minimised well by always predicting "supported": the shipped model | |
| #: abstained on 25% of answerable examples and 12.5% of unanswerable ones, worse | |
| #: than chance in the direction that matters. Set from the measured base rate. | |
| unsupported_weight: float = 5.0 | |
| #: Penalty on the mean pairwise cosine of a graph's final node states. Long | |
| #: runs collapsed every node to one direction (mean cosine 0.03 -> 0.98 over | |
| #: ~100K steps, artifacts/retrain-20260928-full-epoch/freeze.json), which | |
| #: leaves answer ranking and unseen-ontology transfer nothing to work with. | |
| #: 0 disables it and keeps older configs' objective unchanged. | |
| node_spread: float = 0.0 | |
| class TrainingConfig: | |
| precision: Literal["fp32", "fp16", "bf16"] = "fp32" | |
| link_head: Literal["cosine", "mlp"] = "cosine" | |
| gradient_checkpointing: bool = False | |
| batch_size: int = 8 | |
| gradient_accumulation: int = 1 | |
| learning_rate: float = 3e-4 | |
| warmup_ratio: float = 0.03 | |
| weight_decay: float = 0.01 | |
| max_steps: int = 1000 | |
| seed: int = 7 | |
| distributed_backend: str = "nccl" | |
| checkpoint_every: int = 0 | |
| evaluation_every: int = 0 | |
| use_fsdp: bool = False | |
| deterministic: bool = True | |
| grad_clip: float = 1.0 | |
| #: Examples backpropagated together in one autograd invocation. Larger is | |
| #: faster but holds more activations; 1 reproduces per-example backward. | |
| backward_chunk_size: int = 8 | |
| #: When > 0, each training example gets a freshly renamed ontology with this | |
| #: probability. This forces the model to learn relational semantics from | |
| #: type signatures and structure rather than memorizing relation names. | |
| #: 0=no renaming, 1=always rename. Default 0.0 gives the original behavior. | |
| rename_probability: float = 0.0 | |
| #: Run each backward chunk through ``GraphModel.forward_batch`` (one pass of | |
| #: every network over the chunk) instead of one forward per example. The | |
| #: objective is the same; per-example forwards left the GPU ~85% idle. | |
| batched_forward: bool = True | |
| #: Learning-rate multiplier for the query encoder alone. Its first | |
| #: self-attention layer was the source of the gradient blow-up that stopped | |
| #: the v3 full-epoch run at step 183,000 while the rest of the model was | |
| #: stable. 1.0 keeps the optimizer exactly as before. | |
| query_encoder_lr_scale: float = 1.0 | |
| #: Stop training the query encoder. At 0.1x learning rate its gradients | |
| #: still grew (median 88 -> 234 over 15K steps after step 183,500), always | |
| #: concentrated in its input projection and first attention layer. | |
| freeze_query_encoder: bool = False | |
| def __post_init__(self) -> None: | |
| # Literal types are not enforced at runtime; a typo such as "BF16" or | |
| # "gin" would otherwise be accepted and silently mapped to a default. | |
| if self.precision not in ("fp32", "fp16", "bf16"): | |
| raise ValueError( | |
| f"precision must be fp32, fp16 or bf16: {self.precision!r}" | |
| ) | |
| if self.distributed_backend not in ("nccl", "gloo"): | |
| raise ValueError( | |
| f"distributed_backend must be nccl or gloo: {self.distributed_backend!r}" | |
| ) | |
| if self.use_fsdp: | |
| # Refused here, not only when distributed: a single-process run would | |
| # otherwise train without FSDP and say nothing. | |
| raise ValueError("use_fsdp is not implemented; set it to false") | |
| if ( | |
| self.batch_size <= 0 | |
| or self.gradient_accumulation <= 0 | |
| or self.max_steps <= 0 | |
| ): | |
| raise ValueError( | |
| "batch_size, gradient_accumulation, and max_steps must be positive" | |
| ) | |
| if ( | |
| self.learning_rate <= 0 | |
| or self.weight_decay < 0 | |
| or not 0.0 <= self.warmup_ratio <= 1.0 | |
| ): | |
| raise ValueError("invalid optimizer or warmup configuration") | |
| if self.checkpoint_every < 0 or self.evaluation_every < 0: | |
| raise ValueError("checkpoint/evaluation intervals cannot be negative") | |
| if self.grad_clip <= 0: | |
| raise ValueError("grad_clip must be positive") | |
| if self.backward_chunk_size <= 0: | |
| raise ValueError("backward_chunk_size must be positive") | |
| class GraphConfig: | |
| model_version: str = "graph-smoke-v1" | |
| target_parameters: int = 100_000_000 | |
| hidden_dim: int = 128 | |
| graph_input_dim: int = 8 | |
| edge_input_dim: int = 8 | |
| num_node_types: int = 64 | |
| num_relations: int = 128 | |
| num_relation_types: int = 128 | |
| vocab_size: int = 16_384 | |
| max_query_tokens: int = 128 | |
| graph_layers: int = 3 | |
| graph_encoder: Literal["sage", "gat"] = "sage" | |
| use_gat: bool = False | |
| #: Identity skip connection around each message-passing layer. Without it | |
| #: eight layers of mean aggregation oversmooth: node states converged to a | |
| #: shared direction on long runs. Off by default so existing checkpoints | |
| #: keep computing what they were trained to compute. | |
| graph_residual: bool = False | |
| #: Mark the query's start entities in the graph encoder's input: a learned | |
| #: start/other embedding plus the query state on start nodes, so node states | |
| #: are computed relative to where the question begins. Synthetic entity | |
| #: names are random identifiers, so without a marker nodes that play the same | |
| #: structural role are indistinguishable and long runs collapsed them together. | |
| start_marking: bool = False | |
| #: Soft cap on attention scores in PAT and the reasoning core's graph | |
| #: cross-attention (``graph.utils.soft_cap``); 0 disables it. | |
| attention_logit_cap: float = 0.0 | |
| #: Normalise attention queries and keys in PAT and the graph | |
| #: cross-attention (``graph.utils.qk_normalize``). | |
| attention_qk_norm: bool = False | |
| num_reasoning_layers: int = 2 | |
| num_reasoning_heads: int = 4 | |
| #: FFN expansion inside the reasoning core. Exposed because it is the main | |
| #: lever for hitting the section 6 reasoning-transformer budget without | |
| #: changing depth or width. | |
| reasoning_ffn_multiplier: int = 4 | |
| dropout: float = 0.1 | |
| max_hops: int = 8 | |
| beam_width: int = 8 | |
| candidate_width: int = 32 | |
| relation_top_k: int = 16 | |
| answer_top_k: int = 8 | |
| #: Largest count the aggregate head can emit. | |
| max_count: int = 32 | |
| #: Cap on supervised positive/negative paths per example, bounding the PAT | |
| #: and selection batch built for one training step. | |
| max_supervised_paths: int = 8 | |
| #: How many reachable paths to scan when topping up selection negatives from | |
| #: the distribution the beam actually produces. | |
| reachable_negative_pool: int = 256 | |
| #: Negatives per positive in the selection batch. This sets the *base rate* | |
| #: the threshold is fitted against, and it has to match inference: the beam | |
| #: surfaces roughly one gold path in six, so training on a balanced set | |
| #: produced a threshold that kept ten candidates when the gold set held two. | |
| selection_negative_ratio: int = 6 | |
| #: Train path selection on the candidates the beam actually produces. | |
| #: | |
| #: Measured WORSE and therefore off by default. At 200 steps on a 128-dim | |
| #: probe it lost on every metric (F1 0.254 vs 0.310, selection precision 0.149 | |
| #: vs 0.195, renamed-ontology PAT 0.429 vs 0.741). The cause looks like cold | |
| #: start: early in training the beam is near-random, so its candidate set holds | |
| #: almost no gold path and the objective sees a degenerate, low-diversity | |
| #: signal. Kept because it is the right idea once a competent policy exists to | |
| #: bootstrap from -- worth revisiting as a late-training phase, not from step 0. | |
| on_policy_selection: bool = False | |
| #: Size the selected set from the count head's prediction instead of a | |
| #: threshold. The threshold is uncalibrated against the inference candidate | |
| #: distribution, so it kept ~9 paths where the gold set held ~2. | |
| count_guided_selection: bool = True | |
| #: Semantic text features for entity, relation and query strings. ``None`` keeps | |
| #: the original behaviour: every string becomes a one-hot at a hashed index, so two | |
| #: relations with the same meaning and different names get *orthogonal* features and | |
| #: an unseen relation string carries no meaning at all. That is the measured cause | |
| #: of the model's transfer failure -- unseen-ontology exact match 0.150 with PAT at | |
| #: 0.557 against a 0.5 floor, and losing to a released 34M model on real typed KGs | |
| #: whose one advantage is reading strings through a pretrained encoder. | |
| #: Either "hashed-bag" (dependency-free, shares substrings) or | |
| #: "sentence-transformer" (frozen pretrained, shares meaning). | |
| #: Scalar evidence signals fed to the verification trunk. The abstention decision | |
| #: was measured worse than chance -- abstaining on 25% of answerable examples and | |
| #: 12.5% of unanswerable ones -- because the pooled evidence the support head reads | |
| #: is built from gold positives during training and from the beam's own candidates | |
| #: at inference. These signals (max and mean PAT query-compatibility, and the | |
| #: candidate count) are computed identically in both regimes, and PAT compatibility | |
| #: is already a well-trained discriminator (0.959). Set to 3 to enable; it adds | |
| #: weights, so a checkpoint trained without it cannot be loaded with it. | |
| verification_signals: int = 0 | |
| #: Soft bound on the verification and contradiction logits: cap * tanh(logit / cap). | |
| #: 0 disables it. Unbounded logits let one confidently wrong example produce a | |
| #: binary cross-entropy of 40 to 116, which destabilizes training late in the run. | |
| verification_logit_cap: float = 0.0 | |
| #: Subtract a running mean of the verifier state before the heads read it. The | |
| #: state shares a large common direction across examples (mean pairwise cosine | |
| #: 0.69 at the collapse), which leaves the support head with no per-example signal. | |
| #: Uses running statistics at inference, so it has no train/test leak. | |
| verification_center: bool = False | |
| #: Report the trained confidence head's probability as the result confidence | |
| #: instead of sigmoid(best path score). The head is trained with a Brier loss | |
| #: against answer correctness; path scores are trained for ranking and are not | |
| #: probabilities, which is why path-score confidence calibrates poorly. | |
| confidence_from_head: bool = False | |
| #: Support probability below which the model abstains rather than answering. | |
| #: | |
| #: The evidence signals above make abstention work in-distribution, and they make it | |
| #: *distribution-sensitive*: they are built from PAT query-compatibility, which | |
| #: saturates low on relation vocabulary the model has never seen. On real typed KGs | |
| #: the gate therefore fired on 100% of queries -- every real-KG metric went to | |
| #: 0.000, even when the gold relation path was supplied. A caller working outside | |
| #: the training distribution should lower this, or pass ``abstain=False`` to | |
| #: ``reason`` and rely on the deterministic verifier alone. | |
| abstention_threshold: float = 0.5 | |
| text_features: str | None = None | |
| text_feature_dim: int = 384 | |
| text_feature_model: str | None = None | |
| #: Depth and width of the shared neural verification trunk (PLAN section 20). | |
| verifier_layers: int = 2 | |
| verifier_width_multiplier: int = 4 | |
| pat: PATConfig = field(default_factory=PATConfig) | |
| loss: LossConfig = field(default_factory=LossConfig) | |
| training: TrainingConfig = field(default_factory=TrainingConfig) | |
| def __post_init__(self) -> None: | |
| if self.graph_encoder not in ("sage", "gat"): | |
| raise ValueError( | |
| f"graph_encoder must be sage or gat: {self.graph_encoder!r}" | |
| ) | |
| if ( | |
| self.hidden_dim <= 0 | |
| or self.graph_layers <= 0 | |
| or self.graph_input_dim <= 0 | |
| or self.edge_input_dim <= 0 | |
| ): | |
| raise ValueError("model dimensions and graph layer count must be positive") | |
| if ( | |
| self.num_node_types <= 0 | |
| or self.num_relations <= 0 | |
| or self.num_relation_types <= 0 | |
| ): | |
| raise ValueError("node and relation vocabularies must be positive") | |
| if self.vocab_size <= 0 or self.max_query_tokens <= 0: | |
| raise ValueError("query vocabulary and token limit must be positive") | |
| if self.pat.hidden_dim != self.hidden_dim: | |
| raise ValueError( | |
| "PAT hidden_dim must match GRM hidden_dim for residual fusion" | |
| ) | |
| if self.pat.num_heads <= 0 or self.hidden_dim % self.pat.num_heads: | |
| raise ValueError("hidden_dim must be divisible by PAT num_heads") | |
| if self.num_reasoning_heads <= 0 or self.hidden_dim % self.num_reasoning_heads: | |
| raise ValueError("hidden_dim must be divisible by reasoning heads") | |
| if self.max_hops <= 0 or self.beam_width <= 0: | |
| raise ValueError("max_hops and beam_width must be positive") | |
| if ( | |
| self.candidate_width <= 0 | |
| or self.relation_top_k <= 0 | |
| or self.answer_top_k <= 0 | |
| ): | |
| raise ValueError("reasoning widths must be positive") | |
| if self.reasoning_ffn_multiplier <= 0: | |
| raise ValueError("reasoning_ffn_multiplier must be positive") | |
| if self.max_count <= 0 or self.max_supervised_paths <= 0: | |
| raise ValueError("max_count and max_supervised_paths must be positive") | |
| if self.reachable_negative_pool <= 0: | |
| raise ValueError("reachable_negative_pool must be positive") | |
| if self.selection_negative_ratio <= 0: | |
| raise ValueError("selection_negative_ratio must be positive") | |
| if self.verifier_layers <= 0 or self.verifier_width_multiplier <= 0: | |
| raise ValueError("verifier depth and width multiplier must be positive") | |
| if not 0.0 <= self.dropout < 1.0: | |
| raise ValueError("model dropout must be in [0, 1)") | |
| def to_dict(self) -> dict[str, Any]: | |
| return asdict(self) | |