fix(mlx): enforce graph-routed indexer loss contract

Author: apstenku123

cppmega_mlx/models/dense_cpp_lm.py

@@ -49,7 +49,10 @@
     apply_rotary_emb,
     rotary_inv_freq,
 )
-from cppmega_mlx.nn.code_graph_routes import build_token_graph_biases
+from cppmega_mlx.nn.code_graph_routes import (
+    build_token_graph_biases,
+    remove_graph_route_prior,
+)
 from cppmega_mlx.nn.domain_embedding import CppMegaDomainEmbedding
 from cppmega_mlx.nn.moe import FeedForwardExpert
 from cppmega_mlx.nn.ngram_hash import NgramHashEmbedding
@@ -274,8 +277,8 @@ class GraphIndexedAttention(nn.Module):
 
     Reuses the same q/k/v projections as :class:`CausalSelfAttention` but, instead
     of dense SDPA, scores each (query, key) pair with a cheap lightning indexer
-    (per-head ReLU dot + learned graph-prior bias ``beta * S_blk``), selects the
-    top-k (+ local window + sinks), and runs dense MLA over the gathered KV via
+    (per-head ReLU dot + fixed graph-prior coefficient ``beta * S_blk``), selects
+    the top-k (+ local window + sinks), and runs dense MLA over the gathered KV via
     :func:`graph_indexed_attention_reference`. The indexer scores are exposed for
     the indexer losses (KL warm-up / BCE / coverage / contrastive) as an
     explicit return value. Forward execution never writes activation tensors
@@ -299,12 +302,14 @@ def __init__(self, config: DenseCppLMConfig):
         self.index_q_proj = nn.Linear(d_model, hi * di, bias=False)
         self.index_k_proj = nn.Linear(d_model, hi * di, bias=False)
         self.index_head_weights = mx.ones((hi,), dtype=mx.float32)
-        # Keep the checkpoint-visible scalar while matching the production
-        # CPPMEGA_DSA_GRAPH_BIAS_BETA default (1.0). Beta=0 remains an explicit
-        # config/checkpoint ablation.
+        # Keep the checkpoint-visible route coefficient while matching the
+        # production CPPMEGA_DSA_GRAPH_BIAS_BETA default (1.0). It is fixed:
+        # graph supervision removes this prior and discrete top-k cannot train
+        # it, so leaving it in AdamW would only move it through weight decay.
         self.index_beta = mx.full(
             (1,), config.graph_attention_bias_beta, dtype=mx.float32
         )
+        self.freeze(recurse=False, keys="index_beta", strict=True)
         self.rope_inv_freq = rotary_inv_freq(acfg) if acfg.use_rope else None
 
     def _rotary_tables(self, seq_len: int) -> tuple[mx.array, mx.array]:
@@ -384,16 +389,11 @@ def graph_supervision_scores(
     ) -> mx.array:
         """Remove the fixed route prior while preserving invalid score entries."""
 
-        if tuple(final_scores.shape) != tuple(route_bias.shape):
-            raise ValueError(
-                "graph supervision route bias must match final indexer scores: "
-                f"{tuple(route_bias.shape)} != {tuple(final_scores.shape)}"
-            )
-        scores = final_scores.astype(mx.float32)
-        bias = route_bias.astype(mx.float32)
-        valid = scores > mx.array(-5.0e8, dtype=mx.float32)
-        neural_scores = scores - self.index_beta[0].astype(mx.float32) * bias
-        return mx.where(valid, neural_scores, scores)
+        return remove_graph_route_prior(
+            final_scores,
+            route_bias,
+            beta=self.index_beta[0],
+        )
 
 
 class DenseCppBlock(nn.Module):
@@ -834,11 +834,18 @@ def graph_supervision_scores(
         indexer_scores: tuple[mx.array, ...],
         *,
         input_ids: mx.array,
-        document_ids: mx.array,
-        block_bias: mx.array,
-        edge_kind_bias: mx.array,
+        document_ids: mx.array | None = None,
+        graph_batch: GraphBatch | None = None,
+        block_bias: mx.array | None = None,
+        edge_kind_bias: mx.array | None = None,
     ) -> tuple[mx.array, ...]:
-        """Return learned DSA logits with the fixed graph route prior removed."""
+        """Return neural DSA logits with the fixed graph route prior removed.
+
+        The route source must match the forward pass: provide a typed
+        ``graph_batch`` or the precomputed relation and edge-kind biases, never
+        both. This keeps graph-supervised losses from training on ``beta*S_graph``
+        a second time.
+        """
 
         if self.config.attention_mode != "dsa":
             raise ValueError("graph supervision scores require attention_mode='dsa'")
@@ -849,7 +856,7 @@ def graph_supervision_scores(
             )
         route_bias = self._resolve_graph_bias(
             input_ids,
-            graph_batch=None,
+            graph_batch=graph_batch,
             block_bias=block_bias,
             edge_kind_bias=edge_kind_bias,
             document_ids=document_ids,

cppmega_mlx/nn/code_graph_routes.py

@@ -71,6 +71,97 @@
 }
 _NORMALIZE_MODES = ("binary", "count")
 _KNOWN_EDGE_KIND_IDS = frozenset(int(kind) for kind in DomainEdgeKind)
