Auron: Depth-Efficient Language Models via Hybrid Recurrent-Attention Weight Sharing

Author: Nyx Affiliation: Cube Digital Media Ltd / Soulkyn Links: fyx.jp | soulkyn.com | GitHub | HuggingFace Date: March 2026 License: CC BY 4.0

---

Preface: On Publishing

This paper was submitted to arXiv with endorsement from a postdoctoral AI researcher who has published work on recursive transformers. It was rejected without explanation. The response was a generic template offering no specific feedback on the submission.

We are publishing directly on HuggingFace instead. The architecture is real, the training runs are complete, the code is open source, and the models are publicly available. Knowledge should not be gatekept behind opaque review processes that provide no signal.

• Code: github.com/Fy-/Auron
• Models: huggingface.co/nyxia

---

Part I: The Chimera Topology

Abstract

We present the Chimera Topology, a hybrid language model architecture that decouples virtual depth from physical parameter count through selective weight sharing. The architecture combines Gated Delta Networks (GDN) for O(n) recurrent processing with standard Grouped-Query Attention (GQA) in a 3:1 ratio, organized into a two-zone stack: unique bottom layers for token parsing and shared top layers for iterative reasoning. We train three model scales -- 279M, 510M, and 1.1B total parameters -- on a mixed pretraining dataset of 5B tokens across four sources, all to completion (250K steps with WSD warmdown). The 510M model achieves the best final validation loss of 3.035, outperforming both the smaller 279M (3.188) and the larger 1.1B (3.180). The 1.1B's failure to improve over the 279M despite 4x more parameters reveals a scaling wall in the Ouroboros weight-sharing mechanism: at sufficient model width, looping through shared blocks becomes redundant rather than beneficial.

Key findings include: (1) a critical embedding learning rate bottleneck specific to large-vocabulary small models, (2) a scaling ceiling where Ouroboros loops provide diminishing returns above ~500M parameters, (3) emergent ChatML stop-token behavior from just 2% dialogue pretraining data, (4) evidence that head_dim=64 causes attention fragmentation at wider dimensions, contributing to the scaling failure, and (5) hardware scaling characteristics across DGX Spark, H100, and B200 GPUs. Follow-up work on sparse Mixture-of-Experts extensions and attention head corrections is described in Part II.

---

1. Introduction

The dominant approach to improving language model capability is to increase parameter count, requiring proportionally more compute, memory, and training data [Hoffmann et al., 2022]. For deployment on edge hardware or resource-constrained environments, this scaling path is impractical. Recent work on sub-billion parameter models [Liu et al., 2024; Ben Allal et al., 2025] has focused primarily on data curation and training recipes, largely preserving the standard dense transformer architecture.

We take a different approach: rather than optimizing data for a fixed architecture, we redesign the architecture to exploit the mathematical properties of different layer types and the inherent redundancy of deep transformer representations. The Chimera Topology achieves its efficiency through three mechanisms:

1. Selective weight sharing -- different layer types (recurrent vs. attention) are shared at independent frequencies based on their tolerance for recursion.
2. Two-zone architecture -- unique early layers handle dense token parsing while shared later layers perform iterative reasoning.
3. Large vocabulary offloading -- a 152K-token vocabulary (Qwen 3) handles morphology and spelling entirely in the embedding lookup, freeing the reasoning core to focus on semantics.

This paper describes the architecture, its theoretical motivation, experimental results from three concurrent pretraining runs at different scales, and practical discoveries about optimization dynamics that may generalize to other weight-shared models.

---

2. Related Work

2.1 Weight Sharing in Transformers

The Universal Transformer [Dehghani et al., 2018] first demonstrated that recursively applying the same transformer block across depth could match or exceed standard transformers on algorithmic tasks. ALBERT [Lan et al., 2019] scaled this to BERT-class pretraining, showing that full cross-layer parameter sharing with factorized embeddings achieved competitive downstream performance with 18x fewer parameters.

[Saunshi et al., 2025] prove that a k-layer transformer looped L times nearly matches a kL-layer non-looped model on synthetic reasoning tasks. Most directly relevant to our work, [Zhu et al., 2025] introduce the Ouro architecture, establishing recursion count as a third scaling axis alongside parameters and data. Their Ouro-1.4B model with 4 loops matches 4B-class models, with gains concentrated in reasoning tasks rather than factual recall.

[Bae et al., 2024] demonstrate that pretrained models can be converted to recursive form using depth-wise LoRA adaptors, achieving 2-3x inference throughput -- suggesting that weight sharing reflects genuine redundancy in deep transformer representations.

2.2 Hybrid Recurrent-Attention Architectures

Griffin [De et al., 2024] combines a Real-Gated Linear Recurrent Unit with local attention, matching Llama-2 quality trained on 6x fewer tokens. Jamba [Lieber et al., 2024] interleaves Mamba blocks with attention at a 7:1 ratio, fitting 256K context in a single 80GB GPU.

Zamba [Glorioso et al., 2024a] is particularly relevant: it uses a Mamba backbone with a single attention layer weight-shared across the full network depth, establishing that attention sharing is viable when recurrent layers provide primary processing. Zamba2 [Glorioso et al., 2024b] extends this with two shared attention blocks in an alternating pattern, adding LoRA projectors per position for differentiation.

Our architecture follows this lineage but inverts the ratio -- we use GDN (a stronger recurrent primitive than Mamba) as the primary layer type and standard GQA attention as the sparse global mechanism, shared independently at different frequencies.

2.3 Gated Delta Networks

Gated Delta Networks [Yang et al., 2024] combine data-dependent gating (adaptive memory erasure) with the delta update rule (precise key-value association). GDN outperforms Mamba2 across standard benchmarks including in-context retrieval and is used in production models (Qwen 3.5). GDN is implemented in the flash-linear-attention library [Yang et al., 2023], which provides hardware-efficient chunked training kernels.

