8 papers
Proximity-Based Multi-Turn Optimization: Practical Credit Assignment for LLM Agent Training
Yangyi Fang, Jiaye Lin, Xiaoliang Fu +4
Multi-turn LLM agents are becoming pivotal to production systems, spanning customer service automation, e-commerce assistance, and interactive task management, where accurately dis…
Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectory
Runa Eschenhagen, Anna Cai, Tsung-Hsien Lee +1
Optimizers leveraging the matrix structure in neural networks, such as Shampoo and Muon, are more data-efficient than element-wise algorithms like Adam and Signum. While in specifi…
How to Allocate, How to Learn? Dynamic Rollout Allocation and Advantage Modulation for Policy Optimization
Yangyi Fang, Jiaye Lin, Xiaoliang Fu +6
Reinforcement Learning with Verifiable Rewards (RLVR) has proven effective for Large Language Model (LLM) reasoning, yet current methods face key challenges in resource allocation…
Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent
Hiroki Naganuma, Shagun Gupta, Youssef Briki +4
To maximize hardware utilization, modern machine learning systems typically employ large constant or manually tuned batch size schedules, relying on heuristics that are brittle and…
Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs
Aaron Defazio, Konstantin Mishchenko, Parameswaran Raman +2
We propose Generalized Primal Averaging (GPA), an extension of Nesterov's method that unifies and generalizes recent averaging-based optimizers like single-worker DiLoCo and Schedu…
SNOO: Step-K Nesterov Outer Optimizer - The Surprising Effectiveness of Nesterov Momentum Applied to Pseudo-Gradients
Dominik Kallusky, Vinay Rao, Vishal Nandavanam +1
The rapid development of large language models (LLMs) has driven the demand for more efficient optimization techniques. Among these, the Lookahead family of optimizers employs a tw…
Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner
Runa Eschenhagen, Aaron Defazio, Tsung-Hsien Lee +2
The recent success of Shampoo in the AlgoPerf contest has sparked renewed interest in Kronecker-factorization-based optimization algorithms for training neural networks. Despite it…
A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale
Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki +5
Shampoo is an online and stochastic optimization algorithm belonging to the AdaGrad family of methods for training neural networks. It constructs a block-diagonal preconditioner wh…