Muon for Phase Retrieval III — Adaptive power and the moving front

This is the conditional layer of the story. Part I fixes (a), and Part II shows what the spectra do. Here the exponent itself becomes a control (a(t)). The reduced equations are explicit, but global optimality among all online algorithms is not claimed.

The active spectral front moves. Early in training the visible directions are strong. Later the relevant modes are weaker and closer to the bulk. A fixed exponent can be a good compromise, but it is not the natural object if the front itself changes.

The point of variable Muon-(a(t)) is to turn the exponent into a control on the same Volterra state: [ \dot X_t=b(t,X_t,a_t). ] Here (X_t) contains captured masses, residual masses, bulk transforms, finite outlier roots, residues and BBP margins. The symbol (b) is the vector field of this reduced state. In the scalar mode coordinates, one component is [ b_i(t,X,a) = -2\eta\,\delta_i^{(a)}Q^{(a-1)/2}q_i^2 + \eta^2\nu_i^{(a)}Q^a . ] So (b) is not another algorithm. It is what the Paquette/Volterra state does if the current exponent is (a).

The local spectral cost

At time (t), the gradient channel has a bulk law and a finite signal front. Linearized AMP/VAMP says that the locally optimal spectral denoiser is the likelihood-ratio filter [ R_t(\sigma)=\frac{d\mu_{S,t}}{d\nu_{B,t}}(\sigma), ] where (\nu_{B,t}) is the bulk singular-value law and (\mu_{S,t}) is the signal/front measure. A filter (f) is judged by [ \mathcal H_t(f)= \frac{\langle f,R_t\rangle_{\nu_B}^2} {\langle f^2\rangle_{\nu_B}}. ] Cauchy–Schwarz proves that the unrestricted maximizer is (f_{\rm opt}\propto R_t). Muon only allows the one-parameter family [ f_a(\sigma)=\sigma(\sigma^2+\varepsilon^2)^{(a-1)/2}. ] The local cost is therefore [ F_t(a)=-\log \mathcal H_t(f_a). ] This is what (F) means: the loss of projecting the AMP/VAMP local spectral denoiser onto the Muon power family.

The AMP comparison is deliberately first-order. Full AMP contains an Onsager correction. Muon does not literally run AMP. The statement is narrower: on a locally power-law active front, [ \log R_t(\sigma_te^u)=\kappa_t+a_{\rm loc}(t)u+o(1), ] the AMP/VAMP denoiser is locally proportional to (\sigma^{a_{\rm loc}(t)}). In that regime the best scalar Muon exponent is the power-law projection of the local AMP denoiser.

The future action

One-step denoising is not enough. A choice of (a) now changes future BBP margins. If (V(t,X)) is the value of the remaining control problem, then the first-order price of choosing (a) is [ A_t(a)=\nabla_XV(t,X_t)\cdot b(t,X_t,a)-c_t . ] The constant (c_t) is a gauge: it shifts all actions equally, so it cannot change which exponent is selected. Only action differences matter.

If a terminal loss penalizes invisible branches, for example [ \mathcal L_T(X_T)= \sum_{c,i}w_i^c\,{\rm softplus}\epsilon(-\Delta_i^c(X_T)) +\omega Q(T), ] then (A_t(a)) is literally a price for future spectral exits. The linearized sensitivity of a BBP time (t_j^c) to a perturbation at time (s) has the form [ \frac{\delta t_j^c}{\delta a_s} = - \frac{\nabla_X\Delta_j^c(X{t_j^c}) \Phi(t_j^c,s)\partial_ab(s,X_s,a_s)} {d\Delta_j^c(X_t)/dt|_{t=t_j^c}} . ] This is the bridge from HJB to the spectral plots: an adaptive exponent is good only if it improves current denoising without destroying future finite-root margins.

Why Boltzmann

The entropy-regularized HJB policy is [ \pi_t(da)= \frac1{Z_t}\exp[-\beta_t(F_t(a)+A_t(a))]\,da, \qquad a_t=\int a\,\pi_t(da). ] The normalizing constant is [ Z_t=\int \exp[-\beta_t(F_t(u)+A_t(u))]\,du . ] The inverse temperature (\beta_t) controls sharpness. Large (\beta_t) approaches a hard HJB argmin. Smaller (\beta_t) smooths finite-dimensional spectral noise.

Adaptive Muon exponent trajectories over time — The adaptive object is \(a(t)\): fixed powers are horizontal references, while Boltzmann and HJB-style controls move with the active front.

Adaptive Muon spectral plot over time — Spectral consequence of the adaptive rules: time is vertical, spectral value is horizontal, and the goal is synchronized finite-root escape from the bulk.

There is a useful practical simplification. On the region where the Gibbs measure actually puts mass, the action can be close to affine: [ A_t(a)\simeq c_t+\lambda_ta . ] This is not a theorem of global optimality. It is a reduced-model statement: inside the Paquette/Volterra state, and inside the Muon power family, the future can sometimes be summarized by one shadow price (\lambda_t). The gate for this replacement is [ \beta_t \sqrt{\mathrm{Var}_{\pi_t}(A_t-c_t-\lambda_ta)} \ll 1 . ] When this is small, the affine action and the full reduced action choose nearly the same exponent.

What is proved, and what is not

The proved reduced statement is:

the fixed-(a) Volterra state gives the learning curve;
allowing (a) to vary gives a non-autonomous Volterra equation;
the local AMP/VAMP denoiser projects to a Muon power on a locally power-law front;
Boltzmann is the entropy-regularized HJB policy for this reduced state.

The stronger statement is not claimed here: Muon-(a(t)) is not proved to be globally optimal among all online SQ or AMP algorithms. The evidence is more specific and more spectral. The Boltzmann trajectories are attractive because they move smoothly while producing nearly simultaneous BBP exits in the finite-root plots. That is the physical sign of front equalization.

PDF

Download the variable-(a), Boltzmann and HJB companion paper.