Depth Registers Unlock W4A4 on SwiGLU: A Reader/Generator Decomposition

1. Problem Statement

The paper addresses the catastrophic performance degradation observed when applying W4A4 quantization (4-bit weights, 4-bit activations) to SwiGLU-based Large Language Models (LLMs) using naive round-to-nearest (RTN) methods.

• The Phenomenon: In a controlled 300M-parameter model trained on 5B tokens, naive W4A4 causes validation perplexity (PPL) to collapse from 23.6 (FP16) to 1727.
• The Gap: Existing post-training quantization (PTQ) methods (e.g., SmoothQuant, QuaRot, AWQ) recover some performance but often rely on implicit hypotheses about where the error originates. They fail to isolate whether the damage stems from inputs derived from the residual stream or from block-internal computations.
• The Core Question: Which specific input-activation sites in a SwiGLU block dominate the quantization error, and can we theoretically bound these errors based on their source?

2. Methodology: The Reader/Generator Decomposition

The authors propose a theoretical and empirical decomposition of the five trainable linear layers in a SwiGLU block based on their input sources:

A. Theoretical Taxonomy

1. Readers (Residual-Axis Inputs):
   • Layers: qkv (Attention), w1, w3 (MLP).
   • Input Source: The residual stream after RMSNorm.
   • Mechanism: Their input magnitude is directly bounded by any intervention that constrains the residual stream's norm.
2. Generators (Block-Internal Inputs):
   • Layers: o_proj (Attention output), w2 (MLP down-projection).
   • Input Source: Functions of internal block computations.
     • o_proj: Input is a weighted sum of value projections ($\sum \alpha_j V_j$).
     • w2: Input is the bilinear product of the SwiGLU gate: $\text{silu}(w_1(x)) \odot w_3(x)$.
   • Mechanism:
     • o_proj is only loosely bounded by residual constraints (dependent on attention weights and $W_V$).
     • w2 is not tightly bounded by residual constraints. The paper argues that the bilinear product creates a "heavy tail" where a rare large value in one factor multiplied by a typical value in the other creates an outlier not present in either factor individually.

B. The Intervention: Depth Registers + Sink Loss (DR+sink)

To test this decomposition without relying on post-hoc calibration or rotation, the authors introduce a training-time intervention:

• Depth Registers: The residual channel dimension is partitioned into a "semantic" region and a "register" region ($k$ channels).
• Sink Loss: A hinge loss is applied specifically to the register region, penalizing any magnitude exceeding a target threshold ($\tau$).
• Goal: This forces the network to offload outlier magnitudes into the register channels, effectively bounding the magnitude of the semantic residual stream (the "Readers") without using post-hoc scaling or rotation.

3. Key Contributions

1. Operator-Level Decomposition: A formal separation of SwiGLU inputs into Readers (residual-bounded) and Generators (internal-bounded), supported by elementary norm arguments showing that residual-axis interventions cannot bound the bilinear tail of w2.
2. DR+sink Training Probe: A novel training-time method that reduces naive W4A4 PPL from 1727 to 119 (a ~14x improvement) while maintaining matched FP16 PPL and zero-shot capacity.
3. Damage Budget Analysis: Introduction of a "skip-ablation" metric in excess NLL ($\Delta NLL$) to quantify damage contribution. This reveals that once Readers are "rescued," a flat residue of ~0.28 nats remains, dominated entirely by the w2 generator.
4. Empirical Validation of Theory: Demonstrating that orthogonal rotations (like in full QuaRot) fail to close the gap for w2, confirming the theoretical prediction that rotation cannot suppress the heavy tail of a bilinear product.

4. Key Results

A. Performance Metrics (300M Model, 5B Tokens)

• DR+sink alone reduces PPL to 119.
• DR+sink + SmoothQuant achieves 39.9 PPL.
• Per-Linear QuaRot (post-hoc rotation) matches DR+sink performance (~25.4 PPL) on the Reader axis but cannot fix the Generator axis.
• Full QuaRot (adding online w2 rotation) fails to improve further (26.0 PPL), confirming the w2 tail is resistant to orthogonal rotation.

B. Statistical Evidence

• Kurtosis Collapse: DR+sink collapses the excess kurtosis of Readers (e.g., qkv drops from 35.9 to 2.85, near-Gaussian) but leaves Generators unchanged (w2 remains ~1900).
• Damage Decomposition:
  • The "Reader-rescued residue" (damage remaining when Readers are kept in FP16) is flat at ~0.28 nats across all checkpoints.
  • $\Delta \text{remove}(w_2)$ (damage removed by keeping w2 in FP16) is 0.69–0.88 nats, dominating the remaining error.
  • $\Delta \text{remove}(o\_proj)$ is negligible once Readers are clean.

5. Significance and Implications

• Diagnosis of W4A4 Failure: The paper definitively identifies that the primary bottleneck for W4A4 in SwiGLU models is the bilinear w2 input, not the residual stream or attention mechanisms.
• Limitations of Rotation: It provides strong empirical evidence that orthogonal rotations (a staple of modern PTQ like QuaRot) are insufficient for bounding the heavy tails generated by bilinear products. Rotation preserves second-order structure but does not suppress fourth-moment growth in products.
• Future Directions:
  • Since w2 is the irreducible source of error, future solutions should focus on mixed-precision w2, Quantization-Aware Training (QAT) specifically for the MLP down-projection, or product clipping.
  • The "Reader/Generator" taxonomy offers a new framework for analyzing quantization errors in other architectures.
• Caveats: The results are specific to a 300M parameter model trained on 5B tokens with a single seed. While the mechanism is theoretically sound, scaling to 7B+ models requires further validation.

In summary, the paper argues that W4A4 quantization failure in SwiGLU is a "Generator" problem, not a "Reader" problem. While training-time interventions (DR+sink) or post-hoc rotations can fix the residual-stream inputs, the bilinear nature of the SwiGLU down-projection (w2) creates a stubborn error tail that requires targeted, non-rotational solutions.