Reproduce “Torque-Aware Momentum” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Pranshu Malviya, Goncalo Mordido, Aristide Baratin, Reza Babanezhad Harikandeh, Gintare Karolina Dziugaite, Razvan Pascanu, Sarath Chandar
Researchers from Mila, Samsung SAIT AI Lab, and Google DeepMind developed Torque-Aware Momentum (TAM) and its adaptive variant, AdaTAM, which stabilize deep neural network training by applying a damping factor to gradients based on their alignment with previous momentum. This approach leads to improved generalization performance across image classification and large language model fine-tuning, and demonstrates enhanced robustness to distribution shifts in online learning.
Key insights
- By mitigating the influence of "torqued" gradients, the optimizer can achieve more consistent directional updates, leading to improved exploration and better navigation of complex loss landscapes.
- The analogy of an anisotropic friction-like damping term, dependent on the angle between momentum and gradient, provides an intuitive and effective mechanism for stabilizing optimization.
- TAM can effectively serve as a warm-up strategy, guiding models to better-generalizing regions early in training, characterized by lower loss barriers and delayed oscillatory behavior.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.