← Search

Kwangjun Ahn

18 accepted papers

2025

General framework for online-to-nonconvex conversion: Schedule-free SGD is also effective for nonconvex optimization

ICML 2025oral

This work investigates the effectiveness of schedule-free methods, developed by A. Defazio et al. (NeurIPS 2024), in nonconvex optimization settings, inspired by their remarkable empirical success in training neural networks. Specifically, we show that schedule-free SGD achieves optimal iteration co…

Cited by 1SourcePDFScholar
2025

The Belief State Transformer

ICLR 2025poster

We introduce the "Belief State Transformer", a next-token predictor that takes both a prefix and suffix as inputs, with a novel objective of predicting both the next token for the prefix and the previous token for the suffix. The Belief State Transformer effectively learns to solve challenging probl…

2025

Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training

NeurIPS 2025poster

As both model and dataset sizes continue to scale rapidly, conventional pretraining strategies with fixed compute budgets—such as cosine learning rate schedules—are increasingly inadequate for large-scale training. Recent alternatives, including warmup-stable-decay (WSD) schedules and weight averagi…

Cited by 0SourceScholar
2024

Adam with model exponential moving average is effective for nonconvex optimization

NeurIPS 2024poster

In this work, we offer a theoretical analysis of two modern optimization techniques for training large and complex models: (i) adaptive optimization algorithms, such as Adam, and (ii) the model exponential moving average (EMA). Specifically, we demonstrate that a clipped version of Adam with model E…

Cited by 5SourcePDFScholar
2024

Linear attention is (maybe) all you need (to understand Transformer optimization)

ICLR 2024poster

Transformer training is notoriously difficult, requiring a careful design of optimizers and use of various heuristics. We make progress towards understanding the subtleties of training Transformers by carefully studying a simple yet canonical linearized *shallow* Transformer model. Specifically, we…

Cited by 54SourcePDFScholar
2024

Understanding Adam Optimizer via Online Learning of Updates: Adam is FTRL in Disguise

ICML 2024poster

Despite the success of the Adam optimizer in practice, the theoretical understanding of its algorithmic components still remains limited. In particular, most existing analyses of Adam show the convergence rate that can be simply achieved by non-adative algorithms like SGD. In this work, we provide a…

Cited by 13SourcePDFScholar
2023

Learning threshold neurons via edge of stability

NeurIPS 2023poster

Existing analyses of neural network training often operate under the unrealistic assumption of an extremely small learning rate. This lies in stark contrast to practical wisdom and empirical studies, such as the work of J. Cohen et al. (ICLR 2021), which exhibit startling new phenomena (the "edge of…

Cited by 47SourcePDFScholar
2023

Transformers learn to implement preconditioned gradient descent for in-context learning

NeurIPS 2023poster

Several recent works demonstrate that transformers can implement algorithms like gradient descent. By a careful construction of weights, these works show that multiple layers of transformers are expressive enough to simulate iterations of gradient descent. Going beyond the question of expressivity,…

2022

Agnostic Learnability of Halfspaces via Logistic Loss

ICML 2022oral

We investigate approximation guarantees provided by logistic regression for the fundamental problem of agnostic learning of homogeneous halfspaces. Previously, for a certain broad class of “well-behaved” distributions on the examples, Diakonikolas et al. (2020) proved an tilde{Omega}(OPT) lower boun…

Cited by 7SourcePDFScholar
2022

Mirror Descent Maximizes Generalized Margin and Can Be Implemented Efficiently

NeurIPS 2022accept

Driven by the empirical success and wide use of deep neural networks, understanding the generalization performance of overparameterized models has become an increasingly popular question. To this end, there has been substantial effort to characterize the implicit bias of the optimization algorithms…

Cited by 27SourcePDFScholar
2022

Reproducibility in Optimization: Theoretical Framework and Limits

NeurIPS 2022accept

We initiate a formal study of reproducibility in optimization. We define a quantitative measure of reproducibility of optimization procedures in the face of noisy or error-prone operations such as inexact or stochastic gradient computations or inexact initialization. We then analyze several convex o…

Cited by 25SourcePDFScholar
2020

SGD with shuffling: optimal rates without component convexity and large epoch requirements

NeurIPS 2020spotlight

We study without-replacement SGD for solving finite-sum optimization problems. Specifically, depending on how the indices of the finite-sum are shuffled, we consider the RandomShuffle (shuffle at the beginning of each epoch) and SingleShuffle (shuffle only once) algorithms. First, we establish minim…

Cited by 89SourcePDFScholar