Training Ultra-Large Sparse Models for Million-Token Context and Multimodal Input: Key Challenges and Solutions

Background: Why Ultra-Long Context and Ultra-Large Sparse Models Are Hard to Combine

Two trends dominate the current evolution of large models. Context windows keep expanding—from a few thousand tokens in early systems to hundreds of thousands or even millions of tokens. Meanwhile, parameter counts continue to grow, but constrained by compute and memory, more and more systems adopt sparse MoE architectures that activate only a subset of parameters.

When these two trends combine, three classes of problems emerge.

First, attention computation is no longer just “slow to compute”—it becomes “hard to parallelize stably.”
When sequence length reaches hundreds of thousands or millions of tokens, the memory footprint and computational complexity of conventional Softmax attention rise sharply. Linear attention mitigates the complexity problem, but it introduces new bottlenecks in low-precision training, chunked parallelism, and numerical stability.

Second, load balancing across a massive expert pool can easily destabilize training.
If an MoE layer contains hundreds or even nearly a thousand routed experts while each token activates only a small fraction of them, load differences between experts get amplified. A few experts are hit constantly while the rest remain undertrained, ultimately wasting model capacity and destabilizing convergence.

Third, the representational power of deep Transformers cannot be improved by simply stacking more layers.
As model depth grows, information flow between layers slows down and optimization becomes harder. Merely increasing width or depth does not linearly improve long-range reasoning; instead, it adds communication and memory pressure.

A system targeting million-token context, multimodal input, and ultra-large sparse parameters therefore needs coordinated adjustments to attention structure, MoE routing, inter-layer information sharing, and training scheduling.

Overall Architecture: Multimodal Long-Context Design with Sparse Activation

One viable approach is to decouple total MoE parameters from activated parameters. For example, total model size can be designed above the trillion-parameter level while each token activates only tens of billions of parameters, striking a balance between training and inference cost.

To natively support multimodal input such as vision, the encoding side should not simply convert images into text descriptions before feeding them to the language model. Instead, visual signals should be mapped into a representation space shared with text and passed together into the Transformer backbone, so that downstream attention and MoE layers handle all modalities consistently.

For context extension, staged pretraining is more stable than stretching sequence length in one shot. A typical progression starts at 8K, moves to 64K and 256K, and finally reaches the 1M-token level. Each stage requires re-tuning RoPE or the positional encoding scheme together with data-mix optimization; otherwise positional awareness over long sequences tends to degrade.

The benefits of staged extension show up in three ways:

1. The model first builds stable language and vision foundations on short-to-medium contexts;
2. Long-context attention patterns and retrieval behavior can gradually generalize on top of existing representations;
3. The sparse routing mechanism can progressively adapt to more dispersed token distributions as sequences get longer.

Attention Module: Fusing Incremental Attention with Gated Latent Attention

In long-context scenarios, no single attention form satisfies both quality and efficiency. A practical solution is to combine two mechanisms inside the attention module: one that incrementally compresses historical information, and one that performs precise interaction within the current window.

Incremental Attention

Incremental attention can be understood as a form of linear attention with state updates. While processing the current segment, the model not only attends to the current tokens but also maintains a recurrently updatable historical state. Early information in a long sequence then does not need to be recomputed—it participates in current representations as a compressed state.

This structure suits very long-range dependencies and also reduces KV cache pressure during long-sequence inference.

Gated Multi-Head Latent Attention

Latent attention maps keys and values into a low-dimensional latent space, cutting memory usage. Adding a gating mechanism lets the model dynamically decide how much information to retain based on the input content, rather than treating every position equally. This yields a better balance between compression efficiency and information completeness.

Once incremental attention and gated latent attention are fused, the model gains both long-range state memory and precise modeling of local context. The key to a clean fusion is preventing the two attention branches from interfering with each other, which typically calls for independent normalization and gating parameters.

Lower-Bounded Decay Parameterization: Stable Chunked Linear Attention at Low Precision

A frequently overlooked engineering issue with linear attention is that when low-precision floating-point formats are used, the recurrent decay factor can underflow. This is especially problematic in chunked computation, where the numerical ranges of diagonal and off-diagonal blocks can differ enormously.

An effective fix is to give the decay parameter a lower bound—for example, replacing the unbounded negative softplus of log-decay with a lower-bounded, scaled sigmoid:

g = g_min * sigmoid(parameter)

where g_min can be set to -5. With this bound, the reciprocal scaling factors within any small chunk stay in a safe range and no longer fail due to floating-point underflow.

The change carries substantial engineering value because:

• Diagonal and off-diagonal blocks can both go through uniform dense matrix multiplication;
• Special branches in linear attention implementations are reduced;
• High throughput is easier to achieve on GPU Tensor Cores;
• Combined with a no-position-encoding strategy, million-token extrapolation becomes easier.

