Foundations of Efficient LLMs
Status: Completed
Two years spent on a single question: how much of a large model's work is actually necessary? Answering it meant building the machinery first: running language models from one to eight billion parameters on TPU pods, training the smaller ones from scratch, and writing the GSPMD sharding and the static-graph, recompilation-free generation needed to run any of them. What followed were three lines of research, and one habit worth keeping: test a result until it either survives or breaks.
Three lines of research, each reaching into a different part of the model:
- Compressing context. Folding a long context, or a user’s history, into a handful of learned vectors that steer a frozen 8B model, treated as a small Gaussian latent so the steering signal could be sampled rather than fixed.
- Quantifying uncertainty. Giving a model’s token embeddings a learned variance, trained with a Bayesian optimizer (IVON), so a few sampled forward passes turn a single answer into a calibrated confidence.
- Tapering compute. Narrowing a Transformer’s attention and feed-forward width as it deepens, after finding that most of a model’s real work sits in its first and last few layers. Trained from scratch on a TPU v4 pod, it matched a full-width baseline’s perplexity while cutting latency 8.6% and lifting throughput 9.4%.
Several parts of this work were contributed upstream:
- PyTorch/XLA: two changes merged upstream,
booldtype support forfori_loop’srandint(#7632) and rank-0-only deletion in the distributed checkpoint manager (#7296) - Hugging Face Transformers: opened the request for XLA
StaticCachesupport (issue #31126), and the proposedindex_copy_approach (#31118) shipped in #31857