Architecture Design Practices for Ultra-Long Context and Extremely Sparse MoE

1. Problem Background: When Context Length and Expert Count Both Approach Their Limits

Two clear trends dominate current large model development. First, context windows are expanding from thousands of tokens toward millions. Second, model capacity is growing from hundreds of billions to trillions of parameters through sparse Mixture-of-Experts (MoE) architectures. These two trends are not independent in practice. Together, they push training and inference systems toward several critical thresholds in numerical stability and scheduling:

• Compute bottleneck in attention mechanisms: Standard Softmax attention has O(N²) complexity for sequence length N. At million-token contexts, memory and compute costs become prohibitive. Linear attention and state space model variants reduce complexity to O(N), but their recurrent formulations introduce numerical issues during chunked parallelism.
• Load-balancing bottleneck in MoE: When expert counts reach hundreds or nearly a thousand, routing distributions collapse easily. The traditional approach relies on auxiliary losses to force balanced load, but auxiliary losses compete with the main task gradient and degrade model quality.
• Representation bottleneck in deep networks: As layer counts approach one hundred, gradient and representation propagation efficiency between layers declines. The identity pathway of residual connections gradually dilutes, reducing parameter utilization in deeper layers.

The following sections discuss the causes of these three problems and practical engineering solutions. These methods address general architecture-level design patterns rather than implementations specific to any particular product.

2. Numerical Stability in Linear Attention: Lower-Bounded Decay Parameterization

2.1 Why Linear Attention Underperforms on Tensor Cores

Decayed linear attention, such as Delta Net-style structures, is typically written in chunked recurrent form:

S_t = diag(Γ_t) · S_{t-1} + k_t · v_t^T

Here, Γ is the cumulative decay factor. Computing within a chunk requires element-wise reciprocals of Γ to reduce the diagonal block. The problem arises when chunks are long and decay accumulates: Γ values become extremely small. Taking reciprocals in limited-precision formats like BF16 causes direct overflow, producing Inf or NaN.

A common engineering workaround routes the diagonal block through an explicit position-pair computation path, avoiding reciprocal operations. The cost is that this portion cannot use Tensor Core matrix multiply instructions and must run on SIMT element-wise paths, severely reducing throughput and becoming the hot spot of the attention layer.

2.2 Core Idea of Lower-Bounded Decay

The solution is straightforward: if the decay factor has a positive lower bound, its reciprocal naturally has an upper bound, keeping all values in a numerically safe range.

The approach replaces the unbounded negative Softplus formulation of log decay with a lower-bounded scaled Sigmoid form:

g_t = g_min · sigmoid(e_t)      # g_min is a fixed negative value, e.g., -5
Γ_t = exp(cumsum(g_t))          # decay factor lower bound is exp(g_min) > 0

• Lower bound in log domain: g_t is always greater than g_min, so the cumulative sum does not diverge to negative infinity, and Γ_t never approaches zero.
• Upper bound on reciprocal: The upper bound of 1/Γ_t is exp(-g_min), which is safely within BF16 dynamic range.
• Diagonal block becomes matrix-multiply-friendly: With numerical safety guaranteed, diagonal blocks within chunks can uniformly use matrix multiply paths, fully utilizing Tensor Cores.

2.3 Reuse Value and Implementation Notes

The engineering cost of this change is minimal: replace the activation function of the decay parameter and remove the branched position-pair path from the kernel. The benefit is structural—the attention core operator changes from a mixed-precision branched implementation to a pure matrix multiply implementation, typically yielding substantial throughput gains for long sequences. For teams maintaining custom inference frameworks, this is a high-value optimization: the change is confined to the operator internals, does not involve model structure adjustments, and does not affect existing weight loading logic. Decay parameters require recalibration but can be trained from scratch or fine-tuned lightly.

3. Load Balancing in Extremely Sparse MoE: Auxiliary-Loss-Free Quantile Method

3.1 The Inherent Contradiction of Auxiliary Losses

The classic MoE routing approach adds a load-balancing loss term alongside the main loss to encourage uniform token distribution across experts. When expert counts reach hundreds and each token activates very few experts (small k in Top-k), two problems emerge:

1. Gradient conflict: The gradient direction of the auxiliary loss is inconsistent with the language modeling objective. Increasing its weight damages the main task; decreasing it weakens balancing effectiveness. Tuning costs are high.
2. Irreversible collapse: Once an expert receives no tokens in a batch, it receives no gradient, continues to be unselected, and becomes a “dead expert.”

3.2 Quantile Balancing Mechanism

Quantile balancing moves the balancing mechanism from the loss function level down to the routing decision level: instead of penalizing imbalance, it dynamically determines selection thresholds based on the current distribution of routing scores.