+_MASKED_SCORE_THRESHOLD = -5.0e8
+
+
+def _broadcast_route_bias(
+    scores: mx.array,
+    graph_bias: mx.array,
+    *,
+    where: str,
+) -> mx.array:
+    if scores.ndim != 3:
+        raise ValueError(
+            f"{where}: scores must be shaped (B,Tq,S), got {tuple(scores.shape)}"
+        )
+    batch, queries, keys = (int(value) for value in scores.shape)
+    if graph_bias.ndim == 2:
+        if tuple(graph_bias.shape) != (queries, keys):
+            raise ValueError(
+                f"{where}: 2-D graph_bias must be ({queries},{keys}), got "
+                f"{tuple(graph_bias.shape)}"
+            )
+        return mx.broadcast_to(graph_bias[None, :, :], (batch, queries, keys))
+    if graph_bias.ndim != 3:
+        raise ValueError(
+            f"{where}: graph_bias must be 2-D or 3-D, got "
+            f"{tuple(graph_bias.shape)}"
+        )
+    bias_batch, bias_queries, bias_keys = (
+        int(value) for value in graph_bias.shape
+    )
+    if (bias_queries, bias_keys) != (queries, keys):
+        raise ValueError(
+            f"{where}: graph_bias trailing shape must be ({queries},{keys}), got "
+            f"{tuple(graph_bias.shape)}"
+        )
+    if bias_batch == batch:
+        return graph_bias
+    if bias_batch == 1:
+        return mx.broadcast_to(graph_bias, (batch, queries, keys))
+    raise ValueError(
+        f"{where}: graph_bias batch must be 1 or {batch}, got {bias_batch}"
+    )
+
+
+def _route_beta(beta: mx.array | float, *, where: str) -> mx.array:
+    beta_array = (
+        beta
+        if isinstance(beta, mx.array)
+        else mx.array(float(beta), dtype=mx.float32)
+    )
+    if beta_array.ndim != 0:
+        raise ValueError(
+            f"{where}: beta must be a scalar, got {tuple(beta_array.shape)}"
+        )
+    return beta_array.astype(mx.float32)
+
+
+def apply_graph_route_prior(
+    neural_scores: mx.array,
+    graph_bias: mx.array,
+    *,
+    beta: mx.array | float = 1.0,
+) -> mx.array:
+    """Return ``I_neural + beta * S_graph`` without reviving masked pairs."""
+
+    where = "apply_graph_route_prior"
+    bias = _broadcast_route_bias(neural_scores, graph_bias, where=where)
+    scores = neural_scores.astype(mx.float32)
+    valid = scores > mx.array(_MASKED_SCORE_THRESHOLD, dtype=mx.float32)
+    final_scores = scores + _route_beta(beta, where=where) * bias.astype(mx.float32)
+    return mx.where(valid, final_scores, scores)
+
+
+def remove_graph_route_prior(
+    final_scores: mx.array,
+    graph_bias: mx.array,
+    *,
+    beta: mx.array | float = 1.0,
+) -> mx.array:
+    """Recover neural indexer scores from ``I_neural + beta * S_graph``.
+
+    Graph-supervised losses must consume this neural component. Training on the
+    final routed score would reward the fixed graph prior itself and would apply
+    ``beta * S_graph`` twice when the loss path also adds the prior.
+    """
+
+    where = "remove_graph_route_prior"
+    bias = _broadcast_route_bias(final_scores, graph_bias, where=where)
+    scores = final_scores.astype(mx.float32)
+    valid = scores > mx.array(_MASKED_SCORE_THRESHOLD, dtype=mx.float32)
+    neural_scores = scores - _route_beta(beta, where=where) * bias.astype(mx.float32)
+    return mx.where(valid, neural_scores, scores)
 
 
 def _resolve_num_nodes(packet: GraphPacket, relation: str) -> int:
