Architectural Trade-offs in Million-Token MoE Models: From Linear Attention Stability to Sparse Expert Balancing

When model parameters reach the trillion level and context windows expand to the million-token scale, the focus of architecture design shifts from “capability stacking” to “system-level feasibility.” A straightforward example is the attention module: standard multi-head self attention runs stably at an 8K context length. However, when the sequence length grows by two orders of magnitude and the number of experts approaches a thousand, the training bottleneck is no longer algorithmic expressiveness, but rather engineering foundations like numerical precision, operator scheduling, and load distribution.

Based on the public technical roadmap of a class of ultra-long-context multimodal MoE models, this article analyzes three highly reusable architectural designs: Lower-Bound Decay Parameterization for Linear Attention, Auxiliary-Loss-Free Quantile Balancing Routing, and Cross-Layer Block Attention Residuals. Each addresses a real system bottleneck and can be transferred to other large-scale training scenarios.

---

1. Overall Architecture Profile: Why “Hybrid” Instead of “Replacement”

The overall profile of this model class can be summarized as a sparse MoE backbone with approximately 2.8 trillion total parameters and around 104 billion activated parameters per token. It natively supports visual inputs and million-token contexts. On the attention side, it merges two paths: incremental linear attention and gated multi-head latent attention. On the MoE side, it employs a sparse structure with hundreds of routed experts.

There are two key design orientations developers should note.

First, attention is not an either-or choice, but a layered hybrid. The computational complexity of pure softmax attention grows quadratically with sequence length, making memory and compute overheads unmanageable at the million-token scale. Pure linear attention, while linear in complexity, suffers from inherent decay in long-range precise retrieval. Combining them by layer or module is essentially a segmented compromise between “retrieval precision” and “computational feasibility”: retaining full attention in local or specific layers to maintain key information recall, while using linear attention for long-range propagation in the rest.

Second, the sparsity of the MoE is pushed to quite an aggressive level. The 27x difference between total and activated parameters means each token only activates a tiny fraction of experts. Higher sparsity yields better parameter efficiency per unit of compute, but leaves less room for routing error—once load distribution becomes unbalanced, some experts will be overloaded while others barely update, leading to severely uneven training signals. This is exactly the problem the “quantile balancing” mechanism aims to solve.

The phased progression strategy for pre-training context also warrants specific explanation.

Phased Context Extension: 8K → 64K → 256K → 1M

Pre-training directly at the million-token length is practically unfeasible: attention-related memory usage, gradient accumulation scale, and data mixture ratios would spiral out of control. A more robust approach is a step-by-step progression:

1. Short-context phase: Builds foundational language and representation capabilities. Short sequences and high throughput make it suitable for large-scale data processing.
2. Medium-context phase: Introduces longer dependency samples, allowing position-related parameters to adapt gradually.
3. Long-context phase: Focuses on refining long-range retrieval and cross-segment reasoning, while introducing long-document and long-conversation data.
4. Ultra-long-context phase: Concludes with a relatively small amount of high-quality long-sequence data. Combined with reinforcement learning on long-horizon agent tasks in the post-training phase, it translates “ability to read long” into “ability to use long.”

The core benefit of this phased strategy is allowing the model to learn near its “current capability boundary” at each stage. This prevents wasting short samples on padding in ultra-long sequences, and avoids invalid gradients from long samples in early stages where model capability is insufficient.

---

2. Innovation 1: Lower-Bound Decay Parameterization for Linear Attention

This is the technique requiring the least engineering effort but yielding the most direct benefits among the three.

The Pain Point: Why Diagonal Blocks Must Take the “Slow Path”

Linear attention and Delta Net-like structures typically employ chunked (tile) computation during implementation. Intra-block computation requires introducing a cumulative decay factor Γ, and subsequent steps often use its reciprocal for normalization or rescaling. The problem is that the decay factor decreases monotonically with position; in long sequences, Γ becomes extremely small, and its reciprocal quickly overflows under finite precision formats like BF16.

The consequences are cascading: blocks on the diagonal (where query and key fall within the same chunk) cannot be computed using standard dense matrix multiplication and must degrade into an explicit per-position computation path. Non-diagonal blocks, not involving such extreme scaling, can still utilize dense GEMM on Tensor Cores.

This creates a split in implementation: within the same attention operator, some tiles take the high-speed matrix multiplication path, while others take the low-efficiency element-wise path. GPU kernels therefore require branch processing, dragging down occupancy and pipeline efficiency.

The Mechanism: Changing Unbounded Decay to Bounded Decay

The solution is straightforward—since the problem is “reciprocal overflow,” ensure through parameterization that the decay factor does not reach a dangerous range. Specifically, replace the unbounded negative Softplus form in the log domain with a scaled Sigmoid form featuring a lower bound. By fixing a decay lower bound (e.g., at a magnitude of -5), the reciprocal scaling factor for any fixed-size tile is locked within a safe numerical range:

Traditional:  log_decay = -softplus(x)          # Unbounded, approaches extreme min in long sequences
Improved:     log_decay = g_min * sigmoid(x)    # g_min is a fixed lower bound, reciprocal is bounded

