Efficient ML & AI Infrastructure
In my experience, most efficiency problems start with an assumption that quietly stopped being true. This page follows a model from design to serving and looks at the assumption each piece of our work takes apart.
Efficiency is usually framed as a trade: give up a little quality to save compute or memory. The bigger wins I've seen come from somewhere else. A method fails, or wastes resources, because it rests on an assumption that made sense when it was designed and no longer holds. The tokenizer can be fixed before you scale. Every part of a network handles compression the same way. Every cached token is equally important. Waiting for synchronization is free. GPUs don't fail. The strongest teacher teaches best.
Spotting the broken assumption is usually the hard part. Once the diagnosis is right, the fix often turns out to be simple. I've organized this page by where in a model's life each assumption lives. The quantization work in the second section received an ACL 2022 Outstanding Paper Award.
On this page: Design · Compress · Remember · Compute and serve · Transfer · What ties it together · Open questions
1. Design: the tokenizer is not a free choice
Scaling laws tell you how loss changes as you add parameters and data. They usually say nothing about the vocabulary, and you can see the result in practice: models of similar size ship with wildly different vocabulary sizes. In Scaling Laws with Vocabulary (NeurIPS), we treated vocabulary size as a scaling dimension in its own right. We separated vocabulary parameters from the rest of the network, measured data in characters instead of tokens so that different tokenizers could be compared fairly, and used a loss that doesn't depend on vocabulary size. Three independent ways of estimating the best size, from IsoFLOPs analysis, from derivatives, and from fitting the loss directly, all pointed the same way: most models under-invest in their vocabulary.
Why is there a best size at all? It comes down to a trade-off. A bigger vocabulary splits text into fewer tokens and can represent more concepts directly, but when compute is limited, the embeddings for rare tokens never get trained well. As compute grows, the balance tips toward larger vocabularies, though they should grow more slowly than the rest of the model. How much data you have matters as well: with little data, a smaller vocabulary avoids overfitting rare tokens; with far more data than is compute-optimal, a larger one pays off, keeping in mind that a bigger vocabulary also makes inference more expensive. One picture I like: under a fixed compute budget, the best vocabulary is the one that lets the model read the most text. Go smaller and the same text costs more tokens; go bigger and the compute goes into the embedding and output layers instead. The interactive calculator turns all of this into a number you can use.
2. Compress: not every part of a model is alike
Quantization methods that worked well for BERT-style encoders fell apart on GPT-style generators, and for a while it wasn't clear why. In Compression of Generative Pre-trained Language Models via Quantization (ACL 2022 Outstanding Paper), we traced it to two problems. At low precision, the word embeddings collapse until they look almost the same, so a model that has to pick the next token can no longer tell its options apart. And the weights in different modules are spread out very differently, so no single quantizer setting suits all of them. The fixes follow from the diagnosis: token-level contrastive distillation keeps each token's representation close to the full-precision model's and distinct from other tokens, and module-wise dynamic scaling gives each module its own range. Structured pruning applies the same idea to removing parts of a generative model, learning which pieces can go.
Vision-language models make this unevenness even easier to see, because images and text carry information so differently. UPop searches for a smaller network across the vision and language parts together, so the model learns how much to prune from each instead of us guessing, and it alternates between searching and retraining. CrossGET merges redundant tokens while the model runs, letting one modality guide which tokens get matched in the other.
3. Remember: not every cached token matters, but keep a way back
In long-context inference, the KV cache, which holds keys and values for every earlier token, soon takes up most of the memory. Trimming it evenly is wasteful because layers behave differently. D2O gives each layer its own budget depending on how focused its attention is, and keeps a way to bring back or merge tokens it dropped if they resemble ones it kept, all without any fine-tuning. UNComp looks at the same question in terms of information: it uses matrix entropy as a measure of uncertainty to find the parts that can be compressed, and the same analysis points to the heads and layers that do the retrieving.
Diffusion language models have their own version of the problem, since they re-read the entire prompt at every denoising step. Prefilling-dLLM caches the prompt once, in chunks, and then only uses the chunks and tokens relevant to the question, with no retraining. One small detail turned out to matter a lot: putting a beginning-of-sequence token at the start of every chunk gives attention something stable to hold on to, and with that in place, the sparse version can even beat full attention.
4. Compute and serve: coordination costs more than arithmetic
At scale, the bottleneck is usually waiting, and the biggest risk is usually losing state. Test-time scaling is a good example of the first. It improves answers by spending more compute while generating, often with a small model drafting and a large one checking. When we looked at where the time actually went, in ATTS, the main cost was synchronization, not arithmetic. ATTS replaces those synchronous checks with asynchronous inference, using conformal prediction, calibrated on the fly, to keep statistical control over how often drafts get rejected. It scales in both directions at once: longer reasoning and more parallel attempts.
Serving has its own hidden assumption: that GPUs don't fail. With tensor parallelism across several GPUs, losing one can take the whole service down. AnchorTP keeps the weights and KV caches alive in a separate process that outlives the crash, re-splits the model across the GPUs that are left, and plans the smallest amount of data movement needed to get back up.
Infrastructure for training agents. Both lessons come back, more sharply, when you train agents with RL inside their real harness. In Lego-RL, almost all of a training attempt is the agent at work: reading files, editing code, running tests in a sandbox. Attempts are long and very uneven, so a synchronous batch always waits for its slowest member. Asynchronous scheduling lets a new attempt start as soon as a slot opens. When the weights update mid-attempt, a partial rollout pauses generation, loads the new weights and carries on from the cached prefix on the same machine, so a long attempt isn't thrown away. Starting a sandbox is mostly spent building its image, so prebuilt task images, a mounted agent runtime and lazy image pulls remove most of that wait. There's one place where we deliberately chose the slower option: the grading tools are baked into every image, because a reward you can reproduce is worth more than one you get a little faster. A lot of agent RL, it turns out, is an efficiency and reliability problem.
5. Transfer: the strongest teacher is not the best teacher
Distillation works best when the signal fits what the student can learn, not when the teacher is as strong as possible. Take the choice between forward and reverse KL divergence. The common story is that one chases the main mode of the teacher and the other spreads out to cover everything. For LLMs, we found that neither really happens: the two share the same optimum. The real difference shows up early in training, when forward KL pays attention to the teacher's most likely tokens and reverse KL to the long tail. Adaptive KL simply weights the two accordingly. PerSyn starts from the observation that stronger models are not always better teachers, and sends each prompt to the teacher a particular student learns from best before any data is generated, instead of having every teacher answer everything and filtering afterwards.
Agent training arrives at the same conclusion from a different direction. Grounded trajectories from weaker agents teach better than answers from stronger ones, and text feedback from a black-box teacher can guide a student without ever exposing the teacher's probabilities.
What ties it together
- Find the broken assumption before you optimize. Every result here started with asking why an existing method failed, not with trying a faster variant. Once the cause was clear, the fix was usually small.
- Unevenness is the rule. Layers, modules, modalities, tokens and teachers all differ. Giving each its own budget, per layer, module, modality or prompt, keeps paying off in compression, caching and distillation alike.
- Keep a way back. Dropped tokens that can be recalled, state that survives a crash, attempts that survive a weight update: these are what make aggressive optimization safe.
- At scale, waiting is the real cost. In test-time scaling, serving and agent training, the biggest gains came from removing waiting, not from faster arithmetic.
- Efficiency shapes what a model can learn. Rare-token embeddings that never get enough compute, teacher outputs a student can't absorb: efficiency isn't only about cost. It decides what a model is able to learn at all.
Open questions
- How should vocabulary be chosen once inference cost, multilingual data and code are all part of the goal?
- What should long-context memory keep for agents, whose context is mostly tool output rather than prose?
- How much faster can agent RL get, when most of the time goes into acting rather than computing gradients?
- Can the fault-tolerance ideas from serving make long training runs just as resilient?
Papers (12)
-
Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models
Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Boyuan Zheng, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong
Prefilling-dLLM separates prefilling from decoding without training: it caches prefix KV representations by chunk, selects query-relevant chunks and tokens, and reuses the sparse cache to accelerate long-context diffusion language model inference.
-
ATTS: Asynchronous Test-Time Scaling via Conformal Prediction
Jing Xiong, Qiujiang Chen, Fanghua Ye, Zhongwei Wan, Chuanyang Zheng, Chenyang Zhao, Hui Shen, Alexander Hanbo Li, Chaofan Tao, Haochen Tan, Haoli Bai, Lifeng Shang, Lingpeng Kong, Ngai Wong
ATTS is an asynchronous, conformal-prediction-based test-time scaling method that cuts synchronization overhead and accelerates scaling with statistical guarantees and no accuracy loss.
-
UNCOMP: Can Matrix Entropy Uncover Sparsity? - A Compressor Design from an Uncertainty-Aware Perspective
Jing Xiong, Jianghan Shen, Fanghua Ye, Chaofan Tao, Zhongwei Wan, Jianqiao Lu, Xun Wu, Chuanyang Zheng, Zhijiang Guo, Min Yang, Lingpeng Kong, Ngai Wong
UNComp uses uncertainty (via truncated matrix entropy) to uncover KV-cache sparsity and perform adaptive compression for long-context inference, greatly shrinking cache while boosting throughput.
-
Find Your Optimal Teacher: Personalized Data Synthesis via Router-Guided Multi-Teacher Distillation
Hengyuan Zhang, Shiping Yang, Xiao Liang, Chenming Shang, Yuxuan Jiang, Chaofan Tao, Jing Xiong, Hayden Kwok-Hay So, Ruobing Xie, Angel X Chang, Ngai Wong
PerSyn "routes then generates" by assigning each prompt to its best teacher via a router that considers student learnability and teacher quality, improving synthetic-data distillation efficiency.
-
AnchorTP: Resilient LLM Inference with State-Preserving Elastic Tensor Parallelism
Wendong Xu, Chujie Chen, He Xiao, Kuan Li, Jing Xiong, Chen Zhang, Wenyong Zhou, Chaofan Tao, Yang Bai, Bei Yu, Ngai Wong
AnchorTP provides state-preserving elastic tensor parallelism with a GPU-resident daemon and bandwidth-aware migration planning, enabling fast recovery from single-GPU failures.
-
Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies
Chaofan Tao, Qian Liu, Longxu Dou, Niklas Muennighoff, Zhongwei Wan, Ping Luo, Min Lin, Ngai Wong
The paper derives compute-optimal vocabulary scaling laws and shows larger models benefit from larger vocabularies, improving downstream accuracy at fixed FLOPs.
-
Dynamic Discriminative Operations for Efficient Long-Context Inference of Large Language Models
Zhongwei Wan, Xinjian Wu, Yu Zhang, Yi Xin, Chaofan Tao, Zhihong Zhu, Xin Wang, Siqi Luo, Jing Xiong, Longyue Wang, Mi Zhang
D2O dynamically compresses KV cache at layer and token levels with allocation and recall/merge compensation, saving memory and improving long-context throughput without fine-tuning.
-
CrossGET: Cross-Guided Ensemble of Tokens for Accelerating Vision-Language Transformers
Dachuan Shi, Chaofan Tao, Anyi Rao, Zhendong Yang, Chun Yuan, Jiaqi Wang
CrossGET accelerates vision-language Transformers by cross-modally matching and ensembling tokens with efficient soft matching, reducing compute with minimal accuracy loss.
-
Rethinking Kullback-Leibler Divergence in Knowledge Distillation for Large Language Models
Taiqiang Wu, Chaofan Tao, Jiahao Wang, Runming Yang, Zhe Zhao, Ngai Wong
It shows forward and reverse KL behave similarly in LLM distillation with enough training but emphasize head vs. tail early; Adaptive KL combines them to improve generation quality and diversity.
-
Structured Pruning for Efficient Generative Pre-trained Language Models
Chaofan Tao, Lu Hou, Haoli Bai, Jiansheng Wei, Xin Jiang, Qun Liu, Ping Luo, Ngai Wong
It proposes structured pruning across multiple generative-PLM components using learnable masks and sparse training to extract size-flexible subnets with significant speedups.
-
Compression of Generative Pre-trained Language Models via Quantization
Chaofan Tao, Lu Hou, Wei Zhang, Lifeng Shang, Xin Jiang, Qun Liu, Ping Luo, Ngai Wong
It adapts quantization to generative PLMs （e.g. GPT-2, BART）, addressing embedding homogenization with token-level contrastive distillation and module-wise dynamic scaling to achieve 13–14× compression at near full precision.
-
UPop: Unified and Progressive Pruning for Compressing Vision-Language Transformers
Dachuan Shi, Chaofan Tao, Ying Jin, Zhendong Yang, Chun Yuan, Jiaqi Wang
UPop is a unified progressive pruning framework that searches multimodal subnets in a continuous space and progressively retrains to achieve high compression for vision-language Transformers.