Optimizing Large Model Architectures: Key Techniques for Linear Attention, MoE Load Balancing, and Deep Transformers

In the engineering practice of large-scale language models, three core challenges persist: the computational efficiency of attention mechanisms, load balancing in Mixture-of-Experts (MoE) architectures, and representation degradation in deep networks. This article analyzes three highly reusable techniques: lower-bounded decay parameterization for linear attention, auxiliary-loss-free dynamic load balancing, and block-level attention residuals. These techniques target operator efficiency, training stability, and deep representation capability, respectively, with manageable engineering changes suitable for rapid validation and migration in existing training and inference frameworks.

1. Lower-Bounded Decay Parameterization: Unlocking Tensor Core Compute for Linear Attention

1.1 Background: The Diagonal Block Challenge in Linear Attention

Linear attention and its variants (such as incremental network structures) rewrite attention computation into a recurrent form, reducing sequence-level quadratic complexity to linear complexity. In practical implementations, chunk-based parallel strategies are commonly used: sequences are divided into fixed-length chunks, with dense matrix multiplication performed within each chunk and information passed between chunks via recurrent state updates.

However, this approach faces a numerical stability issue during intra-chunk computation. Each positional pair within a chunk must be multiplied by a cumulative decay factor, typically expressed in exponential form. When the logarithm of the decay rate is unbounded negative, the reciprocal of the cumulative decay factor becomes prone to severe overflow or underflow under limited precision (e.g., BF16). As a result, diagonal blocks—representing the contribution of each position to itself—cannot directly use dense matrix multiplication and must fall back to explicit position-pair elementwise computation, failing to leverage tensor core matrix multiply acceleration instructions.

The essence of this problem is that the special handling of diagonal blocks stems from the unconstrained numerical range of the decay factor, not from any inherent algorithmic requirement for special treatment.

1.2 Technical Mechanism: From Unbounded Decay to Lower-Bound Constraint

The solution is to reformulate the logarithm of the decay rate from an unbounded form to a lower-bounded parameterization. Specifically, the original unbounded negative Softplus form is replaced with a scaled Sigmoid form with a fixed lower bound:

g = g_min * Sigmoid(e_A * h_z * t_h)

Here, g_min takes a fixed value (e.g., -5), and e_A, h_z, and t_h are intermediate variables that are either learnable or controlled by hyperparameters. The core effect of this change is that within any fixed-length chunk (e.g., 16 tokens), the reciprocal scaling factor of the cumulative decay factor is locked within a safe numerical range, preventing overflow or underflow under BF16 precision, regardless of the decay rate magnitude.

From a mathematical perspective, the lower-bound constraint ensures that the ratio of decay factors between any positional pairs within a chunk remains within a representable range, making the numerical behavior of diagonal blocks consistent with non-diagonal blocks.

1.3 Engineering Benefits: Unified Computation Path for Causal Blocks

The direct engineering benefits of this parameterization change include:

• Elimination of special-case diagonal block handling: All causal blocks (both diagonal and non-diagonal) can directly leverage dense tensor core matrix multiplication, eliminating the need to maintain a separate position-pair computation path.
• Significant throughput improvement: Since diagonal blocks no longer use the inefficient elementwise path, the overall computational density of the attention operator increases, leading to higher tensor core utilization.
• Native support for long-context extrapolation without positional encoding: The numerical stability introduced by lower-bounded decay parameterization enables stable attention computation at extremely long contexts (up to million-token scale) without relying on explicit positional encoding.

1.4 Reuse Recommendations

The value of this technique lies in its minimal engineering overhead with substantial benefits. For any training framework using linear attention or incremental network structures, modifying the activation function form of the decay parameter and fixing the lower bound enables more efficient operator implementations at the CUDA or CUTLASS level. Key considerations:

• The choice of lower-bound value should be tuned in conjunction with the specific precision format (BF16/FP16) and chunk size to ensure numerical safety without sacrificing model expressiveness.
• Chunk size selection affects the tightness of the lower-bound constraint; experiments in the range of 16 to 64 are recommended.
• This approach is fully compatible with existing chunk-based parallel strategies and can be adopted as a plug-and-play replacement component.

2. Auxiliary-Loss-Free Dynamic Load Balancing: Efficient Training for Large-Scale MoE

2.1 Background: The Load Balancing Dilemma in MoE

Mixture-of-Experts (MoE) architectures replace feed-forward layers with multiple expert networks, with a routing mechanism activating a subset of experts. This allows model parameters to scale significantly while keeping computational cost manageable. However, the inherent nature of the routing mechanism leads to load imbalance: some experts are over-activated while others receive insufficient training. This not only wastes computational resources but also degrades expert proficiency.

