Quantization-Aware Pre-Training with Constrained Empirical Weight Distribution
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 onstrained mpirical eight disribution (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
| Per-Tensor | Per-Channel | TGCS (Ours) |
|---|---|---|
| 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}) |
| STE | PQN | CEWT | DL | IF | CGA | PI | ||
| MSE (unit: 1e-3) | Good HP | 81±0 | 87±0 | 45±0 | 30±0 | 59±0 | 53±0 | 54±0 |
| Bad HP | 375±1 | 2080±12 | 332±1 | 1847±10 | ||||
| MSE (unit: 1e-3) | Good HP | 67±0 | 4±0 | 36±0 | 34±0 | 60±0 | 55±0 | 48±0 |
| Bad HP | 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 |
| Bad HP | 0±0 | 0±0 | 1±0 | 0±0 |
| Model | Q-B | MuonH | Muon-CEWT | AdamH | Adam-CEWT | LionH | Lion-CEWT | SSO | SSO-CEWT | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 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 |
Appendix figures & tables19 assets
Supplementary material from the paper’s appendix.
Appendix
| Name | Lecun | PyTorch | Spectral | GPT-2 (non-residual) | GPT-2 (residual) |
|---|---|---|---|---|---|
| Init Distribution | Normal | Uniform | Normal | Normal | Normal |
| Init Mean | |||||
| Init Std | |||||
| Frobenius norm | |||||
| AdamH Perp | 31.78 | 31.52 | 31.39 | 30.91 | |
| Adam-CEWT Perp | 30.12 | 30.22 | 29.82 | 29.77 | |
| 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 |
| STE | HTGE | JacQuant-Probe | EWGS | ||
|---|---|---|---|---|---|
| MSE (unit: 1e-3) | Default | 81±0 | 53±0 | 63±0 | 61±0 |
| With CEWT | 45±0 | 46±0 | 52±0 | 51±0 | |
| RBM (unit: 1e-3) | Default | 573±1 | 177±1 | 349±1 | 324±1 |
| 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 |
| With CEWT | 17±0 | 30±0 | 37±0 | 13±0 |
| 1 | 2 | 4 | 8 | 16 | 32 | 64 | 128 | 256 | 512 | 1024 | ||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| MSE (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 (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 (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 (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 |
| Iterations = | 1 | 3 | 6 | 12 | 25 | 50 | 100 | 200 | |
|---|---|---|---|---|---|---|---|---|---|
| MSE (unit: 1e-3) | STE | 1847±10 | 1609±9 | 1277±7 | 747±5 | 244±1 | 90±0 | 82±0 | 81±0 |
| MSE (unit: 1e-3) | CEWT | 1839±9 | 1625±9 | 1343±7 | 859±6 | 281±1 | 107±0 | 63±0 | 45±0 |
| RBM (unit: 1e-3) | STE | 11±0 | 12±0 | 15±0 | 19±0 | 27±0 | 40±0 | 64±0 | 108±1 |
| RBM (unit: 1e-3) | CEWT | 10±0 | 10±0 | 10±0 | 10±0 | 10±0 | 10±0 | 10±0 | 10±0 |
| 2 | 4 | 8 | 16 | 32 | 64 | 128 | 256 | 512 | 1024 | ||
|---|---|---|---|---|---|---|---|---|---|---|---|
| MSE (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 (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 (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 (unit: 1e-3) | CEWT | 0±0 | 0±0 | 13±34 | 11±9 | 11±4 | 10±1 | 10±0 | 10±0 | 10±0 | 10±0 |
| Row | Weight Quantizer | Optimizer | Perplexity | RBM | Notes | ||
|---|---|---|---|---|---|---|---|
| Type | HT | TGCS | |||||
| 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 |
| Quantizer | Bits W/A | AdamH + STE | AdamH + PQN | Adam-CEWT | |||
|---|---|---|---|---|---|---|---|
| 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 |
| 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 |
| M-S | Quant-Bits | Optimizer | HSWAG ( ) | ARC-e ( ) | ARC-c ( ) | PiQA ( ) | Winog ( ) | Wiki ( ) | Pile 10k ( ) |
|---|---|---|---|---|---|---|---|---|---|
| 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 |
| 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 |
| 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 |
| 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 |
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