Key Techniques for Training and Inference of Long-Context Large Models

1. Background: The Dual Challenge of Million-Token Contexts and Sparse Large Models

As agents, long-document understanding, and codebase-level reasoning continue to advance, models are required to process increasingly long input sequences. Traditional dense attention scales quadratically in both computation and memory as sequence length grows, making it impractical to extend directly to hundreds of thousands or even millions of tokens. At the same time, to increase model capacity under limited compute budgets, ultra-large-scale sparse expert models have become a common choice. However, as the number of experts grows, load balancing and numerical stability during training become more pronounced problems.

Based on the content of a technical report, this article covers four key dimensions: numerical underflow and computational efficiency in linear attention, load balancing for ultra-large-scale expert routing, information reuse and memory optimization in deep attention, and long-horizon reasoning post-training for agents. These technical points are not exclusive capabilities of any single model; rather, they are common engineering challenges in the current direction of long-context sparse large models. Understanding these problems and their solutions can help developers avoid pitfalls when designing or training similar architectures.

2. Numerical Stabilization and Computational Efficiency in Linear Attention

2.1 Problem Origin: Numerical Underflow of the Decay Factor

Linear attention reduces complexity by decomposing attention computation into a chunkwise recurrent form. Within a chunk, the influence of historical information on the current position still needs to be handled. This influence is typically modeled through a cumulative decay factor: earlier information should have diminishing influence over time. In implementation, the decay factor participates in matrix operations, and especially in diagonal chunks, the reciprocal of the decay factor must be computed.

In conventional approaches, decay parameters are often generated through an unbounded log-decay function. When sequences are long or decay is fast, this reciprocal can become extremely large or extremely small under certain numerical precisions, exceeding the safe representable range of finite-precision floating-point numbers. For example, when using 16-bit floating point for intermediate computation, once a value falls outside the representable interval, overflow or underflow occurs, leading to training instability or abnormal inference results.

2.2 Cascading Impact on Computational Efficiency

The direct consequence of numerical instability is that diagonal chunks cannot use the same efficient dense matrix multiplication kernels as other chunks. Developers are forced to write a separate computational path for diagonal chunks, explicitly computing attention weight pairs position by position, which significantly reduces the throughput of the entire attention module. Especially on modern GPUs, matrix multiplication units are far faster than element-wise operations, and once forced back to a position-by-position path, the performance loss is very noticeable.

Therefore, the essence of the problem is not “whether it can be computed,” but “whether all chunks can follow the same efficient matrix multiplication path.” This requires that the reciprocal of the decay factor remains numerically safe across all practical chunk ranges.

2.3 Solution: Decay Parameterization with a Lower Bound

One effective approach is to impose an explicit lower bound on the decay parameter. Specifically, the original unbounded log-decay function is replaced with a bounded scaling form. By adding a fixed lower bound to a scaled nonlinear function output, the decay coefficient is always kept within a controllable range.

An illustrative representation is as follows:

Original form:
decay = -softplus(param)          # unbounded, can be very small
Improved form:
decay = fixed_lower_bound * sigmoid(projected_param)   # explicit lower bound

In the improved form, regardless of how the input parameter changes, the decay coefficient never falls below the preset lower bound. This way, within any fixed-length chunk, the numerical range of the decay factor’s reciprocal is constrained to a safe interval. From an engineering perspective, this guarantees that under 16-bit precision, no overflow occurs within a chunk of 16 positions.

2.4 Practical Benefits

This change appears simple, but its engineering benefits are significant. All causal chunks—whether diagonal or non-diagonal—can uniformly invoke dense matrix multiplication kernels. There is no longer a need to maintain a special branch for diagonal positions in the code, greatly simplifying kernel implementation. At the same time, because matrix multiplication units are fully utilized, the overall throughput of the attention module improves noticeably. Furthermore, since explicit positional encoding is no longer relied upon, the model can directly extrapolate to longer context lengths after training without retraining position embeddings.

This approach can be generalized to any scenario that implements linear attention in a chunkwise recurrent form: when designing the decay function, first ask clearly “under the target numerical precision, is the reciprocal of the decay factor always safe within a single chunk?” If not, the parameterization should be adjusted first rather than adding computational branches for special cases.

3. Load Balancing for Ultra-Large-Scale Expert Routing

3.1 Problems with Traditional Auxiliary Losses

In mixture-of-experts models, the router is responsible for assigning each input position to several experts. If the routing strategy is driven entirely by the training signal, severe load imbalance typically occurs: a few experts are selected frequently while the rest are barely activated. This not only wastes compute but also reduces the effective utilization of model capacity.

A common solution is to introduce an auxiliary loss function that explicitly penalizes load imbalance in the training objective. However, auxiliary losses have two clear drawbacks: first, the loss weight must be tuned manually—too large and it affects main-task convergence, too small and load imbalance persists; second, auxiliary losses are indirect constraints, and the model may satisfy superficial balance through unexpected means while the actual routing distribution remains unhealthy.

3.2 Quantile-Based Balancing Mechanism

Another approach is to impose constraints directly on the routing selection process itself, rather than adjusting indirectly through a loss function. The core idea of the quantile-based balancing mechanism is to compute statistics on the score distribution of each expert’s inputs and set dynamic thresholds based on quantiles, so that each expert is selected roughly the same number of times per batch.

Specifically, for each expert, the system computes in real time the distribution of scores assigned to it by all positions in the current batch, then determines a quantile as that expert’s dynamic threshold based on the target activation count. Only positions scoring above this threshold have a chance of being routed to that expert. Because the threshold is dynamically adjusted, different experts can maintain similar selection counts even when their raw score distributions differ greatly.

The advantages of this approach are: no additional auxiliary loss term is needed, keeping the training objective clean; load balancing is a hard constraint directly reflected in each routing decision; and implementation can be done with batched parallel statistics at relatively low overhead.