The core flow can be summarized as:

1. Compute global statistics for each expert's routing scores
2. Calculate quantile thresholds according to target load ratios
3. Use quantiles as dynamic per-expert gates; redistribute overflow
4. During backpropagation, routing gradients come only from the main task loss

The advantages:

• No auxiliary loss: The main loss is the sole optimization objective, eliminating gradient competition.
• Fast response: Thresholds update in real time per batch, quickly correcting distribution drift rather than relying on slow gradient accumulation.
• Dead expert avoidance: The quantile mechanism guarantees each expert a clear lower bound on selection probability, preventing permanent exclusion after a single missed batch.

3.3 Implementation Considerations

• Statistics overhead: Quantile computation requires cross-device aggregation of routing scores. Communication volume scales with expert count and batch token count. Design asynchronous or sharded statistics to avoid creating a new synchronization bottleneck.
• Threshold jitter: Inter-batch distribution fluctuations cause threshold jitter, leading to unstable expert selection. In practice, apply sliding averages or momentum to thresholds to balance responsiveness and stability.
• Relationship to capacity factor: Quantile balancing essentially replaces the static capacity factor with a dynamic version. The existing overflow drop logic should still be retained as a fallback for extreme cases.

4. Representation Bottleneck in Deep Transformers: Block-Level Attention Residuals

4.1 Representation Decay from Depth

Residual connections make deep networks trainable, but as layer count increases, deep block inputs become dominated by shallow-layer outputs. The incremental information contributed by new layers diminishes. Representation propagation efficiency between layers declines, manifesting as: smaller deep-layer gradients in late training, nearly frozen deep-layer parameters during fine-tuning, and deep-layer computation that “runs but does nothing” during inference.

4.2 Mechanism of Block-Level Attention Residuals

Block Attention Residuals introduce attention-style cross-layer aggregation on the residual pathway, allowing each layer to selectively read information from historical layers instead of relying solely on layer-by-layer identity accumulation.

Key points:

• Grouped management: Organize several layers into a block. Standard residuals remain within blocks, while attention residuals exchange information between blocks, avoiding the memory and communication explosion of global attention.
• Learnable weights: Cross-layer read weights are determined by attention scores. The model can decide how much information to take from which layer, preventing shallow features from being unconditionally carried to the top.
• Reduced memory and communication: Compared to attending over all historical layers, block-level aggregation significantly reduces cached intermediate activations and lowers cross-stage communication in pipeline parallelism.

4.3 Applicability Assessment

Attention residuals are not necessary for all models. If model depth is within a few dozen layers, standard residuals are typically sufficient. The technique becomes worthwhile when depth approaches or exceeds one hundred layers and training shows clear deep-layer gradient decay. Before implementation, run diagnostics: compare gradient norms and activation variances across layers to confirm that the bottleneck indeed lies in inter-layer propagation rather than other factors.

5. Overall Engineering Recommendations

Considering these three techniques within a single training system, several general engineering principles apply:

1. Prioritize numerical safety over theoretical elegance: The value of changes like lower-bounded decay lies in making the implicit constraint of numerical safety explicit in parameterized design. Any operator involving recurrence, accumulation, or reciprocals should undergo boundary value testing in BF16 first.
2. Push balancing mechanisms downward: Load balancing at the routing decision layer is more direct, converges faster, and is easier to debug than at the loss function layer. Loss functions should remain as close to a single objective as possible.
3. Locate bottlenecks before optimizing: Deep representation bottlenecks, communication bottlenecks, and operator bottlenecks present similarly—all appear as slow training—but require completely different solutions. Use profiling and gradient statistics to identify the real bottleneck before deciding which structural change to introduce.
4. Progress context length in stages: Training that gradually extends from short windows to ultra-long ones is essentially curriculum learning applied to the sequence length dimension. Directly training on ultra-long sequences makes attention distributions difficult to converge. Staged progression, such as 8K to 64K to 256K to 1M, allows the model to fully converge at each length before transferring to the next stage.

6. Summary

Ultra-long context and ultra-large sparse MoE are two of the most engineering-sensitive directions in current large model architecture evolution. Their technical challenges often lie not in feasibility but in stability under limited precision, bandwidth, and memory constraints. Lower-bounded decay parameterization, quantile balancing, and block-level attention residuals address three constraint dimensions: numerical safety, scheduling balance, and cross-layer propagation. Their common thread is transforming stability that was previously implicit and maintained through parameter tuning into stability that is explicit and guaranteed by structure. For teams building their own large model infrastructure, these patterns have low reuse costs and clear benefits, making them worth considering early in architecture design.