Architecture Design for Ultra-Long Context and Sparse Expert Models

When model scales cross the trillion-parameter threshold and context windows expand to the million-token level, system bottlenecks rarely stem from single-point compute power. Instead, they arise from a series of coupled engineering constraints: the numerical stability of linear attention operators on tensor cores, load imbalance caused by routing across massive numbers of experts, representation dilution across deep network layers, and the memory and communication overhead brought by ultra-long sequences. Drawing on architectural designs disclosed in public technical reports, this article analyzes these constraints, focusing on three highly reusable technical solutions and outlining a multi-stage training organization strategy.

1. Background: Three Coupled Challenges

Before diving into specific solutions, it is crucial to clarify the structural contradictions these models face.

Challenge 1: Parallel Efficiency vs. Numerical Precision in Linear Attention. Linear attention and its variants replace the full attention matrix with recurrent states, reducing complexity from quadratic to linear relative to sequence length. However, in the chunked computation paradigm, the reciprocal of the intra-chunk cumulative decay factor is highly prone to overflow or underflow under finite precision. Once numerical stability is lost, the diagonal blocks must fall back to an explicit position-by-position computation path, failing to leverage dense matrix multiplications and losing throughput advantages.

Challenge 2: Expert Expansion vs. Routing Balance. When the number of routing experts reaches hundreds or even nearly a thousand, and the activation ratio is pushed to an extreme low, traditional auxiliary loss balancing strategies begin to fail. The auxiliary loss itself interferes with the main task’s gradients, converges slowly at a massive expert scale, and easily leads to a collapse phenomenon where a few experts are repeatedly activated while others remain idle.

Challenge 3: Network Depth vs. Representation Preservation. In deep Transformers, while residual paths mitigate vanishing gradients, the homogenization and information dilution of inter-layer representations persist. In ultra-long context scenarios, this dilution is amplified, manifesting as the ineffective utilization of remote information in deep layers.

These three challenges are interdependent: numerical precision limits operator selection for attention modules; operator selection impacts memory and communication budgets; and these budgets inversely constrain the viable range of expert scale and network depth. A coordinated design is therefore essential.

2. Lower-Bounded Decay Parameterization: Returning Diagonal Blocks to Dense Matrix Multiplication

2.1 Identifying the Pain Point

In chunked recurrent linear attention, intra-chunk computation requires introducing a cumulative decay factor. If the decay is parameterized logarithmically and its value is unbounded, the dynamic range of the scaling coefficient becomes extremely wide when performing reciprocal scaling within the chunk. Taking BF16 as an example, its representable range is limited. In extreme cases, the scaling factor directly overflows to infinity or underflows to zero, rendering the result unusable.

A common engineering compromise is to route diagonal blocks through a separate position-by-position computation path, bypassing matrix multiplication. While numerically safe, this path breaks operator consistency—diagonal and non-diagonal blocks use different computation graphs, increasing kernel branching and making unified pipeline optimization difficult for compilers.

2.2 Mechanism Design

The solution is to change the logarithmic decay from an unbounded value to one with a lower bound. Specifically, an unbounded transformation is replaced with a scaled sigmoid function, and a lower bound constant is fixed. This ensures that the reciprocal scaling factor is locked within a safe numerical range within any fixed-length chunk (e.g., 16 tokens).

A simplified form illustrates this parameterization approach:

# Pseudo-code: Lower-bounded logarithmic decay parameterization
G_MIN = -5.0
def log_decay(e_A_h_z):
    # Map unbounded input to the (G_MIN, 0) interval
    return G_MIN * sigmoid(e_A_h_z)
def block_scale(log_decay_values):
    # Intra-chunk cumulative decay, its reciprocal range is controlled due to the lower bound
    cumulative = cumsum(log_decay_values)
    return exp(-cumulative)  # Numerical range is predictable, avoiding overflow/underflow

The key lies not in the function form itself, but in the fact that introducing a lower bound shifts the numerical range from “unpredictable” to “predefinable.” Once the range is definable, kernel implementations can safely invoke dense matrix multiplications uniformly for all causal blocks (both diagonal and non-diagonal), eliminating special branches.

2.3 Benefit Analysis

The engineering value of this modification is reflected in three aspects:

• Operator Unification: Eliminates the special path for diagonal blocks. All blocks under causal masking use the same computation pipeline, reducing kernel branching and improving occupancy.
• Controllable Precision: The scaling factor consistently falls within a safe range, significantly reducing numerical anomalies during low-precision training without needing to elevate precision specifically for diagonal blocks.
• Extrapolation-Friendly: Because the numerical behavior does not rely on explicit positional encoding, the model possesses more natural extrapolation capabilities on contexts exceeding training lengths, which is critical for million-level contexts.

