Route side-channel training metadata

Author: apstenku123

cppmega_mlx/data/batch.py

@@ -8,6 +8,14 @@
 import mlx.core as mx
 import numpy as np
 
+SideChannelDropoutPolicy = Mapping[str, float]
+
+_SIDE_CHANNEL_FAMILY_FIELDS: Mapping[str, tuple[str, ...]] = {
+    "platform": ("platform_ids",),
+    "syntax": ("ast_depth_ids", "sibling_index_ids", "node_type_ids"),
+    "structure": ("structure_ids", "dep_levels"),
+}
+
 
 @dataclass(frozen=True)
 class LMTokenBatch:
@@ -24,6 +32,7 @@ class LMTokenBatch:
     sibling_index_ids: mx.array | None = None
     node_type_ids: mx.array | None = None
     platform_ids: mx.array | None = None
+    side_channels: Mapping[str, Mapping[str, mx.array]] | None = None
     metadata: Mapping[str, Any] | None = None
 
     def __post_init__(self) -> None:
@@ -105,6 +114,8 @@ def __post_init__(self) -> None:
             mx.eval(has_negative)
             if bool(has_negative.item()):
                 raise ValueError("platform_ids must be non-negative")
+        if self.side_channels is not None:
+            _validate_side_channel_map(self.side_channels)
         if self.metadata is not None and not isinstance(self.metadata, Mapping):
             raise ValueError("metadata must be a mapping when provided")
 
@@ -163,6 +174,62 @@ def model_kwargs(self) -> dict[str, mx.array]:
             )
         return kwargs
 
+    def side_channel_map(self) -> dict[str, dict[str, mx.array]]:
+        """Return side-channel tensors grouped by family and column name."""
+
+        out: dict[str, dict[str, mx.array]] = {
+            family: {}
+            for family in _SIDE_CHANNEL_FAMILY_FIELDS
+        }
+        for family, fields in _SIDE_CHANNEL_FAMILY_FIELDS.items():
+            for field_name in fields:
+                value = getattr(self, field_name)
+                if value is not None:
+                    out[family][field_name] = value
+        if self.side_channels is not None:
+            for family, columns in self.side_channels.items():
+                out.setdefault(family, {}).update(dict(columns))
+        return {family: columns for family, columns in out.items() if columns}
+
+    def with_side_channel_dropout(
+        self,
+        policy: SideChannelDropoutPolicy | None,
+        *,
+        seed: int | None = None,
+        training: bool = True,
+    ) -> "LMTokenBatch":
+        """Apply shape-stable family dropout by zeroing selected channels."""
+
+        if not training or not policy:
+            return self
+        rates = _validate_side_channel_dropout_policy(policy)
+        if not rates:
+            return self
+        rng = np.random.default_rng(seed)
+        dropped = {
+            family
+            for family, rate in rates.items()
+            if rate >= 1.0 or (rate > 0.0 and float(rng.random()) < rate)
+        }
+        if not dropped:
+            return self
+
+        kwargs = self.as_dict(include_metadata=True, include_side_channels=True)
+        for family in dropped:
+            for field_name in _SIDE_CHANNEL_FAMILY_FIELDS.get(family, ()):
+                value = kwargs.get(field_name)
+                if value is not None:
+                    kwargs[field_name] = mx.zeros_like(value)
+        if self.side_channels is not None:
+            kwargs["side_channels"] = {
+                family: {
+                    name: mx.zeros_like(value) if family in dropped else value
+                    for name, value in columns.items()
+                }
+                for family, columns in self.side_channels.items()
+            }
+        return LMTokenBatch(**kwargs)
+
     def _target_aligned(self, value: mx.array) -> mx.array:
         if tuple(value.shape) == tuple(self.targets.shape):
             return value
@@ -206,7 +273,12 @@ def training_metadata(self) -> dict[str, Any]:
 
         return {} if self.metadata is None else dict(self.metadata)
 
-    def as_dict(self, *, include_metadata: bool = False) -> dict[str, Any]:
+    def as_dict(
+        self,
+        *,
+        include_metadata: bool = False,
+        include_side_channels: bool = False,
+    ) -> dict[str, Any]:
         data: dict[str, Any] = {"tokens": self.tokens}
         if self.target_tokens is not None:
             data["target_tokens"] = self.target_tokens
@@ -219,6 +291,8 @@ def as_dict(self, *, include_metadata: bool = False) -> dict[str, Any]:
         data.update({k: v for k, v in self.structure_fields().items() if v is not None})
         if self.platform_ids is not None:
             data["platform_ids"] = self.platform_ids
+        if include_side_channels and self.side_channels is not None:
+            data["side_channels"] = self.side_channels
         if include_metadata and self.metadata is not None:
             data["metadata"] = self.metadata
         return data
@@ -227,6 +301,41 @@ def as_dict(self, *, include_metadata: bool = False) -> dict[str, Any]:
 _DOCUMENT_ID_ALIASES = ("document_ids", "doc_ids", "packing_document_ids")
 
 