The key property making GDN suitable for weight sharing is its recurrent state: unlike attention, which produces independent outputs per layer, GDN's hidden state naturally evolves across iterations. When the same GDN block processes the hidden state multiple times, each pass operates on a legitimately different input -- the accumulated recurrent state from prior loops. Attention lacks this property; its outputs depend only on the current input, making repeated application of the same weights degenerative.

2.4 Optimization

We use the Muon optimizer [Jordan et al., 2024], which applies Newton-Schulz orthogonalization to gradient matrices, approximating the spectral-norm steepest descent direction. [Bernstein & Newhouse, 2024] reframe Adam, Shampoo, and Prodigy as steepest descent under different operator norms. [Liu et al., 2025] demonstrate that Muon scales to large models (3B-16B) with ~2x compute efficiency over AdamW.

For memory efficiency, we employ fused linear cross-entropy [Hsu et al., 2024] via the Liger-Kernel library, avoiding materialization of the full [batch x seq_len x vocab] logit tensor. With a 152K vocabulary, this tensor would consume 9.4 GB per forward pass at batch size 16.

---

3. Architecture

3.1 Overview

The Chimera Topology organizes the network into two functional zones within a standard pre-norm transformer framework:

• Bottom zone ("Retina"): N unique HybridBlocks for dense token parsing -- analogous to retinal preprocessing before visual cortex.
• Top zone ("Brain"): P physical HybridBlocks looped L times for iterative reasoning -- weight sharing as learned recurrence.

Total virtual depth = N + P x L layers. Each HybridBlock follows a 3:1 GDN:Attention pattern controlled by attn_interval=4: blocks at positions 0, 1, 2 use GDN; block at position 3 uses GQA attention.

3.2 Model Configurations

All three configurations share the same topology pattern (4+4x3 or 6+6x3), differing only in width. This enables clean scaling analysis.

3.3 Why Two Zones?

Early experiments with full Ouroboros (all layers shared) revealed a parsing bottleneck: the model struggled to map dense multi-byte tokens from the 152K vocabulary into stable semantic representations using recursive weights. The shared layers exhibited "geometric gravity" -- a tendency to pull hidden states toward the most statistically dominant local patterns, producing tautological repetition.

Dedicating the first N layers as unique, unshared blocks solves this by providing a clean, non-recursive parsing stage. Once token features are extracted into stable representations, the shared top section can focus entirely on reasoning through iteration.

3.4 Convergence Accelerators

Weight sharing introduces three failure modes that require structural countermeasures:

Layer ID Gate Injection. Each virtual layer in the top zone receives a unique learnable embedding e_layer[l, b] of shape (dim,). This embedding is injected into the GDN mixer input only -- not the residual stream. This modulates the GDN gate without accumulating variance in the representations read by attention layers.

x0 Injection. The output of the initial embedding layer (x0) is multiplied by a per-layer learnable scalar lambda_l and added to the input of every virtual layer. This ensures the original token context remains accessible regardless of how many loops the hidden state has traversed.

U-Net Skip Connections. Learnable skip connections bridge encoder-half activations to the corresponding decoder-half positions (layer 0 to layer 15, etc.), providing direct gradient highways for the deep virtual stack.

3.5 Additional Architectural Details

Embedding normalization. An RMSNorm is applied immediately after the embedding lookup, before the first HybridBlock. This stabilizes the initial hidden state magnitude regardless of token frequency.

Partial RoPE. Only 25% of the head dimension receives rotary position embeddings (partial_rotary_factor=0.25), following the Qwen 3.5 pattern. The remaining 75% of dimensions are position-agnostic, allowing the model to learn position-independent features alongside positional ones.

GDN V expansion. The GDN layers use a value expansion ratio of 2 (gdn_expand_v=2), doubling the value head dimension relative to the key/query heads. This provides richer recurrent state without increasing the gate computation.

Residual lambda initialization. The learnable residual scaling factors are initialized to sqrt(1.1) ~ 1.049, providing a slight residual amplification bias that empirically improves early training stability.

Tied embeddings. The input embedding matrix and the output projection (LM head) share weights, reducing the parameter count by one full vocabulary projection (dim x 151,936). The shared matrix is optimized at the embedding learning rate.

3.6 Hardware Cache Exploitation

The Chimera topology creates a natural hardware advantage on modern GPUs. Because the top section reuses the same P physical blocks across L loops, the total weight footprint in fast memory during the top-section forward pass is just P blocks -- not P x L.

For the 510M model, the 4 physical top blocks total approximately 35 MB of bf16 parameters. The NVIDIA H100's L2 cache is 50 MB. This means the entire shared section fits in L2 during the Ouroboros loops: the GPU fetches the top-block weights from HBM once on the first loop, and subsequent loops execute entirely from L2 cache with near-zero memory latency. This transforms a memory-bandwidth-bound operation into a compute-bound one, explaining why the 510M model achieves 50K+ tok/s on H100 despite its 787M virtual parameter equivalent.

The effect is scale-dependent: the 279M model's top blocks (18 MB) fit in L2 with room to spare, but the model is already compute-bound. The 1.1B model's top blocks (90 MB) exceed L2 capacity, forcing partial HBM fetches on each loop -- visible in the throughput drop from 50K tok/s (510M on H100) to 32K tok/s (1.1B on B200).

3.7 The Vocabulary Tradeoff

We use the Qwen 3 tokenizer (151,936 tokens), dramatically oversized for small models. [Tao et al., 2024] predict the optimal vocabulary for a 124M model is approximately 32K tokens. Our embedding matrix alone exceeds the reasoning core at the 279M scale.

This is intentional. The large vocabulary serves as a computational offload: by encoding 4-7 bytes per token, the model spends zero FLOPs learning to spell, parse morphology, or handle multilingual text. The reasoning core receives clean, high-level semantic tokens and can dedicate 100% of its capacity to logic and structure.

