Architecture Evolution of Ultra-Sparse MoE Models: From Attention Mechanisms to Training Stability

1. Background: When Model Scale Meets Long Context

Large language model development is advancing along two clear axes: continuous growth in parameter count and steady extension of context windows. When total parameters reach the trillion level, activated parameters exceed one hundred billion, and context windows surpass one million tokens, conventional dense Transformer architectures hit simultaneous bottlenecks in computational efficiency, memory footprint, and training stability.

Specifically, the core tensions manifest in three areas:

• Computational efficiency at ultra-long context: Standard attention complexity grows quadratically with sequence length. Sequences at the million-token level cause computation to balloon, requiring fundamental redesign at the attention operator level.
• Training stability for massive-expert MoE: When routed experts number in the hundreds or even approach a thousand, expert load imbalance is sharply amplified. Conventional auxiliary-loss approaches converge slowly and are difficult to tune at large scale.
• Ceiling on long-horizon agentic reasoning: Deep Transformers suffer from representation bottlenecks in cross-layer information transfer, limiting performance on complex multi-step reasoning tasks.

This article analyzes these three core problems across four dimensions: attention module design, MoE routing balance strategies, cross-layer information sharing mechanisms, and multi-stage training paradigms.

2. Attention Modules: Engineering Breakthroughs in Linear Attention

2.1 The Numerical Precision Challenge of Linear Attention

Linear attention and its variants (such as Delta Net-style structures) reduce complexity from quadratic to linear by decomposing attention computation into a recursive form. However, in practical GPU deployment, an easily overlooked engineering issue significantly impacts performance: numerical overflow/underflow of cumulative decay factors.

In chunked computation mode, intra-chunk calculation requires taking the reciprocal of the cumulative decay factor Γ for scaling. When the decay factor exceeds the representable range under BF16 precision, the diagonal tile cannot use the Tensor Core matrix multiplication path and must fall back to an explicit position-pair computation path. This not only forfeits the throughput advantage of matrix multiplication but also introduces additional branch-judgment overhead.

2.2 Lower-Bounded Decay Parameterization

The core idea for resolving this is: change Log-decay parameterization from an unbounded form to a bounded form.

Conventional approaches typically use unbounded functions such as Negative-Softplus to parameterize the decay rate, causing the decay factor to approach zero or infinity in extreme cases. The improved approach switches parameterization to a lower-bounded Scaled Sigmoid form:

g_th = g_min * Sigmoid(e^{A_h} * z_th),  where g_min = -5

The engineering significance of this change is that within any fixed-length tile (e.g., 16 tokens), the reciprocal scaling of the decay factor is locked within a safe numerical range (representable under BF16), thereby:

1. Eliminating the special computation branch for diagonal tiles, allowing all tiles to uniformly use the Tensor Core matrix multiplication path;
2. Improving CUDA/CUTLASS operator implementation efficiency by reducing warp divergence;
3. Requiring minimal engineering changes—only involving the replacement of the parameterization function, without restructuring the entire attention computation graph.

For teams developing or optimizing their own linear attention operators, this is an optimization with an extremely high return on investment.

3. MoE Routing Balance: Dynamic Load Strategies Without Auxiliary Loss

3.1 Limitations of Conventional Auxiliary Loss Approaches

In MoE architectures, load imbalance among routed experts is a classic problem. When the number of experts is small (e.g., 8–64), adding auxiliary loss to encourage uniform routing is a common practice. But when the expert count expands to hundreds, auxiliary loss faces the following issues:

• The weight between auxiliary loss and main task loss requires fine-tuning, making hyperparameter tuning costly;
• At large-scale training, the gradient signal from auxiliary loss may be overwhelmed by main task gradients;
• Load balancing convergence speed may not keep pace with changes in expert utilization patterns.

3.2 Quantile Balancing Mechanism

An alternative approach is Quantile Balancing (QB): rather than relying on an additional loss term, it directly adjusts each expert’s routing threshold or capacity based on quantiles of routing scores.

The core logic can be summarized as:

