FlowAdam: Implicit Regularization via Geometry-Aware Soft Momentum Injection
Authors: Devender Singh, Tarun Sheel
Organizations: Department of Computer Science Memorial University of Newfoundland St. John’s, NL, Canada · Department of Mathematics and Statistics Memorial University of Newfoundland St. John’s, NL, Canada
Adaptive moment methods such as Adam use a diagonal, coordinate-wise preconditioner based on exponential moving averages of squared gradients. This diagonal scaling is coordinate-system dependent and can struggle with dense or rotated parameter couplings, including those in matrix factorization, tensor decomposition, and graph neural networks, because it treats each parameter independently. We introduce FlowAdam, a hybrid optimizer that augments Adam with continuous gradient-flow integration via an ordinary differential equation (ODE). When EMA-based statistics detect landscape difficulty, FlowAdam switches to clipped ODE integration. Our central contribution is Soft Momentum Injection, which blends ODE velocity with Adam's momentum during mode transitions. This prevents the training collapse observed with naive hybrid approaches. Across coupled optimization benchmarks, the ODE integration provides implicit regularization, reducing held-out error by 10-22% on low-rank matrix/tensor recovery and 6% on Jester (real-world collaborative filtering), also surpassing tuned Lion and AdaBelief, while matching Adam on well-conditioned workloads (CIFAR-10). MovieLens-100K confirms benefits arise specifically from coupled parameter interactions rather than bias estimation. Ablation studies show that soft injection is essential, as hard replacement reduces accuracy from 100% to 82.5%.
Figures & tables
Fig. 1: Ablation: Hard replacement (green) induces a state mismatch by overwriting Adam’s momentum buffer, causing a late-stage instability (final accuracy 82.5%). FlowAdam with soft injection (blue) blends the ODE velocity with the existing momentum, avoiding the mismatch and recovering stable convergence (100.0%).
Scenario
Adam
AdamW*
FlowAdam
Improv.
Small Dense
0.098±0.001
0.115±0.003
0.088±0.001
10.2%
Medium Mod.
0.129±0.001
0.181±0.004
0.111±0.001
14.0%
Larger Sparse
0.303±0.005
0.640±0.015
0.237±0.001
21.8%
*Best AdamW across λ∈{10−5,10−3,10−2} . Improv. vs. Adam.
TABLE I: Matrix Completion: Test RMSE (5 seeds). Mode B.
Fig. 2: Matrix Completion: Larger Sparse scenario (400 × 500, rank=20, 15% observed). Left: Training RMSE (log scale). Right: Test RMSE (linear scale). Representative single-seed run. Despite higher training RMSE (left), FlowAdam achieves 23.0% lower test RMSE (right; 5-seed mean: 21.8%, Table I ), demonstrating implicit regularization.
Scenario
Adam
FlowAdam
Improv.
Small Heavy
0.736±0.053
0.619±0.031
15.9%
Medium Heavy
0.955±0.020
0.822±0.016
13.9%
Large Heavy
0.938±0.054
0.819±0.014
12.7%
TABLE II: Robust Matrix Factorization: Test RMSE (5 seeds). Mode B.
Scenario
Adam
AdamW
FlowAdam
Improv.
Small Sparse
0.071±0.002
0.083±0.008
0.063±0.001
11.3%
Medium Sparse
0.068±0.002
0.105±0.027
0.060±0.001
11.8%
Larger Sparse
0.055±0.001
0.067±0.005
0.049±0.001
10.9%
TABLE III: Tensor Completion: Test RMSE (5 seeds). Mode B.
Scenario
Adam
FlowAdam
Improv.
Medium Strong
0.770±0.008
0.793±0.009
3.0%
Large Moderate
0.730±0.009
0.758±0.014
3.8%
Larger Challenge
0.698±0.013
0.723±0.011
3.6%
Easy Calib.
0.941±0.003
0.942±0.003
–
TABLE IV: GNN Link Prediction: Test AUC (5 seeds). Mode B.
Optimizer
Target RMSE
Median
Improv.
Adam
0.182±0.080
0.208
–
AdamW
0.182±0.080
0.208
–
FlowAdam
0.144±0.070
0.193
20.9%
TABLE V: Inverse Kinematics: Target RMSE (5 seeds). Mode B.
Optimizer
Test RMSE
Improv.
Adam (1K steps)
4.731±0.014
–
Adam Ext (4,350 grad evals)
5.007±0.019
−5.8%
FlowAdam (1K steps)
4.427±0.009
6.4%
TABLE VI: Jester: Test RMSE (5 seeds). Mode B.
Optimizer
Test RMSE
Std
vs. FlowAdam
Lion
4.271
0.008
+2.0%
AdaBelief
4.495
0.022
+7.4%
L-BFGS
5.083
0.017
+21.4%
FlowAdam
4.186
0.013
–
TABLE VII: Comparison with Alternative Optimization Strategies (Jester). All methods tuned via grid search (5 seeds).
Fig. 3: Compute-matched comparison on Medium Matrix Completion. Left: RMSE vs. gradient evaluations. Right: RMSE vs. wall-clock time. Horizontal dotted lines indicate converged test RMSE.
Fig. 4: Sensitivity analysis on matrix completion (Mode B, 5 seeds). Left: Performance is robust across γ∈[0.1,0.9] ; asterisk marks default γ=0.5 . Right: FlowAdam (fixed λ=10−5 ) vs. Adam across regularization strengths. The x-axis shows regularization configurations: Adam at 0.5× , 1× (matched), and 2× the regularization strength of FlowAdam, plus FlowAdam at 0.5× regularization. FlowAdam outperforms Adam at all tested configurations.