Organizations: Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) · King Abdullah University of Science and Technology (KAUST) · Advanced Micro Devices, Inc. (AMD) · Carnegie Mellon University (CMU)
Quantization-aware training (QAT) is the most reliable way to obtain accurate low-bit large language models (LLMs), but training with AdamW keeps two dense FP32 states for every weight matrix. We find that low-rank memory-efficient optimizers such as GaLore and LDAdam degrade substantially in QAT. A considerable part of the QAT momentum lies outside its top-r singular subspace, and projection-based methods discard this part at every step. We further observe that this residual remains useful when only its sign is kept. Based on these observations, we propose ResQAdam. It keeps a single BF16 momentum, applies Adam only to the r x r coordinates of the principal subspace, and adds a normalized sign update computed from the residual. We evaluate ResQAdam with two from-scratch QAT methods and three QAT fine-tuning pipelines, with weights quantized down to 2 bits. It reduces optimizer-state memory by 71% to 74% relative to AdamW, stays close to AdamW in most settings, and lowers training throughput by less than 2%.
Figures & tables
Figure 1: WikiText-2 perplexity versus optimizer state memory for different optimizers under two QAT paradigms. The x-axis reports the total optimizer state size of the full model (in GB). All experiments use W2A16 quantization with group size 128, and projection-based optimizers use rank 64. For a fair comparison, the learning rate of each optimizer is tuned via grid search and the best result is reported.
Figure 2: Training loss trajectories of ResQAdam and AdamW during from-scratch QAT under various quantization configurations on the 430M-parameter LLM. The values following the optimizer names in the legend indicate the average memory footprint of the optimizer states.
Methods
Size
Q-Cfg
WikiText-2 PPL (AdamW / Ours)
C4 PPL (AdamW / Ours)
Avg. PPL (AdamW / Ours)
Opt. State Size (AdamW / Ours)
QuEST
200M
W2A2
50.56 / 50.98
26.99 / 27.41
38.78 / 39.20
1.63 / 0.44 GB ↓ 73.0%
W2A16
44.04 / 42.80
22.92 / 23.41
33.48 / 33.11
W4A4
40.87 / 41.93
22.06 / 22.46
31.46 / 32.20
CAGE
W2A2
51.15 / 52.32
27.19 / 27.65
39.17 / 39.99
W2A16
44.80 / 44.44
23.18 / 23.73
33.99 / 34.09
W4A4
41.22 / 42.68
22.21 / 22.61
31.71 / 32.66
Table 1: Language modeling performance comparison across different from-scratch QAT methods. Values in each cell are formatted as AdamW / ResQAdam (Ours). Q-Cfg denotes the quantizer configuration, where W x and A x indicate x -bit weight and activation quantization, respectively. Opt. State Size reports average optimizer state size for each optimizer during the training, computed only over non-embedding and non-LM-head parameters.
Method
Model
Q-Cfg
Avg. Task Acc. (AdamW / Ours)
Avg. PPL (AdamW / Ours)
Opt. State Size (AdamW / Ours)
Basic QAT
Llama 3.2-1B
W2A16G128
47.56 / 47.36
16.92 / 17.73
7.78 / 2.16 GB ↓72.2%
W4A4G128
51.29 / 51.29
15.12 / 14.91
LLM-QAT
W2A16C
38.26 / 39.37
43.73 / 35.05
7.78 / 2.16 GB ↓72.2%
W4A16C
51.16 / 52.00
14.03 / 14.30
W4A4C
38.70 / 38.53
40.91 / 39.21
EfficientQAT (Block-AP)
W2A16G128
35.14 / 35.43
351.78 / 379.52
494 / 141 MB/block ↓71.5%
Table 2: Performance and memory footprint comparison between ResQAdam and standard AdamW across diverse QAT fine-tuning frameworks. Values in each cell are formatted as AdamW / ResQAdam (Ours). Q-Cfg denotes the quantizer configuration: W x and A x indicate x -bit weight and activation quantization, respectively; G128 represents group-wise quantization with a group size of 128; and C denotes channel-wise quantization. Opt. State Size reports average optimizer state size of AdamW / ResQAdam during the training , computed only over non-embedding and non-LM-head parameters.
Figure 3: Component ablation of ResQAdam across different quantization configurations. The y-axis reports the average perplexity on C4 and WikiText-2.
AdamW / ResQAdam
Llama 3.2-1B
Qwen2.5-1.5B
Qwen3 4B-Base
Opt. State Memory (GB)
7.78 / 2.16 ( ↓ 72.2%)
10.48 / 3.01( ↓ 71.3%)
29.07 / 7.78 ( ↓ 73.2%)
Total Peak Memory (GB)
33.33 / 28.24 ( ↓ 15.3%)
43.93 / 37.19 ( ↓ 15.3%)
99.14 / 79.77 ( ↓ 19.5%)
Throughput (# token/s)
29.36k / 28.94k ( ↓ 1.4%)
18.31k / 18.19k ( ↓ 0.6%)
9.64k / 9.56k ( ↓ 0.8%)
Table 3: Comparison of optimizer-state memory, total peak memory, and training throughput across different models. Values in each cell are formatted as AdamW / ResQAdam , with relative changes in parentheses.
Training Paradigm
Q-Cfg
ResQAdam
GaLore-AdamW
GaLore-AdamW8bit
LDAdam
AdaFactor
From-Scratch QAT
W2A16G128
20.05
30.61
30.39
28.57
33.01
W4A4G128
18.89
32.93
32.45
28.23
32.84
QAT Fine-Tuning
W2A16G128
16.89
19.25
20.32
19.59
27.79
W4A4G128
15.05
15.30
15.64
15.42
14.80
Table 4: Perplexity comparison of memory-efficient optimizers under diverse QAT paradigms and quantization configurations. Projection-based optimizers use rank 64. For a fair comparison, all optimizers are tuned with a grid search over learning rates, and the best result is reported.
Figure 4: Sensitivity of ResQAdam to the rank r and the SVD refresh interval τ on a 297M model trained from scratch. The y-axis is the average of WikiText-2 and C4 perplexity.
Figure 5: Energy ratio of the principal subspace ( ρt∥ ) during from-scratch QAT on a 297M Llama-style model. The curves represent different weight matrices within a middle transformer layer across W2A16, W4A4, and W2A2 quantization configurations.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Training loss trajectories of ResQAdam and AdamW during from-scratch QAT under various quantization configurations on the 200M-parameter LLM. The values following the optimizer names in the legend indicate the average memory footprint of the optimizer states.
School of Intelligent Systems Engineering Shenzhen Campus of Sun Yat-sen University Shenzhen, Guangdong 518107, China · Department of Frontier Research Peng Cheng Laboratory Shenzhen, Guangdong 518055, China