QuaRot and SpinQuant - rotations vs outliers

QuaRot and SpinQuant - rotations vs outliers#

Track 05 - Serving · Notebook 07 · Runtime: ≈30 s on CPU

Prerequisites: 05_serving/05 (GPTQ/AWQ), 05_serving/06 (SmoothQuant).

Papers:

Ashkboos et al. 2024, QuaRot (2404.00456).

Liu et al. 2024, SpinQuant (2405.16406).

What#

SmoothQuant (notebook 06) scales channels to reduce outliers. QuaRot and SpinQuant instead rotate the hidden space: multiply activations and weights by an orthogonal matrix Q (and its inverse Q^-1) so (Qx) · (Q^-1 W)^T = x · W^T exactly.

After rotation, outliers mix with many non-outlier channels and the resulting distribution becomes more Gaussian - easier to quantise. Unlike SmoothQuant, the rotation handles both positive and negative outliers and spreads them across the whole hidden dim.

Hadamard matrices are the canonical choice: entries ±1/√d, H·Hᵀ = I. For d = 2^k they exist and can be applied in O(d log d) via the fast Hadamard transform.

We implement: random-rotation baseline, Hadamard rotation, compare quantisation MSE of a linear layer with/without rotation.

from llm_systems_cookbook.nb import bootstrap

import numpy as np
import torch

s = bootstrap("05_serving_07_quarot_spinquant_rotations")

Hadamard matrix#

For d = 2^k, Sylvester’s construction gives a d × d matrix whose entries are ±1/√d with H Hᵀ = I.

H_1 = [[1]]
H_{2n} = (1/√2) * [[H_n,  H_n],
                  [H_n, -H_n]]

def hadamard(d: int) -> torch.Tensor:
    '''Return a d x d orthonormal Hadamard matrix (d must be a power of 2).'''
    assert d > 0 and (d & (d - 1)) == 0, "d must be a power of 2"
    H = torch.ones(1, 1)
    while H.shape[0] < d:
        H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0)
    return H / np.sqrt(d)


H16 = hadamard(16)
print(f"H16 shape = {H16.shape}")
ortho_err = (H16 @ H16.T - torch.eye(16)).abs().max().item()
print(f"max abs(H·Hᵀ - I) = {ortho_err:.3e}")
s.check("hadamard_is_orthonormal", lambda: ortho_err < 1e-6, msg=f"ortho err={ortho_err:.3e}")

Test layer and outlier activations#

A 1024 → 1024 linear layer with Gaussian weights. Activations have three outlier channels (12×, 8×, 6×) - the same realistic pattern used in notebook 06.

torch.manual_seed(0)
IN, OUT = 1024, 1024
W = torch.randn(OUT, IN) * 0.05
x = torch.randn(32, IN) * 0.5
outlier_channels = [17, 123, 456]
outlier_mags = [12.0, 8.0, 6.0]
for c, m in zip(outlier_channels, outlier_mags, strict=True):
    x[:, c] *= m

print(f"max |x| per channel: top 5 = {sorted(x.abs().amax(dim=0).tolist(), reverse=True)[:5]}")

Quantise with and without rotation#

Pipeline:

Baseline: int4 RTN on W and x directly.
Rotated: compute H = hadamard(IN); transform x' = x · H, W' = W · H. The outer-product W · xᵀ = W' · x'ᵀ (because H Hᵀ = I). Quantise W' and x'.

With rotation, the outlier mass is smeared across every channel. Quantisation error per element stays bounded.

def rtn_int4_grouped(t: torch.Tensor, group: int = 64) -> torch.Tensor:
    flat = t.reshape(-1)
    pad = (-flat.numel()) % group
    if pad:
        flat = torch.cat([flat, flat.new_zeros(pad)])
    g = flat.reshape(-1, group)
    scale = g.abs().amax(dim=1, keepdim=True) / 7 + 1e-9
    q = torch.round(g / scale).clamp(-8, 7)
    out = (q * scale).reshape(-1)[: t.numel()].reshape(t.shape)
    return out


# Baseline: quantise W, x directly.
W_q_naive = rtn_int4_grouped(W)
x_q_naive = rtn_int4_grouped(x)
y_ref = x @ W.T
y_naive = x_q_naive @ W_q_naive.T
mse_naive = (y_ref - y_naive).pow(2).mean().item()

# Rotated: multiply by Hadamard on the input-dim axis before quantising.
H = hadamard(IN)
x_rot = x @ H
W_rot = W @ H
W_q_rot = rtn_int4_grouped(W_rot)
x_q_rot = rtn_int4_grouped(x_rot)
# Because H Hᵀ = I:  x_rot · W_rotᵀ = x · H · Hᵀ · Wᵀ = x · Wᵀ.
y_rot = x_q_rot @ W_q_rot.T
mse_rot = (y_ref - y_rot).pow(2).mean().item()

