Training Ultra-Long-Context LLMs: Sparse Attention Stability, Loss-Free MoE Routing, and Agentic RL Post-Training

1. Background: The Dual Challenge of Ultra-Long Context and Sparse Expert Models

In recent years, large language models have continued to scale along two dimensions: context window length and parameter count. Million-token context capabilities allow models to process entire code repositories, lengthy meeting transcripts, or complete collections of research papers, while sparse mixture-of-experts (MoE) architectures support trillion-parameter models at a much lower inference cost. When these two technical directions are combined deeply, however, several deeper problems emerge.

The compute bottleneck of ultra-long context. Traditional self-attention scales quadratically with sequence length. Even with linear attention or chunked computation, in low-precision training the numerical overflow of recursive decay factors can degrade the computation graph into per-position operations, making it impossible to exploit matrix-multiplication acceleration units effectively.

The load balancing dilemma of massive expert counts. As the number of routed experts grows from dozens to hundreds or nearly a thousand, classic auxiliary-loss methods suffer from both gradient conflicts and lagging load adjustment. Expert load distributions can become extremely imbalanced, leaving some experts undertrained for long periods.

Sparse training signals in long-horizon agent reasoning. In agentic settings, the model must maintain goal consistency over extremely long trajectories, completing multi-step tool calls and environment interactions. Supervision is highly sparse, and conventional training methods struggle to propagate gradients across such long decision chains.

These problems intertwine: the longer the context, the more attention states must be carried across layers; the more experts, the harder it is to keep training stable; the longer the reasoning chain, the more carefully post-training must be designed. A complete solution has to address numerical stability, load balancing, and long-horizon credit assignment simultaneously.

2. KDA Lower-Bounded Decay Parameterization: Keeping Linear Attention on Tensor Cores

2.1 The Root of the Problem

Linear attention and Delta networks compress historical information into a recursive state, avoiding explicit storage of the attention matrix. In chunked parallel computation, each chunk requires the reciprocal of the cumulative decay factor. Let the decay parameter be γ: as γ approaches zero, its reciprocal overflows quickly in BF16 precision. Once values exceed BF16’s representable range, the framework has to abandon the matmul path and fall back on explicit per-position computation, which directly cuts training throughput by more than an order of magnitude.

Specifically, conventional approaches parameterize the log-decay with an unbounded negative Softplus. This means the decay factor has no guaranteed lower bound and can theoretically approach zero arbitrarily closely. After layer normalization in deep stacks, the decay parameters of certain dimensions can be pushed toward extremely small values, and the problem tends to erupt in the middle and later stages of training.

2.2 A Mathematical Fix with a Lower Bound

The core improvement replaces the unbounded parameterization of the log-decay with a scaled Sigmoid that has a fixed lower bound. Setting the lower bound to a conservative negative constant—for example, -5—locks the reciprocal scaling factor of any 16-token chunk into a predictable, safe interval. The change is mathematically simple but carries substantial engineering value: it guarantees that chunked diagonal computation remains usable throughout large-scale parallel training.

The key lies in choosing the bound. A bound set too high limits the model’s ability to forget, causing the recursive state to accumulate irrelevant information; a bound set too low fails to fully eliminate overflow risk. In practice, by profiling the distribution of decay parameters across attention heads at different depths, one can identify a lower-bound interval that balances expressiveness and numerical safety.

2.3 Implementation-Level Payoffs

In terms of tensor core utilization, this parameterization change raises the coverage of the chunked matmul path in the linear attention operator from less than half to nearly complete. Under a mixed-precision training framework, there is no longer any need to write inefficient per-element CUDA kernels specifically for diagonal blocks; the entire forward and backward pass can run through a unified, efficient matmul pipeline. For pretraining workloads involving ultra-long sequences, this change translates directly into a significant reduction in end-to-end training time.

3. Gated Multi-Head Latent Attention and Cross-Layer Attention Residuals

3.1 Why Gated Multi-Head Latent Attention

In long-context settings, the attention module must balance information retention against computational efficiency. Gated multi-head latent attention borrows the idea of low-rank compression: key-value pairs are projected into a low-dimensional latent space, state updates are performed in the latent dimension, and an output projection restores the original dimension. The gating mechanism gives the model fine-grained control over information flow—learning when to write new information, when to preserve the existing state, and when to perform a read.

