BMuon: Budgeting Muon's Spectral Geometry for Structured Compressibility

1 Introduction

• Structured pruning removes whole computational units, and MLP hidden channels are a natural compression target because feed-forward modules can account for roughly two thirds of Transformer parameters.Muon’s spectral geometry alone does not prevent sharp degradation under structured MLP pruning; the reported ViT-S/16 comparison contrasts collapse after 30% pruning with BMuon retaining near-unpruned accuracy after 70%.Source: §1 · p. 2
• The BMuon oracle is normalized steepest descent under a norm family that recovers Muon at specified endpoints and includes structured channel-sparse geometries.At α = 1, the unit ball is the convex hull of operator-norm-bounded updates supported on at most K coupled hidden channels.Source: §1 · p. 3
• Across CIFAR-100, ImageNet-1K, and a 1B-parameter LLaMA-style model, BMuon improves structured MLP-channel compressibility over Muon; learned allocations also guide pruning and identify useful non-uniform architectures.The non-uniform architectures outperform matched uniform-width architectures when retrained from scratch.Source: §1 · pp. 3–4
• Because solving the exact global allocation problem at every training step is expensive, the paper derives a minorization–maximization surrogate for a tractable allocation update.The practical procedure follows the allocation update with the corresponding weighted Muon update.Source: §1 · pp. 3–4
• BMuon couples each MLP channel’s up- and down-projection slices and reweights Muon’s spectral geometry so channels across blocks compete for one global hidden-channel budget.The budget asks where update capacity should be spent to achieve the greatest first-order loss decrease.Source: §1 · pp. 2, 4
• With a sufficiently scarce budget, every minimum-cost BMuon factorization is width-minimal, whereas Muon’s minimum cost can also be attained with redundant hidden width.The result concerns exact two-matrix factorizations of the same linear predictor.Source: §1 · pp. 3–4

2.1 Setup and Muon Preliminaries

• Muon optimizes coupled MLP matrix updates using spectral-norm constraints, while pruning removes each hidden channel’s incoming and outgoing weights together.Global pruning ranks channels across blocks jointly, so each block’s retained width depends on the shared ranking.Source: §2.1 · pp. 4–5

2.2 The BMuon Oracle

• BMuon jointly chooses a shared-budget channel allocation and update direction to maximize alignment with the momentum.For a fixed allocation, channel-wise reweighting couples the corresponding row of the up-projection update and column of the down-projection update.Source: §2.2 · p. 5

Allocation curvature.

• For 0 < α ≤ 1/2, the allocation objective is concave; at α = 1 it is convex, and for 1/2 < α < 1 it is generally neither.With 0 < α ≤ 1/2, optimal allocations may be fractional; at α = 1 and integer K, a maximizer can be chosen to select exactly K channels.Source: Allocation curvature · p. 6

2.3 Induced Optimization Geometry

• The BMuon oracle is normalized steepest descent under the norm Ωα,K, whose unit ball is the convex hull of updates feasible under the shared budget.This induced norm is the Minkowski functional of that convex unit ball.Source: §2.3 · p. 6

2.4 From Muon to Sparse Geometries

• The norm family interpolates between Muon’s spectral geometry and structured channel-sparse geometries.It recovers Muon when α = 0 or K = N; at α = 1, the unit ball is the convex hull of operator-norm-bounded updates supported on at most K coupled hidden channels.Source: §2.4 · pp. 6–7

• For linear two-factor blocks realizing nonzero matrices Xℓ, each exact factorization requires at least rank(Xℓ) active channels in block ℓ.The analysis uses blocks with hidden widths mℓ ≥ rank(Xℓ).Source: §3 · p. 7
• When K ≤ d_eff(α, X), every minimum-cost BMuon factorization has active width rank(Xℓ) in every block.Both representation-cost infima are attained.Source: §3 · p. 7
• Muon’s minimum representation cost can be attained at full hidden width, whereas BMuon penalizes redundant width when K < d_eff(α, X).The budget raises representation cost only below this threshold, and d_eff(α, X) is nondecreasing in α.Source: §3 · p. 8

• BMuon replaces the expensive exact allocation at every optimizer step with a minorization–maximization surrogate for fixed momenta and exponent α.After the allocation step, polar factors are recomputed using the updated allocation to form the parameter update.Source: §4 · p. 8
• For 0 < α < 1, the floored surrogate is strictly concave and has a unique maximizer when all channel scores are positive.The KKT conditions yield a clipped allocation update, and the MM step introduces no additional step-size hyperparameter.Source: §4 · p. 8
• For 0 < α ≤ 1/2, repeated MM steps converge in objective value to the global optimum over the floored budget set.Each step is monotone, and every accumulation point of the allocation sequence is globally optimal.Source: §4 · p. 8
• Experiments use α = 1/2, update allocations once every five optimizer steps, and report no measurable increase in overall training time on ImageNet-1K or the 1B-parameter LLaMA model.The practical update combines finite Newton–Schulz polar approximations with one MM allocation step per update period.Source: §4 · p. 9

