Using orthogonal matrices for better quantization - The math

Orthogonal Matrices and Quantization

Sources:

• QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs

• SliceGPT: Compress Large Language Models by Deleting Rows and Columns

• QuIP: 2-Bit Quantization of Large Language Models With Guarantees

• o3

TLDR: It reduces outliers.

Why we care about orthogonal matrices

Let $Q$ denote an orthogonal matrix: $Q^{\top }Q=Q{Q}^{\top }=I.$ Note that multiplying a vector $x$ by $Q$ does not change the norm of the vector, since $\Vert Qx\Vert =\sqrt{x^{\top }{Q}^{\top }Qx}=\sqrt{x^{\top }x}=\Vert x\Vert .$

In this work, the dimensions of $Q$ will always match the embedding dimension of the transformer $D$.

1. Setup

We work with a single linear layer

$Y=XWX\in {\mathbb{R}}^{B\times d},W\in {\mathbb{R}}^{d\times m}$

and a uniform symmetric $b$-bit scalar quantiser

$Q_b(z)=\text{clip}\left(\lfloor z/\Delta \rceil ,-2^{b-1},2^{b-1}-1\right)\Delta ,\Delta =\frac{2\alpha }{2^b-1},$

where the clipping bound $\alpha$ is normally chosen as the largest absolute entry of the tensor that is being quantised. $\Delta$ defines the quantization cell width.

We rotate both weights and activations once, with a fixed orthogonal $Q\in {\mathbb{R}}^{d\times d}$ ($Q^{\top }Q=I$):

$X_{\text{rot}}=XQ,W_{\text{rot}}=Q^{\top }W.$

Because $Q^{\top }Q=I$,

$X_{\text{rot}}{W}_{\text{rot}}=XQ{Q}^{\top }W=XW,$

so the network is mathematically unchanged before quantisation.

2. Quantisation noise before and after the rotation

Write the quantisation error of a tensor $T$ as

${\epsilon }_T=Q_b(T)-T$

In general, the element-wise quantization MSE of the tensor T is equal to $\frac{1}{|T|}E[\Vert {\epsilon }_T{\Vert }_F^2]=\frac{{\Delta }^2}{12}$.

A small proof is that, consider one quantization cell of width $\Delta$, the error $e$ on that cell is the difference between the original sample and the reconstruction point (the cell’s midpoint). We thus assume the error follows $e\approx \text{Unif}[-\frac{\Delta }{2},\frac{\Delta }{2}]$. Then, $E[e^2]=\frac{{\Delta }^2}{12}$, following the variance of a uniform variable.

We focus on $\tilde{Y}=\tilde{X}\tilde{W}$, the resulting output after having quantized both activations and inputs.

Keeping only first–order terms (the usual high-resolution approximation):

• Without rotation:

• $\tilde{Y}=(X+{\epsilon }_X)(W+{\epsilon }_W)\approx XW+X{\epsilon }_W+{\epsilon }_{X}W$

• With rotation:

• $\tilde{Y}=(X_{\text{rot}}+{\epsilon }_{X_{\text{rot}}})(W_{\text{rot}}+{\epsilon }_{W_{\text{rot}}})\approx XW+X{\epsilon }_{W_{\text{rot}}}+{\epsilon }_{X_{\text{rot}}}W$

Hence the output MSE:

$E[\Vert \tilde{Y}-XW{\Vert }_F^2]\propto \underset{\text{weight noise}}{\underbrace{E[\Vert X{\epsilon }_W{\Vert }_F^2]}}+\underset{\text{activation noise}}{\underbrace{E[\Vert {\epsilon }_{X}W{\Vert }_F^2]}}$

The rotation will help because it reduces the variance of every coordinate in ${\epsilon }_W$ and ${\epsilon }_X$; that lets us pick a smaller $\alpha$ ($\Rightarrow$ smaller $\Delta$ $\Rightarrow$ smaller noise $\epsilon$). This is what we will derive next.

3. Why an orthogonal rotation shrinks the required dynamic range

Consider a single length-$d$ vector $v$ (a row of $X$ or a column of $W$).

The quantiser step is $\Delta \propto \alpha ={\max }_{1\le j\le d}|v_j|$.

Let $v^{\prime }=vQ$ (a mere change of basis).

Because $Q$ is orthogonal, $\Vert v^{\prime }{\Vert }_2=\Vert v{\Vert }_2$.

Yet its max-entry is, with high probability, vastly smaller because of the following lemma:

Lemma (“cube-to-sphere” effect)

Let $v\in {\mathbb{R}}^d$ be any non-zero vector and let $Q\in O(d)$ be drawn uniformly at random (Haar measure). Write

$z=vQ=\Vert v{\Vert }_2\frac{g}{\Vert g{\Vert }_2},g\sim \mathcal{N}(0,I_d),$

so $z$ is a random point on the ordesphere of radius $\Vert v{\Vert }_2$. Then with probability $1-o(1)$ (thanks to Proposition 1.1 from Luh et O’Rourke),

$\Vert z{\Vert }_{\infty }\le C\Vert v{\Vert }_2\sqrt{\frac{\log d}{d}},$

and the expectation of the maximum coordinate satisfies

$\mathbb{E}[\Vert z{\Vert }_{\infty }]=\Theta (\Vert v{\Vert }_2\sqrt{\frac{\log d}{d}}).$

(The constant $C$ is universal and usually quoted between 2 and 3 in the literature.)

Why is it true?