This design complements KDA: KDA addresses the numerical stability of recursive decay, while gating addresses the semantic flexibility of information flow. Together, the attention module retains the efficient recurrence of linear attention while preserving the gated structure’s ability to respond sensitively to key events in a sequence.

3.2 Block-Level Attention Residuals: A New Channel for Cross-Layer Information Sharing

In a conventional Transformer, residual connections operate only within a single layer, and cross-layer information flow depends entirely on each layer’s output feeding the next. When networks reach hundreds of layers, deep representations tend to degrade: low-level features captured in shallow layers gradually blur through repeated nonlinear transformations, and deep layers struggle to maintain both local precision and global consistency.

The core idea of block-level attention residuals is to group several adjacent layers into a block and share a single attention-state residual within it. This yields two direct benefits. First, deeper layers can access the attention states of early layers in the block through the residual channel, mitigating deep representation degradation. Second, the attention state only needs to be computed and stored once per block, with each layer updating it through residual increments—substantially reducing memory footprint and cross-device communication. In distributed training, the reduction in communication is often more critical than the reduction in computation, because communication bandwidth is usually the primary bottleneck in cluster training.

4. Stable LatentMoE: Stable Routing for Nearly a Thousand Experts

4.1 Latent-Space Routing as an Architectural Choice

When the expert count reaches the hundreds or nearly a thousand, making routing decisions directly in the raw hidden space runs into the curse of dimensionality and wasted computation. Latent-space routing compresses routing decisions onto a latent vector far smaller than the hidden dimension: the routing network only processes low-dimensional inputs, sharply reducing overhead. The low-dimensional latent space also helps the routing network learn more abstract and more stable expert-selection patterns.

The challenge of this architecture is that the latent projection introduces additional parameters, so initialization and gradient flow must be designed carefully to prevent the latent space from collapsing into a constant vector during the early rounds of training. In practice, orthogonal or unit-scale initialization is commonly used, with a light weight-decay constraint applied to the latent projection matrix over the first few thousand steps.

4.2 Quantile-Based Balancing: Auxiliary-Loss-Free Load Balancing

Classic auxiliary loss imposes a differentiable uniformity constraint on routing probabilities, encouraging all experts to be selected evenly. When the expert count is very large, however, the auxiliary loss gradient conflicts significantly with the main task gradient—the model has to compromise between “pick the most suitable expert” and “spread the load evenly,” and under extreme sparsity the two are often irreconcilable. In addition, the auxiliary loss coefficient requires manual tuning and is highly sensitive to training scale.

Quantile-based balancing takes a fundamentally different approach. It does not modify the training objective; instead, it dynamically adjusts each expert’s routing threshold during the routing decision itself. Concretely, at every training step the system measures each expert’s actual load distribution, computes quantile statistics, then raises the routing threshold for overloaded experts and lowers it for underloaded ones. This correction is injected as a non-differentiable discrete operation and produces no additional gradient signal, completely sidestepping the gradient conflict problem.

The implementation hinges on controlling the correction frequency and magnitude. Correcting too frequently leaves the routing network no time to adapt to the new threshold distribution, while overly large corrections break the consistency of routing decisions. Introducing momentum smoothing keeps threshold updates at a measured pace—fast enough to rebalance load, gentle enough not to disturb learning stability.

4.3 The Stabilizing Role of the SiTU-GLU Activation

In a sparse expert network, each expert receives only a subset of tokens, so its activation distributions shift as routing patterns fluctuate. This distribution drift amplifies layer by layer in deep networks and can eventually cause the outputs of some experts to diverge numerically. Addressing this requires an activation function that is insensitive to input scale and maintains stable outputs under extreme sparsity.

SiTU-GLU combines the self-normalizing property of SiTU with the nonlinear expressiveness of a gated linear unit (GLU). The SiTU component keeps the output scale essentially unchanged when the input distribution shifts, while the GLU component supplies bounded gating signals that cap output magnitude. Together, they keep expert activations within a reasonable numerical range even under drastic routing fluctuations, easing the burden on layer normalization and improving overall training robustness.

