Implicit Bias of Per-sample Adam on Separable Data: Departure from the Full-batch Regime

AI summary

Authors: Beomhan Baek (Seoul National University), Minhak Song (KAIST), Chulhee Yun (KAIST) — arXiv:2510.26303v3 [cs.LG], CC BY 4.0

Overview

Research area: Optimization theory for machine learning — specifically the implicit bias of adaptive gradient methods (Adam, Signum) in linear classification on linearly separable data.

Technical level: Advanced. The paper is a fully theoretical treatment built on convergence proofs, margin-maximization duality, and fixed-point arguments. The intuition is accessible, but the machinery (parametric quadratic programs, dual fixed points, learning-rate schedules) assumes graduate-level optimization background.

Scope in one sentence: The paper characterizes the direction to which incremental (one-sample-per-step) Adam converges on separable logistic regression, showing it can leave the ℓ∞-max-margin bias of full-batch Adam, and contrasts this with Signum, which keeps the ℓ∞-max-margin bias under any batch size.

What This Paper Is About

When a linear classifier is trained with logistic or exponential loss on data that can be perfectly separated, the weights grow without bound, and what matters is the direction they point in — the decision boundary. Prior work showed that full-batch gradient descent picks the ℓ2-max-margin direction, while full-batch Adam picks the ℓ∞-max-margin direction. Almost all of that theory, however, assumes full-batch updates, whereas real training uses mini-batches. This paper asks whether Adam's ℓ∞ bias survives when gradients come from one sample at a time, and finds that it does not: per-sample Adam can converge to the ℓ2-max-margin direction, or to a new data-dependent direction that is neither ℓ2 nor ℓ∞.

Key Contributions

1. An epoch-wise approximation of incremental Adam. The authors show that despite Adam's momentum-based updates, the epoch-wise dynamics of Inc-Adam (one sample per step, processed cyclically) can be approximated by a recurrence depending only on the current iterate, making asymptotic analysis tractable (Proposition 2.5).

2. A sharp full-batch vs. mini-batch contrast on Scaled Rademacher (SR) data. On SR data — where every coordinate of each data point has the same magnitude — incremental Adam provably converges directionally to the ℓ2-max-margin solution (Theorem 3.3), whereas full-batch Adam converges to the ℓ∞-max-margin solution.

3. A general data-adaptive characterization via a proxy algorithm. For arbitrary datasets, the authors introduce AdamProxy, a uniform-averaging proxy valid as β2 → 1, and prove its convergence direction is the optimizer of a parametric margin-maximization problem under a data-adaptive Mahalanobis norm whose covariance matrix is pinned down by a dual fixed-point equation (Theorem 4.8).

4. A counterpoint: Signum keeps ℓ∞ bias. The paper proves that Signum (SignSGD with momentum, Bernstein et al., 2018) converges to the ℓ∞-max-margin direction for any batch size, provided momentum is sufficiently close to 1 (Theorem 5.1), generalizing the full-batch result of Fan et al. (2025).

Main Findings

• Mini-batch Adam loses the ℓ∞ bias. In an experiment on 10 data points drawn from the 50-dimensional standard Gaussian, full-batch Adam with (β1, β2) = (0.9, 0.95) converges to the ℓ∞-max-margin solution, while mini-batch variants with a batch size of 1 converge to a different direction that is closer to the ℓ2-max-margin solution (Figure 1).

• On SR data, incremental Adam provably goes to ℓ2. Under separability, nonzero coordinates, and a polynomial learning-rate schedule η_t = (t+2)^(−a) with a ∈ (2/3, 1], and assuming the loss converges to 0, Inc-Adam with β1 ≤ β2 converges directionally to the unique ℓ2-max-margin solution of the SR data (Theorem 3.3). The result holds for any momentum hyperparameters with β1 ≤ β2, because SR data removes the coordinate adaptivity that momentum would normally shape.

• The contrast is confirmed empirically on a concrete SR dataset. On x₀ = (1,1,1,1), x₁ = (2,2,2,−2), x₂ = (3,3,−3,−3), and x₃ = (4,−4,4,−4), mini-batch Adam variants with batch size 1 consistently converge to the ℓ2-max-margin direction while full-batch Adam converges to the ℓ∞ one (Figure 2).

