Linear Attention and MoE Optimization Techniques for Efficient LLM Inference

As context windows of large language models rapidly expand from tens of thousands of tokens to 1M and beyond, the computational and memory overhead of conventional attention mechanisms has become a core bottleneck for inference deployment. At the same time, ultra-large sparse models face challenges in load balancing, numerical stability, and deep representation degradation during both training and inference—factors that directly determine whether a model can be reliably deployed in real-world workloads. This article provides a systematic review and engineering-oriented analysis of several high-value, reusable techniques across three directions: linear attention, Mixture-of-Experts (MoE), and deep Transformers.

1. Background: New Challenges from Long Context and Sparsity

In long-context scenarios, the computational complexity of standard self-attention scales quadratically with sequence length. Even with engineering approaches such as chunking and caching, beyond a certain length, problems such as insufficient GPU memory and unpredictable latency persist. As a result, linear attention and its variants have regained attention. The core idea is to transform attention computation from explicit position-pair interactions into a form of “state propagation + matrix multiplication,” so that inference cost grows linearly with sequence length.

However, linear attention is not trivial to implement in engineering. Because it involves recursive decay factors, numerical overflow and underflow easily occur during low-precision training or inference. In particular, when performing blocked matrix multiplication on GPU Tensor Cores, certain diagonal blocks must fall back to position-by-position computation paths, preventing full utilization of hardware acceleration.

On the other hand, MoE structures increase model capacity by introducing a large number of expert networks while controlling computation through sparse activation. However, imbalanced expert load has long been a persistent challenge for this architecture. Conventional approaches typically introduce auxiliary loss functions to enforce load balancing, but this conflicts with the primary training objective and incurs high tuning costs in large-scale distributed training.

In addition, deep Transformers commonly exhibit representation degradation: as the number of layers increases, the effective signal in the residual stream gradually weakens, and the contribution of deeper modules to the final output declines. How to break this bottleneck without significantly increasing memory and communication overhead is an important direction in model architecture design.

2. Core Optimization 1: Lower-Bounded Decay Parameterization

1. The Nature of the Problem

The recursive form commonly used in linear attention can be abstracted as:

S_t = S_{t-1} * decay + V_t * K_t^T
O_t = S_t * Q_t

Here, decay is typically generated by a parameterized function. If the decay factor approaches 0 in certain intervals, its reciprocal amplifies sharply under finite-precision representation, causing numerical overflow in low-precision formats such as BF16. In blocked computation, this means certain blocks cannot safely use matrix multiplication acceleration and must take the slower element-wise path.

2. Improvement Approach

The core idea of lower-bounded decay parameterization is to set a reliable lower bound for the decay factor, thereby constraining the reciprocal scaling factor within any local time range to a safe numerical interval. Specifically, the previously unbounded negative softplus form is replaced with a lower-bounded scaled sigmoid form, for example:

g_t = g_min * sigmoid(e_score_t)

where g_min is a fixed constant. This approach does not change the overall recursive structure but makes the numerical range within any small block controllable.

3. Engineering Benefits

The greatest value of this change lies at the engineering implementation level. When all blocks satisfy numerical safety conditions, both diagonal and non-diagonal blocks in the causal sequence can uniformly follow the dense matrix multiplication path, eliminating special-case branch checks. This not only improves implementation efficiency but also makes the code easier to optimize on low-level compute libraries such as CUDA/CUTLASS. Meanwhile, this parameterization is more friendly to long-sequence extrapolation without positional encoding, reducing additional adaptation costs during long-context deployment.

3. Core Optimization 2: Auxiliary-Loss-Free MoE Load Balancing

1. Limitations of Traditional Approaches

In large-scale MoE training, the gating network may tend to route a large number of tokens to a few experts, causing some experts to be overloaded while others sit idle. Traditional methods typically introduce an additional load-balancing loss term weighted into the primary loss function. However, this approach has several notable problems:

• The load-balancing loss competes with the primary task loss, requiring careful weight tuning;
• Tuning costs rise significantly for ultra-large models;
• The addition of auxiliary losses may negatively impact final model quality.

2. Basic Idea of Quantile Balancing

The goal of Quantile Balancing is to dynamically adjust routing strategies directly based on quantile information from current gating outputs, rather than relying on additional loss functions. The general process is:

1. Count the distribution of tokens assigned to each expert in the current batch;
2. Compute key quantile thresholds;
3. Recalibrate gating scores based on quantiles, suppressing scores for overloaded experts and boosting scores for underloaded experts.

This mechanism takes effect in real time during training without modifying the primary loss function or introducing additional hyperparameters.

3. Engineering Significance

For teams with large-scale distributed training needs, auxiliary-loss-free load balancing reduces the conflict between training objectives and system efficiency. It is particularly suitable for scenarios with many experts and large per-step batch sizes, where load imbalance is more pronounced and manual tuning is nearly infeasible.

4. Core Optimization 3: Block Attention Residuals

1. Representation Degradation in Deep Transformers

As model depth increases, although residual connections ensure gradient propagation, they also introduce signal dilution. In certain deep structures, the effective contribution of intermediate layers gradually decreases, producing a phenomenon of “deep redundancy.” This means that even with more stacked layers, model performance does not necessarily improve proportionally, while memory and communication overhead continue to grow.

2. The Idea Behind Block Attention Residuals

The core of Block Attention Residuals is to add a learnable residual pathway for attention output at the boundaries of “blocks” composed of several consecutive layers. Unlike standard residuals, this pathway passes information directly from the block input to the block output, bypassing multiple intermediate sublayers. The structure can be illustrated as:

# Standard residual
x = x + sublayer(x)
# Block-level residual
block_input = x
for layer in block:
    x = layer(x)
x = x + block_attn_residual(block_input)

This design allows deeper modules to no longer rely solely on residual signals propagated layer by layer, but instead to access high-quality representations from the block’s starting point.

3. Benefits Analysis

The main benefits of block attention residuals are reflected in three aspects:

• Improved representation quality: Mitigates deep signal degradation, making the actual contribution of deep network layers more uniform;
• Reduced memory and communication: Because attention computation can be merged and cached at the block level, the need to store intermediate states is reduced;
• Improved training stability: The block-level residual pathway provides a more direct gradient propagation path, facilitating convergence of deep models.

5. Combined Application and Deployment Recommendations

The three techniques above are not mutually exclusive and can be used together within the same model architecture. For example, a long-context sparse large model can simultaneously adopt:

• Linear attention based on lower-bounded decay parameterization to improve long-sequence inference throughput;
• Quantile balancing routing to ensure stable MoE expert load;
• Block attention residuals to improve the actual representation capability of deep models.

For engineering deployment, the following order is recommended:

1. First implement numerically safe blocked linear attention, as this is the foundation for long-context inference;
2. Then introduce auxiliary-loss-free load balancing to improve sparse training efficiency;
3. Finally, validate the actual benefits of block attention residuals on deep structures, since this change has a relatively larger impact on model architecture.

It should be emphasized that each optimization should be validated against specific model scale, context length, hardware configuration, and business latency requirements. For example, when sequence length has not yet reached a critical threshold, the benefits of linear attention may not be obvious; when the number of experts is small, traditional auxiliary loss approaches may already be sufficient.

6. Summary

Long context and sparsity are two important paths toward efficient LLM inference, but both face multiple engineering challenges at the numerical, load, and structural levels. The three categories of techniques discussed in this article offer actionable optimization approaches from three angles: numerical stability of attention operators, load balancing of MoE routing, and representation preservation in deep Transformers. Their common characteristics are: relatively focused changes, strong engineering verifiability, and good compatibility with mainstream training and inference infrastructure. For technical teams building or optimizing long-context large models, these methods can serve as important references for model architecture selection and inference performance tuning.