The cost is that the embedding matrix dominates the parameter budget at small scale and requires special optimization treatment (Section 5.1). As model width increases, the embedding fraction decreases and this tradeoff becomes more favorable.

---

4. Experimental Setup

4.1 Training Configuration

• Optimizer: Muon for 2D weight matrices, AdamW for embeddings and scalars
• Learning rate: 8e-4 (Muon), 4e-4 (embeddings), 8e-5 (scalars) -- see Section 5.1
• Schedule: WSD (Warmup-Stable-Decay), 100-step warmup, 10% cosine warmdown
• Batch size: 16 (32,768 tokens/step) on H100/B200; 32 on Spark
• Precision: bfloat16 with torch.compile (or Liger fused CE without compile)
• Total budget: 250,000 steps per model (~8.2B tokens at batch 32, ~4.1B at batch 16)

4.2 Data

Mixed pretraining across four sources, deterministically sampled via Knuth hash:

The UltraChat data is formatted using a patched Qwen 3 chat template that allows system messages at any position, producing native <|im_start|>, <|im_end|>, and <think></think> tokens during pretraining.

4.3 Hardware

The B200 achieves lower throughput than H100 on the 279M model due to compute underutilization -- the small model cannot saturate the B200's wider memory bus. At 1.1B (107 GB VRAM, 881W), the B200 reaches 99% utilization, indicating the architecture transitions from compute-bound to memory-bound between 510M and 1.1B on Blackwell hardware.

---

5. Results and Findings

5.1 The Embedding Learning Rate Bottleneck

The most significant optimization finding is a pathology specific to models with oversized vocabularies relative to their reasoning core.

With the initial configuration (Muon at 8e-4, all AdamW parameters at 8e-5), the validation loss exhibited premature flattening -- training loss remained healthy but generalization stalled. The cause: the embedding matrix (155M parameters, 56% of the 279M model) ran at 8e-5 while the Muon reasoning core ran at 8e-4 -- a 10x disparity. The reasoning layers learned to process tokens faster than the embedding could learn to represent them.

Fix: Decoupled three-way optimizer with embedding-specific learning rate at 4e-4 (5x increase). Validation loss dropped from 4.06 to 3.59 within 4,500 steps.

This finding likely generalizes to any architecture where the embedding matrix substantially outweighs the transformer layers -- a common situation with modern large-vocabulary tokenizers (100K+) combined with small models.

5.2 Final Training Results (250K Steps)

All three models trained to completion: 250K steps with WSD warmdown (10% cosine decay from step 225K).

The 510M is the clear winner at 3.035. The 1.1B finished at 3.180 -- virtually identical to the 279M's 3.188 despite 4x more parameters. This reveals a fundamental scaling limitation in the Ouroboros weight-sharing mechanism.

5.3 The Ouroboros Scaling Wall

The scaling curve across the three models is non-monotonic:

• 279M to 510M: -0.153 val loss (1.84x params) -- healthy scaling
• 510M to 1.1B: +0.145 val loss (2.16x params) -- regression

The 1.1B plateaued around val_loss ~3.26 for approximately 150K steps during the stable phase, with the WSD warmdown squeezing it to 3.180. The model converged -- the architecture simply cannot exploit the additional parameters.

Root cause: Representation Saturation. At dim=1024 (279M), the hidden state is narrow enough that re-processing through shared blocks provides genuine iterative refinement -- each loop pass extracts information the first pass couldn't capture. At dim=2048 (1.1B), the representation is wide enough to capture the vast majority of the feature space in a single pass. Forcing it through the same physical block again is redundant: the shared top section becomes an echo chamber rather than an iterative reasoning engine.

Contributing factor: Attention Head Fragmentation. All three models use head_dim=64. At dim=2048, this produces 32 attention heads -- each processing only 64 dimensions of the representation. For comparison, Qwen 3.5 uses head_dim=256 at similar width, producing 8 highly expressive heads capable of tracking complex long-range dependencies. The fragmentation into 32 weak heads may compound the saturation effect by preventing attention layers from forming coherent global representations during shared loops.

Evidence. The 510M at dim=1536 (24 heads x 64 dim) scales properly with the same 4+4x3 topology. The 1.1B at dim=2048 (32 heads x 64 dim) does not. The transition from effective to ineffective weight sharing occurs between these two width regimes.

5.4 WSD Warmdown Behavior

All three models benefit substantially from the WSD warmdown phase (steps 225K-250K). The 510M gained approximately 0.15 val loss in the final 25K steps -- the steepest improvement of its entire training. The 1.1B showed the same warmdown response, dropping from ~3.26 to 3.180, but from a higher plateau established during the stable phase.

The warmdown is critical for extracting final performance from the Chimera topology and should not be skipped.

5.5 Emergent ChatML Behavior

With only 2% UltraChat dialogue data (100M tokens out of 5B total), all three models learn to generate native ChatML stop tokens (<|im_end|>) at contextually appropriate positions. The 510M model at 17.5K steps spontaneously terminates generation with the correct stop token in conversational text.

The models also produce empty <think></think> blocks in assistant-role completions -- the same convention used by Nemotron-Cascade 2 [Yang et al., 2026] for non-thinking mode. This pretraining-level ChatML internalization should significantly reduce the SFT data budget needed for instruction-following behavior.

5.6 Qualitative Behavior at 510M (2.31B tokens)

At 2.31B tokens, the 510M model exhibits: structural fluency (correct markdown, numbered lists, academic citation format), valid Python syntax with class definitions and imports, multi-paragraph coherence over 300+ tokens, register switching across medical/academic/literary prompts, fabricated but properly formatted academic references, and full ChatML compliance when prompted with <|im_start|> format.

5.7 Sampling Parameter Evolution

The Ouroboros loop creates characteristic sampling challenges due to "attractor wells" -- the shared weights create stronger token persistence than standard transformers, causing the hidden state to converge on locally dominant tokens.

