Practical Efficiency of Muon for Pretraining

• URL: http://arxiv.org/abs/2505.02222v4
• Date: Tue, 20 May 2025 01:04:35 GMT
• Title: Practical Efficiency of Muon for Pretraining
• Authors: Essential AI, :, Ishaan Shah, Anthony M. Polloreno, Karl Stratos, Philip Monk, Adarsh Chaluvaraju, Andrew Hojel, Andrew Ma, Anil Thomas, Ashish Tanwer, Darsh J Shah, Khoi Nguyen, Kurt Smith, Michael Callahan, Michael Pust, Mohit Parmar, Peter Rushton, Platon Mazarakis, Ritvik Kapila, Saurabh Srivastava, Somanshu Singla, Tim Romanski, Yash Vanjani, Ashish Vaswani,
• Abstract summary: We find that Muon is more effective than AdamW in retaining data efficiency at large batch sizes.<n>We present a simple algorithm that accounts for all sources of error in muP while introducing only a modest overhead in resources.
• Score: 13.914926836677648
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We demonstrate that Muon, the simplest instantiation of a second-order optimizer, explicitly expands the Pareto frontier over AdamW on the compute-time tradeoff. We find that Muon is more effective than AdamW in retaining data efficiency at large batch sizes, far beyond the so-called critical batch size, while remaining computationally efficient, thus enabling more economical training. We study the combination of Muon and the maximal update parameterization (muP) for efficient hyperparameter transfer and present a simple telescoping algorithm that accounts for all sources of error in muP while introducing only a modest overhead in resources. We validate our findings through extensive experiments with model sizes up to four billion parameters and ablations on the data distribution and architecture.

Related papers

• Muon is Scalable for LLM Training [50.68746986439438]
  We introduce Moonlight, a Mixture-of-Expert (MoE) model trained with 5.7T tokens using Muon.<n>Our model improves the current frontier, achieving better performance with much fewer training FLOPs compared to prior models.<n>We open-source our distributed Muon implementation that is memory optimal and communication efficient.
  arXiv  Detail & Related papers  (2025-02-24T09:12:29Z)
• Adaptive Data Exploitation in Deep Reinforcement Learning [50.53705050673944]
  We introduce ADEPT, a powerful framework to enhance the **data efficiency** and **generalization** in deep reinforcement learning (RL)<n>Specifically, ADEPT adaptively manages the use of sampled data across different learning stages via multi-armed bandit (MAB) algorithms.<n>We test ADEPT on benchmarks including Procgen, MiniGrid, and PyBullet.
  arXiv  Detail & Related papers  (2025-01-22T04:01:17Z)
• EPS-MoE: Expert Pipeline Scheduler for Cost-Efficient MoE Inference [49.94169109038806]
  This paper introduces EPS-MoE, a novel expert pipeline scheduler for MoE that surpasses the existing parallelism schemes.<n>Our results demonstrate at most 52.4% improvement in prefill throughput compared to existing parallel inference methods.
  arXiv  Detail & Related papers  (2024-10-16T05:17:49Z)
• Pareto Low-Rank Adapters: Efficient Multi-Task Learning with Preferences [49.14535254003683]
  We introduce PaLoRA, a novel parameter-efficient method that addresses multi-task trade-offs in machine learning.<n>Our experiments show that PaLoRA outperforms state-of-the-art MTL and PFL baselines across various datasets.
  arXiv  Detail & Related papers  (2024-07-10T21:25:51Z)
• Towards Efficient Pareto Set Approximation via Mixture of Experts Based Model Fusion [53.33473557562837]
  Solving multi-objective optimization problems for large deep neural networks is a challenging task due to the complexity of the loss landscape and the expensive computational cost. We propose a practical and scalable approach to solve this problem via mixture of experts (MoE) based model fusion. By ensembling the weights of specialized single-task models, the MoE module can effectively capture the trade-offs between multiple objectives.
  arXiv  Detail & Related papers  (2024-06-14T07:16:18Z)
• SPT: Fine-Tuning Transformer-based Language Models Efficiently with Sparsification [14.559316921646356]
  Fine-tuning Transformer-based models for downstream tasks has long running time and high memory consumption. We propose the SPT system to fine-tune Transformer-based models efficiently by introducing sparsity. SPT consistently outperforms well-optimized baselines, reducing the peak memory consumption by up to 50% and accelerating fine-tuning by up to 2.2x.
  arXiv  Detail & Related papers  (2023-12-16T07:44:52Z)
• Critical Bach Size Minimizes Stochastic First-Order Oracle Complexity of Deep Learning Optimizer using Hyperparameters Close to One [0.0]
  We show that deep learnings using small constant learning rates, hyper parameters close to one, and large batch sizes can find the model parameters of deep neural networks that minimize the loss functions. Results indicate that Adam using a small constant learning rate, hyper parameters close to one, and the critical batch size minimizing SFO complexity has faster convergence than Momentum and gradient descent.
  arXiv  Detail & Related papers  (2022-08-21T06:11:23Z)
• MoEfication: Conditional Computation of Transformer Models for Efficient Inference [66.56994436947441]
  Transformer-based pre-trained language models can achieve superior performance on most NLP tasks due to large parameter capacity, but also lead to huge computation cost. We explore to accelerate large-model inference by conditional computation based on the sparse activation phenomenon. We propose to transform a large model into its mixture-of-experts (MoE) version with equal model size, namely MoEfication.
  arXiv  Detail & Related papers  (2021-10-05T02:14:38Z)
• CPM-2: Large-scale Cost-effective Pre-trained Language Models [71.59893315671997]
  We present a suite of cost-effective techniques for the use of PLMs to deal with the efficiency issues of pre-training, fine-tuning, and inference. We introduce knowledge inheritance to accelerate the pre-training process by exploiting existing PLMs instead of training models from scratch. We implement a new inference toolkit, namely InfMoE, for using large-scale PLMs with limited computational resources.
  arXiv  Detail & Related papers  (2021-06-20T15:43:54Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.