Adam Optimizer's Implicit Bias Depends on Learning Rate Schedule: New Research Reveals

In the field of deep learning optimization, the Adam optimizer is widely popular due to its adaptive learning rates, but the mechanism of its implicit bias has long been debated. A recent paper, "Watch Your Step: Adam's Margin Maximization Bias Depends on the Learning Rate Schedule," delves into Adam's implicit bias in binary classification tasks with homogeneous models and reveals a key finding: Adam's implicit bias strongly depends on the learning rate schedule. This result challenges the previous understanding that Adam always tends to maximize the ℓ∞ margin.

Background: Similarity Between Adam and Sign Gradient Descent

Prior research has shown that on homogeneous models, Adam behaves similarly to sign gradient descent and exhibits an ℓ∞ margin maximization implicit bias. Specifically, when the learning rate schedule is strongly decaying, i.e., η(t) = o(t^{1/L-1}), where L is the homogeneity degree, Adam is proven to converge to a solution that maximizes the ℓ∞ margin. However, such strongly decaying schedules are uncommon in practice; more common are constant learning rates or weakly decaying schedules (e.g., polynomial decay t^α).

This paper focuses on this more common complementary regime—constant and weakly decaying learning rate schedules—and reaches a conclusion starkly different from prior work.

Homogeneous models refer to models satisfying (f(c\mathbf{w}) = c^L f(\mathbf{w})) for any scalar (c>0), where (L) is called the homogeneity degree. Common deep linear networks and ReLU networks (ignoring bias terms) possess this property. Homogeneity makes the scaling behavior of the loss function predictable, thereby facilitating analysis of the asymptotic direction of the optimizer's implicit bias. For example, for an (L)-homogeneous model, the gradient satisfies (\nabla f(c\mathbf{w}) = c^{L-1} \nabla f(\mathbf{w})), a property that plays a key role in analyzing the limiting behavior of Adam's preconditioner.

Core Finding: No Trajectory-Independent Margin Maximization Bias

For polynomial schedules η(t) = t^α and homogeneity degree L, the paper proves that when α > 1/L - 1/2, Adam in deep linear networks does not have a trajectory-independent margin maximization bias (in the KKT sense). Specifically:

• First, under this setting, there is no nontrivial, locally Lipschitz convex implicit regularizer, including any squared norm.
• Second, on width-1 networks, the result holds without the convexity assumption, meaning that any implicit regularizer for wide networks must be trivial when restricted to width-1 networks.

This finding indicates that under common weakly decaying schedules, Adam does not drive the model toward maximum-margin solutions as strongly decaying schedules do. Instead, its implicit bias is trajectory-dependent.

Implicit regularizer refers to some norm or function that the optimization algorithm implicitly minimizes during training, determining which type of solution the algorithm prefers. If the optimizer has a trajectory-independent implicit regularizer (R(\mathbf{w})), its converged solution can be viewed as the solution that minimizes (R(\mathbf{w})) under the constraint of zero training loss, independent of initialization. Conversely, if the implicit bias is trajectory-dependent, the direction of the final solution will vary with initialization or optimization path. The paper proves that under weakly decaying schedules, Adam has no nontrivial locally Lipschitz convex implicit regularizer, meaning its preference cannot be characterized by a fixed convex function, thereby challenging the simple picture that "Adam always maximizes the margin."

Positive Characterization: Trajectory-Dependent Weighted-ℓ2 Margin Maximization

Although there is no trajectory-independent margin maximization bias, the paper still provides a positive characterization of Adam's implicit bias: in this regime, for any L-homogeneous model, Adam maximizes a trajectory-dependent weighted-ℓ2 margin. The fundamental reason behind this mechanism is that Adam's preconditioner converges to a trajectory-dependent vector, which also explains why the limiting parameter direction depends on initialization.

In other words, Adam's implicit bias is not fixed but varies with the optimization trajectory and initialization. This provides a more nuanced perspective for understanding Adam's behavior.

Preconditioner in Adam is embodied by the square root of the second-moment estimate (\sqrt{\mathbf{v}_t}), which scales the gradient coordinate-wise. Under strongly decaying schedules, the preconditioner converges to a fixed vector (e.g., the all-ones vector), thereby inducing ℓ∞ margin maximization; under constant or weakly decaying schedules, the preconditioner converges to a vector that depends on the trajectory and initialization, causing the weights in the weighted ℓ2 margin to also be trajectory-dependent. This limiting behavior of the preconditioner is the core mechanism for understanding how Adam's implicit bias changes with the schedule.

Intermediate Regime: Conditional ℓ∞ Bias

In the intermediate regime 1/L - 1 ≤ α ≤ 1/L - 1/2, the paper proves an ℓ∞ bias, but this conclusion relies on the condition that the gradient direction converges to a full-support vector. This indicates that as the learning rate decay rate varies, Adam's implicit bias transitions from trajectory-dependent weighted-ℓ2 margin maximization to ℓ∞ margin maximization, but the latter requires additional conditions.

Independent Contribution: Loss Decay Rate of Full-Batch Adam

As an independently interesting contribution, the paper also establishes a loss decay rate for full batch Adam on general homogeneous models, whose leading order depends only on the second-moment parameter. This result helps understand Adam's convergence behavior during training and provides a theoretical tool for future research.

Practical Implications and Conclusion

This research reminds us that the choice of learning rate schedule not only affects convergence speed but also profoundly influences the implicit bias induced by the optimizer. For practitioners, when using Adam with a constant or weakly decaying learning rate, one should not expect the model to automatically tend toward maximum-margin solutions; instead, the properties of the solution may depend more on initialization and training trajectory. This finding is crucial for understanding the generalization performance of deep learning and provides new theoretical foundations for the design of optimization algorithms.

Future work can further explore the implicit biases of other optimizers (such as AdaGrad, RMSProp) under different learning rate schedules, as well as the manifestation of these theoretical results in more complex models (such as convolutional networks and Transformers).

One-line Summary

Adam's implicit bias depends on the learning rate schedule: under weak decay, there is no trajectory-independent margin maximization, but rather trajectory-dependent weighted-ℓ2 margin maximization.

Paragraph Summary

The paper "Watch Your Step: Adam's Margin Maximization Bias Depends on the Learning Rate Schedule" challenges the understanding that Adam always maximizes the ℓ∞ margin. The authors prove that for polynomial learning rate schedules η(t)=t^α, when α > 1/L - 1/2, Adam in deep linear networks has no trajectory-independent margin maximization bias and no nontrivial convex implicit regularizer. Instead, Adam maximizes a trajectory-dependent weighted-ℓ2 margin, rooted in the preconditioner converging to a trajectory-dependent vector, causing the limiting parameter direction to depend on initialization. In the intermediate regime 1/L - 1 ≤ α ≤ 1/L - 1/2, the paper proves a conditional ℓ∞ bias that depends on the gradient direction converging to a full-support vector. Additionally, the paper establishes the loss decay rate of full-batch Adam on general homogeneous models, whose leading order depends only on the second-moment parameter. These results profoundly reveal the impact of learning rate schedules on the optimizer's implicit bias.