LLM Compression with Jointly Optimizing Architectural and Quantization choices
This work addresses the challenge of deploying large language models on edge devices by providing a more efficient compression method that jointly optimizes architecture and quantization.
The paper introduces a differentiable NAS framework that jointly optimizes architectural configurations and mixed-precision quantization for LLMs, achieving up to 1.4x faster inference or 6% higher accuracy on reasoning tasks compared to sequential baselines.
Deploying large language models (LLMs) is challenging due to their significant memory and computational requirements. While some methods address this by developing small or tiny language models from scratch, these approaches demand extensive GPU training. Compressing pre-trained LLMs for edge devices offers a compelling alternative. Beyond pruning and quantization, Neural Architecture Search (NAS) enables effective compression, yet prior NAS approaches often limit the search space and decouple architecture from quantization. We introduce a differentiable NAS framework that explores the entire space and jointly optimizes architectural configurations alongside mixed-precision quantization for linear layers of LLMs. Experiments demonstrate superior accuracy-latency trade-offs: our models achieve up to 1.4x faster inference than sequential NAS-then-quantization baselines at comparable accuracy, or up to 6% higher average accuracy across seven reasoning tasks at equivalent latency.