How Memory in Optimization Algorithms Implicitly Modifies the Loss
This provides theoretical insights into optimization dynamics for deep learning practitioners, though it is incremental as it builds on existing memory analysis.
The paper tackles the problem of understanding how memory in optimization algorithms affects loss functions, by developing a technique to approximate memory-based algorithms with memoryless ones and identifying implicit regularization effects. As a result, it explains that Lion lacks the implicit anti-regularization found in AdamW, correlating with Lion's better generalization performance.
In modern optimization methods used in deep learning, each update depends on the history of previous iterations, often referred to as memory, and this dependence decays fast as the iterates go further into the past. For example, gradient descent with momentum has exponentially decaying memory through exponentially averaged past gradients. We introduce a general technique for identifying a memoryless algorithm that approximates an optimization algorithm with memory. It is obtained by replacing all past iterates in the update by the current one, and then adding a correction term arising from memory (also a function of the current iterate). This correction term can be interpreted as a perturbation of the loss, and the nature of this perturbation can inform how memory implicitly (anti-)regularizes the optimization dynamics. As an application of our theory, we find that Lion does not have the kind of implicit anti-regularization induced by memory that AdamW does, providing a theory-based explanation for Lion's better generalization performance recently documented.