1. Gaussian picture – A random rotation turns the fixed vector $v$ into a direction whose coordinates behave as if they were i.i.d. $\mathcal{N}(0,\Vert v{\Vert }_2^2/d)$.

Formal step: because $g_i\sim \mathcal{N}(0,1)$ are independent, $z_i=\Vert v{\Vert }_2{g}_i/\Vert g{\Vert }_2$. For large $d$ we have $\Vert g{\Vert }_2\approx \sqrt{d}$ with high probability.

2. Order statistics of Gaussians – For $d$ independent $\mathcal{N}(0,{\sigma }^2)$ variables, we can bound its maximum expected value

$\mathbb{E}[{\max }_{1\le i\le d}|X_i|]\le \sigma \sqrt{2\log d}.$

3. Put together – Substitute $\sigma =\Vert v{\Vert }_2/\sqrt{d}$:

$\mathbb{E}[\Vert z{\Vert }_{\infty }]\le \frac{\Vert v{\Vert }_2}{\sqrt{d}}\sqrt{2\log d}=\Vert v{\Vert }_2\sqrt{\frac{2\log d}{d}}.$

Consequence:

$\frac{{\alpha }_{\text{rot}}}{{\alpha }_{\text{plain}}}\approx \sqrt{\frac{2\log d}{d}}\ll 1\Rightarrow {\Delta }_{\text{rot}}\ll {\Delta }_{\text{plain}}.$

Because the element-wise quantisation MSE is $\frac{1}{12}{\Delta }^2$, the overall error falls by roughly the same factor.

Even a structured fast transform (Hadamard, DCT, butterfly, Householder chains, Givens cascades, …) is “sufficiently random” for typical correlated neural activations and weights: it still spreads energy almost uniformly, so the dynamic-range argument largely holds in practice.

4. Putting it all together

• Rotation preserves function value (exact linear algebra identity).

• Max-entry shrinks (cube-to-sphere), so the clipping range $\alpha$ and step size $\Delta$ shrink.

• Quantisation noise power $\propto {\Delta }^2$ therefore falls.

• Because we rotate both activations and weights, the first-order output error is reduced on both terms $X{\epsilon }_W$ and ${\epsilon }_{X}W$.

• Fast structured orthogonals (e.g. Hadamard matrices) give most of the benefit with $O(d\log d)$ or even $O(d)$ FLOPs instead of the $O(d^2)$ cost of a dense rotation.

Hence: applying one properly chosen (but inexpensive) orthogonal matrix before a single quantisation pass makes the tensors look “rounder”, dramatically lowering the dynamic range that the scalar quantiser must cover, and thereby reducing the quantisation error—all while leaving the original network computation untouched.

How to use orthogonal matrices in a classic residual NN (purely offline)

• The computational invariance theorem in SliceGPT states that the weights and between-block activations of an RMSNorm transformer can be transformed using an orthogonal matrix with no change to the model output.

• Here’s the notation from SliceGPT (note that this is post-norm):

Here is the proof for applying orthogonal matrices in a function-preserving way:

• First fact: $\text{RMSNorm}(XQ)Q^T=\text{RMSNorm}(X)Q{Q}^T=\text{RMSNorm}(X_l)$

• Second fact: For a standard post-norm residual block, the following network is equivalent to the one describe above:

  • Assume the input-embedding is right-multiplied i.e. ${\tilde{W}}_{embd}=W_{embd}Q$
  • We will assume the following invariant,
    • for each layer, the input $\tilde{X}$ is equal to the original layer input right-multiplied by $Q$ i.e. $\tilde{X}=XQ$
  • Then we apply the following transforms to each layer:
    • $\tilde{Z}\leftarrow {\sigma }_l(\tilde{X}(Q^T{W}_{in}^l)+b_{in}^l)W_{out}^{l}Q+b_{out}^{l}Q$
      • note that $\tilde{Z}=ZQ$
    • ${\tilde{X}}_{new}\leftarrow \text{RMSNorm}(\tilde{X}+\tilde{Z})$
      • which is equal to $\text{RMSNorm}((X+Z)Q)=\text{RMSNorm}(X+Z)Q$
      • this is the expected output $X$ right-mutliplied.
  • Finally, we “reverse” the right-transform by left-multiplying the LM head
    • ${\tilde{W}}_{head}=Q^T{W}_{head}$

• The general rules can be summarized as:

  • use RMSNorm
  • The embedding matrix should be right-mutliplied
  • The input featurizer weights should be left-multiplied (e.g. QKV projections, gate and up projection in SwiGLU)
    • We can absorb the RMSNorm diagonals scaling parameters $(\alpha )$ into the input featurizer pre-processing.
  • The output featurizer weights should be right-multiplied (usually just a linear layer)
  • The LM head matrix should be left-mutliplied

Why we care about rotation/Hadamard matrices

A Hadamard matrix is an orthogonal matrix with entries drawing from $\{+1,-1\}$. A Walsh-Hadamard matrix is a square matrix of size $d=2^n$, with

$${\mathbf{H}}_2=\frac{1}{\sqrt{2}}\left[\begin{array}{cc}1& 1\\ 1& -1\end{array}\right]\text{and}{\mathbf{H}}_{2^n}={\mathbf{H}}_2\otimes {\mathbf{H}}_{2^{n-1}}.$$

These identities give rise to the Walsh-Hadamard transform, which computes the matrix-vector product $\mathbf{Hx}$ in $\mathcal{O}(d{\log }_2(d))$ operations.

We care because it’s fast to apply basically :)