Drop-Muon: Optimization & Muonium Beam Advances

• Drop-Muon is a dual-concept innovation combining a non-Euclidean randomized deep learning optimizer and a precision apparatus for muonium free-fall experiments.
• The optimization method employs selective layer updates with mirror descent and momentum to achieve faster convergence and lower computational cost.
• In beam physics, Drop-Muon utilizes phase-space compression to generate high-brightness, low-emittance slow muon beams for fundamental gravitational tests.

A Drop-Muon refers to two distinct concepts in scientific research: (1) a non-Euclidean randomized progressive training method for optimization in deep learning, and (2) a precision apparatus for muonium free-fall experiments enabled by phase-space compressed slow muon beams. The first is an algorithmic framework that provides theoretical and empirical advances in neural network training. The second leverages low-emittance muon sources to enable tests of fundamental physical symmetries via gravitational free-fall of muonium. Each is characterized by selective, highly efficient operations—layer-wise parameter updates in optimization, and phase-space reduction in beam design. The following sections delineate these dual meanings with technical rigor.

1. Drop-Muon in Deep Learning Optimization

The Drop-Muon algorithm is a modification of recent non-Euclidean layer-specific optimizers (notably Muon, as well as Scion and Gluon) that, contrary to conventional practice, does not update all network layers each optimization step. Instead, Drop-Muon “drops” some layer blocks—temporarily freezing their parameters—and updates only a sampled sub-network at each iteration according to a randomized schedule. Despite its simplicity, Drop-Muon delivers both strong theoretical guarantees and significant empirical gains in deep neural network training (Gruntkowska et al., 2 Oct 2025).

Algorithm Overview:

Given a model decomposed into $b$ parameter blocks $X = [X_1,\dots, X_b]$ (e.g., layers), Drop-Muon operates as follows:

• At step $k$, a random subset $S^k \subseteq \{1,\dots, b\}$ is sampled, according to a user-defined distribution $D$ (often via Randomized Progressive Training, RPT).
• Only blocks $i \in S^k$ are updated; others are held constant.
• For active blocks, the method applies a layer-wise mirror-descent update over arbitrary norm balls, using the corresponding dual norm and sharp operator.
• Momentum is retained per block: $$M_i^k = (1-\beta_i)M_i^{k-1} + \beta_i \nabla_i f(X^k; \xi^k)$$, with $X_i^{k+1} = X_i^k - t_i^k (M_i^k)^\sharp$.

A particularly efficient instantiation is RPT, where the sampled set $S^k$ always consists of a contiguous stack of layers deepest in the network, aligning with practical backpropagation computations.

2. Theoretical Guarantees and Convergence Analysis

Drop-Muon admits rigorous convergence guarantees under layer-wise smoothness and generalized $(L^0, L^1)$-smoothness, in both deterministic and stochastic settings. This marks the first such theoretical results for randomized progressive training in the non-smooth and stochastic regime (Gruntkowska et al., 2 Oct 2025).

Proven Regimes:

• Block-wise Lipschitz smoothness: Guarantees $O(1/K)$ convergence of weighted dual-normed gradients.
• Generalized $(L^0, L^1)$-smoothness: Achieves $O(1/\sqrt{K})$ rates.
• Stochastic gradients with bounded variance: Establishes $O(K^{-1/4})$ convergence in expectation.

Central to these results are (a) non-Euclidean descent lemmas accounting for curvature variability, (b) recursions for momentum-induced bias, and (c) adaptive step-size schedules per layer, often inversely scaled to local smoothness constants.

3. Computational Complexity and Cost Analysis

Standard full-network update methods, such as Muon or SGD, update all $b$ layers each iteration, leading to per-step cost $c_{\mathrm{ov}} + \sum_{i=1}^b (c_i + c_i^\sharp)$, with $c_i$ covering forward/backward compute and $c_i^\sharp$ the cost of computing mirror-descent sharp operators. Drop-Muon, by updating only a (typically proper) subset $S^k$, reduces the expected per-iteration compute.

The expected total cost to reach an $\varepsilon$-accurate solution is

$\mathbb{E}[K] \cdot \mathbb{E}[\mathrm{cost}(S)]$

with $K = O(1/\varepsilon)$ or $O(1/\varepsilon^2)$ depending on smoothness assumptions.