MoE Layers: Stable Routing and Load Balancing at High Expert Counts

Stable Latent MoE Architecture

When the expert count reaches the hundreds, the conventional implementation—every expert as an independent feed-forward network—drives memory usage too high. A latent expert mechanism can help: map expert weights into a low-dimensional latent space and expand them at computation time. This preserves model capacity while easing storage pressure.

In addition, the more experts there are, the more routing stability matters. If routing scores fluctuate sharply, the same token may be assigned to entirely different expert combinations at different training steps, preventing experts from developing stable specializations.

Quantile-Based Balancing

Conventional MoE training often relies on auxiliary losses to encourage load balancing, but auxiliary losses alter the primary objective and can easily cause routing degradation in large-scale sparse models. An alternative is to adjust dynamically based directly on quantiles of the routing scores.

The quantile-balancing procedure works as follows:

1. Track the routing-score distribution of each expert;
2. Compute quantile thresholds;
3. Cap or redistribute routing scores above the thresholds;
4. Introduce no additional auxiliary loss.

The approach is simple to implement and stable to train, and it maintains high expert utilization even at the scale of hundreds of billions of activated parameters.

Activation Function Design

In sparse expert networks, the choice of activation function also affects training stability. Using an activation structure with gated linear unit (GLU) characteristics keeps some dimensions on a linear path while others undergo nonlinear transformation. This preserves gradient flow while strengthening the nonlinear expressiveness of experts.

Inter-Layer Information Sharing: Attention Residuals

Deep Transformers commonly hit a representational bottleneck: the information difference between high-layer and low-layer outputs shrinks, and optimization settles into a slow region. One remedy is to add residual connections between attention modules of different layers, allowing part of the attention output to propagate across layers.

This “block-level attention residual” differs from an ordinary residual connection. Ordinary residuals operate only between adjacent layers, whereas block-level residuals let the attention output of one layer directly participate in information aggregation in several subsequent layers.

The benefits include:

• Mitigating gradient vanishing in deep models;
• Reducing the cost of high-level attention redundantly re-modeling low-level details;
• Cutting part of the communication volume in distributed training;
• Improving hierarchical representations for long-range reasoning.

During implementation, the residual proportion matters. Rather than adding outputs directly, a learnable scaling coefficient typically controls how much cross-layer information contributes.

Training Strategy: Multi-Stage Pretraining and Long-Horizon Agent RL

Multi-Stage Context Pretraining

Extending context from 8K to 1M tokens is not just a configuration change; data ordering, batch size, and learning rate must be adjusted in sync. Short-context stages should prioritize consolidating basic semantics, factual knowledge, and vision-language alignment; long-context stages should progressively introduce long documents, multi-document linkage, cross-section reasoning, and long time-series data.

Agent-Style Long-Horizon Reinforcement Learning

Once the model reaches 1M-token context, agent-style reinforcement learning tasks can be introduced—for example, having the model complete multi-step decision-making, tool calling, information retrieval, and path planning within long environments. The learning objective should cover not only whether the final answer is correct, but also whether intermediate steps are effective, whether retrieval is repeated unnecessarily, and whether key context is missed.

This kind of training improves the model’s practical reasoning ability over ultra-long contexts, rather than merely keeping perplexity low on synthetic long sequences.

Combining Multimodality with Long Context

Visual information plays more than an “extra input” role in long contexts. For instance, when images, tables, or diagrams are inserted into long documents, the model must learn cross-modal localization, referencing, and comparison. Multimodal data therefore should not appear only in the short-context stage; it should run through the entire training cycle.

Engineering Takeaways and Recommendations

For teams building or evaluating ultra-long-context sparse models, the following points are a good starting point:

Attention implementation: Verify the chunked-computation stability of linear attention under low precision first, and consider imposing a lower bound on decay parameters.

MoE routing: When expert counts are high, avoid over-reliance on auxiliary losses and evaluate quantile-based dynamic balancing instead.

Inter-layer structure: Experiment with block-level attention residuals in deep models and observe the impact on long-range reasoning and memory usage.

Training schedule: Adopt staged context extension instead of training on ultra-long sequences directly, and re-evaluate the data mix and positional encoding strategy at every stage.

Evaluation: Beyond perplexity, add tasks such as long-document retrieval, multi-step reasoning, and cross-modal localization to reflect real-world usage.

Overall, combining ultra-large sparse models with million-token context is no longer just a matter of stacking parameters. It is a systems-engineering effort spanning numerical computation, routing strategy, inter-layer communication, and training scheduling. Only by co-optimizing attention, MoE, and training strategy can teams obtain stable, usable long-range multimodal capability at a controllable cost.