The direct effect is that diagonal blocks no longer require special branching. All causal tiles (diagonal or non-diagonal) can uniformly invoke dense Tensor Core GEMM. Kernel implementation is significantly simplified, and throughput improvements stem from two sources: instruction efficiency from branch elimination, and better parallelism from a unified dense path.

Incidental Benefit: Long-Context Extrapolation Without Positional Encoding

Bounded decay also brings a structural benefit. Once decay behavior is constrained within a stable range, the model’s reliance on absolute positions is further weakened, naturally supporting long-context direct extrapolation under a No Positional Encoding (NoPE) setting. For scenarios requiring the inference of a 256K-trained model up to 1M tokens, this reduces the cost of additional position interpolation or recalibration.

Reuse Advice: If your system already uses linear attention or state-space-like structures, the code changes for this modification are typically minimal. However, two things need rigorous validation: the impact of the lower bound value on long-range recall, and the empirical error boundary of reciprocal calculations under BF16. It is recommended to compare the sensitivity of “lower bound ±1” in ablation studies.

---

3. Innovation 2: Auxiliary-Loss-Free Quantile Balancing Routing

The Pain Point: Load Imbalance in Ultra-Sparse MoE

Traditional MoE commonly uses auxiliary loss to encourage expert load balancing. However, auxiliary loss inherently conflicts with the main task objective: tune the coefficient too high, and language modeling quality drops; tune it too low, and load imbalance occurs. Under the setting of nearly a thousand experts and an extremely low activation ratio, this tug-of-war becomes more pronounced—a few experts are repeatedly selected, while the rest lack effective gradients for long periods. The number of parameters actually participating in training falls far below the nominal value.

The Mechanism: Replacing Soft Constraints with Quantile Thresholds

The concept of Quantile Balancing replaces “soft incentives” with “hard allocation constraints.” It performs quantile statistics on the routing scores of each expert and determines token assignment based on quantile thresholds, ensuring the number of tokens received by each expert statistically approaches the target quota. Since the allocation is determined by quantiles, balancing no longer relies on loss terms to “induce” it, allowing the auxiliary loss to be removed entirely.

The significance for training stability lies in:

• Eliminating objective conflicts: The main loss function solely handles language modeling and multimodal alignment, unpolluted by balancing terms.
• Dynamic adaptation: Quantiles update in real-time during training, keeping pace with routing distribution drift, unlike fixed coefficients that require repeated tuning.
• Greater stability under extreme sparsity: When the number of activated experts is very small, the impact of any single routing error is amplified. Hard constraints suppress long-tail imbalance better than soft incentives.

Reuse Advice: During implementation, pay attention to the update frequency and statistical window of the quantile estimation. A window that is too short will cause routing jitter, while one too long will cause sluggish responses. It is recommended to tune this alongside a routing temperature annealing strategy and monitor the entropy of the expert hit distribution during early training phases.

---

4. Innovation 3: Cross-Layer Block Attention Residuals

The Pain Point: Representation Bottlenecks in Deep Transformers

As layer depth increases, information in the residual stream is gradually “diluted.” Each layer only makes incremental modifications, making it difficult for deep networks to effectively retrieve shallow features, forming a deep representation bottleneck. A common mitigation is adding cross-layer connections, but the memory and communication overhead of full connections is unacceptable in large-scale distributed training.

The Block Attention Residual (Block AttnRes) approach limits cross-layer information sharing to the “block” granularity: grouping several layers into a block, maintaining shared residual pathways within the block, and connecting blocks sparsely. This breaks the purely sequential depth bottleneck while keeping memory and communication overheads within budget.

The key engineering trade-off is block size. A larger block means a longer shared pathway and more pronounced cross-layer benefits, but also higher peak memory per block; a smaller block is the opposite. In practice, block size is often chosen in alignment with pipeline parallelism partition boundaries, ensuring that residual sharing communication stays within the same pipeline stage to avoid cross-node synchronization.

---

5. Shared Methodology of the Three Techniques

Viewed together, these three techniques share a single design philosophy: rather than pursuing the limits of algorithmic expressiveness, they first remove hard constraints at the system level.

• Lower-bound decay resolves numerical precision constraints—ensuring operators can consistently take the dense path.
• Quantile balancing resolves training signal constraints—making gradient distribution controllable under ultra-sparse structures.
• Block residuals resolve memory and communication constraints—ensuring depth benefits are not consumed by overhead.

For developers, the transferable lesson is: before pushing model scale to the next order of magnitude, first identify the “hard bottleneck” of the current system—whether it is precision, load, or communication—and then make targeted modifications at the parameterization or scheduling level. A common characteristic of such modifications is that they require minimal code changes but necessitate solid ablation experiments to confirm boundary conditions. Specifically, the three hyperparameters—the lower bound value, the quantile window, and the block size—have sensitivities that directly determine whether the modification can converge stably in real training.

Building upon this foundation, adding phased context pre-training and post-training for long-horizon agent tasks translates architectural feasibility into actual long-context task performance. Architecture determines the upper limit, and training strategies determine whether that limit can be reached; neither is dispensable.