Analysis shows that full-network updates are optimal only under the exceptional case where the first layer’s curvature constant dominates all others—a highly non-generic condition. For almost all practical scenarios, randomized partial updates are strictly superior in expected wall-clock compute (Gruntkowska et al., 2 Oct 2025). Optimal layer sampling distributions $p_s$ can be computed explicitly and tuned to match inverse curvature profiles.

4. Empirical Performance and Experimental Verification

Extensive experiments on three-layer CNNs across MNIST, Fashion-MNIST, and CIFAR-10 datasets demonstrate the empirical advantages of Drop-Muon. Key findings include:

• Up to $1.4\times$ reduction in wall-clock time to reach specified accuracy on MNIST benchmarks.
• On Fashion-MNIST, Drop-Muon attains $\approx1.2\times$ speedup up to 95% training accuracy.
• For CIFAR-10, epoch-shift Drop-Muon reaches 90% training accuracy faster than full-network Muon, but gains are smaller for lower accuracy thresholds, underscoring the importance of problem-specific tuning.
• Variance across random initializations is higher early in training due to random layer selection; this effect diminishes as convergence is approached.

Empirical configurations use consistent hyperparameters with Muon baselines—including momentum 0.5 and spectral norm constraints—allowing direct attribution of speedup to partial-layer updating rather than optimizer-specific tuning (Gruntkowska et al., 2 Oct 2025).

5. Practical Implementation and Guidance

Drop-Muon is designed for ease of implementation, requiring only modest changes to standard training loops. It is compatible with existing backpropagation infrastructures by leveraging RPT, which aligns randomized cut-off indices with the propagation of gradients during the backward pass.

Practical guidance includes:

• Applying RPT or epoch-shift scheduling to economize on backward passes over nearly-converged shallow layers.
• Maintaining block-wise momentum buffers and mirror descent geometries.
• Tuning sampling distributions and step sizes jointly; analytical results prescribe inverse-curvature scaling for optimality.
• Employing Drop-Muon when hardware is bottlenecked by gradient computation, as substantial time can be saved by reducing the amount of necessary computation per iteration.

A plausible implication is that, given the ubiquity of non-uniform layerwise curvature in modern deep neural networks, randomized progressive training should become the default rather than updating all layers indiscriminately.

6. Drop-Muon in Muonium Free-Fall and Beam Physics

In a distinct context, Drop-Muon also denotes the apparatus and methodology for producing high-brightness, phase-space compressed slow muon beams for precision muonium free-fall experiments. The muCool concept developed at the Paul Scherrer Institute provides a practical pathway for this (Belosevic et al., 2019).

Key Elements:

• Phase-Space Compression: Achieves compression factor $C_{\rm tot}\simeq10^{10}$ at $10^{-3}$ efficiency, transforming a standard $N_{\rm in} \approx 2\times10^4$ muons/s into $N_{\rm out} \approx 20$ slow $\mu^+$/s.
• Field Configurations: Utilizes a 5 T solenoidal magnetic field, crossed electric fields (up to 1 kV/cm), and controlled helium gas density gradients (temperatures 4–12 K) to bunch muons longitudinally and transversely.
• Beam Properties: Yields sub-mm transverse size and mm-scale longitudinal emittance, eV-range mean energy, and $\sim4\,\mu$s pulsed timing suitable for gravity measurement.
• Muonium Formation: Stopped in thin silica layers, typical muonium (Mu) formation efficiencies are on the order of 10%. The apparatus can deliver $\sim2$ Mu/s into a field-free region with minimal background, enabling gravity experiments at the 1% level with $\sim10^4$ event statistics.

This approach underpins future tests of fundamental symmetries and gravity with muonium, representing a major advance in slow muon beam delivery and precision measurement capability (Belosevic et al., 2019).

7. Limitations, Challenges, and Prospects

Current limitations of Drop-Muon optimization include increased variance early in training due to randomized sampling and the necessity to tune sampling policies and step sizes for optimal performance. Hardware constraints and non-uniform cost ratios across layers or architectures may influence empirical outcomes.

In the beam physics context, confirmation of the merged-stage compression and extraction efficiency remains pending. The flux of slow muons, currently limited to few tens per second, constrains statistical power in gravity experiments. Prospective improvements include enhanced primary beams, optimized cell geometry, and superior extraction optics to boost both efficiency and output.

A plausible implication is that further development of both algorithmic strategies for progressive training and hardware upgrades for muon beamlines will expand the applicability and precision of Drop-Muon concepts across computational and physical domains.