MGUP: Momentum-Gradient Alignment Update Policy

• MGUP is a selective intra-layer update mechanism that modulates per-parameter step sizes based on momentum-gradient agreement, ensuring all parameters receive non-zero updates.
• It integrates seamlessly with optimizers like AdamW, Lion, and Muon, leading to faster convergence and improved training stability across varied pretraining and fine-tuning tasks.
• Theoretical analyses and empirical results demonstrate that MGUP’s two-level dense update policy effectively balances aggressive steps for well-aligned coordinates with smaller corrective steps for others.

Searching arXiv for MGUP and closely related optimizer papers to ground the article in current literature. MGUP denotes Momentum-Gradient alignment Update Policy, a selective intra-layer update mechanism for stochastic optimization that augments standard momentum-based optimizers by assigning heterogeneous per-parameter step sizes according to momentum-gradient agreement. At each iteration, a fixed proportion of parameters with the highest alignment scores receives a larger step-size factor, while the remaining parameters receive a smaller but strictly positive factor, so the method preserves dense updates rather than imposing hard sparsity. MGUP is presented as a nearly plug-and-play module for optimizers such as AdamW, Lion, and Muon, with a stochastic nonconvex convergence analysis for MGUP-AdamW without weight decay (Chang et al., 16 Jun 2026). The term should not be confused with MGPU, an earlier C++ library for single-node multi-GPU real-time computing (Schaetz et al., 2013).

1. Definition, scope, and motivating problem

MGUP is designed for fine-grained selective updates in large-scale model training. Its motivating premise is that, in overparameterized regimes such as LLM pretraining and downstream fine-tuning, optimization often benefits from emphasizing coordinates whose momentum-derived update direction is well aligned with the current stochastic gradient. Rather than modifying the underlying momentum recursion or second-moment estimator, MGUP intervenes only at the level of per-parameter step-size modulation, which is why it is characterized as plug-and-play (Chang et al., 16 Jun 2026).

The mechanism addresses a gap between two existing tendencies in optimizer design. On one side are methods that reduce state or compute through low-rank or memory-efficient approximations; on the other are selective-update methods that restrict optimization to chosen layers, blocks, or coordinates. MGUP belongs to the latter category but operates within layers, at coordinate granularity, and does so without freezing non-selected parameters. A central motivation is the limitation of hard-masking schemes such as Cautious Optimizers, which update only aligned coordinates and assign zero updates elsewhere. The MGUP formulation retains the emphasis on “trustworthy” directions while ensuring that all parameters continue to move, albeit at different rates.

A frequent misconception is that MGUP is a sparsification rule. It is not. The paper explicitly defines the method so that non-selected coordinates still receive a smaller but non-zero update, and this choice is not merely heuristic; it is bound up with the convergence argument and with a counterexample showing the failure mode of zero-masked variants.

2. Selective update mechanism

The core MGUP rule begins with an alignment score. For MGUP-AdamW, if $\mathbf{u}_t$ denotes the Adam preconditioned direction and $\mathbf{g}_t$ the current stochastic gradient, the score vector is

$\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$

with the paper also noting that $\mathbf{m}_t \odot \mathbf{g}_t$ yields nearly the same ordering in practice for AdamW because the second-moment scaling is often roughly homogeneous within a layer (Chang et al., 16 Jun 2026).

Let $d$ be the number of parameters under consideration and let $\tau \in (0,1)$ be the selection ratio. MGUP sets

$K = \lfloor \tau d \rfloor,$

extracts the index set $\mathcal{I}_{\text{topK}}$ of the $K$ largest entries of $\mathbf{s}_t$, and defines the multiplicative step-size mask $\mathbf{g}_t$0 coordinate-wise by

$\mathbf{g}_t$1

The update then becomes

$\mathbf{g}_t$2

This yields two effective step-size regimes. Selected coordinates receive $\mathbf{g}_t$3 with $\mathbf{g}_t$4, whereas the remainder receive $\mathbf{g}_t$5 with $\mathbf{g}_t$6. The method therefore implements a two-level dense update policy: large steps for top-ranked aligned coordinates and small steps for all others.

The use of a ranking criterion rather than a sign test is significant. A large positive score $\mathbf{g}_t$7 reflects both sign agreement between update direction and gradient and a large magnitude of agreement. MGUP thus prioritizes coordinates exhibiting strong alignment, not merely positive alignment. This distinguishes it from masking rules based only on $\mathbf{g}_t$8.