print(f"naive int4           MSE = {mse_naive:.4e}")
print(f"hadamard-rotated int4 MSE = {mse_rot:.4e}")
print(f"improvement           = {mse_naive / mse_rot:.1f}x less error")

s.check(
    "rotation_reduces_quant_error",
    lambda: mse_rot < mse_naive,
    msg=f"naive={mse_naive:.3e}  rot={mse_rot:.3e}",
)
s.check(
    "rotation_at_least_1_5x_improvement",
    lambda: mse_naive / max(mse_rot, 1e-12) >= 1.5,
    msg=f"ratio = {mse_naive / max(mse_rot, 1e-12):.2f}x",
)

SpinQuant: learned rotation#

Hadamard is a fixed rotation. SpinQuant learns the rotation by minimising post-quantisation MSE via gradient descent on the orthogonal group (using a Cayley or exponential parameterisation).

Here we run 50 gradient steps on a rotation parameterised as Q = exp(A - Aᵀ) (any skew-symmetric A gives an orthogonal Q). It should do at least as well as Hadamard on this specific layer.

def exp_skew(A: torch.Tensor) -> torch.Tensor:
    '''Matrix exponential of skew-symmetric (A - Aᵀ) - gives orthogonal.'''
    S = A - A.T
    # Rough approximation for small S: I + S + S² / 2.
    return torch.eye(S.shape[0]) + S + (S @ S) / 2


def spinquant_loss(A: torch.Tensor, W: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    Q = exp_skew(A)
    x_rot = x @ Q
    W_rot = W @ Q
    W_q = rtn_int4_grouped(W_rot)
    x_q = rtn_int4_grouped(x_rot)
    y_hat = x_q @ W_q.T
    y = x @ W.T
    return (y - y_hat).pow(2).mean()


A = torch.zeros(IN, IN, requires_grad=True)
opt = torch.optim.SGD([A], lr=0.005)
for step in range(30):
    loss = spinquant_loss(A, W, x)
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step in (0, 29):
        print(f"step {step:>2}  loss = {loss.item():.4e}")

mse_spin = spinquant_loss(A, W, x).item()
print(f"SpinQuant learned rotation MSE = {mse_spin:.4e}")

s.check(
    "spinquant_learns_non_identity",
    lambda: A.abs().sum().item() > 0.01,
    msg=f"|A|_sum = {A.abs().sum().item():.3f}",
)
s.check(
    "spinquant_not_catastrophically_worse_than_hadamard",
    lambda: mse_spin <= mse_rot * 3,
    msg=f"hadamard={mse_rot:.3e}  spin={mse_spin:.3e}",
)

Activation distribution: before vs after Hadamard#

Two histograms of the same activation tensor. Pre-rotation, the long tail from the three outlier channels pulls the axis out to ±20; the quantiser must waste codepoints on that tail. Post-rotation, the outlier mass gets smeared across every channel and the distribution collapses toward Gaussian - a much easier target for int4 RTN.

import matplotlib.pyplot as plt

pre  = x.flatten().numpy()
post = x_rot.flatten().numpy()
lim  = float(max(abs(pre).max(), abs(post).max())) * 1.05

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.6), sharey=True)
ax1.hist(pre, bins=80, range=(-lim, lim), color="tab:red",
         alpha=0.75, edgecolor="black", lw=0.3)
ax1.set_title(f"before rotation (kurtosis={(((pre - pre.mean()) / pre.std()) ** 4).mean():.1f})")
ax1.set_xlabel("activation value")
ax1.set_ylabel("count (log)")
ax1.set_yscale("log")
ax1.axvline(-7 * pre.std(), color="black", ls=":", alpha=0.4)
ax1.axvline( 7 * pre.std(), color="black", ls=":", alpha=0.4)

ax2.hist(post, bins=80, range=(-lim, lim), color="tab:green",
         alpha=0.75, edgecolor="black", lw=0.3)
ax2.set_title(f"after Hadamard (kurtosis={(((post - post.mean()) / post.std()) ** 4).mean():.1f})")
ax2.set_xlabel("activation value")
ax2.set_yscale("log")

fig.suptitle("rotation smears outliers -> tighter tails -> easier to quantise")
fig.tight_layout()
plt.show()

Exercises#

Randomised Hadamard (the QuaRot variant): multiply by a signed diagonal D with ±1 entries before H. The randomisation de-correlates outliers across rows. Compare MSE.
Fast Hadamard transform. scipy.linalg.hadamard(d) / √d works for small d, but for d = 8192 you want the O(d log d) FWHT. Implement it recursively and time vs dense multiplication.
Outlier-free check. After rotation, histogram the activations and confirm they look more Gaussian (kurtosis closer to 3).

References#

QuaRot paper §3 for the online Hadamard-rotation layer.
SpinQuant paper §3 for the learned rotation on the orthogonal group.

s.summary()
s.save()