🚨 Excited to share our latest paper: “Rethinking Muon Beyond Pretraining: Spectral Failures and High-Pass Remedies for VLA and RLVR”
▸ Paper: https://t.co/qfe8itEWxP
▸ Blog: https://t.co/WMG33IwVrs
▸ Code: https://t.co/C85czsTz74
⚠️ The Question
▸ Muon is a matrix-aware optimizer widely used in LLM pretraining. Through Newton–Schulz (NS) iterations, it pushes singular values of the momentum matrix toward 1, a form of uniform spectral whitening.
▸ LLM pretraining is only one specific setting (text modality, classification loss, supervised learning). We ask whether uniform whitening remains appropriate when we change the modality, the loss, or the learning paradigm.
🔍 Key Findings
Muon exhibits spectral failure in two representative regimes.
▸ VLA training (modality extended to language / vision / action, loss changed from classification to regression): the rank of the action gradient is much lower than those of vision and language. Uniform whitening amplifies noise-dominated tail directions to the scale of informative ones, contaminating the update.
▸ RLVR (learning paradigm changed from supervised learning to reinforcement learning): RL gradients have substantially lower SNR than SFT gradients. Uniform whitening amplifies noise and can cause the policy to collapse within very few training steps.
💡 Method
We propose Pion (sPectral hIgh-pass Optimization on momeNtum), which changes only the polynomial coefficients inside Muon's NS to obtain a high-pass NS iteration.
▸ Pion realizes a spectral high-pass effect: anchoring informative leading singular values near 1, and suppressing noise-dominated tail singular values toward 0.
▸ Pion provides two modes that determine how the high-pass NS iteration is applied: a default mode that applies it to the entire layer matrix, and a per-head mode that applies it independently to each attention head. The per-head mode is critical for RLVR, as it preserves the per-head heterogeneity inherited from pretrained or SFT models.
✅ Results
▸ In simulated VLA, Pion outperforms AdamW and Muon on both ℓ1-regression and flow-matching architectures, e.g., reaching 100% on LIBERO Object in 1,500 steps vs. 32.2% for AdamW. On a real Franka Research 3, Pion also leads on grasp-and-place tasks.
▸ In RLVR (Qwen3-1.7B / 4B × GRPO / GMPO × MATH / GSM8K), Muon collapses across all settings, while Pion trains stably and outperforms AdamW.
📌 Takeaway
▸ Muon’s uniform spectral whitening is highly effective for LLM pretraining, but it can be the wrong inductive bias in low-rank or low-SNR regimes.
▸ Pion replaces uniform whitening with controllable spectral high-pass optimization, making Muon-style optimizers better suited to VLA training and RLVR at zero additional per-step cost.
🤝 Acknowledgements
I am sincerely grateful to my advisor @sijialiu17, and to our collaborators @Mingyi552237 (UMN), Gaowen Liu and Ramana Rao Kompella (Cisco) for their valuable discussions and support.
🎯 Our EMNLP 2025 Main paper
“Reasoning Model Unlearning: Forgetting Traces, Not Just Answers, While Preserving Reasoning Skills” goes live soon!
Catch us on Wednesday in Suzhou at #EMNLP2025 🇨🇳
🔗 Paper Link: https://t.co/6zm3EYbfHk
🏡 Project Page Link: https://t.co/1bBgTJ8qRW
🗓 November 5, 11:00–12:30 CST (UTC+8)
📍 Hall C, Section 2, 500-Main
🧍 I won’t be there in person — but feel free to chat with my co-authors!
🧠 The Problem
You’ve erased sensitive answers from your LRM.
But the reasoning traces, the step-by-step “thoughts” that led there, still remain.
Even after unlearning, the model can reconstruct or re-infer forgotten answers through these traces.
So the question is:
👉 Can we truly forget reasoning traces, while preserving the model’s reasoning ability?
🎯 Our Solution: R²MU (Reasoning-aware Representation Misdirection for Unlearning)
We go beyond answer-level forgetting and target the reasoning process itself.
R²MU suppresses sensitive reasoning traces while maintaining general reasoning competence.
Through representation misdirection, the model unthinks unsafe reasoning paths, while CoT supervision preserves valid reasoning skills.
⚙️ How it Works
🔄 Unthinking Loss: misaligns hidden representations of sensitive reasoning traces with randomized features.
💡 Reasoning Preservation: uses CoT datasets (like LIMO) to retain problem-solving ability.
✅ R²MU erases reasoning traces — not just answers.
✅ Preserves general reasoning and utility across diverse benchmarks.
✅ Achieves the lowest reasoning-trace leakage (RT-UA ↓) on unlearning benchmark WMDP and LRM safety benchmark STAR-1, while maintaining top reasoning accuracy on AIME, MATH-500, and GPQA.
👥 With amazing collaborators from MSU: @ChongyuFan ,@zyh2022 , @jia_jinghan , and my advisor @sijialiu17 .
🙏 Grateful to our IBM collaborators from @MITIBMLab : @NathalieBaraca1 , Dennis Wei, @p_ram_p.
Although I won’t be attending in person, feel free to stop by our poster on Wed, July 16 @ 4:30 pm PT (East Exhibition Hall A-B, E-2803) to chat with my co-authors about the work!
3. We integrate unlearning methods with smoothness optimization and validate their effectiveness through extensive experiments (e.g., on WMDP and MUSE) against both relearning and jailbreaking attacks.
2. Beyond SAM, we explore a suite of smoothness optimization techniques, including gradient penalty, curvature regularization, randomized smoothing, and weight averaging.
🔬Our key insight:
1. We frame relearning as a min–max adversarial game over model weights, and connect it to Sharpness-Aware Minimization (SAM)—which can flatten the loss landscape to make forgetting stick.
🎯 What problem do we tackle?
Current LLM unlearning methods aim to erase sensitive data—but remain vulnerable to relearning attacks, where even a few examples can resurrect forgotten knowledge. So how can we make unlearning truly robust?