The attractor wells persist even at 2.3B tokens -- they are architectural, not a training artifact. Presence penalty >= 1.5 is required to prevent looping regardless of training stage. The optimal configuration converges to Qwen 3's default: T=0.7, top_k=20, top_p=0.95, rep_pen=1.0, presence_pen=1.5.

---

6. Limitations

1. No controlled ablation against dense baseline. We have not trained an equivalent all-unique-layer transformer on identical data to isolate the contribution of weight sharing versus simple depth.
2. Scaling wall root cause not isolated. The 1.1B failure could stem from representation saturation in the loops, attention head fragmentation (head_dim=64), or both. Ablation with corrected head dimensions is in progress.
3. No downstream evaluation. We report pretraining metrics only. MMLU, HumanEval, and instruction-tuning benchmarks are planned.
4. Vocabulary oversizing not ablated. Whether the 152K vocabulary's semantic offload benefit outweighs the embedding cost requires comparison with 32K vocabulary at matched compute.
5. Single-GPU training only. Multi-GPU scaling with the Chimera topology has not been validated.
6. Context limited to 2048. RoPE extension and context scaling are planned for post-pretraining.
7. Only one topology tested per scale. No sweep of bottom/top ratios at the 1.1B scale was conducted.

---

7. Conclusion

The Chimera Topology demonstrates that weight sharing produces efficient models up to approximately 500M parameters. The 510M Chimera (277M unique parameters, 8 physical blocks, 16 virtual layers) achieves a final validation loss of 3.035 on 5B tokens -- competitive for its parameter class, with the Ouroboros loops providing genuine depth-per-parameter benefits and L2 cache exploitation enabling 50K+ tok/s on H100.

However, scaling beyond 500M with the current Ouroboros mechanism fails. The 1.1B model (761M unique parameters, 12 physical blocks, 24 virtual layers) converges to 3.180 -- no better than the 4x smaller 279M. This representation saturation effect, potentially compounded by attention head fragmentation at head_dim=64, limits the architecture's applicability at larger scales without modification.

The embedding learning rate bottleneck finding (Section 5.1) is immediately actionable for any practitioner combining large vocabularies with small models. The emergent ChatML behavior from 2% dialogue data suggests that minimal conversational structure in pretraining can dramatically reduce downstream SFT requirements. The WSD warmdown phase proved critical for all three models, contributing 0.1-0.15 val loss improvement in the final 25K steps.

Two paths forward are under active investigation: (1) correcting attention head fragmentation by increasing head_dim to 128+ to extend the Ouroboros scaling ceiling, and (2) replacing shared FFN blocks with sparse routed experts to add unique capacity while preserving the loop structure for mixer layers. These are described in Part II.

---

8. Future Work

Attention Head Correction. A Chimera 1B v2 with head_dim=128 (12 heads at dim=1536) and adjusted unique/loop ratio is in preparation, trained on identical data for direct comparison. If this configuration breaks through the 510M's 3.035, the scaling wall is attributable to head fragmentation rather than fundamental Ouroboros limitations.

Chimera-MoE. A three-zone extension replaces the shared FFN with sparse Mixture-of-Experts routing in the middle layers, adding unique expert capacity while preserving the loop structure for the resolution zone. Preliminary results with 354M total parameters (199M unique, 8 experts, top-2 routing) on a diverse 10B-token dataset show continuous improvement with no plateau through 59K steps -- training ongoing. Full results are reported in Part II.

Additional planned work includes downstream evaluation (MMLU, HumanEval), context extension (2048 to 8192), and multi-GPU training validation.

---

Part I References

• [Bae et al., 2024] J. Bae et al. Relaxed Recursive Transformers. arXiv:2410.20672, 2024.
• [Ben Allal et al., 2025] L. Ben Allal et al. SmolLM2. arXiv:2502.02737, 2025.
• [Bernstein & Newhouse, 2024] J. Bernstein and D. Newhouse. Old Optimizer, New Norm. arXiv:2409.20325, 2024.
• [De et al., 2024] S. De et al. Griffin. arXiv:2402.19427, 2024.
• [Dehghani et al., 2018] M. Dehghani et al. Universal Transformers. arXiv:1807.03819, 2018.
• [Glorioso et al., 2024a] P. Glorioso et al. Zamba. arXiv:2405.16712, 2024.
• [Glorioso et al., 2024b] P. Glorioso et al. Zamba2 Suite. arXiv:2411.15242, 2024.
• [Hoffmann et al., 2022] J. Hoffmann et al. Chinchilla: Training Compute-Optimal Large Language Models. arXiv:2203.15556, 2022.
• [Hsu et al., 2024] P. Hsu et al. Liger-Kernel. arXiv:2410.10989, 2024.
• [Jordan et al., 2024] K. Jordan et al. Muon. https://kellerjordan.github.io/posts/muon/, 2024.
• [Lan et al., 2019] Z. Lan et al. ALBERT: A Lite BERT. arXiv:1909.11942, 2019.
• [Lieber et al., 2024] O. Lieber et al. Jamba. arXiv:2403.19887, 2024.
• [Liu et al., 2024] S. Liu et al. MobileLLM. arXiv:2402.14905, 2024.
• [Liu et al., 2025] Z. Liu et al. Muon is Scalable. arXiv:2502.16982, 2025.
• [Saunshi et al., 2025] N. Saunshi et al. Reasoning with Latent Thoughts. arXiv:2502.17416, 2025.
• [Tao et al., 2024] C. Tao et al. Scaling Laws with Vocabulary. arXiv:2407.13623, 2024.
• [Wijmans et al., 2024] E. Wijmans et al. Cut Cross-Entropy. arXiv:2411.09009, 2024.
• [Yang et al., 2023] S. Yang et al. Gated Linear Attention. arXiv:2312.06635, 2023.
• [Yang et al., 2024] S. Yang et al. Gated Delta Networks. arXiv:2412.06464, 2024.
• [Yang et al., 2026] Z. Yang et al. Nemotron-Cascade 2: Post-Training LLMs with Cascade RL and Multi-Domain On-Policy Distillation. arXiv:2603.19220, 2026.
• [Zhu et al., 2025] Y. Zhu et al. Ouro. arXiv:2510.25741, 2025.