5.1 Post-Training Compression

• At 80% pruning, BM-µ retains 70.57% CIFAR-100 test accuracy versus 60.84% for BM-mag, while Muon, Muon+GL, and NuMuon fall to 1.93%, 6.01%, and 3.84%, respectively.Figure 2 reports mean test accuracy over three seeds after post-training compression without recovery; under magnitude pruning, BMuon reaches 71.46% dense accuracy versus Muon’s 71.49%.Source: Figure 2 · p. 10, §5.1 · p. 10

5.2 Gradual Structured Pruning

• 1.78 and 3.18 percentage points on CIFAR-100, and 3.45 and 4.53 points on ImageNet-1K: BMuon-mag’s gains over Muon-mag at 90% and 95% pruning, respectively.Using µ as the pruning score adds gains of 2.46 and 1.03 points on CIFAR-100, and 0.35 and 1.27 points on ImageNet-1K; discovered architectures also outperform nearly channel-matched uniform-width models in all four settings.Source: Table 1 · p. 11, §5.2 · pp. 11–12

5.3 Structured Compressibility at Language-Model Scale

• 10.92 versus 12.84: Muon’s unpruned validation perplexity is lower than BMuon’s, but BMuon preserves its unpruned performance through 80% pruning, reaching 12.86 with activation ranking and 12.85 with µ ranking.Figure 3 reports validation perplexity after global structured pruning of a 1B-parameter LLaMA-style model trained on FineWeb; BMuon outperforms Muon at moderate and high pruning levels under shared magnitude or activation criteria.Source: Figure 3 · p. 12, §5.3 · p. 12

6 Conclusion

• Across vision and language models, BMuon is more robust than Muon to magnitude-based MLP-channel pruning, and learned allocations provide an effective pruning score.The learned allocation also identifies useful non-uniform MLP architectures when retrained from scratch.Source: §6 · pp. 12–13
• The paper leaves open BMuon’s implicit bias beyond linear settings, approximation guarantees, budgets for other Transformer components, and the quality–robustness tradeoff at larger scales.Experiments focus on MLP hidden channels, and language-model experiments observe a tradeoff between dense-model quality and pruning robustness.Source: §6 · p. 13

Appendix A. Muon Algorithm

• Muon applies momentum followed by approximate matrix orthogonalization; its geometric results hold for arbitrary momentum matrices, independent of the momentum variant.The appendix describes five Newton–Schulz iterations as an approximate orthogonalization and states that parameters not optimized by Muon are optimized separately with AdamW.Source: Appendix A · p. 18

Appendix B. BMuon Oracle and Geometry

• Appendix B provides derivations for the BMuon allocation oracle and update, its gated-MLP extension, allocation-curvature results, and induced norm geometry.The derivations underlie the oracle and geometry introduced in the main text.Source: Appendix B · p. 19

B.1 Derivation of the BMuon Oracle

• For a fixed allocation, the oracle optimizes update directions independently across matrices, then maximizes the resulting value over feasible allocations to obtain a joint solution.A maximizing allocation exists, and polar directions attain the matrixwise optima.Source: §B.1 · pp. 19–20

B.2 Extension to Gated MLPs

• For SwiGLU, each hidden-channel allocation is shared across three coupled matrix components, while the same global allocation constraint applies.The appendix states that this is the only change from the two-matrix case.Source: §B.2 · p. 20

B.3 Allocation Curvature

• The allocation objective is concave for all momenta exactly when 0 < α ≤ 1/2, convex for all momenta at α = 1, and generally neither in between.At α = 1 with integer K, an optimal allocation can be chosen binary, selecting exactly K channels.Source: §B.3 · pp. 21–22

B.4 Norm Geometry of BMuon

• For 0 < α ≤ 1, BMuon induces a norm whose unit ball is the convex hull of its admissible directions, with Ωα,K bounded between the Muon norm and (N/K)^α times that norm.The appendix establishes that the unit ball is compact, convex, centrally symmetric, and contains the origin in its interior.Source: §B.4 · pp. 22–23

B.5 Proofs for the Geometric Special Cases

• When K = N or α = 0, the induced BMuon norm equals Muon’s product norm.The appendix states these geometric special cases directly.Source: §B.5 · p. 25
• At α = K = 1, the appendix derives an explicit norm expression and its dual using the Euclidean norms of single nonzero rows or columns.The nuclear norm of each such single-row or single-column matrix equals the Euclidean norm of that row or column.Source: §B.5 · pp. 26–27

Appendix C. Compact-Representation Theory

