Foundations of Efficient LLMs
Status: Completed
For two years, I explored how much of a large model’s work is actually necessary. I ran models from one to eight billion parameters on TPU pods, trained the smaller ones from scratch, and wrote GSPMD sharding code and static-graph generation without recompilation. That infrastructure supported experiments in context compression, uncertainty estimation, and width tapering. One habit carried through all three: test a result until it either survives or breaks.
I approached these experiments through the connection between prediction and compression. Cross-entropy measures predictive coding cost; KL divergence measures the extra cost of using a mismatched distribution. That perspective helped frame three questions about what a model needs to retain, how certain its predictions are, and where it needs to spend compute.
- Compressing context. Long contexts and user histories were folded into a handful of learned vectors to steer a frozen 8B model. A small Gaussian latent representation allowed the steering signal to be sampled rather than fixed.
- Quantifying uncertainty. This experiment gave token embeddings a learned variance and trained them with the Bayesian optimizer IVON. A few sampled forward passes then produced a calibrated confidence estimate for an answer.
- Tapering compute. The models studied here concentrated much of their useful work in the first and last few layers. A Transformer with attention and feed-forward width tapered with depth was trained from scratch on a TPU v4 pod. It reached comparable perplexity to a full-width baseline at the same training budget, with 8.6% lower latency and 9.4% higher throughput.
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
The long-context work that followed, SpotAttention and Declarative Attention, carried these questions into attention.