Momentum as Residual-Driven Multiplier Correction for Deep Learning Optimization

The paper proposes an ADMM-Inspired Momentum (AIM) framework based on residual-penalty variable splitting, which interprets momentum as a multiplier-like correction driven by the splitting residual. AIM recovers the exponential moving average of gradients from an ADMM-style multiplier update and separates two mechanisms that are usually intertwined in practical optimizers: the residual penalty determines the update geometry, whereas the approximation of the objective-related subproblem determines the acceleration form. Building on AIM, the paper proposes Relativistic Adaptive gradient Descent with Accelerated Residual (RADAR), which combines relativistic adaptive geometry, decoupled residual correction, and second-order momentum filtering to improve the update direction and momentum estimation. The paper establishes stochastic convergence through a variance-perturbed Lyapunov drift analysis. Experiments on supervised vision learning, language modeling, and reinforcement learning show that RADAR achieves consistent improvements over strong adaptive optimizer baselines.

The contributions are summarized as follows:

• The AIM framework is proposed, an ADMM-inspired residual-penalty splitting formulation for momentum-based optimization. By introducing an auxiliary descent variable and a splitting constraint, AIM recovers the exponential moving average of gradients from an ADMM-style multiplier update, rather than treating momentum as a standalone filtering heuristic. In the Euclidean case, the splitting residual is further shown to be proportional to the gradient–momentum mismatch, providing a residual-correction interpretation of momentum that complements the conventional low-pass filtering view.

• AIM is used to disentangle two mechanisms that are often intertwined in practical optimizers: update geometry and acceleration form. The residual penalty determines the descent geometry, recovering Euclidean momentum descent, Adam/NAdam-type adaptive diagonal preconditioning, and Muon-type matrix-norm updates. The approximation of the objective-related subproblem determines the acceleration form, distinguishing heavy-ball momentum from Nesterov-type momentum. This separation explains multiple optimizer families through a unified residual-geometry and residual-correction mechanism.

• RADAR is proposed, a new optimizer derived from the AIM design principle and the structure-preserving motivation of RAD. RADAR combines relativistic adaptive geometry for speed-limited preconditioning, decoupled residual correction for acceleration-like parameter refinement, and second-order momentum filtering for improved momentum estimation. Its stochastic convergence is established via a variance-perturbed Lyapunov drift analysis, and experiments on supervised vision learning, language modeling, and reinforcement learning tasks validate its effectiveness.

The AIM framework introduces an auxiliary descent variable y and couples it with the network parameter θ through the equality constraint y − θ = 0. The residual-penalty augmented Lagrangian is defined as Lρ(θ, y, m) = L(θ) + ⟨m, y − θ⟩ + (ρ/2)ψ(y − θ), where m is a multiplier-like variable, ρ > 0, and ψ is a proper closed convex residual penalty. The y-subproblem contains no objective term and therefore specifies the descent geometry induced by ψ, while the θ-subproblem contains L(θ) and determines how objective-gradient information corrects the tentative descent point yk+1. The multiplier update is a relaxed ADMM-style residual correction applied to the splitting residual, using the generalized residual subgradient induced by ψ. In the idealized case where the θ-subproblem is solved exactly, the multiplier update yields mk+1 = β1 mk + (1 − β1)∇L(θk+1), recovering the standard exponential moving average form of momentum. For the Euclidean penalty ψ(r) = ∥r∥2, the splitting residual has two roles: it drives the ADMM-style multiplier update and, at the same time, measures the scaled mismatch between the current objective gradient and the previous momentum estimate, giving momentum a residual-correction interpretation.