5. Multi-Stage Context Extension and Agentic RL Post-Training

5.1 Why Context Extension Must Be Progressive

Jumping straight from a short context to a million tokens confronts the model with a massive distribution shift: the generalization range of positional encodings, the adaptation speed of attention decay parameters, and the data sampling strategy all need to adjust in sync. A one-shot switch typically leaves the model unstable on long-sequence tasks while its short-range dependency capabilities degrade noticeably.

Staged progression lets the model adapt to length gradually. Pretraining advances along an 8K → 64K → 256K → 1M ladder, with each stage subdivided into sub-stages. At each stage transition, the data mix tilts toward longer contexts while a portion of short-sequence samples is retained to prevent forgetting. The learning rate receives a small warm-up at each transition, helping the model move smoothly onto the new length distribution.

A further benefit of this progressive strategy is debugging efficiency. When training shows numerical instability or capability regression at a certain stage, you know exactly which length interval the problem lies in, so you can adjust attention parameters or data ratios in a targeted way instead of blindly searching a huge space.

5.2 Data Construction for Million-Token Contexts

Ultra-long training data cannot be produced by simply concatenating short documents into long sequences. Genuinely valuable long data must have intrinsic coherence and cross-segment dependencies; otherwise the model learns to “ignore distant information” rather than “use distant information.” Effective synthesis approaches include: organizing mutually referencing source files from a code repository according to the dependency graph; weaving multi-turn dialogues together with external knowledge retrieval results into continuous reasoning chains; and transforming the citation networks of research literature into structured long-range reference tasks.

Data quality verification matters just as much. It is worth designing dedicated “long-range dependency probes”—place key information at different positions in a sequence and evaluate whether the model can accurately recall and use it after spans of thousands or even tens of thousands of tokens. These probes should be run periodically during training as a dynamic monitoring metric for context capability.

5.3 Long-Horizon Credit Assignment in Agentic RL

In agentic scenarios, a single episode of reasoning may involve dozens of decisions: calling a tool, parsing the result, revising the plan, calling another tool. The final success or failure signal must be attributed to every step along the chain. When reasoning chains stretch to hundreds of thousands of tokens, traditional temporal-difference methods face severe credit assignment ambiguity.

An effective practice combines process supervision with outcome supervision: introduce process reward signals at key nodes—such as the format correctness of tool calls and the logical consistency of intermediate reasoning steps—while retaining the sparse reward for final task completion. The density of process rewards can be gradually reduced as training progresses, eventually transitioning the model to learning from outcome signals alone. This dense-to-sparse curriculum effectively mitigates the training instability of long-horizon reasoning.

6. Engineering Recommendations and Takeaways

For teams working on large-scale sparse models and ultra-long context, the following practices have high reuse value:

Put numerical stability first. In any low-precision training framework, verify the numerical boundaries of all attention and recurrence operators before scaling up. Lower-bound constraints, scaling factor clipping, and gating caps are inexpensive safeguards with reliable payoffs. Don’t wait for a NaN to start investigating.

Decouple load balancing from the training objective. Prefer load balancing mechanisms that introduce no extra gradients. Auxiliary-loss-free schemes offer clear advantages in tuning effort and training stability at very large scale.

Extend context progressively. Long-context adaptation is a smooth migration, not a one-time capability switch. Establishing independent evaluation baselines for each stage makes it much faster to locate the source of any capability regression.

Invest in cross-layer information sharing. As depth grows, mechanisms such as cross-layer attention residuals markedly improve deep representation quality while cutting memory and communication costs—a benefit that is especially visible in multi-node, multi-GPU training.

Overall, combining ultra-long context with ultra-large sparse expert models is becoming a clear direction in the evolution of large language models. Numerical stability, load balancing, and long-horizon credit assignment are the three pillars that make this direction engineering-feasible. Looking ahead, as long-context caching strategies at inference time, expert offloading techniques, and agentic training paradigms mature further, sparse LLMs with million-token contexts are poised to deliver greater practical value in code intelligence, research assistance, and complex decision-making.