Papers for

large ai model trainers

Papers whose findings have a practical use for this group, as judged from the abstract. Open a paper to read what it means in practice.

EnergyCal reduces training loss spikes in sparse moes with muon optimizers

EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training

Abstract: Sparse mixture-of-experts (MoE) language-model training can exhibit large loss spikes with Muon-family optimizers. We introduce EnergyCal, a lightweight rule that attenuates routed-expert updates according to the normalized squared Frobenius norm of their finite-step Newton-Schulz directions. EnergyCal leaves the spectral transform unchanged and adds no persistent optimizer state. We evaluate it on a 4B-parameter MoE model. Against a development-selected, lower-learning-rate MuonH recipe, EnergyCal-MuonH reduces mean thresholded loss-spike burden by 99.51% and final geometric-mean perplexity by 4.06%, with a modest throughput cost. With plain Muon at matched nominal learning rate and weight decay, the reduction in mean spike burden averages 95.98% across the two tested exponents, accompanied by small endpoint improvements. These findings support stabilization both with and without Hyperball in the evaluated setting.

Tue 27 OctMachine LearningArtificial IntelligenceComputation and Language

The gist

Training large language models with many experts can sometimes cause sudden jumps in error, making the learning unstable. The authors propose EnergyCal, a simple method that adjusts updates during training to avoid these jumps without adding extra memory or changing key calculations. They tested EnergyCal on a 4-billion-parameter model and found it drastically reduces error spikes and improves final accuracy with only a small speed cost. This makes training these complex models more stable and efficient.

Open → 2610.37062v1

KITE method improves large language model scaling efficiency

KITE: KV-Invariant Transformer Expansion for Efficient Agentic LLM Scaling

Abstract: Scaling a language model is not only a question of final quality: the architectural choice determines how much computation is spent during training, prompt processing, and autoregressive decoding to achieve certain model quality. An ideal model architecture should lower all above computation costs to facilitate scaling to a larger model, while ensure the larger model indeed outperforms smaller baselines. We introduce KV-Invariant Transformer Expansion (KITE), a scaling paradigm that achieves this goal. It trains the model from a smaller size to a larger size (i.e., saving training costs via upcycling), while places newly added parameters in regions that do not affect attention KV. Consequently, during inference, prefilling KV only relies on the smaller part of the model, so the inference costs are saved. As a concrete instantiation, we present Step Scale Transformer (SST), a two-tower decoder in which one tower produces KV and the other reads them. At comparable cumulative training compute, SST, a 67B MoE model with 2.15B active body parameters per decode token, achieves lower training loss than 47B and 63B MoE Transformers with 1.48B and 2.02B active body parameters, respectively, while reducing estimated inference cost by 6.7% and 31.6%.

Wed 23 SeptMachine LearningArtificial Intelligence

The gist

Training bigger language models usually costs a lot of computer power and time. The authors introduce KITE, a way to grow a model from a smaller size by adding new parts that don’t slow down how the model understands context. This means training can be cheaper and the model runs faster when answering questions. Their specific design, the Step Scale Transformer, showed better learning with less inference cost compared to similar models.

Open → 2609.27294v1

Asynchronous reinforcement learning improves with stale critic correction

BRACE: Anchored Bellman-Residual Correction for Stale Critics in Asynchronous RL

Abstract: Asynchronous reinforcement learning has become the standard way to scale training for language models, but the resulting policy lag biases the critic toward the stale behavior policy. Existing work on asynchronous LLM training corrects the actor and leaves this bias unaddressed, while the off-policy value correction of classical RL does not carry over to long-horizon agentic tasks, since a short correction horizon leaves the regression target free of the reward and a long one lets the product of importance ratios drift exponentially with the trajectory length. We propose BRACE, an anchored Bellman-residual correction for stale value models. BRACE bounds the correction horizon to a prefix of policy tokens and anchors a constant-weight Monte-Carlo tail beyond it, which separates policy correction from reward propagation. BRACE improves mean@1 on BrowseComp-Plus by $2.4\%$ over the strongest baseline, runs $2.46\times$ faster per step than synchronous training, and remains stable $50$ updates off-policy.

Wed 9 SeptMachine LearningArtificial Intelligence

The gist

Training language models with asynchronous reinforcement learning can cause delays that make the system's value estimates outdated, leading to less accurate learning. The authors propose a method called BRACE that fixes these outdated value estimates by combining a limited lookahead correction with a steady long-term value estimate. This approach separates adjusting the policy from tracking rewards, making training both faster and more stable. Their tests show that BRACE improves performance on a browsing task while running significantly faster than traditional methods.

Open → 2609.09783v1