Research map / Architectures and efficiency / RNNs and normalization

When, Where and Why to Average Weights?

ICML 2025 · Architectures and efficiency › RNNs and normalization · approach: Efficient network compression and training · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Benchmarks checkpoint weight averaging (LAWA and EMA) on the AlgoPerf suite and a 124M language model, finding it speeds up training (about 12% fewer GPU-hours) and mildly improves generalization, but does not fully replace learning rate decay.

Problem

Whether weight averaging along the training trajectory can reduce training time, improve generalization, and replace learning rate decay for modern deep learning optimizers, evaluated fairly against heavily tuned baselines.

Approach

Evaluates LAWA and EMA weight averaging as add-ons to strong, tuned optimizers (NadamW, Distributed Shampoo, AdamW) across the AlgoPerf workloads, using time-to-target and best-validation metrics, and studies averaging as a proxy for shorter LR decay.

Key results

• LAWA and EMA on NadamW reach the AlgoPerf validation targets in on average about 82% of NadamW's steps, an estimated 12% reduction in GPU-hours (612 to 550/541).
• In the language modeling setting, LAWA and EMA shorten steps-to-target from 3200 (AdamW) to 2240 and 2200, roughly 30% faster.
• Weight averaging alone generally fails to match LR annealing on AlgoPerf, except on Criteo1TB and OGBG where it reaches the target without decaying the learning rate.

Measured effects reported in the paper

Surprising findings

• Averaging alone does not replace LR decay on most AlgoPerf workloads, and gains vanish when the learning rate is fully annealed to zero.
• EMA with update frequency nu=1 performed poorly for beta <= 0.99 in the AlgoPerf experiments, contrary to prior reports.

Limitations

• Analysis restricted mainly to NadamW and AdamW baselines, with Shampoo tested only on three workloads.
• Averaging schemes do not influence updates, so they cannot fully replace LR annealing in the explored variants.
• Language modeling experiments are small scale (124M parameters, 5B tokens).

Experimental setup

Datasets: Criteo 1TB FastMRI ImageNet-1k OGBG WMT LibriSpeech FineWebEdu

Benchmarks: AlgoPerf FineWebEdu held-out loss

Metrics: steps to validation target GPU-hours WER loss MAP SSIM BLEU accuracy validation perplexity

Baselines: NadamW AdamW Distributed Shampoo LAWA EMA

Builds on: AlgoPerf LAWA SWA Polyak-Ruppert averaging Distributed Shampoo schedule-free optimization

Closest papers by meaning

• Why Do We Need Weight Decay in Modern Deep Learning? (NeurIPS 2024 · similarity 0.86)
• Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations (NeurIPS 2024 · similarity 0.86)
• Accelerating neural network training: An analysis of the AlgoPerf competition (ICLR 2025 · similarity 0.85)
• Gradient-Weight Alignment as a Train-Time Proxy for Generalization in Classification Tasks (NeurIPS 2025 · similarity 0.85)
• Feature Averaging: An Implicit Bias of Gradient Descent Leading to Non-Robustness in Neural Networks (ICLR 2025 · similarity 0.85)
• How to set AdamW's weight decay as you scale model and dataset size (ICML 2025 · similarity 0.85)