3.3 Significance for Training Stability

When the number of experts reaches hundreds or even close to a thousand, the difficulty of tuning traditional auxiliary losses rises substantially. By turning load balancing from a “soft constraint” into a “hard constraint,” quantile-based balancing significantly reduces the burden of hyperparameter tuning. More importantly, it avoids interference with the model’s primary objective caused by auxiliary losses, making training curves more stable.

For engineering practice, one principle worth considering is: when the model scale becomes large enough, the controlling power of indirect soft constraints declines, and direct constraint mechanisms should be prioritized. Quantile-based balancing is one such direct constraint implementation in the expert routing scenario.

4. Information Reuse and Memory Optimization in Deep Attention

4.1 Representation Bottleneck in Deep Models

As the number of layers increases, intermediate representations produced by attention modules may gradually become diluted or homogenized as they propagate across deep layers. Each layer’s attention computation extracts information anew from the previous layer’s output, but deep layers often struggle to retain important patterns from early layers.

One intuitive improvement is to introduce residual connections: pass attention outputs from shallower layers directly to deeper layers and add them to the current layer’s attention output. This way, deep layers do not need to rely entirely on information passed layer by layer but can directly leverage some representations from early layers.

4.2 Implementation Strategy for Block-Level Attention Residuals

In practice, directly passing complete attention matrices across layers increases memory and communication overhead. A more efficient approach is to apply residual connections at the block level. Specifically, several layers are grouped together, the attention output blocks of selected layers within the group are cached, and residual additions are performed at corresponding positions in subsequent layers.

This approach yields several benefits: first, it breaks the information decay of layer-by-layer propagation and mitigates representation degradation in deep layers; second, because only residual information for some blocks needs to be cached, memory overhead remains controllable; third, in distributed training, the communication volume for cross-layer residuals is lower than repeatedly transmitting complete attention matrices, reducing communication pressure.

4.3 Difference from Conventional Residual Connections

Conventional residual connections typically operate on feedforward network outputs or the overall output of an attention layer. Block-level attention residuals operate at a finer granularity on local regions of the attention output matrix and are shared across multiple layers. This design allows the model to introduce cross-layer information for specific position ranges at different depths, rather than handling everything uniformly on a global basis.

From an engineering perspective, block-level residuals require maintaining additional caching and indexing logic, but as long as block size is chosen reasonably, the extra overhead can be kept within a small range. For ultra-deep model architectures, the benefits of this design typically outweigh its implementation cost.

5. Long-Horizon Reasoning Post-Training for Agents

5.1 Why Dedicated Long-Horizon Post-Training Is Needed

The pretraining stage primarily improves the model’s ability to model long contexts, but agent tasks also require the model to maintain a consistent reasoning chain across multi-step interactions, remember key information from early steps, and backtrack to correct course when necessary. This is not merely a matter of “reading long” but of “using correctly.”

Therefore, after pretraining is complete, long-horizon agent training data must be specifically constructed for reinforcement learning post-training. The length of this data typically reaches hundreds of thousands or even millions of tokens, containing multiple rounds of tool calls, environment feedback, intermediate reasoning, and final decisions. The model must learn long-term memory and strategy backtracking on this data.

5.2 Training Data and Task Design

Long-horizon agent post-training data is typically not simply concatenated documents but interaction trajectories with a clear task structure. A single trajectory may include: an initial goal, multiple action steps, environment returns, intermediate reflections, alternative approach attempts, and a final result. The training objective is to teach the model, within long trajectories, when to trust early information, when to abandon prior conclusions in light of new evidence, and how to avoid losing key goals when the context window is filled with large amounts of intermediate steps.

This kind of data is difficult to construct because the long-term dependencies of the task must be genuine rather than easily solvable through local information. Data quality directly affects post-training effectiveness.

5.3 Difference from Ordinary Reinforcement Learning Post-Training

Ordinary reinforcement learning post-training typically targets short-horizon question answering or simple tasks, where reward signals are easy to obtain. Long-horizon agent post-training must handle sparse rewards, partially observable states, and uncertainty in the action space. The model needs to maintain a correct policy over a very long time span, which places higher demands on credit assignment mechanisms.

In practice, long trajectories usually need to be decomposed into multiple evaluable intermediate nodes, with local rewards or process supervision signals set at these nodes to alleviate the sparse reward problem. At the same time, training compute costs increase substantially, requiring systematic optimization of sampling, replay, and gradient computation.

6. Summary of Reusable Engineering Lessons

Based on the above analysis, several broadly applicable lessons can be distilled from this technical report:

First, prioritize solving numerical stability problems through parameterization design. Rather than adding special branches in compute kernels to handle extreme cases, redesign the parameter form so that all computations follow a unified, efficient path. Bounded decay parameterization is a typical example.

Second, large-scale sparse models should move from soft constraints to hard constraints. As model scale grows, the regulatory power of auxiliary losses declines, and imposing balancing constraints directly at the routing selection level is more reliable.

Third, deep models need to pay attention to cross-layer information reuse. Pure layer-by-layer forward propagation leads to information decay, and block-level attention residuals achieve cross-layer information sharing at relatively low cost, making them a worthwhile architectural choice.

Fourth, long-context capability does not equal long-horizon reasoning capability. Pretraining solves the “reading in” problem, while post-training solves the “using” problem. Long-horizon agent post-training requires dedicated trajectory data and training strategies.

These lessons are not unique to any single vendor or model but represent common technical threads in the current direction of long-context sparse large models. Understanding the design motivations and implementation constraints behind them is more valuable than memorizing specific parameters. For teams building or optimizing similar architectures, this content can serve as a reference starting point for architecture reviews and engineering iterations.