I am a research scientist at Google working on distributed learning methods, especially for large-scale models. I received a Ph.D. in applied mathematics from the University of Wisconsin-Madison, and did a postdoc with Dimitris Papailiopoulos.

What do I work on?

I currently work on communication-efficient LLM training methods like DiLoCo. Earlier this year we introduced scaling laws for DiLoCo, which validated DiLoCo’s scaling behavior, showed a myriad of other benefits incurred by using the method, and gave concrete hyperparameter scaling laws and best practices to make it easy to use at scale. Check out a video walkthrough of the paper. We also put out work on Streaming DiLoCo, which greatly reduces peak bandwidth of LLM training.

I spend a lot of time programming in JAX. I helped develop DrJAX, a library for building MapReduce-style algorithms in JAX, which we use to speed up DiLoCo training. I also developed Dataset Grouper, a library for creating fast group-structured dataset pipelines, which can greatly accelerate research in so-called “heterogeneous” distributed settings.

Previously, I worked on federated learning. Among other things, I co-developed the popular FedAdam method, established best practices for scaling federated learning to large cohort settings, and helped write an (at the time) authoritative field guide on federated optimization methods. In 2023, I led work on applying federated learning algorithms to foundation models, with a focus on accelerating requisite dataset pipelines.

I have also dabbled in privacy-oriented work. I am particularly interested in the scalability of formal privacy methods (ie. differential privacy). I spearheaded work on applying user-level privacy to LLMs and more recently was involved in a project that developed scaling laws for differentially private language model training.

What else?

I am an erstwhile mathematician. My first paper in graduate school developed algorithms for efficiently generating random factored ideals in number fields. If this means anything to you, please reach out. Thinking in this way eventually led me to machine learning and optimization, and so here I am. In my spare time, I foster animals, bake, and participate in a Flyball team.