Quantization-Aware Pre-Training with Constrained Empirical Weight Distribution

Ningfeng Yang; Tor M. Aamodt

cs.LGSep 26, 2026

Quantization-Aware Pre-Training with Constrained Empirical Weight Distribution

Authors: Ningfeng Yang, Tor M. Aamodt

Organizations: University of British Columbia

Abstract

Quantization-Aware Pre-Training (QAPT) can increase the inference efficiency of DNNs, but a problematic behaviour known as rounding boundary weight oscillation can introduce detrimental noise into the training process and significantly reduce convergence speed. While existing methods can reduce this detrimental noise, they either introduce additional hyperparameters or memory overhead, or cannot consistently improve model accuracy. In this work, we propose optimization with $\textbf{C}$ onstrained $\textbf{E}$ mpirical $\textbf{W}$ eight dis $\textbf{T}$ ribution (CEWT), the first hyperparameter-free memory-overhead-free oscillation suppression method that consistently improves QAPT performance: an optimizer post-update step that projects weights to the nearest point in weight space whose empirical distribution (histogram) matches a zero-mean Gaussian. Our key insight is many quantizers are designed with the implicit assumption that the to-be-quantized data are permutations of samples from a zero-mean Gaussian, and this assumption is not true during QAPT. By enforcing the zero-mean Gaussian prior as a hard constraint, CEWT can suppress this detrimental noise. Empirical results on various combinations of SOTA quantizers and hypersphere optimizers suggest, that with a geomean increase of 4% in training time, CEWT can consistently reduce the pre-training perplexity (by an average of 2.5 and up to 21 points) of low-precision (down to 1-bit activations and weights and up to 610M parameters) LLaMA/GPT models without introducing any hyperparameters or storage overhead. Code is available at https://github.com/1733116199/cewt

Figures & tables

Panel of Figure 1 — Figure 1: Effects of CEWT on pre-round (BBQ) and pre-clip (QuEST/LSQ) weight distributions of a LLaMA-95M with 2-bit weights and activations pre-trained on 3 billion C4 tokens. See Algorithms 8 , 9 , and 10 from definitions of pre-round and pre-clip weights.

Panel of Figure 2 — Figure 2: Visualization of CEWT (Equation 4 ). Notations like vec $(\cdot)$ and $\cdot^{(t)}$ are omitted. The example $\bar{W}$ we use here is the weights of the STE experiment in Table 2 at iteration 60. The $s$ in sub-figures 2(a) and 2(d) is defined in Section D .

Per-Tensor	Per-Channel	TGCS (Ours)
$s_{q}=\text{quant\_scale}(\text{vec}(W))$	$s_{qi}=\text{quant\_scale}(W_{i})$	$s_{q}=\text{quant\_scale}(\text{vec}(W))$
$s_{d}=\text{dequant\_scale}(\text{vec}(W))$	$s_{di}=\text{dequant\_scale}(W_{i})$	$s_{di}=\text{dequant\_scale}(W_{i})$
$Q={\color[rgb]{0.5,0.25,0}s_{d}}\cdot\text{clipround}(1/{\color[rgb]{0.5,0.25,0}s_{q}}\cdot W)$	$\forall i,Q_{i}={\color[rgb]{0,0,0.5}s_{di}}\cdot\text{clipround}(1/{\color[rgb]{0,0,0.5}s_{qi}}\cdot W_{i})$	$\forall i,Q_{i}={\color[rgb]{0,0,0.5}s_{di}}\cdot\text{clipround}(1/{\color[rgb]{0.5,0.25,0}s_{q}}\cdot W_{i})$

Table 1: Comparison between per-tensor, per-channel, and TGCS quantization. Example definitions of quant_scale, dequant_scale, and clipround are found in Algorithms 8 , 9 , and 10 .

		STE	PQN	CEWT	DL	IF	CGA	PI
