docs/quantization.md

Quantization Deep Dive

Overview

We explored multiple quantization strategies for LeWM. This document records what was tried, why it worked or didn't, and the engineering trade-offs.

Formats Evaluated

Format	Bits/weight	Compression	Cos vs f32	Status
INT8 (encoder)	8	4x	0.9999	Production
Q4 (predictor)	4	2x more	0.998	Production
INT8+Q4 (full)	mixed	5x total	0.999	Production
Q4 encoder only	4	4x	0.93	Rejected
Ternary {-1,0,+1}	2	8x	~0.85	Rejected
WANDA 20%	sparse Q4	20% fewer	~0.99	Experimental
WANDA 40%	sparse Q4	40% fewer	~0.97	Experimental

INT8 Per-Channel (Encoder)

Method

f32_weight [out, in] → per-output-channel quantization
→ int8_weight [out, in]
→ per-output-channel scales [out]

At inference:
  result = input @ dequant(weight) × scales
         = input @ (int8_weight + zero_point) × scales
         = (input @ int8_weight) × scales   // zero_point = 0 for symmetric

Why It Works

1. Per-channel preserves channel-level statistics. Each output channel gets its own scale.
2. Symmetric (zero_point = 0) avoids the overhead of asymmetric quantization.
3. INT8 GEMV on PIE SIMD: 16-wide multiply-accumulate per cycle.
4. Encoder activations have predictable dynamic range after LayerNorm.

Quality

Layer	cos vs f32
patch_embed	0.9999
encoder layer 0	0.9999
encoder layer 5	0.9998
encoder.proj	0.9999
Total	0.9999

Engineering Notes

• Activation quantization is dynamic (per-row at inference time), not static.
• QKV shared quantization: same normalized input quantized once for Q, K, V.
• Scales stored as f32 (4 bytes per channel) — negligible overhead.

Q4 Block (Predictor)

Method

f32_weight [out, in] → per-32-element-block quantization
→ nibble-packed weight data
→ per-block f16 scales

At inference:
  for each block of 32:
    unpack nibbles → int8 [-8, 7]
    dot = simd_dot(input_block, unpacked)
    result += dot × block_scale

Why It Works

1. Per-block (32 elements) matches the predictor's adaLN normalization.
2. adaLN modulation provides implicit normalization — weights don't need per-channel precision.
3. Smaller layers = less error accumulation. 4-layer predictor vs 6-layer encoder.
4. PIE SIMD handles the nibble unpack + dot in tight loops.

Why Full Q4 (Encoder + Predictor) Fails

When we skip INT8 and quantize the encoder to Q4:

Layer	INT8 cos	Q4 cos
encoder	0.9999	0.93
predictor	0.998	0.998
Total	0.999	0.93

Root cause: ViT encoder has high dynamic range in intermediate activations. The 32-element block granularity doesn't align with the encoder's channel statistics. INT8's per-channel precision is essential.

Engineering Notes

• Nibbles decode as value - 8 → range [-8, 7]
• Block scales stored as f16 (2 bytes per block) — 64 bytes per 32×32 block
• Zero weights are rare in Q4 (<1%) — skip-zero optimization not implemented

Ternary ({-1, 0, +1})

Method

f32_weight → hard threshold
→ +1 if w > +tau
→ -1 if w < -tau
→  0 otherwise
→ bit-packed ternary

At inference:
  result = input @ ternary_weight
         = sum(sign(w_i) × x_i)   // pure addition/subtraction

Why It Fails

Metric	Q4	Ternary
Cos vs f32	0.998	~0.85
Compression	8x vs f32	16x vs f32

Root cause: adaLN generates 6 modulation vectors (scale1, shift1, gate1, scale2, shift2, gate2) that multiply and add to the normalized activations. The magnitudes of these modulation vectors matter — ternary destroys them.

What was tried:

• Various thresholds (τ = 0.5σ, σ, 2σ)
• STE (straight-through estimator) during fine-tuning
• Mixed ternary (ternary weights + fp32 scales)

None recovered quality sufficiently.

WANDA Pruning

Method

1. Forward pass on calibration set → collect activations
2. Compute WANDA score: s(w) = |w| × ||a||
3. Sort scores, prune bottom N%
4. Fine-tune 1 epoch
5. Re-quantize remaining weights to Q4

Results

Model	Pruned	Size	Cos	Notes
Baseline Q4	0%	23.6 MB	0.998	Reference
WANDA 20%	20%	22.0 MB	~0.99	Pruned, no fine-tune
WANDA 40%	40%	25.1 MB	~0.97	Bitmap overhead exceeds savings

Engineering Notes

• 40% pruned model has larger binary than 20% due to bitmap overhead
• Skip-zero GEMV needs hardware support (not implemented in PIE SIMD)
• Would benefit from fine-tuning after pruning

Shared Lessons

1. Per-channel > per-block for encoders: Encoders have high per-channel variance. INT8's per-channel precision beats Q4's per-block.

2. Predictors are quantization-friendly: adaLN provides implicit normalization. Predictors can use Q4 with minimal quality loss.

3. Architecture changes beat quantization: hybrid_ALAL (3.0M params) achieves similar quality to slim_96d (10.2M params) at INT8+Q4. Architecture > precision.

4. Epoch 1 models have headroom: All current slim/epoch-1 models will improve with longer training. The quality comparisons should be re-run at convergence.

5. Hardwired is the limit: Q4 weights decompose to shift-add operations. Zero multiplications, zero memory fetches. That's the theoretical floor.