• For general data, the limit is neither ℓ2 nor ℓ∞. The AdamProxy limit direction takes the form ŵ ∝ (Σᵢ cᵢxᵢ) / sqrt(Σᵢ cᵢ²xᵢ²), with cᵢ = 0 outside the set of support vectors (Lemma 4.5). The coefficients c are simultaneously the parameters of a parametric margin problem and its dual variables, and the direction is identified at the fixed point where these two roles coincide (Theorem 4.8).

• Two examples show the framework reduces to known classifiers. For SR data, the Mahalanobis matrix M(c) collapses to a scaled identity for every c on the simplex, so the parametric problem becomes the standard SVM formulation and the limit is the ℓ2-max-margin solution (Example 4.9). For shifted-diagonal data, the ℓ∞-max-margin solution is a candidate limit for AdamProxy; empirically, all mini-batch Adam variants with batch size 1 converge to it on the dataset with δ = 0.1 (Example 4.11, Figure 4).

• Gaussian data is validated against the fixed point. On the same Gaussian dataset as Figure 1, both AdamProxy and mini-batch Adam variants with batch size 1 converge to the fixed-point solution computed by the fixed-point iteration (Example 4.10, Figure 3).

• Signum is batch-size invariant. For any batch size b and momentum β ∈ (1−ε, 1), Inc-Signum satisfies liminf of the normalized minimum margin ≥ γ∞ − δ, where γ∞ = max over ‖w‖∞ ≤ 1 of min over i of wᵀxᵢ, D = max over i of ‖xᵢ‖₁, and ε = (1/(2D · (N/b)(N/b−1))) · min{δ, γ∞/2} when b < N, with ε = 1 when b = N (Theorem 5.1).

• Det-Adam is essentially sign gradient descent. Under a stated condition, the k-th coordinate update of full-batch Adam equals −η_t(sign(∇L(w_t)[k]) + ε_t) with ε_t → 0 (Proposition 2.4), consistent with earlier observations (Balles and Hennig, 2018; Zou et al., 2023). Inc-Adam instead behaves like weighted, preconditioned gradient descent (Proposition 2.5), with the residual vector bounded as O(r^(−a/2)) for polynomial schedules.

• A key tension: the preconditioner. The discrepancy arises because Inc-Adam's preconditioner tracks the sum of squared mini-batch gradients, which diverges from the squared full-batch gradient — so the behavior depends on both the momentum parameters and the current iterate, unlike the full-batch case.

Methodology in Plain English

The authors study binary linear classification with logistic and exponential loss on data that is linearly separable, using a setup where the weights diverge and only their direction matters. They analyze Inc-Adam: one data point per step, cycling through indices 0 to N−1 each epoch. Three steps structure the argument.

First, they prove an approximation: the messy, history-dependent Adam update can, epoch by epoch, be replaced by an expression that depends only on the current iterate, up to an error term that vanishes. For full-batch Adam this reduces to sign gradient descent; for incremental Adam it becomes a weighted, preconditioned gradient step.

Second, they exploit structured data. Scaled Rademacher data — where each data point's coordinates all have equal magnitude — kills the coordinate-wise adaptivity in the denominator, leaving weighted normalized gradient descent with weights bounded between two positive constants. A careful analysis of this weighted flow shows the direction is the ℓ2-max-margin solution.

Third, for general data they move to a simpler proxy valid as β2 → 1, a limit justified by the fact that β2 is usually chosen near 1 in practice. The proxy update is the full-batch gradient divided by the square root of the sum of squared per-sample gradients. To describe its limit direction they set up a parametric optimizer: minimize a Mahalanobis-norm margin, where the diagonal matrix depends on coefficients c on the probability simplex. The trick is that c plays two roles at once — the parameter defining the problem and the dual solution the problem returns — so the answer is found at a fixed point of the map from c to its normalized dual solution. They prove this map is continuous on a compact simplex and hence has at least one fixed point, and that the proxy's limit direction is the primal solution there. A numerical fixed-point iteration is used to compute the prediction in experiments; the paper explicitly states that a formal convergence guarantee for this iteration is not established.

