"""Matched dense causal-decoder controls for the Keystone Metal contest. Keystone's claim is *stored-parameter* efficiency, so any comparison is only meaningful against a conventional dense decoder holding the same number of trainable scalars, reading the same tokenizer, and seeing the same batches. This module supplies that control: a tied dense token table, full-rank multi-head causal attention, and dense gated channel MLPs. There is no weight sharing, low-rank factorization, mixture routing, or quantization here — that is the point of the control. Two shapes are published rather than one. A single hand-picked control shape can silently flatter the candidate, so the contest trains every shape and the *best* dense result is used as the control (strongest-control policy). """ from __future__ import annotations from dataclasses import asdict, dataclass import math from typing import Any import mlx.core as mx import mlx.nn as nn from mlx.utils import tree_flatten import numpy as np try: from keystone_mlx.model import CenteredUnitNorm, DenseTiedTokens, deterministic_coordinates except ModuleNotFoundError: from model import CenteredUnitNorm, DenseTiedTokens, deterministic_coordinates @dataclass(frozen=True) class DenseControlConfig: """Immutable specification for one matched dense decoder control. Args: vocab_size: Token IDs accepted by the tied dense interface. width: Residual stream width. blocks: Number of independent (non-shared) decoder blocks. head_dim: Channels per attention head; ``width`` must divide evenly. mlp_width: Hidden width of each block's dense gated MLP. """ vocab_size: int = 2_048 width: int = 128 blocks: int = 6 head_dim: int = 32 mlp_width: int = 148 def __post_init__(self) -> None: if self.vocab_size < 2: raise ValueError("vocab_size must be at least two") if self.width < 2 or self.blocks < 1 or self.mlp_width < 1: raise ValueError("width, blocks, and mlp_width must be positive") if self.head_dim < 1 or self.width % self.head_dim != 0: raise ValueError("head_dim must divide width exactly") # The 1M shapes are matched to Keystone-1M-2K's exact 999,744 stored scalars. # `deep` is 999,680 (0.0064% under) and `wide` is 999,360 (0.0384% under). DENSE_CONTROL_DEEP = DenseControlConfig(vocab_size=2_048, width=128, blocks=6, head_dim=32, mlp_width=148) DENSE_CONTROL_WIDE = DenseControlConfig(vocab_size=2_048, width=160, blocks=4, head_dim=32, mlp_width=135) # 600K M-SimOW scale-down: preserve the 1M deep model's six blocks and 2K # tied vocabulary while narrowing the residual stream to 96. This geometry # lands on exactly 600,000 stored parameters. DENSE_CONTROL_600K_DEEP = DenseControlConfig(vocab_size=2_048, width=96, blocks=6, head_dim=32, mlp_width=104) # The 10M pair is matched to Trellis-10M-2K's 10,001,408 scalar budget: # deep = 10,000,256 (-0.0115%); wide = 9,996,800 (-0.0461%). Both retain the # same 2K tokenizer interface and are inside the predeclared 0.05% tolerance. DENSE_CONTROL_10M_DEEP = DenseControlConfig(vocab_size=2_048, width=448, blocks=8, head_dim=64, mlp_width=246) DENSE_CONTROL_10M_WIDE = DenseControlConfig(vocab_size=2_048, width=512, blocks=5, head_dim=64, mlp_width=481) # Under-3M deep geometry for the Ember-inspired XSA recurrent scale gate. # It preserves the 10 physical / 14 effective-block layout with a healthy # 4.22x gated MLP and totals exactly 2,996,480 stored parameters. DENSE_CONTROL_3M_DEEP = DenseControlConfig(vocab_size=2_048, width=128, blocks=10, head_dim=32, mlp_width=540) # Exact 2.9M scale-up: 2,896,640 stored parameters at the 2K vocabulary. DENSE_CONTROL_2P9M_DEEP = DenseControlConfig(vocab_size=2_048, width=128, blocks=10, head_dim=32, mlp_width=514) # Matched-budget wider recurrent-XSA candidate. At 2,893,440 stored parameters, # it trades four physical blocks for a wider residual stream while retaining # ten effective transformations under XSA's recurrent middle-block layout. DENSE_CONTROL_2P9M_WIDE = DenseControlConfig(vocab_size=2_048, width=192, blocks=6, head_dim=32, mlp_width=466) # The 5M scale-up keeps the same tied dense-deep family while increasing width # and depth together: 5,001,728 stored parameters at the 2K vocabulary. DENSE_CONTROL_5M_DEEP = DenseControlConfig(vocab_size=2_048, width=256, blocks=8, head_dim=64, mlp_width=386) DENSE_CONTROL_SHAPES = { "1m-2k": {"deep": DENSE_CONTROL_DEEP, "wide": DENSE_CONTROL_WIDE}, "600k-2k": {"deep": DENSE_CONTROL_600K_DEEP}, "10m-2k": {"deep": DENSE_CONTROL_10M_DEEP, "wide": DENSE_CONTROL_10M_WIDE}, "3m-2k": {"deep": DENSE_CONTROL_3M_DEEP}, "2p9m-2k": {"deep": DENSE_CONTROL_2P9M_DEEP, "wide": DENSE_CONTROL_2P9M_WIDE}, "5m-2k": {"deep": DENSE_CONTROL_5M_DEEP}, } def control_for_shape(shape: str, model_size: str = "1m-2k") -> DenseControlConfig: """Select a published dense control shape at one fixed budget.""" try: return DENSE_CONTROL_SHAPES[model_size][shape] except KeyError as error: choices = ", ".join(sorted(DENSE_CONTROL_SHAPES)) raise ValueError(f"unknown dense control selection {shape!r} at {model_size!r}; budgets: {choices}") from error class DenseCausalAttention(nn.Module): """Full-rank multi-head causal attention with an explicit additive mask.""" def __init__(self, width: int, head_dim: int) -> None: super().__init__() if width % head_dim != 0: raise ValueError("head_dim must divide width exactly") self.heads = width // head_dim self.head_dim = head_dim self.query = nn.Linear(width, width, bias=False) self.key = nn.Linear(width, width, bias=False) self.value = nn.Linear(width, width, bias=False) self.output = nn.Linear(width, width, bias=False) def __call__(self, hidden: mx.array) -> mx.array: """Attend only to the current position and its strict prefix. Args: hidden: States shaped ``[batch, sequence, width]``. Returns: Context shaped like ``hidden``. Position ``i`` cannot read ``> i``. Side effects: Allocates per-head query/key/value and mask intermediates. """ if hidden.ndim != 3: raise ValueError("hidden must have shape [batch, sequence, width]") batch, length, width = hidden.shape def split_heads(values: mx.array) -> mx.array: return values.reshape(batch, length, self.heads, self.head_dim).transpose(0, 2, 1, 3) queries = split_heads(self.query(hidden)) keys = split_heads(self.key(hidden)) values = split_heads(self.value(hidden)) scores = (queries @ keys.transpose(0, 1, 3, 2)) * (self.head_dim**-0.5) positions = mx.arange(length) causal_block = mx.where( positions[:, None] >= positions[None, :], mx.array(0.0, dtype=scores.dtype), mx.array(-1e9, dtype=scores.dtype) ) weights = mx.softmax(scores + causal_block[None, None, :, :], axis=-1) attended = (weights @ values).transpose(0, 2, 1, 3).reshape(batch, length, width) return self.output(attended) class DenseGatedMLP(nn.Module): """Dense gated channel MLP with independent, unshared matrices.""" def __init__(self, width: int, mlp_width: int) -> None: super().__init__() self.expand = nn.Linear(width, 2 * mlp_width, bias=False) self.contract = nn.Linear(mlp_width, width, bias=False) def __call__(self, state: mx.array) -> mx.array: """Apply one gated dense channel transform.""" content, gate = mx.split(self.expand(state), 2, axis=-1) return self.contract(nn.silu(content) * mx.sigmoid(gate)) class DenseControlBlock(nn.Module): """One pre-normalized attention plus gated-MLP residual block.""" def __init__(self, config: DenseControlConfig) -> None: super().__init__() self.attention_norm = CenteredUnitNorm(config.width) self.attention = DenseCausalAttention(config.width, config.head_dim) self.mlp_norm = CenteredUnitNorm(config.width) self.mlp = DenseGatedMLP(config.width, config.mlp_width) def __call__(self, hidden: mx.array) -> mx.array: """Run attention and channel mixing with residual connections.""" hidden = hidden + self.attention(self.attention_norm(hidden)) return hidden + self.mlp(self.mlp_norm(hidden)) class DenseControlLM(nn.Module): """Conventional dense decoder used as Keystone's parameter-matched control.""" def __init__(self, config: DenseControlConfig = DENSE_CONTROL_DEEP) -> None: super().__init__() self.config = config self.tokens = DenseTiedTokens(config.vocab_size, config.width) self.blocks = [DenseControlBlock(config) for _ in range(config.blocks)] self.final_norm = CenteredUnitNorm(config.width) def __call__(self, token_ids: mx.array) -> mx.array: """Compute full-vocabulary causal next-token logits. Args: token_ids: Integer IDs shaped ``[batch, sequence]``. Returns: A ``[batch, sequence, vocab_size]`` logit tensor. Side effects: Allocates coordinate and per-block intermediate arrays. """ if token_ids.ndim != 2: raise ValueError("token_ids must have shape [batch, sequence]") if token_ids.shape[1] < 1: raise ValueError("token_ids must include at least one position") hidden = self.tokens.encode(token_ids) coordinates = deterministic_coordinates(token_ids.shape[1], self.config.width) hidden = hidden + mx.expand_dims(coordinates, axis=0) for block in self.blocks: hidden = block(hidden) return self.tokens.decode(self.final_norm(hidden)) def parameter_accounting(config: DenseControlConfig = DENSE_CONTROL_DEEP) -> dict[str, int]: """Return the exact stored-parameter budget, counting the tied table once.""" vocabulary = config.vocab_size * config.width attention = config.blocks * 4 * config.width * config.width mlp = config.blocks * 3 * config.width * config.mlp_width normalization = (config.blocks * 2 + 1) * 2 * config.width total = vocabulary + attention + mlp + normalization return { "dense_tied_vocabulary_interface": vocabulary, "independent_attention": attention, "independent_gated_mlps": mlp, "normalization": normalization, "total": total, } def parameter_count(model: nn.Module) -> int: """Count trainable MLX scalars, including each tied array exactly once.""" return sum(int(parameter.size) for _, parameter in tree_flatten(model.trainable_parameters())) def future_token_causality_error( model: DenseControlLM, sequence_length: int = 11, batch_size: int = 2, seed: int = 23, ) -> float: """Return the largest changed prefix logit after deterministic suffix edits.""" if sequence_length < 3 or batch_size < 1: raise ValueError("sequence_length must be at least three and batch_size must be positive") generator = np.random.default_rng(seed) values = generator.integers(0, model.config.vocab_size, size=(batch_size, sequence_length), dtype=np.int32) boundary = sequence_length // 2 mutated = values.copy() mutated[:, boundary:] = (mutated[:, boundary:] + 1) % model.config.vocab_size original = model(mx.array(values)) changed = model(mx.array(mutated)) mx.eval(original, changed) return float(mx.max(mx.abs(original[:, :boundary] - changed[:, :boundary]))) def self_check(config: DenseControlConfig = DENSE_CONTROL_DEEP, reference_total: int = 999_744) -> dict[str, Any]: """Validate accounting, shape, causality, and the match tolerance. Args: config: Dense control shape to validate. reference_total: Stored-parameter budget the control must match. Returns: A JSON-serializable diagnostic record including the signed mismatch. Side effects: Builds the model and evaluates two Metal inference passes. """ model = DenseControlLM(config) formula = parameter_accounting(config) observed = parameter_count(model) if observed != formula["total"]: raise RuntimeError(f"formula counted {formula['total']}, model stores {observed}") relative = abs(observed - reference_total) / reference_total if relative > 5e-4: raise RuntimeError(f"control stores {observed} scalars, {relative:.4%} from reference {reference_total}") probe = model(mx.zeros((2, 11), dtype=mx.int32)) mx.eval(probe) if tuple(probe.shape) != (2, 11, config.vocab_size): raise RuntimeError(f"unexpected logit shape {tuple(probe.shape)}") causal_error = future_token_causality_error(model) if causal_error != 0.0: raise RuntimeError(f"future-token leakage: {causal_error}") return { "architecture": "DenseControl", "config": asdict(config), **formula, "reference_total": reference_total, "relative_mismatch": relative, "causal_forward_shape": list(probe.shape), "future_token_causality_error": causal_error, }