Architecture Breakdown of a Trillion-Parameter Sparse Multimodal Model

Introduction: Three Core Tensions in Architecture Design

The design goals of this class of models can be summarized as the reconciliation of three fundamental contradictions:

1. Ultra-long context vs. computational efficiency: With million-token sequences, standard full attention incurs quadratic memory and compute costs. Linear attention variants reduce complexity to near-linear, but introduce numerical precision and hardware utilization challenges.
2. Extreme sparsity vs. training stability: With total parameters in the trillions and activated parameters in the hundreds of billions, the number of routed experts is massive and the activation ratio is extremely low. This makes expert load balancing difficult; traditional auxiliary loss functions converge slowly and are sensitive to hyperparameter tuning at this scale.
3. Deep stacking vs. representation degradation: Increasing network depth to boost capability leads to information decay across layers, gradient propagation difficulties, and higher memory/communication overhead.

The following sections address each of these tensions through four key design mechanisms: lower-bounded decay parameterization for linear attention, gated multi-head latent attention with block-level attention residuals, stable latent expert routing with quantile-based load balancing, and a multi-stage training strategy combining long-context pretraining with agentic reinforcement learning.

Attention Layer: Numerical Stabilization and Hardware Affinity for Linear Attention

The Diagonal Block Problem in Linear Attention

Linear attention and Delta Net-style methods achieve near-linear complexity through recurrent state updates, typically computed in chunks. Within a chunk, the inverse of the cumulative decay factor Γ participates in the computation of the diagonal block. When decay accumulates significantly, this inverse is prone to overflow or underflow in lower precision formats like BF16. To avoid numerical issues, implementations often force the diagonal block to take an explicit position-pair path, bypassing dense matrix multiplication and thus underutilizing Tensor Cores.

Lower-Bounded Decay Parameterization

An effective solution is to change the log-decay from an unbounded negative Softplus form to a scaled Sigmoid form with a fixed lower bound (e.g., a constant around -5). This constrains the decay parameter to a safe numerical range. The practical effect is that within any fixed-size tile (e.g., 16 tokens), the inverse scaling factor remains within a stable interval.

# Illustration: lower-bounded decay parameterization
def lower_bounded_decay(e_A_h_z_th, g_min=-5.0):
    # e_A_h_z_th: unconstrained decay logits
    return g_min * torch.sigmoid(e_A_h_z_th)

The engineering value of this change is significant:

• Eliminates special branches: Both diagonal and off-diagonal blocks can use dense Tensor Core GEMM operations directly, removing the need for a separate position-pair path for the diagonal block and simplifying kernel implementation.
• Improved throughput: All causal tiles utilize a unified, high-throughput matrix multiplication path, avoiding warp divergence and extra memory access from branching.
• Natural support for extrapolation without position encodings: Since the decay mechanism inherently encodes relative position information, the model can extrapolate to million-token contexts without explicit position encodings (NoPE), reducing distribution shift issues.

Notably, this change requires minimal code modification, making it a “low-change, high-benefit” optimization suitable for incremental adoption in existing linear attention implementations.

Cross-Layer Information Sharing: Gated Latent Attention and Block-Level Attention Residuals

Gated Multi-Head Latent Attention

Within the attention module, a gating mechanism is used to selectively aggregate multi-head latent representations, increasing expressive capacity without significantly increasing parameter count. The low-rank compression of the latent representations reduces KV cache overhead, while the gating allows the model to dynamically adjust the contribution of each head based on input content, balancing efficiency and flexibility. For million-token contexts, effective compression of the KV cache is a prerequisite for inference within limited memory.

Block-Level Attention Residuals

A common issue in deep Transformers is the depth representation bottleneck: as layers increase, outputs of adjacent layers become highly similar, reducing cross-layer information flow. Block AttnRes (Block Attention Residuals) addresses this by introducing residual-style attention aggregation between blocks composed of several layers, enabling representation sharing and recombination across depths.

The benefits are twofold:

• Mitigates representation degradation: Cross-block connections provide shorter paths for gradient and information flow, maintaining representational diversity in deep networks.
• Reduces memory and communication overhead: Organizing residuals at the block level, rather than per-layer full residuals, reduces the intermediate activations that need to be retained and the amount of cross-device communication, which is particularly critical for large-scale distributed training.

Expert Layer: Stability Design for Ultra-Sparse Routing