Why This Matters

The paper shows that the widely taught picture — "Adam prefers ℓ∞ geometry, gradient descent prefers ℓ2" — is incomplete. Practically every modern training run uses mini-batches, and the result says the batching scheme, not just the optimizer name, determines which of the many interpolating solutions is selected. This is one of the first theoretical accounts of a data-dependent implicit bias for Adam at batch size 1, and it sharpens the message from Nacson et al. (2019) that batch size does not disturb the ℓ2 bias of SGD under Adam-style adaptivity: for adaptive methods, it does.

Potential application areas (note: the paper itself reports no applied experiments; these follow from its implications):

• Large language model pretraining, where Adam is the de facto optimizer and mini-batches are universal — relevant to claims that Adam's ℓ∞ geometry drives its empirical success (Xie et al., 2025).
• Linear or last-layer probing and classification on cleanly separable features, where the choice of Adam vs. Signum controls whether an ℓ2- or ℓ∞-style solution is reached.
• Optimizer selection for resource-constrained training, where small batch sizes are forced by memory and the paper gives a reason to prefer Signum-style sign updates if ℓ∞-type solutions are desired.
• Debugging unexpected generalization behavior in adaptive optimizers, since the paper supplies a concrete procedure to compute the predicted limit direction for a given dataset.

Industry relevance: Teams that tune batch size and β2 for convergence or generalization now have a theoretical reason to expect the target of optimization to shift with those choices, not merely the speed of getting there.

Future Directions

1. Extending beyond β2 → 1. The fixed-point framework is derived for the β2 → 1 limit; the authors state in Appendix A (Section A.2) that it does not directly extend to general β2 < 1, and leave that generalization open.

2. Establishing convergence of the fixed-point iteration. The paper notes that the complexity of the map T makes a formal convergence guarantee for Algorithm 3 difficult, and explicitly leaves a rigorous analysis for future work.

3. General batch sizes. The analysis assumes batch size 1. Appendix A provides only a preliminary analysis of how batch size and momentum hyperparameters affect mini-batch Adam's implicit bias, leaving the intermediate regime largely uncharacterized.

4. Closing the gap between Inc-Adam and true stochastic Adam. The paper positions Inc-Adam as a faithful surrogate whose directional convergence aligns with batch-size-1 Adam under sampling with replacement and random reshuffling; a direct analysis of those sampling schemes would test how far the surrogate reaches.

Target Audience

Optimization and learning-theory researchers working on implicit bias, edge-of-stability analyses, and the theory of adaptive gradient methods; graduate students who have read Soudry et al. (2018) and Zhang et al. (2024a) and want the mini-batch continuation; and practitioners with a strong theoretical bent who choose between Adam and Signum and want to know what direction their optimizer is actually heading toward. Readers seeking empirical benchmark numbers, dataset-scale evaluations, or guidance on non-separable or deep-network settings will not find them here — the paper is a theory contribution with small illustrative experiments only.

Authors’ abstract

Adam [Kingma &amp; Ba, 2015] is the de facto optimizer in deep learning, yet its theoretical understanding remains limited. Prior analyses show that Adam favors solutions aligned with $\ell_\infty$-geometry, but these results are restricted to the full-batch regime. In this work, we study the implicit bias of incremental Adam (using one sample per step) for logistic regression on linearly separable data, and show that its bias can deviate from the full-batch behavior. As an extreme example, we construct datasets on which incremental Adam provably converges to the $\ell_2$-max-margin classifier, in contrast to the $\ell_\infty$-max-margin bias of full-batch Adam. For general datasets, we characterize its bias using a proxy algorithm for the $β_2 \to 1$ limit. This proxy maximizes a data-adaptive Mahalanobis-norm margin, whose associated covariance matrix is determined by a data-dependent dual fixed-point formulation. We further present concrete datasets where this bias reduces to the standard $\ell_2$- and $\ell_\infty$-max-margin classifiers. As a counterpoint, we prove that Signum [Bernstein et al., 2018] converges to the $\ell_\infty$-max-margin classifier for any batch size. Overall, our results highlight that the implicit bias of Adam crucially depends on both the batching scheme and the dataset, while Signum remains invariant.