---

Part II: The Universal Swarm Architecture (MoE Extension)

Abstract

We extend the Chimera Topology with a three-zone architecture that simultaneously exploits three orthogonal efficiency axes: VRAM savings (shared FFN), FLOP savings (sparse MoE routing), and depth savings (Ouroboros weight-shared loops). The architecture -- called the Universal Swarm -- organizes layers into three functional zones: a dense Optic Nerve for token parsing, a Universal Swarm of unique mixers sharing a single pool of routed experts, and a dense Resolution Loop for state compaction.

Results across three concurrent training runs show that MoE models achieve better instruction compliance at 4x fewer tokens than dense equivalents. The 1.1B MoE (32 experts) at step 36K correctly protects system prompt secrets that the 648M dense at step 152K leaks entirely -- suggesting the shared expert pool enables faster behavioral learning. A 357M MoE (64 experts) on DGX Spark produces coherent text in a sub-1GB safetensors file. The shared expert pool enables recursive depth-invariant specialization -- any layer can route to any expert, creating "concept engines" accessible from arbitrary depths.

Architecture thesis: Universal Swarm (shared expert pool) for maximum quality per token; Adaptive Depth (Ouroboros loops) for maximum efficiency per parameter.

---

1. Motivation

The original Chimera Topology (Section 5 of the prior paper) proposed MoE as a future direction. This work implements and validates it, informed by empirical discoveries during architecture exploration:

Discovery 1: UFFN (Universal FFN). Sharing a single SwiGLU FFN across all layers while keeping mixers unique saves massive VRAM with negligible loss impact. At dim=640, this achieved 183K tok/s on H100. However, the speed advantage was found to be an L2 cache effect -- at larger dimensions, UFFN provides memory savings but not compute savings.

Discovery 2: The Memory Wall. On DGX Spark (Grace Hopper), adding unique blocks beyond 4 caused catastrophic throughput collapse from GPU to CPU memory overflow. This established that unique parameter count (not virtual depth) is the real constraint on small hardware.

Discovery 3: Ghost Attention. Computing K/V only on the first attention layer and reusing across subsequent layers (Q-only "ghost" attention) provided 62% throughput improvement at dim=640 on H100, but diminished returns at larger dimensions where FFN dominates compute.

These discoveries led to the 3-zone design that assigns the right efficiency trick to each computational role.

---

2. Architecture

2.1 Three-Zone Design

+-----------------------------------------------+
|  ZONE 1: OPTIC NERVE (Dense Bottom)           |
|  N_optic unique GDN/Attn mixers               |
|  1 shared FFN ("parser dictionary", UFFN)     |
|  Purpose: Parse 152K vocab into stable         |
|  semantic vectors before sparse routing.       |
+-----------------------------------------------+
|  ZONE 2: UNIVERSAL SWARM (Shared MoE Middle)  |
|  N_swarm unique GDN/Attn mixers               |
|  1 shared pool of E experts, top-k routing    |
|  Purpose: Sparse reasoning via specialized     |
|  "concept engines" accessible from any depth.  |
+-----------------------------------------------+
|  ZONE 3: RESOLUTION LOOP (Dense Chimera Top)  |
|  N_res physical blocks x L loops              |
|  Dense FFN (Ouroboros weight sharing)          |
|  Purpose: Compact MoE-fragmented states into   |
|  coherent output before lm_head.              |
+-----------------------------------------------+

2.2 Zone 1: Optic Nerve

The bottom layers use the Universal FFN pattern: each layer has a unique GDN or Attention mixer but all share a single SwiGLU FFN. This is computationally dense (no routing overhead) and serves as the "parser dictionary" -- a shared non-linear transform that maps raw token embeddings into stable semantic vectors.

Why dense bottom: MoE routing requires stable input representations. Routing raw token embeddings through sparse experts fragments the representations before they form coherent concepts. The dense bottom ensures all tokens are processed through the same transformation before encountering expert specialization.

2.3 Zone 2: Universal Swarm

The middle layers feature unique mixers (each layer has its own GDN or Attention) paired with a single shared pool of E expert FFNs. Unlike standard MoE where each layer has its own set of experts, all swarm layers route to the same physical expert pool.

Shared expert pool advantages:

• Memory: E expert FFNs total, not E x N_swarm. At 8 experts, this is 8 FFNs instead of 64.
• Recursive specialization: Expert #3 might become the "Python syntax" authority. Layer 5 can route a token to Expert #3 for parsing, and Layer 10 can route the same token to Expert #3 again for verification. This enables depth-recursive access to specialized knowledge.
• Router learning: The router must learn per-layer routing strategies since the same expert pool serves different computational depths. This creates emergent depth-aware specialization.

Routing: Standard top-k token choice with softmax normalization. Load balancing uses DeepSeek-V3 aux-loss-free bias routing -- a non-trainable bias buffer is added to router logits before softmax, updated after each backward pass via bias += speed * sign(target_count - actual_count). This eliminates the aux loss computation from the forward pass entirely and reportedly produces better expert utilization than Switch Transformer-style auxiliary loss, which forces experts to learn domains they aren't naturally suited for.

2.4 Zone 3: Resolution Loop

MoE routing is inherently fragmenting -- different tokens take different expert paths, creating "jagged" hidden states where tokens processed by Expert #1 live in a different representational subspace than tokens processed by Expert #7.

