Compute-Optimal Quantization-Aware Training

ICLR 2026 Conference Submission22492 Authors

Keywords: quantization-aware training, QAT, neural network quantization, compute optimization, scaling laws, large language models, LLMs, model compression, compute budget allocation, training efficiency, model optimization, quantized neural networks, efficient deep learning

TL;DR: The optimal fraction of quantization-aware training compute (vs. pretrain stage) increases with total compute budget. We derive scaling laws to predict optimal allocation and model loss, enabling higher-quality model training with the same compute.

Abstract: Quantization-aware training (QAT) is a leading technique for improving the accuracy of quantized neural networks. Previous work has shown that decomposing training into a full-precision (FP) phase followed by a QAT phase yields superior accuracy compared to QAT alone. However, the optimal allocation of compute between the FP and QAT phases remains unclear. We conduct extensive experiments with various compute budgets, QAT bit widths, and model sizes from 86.0M to 2.2B to investigate how different QAT durations impact final performance. We demonstrate that, contrary to previous findings, the loss-optimal ratio of QAT to FP training increases with the total amount of compute. Moreover, the optimal fraction can be accurately predicted for a wide range of model sizes and quantization widths using the tokens-per-parameter-byte statistic. From experimental data, we derive a loss scaling law that predicts both optimal QAT ratios and final model performance across different QAT/FP compute allocation strategies and QAT bit widths. We use the scaling law to make further predictions, which we verify experimentally, including which QAT bit width is optimal under a given memory constraint and how QAT accuracy with different bit widths compares to full-precision model accuracy. Additionally, we propose a novel cooldown and QAT fusion approach that performs learning rate decay jointly with quantization-aware training, eliminating redundant full-precision model updates and achieving significant compute savings. These findings provide practical insights into efficient QAT planning and enable the training of higher-quality quantized models with the same compute budget.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 22492