Anytime Pretraining: Horizon-Free Learning-Rate Schedules with Weight Averaging
Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade
arXiv 2026
Publications
Year
Topic
74 publications
2026
Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models
Samy Jelassi, Mujin Kwun, Rosie Zhao, Yuanzhi Li, Nicolo Fusi, Yilun Du, Sham Kakade, Carles Domingo-Enrich
arXiv 2026Prescriptive Scaling Reveals the Evolution of Language Model Capabilities
Hanlin Zhang, Jikai Jin, Vasilis Syrgkanis, Sham Kakade
arXiv 2026Scaling Reward Modeling without Human Supervision
Jingxuan Fan, Yue Li, Zhenting Qi, Dinghuai Zhang, Kianté Brantley, Sham Kakade, Hanlin Zhang, Hanlin Zhang
arXiv 2026Stop Training for the Worst: Progressive Unmasking Accelerates Masked Diffusion Training
Jaeyeon Kim, Jonathan Geuter, David Alvarez-Melis, Sham Kakade, Sitan Chen
arXiv 2026Weight Decay Improves Language Model Plasticity
Tessa Han, Sebastian Bordt, Hanlin Zhang, Sham Kakade
arXiv 2026
2025
The Role of Sparsity for Length Generalization in LLMs
Noah Golowich, Samy Jelassi, David Brandfonbrener, Sham Kakade, Eran Malach
International Conference on Machine Learning 2025Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, Sitan Chen
International Conference on Machine Learning 2025SOAP: Improving and Stabilizing Shampoo using Adam for Language Modeling
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
International Conference on Learning Representations 2025A Simplified Analysis of SGD for Linear Regression with Weight Averaging
Alexandru Meterez, Depen Morwani, Costin-Andrei Oncescu, Jingfeng Wu, Cengiz Pehlevan, Sham Kakade
arXiv 2025Adam or Gauss-Newton? A Comparative Study In Terms of Basis Alignment and SGD Noise
Bingbin Liu, Rachit Bansal, Depen Morwani, Nikhil Vyas, David Alvarez-Melis, Sham Kakade
arXiv 2025Any-Order Flexible Length Masked Diffusion
Jaeyeon Kim, C. Lee, Carles Domingo-Enrich, Yilun Du, Sham Kakade, Timothy Ngotiaoco, Sitan Chen, M. Albergo
arXiv 2025Cognitive models can reveal interpretable value trade-offs in language models
Sonia K. Murthy, Rosie Zhao, Jennifer Hu, Sham Kakade, Markus Wulfmeier, Peng Qian, Tomer D. Ullman
arXiv 2025Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants
Depen Morwani, Nikhil Vyas, Hanlin Zhang, Sham Kakade
arXiv 2025Discovering Hierarchical Latent Capabilities of Language Models via Causal Representation Learning
Jikai Jin, Vasilis Syrgkanis, Sham Kakade, Hanlin Zhang
arXiv 2025Distributional Scaling Laws for Emergent Capabilities
Rosie Zhao, Tian Qin, David Alvarez-Melis, Sham Kakade, Naomi Saphra
arXiv 2025Echo Chamber: RL Post-training Amplifies Behaviors Learned in Pretraining
Rosie Zhao, Alexandru Meterez, Sham Kakade, Cengiz Pehlevan, Samy Jelassi, Eran Malach
arXiv 2025EvoLM: In Search of Lost Language Model Training Dynamics
Zhenting Qi, Fan Nie, Alexandre Alahi, James Zou, Hima Lakkaraju, Yilun Du, Eric P. Xing, Sham Kakade, Hanlin Zhang
arXiv 2025Fine-Tuning Masked Diffusion for Provable Self-Correction
Jaeyeon Kim, Seunggeun Kim, Taekyun Lee, David Z. Pan, Hyeji Kim, Sham Kakade, Sitan Chen
arXiv 2025GQ-VAE: A gated quantized VAE for learning variable length tokens
Theo Datta, Kayla Huang, Sham Kakade, David Brandfonbrener
arXiv 2025In Good GRACEs: Principled Teacher Selection for Knowledge Distillation
Abhishek Panigrahi, Bing Liu, Sadhika Malladi, Sham Kakade, Surbhi Goel
arXiv 2025LOTION: Smoothing the Optimization Landscape for Quantized Training
Mujin Kwun, Depen Morwani, Chloe Huangyuan Su, Stephanie Gil, Nikhil Anand, Sham Kakade
arXiv 2025Random Scaling of Emergent Capabilities
Rosie Zhao, Tian Qin, David Alvarez-Melis, Sham Kakade, Naomi Saphra
arXiv 2025Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling
Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, Sham Kakade
arXiv 2025Selective Underfitting in Diffusion Models
Kiwhan Song, Jaeyeon Kim, Sitan Chen, Yilun Du, Sham Kakade, Vincent Sitzmann
arXiv 2025Soup to go: mitigating forgetting during continual learning with model averaging
Anat Kleiman, G. Dziugaite, Jonathan Frankle, Sham Kakade, M. Paul
arXiv 2025The Emergence of Complex Behavior in Large-Scale Ecological Environments
Joseph Bejjani, Chase Van Amburg, Chengrui Wang, Chloe Huangyuan Su, Sarah M. Pratt, Yasin Mazloumi, N. Khoshnevis, Sham Kakade, Kianté Brantley, Aaron Walsman
arXiv 2025The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
Natalie Abreu, Nikhil Vyas, Sham Kakade, Depen Morwani
arXiv 2025The Role of Sparsity for Length Generalization in Transformers
Noah Golowich, Samy Jelassi, David Brandfonbrener, Sham Kakade, Eran Malach
arXiv 2025
2024
LoRA Soups: Merging LoRAs for Practical Skill Composition Tasks
Akshara Prabhakar, Yuanzhi Li, Karthik Narasimhan, Sham Kakade, Eran Malach, Samy Jelassi
International Conference on Computational Linguistics 2024Q-Probe: A Lightweight Approach to Reward Maximization for Language Models
Kenneth Li, Samy Jelassi, Hugh Zhang, Sham Kakade, M. Wattenberg, David Brandfonbrener
International Conference on Machine Learning 2024Repeat After Me: Transformers are Better than State Space Models at Copying
Samy Jelassi, David Brandfonbrener, Sham Kakade, Eran Malach
International Conference on Machine Learning 2024A New Perspective on Shampoo's Preconditioner
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, Lucas Janson
International Conference on Learning Representations 2024CoLoR-Filter: Conditional Loss Reduction Filtering for Targeted Language Model Pre-training
David Brandfonbrener, Hanlin Zhang, Andreas Kirsch, J. Schwarz, Sham Kakade
Neural Information Processing Systems 2024DataComp-LM: In search of the next generation of training sets for language models
Jeffrey Li, Alex Fang, G. Smyrnis, Maor Ivgi, Matt Jordan, S. Gadre, Hritik Bansal, E. Guha, Sedrick Scott Keh, K. Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean-Pierre Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng-Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, G. Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, K. Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke S. Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander Toshev, Stephanie Wang, Dirk Groeneveld, Luca Soldani, Pang Wei Koh, J. Jitsev, Thomas Kollar, A. Dimakis, Y. Carmon, Achal Dave, Ludwig Schmidt, Vaishaal Shankar
Neural Information Processing Systems 2024Deconstructing What Makes a Good Optimizer for Language Models
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, Sham Kakade
International Conference on Learning Representations 2024Eliminating Position Bias of Language Models: A Mechanistic Approach
Ziqi Wang, Hanlin Zhang, Xiner Li, Kuan-Hao Huang, Chi Han, Shuiwang Ji, Sham Kakade, Hao Peng, Heng Ji
International Conference on Learning Representations 2024Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond
Costin-Andrei Oncescu, S. Purandare, Stratos Idreos, Sham Kakade
International Conference on Learning Representations 2024Follow My Instruction and Spill the Beans: Scalable Data Extraction from Retrieval-Augmented Generation Systems
Zhenting Qi, Hanlin Zhang, Eric P. Xing, Sham Kakade, Hima Lakkaraju
International Conference on Learning Representations 2024How Does Critical Batch Size Scale in Pre-training?
Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, Sham Kakade
International Conference on Learning Representations 2024Loss-to-Loss Prediction: Scaling Laws for All Datasets
David Brandfonbrener, Nikhil Anand, Nikhil Vyas, Eran Malach, Sham Kakade
Trans. Mach. Learn. Res. 2024Mixture of Parrots: Experts improve memorization more than reasoning
Samy Jelassi, Clara Mohri, David Brandfonbrener, A. Gu, Nikhil Vyas, Nikhil Anand, David Alvarez-Melis, Yuanzhi Li, Sham Kakade, Eran Malach
International Conference on Learning Representations 2024Transcendence: Generative Models Can Outperform The Experts That Train Them
Edwin Zhang, Vincent Zhu, Naomi Saphra, Anat Kleiman, Benjamin L. Edelman, Milind Tambe, Sham Kakade, Eran Malach
Neural Information Processing Systems 2024SOAP: Improving and Stabilizing Shampoo using Adam
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
arXiv 2024Universal Length Generalization with Turing Programs
Kaiying Hou, David Brandfonbrener, Sham Kakade, Samy Jelassi, Eran Malach
arXiv 2024
2023
A Study on the Calibration of In-context Learning
Hanlin Zhang, Yi-Fan Zhang, Yaodong Yu, Dhruv Madeka, Dean Foster, Eric P. Xing, Hima Lakkaraju, Sham Kakade
North American Chapter of the Association for Computational Linguistics 2023Beyond Implicit Bias: The Insignificance of SGD Noise in Online Learning
Nikhil Vyas, Depen Morwani, Rosie Zhao, Gal Kaplun, Sham Kakade, B. Barak
On Provable Copyright Protection for Generative Models
Nikhil Vyas, Sham Kakade, B. Barak
International Conference on Machine Learning 2023Feature emergence via margin maximization: case studies in algebraic tasks
Depen Morwani, Benjamin L. Edelman, Costin-Andrei Oncescu, Rosie Zhao, Sham Kakade
International Conference on Learning Representations 2023Pareto Frontiers in Deep Feature Learning: Data, Compute, Width, and Luck
Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, Cyril Zhang
Neural Information Processing Systems 2023Pareto Frontiers in Neural Feature Learning: Data, Compute, Width, and Luck
Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, Cyril Zhang
Provable Copyright Protection for Generative Models
Nikhil Vyas, Sham Kakade, Boaz Barak
arXiv 2023
2022
Capacity of Group-invariant Linear Readouts from Equivariant Representations
Matthew Farrell, Blake Bordelon, Shubhendu Trivedi, Cengiz Pehlevan
ICLR 2022Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit
Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, Cyril Zhang
NeurIPS 2022Inductive Biases and Variable Creation in Self-Attention Mechanisms
Benjamin Edelman, Surbhi Goel, Sham Kakade, Cyril Zhang
ICML 2022Learning Curves for SGD on Structured Features
Blake Bordelon, Cengiz Pehlevan
ICLR 2022Neural Networks as Kernel Learners: The Silent Alignment Effect
Alexander Atanasov, Blake Bordelon, Cengiz Pehlevan
ICLR 2022Contrasting random and learned features in deep Bayesian linear regression
Jacob A. Zavatone-Veth, William L. Tong, Cengiz Pehlevan
Manuscript 2022Deconstructing Distributions: A Pointwise Framework of Learning
Gal Kaplun, Nikhil Ghosh, Saurabh Garg, Boaz Barak, Preetum Nakkiran
Manuscript 2022
2021
Asymptotics of Representation Learning in Finite Bayesian Neural Networks
Jacob A. Zavatone-Veth, Abdulkadir Canatar, Benjamin S. Ruben, Cengiz Pehlevan
NeurIPS 2021Depth induces scale-averaging in overparameterized linear Bayesian neural networks
Jacob A. Zavatone-Veth, Cengiz Pehlevan
55th Asilomar Conference 2021Out-of-Distribution Generalization in Kernel Regression
Abdulkadir Canatar, Blake Bordelon, Cengiz Pehlevan
NeurIPS 2021Revisiting Model Stitching to Compare Neural Representations
Yamini Bansal, Preetum Nakkiran, Boaz Barak
NeurIPS 2021
2020
Deep Double Descent: Where Bigger Models and More Data Hurt
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, Ilya Sutskever
ICLR 2020For Self-supervised Learning, Rationality Implies Generalization, Provably
Yamini Bansal, Gal Kaplun, Boaz Barak
ICLR 2020Learning From Strategic Agents: Accuracy, Improvement, and Causality
Yonadav Shavit, Benjamin Edelman, Brian Axelrod
ICML 2020The Deep Bootstrap: Good Online Learners are Good Offline Generalizers
Preetum Nakkiran, Behnam Neyshabur, Hanie Sedghi
ICLR 2020Distributional Generalization: A New Kind of Generalization
Preetum Nakkiran, Yamini Bansal
Manuscript 2020
2019
Computational Limitations in Robust Classification and Win-Win Results
Akshay Degwekar, Preetum Nakkiran, Vinod Vaikuntanathan
COLT 2019SGD on Neural Networks Learns Functions of Increasing Complexity
Preetum Nakkiran, Gal Kaplun, Dimitris Kalimeris, Tristan Yang, Benjamin L. Edelman, Fred Zhang, Boaz Barak
Adversarial Robustness May Be at Odds With Simplicity
Preetum Nakkiran
Manuscript 2019Minnorm training: an algorithm for training over-parameterized deep neural networks
Yamini Bansal, Madhu Advani, David D Cox, Andrew M Saxe
Manuscript 2019More Data Can Hurt for Linear Regression: Sample-wise Double Descent
Preetum Nakkiran
Manuscript 2019
2018
On the Information Bottleneck Theory of Deep Learning
Andrew Michael Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Brendan Daniel Tracey, David Daniel Cox
ICLR 2018