Nesterov-type acceleration arises from the approximation of the θ-subproblem. The y-subproblem first produces a tentative descent point along the momentum direction, whereas the θ-subproblem further corrects this point using objective-gradient information. Directly setting θk+1 = yk+1 gives heavy-ball momentum, while keeping a first-order explicit approximation of this correction leads to a Nesterov-type update. Theorem 1 states that with ψ(r) = ∥r∥2, ρ1 = 1/η, and ρ2 = 1/(η(1 − β1)), if the θ-subproblem is approximated by θk+1 = yk+1 − η(1 − β1)(∇L(θk) − mk), the resulting update recovers a Nesterov-type momentum method up to a change of variables. If the simpler approximation θk+1 = yk+1 is used instead, the correction term vanishes and the update reduces to heavy-ball momentum.

Adaptive geometry arises from the y-subproblem. For Adam-type methods, using the weighted residual penalty ψ(r) = ∥r∥2Qk changes the y-subproblem into the preconditioned momentum step yk+1 = θk − ηQ−1k mk, where Qk = Diag(√vk + ϵ). Under this adaptive geometry, the direct approximation θk+1 = yk+1 recovers Adam without bias correction, while the Nesterov-type approximation recovers NAdam without bias correction. The same geometric view also extends beyond coordinate-wise adaptive preconditioning; replacing the adaptive diagonal geometry with a matrix-level steepest-descent geometry yields a Muon-type update under the direct approximation of matrix-valued variables.

RADAR instantiates the two AIM design axes with a relativistic adaptive residual geometry and a decoupled residual correction mechanism. The relativistic adaptive geometry matrix is defined as Rk+1:= Diag(√(δ2vk+1 + ζ)), where δ > 0 is the speed coefficient and ζ ∈ (0, 1] is the symplectic factor. Given the filtered momentum estimate mk+1, the AIM y-subproblem gives the tentative descent point yk+1 = θk − ηR−1k+1mk+1. The decoupled residual correction takes the preconditioned form R−1k+1(mk+1 − gk), with an independent correction coefficient l, giving the update θk+1 = θk − ηR−1k+1mk+1 − lR−1k+1(gk − mk+1). When l = 0, RADAR reduces to a relativistic adaptive momentum update. When l > 0, the residual term provides an acceleration-like correction. Second-order momentum filtering augments the standard first-order exponential moving average with a gradient-difference term: mk+1 = β1 mk + (1 − β1)gk + γ(gk − gk−1), where γ ≥ 0 controls the strength of the filtering. When γ = 0, this reduces to the standard first-order filtering formulation.

The convergence analysis uses a variance-perturbed Lyapunov drift argument that jointly controls the objective value, the momentum residual, and the successive parameter displacement. Under Assumptions 1–3 (smoothness and lower boundedness, unbiased stochastic gradient and bounded variance, bounded adaptive geometry), Lemma 1 provides a Lyapunov drift bound: EVk+1 − EVk ≤ −d1E∥rk+1∥2 − d2E∥∇L(θk)∥2 + Cvσ2(1/Bk + 1/Bk−1), where d1, d2 > 0 are the descent coefficients for the momentum residual and the stationarity measure, and Cv > 0 quantifies the amplification of stochastic-gradient variance. Theorem 2 establishes sublinear convergence: R(T) ≤ (EV0 − L*)/(d2T) + (Cvσ2/(d2T))Σ(1/Bk + 1/Bk−1). Therefore, with a fixed mini-batch size, RADAR converges to a variance-controlled stationary neighborhood. If full gradients are used or if the batch-size schedule makes the accumulated variance term uniformly bounded, RADAR achieves an O(1/T) average stationarity bound.

