"""Configuration objects for reproducible GRM model construction.""" from __future__ import annotations from dataclasses import asdict, dataclass, field from typing import Any, Literal @dataclass(frozen=True) 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)") @dataclass(frozen=True) 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 @dataclass(frozen=True) 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") @dataclass(frozen=True) 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)