Traditional solutions introduce an auxiliary loss—a regularization term added to the training objective to penalize load imbalance. But auxiliary losses have notable drawbacks: the weight coefficient requires manual tuning (too large harms main task performance, too small provides insufficient balancing), and the auxiliary loss can conflict with the main task objective, leading to training instability.

2.2 Technical Mechanism: Quantile Balancing Strategy

Quantile Balancing offers an auxiliary-loss-free dynamic load balancing approach. The core idea: during the routing decision stage, dynamically adjust routing thresholds based on the load distribution across experts within the current batch, so that the number of tokens received by each expert tends toward balance.

Specifically, instead of using raw routing scores for Top-K selection, dynamic thresholds are determined based on the quantiles of the score distribution. This allows routing decisions to adaptively account for current load conditions without introducing additional terms into the loss function.

The key advantages of this strategy:

• No auxiliary loss: The training objective remains unchanged, avoiding conflicts between the main task and balancing objectives.
• Rapid dynamic adjustment: Routing can be adjusted in real-time based on current batch load at every training step.
• Scalability: Suitable for scenarios with extremely large numbers of experts; computational overhead scales linearly (not quadratically) with the number of experts.

2.3 Engineering Benefits and Reuse Recommendations

In practice, this strategy delivers:

• Significantly improved expert utilization, reducing the “dead expert” phenomenon.
• Enhanced training stability without the need for fine-tuning auxiliary loss weights.
• Low integration cost into existing MoE training frameworks—only the routing decision module needs replacement.

When reusing this technique:

• Consider the update frequency and smoothing strategy for quantile thresholds to avoid drastic load fluctuations.
• Coordinate with the capacity factor to prevent token dropping in extreme cases.
• Consider looser balancing constraints during early training, gradually tightening as training progresses.

3. Block-Level Attention Residuals: Breaking Representation Bottlenecks in Deep Transformers

3.1 Background: Representation Degradation with Increased Depth

As the number of Transformer layers increases, deep networks face representation degradation: useful features learned in shallow layers become diluted or overwritten during forward propagation, making it difficult for deep layers to effectively utilize shallow information. While traditional residual connections mitigate gradient vanishing, at extreme depths, information accumulation along the residual path can still lead to representation homogenization.

Additionally, standard attention mechanisms in deep networks require storing full key-value caches, with memory and communication overhead growing linearly with layer count—becoming a bottleneck in long-context scenarios.

3.2 Technical Mechanism: Block-Level Attention Residuals

The core idea of Block Attention Residuals is to partition Transformer layers into blocks and introduce attention residual connections at the block level, rather than only at the layer level with standard residuals. Specifically, the attention output within each block interacts with block-level historical representations, allowing deeper blocks to selectively reuse representation information from shallower blocks.

The effects of this mechanism:

• Breaking the deep representation bottleneck: Block-level residuals provide more direct propagation paths for shallow information, preventing dilution across layers.
• Reduced memory overhead: Through block-level key-value cache sharing or compression, redundant storage is reduced.
• Lower communication overhead: In distributed training, block-level structure reduces the frequency and volume of cross-device communication.

3.3 Engineering Benefits and Reuse Recommendations

The engineering value of this technique:

• Stronger representation capability at equivalent depth, with improved downstream task performance.
• Lower memory footprint, supporting longer context windows.
• Reduced communication overhead, improving distributed training efficiency.

Reuse recommendations:

• Block size should balance representation reuse benefits against computational overhead; 4 to 8 layers per block is generally appropriate.
• A gating mechanism can be applied to block-level residual connections, allowing the model to adaptively determine the degree of reuse.
• This approach is compatible with existing residual connections and can be introduced as an incremental improvement.

4. Summary and Outlook

The three techniques analyzed in this article address three critical dimensions of large-scale model architectures:

Technical Direction	Problem Addressed	Core Mechanism	Reuse Value
Lower-Bounded Decay Parameterization	Numerical instability in diagonal blocks of linear attention	Lower-bounded scaled Sigmoid constraining decay factors	High; minimal engineering changes
Quantile Balancing	MoE load imbalance with difficult-to-tune auxiliary losses	Dynamic quantile-based threshold adjustment for routing	High; no auxiliary loss required
Block-Level Attention Residuals	Representation degradation and memory/communication overhead in deep networks	Block-level attention residual connections	Medium-High; requires block size tuning

A common characteristic of these three techniques is that they do not depend on hardware-specific or framework-exclusive features; instead, they optimize from the algorithmic and numerical levels, providing broad applicability. For teams building or optimizing large-scale model training and inference systems, it is recommended to validate these techniques in the following scenarios:

1. Long-sequence modeling tasks using linear attention or incremental network structures.
2. Large-scale MoE training with dozens or more experts.
3. Deep Transformer architectures exceeding 48 layers.

By combining these techniques, simultaneous improvements in operator efficiency, training stability, and representation capability can be achieved, providing a solid technical foundation for efficient training and deployment of ultra-large-scale models.