With minimal modification and direct benefits, this is the most reusable component of this design.

3. Quantile Balancing: Auxiliary-Loss-Free Expert Load Dispatching

3.1 Limitations of Traditional Solutions

Load balancing in Mixture of Experts (MoE) models has long relied on auxiliary losses: appending a penalty term alongside the main task loss to punish routing distribution skew. This approach is effective with a small number of experts but exposes two issues under settings of nearly a thousand experts and extremely low activation ratios: the auxiliary loss weight requires fine-tuning (too large harms the main task, too small fails to constrain routing), and gradient signals are diluted among numerous experts, slowing convergence.

3.2 Quantile-Based Dynamic Adjustment

The alternative is to avoid introducing additional loss terms and instead perform dynamic calibration directly during the routing scoring or capacity allocation phase. The core idea is to maintain online statistics on each expert’s cumulative activation volume or score distribution, determine a dynamic threshold based on quantiles, and redistribute excess load to low-load experts.

The advantages of this mechanism include:

• No Interference with Main Gradient: Balancing is executed through forward-pass routing decisions, without injecting additional terms into backpropagation.
• Fast Response: Adjustments based on online statistics eliminate the need to wait for auxiliary losses to converge slowly over multiple iterations.
• Adaptation to Extreme Sparsity: Under extremely low activation ratios, the quantile method protects tail-end experts more directly, effectively preventing expert collapse.

Implementation requires attention to the update frequency and sliding window length of the statistics, which together determine the sensitivity and stability of the balancing mechanism.

4. Block-Wise Attention Residual: Mitigating Deep Layer Representation Dilution

4.1 The Core Problem

Standard residual connections directly add each layer’s output to its input, working well in shallow networks. In extremely deep networks, however, this leads to gradual representation homogenization, reducing the proportion of incremental information between layers. Ultra-long contexts exacerbate this: remote information must traverse more layers to impact the output, undergoing repeated dilution along the way.

4.2 Block-Wise Cross-Layer Sharing

The concept of block-wise attention residuals involves dividing the network into blocks of several layers and establishing cross-layer attention-style information pathways at the block level. This allows layers within a block to directly access the representation at the block’s entrance, rather than relying solely on layer-by-layer residual accumulation.

The benefits are twofold:

• Representation Quality: Layers within a block can directly “look back” at the block’s entrance state. Deep layers no longer rely entirely on sequential transmission, preserving remote information more intact.
• Resource Overhead: Because cross-layer pathways are established at the block granularity rather than connecting every two layers, both memory footprint and communication volume are kept within acceptable limits.

Designing this requires a trade-off in block size: a larger block means sparser cross-layer pathways and lower overhead, but fewer opportunities for direct access for intra-block layers; a smaller block reverses this dynamic.

5. Training Organization: Multi-Stage Context Progression and Agent Post-Training

Beyond architecture, the organization of the training pipeline is equally critical. Context window expansion is typically not done in a single step but progressed in stages: starting with a shorter window and gradually extending to longer ones. This progressive strategy allows the model to converge stably on short sequences first, gradually adapting to the positional and attention distributions of longer sequences, avoiding the convergence difficulties of training directly on ultra-long sequences from the start.

Following pre-training, the post-training phase tailored for agent scenarios is also noteworthy. Agent tasks are characterized by long-horizon reasoning and multi-step decision-making. A single interaction often spans a massive number of tokens, demanding high preservation of intermediate states. Scaling post-training data to the million-token level allows the model to undergo targeted optimization on real long-horizon task distributions. This complements the architectural capacity for ultra-long contexts: the architecture provides the capacity, and post-training teaches the model how to utilize it.

6. Reusability Summary

Ranking the technical points in this design by reusability value:

Technical Point	Problem Solved	Reusability Value	Modification Cost
Lower-bounded decay parameterization	Numerical overflow in linear attention diagonal blocks	High	Minimal
Quantile balancing	Load imbalance in large-scale expert routing	High	Medium
Block-wise attention residual	Representation dilution in deep networks	Medium-High	Medium

The common characteristic among the three is that they each target a specific engineering bottleneck and can be introduced independently without drastically altering the overall architecture. For teams building ultra-long context or ultra-large sparse expert systems, this “point-wise pluggable” design philosophy is more instructive than copying an entire architecture.

Finally, it must be emphasized that architectural design is never isolated. Numerical precision dictates available operators; available operators dictate memory and communication budgets; and these budgets inversely constrain expert scale and network depth. Only by finding a consistent solution across these constraints can trillion parameters, million-token contexts, and extreme sparsity coexist simultaneously.