Stable Latent Expert Architecture

When the number of routed experts reaches the hundreds and the activation ratio is pushed to extremely low levels, the variance of routing decisions increases significantly. Using latent-space expert representations (i.e., organizing expert parameters in a low-rank latent form) helps control total parameter count while maintaining expert diversity. This is complemented by gated activation functions (such as SiTU-GLU, a smooth gating unit) to improve gradient behavior in the early stages of training.

Quantile-Based Load Balancing

Traditional MoE training relies on auxiliary loss functions to enforce load balancing among experts. However, under ultra-large-scale and extremely sparse conditions, these auxiliary losses have notable drawbacks: sensitivity to weight coefficients, potential conflicts with the primary task loss, and slow convergence.

The Quantile Balancing (QB) mechanism offers an alternative that eliminates the auxiliary loss. It uses the quantiles of routing scores as thresholds to assign tokens, dynamically constraining the number of tokens each expert receives to a target range. Its advantages over auxiliary losses are:

• No extra loss term: It avoids introducing an auxiliary objective that requires hyperparameter tuning and could interfere with the primary language modeling objective.
• Dynamic adaptation: Thresholds adjust in real-time based on the score distribution, making it more robust to distribution drift during training.
• Simple implementation: The core operation is quantile statistics and threshold clipping on routing scores, which is easy to implement in distributed training frameworks.

# Illustration: routing assignment based on quantile threshold
def quantile_balancing_route(scores, capacity_ratio):
    # scores: [num_tokens, num_experts]
    threshold = torch.quantile(scores, 1.0 - capacity_ratio, dim=0)
    mask = scores >= threshold
    return mask

It is important to note that the code above is for illustrative purposes only. A production implementation must account for specific parallelization schemes, communication topology, and capacity factor settings.

Training Pipeline: Multi-Stage Context Expansion and Agentic Post-Training

Multi-Stage Pretraining

Million-token context capability cannot be achieved in a single step. The standard approach is progressive context extension: starting with short sequences of a few thousand tokens and gradually increasing to tens of thousands, hundreds of thousands, and finally up to a million tokens. Each stage corresponds to different data compositions and parallelization strategies:

Stage	Context Length	Primary Objective
Stage 1	8K	Basic language capability and multimodal alignment
Stage 2	64K	Medium-length document understanding and retrieval
Stage 3	256K	Long-document reasoning and cross-paragraph integration
Stage 4	1M	Ultra-long context extrapolation and long-range dependency modeling

The advantage of progressive expansion is that short-sequence stages can use higher parallelism and larger batch sizes for greater throughput, while long-sequence stages focus on optimizing memory usage and communication efficiency. This makes the overall training cost more manageable.

Agentic Reinforcement Learning Post-Training

Following pretraining, reinforcement learning (RL) post-training for Agent scenarios is crucial for enhancing the model’s ability to “act within long contexts.” Unlike general conversational post-training, agentic tasks often involve multi-step tool calls, state tracking, and long-horizon planning, with trajectories potentially reaching hundreds of thousands of tokens. This type of post-training needs to address:

• Credit assignment in long trajectories: How to attribute the final outcome to early decisions under sparse rewards.
• Context budget management: Preserving key intermediate states within a limited window to avoid relevant history being overwhelmed.
• Training stability: Managing variance and sampling efficiency for RL on long sequences.

Summary of Reusable Technical Patterns

The following three design choices offer significant value for reuse across projects:

1. Lower-bounded decay parameterization: Achieves simplified kernel implementation and higher throughput for linear attention operators with minimal code changes, while naturally supporting position-encoding-free long-context extrapolation.
2. Quantile-based load balancing: Replaces auxiliary losses for load balancing in ultra-large-scale sparse MoE models, reducing tuning costs and improving training stability, suitable for scenarios with a high number of experts and a low activation ratio.
3. Block-level attention residuals: Organizes cross-layer residuals at the block level, mitigating representation degradation in deep networks while controlling memory and communication overhead.

These three designs address the core challenges of long-context computational efficiency, training stability for massive expert counts, and deep representation bottlenecks. They constitute the key technical pillars for the engineering implementation of trillion-parameter sparse multimodal models. For teams working on advancing large model architectures, the underlying principle is more instructive than the specific implementations: seek structural solutions at the intersection of numerical precision, hardware affinity, and training stability, rather than relying on post-hoc parameter tuning as a remedy.