MSE $(XQ,XW^{*})$ (unit: 1e-3)	Good HP	81±0	87±0	45±0	30±0	59±0	53±0	54±0
MSE $(XQ,XW^{*})$ (unit: 1e-3)	Bad HP	81±0	87±0	45±0	375±1	2080±12	332±1	1847±10
MSE $(XW,XW^{*})$ (unit: 1e-3)	Good HP	67±0	4±0	36±0	34±0	60±0	55±0	48±0
MSE $(XW,XW^{*})$ (unit: 1e-3)	Bad HP	67±0	4±0	36±0	295±1	2080±12	386±1	1807±9
EMA Oscillation Freq (unit: 1e-4)	Good HP	701±2	9±0	17±0	5±0	20±0	264±5	0±0
EMA Oscillation Freq (unit: 1e-4)	Bad HP	701±2	9±0	17±0	0±0	0±0	1±0	0±0

Table 2: Comparison between weight oscillation suppression methods. We report mean plus minus 2 standard deviations over 5 random seeds. Section C defines EMA Oscillation Frequency and RBM.

Model	Q-B	MuonH		Muon-CEWT		AdamH		Adam-CEWT		LionH		Lion-CEWT		SSO		SSO-CEWT
Model	Q-B	PPL	RBM	PPL	RBM	PPL	RBM	PPL	RBM	PPL	RBM	PPL	RBM	PPL	RBM	PPL	RBM
L-95	L-1	52.02	0.047	50.63	0.006	51.51	0.054	50.27	0.006	64.20	0.033	54.02	0.006	51.23	0.034	52.79	0.007
L-95	Q-1	51.97	0.215	48.66	0.119	52.85	0.240	48.02	0.119	51.33	0.244	48.23	0.119	53.61	0.182	49.47	0.119
L-95	B-1	50.89	0.031	47.48	0.010	51.06	0.034	46.12	0.010	50.70	0.034	46.67	0.010	52.15	0.023	49.05	0.010
L-95	L-2	33.73	0.030	33.42	0.009	34.64	0.032	33.45	0.009	41.21	0.035	34.91	0.009	34.72	0.024	34.22	0.009
L-95	Q-2	31.71	0.081	30.66	0.056	32.30	0.095	30.67	0.056	31.99	0.095	30.68	0.056	32.52	0.070	31.22	0.056

Table 3: Evaluation Perplexity of LLaMA and GPT models pre-trained on C4 with hypersphere optimizers, with and without CEWT. The first two columns are “Model” (L for LLaMA, G for GPT, with a dash followed by the number of parameters in millions) and “Q-B” which stands for Quantizer-Bits (L for LSQ, Q for QuEST, B for BBQ, with a dash followed by the activation and weight precision in number of bits). PPL stands for perplexity, and RBM stands for rounding boundary mass (Equation 15 ) of pre-round weights (Algorithms 8 , 9 , and 10 ).

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Visualization of $L(w)=\frac{1}{2}(\text{round}(w)-0.7)^{2}$ as a function and its negative STE-estimated gradient $0.7-\text{round}(w)$ as a one-dimensional vector field. The direction of an arrow represents its sign (left for negative and right for positive), and the size of an arrow represents its magnitude.

Figure 4: Visualization of $L(w)=\frac{1}{2}(q_{1}-0.7)^{2}+\frac{1}{2}(4q_{1}-q_{2}-2)^{2}$ as a function and its negative STE-estimated gradient $\left[-(q_{1}-0.7)-4\cdot(4q_{1}+q_{2}-2),-(4q_{1}+q_{2}-2)\right]^{\top}$ as a two-dimensional vector field. Note that $q_{1}=\text{round}(w_{1})$ and $q_{2}=\text{round}(w_{2})$ . Darker colored regions have lower loss than brighter regions.

Name	Lecun	PyTorch	Spectral $\mu P$	GPT-2 (non-residual)	GPT-2 (residual)
Init Distribution	Normal	Uniform	Normal	Normal	Normal
Init Mean	$0$	$0$	$0$	$0$	$0$
Init Std	$\sqrt{\frac{1}{n}}$	$\sqrt{\frac{1}{3n}}$	$\frac{\sqrt{m/n}}{\sqrt{m}+\sqrt{n}}$	$0.02$	$\frac{0.02}{\sqrt{2l}}$
Frobenius norm $R$	$\sqrt{m}$	$\sqrt{m/3}$	$\frac{m}{\sqrt{m}+\sqrt{n}}$	$0.02\sqrt{mn}$	$0.02\sqrt{\frac{mn}{2l}}$
AdamH Perp	31.78	31.52	31.39	30.91
Adam-CEWT Perp	30.12	30.22	29.82	29.77

