DOSA: Dynamic Online State Allocation for Adaptive Optimizers via Per-Tensor Sketched Smoothness Tests
Abstract
Which parameters of an adaptive optimizer should retain exact second-moment state, and which can be compressed? Memory-efficient methods (e.g. Adafactor, GaLore, 8-bit Adam) commit at design time. On transformer LM training, AdamW second moments concentrate on a small fraction of coordinates, with one dominant tensor typically carrying most of the adaptive mass. DOSA treats state allocation as an online statistical decision: promote the candidate update rule with the largest lower confidence bound (LCB) on a smoothness-model predicted-descent score. The main result is an LCB-certificate separation theorem: under heavy-tail concentration, any uniform-width sketch meeting the same identification certificate as per-tensor adaptive widths requires asymptotically more memory. We formulate wwo decision rules that are connected by a horizon-parameterized surrogate: the one-step endpoint recovers the predicted-descent score used inside LCB-greedy, while a closed-form anchored-residual blend represents the long-horizon quadratic-surrogate endpoint. A swept blend improves validation perplexity over a factored-Adafactor baseline across model scales and datasets, and the closed-form long-horizon surrogate rule recovers about half the swept gain without tuning on DistilGPT2/WikiText-2.