• Under a sufficiently scarce BMuon budget, every minimum-cost linear factorization is width-minimal, whereas Muon admits a minimum-cost representation at full width.The appendix also gives a complementary weight-decay result: channel norms are controlled by allocation histories, and channels with vanishing allocations are driven to zero.Source: Appendix C · p. 28

C.1 Intrinsic-Width Recovery

• Both representation-cost infima are attained; Muon admits a minimum-cost factorization using all hidden channels, while if K≤d_eff^α(X), every BMuon minimizer uses rank(Xℓ) channels per block.Source: §C.1 · p. 28

C.2 BMuon with Weight-Decay

• Weight decay makes channel norms depend on discounted allocation histories, so sustained under-allocation suppresses a channel and allocations tending to zero drive its norm toward zero.Source: §C.2 · pp. 33–35

Appendix D. MM Allocation Update

• For 0<α≤1/2, Appendix D establishes MM-update monotonicity, convergence in objective value to the global optimum, and global optimality of every accumulation point; it also bounds practical approximation errors and reports update overhead.The appendix further compares the practical update with a certified allocation solve and reports training-time measurements for ImageNet-1K and a 1B-parameter LLaMA model.Source: Appendix D · p. 36

D.1 Minorization and allocation scores

• The allocation-score contributions are nonnegative, and scores are strictly positive when a channel is nontrivial.Source: §D.1 · p. 38

D.2 Closed-form surrogate maximization

• Under the nontrivial-channel condition, the surrogate is strictly concave over the feasible set and has a unique maximizer.Source: §D.2 · p. 38
• The closed-form MM update finds λµ > 0 with H(λµ) = K by bisection, since H decreases continuously from N to Nε and Nε < K < N.Source: §D.2 · p. 39

D.3 Convergence of repeated MM updates

• Repeated MM updates from any feasible start have every accumulation point globally maximize ΦM when the nontrivial-channel condition holds and 0 < α ≤ 1/2.The objective values are nondecreasing and converge.Source: §D.3 · p. 39

D.4 Effect of the allocation floor

• To satisfy the allocation floor, the proof mixes an optimizer over KK with the uniform allocation, producing a feasible allocation in KK,ε.The mixture uses δ = εN/K, with 0 < δ < 1.Source: §D.4 · p. 41

D.5 Approximation of the allocation oracle

• The largest sampled one-update deficits were about 1.2 × 10^-3 on 1-step-refresh states and 1.1 × 10^-3 on certified-refresh states.The similar values indicate that one-step accuracy was not specific to states generated by that update rule.Source: §D.5 · p. 43, Figure 4 · p. 44
• After four optimizer steps with µ held fixed, the certified deficit remained on the order of 10^-3 in the worst-over-sampled-windows summary.This indicates only a small additional error in the allocation objective at the sampled states.Source: §D.5 · p. 43, Figure 4 · p. 44
• Replacing one practical MM update with a certified solve at each refresh left accuracy and post-training compressibility within observed seed-to-seed variability.The paired differences had inconsistent signs across seeds.Source: §D.5 · p. 44, Figure 4 · p. 44
• The end-to-end comparison isolates allocation computation at refreshes; it does not test certified recomputation at every optimizer step.Both runs refreshed µ every five optimizer steps.Source: §D.5 · p. 44

D.6 Computational Overhead

• An allocation update adds one polar evaluation per budgeted matrix, making the amortized count 1.2 per step when updates occur every five steps.Persistent extra optimizer state is one allocation scalar per budgeted channel; other parameters are updated as in Muon.Source: §D.6 · pp. 44–45
• ImageNet-1K epoch times differed by 0.6% at 90% pruning and 0.3% at 95%, comparable to seed variation; the runs did not resolve a BMuon slowdown.The reported means were 92.5±0.4 versus 91.9±1.2 seconds and 92.1±0.5 versus 91.9±0.6 seconds per epoch, respectively.Source: §D.6 · p. 45, Table 2 · p. 46
• For LLaMA, mean step times were 4858 ± 42 ms with BMuon and 4918 ± 22 ms with Muon; the difference was not attributed to the optimizer.The 1.2% difference was comparable to seed-to-seed and normal system-level variation.Source: §D.6 · pp. 45–46, Table 3 · p. 46

Appendix E. Experimental Details

• Appendix E provides experimental setups and additional results for ImageNet-1K and CIFAR-100 post-training pruning, global-budget effects, gradual pruning, and LLaMA/FineWeb experiments.It also includes architecture controls for the gradual-pruning experiments.Source: Appendix E · p. 47

E.1 ImageNet-1K Post-Training Structured Pruning

