AWQ › Salient Weight Preservation

The simplest approach to protecting salient weights would be to keep them in higher precision, creating a mixed-precision representation where important weights use 8 or 16 bits while less important weights use 4 bits. However, mixed-precision schemes complicate hardware implementation and sacrifice some of the memory and speed benefits we seek from quantization. Inference kernels optimized for uniform 4-bit weights cannot efficiently handle a mix of precisions, and the irregular memory access patterns that result from mixed precision degrade performance. AWQ takes a more clever approach: it uses per-channel scaling factors to effectively "hide" more precision in the quantized representation for salient weights while maintaining a uniform bit width throughout the model.

To understand AWQ's approach, recall from our earlier discussion of quantization basics that the quantization error for a weight $w$ scales with the quantization step size. The step size determines how far apart adjacent representable values are in the quantized representation. When we quantize to $b$ bits with scale $s$, the maximum error is approximately $s/2$, since any real value can be at most half a step away from the nearest quantization level.

For a standard per-tensor or per-channel quantization scheme, the scale is determined by the range of values being quantized. Let the minimum and maximum weights in a group be $w_{\min}$ and $w_{\max}$. Then:

$$s_q = \frac{w_{\max} - w_{\min}}{2^b - 1}$$

where:

• $s_q$: the quantization scale (step size), determining the distance between adjacent representable values in the integer grid
• $w_{\max}$: the maximum value in the weight tensor, defining the upper bound of the dynamic range
• $w_{\min}$: the minimum value in the weight tensor, defining the lower bound of the dynamic range
• $b$: the quantization bit width (4 for INT4 quantization)
• $2^b - 1$: the number of quantization intervals, which equals the maximum representable integer with $b$ bits

This formula reveals an important insight: the quantization error is determined by how much of the available integer range each weight occupies. Weights near the extremes of the distribution (close to $w_{\max}$ or $w_{\min}$) effectively receive finer resolution because they use more of the available quantization levels. Weights in the middle of the distribution share the interior levels and receive the same granularity. The range is the control point: changing the quantized range changes the resolution that all weights receive.

Now consider what happens when we multiply a weight by a constant $k > 1$ before quantization, then divide by $k$ after dequantization. The weight's contribution to the output is unchanged because the scaling operations cancel out exactly. But its quantization behavior shifts in a favorable way. The scaled weight $k \cdot w$ occupies proportionally more of the quantization range, effectively getting finer granularity in the discrete representation. When we dequantize and divide by $k$, the quantization error is also divided by $k$, meaning the effective error in the original weight space is reduced by the same factor $k$.

The catch is that scaling up one weight affects the overall range, potentially increasing quantization error for other weights in the same quantization group. If we scale up salient weights without adjusting anything else, the quantization scale $s_q$ increases to accommodate the larger values, and non-salient weights receive coarser quantization as a result. This is where the activation-awareness becomes important. If weight $w_{ij}$ multiplies activations that are on average $k$ times larger than typical, scaling that weight up by $k$ while scaling down the corresponding activation channel by $k$ preserves the computation while reducing quantization error for that salient weight. The inverse scaling on the activation side can be absorbed into preceding layers, making the transformation invisible at inference time.

Mathematically, for an input channel $j$, we can introduce a scaling factor $s_j$ applied to the weights while absorbing the inverse into the incoming activations:

$$\begin{aligned} y_i &= \sum_j w_{ij} \cdot x_j && \text{(standard linear layer)} \\ &= \sum_j (w_{ij} \cdot s_j) \cdot \left(\frac{x_j}{s_j}\right) && \text{(insert scaling factor } s_j \text{)} \\ &= \sum_j \tilde{w}_{ij} \cdot \tilde{x}_j && \text{(scaled weight and activation)} \end{aligned}$$

where:

• $y_i$: the output value for neuron $i$, which remains mathematically unchanged
• $w_{ij}$: the original weight connecting input channel $j$ to output $i$
• $x_j$: the input activation at channel $j$
• $s_j$: the per-channel scaling factor applied to weights in column $j$
• $\tilde{w}_{ij} = w_{ij} \cdot s_j$: the scaled weight that will be quantized
• $\tilde{x}_j = x_j / s_j$: the inverse-scaled activation, which compensates for the weight scaling

The key insight is that we want $s_j$ to be larger for channels with larger typical activations, protecting those salient weights with more precision in the quantized representation. Channels with large activations have their weights scaled up, giving those weights more quantization levels to stand for fine-grained distinctions. Channels with small activations have their weights left mostly unchanged or scaled down slightly. The net effect is a redistribution of quantization precision from where it is wasted (low-activation channels where quantization errors barely affect outputs) to where it matters most (high-activation channels where the same error level would cause significant output distortion).

The inverse scaling $1/s_j$ applied to $\tilde{x}_j$ must be applied somewhere to preserve the original computation. In transformer models, AWQ handles this elegantly by fusing the inverse scaling into the preceding layer's normalization. Since LayerNorm and RMSNorm layers output directly to the input channels of the subsequent linear layer, multiplying the normalization scale by $1/s_j$ for each output channel reaches the needed compensation without adding any runtime overhead. At inference time, the model runs exactly as fast as a standard INT4 model with no extra computations.