Mechanistic origin of the diminishing WSD–SWA gap at large batch sizes

Investigate the mechanism underlying the empirical finding that the performance difference between a constant learning rate with weight averaging and the Warmup–Stable–Decay schedule decreases with dataset size N in the large-batch regime, rigorously testing the hypothesis that the deterministic edge-of-stability phenomenon—where gradient descent fails to learn features in the top subspace—accounts for this behavior, and characterize the conditions under which it occurs.

Background

In the large-batch ablation, the authors observe that constant learning rate with averaging increasingly closes the gap to WSD as N grows. They hypothesize that this may relate to deterministic edge-of-stability dynamics that limit learning of top-subspace features but do not provide a definitive explanation.

They explicitly defer a precise study of this phenomenon to future work, leaving open a targeted analysis to confirm the mechanism and delineate when it arises.