QuaRot

Summary

• Activation quantization!
• Use rotation to remove outliers from the hidden state without changing the output.
• Computational invariance.
• Applied to hidden states (residual), activations of feed forward components, attention mechanism, and KV cache.
• 99% accuracy retain on 4-bit quantized Llama-2-70B.
• Lossless 6 & 8-bit without any calibration data using RTN.

Methodology

Background

Hadamard Transform

• Walsh-Hadamard matrix:

  $${\mathbf{H}}_2=\frac{1}{\sqrt{2}}\left[\begin{array}{cc}1& 1\\ 1& -1\end{array}\right]{\mathbf{H}}_n=\frac{1}{\sqrt{2}}\left[\begin{array}{cc}{\mathbf{H}}_{n-1}& {\mathbf{H}}_{n-1}\\ {\mathbf{H}}_{n-1}& -{\mathbf{H}}_{n-1}\end{array}\right]$$
  • The computation complexity of ${\mathbf{H}}_d\mathbf{x}$ is $O(d{\log }_2(d))$

• Randomized Hadamard matrix:

  $$\mathbf{s}\in {\left\{-1,1\right\}}^d\tilde{\mathbf{H}}=\mathbf{H}\cdot \text{Diag}\left(\mathbf{s}\right)$$
  • $\tilde{\mathbf{H}}$ is also orthogonal.

Incoherence Processing

• Introduced by QuiP.
• A weight matrix $\mathbf{W}\in {\mathbb{R}}^{m\times n}$ is $\mu$-incoherent if $$\max \left(\mathbf{W}\right)\le \mu \cdot \frac{{\Vert \mathbf{W}\Vert }_F}{\sqrt{mn}}$$
  • $\max$ is the element-wise max of the matrix.
    • Is this different from $\Vert \mathbf{W}{\Vert }_{\infty }$?
• A matrix that has high incoherence is hard to quantize: the largest element is an outlier relative to the magnitude of the average element.

Computational Invariance

$$\text{RMSNorm}(\mathbf{X})=\text{RMSNorm}(\mathbf{X}{\mathbf{Q}}^{\top })\mathbf{Q}$$

• No re-scaling happens in the RMSNorm block.
  • In practice, re-scaling is absorbed into the adjacent weight matrices first.

Method

Inter-Block Activation

• Goal: apply incoherence processing on the activations between decoder blocks, i.e., make them easily quantizable by rotating them, $\mathbf{X}\leftarrow \mathbf{XQ}$.
• Merge RMSNorm into all front linear layers ($\text{diag}(\alpha )$).
• Choose a global randomized Hadamard transform $\mathbf{Q}$.
  • All activations between decoder blocks are rotated, therefore needed to rotate back.
• ${\mathbf{W}}_i\leftarrow {\mathbf{Q}}^{\top }\text{diag}(\mathbf{\alpha }){\mathbf{W}}_i,i\in \left\{\text{gate, up, q, k, v}\right\}$
• ${\mathbf{W}}_i\leftarrow {\mathbf{W}}_i\mathbf{Q},i\in \left\{\text{o, down}\right\}$

Input Activation to Down Projection

• Goal: apply incoherence processing on the input activations to down_proj
• Insert an online Hadamard transformation before ${\mathbf{W}}_{\text{down}}$.
• Merge the other half of the Hadamard transformation into ${\mathbf{W}}_{\text{down}}\leftarrow \mathbf{H}{\mathbf{W}}_{\text{down}}\mathbf{Q}$
• Apply quantization to the input of ${\mathbf{W}}_{\text{down}}$

Input Activation to Out Projection

• Goal: apply incoherence processing on the input activations to o_proj.
• ${\mathbf{W}}_v\leftarrow {\mathbf{Q}}^{\top }(\alpha ){\mathbf{W}}_v{\mathbf{H}}_{\text{head}}$
• ${\mathbf{W}}_{\text{out}}\leftarrow \mathbf{H}{\mathbf{W}}_{\text{out}}\mathbf{Q}$
• An additional online Hadamard (Hadamard heads) is required to head-wise rotate output of multi-head attention.

KV-Cache

• Due to positional embedding, Hadamard transformation cannot be merged into weights.
• Online Hadamard after RoPE to query and key, which would be canceled out during attention.
• Since K and V are both rotated, apply quantization and store them into KV cache.