Adam and its variants have dominated neural network optimization for nearly a decade. The core formula is well established: maintain exponential moving averages of the gradient and its square, bias-correct them, and scale the update by the ratio of the first moment to the square root of the second. Francisco Caldas, Ruben Belo, and Cláudia Soares from Universidade NOVA de Lisboa take a different approach, adding a single multiplicative scaling factor derived from the cosine similarity between consecutive gradients. The result, AdamX, is a straightforward modification to the Adam update rule that introduces a gradient-direction signal without adding architectural complexity or new hyperparameters.
The Core Idea: Gradient Consistency as a Scaling Factor
The key mechanism is simple. At each step, AdamX computes the cosine similarity between the current gradient and the previous gradient, both flattened to 1-D vectors. This value, which ranges from -1 (opposite directions) to 1 (identical directions), is exponentiated with a scaling parameter to produce a multiplicative factor. When consecutive gradients point in similar directions, the cosine similarity is high and the exponentiated factor amplifies the update. When gradients flip direction or become orthogonal, the factor suppresses the update.
Formally, the scaling factor is gamma = exp(lambda * cos(g_t, g_{t-1})), where g_t is the current gradient and g_{t-1} is the previous gradient. The parameter lambda controls the sensitivity. When lambda equals zero, gamma collapses to 1 and AdamX reduces to standard Adam. The cosine similarity is computed per-parameter across all elements of the weight tensor, meaning it captures the global direction consistency of updates for each parameter group rather than element-wise agreement.
The update rule itself is a direct modification of Adam. The authors compute bias-corrected first and second moments as usual, then scale the first moment by gamma before dividing by the square root of the second moment plus epsilon. The full update is:
m_hat_t = m_t / (1 - beta1^t)
v_hat_t = v_t / (1 - beta2^t)
gamma_t = exp(lambda * cosine_similarity(g_t, g_{t-1}))
theta_t = theta_{t-1} - lr * gamma_t * m_hat_t / (sqrt(v_tilde_t) + eps)
where v_tilde is a variance rectification term described below. The cosine similarity introduces no new per-parameter state beyond what Adam already maintains. The previous gradient is stored as a single tensor, the same size as any running average.
Variance Rectification for Early Training Stability
A practical problem with Adam during early training is that the second-moment estimate can be noisy and small when the model is far from convergence. This produces large, erratic updates. AdamX addresses this with a variance rectification scheme that replaces the standard second moment with a running maximum.
The rectified denominator is v_tilde_t = max(v_hat_t, v_tilde_{t-1}), where the maximum is taken element-wise. This means the denominator can only increase or stay the same over time, never decrease. In practice, during early training when the second moment is volatile, the rectification prevents the denominator from dropping and producing large spikes. As training progresses and the second moment stabilizes, it eventually dominates and the rectification has no further effect.
The authors note that this is distinct from AMSGrad, which maintains the maximum of past second moments but does not use it as the denominator in the update. In AdamX, the rectified second moment is used directly in the division, making it a hard floor on the effective step size.
Optional Loss Revert Mechanism
AdamX includes an optional mechanism to revert parameter updates when they increase the loss. After the main update step, if a closure function is provided, the optimizer re-evaluates the loss. If the new loss exceeds the old loss, the parameters are reverted to their previous values plus a fraction of the attempted update, controlled by a parameter called cc. When cc is zero (the default), no revert occurs. When cc is 1.0, the revert is partial. When cc is negative, the revert overshoots in the opposite direction.
This mechanism is borrowed from line-search-style optimizers and adds computational cost because it requires an additional forward and backward pass per step. The authors include it as an option for scenarios where training stability is critical, but note that it is not necessary for competitive performance on the benchmarks they evaluate.
Experimental Setup and Baselines
The authors evaluate AdamX on three datasets: MNIST with a simple MLP and CNN, CIFAR-10 with a multi-layer CNN and a CifarNet architecture, and OGBG-MolPCBA with a graph neural network. The baseline optimizers include Adam, AdamW, Adagrad, AMSGrad, RAdam, RMSprop, SGD with momentum, Lion, Yogi, Adan, and AdamHD. All optimizers use the same hyperparameter budget, with performance measured by the number of epochs required to reach predefined accuracy thresholds.
The evaluation protocol uses a fixed hyperparameter budget, meaning all optimizers receive the same search space and the same number of random configurations. This avoids the common practice of tuning one optimizer more carefully than others. Performance is reported in terms of epochs to reach target accuracy, which directly measures convergence speed rather than final accuracy.
Experiments are tracked through Weights and Biases, with the project publicly available. The codebase is organized with separate directories for datasets, models, optimizers, training loops, and configuration files, making it straightforward to reproduce results.
What the Numbers Show
The paper reports that AdamX achieves competitive convergence rates across the benchmark suite. On CIFAR-10 with the ImprovedCNN architecture, AdamX converges to the target accuracy in fewer epochs than Adam, Lion, and Adan. On OGBG-MolPCBA, a molecular property prediction task with graph neural networks, AdamX matches or exceeds the performance of all baselines. On MNIST, the differences are smaller because the task is relatively easy, but AdamX remains competitive.
The variance rectification scheme is particularly effective during early training. The running-max denominator prevents the large initial steps that sometimes cause Adam to diverge on sensitive architectures. The authors show that without rectification, AdamX still outperforms Adam, but with more variance in the early epochs.
The cosine similarity scaling factor produces values that cluster around 1.0 when gradients are consistent (typical of late training) and fluctuate widely when gradients are noisy (typical of early training or difficult loss landscapes). This adaptive behavior is automatic, requiring no manual scheduling of the learning rate or the scaling parameter.
Integration and Practical Considerations
AdamX is implemented as a drop-in replacement for Adam in PyTorch. It accepts the same hyperparameters as Adam (learning rate, betas, epsilon) plus three new parameters: alpha (default 0.99), lambda_exp (default 1.0), and cc (default 0.0). The alpha parameter appears in the code but does not affect the core update rule as currently implemented. The lambda_exp parameter controls the sensitivity of the cosine similarity scaling. The cc parameter controls the loss-revert behavior.
The only additional memory overhead compared to Adam is storing the previous gradient, which is the same size as the gradient tensor itself. For a model with N parameters, this adds N floats of storage, which is negligible for most practical settings. The cosine similarity computation adds a dot product and two norm computations per step, both of which are dominated by the backward pass.
The authors note that the method is model-agnostic and can be integrated into any training pipeline that uses standard PyTorch optimizers. The public repository includes training scripts, configuration files, and experiment tracking through Weights and Biases, making reproduction straightforward.
Limitations and Open Questions
The paper evaluates AdamX on relatively small-scale benchmarks. The largest model is a graph neural network on OGBG-MolPCBA, which has about 200K parameters. Whether the cosine similarity scaling provides the same benefit on transformer-scale models with billions of parameters remains untested. The scaling behavior of cosine similarity across very high-dimensional parameter spaces is an open question.
The variance rectification scheme uses an element-wise maximum, which may interact poorly with parameters that have very different gradient magnitudes. The authors do not ablate the rectification separately from the cosine similarity scaling, so it is difficult to isolate the contribution of each mechanism. The loss-revert mechanism adds computational cost and is not evaluated on the main benchmarks, leaving its practical value unclear.
The paper does not provide theoretical convergence guarantees for AdamX. The cosine similarity scaling factor can be unbounded (exponentiated values grow exponentially with lambda), and while the authors report stable training in practice, there are no formal bounds on the step size or convergence rate. This is consistent with the broader literature on adaptive optimizers, where practical performance often outpaces theoretical understanding.
Read the paper on arXiv