Appendix

Table 4: Robustness of CEWT over different values of

R

m

is the number of output features of a linear layer.

n

is the number of input features of a linear layer.

l

is the number of layers in a model. Experiments are conducted using a LLaMA-95M with weights and activations quantized to 2-bit using BBQ. The model is pre-trained on 3 billion C4 tokens. We report perplexity (Perp).

	Weight Init HP (R)	Weight Decay HP	RB Weight Oscillation HP	# HPs
Adam	Yes	Yes	No	2
AdamH	Yes	No	No	1
Adam + Periodic Interpolation	Yes	Yes	Yes	3
AdamH + Periodic Interpolation	Yes	No	Yes	2
Adam + Dampening Loss	Yes	Yes	Yes	3
AdamH + Dampening Loss	Yes	No	Yes	2

Appendix

Table 5: The number of hyperparameters among weight initialization radius

R

, weight decay, and rounding boundary weight oscillation, for multiple combinations of regular/hypersphere optimizers and oscillation suppression methods. A “yes” means the hyperparameter is used. A “no” means the hyperparameter is not used.

Panel of Figure 5 — Figure 5: Rounding Boundary Mass (RBM), EMA Oscillation Frequency (EMAOF), learning rate (LR), and perplexity of a LLaMA-95M model with 2-bit weights and activations, pretrained on 3B C4 tokens. EMAOF is a more direct indicator of rounding boundary weight oscillation, but it correlates with LR and is not scale-invariant. RBM is scale-invariant, and RBM shows CEWT regulates the number of parameters close to rounding boundary by design .

		STE	HTGE	JacQuant-Probe	EWGS
MSE $(XQ,XW^{*})$ (unit: 1e-3)	Default	81±0	53±0	63±0	61±0
MSE $(XQ,XW^{*})$ (unit: 1e-3)	With CEWT	45±0	46±0	52±0	51±0
RBM $(\text{vec}(1/s\cdot W))$ (unit: 1e-3)	Default	573±1	177±1	349±1	324±1
RBM $(\text{vec}(1/s\cdot W))$ (unit: 1e-3)	With CEWT	10±0	10±0	10±0	10±0
EMA Oscillation Freq (unit: 1e-4)	Default	701±2	245±1	511±2	392±1
EMA Oscillation Freq (unit: 1e-4)	With CEWT	17±0	30±0	37±0	13±0

Appendix

Table 6: Effectiveness of CEWT on different rounding Jacobian estimators.

	$n=$	1	2	4	8	16	32	64	128	256	512	1024
MSE $(XQ,XW^{*})$ (unit: 1e-3)	STE	1±0	1±0	1±0	2±0	3±0	4±0	7±0	13±0	23±0	44±0	81±0
MSE $(XQ,XW^{*})$ (unit: 1e-3)	CEWT	1±0	1±0	1±0	2±0	2±0	3±0	5±0	9±0	14±0	25±0	45±0
RBM $(\text{vec}(1/s\cdot W))$ (unit: 1e-3)	STE	11±0	12±0	15±0	19±0	27±0	40±0	64±0	108±1	186±0	332±1	573±1
RBM $(\text{vec}(1/s\cdot W))$ (unit: 1e-3)	CEWT	10±0	10±0	10±0	10±0	10±0	10±0	10±0	10±0	10±0	10±0	10±0

Appendix

Table 7: Rounding boundary mass and MSE loss for different values of

n

(dataset size).

	Iterations =	1	3	6	12	25	50	100	200
MSE $(XQ,XW^{*})$ (unit: 1e-3)	STE	1847±10	1609±9	1277±7	747±5	244±1	90±0	82±0	81±0
MSE $(XQ,XW^{*})$ (unit: 1e-3)	CEWT	1839±9	1625±9	1343±7	859±6	281±1	107±0	63±0	45±0
RBM $(\text{vec}(1/s\cdot W))$ (unit: 1e-3)	STE	11±0	12±0	15±0	19±0	27±0	40±0	64±0	108±1
RBM $(\text{vec}(1/s\cdot W))$ (unit: 1e-3)	CEWT	10±0	10±0	10±0	10±0	10±0	10±0	10±0	10±0

