Paper deep dive
Scaling Muon for Diffusion Transformers
Chenghao Li, Xiao Han, Xinxin Huang, Wei Liu, Boyang Li, Bing Xiao, Heran Zhang, Juanma Perez Rua, Ke Xu, Kangning Liu, Linjun Kuang, Na Li, Tan Wang, Tian Xie, Wei Peng, Yang Pei, Yifan Xu, Yuanhao Zhai, Yuwei Lin, Zhe Wang, Zihao He, Daniel Li, Junbiao Tang, Ziyang Jiang, Dake Chen
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 94%
Last extracted: 8/26/2026, 3:47:22 AM
Summary
This paper investigates the scaling behavior of the Muon optimizer on large Diffusion Transformers (DiTs) ranging from 1.3B to 15B parameters. While Muon offers generative quality advantages over AdamW, its high computational and communication overhead due to Newton-Schulz iterations and full-momentum materialization limits efficiency. The authors propose 'Periodic Row-wise Muon', which reduces overhead by performing full spectral updates periodically and using low-cost row-wise updates otherwise, alongside a distributed implementation with bucketed all-gather. This approach preserves generative quality while significantly reducing optimizer time, end-to-end step time, and communication volume.
Entities (6)
Relation Signals (5)
Periodic Row-wise Muon → improves → end-to-end training efficiency
confidence 95% · These results show that Periodic Row-wise Muon preserves Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.
Periodic Row-wise Muon → reduces → optimizer time
confidence 93% · It reduces optimizer time by 46.9--54.3%
Periodic Row-wise Muon → reduces → logical communication volume
confidence 93% · It reduces ... logical communication volume by 66.7%
Muon → outperforms → AdamW
confidence 90% · Across all scales, Muon improves the best observed generative quality over AdamW by 12.9--19.1%.
Newton-Schulz iteration → causes → Computation Overhead
confidence 88% · the 5-step Newton--Schulz iteration (NS5) ... introduces substantial computation and communication overhead
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:The matrix-aware optimizer Muon improves large model training by balancing updates across singular directions, yet its scaling behavior and end-to-end efficiency on large Diffusion Transformers (DiTs) remain unclear. We first establish Muon's scaling behavior on DiTs from 1.3B to 15B parameters, showing that its optimization and generative quality advantages over AdamW persist across model scales. However, at scale, the 5-step Newton--Schulz iteration (NS5) performed at every optimization step, together with full-momentum materialization, introduces substantial computation and communication overhead that can offset Muon's step-efficiency advantage. We introduce \emph{Periodic Row-wise Muon}, which performs a full NS5 spectral update once every \(K\) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps. We further co-design a distributed implementation that operates directly on sharded momentum during non-refresh steps and accelerates spectral refreshes through bucketed all-gather and communication--computation overlap. Across all scales, Muon improves the best observed generative quality over AdamW by 12.9--19.1\%. Compared with vanilla Muon, Periodic Row-wise Muon remains within 0.5\% in best generative quality on the 1.3B--4B models and improves it by 4.5\% at 9B. It reduces optimizer time by 46.9--54.3\%, end-to-end step time by 15.7--24.3\%, and logical communication volume by 66.7\%, while reaching its respective best generative quality with 33.7--64.8\% less active training time. These results show that Periodic Row-wise Muon preserves Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.
Tags
Links
- Source: https://arxiv.org/abs/2608.20818v2
- Canonical: https://arxiv.org/abs/2608.20818v2
PDF not stored locally. Use the link above to view on the source site.
Full Text
5,621 characters extracted from source content.
Expand or collapse full text
Skip to main content Search Submit Donate Log in Search arXiv Press Enter to search · Advanced search Computer Science > Machine Learning arXiv:2608.20818v2 (cs) This paper has been withdrawn by Chenghao Li [Submitted on 21 Aug 2026 (v1), last revised 24 Aug 2026 (this version, v2)] Title:Scaling Muon for Diffusion Transformers Authors:Chenghao Li, Xiao Han, Xinxin Huang, Wei Liu, Boyang Li, Bing Xiao, Heran Zhang, Juanma Perez Rua, Ke Xu, Kangning Liu, Linjun Kuang, Na Li, Tan Wang, Tian Xie, Wei Peng, Yang Pei, Yifan Xu, Yuanhao Zhai, Yuwei Lin, Zhe Wang, Zihao He, Daniel Li, Junbiao Tang, Ziyang Jiang, Dake Chen View a PDF of the paper titled Scaling Muon for Diffusion Transformers, by Chenghao Li and 24 other authors No PDF available, click to view other formats Abstract:The matrix-aware optimizer Muon improves large model training by balancing updates across singular directions, yet its scaling behavior and end-to-end efficiency on large Diffusion Transformers (DiTs) remain unclear. We first establish Muon's scaling behavior on DiTs from 1.3B to 15B parameters, showing that its optimization and generative quality advantages over AdamW persist across model scales. However, at scale, the 5-step Newton--Schulz iteration (NS5) performed at every optimization step, together with full-momentum materialization, introduces substantial computation and communication overhead that can offset Muon's step-efficiency advantage. We introduce Periodic Row-wise Muon, which performs a full NS5 spectral update once every \(K\) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps. We further co-design a distributed implementation that operates directly on sharded momentum during non-refresh steps and accelerates spectral refreshes through bucketed all-gather and communication--computation overlap. Across all scales, Muon improves the best observed generative quality over AdamW by 12.9--19.1\%. Compared with vanilla Muon, Periodic Row-wise Muon remains within 0.5\% in best generative quality on the 1.3B--4B models and improves it by 4.5\% at 9B. It reduces optimizer time by 46.9--54.3\%, end-to-end step time by 15.7--24.3\%, and logical communication volume by 66.7\%, while reaching its respective best generative quality with 33.7--64.8\% less active training time. These results show that Periodic Row-wise Muon preserves Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs. Comments: This version has to be withdrawn due to the incomplete publication privacy review required by the company. After completing the privacy review, the newer version will be updated Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computer Vision and Pattern Recognition (cs.CV) Cite as: arXiv:2608.20818 [cs.LG] (or arXiv:2608.20818v2 [cs.LG] for this version) https://doi.org/10.48550/arXiv.2608.20818 Focus to learn more arXiv-issued DOI via DataCite Submission history From: Chenghao Li [view email] [v1] Fri, 21 Aug 2026 07:40:28 UTC (998 KB) [v2] Mon, 24 Aug 2026 18:34:57 UTC (1 KB) (withdrawn) Full-text links: Access Paper: View a PDF of the paper titled Scaling Muon for Diffusion Transformers, by Chenghao Li and 24 other authorsWithdrawn No license for this version due to withdrawn Current browse context: cs.LG < prev | next > new | recent | 2026-08 Change to browse by: cs cs.AI cs.CV References & Citations NASA ADSGoogle Scholar Semantic Scholar export BibTeX citation Loading... BibTeX formatted citation × loading... Data provided by: Bookmark Bibliographic Tools Bibliographic and Citation Tools Bibliographic Explorer Toggle Bibliographic Explorer (What is the Explorer?) Connected Papers Toggle Connected Papers (What is Connected Papers?) Litmaps Toggle Litmaps (What is Litmaps?) scite.ai Toggle scite Smart Citations (What are Smart Citations?) Code, Data, Media Code, Data and Media Associated with this Article alphaXiv Toggle alphaXiv (What is alphaXiv?) Links to Code Toggle CatalyzeX Code Finder for Papers (What is CatalyzeX?) DagsHub Toggle DagsHub (What is DagsHub?) GotitPub Toggle Gotit.pub (What is GotitPub?) Huggingface Toggle Hugging Face (What is Huggingface?) ScienceCast Toggle ScienceCast (What is ScienceCast?) Demos Demos Replicate Toggle Replicate (What is Replicate?) Spaces Toggle Hugging Face Spaces (What is Spaces?) Spaces Toggle TXYZ.AI (What is TXYZ.AI?) Related Papers Recommenders and Search Tools Link to Influence Flower Influence Flower (What are Influence Flowers?) Core recommender toggle CORE Recommender (What is CORE?) IArxiv recommender toggle IArxiv Recommender (What is IArxiv?) Author Venue Institution Topic About arXivLabs arXivLabs: experimental projects with community collaborators arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website. Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them. Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs. Which authors of this paper are endorsers? | Disable MathJax (What is MathJax?) mathjaxToggle(); We gratefully acknowledge support from our major funders, member institutions, , and all contributors. About · Help · Contact · Subscribe · Copyright · Privacy · Accessibility · Operational Status (opens in new tab) Major funding support from