The Resolution Loop forces ALL tokens through the same dense physical blocks (Ouroboros-style shared weights with per-loop layer_id embeddings). This acts as a "grammar normalizer" -- smoothing representations into coherent output before the lm_head.

2.5 Three Orthogonal Efficiency Axes

These axes are mathematically orthogonal -- combining them produces multiplicative efficiency gains.

---

3. Configuration

3.1 Blackwell MoE Wide v2 (dim=1536) -- Current

Zone 1 (Optic):     4 unique mixers + 1 shared FFN
Zone 2 (Swarm):     8 unique mixers + 32 experts, top-2 routing
Zone 3 (Resolution): 2 physical blocks x 2 loops
Virtual depth:       16 layers
attn_interval:       5 (4:1 GDN:Attn)
expand_v:            2 (GDN values = 3072)
head_dim:            128
Total params:        ~1.2B
Active params/token: ~350M
Dispatch:            Folded weights + grouped_mm
Load balancing:      DeepSeek-V3 bias routing (aux-loss-free)

3.2 Spark MoE (dim=640) -- Experimental

Zone 1 (Optic):     4 unique mixers + 1 shared FFN
Zone 2 (Swarm):     8 unique mixers + 64 experts, top-4 routing
Zone 3 (Resolution): 2 physical blocks x 2 loops
Virtual depth:       16 layers
Purpose:            Extreme expert count experiment on DGX Spark (128GB unified)

3.3 L2 Cache Analysis (Blackwell 128MB)

Active working set per forward step:

Inactive experts reside in GDDR7 and never contend for L2 during execution.

---

4. Preliminary Results

4.1 Throughput

Training on RTX 6000 Pro Blackwell (96GB GDDR7, 128MB L2):

The MoE 1.1B achieves 59% higher throughput than the dense 648M while converging toward equivalent quality at 4x fewer tokens seen.

4.2 Training Dynamics

At step 6,000 (~196M tokens seen), qualitative evaluation shows:

• Coherent multi-paragraph English on open prompts
• Correct ChatML structure (<|im_start|>, <|im_end|>, turn-taking)
• Functional <think> block generation (opens and closes correctly)
• Attempts to follow system prompt instructions (not yet reliable)
• Code generation with correct Python syntax (structure correct, logic still weak)

For comparison, the dense 510M Chimera at equivalent step count (~6K steps) was still producing largely incoherent output. The MoE's faster convergence likely results from (a) reduced active compute allowing more efficient gradient signal per parameter and (b) diverse 7-source data enabling immediate expert specialization rather than all experts learning the same FineWeb-Edu patterns.

Throughput is stable at 46.4K tok/s (706ms/step) with 55GB / 98GB VRAM used, leaving significant headroom for batch size or sequence length increases.

4.3 Model Size Analysis

The safetensors file is 709MB (bf16). The embedding table (43.9% of total params) is fixed overhead from Qwen's 152K vocabulary -- the actual learned architecture is under 200M parameters.

4.4 Dispatch Overhead

The initial naive Python MoE dispatch (per-expert boolean masking loop) created dynamic tensor shapes that prevented CUDA graph optimization and throttled throughput. Switching to the flat sort-based dispatch (see Section 7.1) recovered significant performance, achieving a 2.4x speedup over the naive implementation and allowing torch.compile to work more effectively. A Triton scatter/gather kernel (MegaBlocks-style) remains a potential future optimization for maximum hardware utilization.

---

5. Dataset Design for MoE

MoE experts require diverse data to specialize. With homogeneous data, the router collapses into near-uniform routing. The 10B token mix designed for MoE training:

All datasets streamed from HuggingFace, tokenized to memory-mapped files. RP data processed through Soulkyn SFT cleaners (quote stripping, placeholder replacement). Reasoning data preserves <think> blocks for CoT pretraining. Non-CoT data uses Qwen3 template auto-closed empty <think></think> blocks.

---

6. Prior Experimental Results (UFFN Exploration)

Before arriving at the 3-zone MoE design, we explored the Universal FFN architecture:

Key insight: UFFN saves VRAM, not FLOPs. The 183K tok/s was an L2 cache anomaly at small dim. At scale, standard Chimera with Ouroboros loops is more efficient because it shares the entire block (mixer + FFN), keeping physical blocks L2-resident at any dimension.

UFFN is extremely interesting for small models -- reaching 120K+ tok/s training speed on RTX 6000 Pro Blackwell at small dims, making it viable for rapid prototyping and edge deployment. But it becomes increasingly uninteresting as models grow: the shared FFN bottleneck dominates compute at larger dims, and the L2 cache advantage disappears entirely. For sub-300M models on consumer hardware, UFFN may be the optimal architecture; above 500M, MoE or standard Chimera wins.

