Efficient LLM Training and Inference
This line of work asks how to get more capability per FLOP and per byte of memory. It spans compute-optimal vocabulary scaling, quantization and structured pruning of generative language models, KV cache compression for long-context inference, faster test-time scaling, distillation, and fault-tolerant distributed inference. The quantization paper received the ACL 2022 Outstanding Paper Award.
Key results
| Paper | Problem | Result |
|---|---|---|
| Scaling Laws with Vocabulary (NeurIPS 2024) | Compute-optimal vocabulary size | Llama2-70B should use ≥216K tokens (7× its 32K); 32K → 43K lifts ARC-Challenge 29.1 → 32.0 at 2.3e21 FLOPs |
| Compression of Generative PLMs via Quantization (ACL 2022, Outstanding Paper) | Low-bit generative models | 14.4× (GPT-2) and 13.4× (BART) compression, comparable to full precision |
| UNComp (EMNLP 2025) | KV cache compression | Cache reduced to 4.74% of original, 6.4× throughput |
| D2O | KV cache eviction without fine-tuning | More than 3× inference throughput |
| ATTS (ICLR 2026) | Test-time scaling latency | Up to 56.7× speedup in test-time scaling and 4.14× throughput, no accuracy loss |
| Prefilling-dLLM (EMNLP 2026) | Long-context diffusion LLMs | 9.1–28.0× speedup at 8K–32K context, training-free |
What vocabulary size should an LLM use?
Larger than most models use today. By training models from 33M to 3B parameters on up to 500B characters, Scaling Laws with Vocabulary derives the compute-optimal vocabulary size three ways (IsoFLOPs analysis, derivative estimation, and a parametric loss fit), and all three agree that the optimum grows with compute. Most open LLMs are under-provisioned: Llama2-70B should have used at least 216K tokens instead of 32K. Try the interactive demo or the code.
How do you quantize generative language models without breaking generation?
Fix the word embeddings first. Earlier quantization methods that worked for BERT fail on GPT-2 and BART because low-bit word embeddings become homogeneous and weight distributions vary across modules. Token-level contrastive distillation keeps embeddings distinguishable, and module-wise dynamic scaling adapts each quantizer, giving 14.4× and 13.4× compression with near full-precision quality. Follow-up work applies learnable masks for structured pruning of generative PLMs (Findings of ACL 2023).
How can long-context inference be made cheaper?
By deciding which parts of the KV cache matter. D2O allocates cache budgets per layer from attention density and recalls discarded tokens that are similar to kept ones. UNComp uses truncated matrix entropy as an uncertainty signal to find sparsity, which also exposes retrieval heads and layers. Prefilling-dLLM brings prefill/decode separation to diffusion language models by caching prefix chunks once and decoding over only the most relevant ones.
Can test-time scaling be fast?
Yes, if synchronization is removed. ATTS identifies synchronization, not compute, as the main bottleneck in speculative test-time scaling, and replaces it with asynchronous inference guided by conformal prediction. On MATH, AMC23, AIME24 and AIME25, it reaches up to 56.7× speedup while controlling the rejection rate.
What matters in LLM distillation?
With enough training, forward and reverse KL behave similarly, but early on they emphasize the head and tail of the distribution; Adaptive KL (COLING 2024) combines them. PerSyn (ACL 2026) routes each prompt to the teacher a given student can best learn from, and On-policy Verbal Distillation replaces token-level matching with trajectory-level teacher scores.
Efficient vision-language models and resilient serving
CrossGET (ICML 2024) and UPop (ICML 2023) accelerate and compress vision-language Transformers through token ensembling and progressive pruning. AnchorTP (DATE 2026) keeps LLM serving alive through single-GPU failures with state-preserving elastic tensor parallelism.
Papers
-
OVD: On-policy Verbal Distillation
Jing Xiong, Hui Shen, Shansan Gong, Yuxin Cheng, Jianghan Shen, Chaofan Tao, Haochen Tan, Haoli Bai, Lifeng Shang, Ngai Wong
On-policy Verbal Distillation replaces token-level matching with trajectory-level distillation using discrete teacher scores, reducing memory and improving student performance on web QA and math reasoning.
-
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.
-
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.
-
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.
FAQ
What is the compute-optimal vocabulary size for an LLM?
It grows with the compute budget, so larger models deserve larger vocabularies. Scaling Laws with Vocabulary (NeurIPS 2024) predicts that Llama2-70B should have used at least 216K tokens, 7 times its 32K vocabulary. At a fixed 2.3e21 FLOPs, increasing a 3B model's vocabulary from 32K to 43K raised ARC-Challenge accuracy from 29.1 to 32.0.
Which paper by Chaofan Tao received the ACL 2022 Outstanding Paper Award?
Compression of Generative Pre-trained Language Models via Quantization (ACL 2022). It shows that earlier quantization methods fail on generative models because of homogeneous word embeddings and varied weight distributions, and fixes this with token-level contrastive distillation and module-wise dynamic scaling, reaching 14.4x compression on GPT-2 and 13.4x on BART with performance comparable to full precision.
How much can KV cache compression save in long-context inference?
In UNComp (EMNLP 2025), uncertainty-guided compression reduces the KV cache to 4.74% of its original size with a 6.4x throughput gain. D2O, which allocates cache per layer and recalls discarded tokens by similarity, improves inference throughput by more than 3x without fine-tuning.
Is there code or a demo for the vocabulary scaling law?
Yes. Code is at github.com/sail-sg/scaling-with-vocab and an interactive demo that predicts the optimal vocabulary size is on Hugging Face Spaces at huggingface.co/spaces/sail/scaling-with-vocab-demo.
How to cite
@inproceedings{tao2024scaling,
title = {Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies},
author = {Tao, Chaofan and Liu, Qian and Dou, Longxu and Muennighoff, Niklas and Wan, Zhongwei and Luo, Ping and Lin, Min and Wong, Ngai},
booktitle = {Advances in Neural Information Processing Systems},
year = {2024}
}
@inproceedings{tao2022compression,
title = {Compression of Generative Pre-trained Language Models via Quantization},
author = {Tao, Chaofan and Hou, Lu and Zhang, Wei and Shang, Lifeng and Jiang, Xin and Liu, Qun and Luo, Ping and Wong, Ngai},
booktitle = {Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics},
year = {2022}
}Related: Coding agents and agentic RL · Multimodal reasoning and agents · llms.txt