Experiments evaluate RADAR on supervised vision learning, language modeling, and reinforcement learning. For supervised vision learning, RADAR achieves the best test accuracy and test loss in both tasks: on CIFAR-10 with ViT, RADAR improves the best baseline test accuracy from 0.8761 to 0.8876 and reduces the test loss from 0.8026 to 0.7696; on CIFAR-100 with ResNet-50, RADAR also obtains the highest accuracy and lowest loss. For language modeling, RADAR achieves the lowest PPL in both GPT-2 settings: 25.6915 on WikiText-103 pre-training and 21.5462 on WikiText-2 fine-tuning. For reinforcement learning, RADAR achieves the highest average test return on both SAC tasks: on HalfCheetah-v4, RADAR improves the strongest baseline from 8996 to 9569 (a relative gain of 6.37%); on Walker2d-v4, RADAR improves the strongest baseline from 3362 to 3580 (a relative gain of 6.48%) and also achieves the smallest standard deviation. Ablation studies on CIFAR-10 with ViT isolate the effects of the two RADAR-specific components: decoupled residual correction (DRC) and second-order momentum filtering (SF). The full RADAR achieves the lowest training and test losses in the late training stage, showing that the two components jointly improve optimization performance. Removing SF leads to a clear degradation, suggesting that the gradient-difference term improves the momentum estimate and stabilizes the update direction. Removing DRC also weakens the performance, confirming that the residual correction contributes beyond the relativistic adaptive geometry alone. The variant without both components performs worst, further supporting the combined design of RADAR.

Improvements to AI Systems:

1. Adaptive Optimizer with Relativistic Geometry and Residual Correction (RADAR)

• Replace standard Adam/AdamW in training loops with RADAR, which uses a speed-limited preconditioning matrix R k+1 = Diag(sqrt delta squared v k+1 + zeta) to prevent extreme step sizes, and adds a decoupled residual correction term R k+1-1(g k - m k+1) to refine updates.

• Result: Faster convergence and lower final loss on vision (e.g., CIFAR-10 ViT accuracy +1.15%, CIFAR-100 ResNet-50), language modeling (WikiText-103 PPL 25.69 vs. baselines), and reinforcement learning (HalfCheetah +6.37% return).

2. Second-Order Momentum Filtering for Gradient Estimation

• Modify the momentum update to include a gradient-difference term: m k+1 = beta 1 m k + (1-beta 1)g k + gamma(g k - g k-1). This improves momentum accuracy by leveraging gradient change information, reducing oscillation and stabilizing training.

• Result: More reliable gradient direction estimates, leading to smoother loss curves and better generalization, especially in late-stage training.

3. Unified Optimizer Design Framework (AIM) for Task-Specific Geometry

• Use the AIM framework to systematically choose update geometry (Euclidean, diagonal adaptive, or matrix-norm) and acceleration form (heavy-ball vs. Nesterov) based on problem structure—e.g., matrix-level updates for transformer weights, diagonal adaptive for CNNs, and Nesterov for RL policy networks.

• Result: A configurable optimizer that can be tailored to different model architectures, improving training efficiency without manual hyperparameter tuning.

4. Variance-Aware Convergence Guarantees for Stochastic Training

• Apply the variance-perturbed Lyapunov drift analysis to automatically adjust batch sizes or learning rates based on the bounded variance term C v sigma 2(1/B k + 1/B k-1). This ensures convergence to a stationarity neighborhood with an O(1/T) bound when variance is controlled.

• Result: More robust training in noisy environments (e.g., RL with stochastic rewards), with theoretical guarantees of stability and convergence.

5. Residual-Correction Interpretation of Momentum for Debugging and Scheduling

• Use the insight that momentum is a multiplier-like correction proportional to the gradient–momentum mismatch to design adaptive learning rate schedules (e.g., increase rho when residual is large, decrease when small).

• Result: Automatic adjustment of update aggressiveness, reducing sensitivity to initial learning rate and improving performance across diverse tasks.

What the Improved AI System Can Do:

• Train deep models (vision transformers, ResNets, GPT-2, SAC agents) with consistently higher accuracy/lower loss/return than state-of-the-art adaptive optimizers (Adam, AdamW, NAdam, Muon).

• Adapt its optimization strategy to different architectures and data types via the AIM framework, without manual redesign.

• Maintain stable convergence under high stochastic variance (e.g., RL, small batch sizes) with theoretical guarantees.

• Reduce tuning effort by automatically balancing momentum and residual correction, leading to faster experimentation and deployment.