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 ways of asking the same question, and one habit worth keeping: test a result until it either survives or breaks.
The question is older than it looks. The cross-entropy that trains a language model is a code length, so a model that predicts well is a model that compresses well, and KL divergence counts the bits a wrong model wastes. Read that way, learning is compression, and efficiency has a unit: bits.
Three probes, each into a different part of the model, each measuring how much information it actually uses:
- Compressing context. How few vectors can carry a 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. How sure is a model of what it says? Giving its 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. Where in the stack does the work happen? 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
Every probe returned the same answer: the work concentrates. A few vectors can carry a context, a few layers do most of the computing, and the model itself can say how sure it is. The long-context sparsity work that followed, SpotAttention, Attention-Span, and Content-Aware Sparsity, is that finding taken seriously.