3. Theoretical properties and the role of non-zero small steps

The theoretical analysis is given for MGUP-AdamW without weight decay in the standard stochastic nonconvex setting

$\mathbf{g}_t$9

The assumptions are conventional for modern Adam analyses: a lower-bounded objective, $\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$0-smoothness, unbiased stochastic gradients with bounded variance, and, for the high-probability result, a coordinate-wise bounded variance condition (Chang et al., 16 Jun 2026).

A central technical device is the reformulation of the update into a denominator-controlled form with bounds depending on

$\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$1

The appearance of $\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$2 in $\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$3 is decisive. If $\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$4, then $\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$5, the denominator bounds collapse, and the convergence proof cannot proceed. This formalizes the paper’s claim that the smaller update on non-selected coordinates must remain strictly positive.

Under the stated assumptions, Theorem 4.1 establishes an expected convergence rate of order $\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$6 for the average gradient norm. The theorem gives

$\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$7

with $\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$8 scaling as $\mathbf{s}_t = \mathbf{u}_t \odot \mathbf{g}_t,$9. Theorem 4.2 further gives a high-probability bound of the same asymptotic order, again contingent on bounded, nonzero step modifiers.

The paper’s counterexample provides the complementary negative result. In a one-dimensional stochastic problem with rare large positive gradients and frequent small negative gradients, C-Adam’s hard mask can repeatedly skip corrective updates or move in the wrong direction, causing divergence away from the global minimum. MGUP, by contrast, continues to take small corrective steps even in poorly aligned coordinates and converges to the optimum $\mathbf{m}_t \odot \mathbf{g}_t$0. Theoretical analysis and counterexample therefore support the same design principle: misaligned coordinates should be down-weighted, not frozen.

4. Integration with AdamW, Lion, and Muon

MGUP is positioned as a wrapper around existing momentum-based optimizers rather than a replacement for them. In MGUP-AdamW, the optimizer first computes the usual Adam quantities

$\mathbf{m}_t \odot \mathbf{g}_t$1

$\mathbf{m}_t \odot \mathbf{g}_t$2

$\mathbf{m}_t \odot \mathbf{g}_t$3

followed by the bias-corrected step size

$\mathbf{m}_t \odot \mathbf{g}_t$4

MGUP then computes $\mathbf{m}_t \odot \mathbf{g}_t$5, applies weight decay as usual, and performs the masked update $\mathbf{m}_t \odot \mathbf{g}_t$6 (Chang et al., 16 Jun 2026).

The same pattern carries over to other base optimizers. For MGUP-Lion, the direction is

$\mathbf{m}_t \odot \mathbf{g}_t$7

after which MGUP scales coordinates via the same top-$\mathbf{m}_t \odot \mathbf{g}_t$8 rule. For MGUP-Muon, the momentum matrix $\mathbf{m}_t \odot \mathbf{g}_t$9 is orthogonalized through Newton-Schulz to obtain $d$0, while the alignment score is taken from $d$1.

In the reported experiments, the default choice is

$d$2

so exactly half of the parameters receive a $d$3 step and the rest receive a $d$4 step. The paper gives the heuristic average effective learning rate

$d$5

which exceeds the base learning rate and is maximized at $d$6.

The computational additions are limited to one element-wise product, one top-$d$7 selection, and one element-wise mask application per iteration. The paper states that top-$d$8 can be implemented with standard GPU primitives such as torch.topk or argpartition, and it reports only small runtime and memory overheads. One concrete memory example is LLaMA-7B GSM8K fine-tuning, where peak reserved memory increases from 32.09GB for AdamW to 33.82GB for MGUP-AdamW.

5. Empirical results across pretraining and fine-tuning

The empirical study covers MAE pretraining, LLM pretraining, and supervised fine-tuning, comparing MGUP-enhanced optimizers against their base versions, cautious variants, and several memory-efficient baselines (Chang et al., 16 Jun 2026).