Appendix

Table 8: Rounding boundary mass and MSE loss for different number of training iterations.

	$d=$	2	4	8	16	32	64	128	256	512	1024
MSE $(XQ,XW^{*})$ (unit: 1e-3)	STE	398±452	258±352	130±47	137±38	128±37	128±6	122±2	117±1	105±1	81±0
MSE $(XQ,XW^{*})$ (unit: 1e-3)	CEWT	1849±2115	1876±1619	313±281	117±41	95±8	90±4	85±1	79±1	66±0	45±0
RBM $(\text{vec}(1/s\cdot W))$ (unit: 1e-3)	STE	600±671	775±244	756±114	753±52	741±40	733±21	712±5	693±2	661±1	573±1
RBM $(\text{vec}(1/s\cdot W))$ (unit: 1e-3)	CEWT	0±0	0±0	13±34	11±9	11±4	10±1	10±0	10±0	10±0	10±0

Appendix

Table 9: Rounding boundary mass and MSE loss for different values of

d

(model size).

Row	Weight Quantizer			Optimizer	Perplexity	RBM	Notes
Row	Type	HT	TGCS	Optimizer	Perplexity	RBM	Notes
1	QuEST	T	F	AdamH	32.42	0.079	Vanilla QuEST without CEWT
2	QuEST	F	F	Adam-CEWT	31.53	0.071	Introduce CEWT without TGCS and remove HT
3	QuEST	F	T	Adam-CEWT	30.70	0.056	Further Introduce TGCS (Intended usage of this work)
4	QuEST	T	T	Adam-CEWT-HT	30.81	0.056	Introduce HT to both weight quantizer and optimizer
5	BBQ	T	F	AdamH	31.61	0.029	Vanilla BBQ without CEWT

Appendix

Table 10: Evaluation Perplexity of HO, CEWT w.o TGCS, CEWT, and CEWT-HT.

Panel of Figure 6 — Figure 6: Pre-clip (QuEST) and pre-round (BBQ) weight distribution of models (rows) in Table 10 .

Figure 7: Comparing STE, CEWT, and Periodic Interpolation with a period of 1 and with grid search of the interpolation ratio hyperparameter $\alpha$ . Experiments are conducted on a LLaMA-95M with activations and weights quantized to 2-bit (with QuEST) and pre-trained on 3 billion C4 tokens with the AdamH/Adam-CEWT optimizer.

Figure 8: Comparing STE, CEWT, and Dampening Loss with a grid search of the dampening hyperparameter $\alpha$ . Experiments are conducted on a LLaMA-95M with activations and weights quantized to 2-bit (with QuEST) and pre-trained on 3 billion C4 tokens with the AdamH/Adam-CEWT optimizer.

Quantizer	Bits W/A	AdamH + STE		AdamH + PQN		Adam-CEWT
Quantizer	Bits W/A	PPL	RBM	PPL	RBM	PPL	RBM
QuEST	1	52.85	0.240	50.61	0.002	48.02	0.119
BBQ	1	51.06	0.034	49.54	0.001	46.12	0.010
QuEST	2	32.30	0.095	32.14	0.007	30.67	0.056
BBQ	2	31.80	0.023	31.75	0.005	30.18	0.010
QuEST	3	27.27	0.043	27.31	0.010	26.49	0.029

Appendix

Table 11: Comparing STE, PQN, and CEWT.

Model	Size	Optimizer	LR	Bits W/A	Perplexity
LLaMA	95M	AdamW (wd=0.1)	0.0012	32/16	24.76
LLaMA	95M	AdamH	0.01	32/16	23.95
LLaMA	127M	AdamW (wd=0.1)	0.0012	32/16	21.53
LLaMA	127M	AdamH	0.01	32/16	20.86
LLaMA	206M	AdamW (wd=0.1)	0.0006	32/16	17.93
LLaMA	206M	AdamH	0.005	32/16	17.52