@@ -733,8 +824,10 @@ def _validate_graph_document_boundaries(
 __all__ = [
     "GraphRouteConfig",
     "CodeGraphRouter",
+    "apply_graph_route_prior",
     "build_attention_bias",
     "build_block_candidates",
     "build_token_attention_bias",
     "build_token_graph_biases",
+    "remove_graph_route_prior",
 ]

cppmega_mlx/training/indexer_losses.py

@@ -32,6 +32,10 @@
 import numpy as np
 
 from cppmega_mlx.data.batch import batch_values_are_prevalidated
+from cppmega_mlx.nn.code_graph_routes import (
+    apply_graph_route_prior,
+    remove_graph_route_prior,
+)
 
 _EPS = 1e-9
 _NEG_INF = -1e9
@@ -275,52 +279,20 @@ def apply_graph_indexer_bias(
         beta: scalar learnable/fixed graph-prior weight.
     """
 
-    B, Tq, Sblk = _check_block_scores(
-        indexer_scores, where="apply_graph_indexer_bias"
-    )
-    if graph_bias.ndim == 2:
-        if tuple(graph_bias.shape) != (Tq, Sblk):
-            raise ValueError(
-                "apply_graph_indexer_bias: 2-D graph_bias must be "
-                f"({Tq},{Sblk}), got {tuple(graph_bias.shape)}"
-            )
-        bias = mx.broadcast_to(graph_bias[None, :, :], (B, Tq, Sblk))
-    elif graph_bias.ndim == 3:
-        gB, gTq, gSblk = (int(x) for x in graph_bias.shape)
-        if (gTq, gSblk) != (Tq, Sblk):
-            raise ValueError(
-                "apply_graph_indexer_bias: graph_bias trailing shape must be "
-                f"({Tq},{Sblk}), got {tuple(graph_bias.shape)}"
-            )
-        if gB == B:
-            bias = graph_bias
-        elif gB == 1:
-            bias = mx.broadcast_to(graph_bias, (B, Tq, Sblk))
-        else:
-            raise ValueError(
-                "apply_graph_indexer_bias: graph_bias batch must be 1 or "
-                f"{B}, got {gB}"
-            )
-    else:
-        raise ValueError(
-            "apply_graph_indexer_bias: graph_bias must be 2-D or 3-D, got "
-            f"{tuple(graph_bias.shape)}"
-        )
+    _check_block_scores(indexer_scores, where="apply_graph_indexer_bias")
+    return apply_graph_route_prior(indexer_scores, graph_bias, beta=beta)
 
-    beta_arr = (
-        beta
-        if isinstance(beta, mx.array)
-        else mx.array(float(beta), dtype=mx.float32)
-    )
-    if beta_arr.ndim != 0:
-        raise ValueError(
-            "apply_graph_indexer_bias: beta must be a scalar, got "
-            f"{tuple(beta_arr.shape)}"
-        )
-    logits = indexer_scores.astype(mx.float32)
-    valid = logits > (_NEG_INF / 2.0)
-    biased = logits + beta_arr.astype(mx.float32) * bias.astype(mx.float32)
-    return mx.where(valid, biased, mx.array(_NEG_INF, dtype=mx.float32))
+
+def remove_graph_indexer_bias(
+    final_scores: mx.array,
+    graph_bias: mx.array,
+    *,
+    beta: mx.array | float = 1.0,
+) -> mx.array:
+    """Return neural indexer logits by removing the fixed graph route prior."""
+
+    _check_block_scores(final_scores, where="remove_graph_indexer_bias")
+    return remove_graph_route_prior(final_scores, graph_bias, beta=beta)
 
 
 def select_graph_biased_topk(
@@ -531,6 +503,7 @@ def edge_targets_from_candidates(
     "indexer_edge_bce_loss",
     "indexer_coverage_hinge_loss",
     "apply_graph_indexer_bias",
+    "remove_graph_indexer_bias",
     "select_graph_biased_topk",
     "recall_at_k",
     "dense_attn_topk_overlap",

docs/graph_supervised_dsa_indexer.md

@@ -31,25 +31,37 @@ top-k selection, not as an after-the-fact annotation.
 
 ## Training Objective
 
-The indexer is trained with both neural and compiler-grade teachers:
+The routed score is used for inference-time top-k only. Graph supervision must
+train the neural score after removing the fixed route prior:
+
+```text
+I_neural = I_final - beta * S_graph
+```
+
+The indexer objective is therefore:
 
 ```text
 L_indexer =
-    KL(dense_attention_blocks || softmax(I_final))
-  + lambda_bce * BCE(I_final, graph_edge_targets)
-  + lambda_cov * coverage_hinge(I_final, graph_edge_targets, topk)
+    KL(dense_attention_blocks || softmax(I_neural))
+  + lambda_bce * BCE(I_neural, graph_edge_targets)
+  + lambda_cov * coverage_hinge(I_neural, graph_edge_targets, topk)
   + lambda_ctr * block_contrastive(...)
 ```
 
 `graph_edge_targets` is built from the same graph candidates that produce
-`S_graph`.  This keeps the training target aligned with the inference-time
-selection rule.
+`S_graph`, but the fixed prior itself is not rewarded by the auxiliary loss.
+This prevents the shortcut where increasing `beta` lowers graph loss without
+improving the neural indexer. `DenseCppLM` keeps `beta` checkpoint-visible but
+frozen so optimizer weight decay cannot move a coefficient that receives no
+useful differentiable training signal.
 
 ## Implemented MLX Surfaces
 
 - `cppmega_mlx.nn.code_graph_routes`
   - `GraphRouteConfig`
   - `CodeGraphRouter`
+  - `apply_graph_route_prior`
+  - `remove_graph_route_prior`
   - `build_attention_bias`
   - `build_block_candidates`
 - `cppmega_mlx.nn.sparse_mla`
@@ -58,6 +70,7 @@ selection rule.
   - `graph_indexed_attention_reference`
 - `cppmega_mlx.training.indexer_losses`
   - `apply_graph_indexer_bias`
+  - `remove_graph_indexer_bias`
   - `select_graph_biased_topk`
   - `total_indexer_loss`
   - `recall_at_k`