Local-Axiom-AI

/

Sabaki-Preview

Like 1

Hi Local-Axiom-AI team,

Pretraining a custom 221.5M 3:1 hybrid (12 KDA recurrent layers, 4 latent-attention MLA layers) with Top-2/8 MoE from scratch across 8B tokens, and shipping a custom GGUF runtime with honest documentation of the runaway generation failure mode, is fantastic systems work. The sample efficiency against 125M baselines on just 8B tokens confirms the representational power of the 3:1 macro-topology.

Looking at your architectural schedule and the runaway generation issue:

1. Rotary phase mismatch in intermittent MLA:
   Your schedule interleaves 3 KDA blocks with 1 rotary MLA block. KDA updates its recurrent state in continuous coordinate space without positional phase rotation. Injecting RoPE specifically at the MLA layers introduces an alternating phase projection: representations must repeatedly transition between rotation-invariant recurrent updates and frequency-rotated attention manifolds.
   This periodic coordinate torsion makes it difficult for the network to establish sharp delimiter boundaries, which directly contributes to the model failing to emit EOS after completing a response.

2. Residual state retention and attractor basins:
   In a recurrent backbone, finishing an answer leaves non-zero activation energy in the KDA state matrices. In standard attention, trailing tokens can cleanly attend back to a prompt-level boundary to fire an EOS logit. But when 75% of depth is linear recurrence, if the attention layers lack an explicit salience-isolation mechanism, background sequence momentum keeps propagating through the recurrent state. The model cannot dissipate this residual energy, pushing generation into the runaway nonsensical loops you observed.

3. The 22.4% active vocabulary footprint:
   With a 49,152 byte-level BPE vocabulary at hidden size 768, the tied embedding table consumes 37.75M parameters.
   Against your 130.7M active non-embedding parameters, static token lookup accounts for 22.4% of your total forward compute budget per token. Decoupling or factorizing token embeddings would free up substantial parameter capacity to expand your routed expert count without changing inference latency.

4. Resolving runaway dynamics via NoPE and bounded indexing:
   In an open architecture project called Maba v2 (reference release: https://huggingface.co/AndrewThompson1233/maba-v2-architecture), which shares this exact 3:1 macro-balance between linear recurrence and sparse attention:
   We eliminated positional phase conflict entirely by running native NoPE across all layers, relying on recurrent state progression for sequence order.
   To prevent residual runaway loops, we paired the recurrent backbone with MABA-SA (a sparse attention mechanism with anti-dilution indexing). This cleanly isolates delimiter signals from lingering recurrent activation noise, guaranteeing clean EOS termination without runaway wandering.

If you are planning a Sabaki v2 or another post-training pass, checking out the hybrid coordination and state management layout in the Maba v2 repo might provide some practical reference points for stabilizing generation boundaries.

Did the runaway generation behavior emerge primarily during the 4,096-token curriculum stage, or was it already pronounced during the 512 base pretrain?

Best,
Andrew