How Memory in Optimization Algorithms Implicitly Modifies the Loss

TL;DR

This paper introduces a technique to analyze how memory in optimization algorithms affects the loss landscape, revealing differences in implicit regularization properties that influence generalization in deep learning.

Contribution

The authors develop a method to approximate memory-based optimization algorithms with memoryless ones plus correction terms, linking memory effects to implicit regularization.

Findings

Lion lacks the implicit anti-regularization of AdamW due to memory effects.

Memory influences the loss landscape, affecting optimization dynamics and generalization.

The theory explains why Lion may generalize better than algorithms with memory.

Abstract

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.