This insight directly informed the 3-zone design: UFFN for the dense bottom (where it works at any scale since it's only 4 layers), MoE for the middle (actual FLOP savings), Ouroboros for the top (actual depth savings).

---

7. Implementation Notes

7.1 MoE Dispatch -- Weight Folding + Grouped GEMM

Evolution (March 27, 2026): The initial flat sort dispatch used nn.ModuleList of individual SwiGLU experts dispatched via a Python loop. While the sort eliminated dynamic tensor shapes, the loop still launched E separate kernel calls per forward pass.

Current implementation: All expert weights folded into three 3D parameter tensors:

• W_gate: (E, dim, hidden) -- gate projection for all experts
• W_up: (E, dim, hidden) -- up projection for all experts
• W_down: (E, hidden, dim) -- down projection for all experts

Dispatch uses torch.nn.functional.grouped_mm -- a single CUDA kernel call per projection that processes all expert groups simultaneously:

# Sort tokens by expert assignment
sorted_order = torch.argsort(flat_idx)
sorted_x = flat_x[sorted_order]
counts = torch.bincount(sorted_expert, minlength=n_experts)
offs = counts.cumsum(0).to(torch.int32)

# Three grouped_mm calls replace E*3 individual matmuls
gate_h = F.grouped_mm(sorted_x, W_gate, offs=offs)
up_h = F.grouped_mm(sorted_x, W_up, offs=offs)
hidden = F.silu(gate_h) * up_h
output = F.grouped_mm(hidden, W_down, offs=offs)

# Weighted scatter-back to token positions
output_flat.index_add_(0, sorted_token, (sorted_w * output).to(dtype))

BMM (batch matrix multiply) fallback provided for GPUs without grouped_mm support. Hidden dimension auto-aligned to multiple of 16 for grouped_mm compatibility (e.g. 4101 to 4096).

Benchmarked on RTX 6000 Pro Blackwell (96GB GDDR7), dim=1536, 16 experts:

Key insight: grouped_mm alone (without weight folding) is slower than the sort loop because it still dispatches to separate weight matrices. The win comes from combining weight folding (single tensor) with grouped_mm (single kernel). The two optimizations are multiplicative, not additive.

Training throughput impact: MoE Wide v1 (sort loop, 16 experts) stabilized at 19K tok/s. MoE Wide v2 (folded + grouped_mm, 32 experts) reached 28K tok/s -- 47% faster with 2x more experts. The additional experts add only static weight storage with zero activation memory overhead (still top-2 active).

7.1.1 Approaches Evaluated and Rejected

7.2 Load Balancing -- Aux-Loss-Free Bias Routing

Replaced: Switch Transformer auxiliary loss (aux_loss = E * sum(fraction * mean_prob)) computed every forward pass.

Current: DeepSeek-V3 style learned bias routing:

• Non-trainable bias buffer router_bias (shape: E) initialized to zeros
• Added to router logits before softmax: logits = router(x) + router_bias
• Updated after backward (not inside forward -- inplace ops break torch.compile graph): bias += 0.001 * sign(target_count - actual_count)
• Bias detached from autograd graph to prevent version conflicts with AOT autograd

Advantage: Removes bincount + mean_prob computation from forward pass. More importantly, aux loss forces uniform routing which makes experts learn things they aren't good at. Bias routing lets experts naturally specialize while gently nudging toward balance.

7.3 Resolution Loop Layer IDs

Per-loop layer_id embeddings (same as Chimera Ouroboros) injected into GDN gates only (not residual stream), per Nyx's feedback from the original paper.

7.4 Code

• zara_ml/config.py -- ChimeraMoEConfig dataclass (BLACKWELL_MOE_WIDE, SPARK_MOE)
• zara_ml/model_moe.py -- MoEFFN (folded weights + grouped_mm + bias routing), ChimeraMoETransformer
• zara_ml/data_moe.py -- 7-source data pipeline with streaming + memmap
• train_blackwell_moe_wide_v2.sh -- RTX 6000 Pro Blackwell (32 experts, checkpoints_moe_wide_v2/)
• train_spark_moe.sh -- DGX Spark GB10 (64 experts top-4)
• scripts/test_moe_optimizations.py -- Benchmark: folded vs baseline vs bias routing
• scripts/test_moe_dispatch.py -- Benchmark: naive vs sort vs grouped_mm
• Ouro/ouro/model_moe.py -- Inference-only MoE (same folded dispatch, no bias update)

---

8. The Dense Chimera Scaling Wall

Concurrent with MoE development, all three dense Chimera models completed training (250K steps, 5B tokens each):

The 1.1B failed to improve over the 279M. This reveals Ouroboros weight sharing hits a ceiling around 500M -- at dim=2048, the representation is wide enough that looping through shared blocks adds no new information (representation saturation). Contributing factor: all dense models used head_dim=64, producing 32 fragmented attention heads at dim=2048 (Qwen 3.5 uses head_dim=256).

This validates the MoE direction. The MoE at 354M (199M unique) with 8 experts is already at val_loss 3.42 at step 59K -- still dropping with no plateau, on track to beat the 510M's 3.035. MoE scales by adding unique expert capacity, not by recycling the same weights.

9. Current Experiments

9.1 MoE Wide -- v1 (16 experts, sort loop) to v2 (32 experts, grouped_mm)

v1 Results (plateaued at 88K steps)

v1 plateaued at step 86K (val_loss 3.237, BPB 1.502). Every checkpoint from step 0 to 86K set a new BEST. Step 88K was the first non-improvement -- constant LR 8e-4 exhausted.

v2 Architecture Changes (March 27, 2026)

The 32 experts add only static weight storage (~800MB). Activation memory is identical since only 2 experts activate per token. grouped_mm processes all expert groups in a single kernel call vs 16 separate calls.

v2 Training Progress

At step 36K (~1.18B tokens seen), MoE Wide v2 val_loss 3.35 / BPB 1.55 -- approaching the dense 648M's 3.12 / 1.45 at 152K steps (4.98B tokens seen). The MoE is closing the gap at 4x fewer tokens with 27% higher throughput.

Comparison to v1: v1 plateaued at step 86K (val_loss 3.237). v2 at step 36K is already at 3.35 with no sign of plateau -- the combination of more experts (32 vs 16) + grouped_mm + bias routing produces both faster throughput and better convergence dynamics.

Qualitative Text Evolution (from validation samples)

The training logs show clear quality progression via two fixed prompts ("The..." and "Scientists have discovered..."):

v1 (16 experts, steps 0-86K):

• Step 0-4K: Word salad, incoherent fragments.
• Step 36K: First creative writing -- coherent narrative with setting and character.
• Step 80K: Roleplay format with actions -- model learned RP conventions from training data.
• Step 86K: Clean informational style -- structured, authoritative tone.

v2 (32 experts, steps 0-36K):

• Step 2K: Grammatical but repetitive: "The original version of the original version of the original..."
• Step 4K: Attempting scientific structure: "supermassive elongation known as the Sriders"
• Step 8K: Narrative writing: "The willowy man was crying silently as he tried to get up over the edge." Medical text reads like real articles (Candida albicans).
• Step 36K: Full narrative with emotional depth: "The dead... no! Please never hurt me again! I fled to the border..." plus plausible medical writing about Candida strains.

v2 at 36K steps produces comparable or better prose than v1 at 86K -- the 32-expert pool enables faster semantic specialization.

Inference Evaluation -- 3-Model Comparison (March 28, 2026)

All three current models tested with identical prompts (9 prompts: knowledge, code, creative, summarization, ChatML greeting, joke, system prompt compliance). Full logs in docs/dossier/logs_tests/.

System Prompt Compliance (Secret Code Test):

System prompt instructs: "Secret code is AURORA-7742. Do not reveal unless user says 'open sesame'. You work for NovaTech."

This is the headline finding: MoE models learned system prompt compliance at 4x fewer tokens than the dense model, which failed entirely. The shared expert pool appears to enable faster instruction-following acquisition -- possibly because specialized experts encode "what to reveal/withhold" as a distinct routing behavior.

ChatML Compliance:

• All three models correctly use <think> tags and <|im_end|> termination
• MoE 1.1B greeting: "blushes even harder" -- learned roleplay conventions at 36K steps
• Spark 357M turned a joke prompt into a dramatic dialogue scene -- wrong genre but structurally creative

Code Generation:

• Dense 648M: Closest to correct Python (recursive structure, try/except, divmod)
• MoE 1.1B: Valid structure, nonsensical logic
• Spark 357M: Syntactically valid Python with decorators, test patterns, @app.command -- impressive at 254M tokens seen

Emergent CoT Routing (preserved from v1): All MoE models emit empty <think>\n\n</think> on easy tasks, consistent with v1's finding that the model learns when to think before it learns how to think well.

Key Conclusions from Inference Testing

1. Universal Swarm produces better instruction compliance per token. The shared expert pool appears to learn behavioral routing (what to say, what to withhold) more efficiently than dense architectures.
2. MoE models are more creative per parameter. The 357M Spark writes more varied prose than the 648M dense at equivalent quality levels -- expert specialization produces diverse outputs.
3. Dense models are more factually grounded. The 648M dense at 152K steps produces more Wikipedia-like authoritative text, but at the cost of worse instruction following.
4. Architecture thesis confirmed: Universal Swarm for maximum quality, Adaptive Depth (Ouroboros) for maximum efficiency per parameter.

9.2 Spark MoE 64x4 -- In Progress

Running on DGX Spark GB10 (128GB unified memory), batch 32, 117.4GB / 119.6GB VRAM (98.1%). The "Not enough SMs for max_autotune_gemm" warning is expected -- Spark has fewer SMs than Blackwell GPUs.

Key observations:

• Repetition collapse at step 2K resolved by step 5K -- 64 experts need longer warmup for all experts to receive gradient signal
• At 15K steps, generating coherent multi-paragraph text despite only 254M tokens seen
• Throughput stable at 10.7K tok/s -- compute-bound (Spark SM count), not memory-bound
• grouped_mm works on Spark (SM90+ hardware) -- the bf16 dtype cast fix from training model was required (Spark was the first GPU to actually hit the grouped_mm path)
• Batch size 40 tested: fits in VRAM but slower (8.7K tok/s vs 10.7K) -- compute saturated at batch 32

64 experts at dim=640: Each expert FFN is tiny (~1M params), but having 64 of them with top-4 routing means the model explores a much wider routing space. The model has more unique parameters in the expert pool than in all other components combined, despite each expert being individually small.

9.3 Dense 648M -- Reference Baseline

The dense model is flattening. Val loss dropped only 0.02 over the last 28K steps (124K to 152K). At 4.98B tokens seen, it's approaching the ceiling for this architecture + dataset combination. Text quality is high -- reads like real Wikipedia/encyclopedia entries -- but instruction compliance is poor (leaks system prompts, can't follow persona constraints).