Appendix

Table 12: Full-Precision Perplexity and learning rate of LLaMA models in Table 3 . A “-” indicates an experiment omitted due to resource constraint.

M-S	Quant-Bits	Optimizer	HSWAG ( $\uparrow$ )	ARC-e ( $\uparrow$ )	ARC-c ( $\uparrow$ )	PiQA ( $\uparrow$ )	Winog ( $\uparrow$ )	Wiki ( $\downarrow$ )	Pile 10k ( $\downarrow$ )
L-206	None-16	AdamW	30.12	41.29	17.92	65.29	50.67	48.90	539.96
L-206	None-16	AdamH	30.55	43.98	18.17	64.69	49.80	47.20	538.48
L-332	None-16	AdamH	32.69	46.09	19.45	67.14	51.38	37.92	407.21
L-206	BBQ-1	MuonH	27.13	32.70	17.66	56.64	49.57	128.73	2277.11
L-206	BBQ-1	Muon-CEWT	27.43	32.24	18.77	57.89	53.35	115.87	2363.58
L-206	BBQ-1	AdamH	27.12	33.16	17.41	56.80	50.67	133.04	2362.53

Appendix

Table 13: Zero-shot Evaluation of Models (from Table 3 ) Pre-trained on the C4 Dataset.

M-S	Q-B	GPU	MuonH	Muon-CEWT	AdamH	Adam-CEWT	LionH	Lion-CEWT	SSO	SSO-CEWT
L-95	L-1	RTX 5090	2.44	2.45	2.42	2.43	2.42	2.43	3.42	3.42
L-95	Q-1	RTX 5090	2.74	2.85	2.73	2.84	2.72	2.83	3.76	3.83
L-95	B-1	RTX 5090	2.76	2.91	2.74	2.90	2.74	2.90	3.76	3.90
L-95	L-2	RTX 5090	2.44	2.46	2.43	2.44	2.42	2.43	3.42	3.42
L-95	Q-2	RTX 5090	2.75	2.86	2.74	2.84	2.73	2.85	3.75	3.85
L-95	B-2	RTX 5090	2.77	2.92	2.76	2.91	2.75	2.90	3.75	3.90

Appendix

Table 14: Training Time in Hours of experiments in Table 3 .

Exp ID	Model	Epochs	Quantizer	A/W	Optimizer	RBM	Train Loss	Test Accuracy
1	DEIT-Base	300	BBQ	16/2	AdamH	0.018	4.25	60.22
2	DEIT-Base	300	BBQ	16/2	Adam-CEWT	0.010	4.15	61.04
3	DEIT-Base	300	BBQ	16/1	AdamH	0.020	4.54	56.41
4	DEIT-Base	300	BBQ	16/1	Adam-CEWT	0.010	4.39	58.88
5	DEIT-Base	300	QuEST	16/2	AdamH	0.088	4.19	61.32
6	DEIT-Base	300	QuEST	16/2	Adam-CEWT	0.056	4.16	61.18

Appendix

Table 15: Evaluation on Vision Models pre-trained on Imagenet 64x64.

Exp ID	Model	Quantizer	A/W	Optimizer	Ep. 50	Ep. 100	Ep. 150	Ep. 200	Ep. 250	Ep. 299
1	DEIT-Base	BBQ	16/2	AdamH	0.008	0.010	0.012	0.013	0.012	0.018
9	DEIT-Base	BBQ	2/2	AdamH	0.006	0.006	0.006	0.006	0.007	0.011
2	DEIT-Base	BBQ	16/1	AdamH	0.005	0.008	0.009	0.009	0.010	0.020
11	DEIT-Base	BBQ	1/1	AdamH	0.003	0.003	0.003	0.004	0.005	0.013
3	DEIT-Base	QuEST	16/2	AdamH	0.077	0.079	0.079	0.079	0.079	0.088
13	DEIT-Base	QuEST	2/2	AdamH	0.067	0.067	0.068	0.069	0.071	0.081

Appendix

Table 16: RBM of Vision Models at different epochs (Ep.). The experiment IDs in this table refers to the IDs from Table 15 .

Quantization-Aware Training