+def _validate_side_channel_map(
+    side_channels: Mapping[str, Mapping[str, mx.array]],
+) -> None:
+    if not isinstance(side_channels, Mapping):
+        raise ValueError("side_channels must be a mapping when provided")
+    for family, columns in side_channels.items():
+        if not isinstance(family, str) or not family.strip():
+            raise ValueError("side channel family names must be non-empty strings")
+        if not isinstance(columns, Mapping):
+            raise ValueError(f"side channel family {family!r} must map column names")
+        for name, value in columns.items():
+            if not isinstance(name, str) or not name.strip():
+                raise ValueError("side channel column names must be non-empty strings")
+            if not isinstance(value, mx.array):
+                raise ValueError(
+                    f"side channel {family}.{name} must be an mlx array"
+                )
+
+
+def _validate_side_channel_dropout_policy(
+    policy: SideChannelDropoutPolicy,
+) -> dict[str, float]:
+    rates: dict[str, float] = {}
+    for family, rate in policy.items():
+        if not isinstance(family, str) or not family.strip():
+            raise ValueError("side channel dropout family names must be non-empty")
+        value = float(rate)
+        if not 0.0 <= value <= 1.0:
+            raise ValueError(
+                f"side channel dropout for {family!r} must be in [0, 1], got {rate!r}"
+            )
+        rates[family] = value
+    return rates
+
+
 def _document_ids_from_mapping(batch: Mapping[str, Any]) -> Any | None:
     present = [
         alias
@@ -271,6 +380,7 @@ def ensure_lm_batch(batch: LMTokenBatch | Mapping[str, Any] | mx.array) -> LMTok
             sibling_index_ids=batch.get("sibling_index_ids"),
             node_type_ids=batch.get("node_type_ids"),
             platform_ids=batch.get("platform_ids"),
+            side_channels=batch.get("side_channels"),
             metadata=batch.get("metadata"),
         )
     raise TypeError(f"unsupported batch type: {type(batch)!r}")
@@ -331,4 +441,9 @@ def synthetic_token_batch(
     )
 
 
-__all__ = ["LMTokenBatch", "ensure_lm_batch", "synthetic_token_batch"]
+__all__ = [
+    "LMTokenBatch",
+    "SideChannelDropoutPolicy",
+    "ensure_lm_batch",
+    "synthetic_token_batch",
+]

cppmega_mlx/nn/structure_embedding.py

@@ -71,7 +71,12 @@ def __init__(
         self.stacked_emb = nn.Embedding(total_vocab, self.bottleneck_dim)
         self.up_proj = nn.Linear(self.bottleneck_dim, self.hidden_size, bias=False)
         self.stacked_emb.weight = mx.zeros_like(self.stacked_emb.weight)
-        self.up_proj.weight = mx.zeros_like(self.up_proj.weight)
+        # Keep the residual output exactly zero at init via the zero table, but
+        # leave a small non-zero projection so CE gradients reach the table.
+        self.up_proj.weight = self.up_proj.weight * mx.array(
+            0.02,
+            dtype=self.up_proj.weight.dtype,
+        )
         self.component_scales = mx.full(
             (len(self.active_component_names),),
             1.0 / len(self.active_component_names),
@@ -180,4 +185,3 @@ def __call__(
 
 
 CppMegaStructureEmbedding = StructureEmbedding
-

cppmega_mlx/training/compiled.py

@@ -55,6 +55,7 @@
     "ast_depth_ids",
     "sibling_index_ids",
     "node_type_ids",
+    "platform_ids",
 )
 
 CompiledBatch = dict[str, mx.array | None]

cppmega_mlx/training/loss.py

@@ -8,7 +8,11 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from cppmega_mlx.data.batch import LMTokenBatch, ensure_lm_batch
+from cppmega_mlx.data.batch import (
+    LMTokenBatch,
+    SideChannelDropoutPolicy,
+    ensure_lm_batch,
+)
 from cppmega_mlx.nn.structure_embedding import StructureEmbedding
 from cppmega_mlx.training.cut_cross_entropy import (
     DEFAULT_CHUNK_ROWS,
@@ -31,10 +35,18 @@
 def next_token_cross_entropy(
     model: nn.Module,
     batch: LMTokenBatch | Mapping[str, mx.array] | mx.array,
+    *,
+    side_channel_dropout: SideChannelDropoutPolicy | None = None,
+    side_channel_dropout_seed: int | None = None,
 ) -> tuple[mx.array, mx.array]:
     """Return masked next-token CE loss and the number of contributing tokens."""
 
     lm_batch = ensure_lm_batch(batch)
+    lm_batch = lm_batch.with_side_channel_dropout(
+        side_channel_dropout,
+        seed=side_channel_dropout_seed,
+        training=True,
+    )
     document_ids = lm_batch.input_document_ids
     model_kwargs = lm_batch.model_kwargs()
     if document_ids is not None:
@@ -63,6 +75,8 @@ def next_token_cut_cross_entropy(
     *,
     chunk_rows: int = DEFAULT_CHUNK_ROWS,
     eval_chunks: bool = True,
+    side_channel_dropout: SideChannelDropoutPolicy | None = None,
+    side_channel_dropout_seed: int | None = None,
 ) -> tuple[mx.array, mx.array]:
     """Return masked next-token CE via the MLX-native chunked linear path.
 
@@ -73,6 +87,11 @@ def next_token_cut_cross_entropy(
     """
 
     lm_batch = ensure_lm_batch(batch)
+    lm_batch = lm_batch.with_side_channel_dropout(
+        side_channel_dropout,
+        seed=side_channel_dropout_seed,
+        training=True,
+    )
     document_ids = lm_batch.input_document_ids
     hidden_states = _decoder_hidden_states_for_mtp(
         model,