Arain119

/

sophia

Like 1

Hi Andrew — thank you for such a careful read. I went through the Maba v2 repo: it is rare for an architecture release to ship with actual Triton kernels, CPU/XLA dispatch fallbacks, and a 649-test suite — most stop at a diagram. The decoupled-gated delta + centroid indexing design is thoroughly documented and detailed enough to reproduce directly; that kind of engineering rigor is admirable.

On the optimizer split: Muon is not restricted to MLPs here. Grouping is by tensor rank — every non-embedding 2D matrix goes to Muon, including KDA's recurrent transition projections (q/k/v/out and the gate projections). AdamW only takes embeddings/lm_head and the sub-2D parameters — norms, biases, and 1D params like A_log and dt_bias. Newton–Schulz at 5 steps.

On NIAH our reads largely agree: 62.5% with depth decay is consistent with progressive dilution of isolated facts across linear state transitions. The difference is only in attribution — we lean toward state capacity at 1B scale, since the decay begins before the target feature reaches the attention blocks. Which is exactly why your indexer caught my attention: the genuinely open question is whether salient-coordinate routing can preserve the signal when the upstream recurrent state is already lossy. If you have small-scale ablations on indexer vs state quality, that would be a valuable reference.

Best,
Arain

Hi Arain,

That is the exact crux of the problem, and our empirical runs on the 1M NIAH benchmark confirm your diagnosis on state capacity: when depth relies purely on recurrent state compression, isolated token coordinates inevitably suffer collision and progressive dilution across long sequence spans.

The reason the centroid indexer rescues the needle even when upstream recurrent state is lossy comes down to three concrete mechanics in the architecture:

1. Decoupling from the Recurrent State Matrix:
   The indexer (DGIndexer) and sparse attention (MABA-SA) do not read from DGDA's compressed recurrent state S_t. Instead, they project coordinates (q_idx, k_idx) and latent key-values directly from the global residual stream. While sequential recurrent steps experience capacity saturation, token identities persist along the residual highway via skip-connections.

2. Hybrid Anti-Dilution Pooling (Mean + Max):
   Standard chunk pooling using mean(K_b) is disastrous for single-token retrieval: a single needle token inside a 64-token block gets diluted by a factor of 64 against background filler.
   In Maba v2, the indexer computes:

$$c_b=\frac{1}{2}(\text{mean}(K_b)+\max (K_b))-\lambda \cdot \log (1+{\mathrm{\Delta }}_b)c\_b\; =\; \backslash frac\{1\}\{2\}\backslash left(\backslash text\{mean\}(K\_b)\; +\; \backslash max(K\_b)\backslash right)\; -\; \backslash lambda\; \backslash cdot\; \backslash log(1\; +\; \backslash Delta\_b)$$

The element-wise max operation latches onto the sharp coordinate activation of the needle token, preventing it from being washed out by surrounding tokens. In our benchmark with 50 adversarial hard negatives (decoys crafted at 95% cosine similarity), this hybrid pooling consistently pulls the true needle block into Top-32 (Rank #1 out of 15,625 blocks at 1M context, 123 ms scan latency).

3. Empirical Ablation Signal:
   We tracked this exact tension across our ablation ladder (pure_dgda vs no_hca vs full 3:1 Maba):
   Under pure linear recurrence (pure_dgda / pure SSM), retrieval on distant tokens degrades to ~61.3% (Rank #4 / 15,625), confirming your observation that state capacity bounds long-range fidelity.
   Routing the top-32 candidate blocks into uncompressed local attention every 4th layer acts as an exact retrieval anchor. In our multi-hop benchmark across a 640k-token gap (Hop 1 at 128k, Hop 2 at 768k), both coordinates are jointly recovered into Top-32, re-injecting sharp representational signal back into the residual stream for subsequent recurrent blocks.

If you want to test whether this hybrid pooling stabilizes retrieval on Sophia's 1B checkpoints, the DGIndexer module is standalone and plugs directly into standard residual streams. Happy to share the raw ablation logs and test harnesses!

Best,
Andrew

Hi Andrew,

Thanks for the detailed breakdown. The mean + max pooling is a neat fix for single-token dilution inside a block.

One thing I should have mentioned earlier: in Sophia the Gated-MLA layers use full causal attention over the whole 4k context. Their inputs come from our depth-wise attention residuals, which mix the outputs of earlier blocks directly, so the signal doesn't have to pass through the KDA recurrent state alone. So block selection is not really a bottleneck at our scale: every 4th layer can already attend to every token. That's why I don't think an indexer would change much on the current model. Our 62.5% is more likely down to state capacity at 1B.

We won't change the architecture of the released checkpoint, but I'll keep the indexer and the pooling idea in mind for the next generation, especially if we push context well beyond 4k.

Thanks again for the discussion!

Best,
Arain

Hi Arain,

That makes total sense! If your Gated-MLA layers are already running dense attention over the full 4k context via depth-wise residuals, an indexer would definitely be redundant at that length.

You're right: at 4k, chunk selection isn't the bottleneck, so that 62.5% really is just state capacity in the KDA stack at 1B. Sparse indexing only starts paying dividends once you scale past 16k/32k+ where running dense attention every 4th layer becomes too expensive.

Really enjoyed talking shop and digging into the mechanics. Best of luck with Sophia and future scaling! :)

Best,
Andrew