Quantization-Aware Training adapts neural networks during optimization to retain accuracy when weights and activations use low-precision values.

See all 82 papers

Mar 20, 2026

DAQ: Delta-Aware Quantization for Post-Training LLM Weight Compression

We introduce Delta-Aware Quantization (DAQ), a data-free post-training quantization framework that preserves the knowledge acquired during post-training. Standard quantization objectives minimize reconstruction error but are agnostic to the base model, allowing quantization noise to disproportionately corrupt the small-magnitude parameter deltas (

ΔW

) that encode post-training behavior -- an effect we analyze through the lens of quantization as implicit regularization. DAQ replaces reconstruction-based objectives with two delta-aware metrics -- Sign Preservation Rate and Cosine Similarity -- that directly optimize for directional fidelity of

ΔW

, requiring only the base and post-trained weight matrices. In a pilot FP8 study, DAQ recovers style-specific capabilities lost under standard quantization while maintaining general performance.

Xiaoming Yu, Shize Tang, Guanghua Yu +4

Yuanbao & Hunyuan AI Infra Team

Quantization LLM Quantization Neural Network Quantization Post-Training Quantization Weight-Only Quantization

Ask PDF

Oct 6, 2026

Q-PACE: Dynamic Precision Allocation for Quantization-Aware Training

Quantization-aware training (QAT) leverages lower-precision arithmetic to reduce the cost of LLM deployment, but aggressive quantization degrades final model performance. A common remedy is mixed-precision training, in which high precision is assigned to some of the layers to maintain performance while keeping the cost constrained. This approach then requires precision assignments for model layers during training. We provide a new approach, called Q-PACE, consisting of a second-order sensitivity model that predicts the loss increase as a sum of quantization noise MSE weighted by per-layer curvature coefficients. During training, we periodically re-compute these coefficients using perturbations across layers, and re-assign precision. Pretraining and supervised fine-tuning experiments on LLMs of up to 4B parameters show that Q-PACE consistently improves over existing mixed-precision training recipes, and achieves comparable loss at substantially lower total memory budgets. We further find that quantization sensitivity is highly predictable by depth and layer type, and its stability during training allows for infrequent, cheap recalibration.

Alexandra Volkova, Matin Ansaripour, Erik Schultheis +2

ISTA · EPFL, ISTA

LLM Quantization Quantization-Aware Training Quantization Mixed-Precision Quantization Adaptive Quantization Low-Precision Training+2

Ask PDF

Oct 4, 2026

Understanding the Weight Averaging Mechanism in LLM Training for Post-Training Quantization

Large language models (LLMs) are typically pretrained in high precision but increasingly deployed with low-precision post-training quantization (PTQ). Recent studies have shown that using weight averaging during pretraining can improve PTQ performance compared with learning-rate decay, suggesting that it might provide a simple way to improve the pretraining-to-quantization transition. But the mechanism behind weight averaging remains insufficiently explained. This leads to inconsistent and fragile performance gains, thereby preventing practitioners from applying such a technique confidently. As a response, we formulate weight averaging as a trade-off between retaining training progress and improving robustness under perturbation. We further derive a continuous family of averaging kernels that unifies conventional strategies and achieves the Pareto frontier between the two competing goals. Critically, a theoretical framework for performing weight averaging under PTQ is developed. It can be shown that coarser quantization is more susceptible to perturbations, whereas finer quantization could be less affected. Thus, our results could provide unified theoretical guidance for performing weight averaging under different PTQ conditions. Experiments validate both the predicted behavior and the proposed averaging strategy. Code is available at https://github.com/MOFA-LAB/weight-averaging-for-ptq.

Hanzhang Wang, Tianqi Shen, Zonglin Liu +3

City University of Hong Kong · The University of Hong Kong

LLM Quantization Quantization Robustness Post-Training Quantization Neural Network Quantization

Ask PDF

Abstract

Quantization-Aware Training

Explore similar work

DAQ: Delta-Aware Quantization for Post-Training LLM Weight Compression

Q-PACE: Dynamic Precision Allocation for Quantization-Aware Training

Understanding the Weight Averaging Mechanism in LLM Training for Post-Training Quantization