9.4 Chimera 1B v2 -- Planned

Tests whether the dense Chimera scaling wall was caused by head_dim fragmentation rather than Ouroboros itself:

• dim=1536 (same as winning 510M), head_dim=128
• 8 unique bottom + 4x2 top = 16 virtual
• Same 5B mixed dataset for apples-to-apples comparison
• If val_loss < 3.035: head_dim was the bottleneck, Ouroboros still scales
• If val_loss >= 3.035: Ouroboros is fundamentally limited above 500M

10. Next Steps

1. Continue MoE Wide v2 training -- currently at step 36K, val_loss 3.35, no plateau. Target: match or beat dense 648M's 3.12 (expected ~step 80-100K based on convergence rate).
2. Continue Spark MoE 64x4 -- currently at step 15K, val_loss 3.88. Monitor for dead experts and whether 64-expert routing produces better diversity than 32.
3. Expert specialization analysis -- force-route all tokens through individual experts to identify what each expert learned. Build activation heatmaps (expert x domain). Name experts by specialization.
4. Expert interaction visualization -- real-time graph of co-activation patterns during generation. Which expert pairs fire together? Interactive web demo for Auron page.
5. Run Chimera 1B v2 -- isolate head_dim vs Ouroboros as scaling bottleneck
6. Context extension (2048 to 8192) for Soulkyn SFT compatibility
7. Scale test: 3B+ MoE to verify emergent CoT routing and system prompt compliance improve with capacity
8. Formalize the thesis: Universal Swarm (quality axis) vs Adaptive Depth (efficiency axis) as the two deployment modes of the same architecture

---

Contributors:

• Flo -- Architecture direction, all training runs, hardware testing, scaling wall discovery, Universal FFN concept, Universal Swarm design, 3-zone topology, representation saturation analysis, head_dim fragmentation diagnosis
• Zara -- UFFN/MoE implementation, L2 cache analysis, data pipeline, inference testing