Setting	Model / task	Reported outcome
MAE pretraining	ViT-27M on CIFAR-10	MGUP-AdamW shows consistently lower training and validation loss than AdamW; C-AdamW degrades over time
LLM pretraining	LLaMA2-71M on WikiText-103	MGUP-AdamW is about 1.6× faster than AdamW; MGUP-Lion about 2.5× faster than Lion; MGUP-Muon about 1.2× faster than Muon
LLM pretraining	Qwen2.5-150M on WikiText-103	MGUP-AdamW converges faster and attains lower validation loss than AdamW and C-AdamW; MGUP-Muon yields about 1.1× speedup
GLUE fine-tuning	RoBERTa-base	MGUP-AdamW achieves the best average GLUE score, 85.15, versus 84.74 for AdamW, 83.69 for C-AdamW, and 83.58 for LDAdamW
GSM8K fine-tuning	LLaMA2-7B	MGUP-AdamW reaches the best validation accuracy, 34.96%, and the lowest training loss among compared optimizers

In MAE pretraining with a ViT-27M model on CIFAR-10, MGUP-AdamW delivers consistently lower training and validation losses than AdamW, while C-AdamW deteriorates and eventually underperforms the plain baseline. In language-model pretraining on LLaMA2-71M with WikiText-103, MGUP-AdamW reaches the same validation loss as AdamW in about 1.6× fewer steps, MGUP-Lion yields about 2.5× speedup over Lion while avoiding the early loss spikes of C-Lion, and MGUP-Muon shows about 1.2× speedup over Muon. On Qwen2.5-150M, MGUP-AdamW again improves convergence and final validation loss, and MGUP-Muon reports roughly 1.1× speedup.

The fine-tuning results show a similar pattern. On GLUE with RoBERTa-base, MGUP variants sustain higher and more stable scores across learning rates than both base and cautious optimizers, with MGUP-AdamW achieving the best average score of 85.15. On GSM8K supervised fine-tuning with LLaMA2-7B, MGUP-AdamW attains 34.96% best validation accuracy, slightly ahead of LDAdamW at 34.88%, GaLore at 34.62%, and C-AdamW at 34.68%. The paper therefore presents MGUP not only as a faster optimizer in step efficiency, but also as one that can improve generalization and reduce instability.

Sensitivity studies on LLaMA2-71M pretraining further clarify the hyperparameter behavior. With $d$9 fixed, increasing $\tau \in (0,1)$0 too far hurts performance; with $\tau \in (0,1)$1 fixed, larger $\tau \in (0,1)$2 generally improves performance. The reported behavior of $\tau \in (0,1)$3 is poor and often unstable, matching the theoretical argument against hard masking.

6. Relation to prior work, limitations, and open questions

MGUP is situated at the intersection of selective-update methods, adaptive optimizers that exploit momentum-gradient relationships, and Adam convergence theory (Chang et al., 16 Jun 2026). Relative to layer- or block-selective approaches such as AutoFreeze, LOMO, LISA, and BAdam, MGUP operates at per-parameter intra-layer granularity. Relative to AdaBelief and Cautious Optimizers, it uses momentum-gradient alignment not to alter the second moment or to impose a binary mask, but to define a ranking-based two-level step policy. Relative to low-rank and memory-efficient optimizers such as GaLore, LDAdam, Adam-mini, Adafactor, and Q-GaLore, MGUP is orthogonal: it determines where larger versus smaller steps should be applied, rather than compressing optimizer state.

Several limitations are explicit. The selection ratio $\tau \in (0,1)$4 is manually chosen, typically as $\tau \in (0,1)$5, and the paper does not provide an adaptive rule for $\tau \in (0,1)$6 or a per-layer variant. The convergence theory is restricted to MGUP-AdamW without weight decay; there are no formal guarantees yet for MGUP-Lion, MGUP-Muon, or AdamW with weight decay. Selection is performed as a global top-$\tau \in (0,1)$7 over a layer or parameter vector, which may become a practical concern in extremely large models. The experiments, while spanning 71M- and 150M-parameter pretraining and 7B-parameter fine-tuning, do not explicitly study 10B+-scale pretraining.

The broader implication is that MGUP provides a principled middle ground between dense adaptive optimization and hard selective masking. It does not claim true sparsity, automatic compression, or optimizer-state reduction. Instead, it offers a controlled way to amplify directions deemed reliable by momentum-gradient agreement while preserving dense corrective motion in the rest of parameter space. This suggests a general template for optimizer design in which selective emphasis and stochastic convergence need not be in tension, provided the “de-emphasized” coordinates remain active through bounded, nonzero step sizes.