Research
I study optimization, dynamics and learning, with a focus on modern machine learning. I have also worked at the intersection of systems and theory. Below are the themes my group works on now, with our papers on each. The full list is on the publications page.
Optimization for modern deep learning
All papers on this themeWhy optimizers like Adam work, and how to make training faster and more robust: batch sizes, sign and spectral methods, distributed and asynchronous training.
-
Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent
Adaptive batch-size rules rely on the gradient noise scale, which assumes the Euclidean geometry of SGD. We derive the noise scale for sign and spectral methods (Signum, Muon) from their dual norms, and estimate it cheaply from gradients already on each data-parallel worker. On a 160M-parameter Llama model, this matches the loss of fixed batch sizes with up to 66% fewer training steps.
- Pseudo-Asynchronous Local SGD: Robust and Efficient Data-Parallel Training
-
Understanding Adam Requires Better Rotation Dependent Assumptions
Theory often explains Adam with assumptions that do not change when the parameter space is rotated. We show that Adam on transformers gets worse under random rotations, while some structured rotations help, and that existing rotation- dependent assumptions do not explain this. The orthogonality of Adam's update emerges as a promising quantity for a better theory.
-
No Wrong Turns: The Simple Geometry Of Neural Networks Optimization Paths
Why do stochastic first-order methods train non-convex networks so reliably? We measure two quantities along the optimization path, tied to the restricted secant inequality and the error bound, and find that they behave predictably across vision and language tasks, architectures and optimizers. These properties are enough to guarantee linear convergence and to motivate learning-rate schedules like those used in practice.
- Feature learning as alignment: a structural property of gradient descent in non-linear neural networks
- Neural Networks Efficiently Learn Low-Dimensional Representations with SGD
Generalization, out-of-distribution robustness and compositionality
All papers on this themeModels that keep working when the data changes: domain shift, OOD detection, identifiable and compositional representations, and pretraining objectives.
- DisTaC: Conditioning Task Vectors via Distillation for Robust Model Merging
- A Derandomization Framework for Structure Discovery: Applications in Neural Networks and Beyond
-
Beyond Multi-Token Prediction: Pretraining LLMs with Future Summaries
Next-token prediction trains a model to look one step ahead, which limits long- horizon reasoning and planning. We add an auxiliary head that predicts a summary of the long-term future, either handcrafted (a bag of words) or learned with a reverse language model. In 3B and 8B pretraining runs, it beats next-token and multi-token prediction on math, reasoning and coding benchmarks.
- Generalization Measures under Controlled Covariate Shift: A Regime-Aware Benchmark
-
Compositional risk minimization
Classifiers fail when test data contain combinations of attributes never seen together in training. We model each attribute as an additive energy term, train a classifier with that structure, then adjust it for the new combinations. The method, compositional risk minimization, provably extrapolates beyond the training combinations and is more robust than methods built for subpopulation shift.
- An Empirical Study of Pre-trained Model Selection for Out-of-Distribution Generalization and Calibration
Privacy and machine unlearning
All papers on this themeRemoving the influence of data from trained models, with guarantees, and measuring whether unlearning methods actually work.
-
Unlearning with Asymmetric Sources: Improved Unlearning-Utility Trade-off with Public Data
Certified unlearning adds noise, and the noise costs utility, especially for large deletion requests. We show that mixing in public data cuts the cost of unlearning by a factor of O(1/n_pub²), and we account for shifts between the public and private data. This makes it practical to unlearn a constant fraction of the training set while keeping utility.
-
Are we making progress in unlearning? Findings from the first NeurIPS unlearning competition
We ran the first NeurIPS machine unlearning competition, with nearly 1,200 teams, scoring forgetting through a formal definition of unlearning together with model utility. The top entries beat existing algorithms, and their ranking stays stable across variations of the evaluation, so it can be made cheaper. We also analyze where the methods succeed and fail.
Dynamics of games and reinforcement learning
All papers on this themeMin-max optimization, variational inequalities and the learning dynamics of multi-agent and reinforcement learning.
- A Surrogate Perspective on Convergence of Fixed-Target DQN
-
Solving hidden monotone variational inequalities with surrogate losses
Many problems in deep learning, such as min-max games and projected Bellman error minimization, are variational inequalities, and standard gradient methods cycle or diverge on them. We solve them through a sequence of surrogate losses that work with standard optimizers such as Adam, with convergence guarantees under hidden monotonicity. The approach also yields a more compute- and sample-efficient variant of TD(0).
-
A Unified Approach to Reinforcement Learning, Quantal Response Equilibria, and Two-Player Zero-Sum Games
We introduce magnetic mirror descent, one algorithm that serves both as an equilibrium solver and as a reinforcement learning method for two-player zero-sum games. It is the first solver of quantal response equilibria with linear convergence in extensive-form games from first-order feedback, and the first standard RL algorithm competitive with counterfactual regret minimization in tabular settings. As a deep self-play method, it does well on Dark Hex and Phantom Tic-Tac-Toe.
-
LEAD: Min-Max Optimization from a Physical Perspective
Min-max training, as in GANs, stalls because of rotational dynamics. Treating the two players as particles under physical forces, we derive LEAD, a second-order method that damps the rotation. It converges linearly on quadratic games and improves GAN training on CIFAR-10.
- Stochastic Gradient Descent-Ascent and Consensus Optimization for Smooth Games: Convergence Analysis under Expected Co-coercivity
- A Tight and Unified Analysis of Gradient-Based Methods for a Whole Spectrum of Differentiable Games
Earlier work and resources
- Summary of research contributions from my group (PDF, not fully up to date).
- Slides summarizing our work on games and min-max optimization.
- Recent manuscripts on Google Scholar.
- Older project pages: FrogWild! (fast PageRank approximations, VLDB 2015) and Asynchrony begets momentum.
Funding acknowledgements
Special thanks to Intel and NVIDIA for donating access to hardware, and to SigOpt for access to their platform for some of our work.