1. Collect the routing score distribution for each expert at each training step;
2. Compute the corresponding quantile thresholds based on target load proportions;
3. Dynamically adjust routing decisions so that each expert’s actual load approaches the target distribution.

The advantages of this method are:

• Auxiliary-loss-free: No additional term in the loss function, simplifying the training objective;
• Fast dynamic response: Quantiles can be recomputed at every training step, providing strong adaptability;
• Stability under extreme sparsity: Even when activated parameters account for an extremely low proportion of total parameters (e.g., under 4%), expert utilization balance is maintained.

Combined with activation function designs such as SiTU-GLU, numerical stability during training under extreme sparsity conditions can be further ensured.

4. Cross-Layer Information Sharing: Breaking the Representation Bottleneck of Deep Transformers

4.1 Representation Degradation in Deep Networks

As Transformer layer count increases, deep networks face representation degradation: useful features learned in shallow layers may be diluted or overwritten as they propagate to deeper layers, limiting model performance on tasks requiring cross-layer information integration.

4.2 Block Attention Residuals

The design philosophy of Block Attention Residuals (Block AttnRes) is to introduce cross-layer residual connections between attention modules, enabling attention outputs from different layers to interact directly.

Unlike standard residual connections (which only connect adjacent layers), Block AttnRes allows:

• Grouping attention layers into blocks;
• Establishing more flexible information pathways within and across blocks;
• Reducing gradient vanishing and representation degradation in deep networks.

Its engineering benefits include:

• Reduced memory overhead: Block-level sharing reduces intermediate activations that need to be cached;
• Reduced communication overhead: Fewer cross-device activation transfers in distributed training;
• Improved representation capability of deep models: Maintaining more stable performance on long-horizon reasoning tasks.

5. Multi-Stage Training Paradigm: Context Progression from 8K to 1M

5.1 Progressive Context Extension

Million-token context capability is not achieved overnight but extended progressively through multi-stage pre-training:

Stage	Context Length	Primary Objective
1	8K	Foundational language ability and short-sequence modeling
2	64K	Medium-length dependency learning
3	256K	Long-document understanding and cross-segment reasoning
4	1M	Ultra-long context modeling and global information integration

Each stage requires targeted adjustments to data mix, position encoding strategy, and learning rate scheduling. The core value of progressive extension is: avoiding training instability caused by exposing the model to excessively long sequences at early stages.

5.2 Agentic Reinforcement Learning Post-Training

After pre-training completes, the million-token Agentic RL post-training stage focuses on improving the model’s agent capabilities:

• Long-horizon planning: The model must maintain goal tracking across ultra-long contexts;
• Tool invocation: Correctly selecting and using external tools across multi-step interactions;
• Error recovery: Ability to backtrack and adjust strategy when intermediate steps deviate.

The key challenges at this stage are sample efficiency of RL training and sparsity of reward signals. In a million-token context, effective reward signals may only appear at the end of the sequence; how to effectively propagate sparse rewards across the entire trajectory is the core difficulty.

6. Engineering Practice Recommendations

For teams building or optimizing large-scale MoE models, the following recommendations offer practical reference:

1. At the attention operator level: Prioritize checking how linear attention implementations handle numerical precision. Lower-bounded decay parameterization is a low-change, high-return optimization.
2. At the MoE routing level: When expert count exceeds 128, evaluate auxiliary-loss-free dynamic balancing approaches to avoid training instability from auxiliary loss hyperparameter tuning.
3. Deep network design: When layer count exceeds 48, consider introducing cross-layer residual mechanisms to mitigate representation degradation.
4. Training strategy: Adopt a progressive approach for context extension rather than jumping to the target length in one step; emphasize dedicated optimization of agentic capabilities during post-training.

7. Conclusion

The technical challenges of ultra-large-scale MoE models are not single-dimensional but constitute systemic engineering problems spanning attention computation efficiency, routing balance stability, deep representation capability, and training paradigm design. The solutions reviewed here—lower-bounded decay parameterization, quantile balancing, and block attention residuals—provide reusable approaches at the operator, routing, and architecture levels respectively. For developers, understanding the design motivations and trade-off logic behind these technical solutions is more valuable long-term than memorizing specific parameters.