• The ViT-S/16 has 12 Transformer blocks, MLP width 1536, and 18,432 total MLP hidden channels.It uses 224×224 inputs and 16×16 patches.Source: §E.1 · p. 47
• BMuon uses α = 1/2 and anneals its global budget from 18,432 channels to 5,530, approximately 30% of the total, during the first quarter of training.The budget then stays fixed; allocation updates occur every five optimizer steps.Source: §E.1 · p. 47
• After training, channels are globally ranked by coupled weight magnitude and pruned independently at each level, without fine-tuning or recovery.The experiment evaluates Top-1 accuracy on the full 50,000-image validation set over three seeds.Source: §E.1 · pp. 47–48, Table 4 · p. 48

E.2 CIFAR-100 Post-Training Compression

• The six-block CIFAR-100 ViT is trained for 200 epochs; BMuon anneals its 4,608-channel global budget to 922 and refreshes allocations every five optimizer steps.Non-MLP Transformer matrices use ordinary Muon.Source: §E.2 · p. 48

Post-training compression.

• Post-training comparisons prune MLP channels globally without cumulative pruning or fine-tuning, while NuMuon uses truncated SVD with ranks chosen to approximately match parameter counts.BMuon is evaluated using both coupled weight magnitude and its final learned allocation for channel ranking.Source: Post-training compression · p. 49

E.3 Effect of the Global Budget

• On CIFAR-100, after global MLP-channel pruning, magnitude-ranked models remain robust until approximately pruning 1−K/N of channels, then degrade sharply; using final learned allocation µ shows the same pattern, generally more gradual.Source: §E.3 · pp. 49–50, Figure 6 · p. 50
• At K/N = 0.1 and 0.2, unpruned accuracy decreases slightly, trading some dense performance for more restrictive budgets.This is reported for the CIFAR-100 budget ablation.Source: §E.3 · p. 50

Pruning protocol

• Pruning removes only MLP hidden channels, ranks surviving candidates globally across blocks, and enforces at least four channels per block.Pruned channels do not return; training resumes from surviving weights and optimizer state.Source: §E.4 · pp. 50–51

BMuon schedule

• BMuon’s budget is annealed from the current to the next target channel count during the first half of each phase, then held fixed; pruning ends the phase at that target.After pruning, surviving channels’ allocation coordinates are reset to one.Source: §E.4 · p. 51

Experimental setup

• CIFAR-100 uses a six-block ViT with 2,688 prunable channels, while ImageNet-1K uses ViT-S/16 and interleaves physical channel removal with training.The CIFAR-100 model trains for 200 epochs; ImageNet-1K uses a 150-epoch schedule.Source: §E.4 · p. 51

Evaluation

• Final Top-1 accuracy is measured after recovery on the CIFAR-100 test set at epoch 200 and the full ImageNet-1K validation set at epoch 150; means and standard deviations use three seeds.Tables 6 and 7 report the selected per-block widths for CIFAR-100 and ImageNet-1K, respectively.Source: §E.4 · p. 51, Table 6 · p. 52, Table 7 · p. 52

Architecture controls

• At 90% and 95% pruning, BMuon-µ architectures have 269 and 134 total CIFAR-100 channels versus uniform baselines with 270 and 132; on ImageNet-1K, they have 1,843 and 922 versus 1,848 and 924.The uniform baselines use widths 45 and 22 per CIFAR-100 block, and 154 and 77 per ImageNet-1K block; discovered architectures are retrained from scratch with standard Muon and compared with matched-total uniform models.Source: §E.4 · pp. 51–52, Table 6 · p. 52, Table 7 · p. 52

Model and data

• The 17-layer LLaMA-style decoder has 92,837 structural MLP channels and is trained on about 13.1 billion FineWeb tokens.It uses a 2,048 hidden dimension, SwiGLU width 5,461, context length 4,096, and the LLaMA-2 tokenizer.Source: §E.5 · p. 53

Training

• BMuon updates the three SwiGLU projection matrices in each block, while attention matrices continue to use standard Muon.The Muon baseline uses Muon for both attention and MLP matrices.Source: §E.5 · p. 53

BMuon configuration

• BMuon uses a fixed budget K = 18,567, approximately 20% of the 92,837 channels, with one MM allocation iteration every five optimizer steps.The allocation uses α = 1/2 and a numerical floor of 10^-6; no budget schedule is used.Source: §E.5 · p. 53

Post-training pruning

• Post-training pruning removes only SwiGLU hidden channels, ranks channels globally across blocks, and retains at least 10 channels per layer.Each pruning ratio is applied directly to a restored unpruned checkpoint; no post-pruning recovery is performed.Source: §E.5 · p. 53, Table 8 · p. 54

Evaluation

• Pruned models are evaluated without recovery on 5,046,272 validation tokens, with activation statistics accumulated in FP32 over 1,048,576 calibration tokens.Table 8 reports validation perplexity as mean ± sample standard deviation over three seeds.Source: §E.5 · p. 53, Table 8 · p. 54