Quantisation advanced 7 min read 12 flashcards

AWQ and Activation-Aware Scaling

The observation that weight importance is determined by activation magnitude rather than weight magnitude, and how a per-channel rescaling protects the important weights without keeping any of them in higher precision.

Not all weights matter equally, and the usual proxy for importance is wrong. A weight with a large magnitude that multiplies a consistently small activation contributes little to the output; a small weight multiplying a large activation contributes a lot. Activation-aware weight quantisation builds on that observation and turns it into a method that needs no mixed precision and no backpropagation.

The observation

AWQ (Lin et al., 2024, MLSys) reports that protecting roughly one percent of weight channels, selected by the average magnitude of the activations they multiply, recovers most of the accuracy lost to 4-bit quantisation. Selecting the same fraction by weight magnitude does substantially worse, and selecting at random does worse still. Importance follows the activation, not the weight.

The obvious implementation, keeping that one percent in fp16, works and is unattractive. Mixed-precision tensors mean irregular memory layouts, kernels that handle two formats, and a hardware-unfriendly access pattern that gives back much of the speedup.

Scaling instead of protecting

The trick is that quantisation error is relative to the scale. Multiply a weight channel by \(s > 1\) before quantising and divide the corresponding activation channel by \(s\), and the product is mathematically unchanged while the weight now occupies a larger share of its quantisation grid, so its relative rounding error falls by roughly \(s\).

\[W X = \left(W \cdot \operatorname{diag}(s)\right)\left(\operatorname{diag}(s)^{-1} X\right)\]

Apply a per-input-channel scale \(s\), chosen larger for channels with large average activation magnitude, and the important weights get finer effective resolution while everything stays uniformly 4-bit. Crucially, the activation-side division can be folded into the preceding layer's weights, usually the preceding normalisation's scale parameters, so at inference there is no extra operation at all.

The scale is parameterised as \(s = \bar{a}^{\alpha}\) where \(\bar{a}\) is the per-channel mean activation magnitude and \(\alpha \in [0,1]\) is found by a small grid search per layer, minimising reconstruction error. At \(\alpha = 0\) the method reduces to plain quantisation; at \(\alpha = 1\) it over-corrects and inflates the error on the unimportant channels enough to hurt.

Why this is attractive in practice

It requires no backpropagation, only forward passes to collect activation statistics, so it runs in minutes rather than hours. It preserves a uniform data layout, so the kernels are simple and fast. It does not overfit the calibration set the way a reconstruction-minimising method can, because it only uses the first-order statistic of activation magnitude rather than fitting a full Hessian, which is why AWQ tends to hold up better than GPTQ when calibration and deployment distributions differ.

When it breaks

Scaling moves the outlier problem, it does not remove it. Dividing an activation channel by \(s\) shrinks it, which is fine when activations stay in fp16. If activations are also quantised, that division consumes activation grid range, and the two objectives conflict directly. AWQ is a weight-only method for this reason, and combining it with activation quantisation requires a scheme such as rotation that redistributes energy rather than trading it between tensors.

Folding the inverse scale is not always possible. It requires a preceding operation whose weights can absorb a per-channel scale, typically a LayerNorm or a linear layer. Where the previous operation is a residual add or a non-linearity applied elementwise across a shared tensor, the fold has nowhere to go and the division becomes a real runtime cost.

Grid-searching alpha per layer is a small overfit. The search minimises reconstruction error on the calibration set, and the optimum is fairly flat, so a value found on 128 samples transfers acceptably. It is still a fitted hyperparameter per layer, and reporting results without noting that the search used the evaluation domain's data is a subtle form of leakage.

Group size interacts with the scaling. AWQ's benefit is largest when the group is large enough that outliers within a group would otherwise dominate its scale. At group size 32, per-group scales already absorb much of the same variation, and the additional gain from activation-aware scaling narrows, which makes the two techniques partly redundant rather than additive.

Check yourself

12 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track