NUPES: Non-Uniform Quantization
via Power Exponent Search

Chapter 1: The Key Insight

If the grid is in the wrong place, there are two ways to fix it. The first is to move the grid — bunch the levels near zero where the weights live, and spread them out in the tails where weights are rare. That is non-uniform quantization.

The naive way to do this is a lookup table (codebook): you store 16 arbitrary float levels and snap each weight to the nearest one. It fits the distribution perfectly — but it breaks the math. A multiply w × x on a hardware integer unit assumes the integer codes are scalars you can multiply. If your codes are arbitrary table indices, every multiply needs a lookup, which kills the speed you bought.

That is the PowerQuant idea, from the same authors' earlier paper. It already beat everything else on data-free quantization. But it had two stubborn weaknesses, and fixing them is what NUPES is:

So in one breath: NUPES = PowerQuant's bendable grid + the ability to relearn every quantized weight + the ability to learn the bend itself by gradient descent, all numerically stabilized. We will build each piece from the ground up.

Chapter 2: Uniform Quantization — and its two pitfalls

Before we bend anything, let us nail down the baseline so the bending makes sense. Uniform quantization takes a real number x, divides by a scale s, rounds to the nearest integer, and clips into the representable range. For b bits the range is the integers from −2^b−1 to 2^b−1−1 — for 4 bits, that is −8 to 7, sixteen levels.

Read every symbol. X is the whole weight tensor. s(X) is the scale — the size of one quantization step in real units. It is set so the largest-magnitude weight lands exactly on the top integer (here, 7). round(·) snaps to the nearest integer; Q(x) is the resulting integer code. To recover an approximate float you go back: Q⁻¹(Q(x)) = s(X) · Q(x). It will not equal the original x — that gap is the quantization error.

Worked example — quantize three real weights to int4

Say a tensor's largest magnitude is max|x| = 2.0. Then s = 2.0 / 7 ≈ 0.2857. Now quantize:

Look at row two. The weight 0.050 — a perfectly real, common, near-zero value — gets crushed to the integer 0. It vanishes. With only ~3 levels covering the dense spike near zero, tiny weights are systematically annihilated. That is pitfall number one.

Pitfall 2 — outliers stretch the grid

Now suppose a single weight is 17.0 while everything else lives in [−1, 1]. The scale becomes s = 17 / 7 ≈ 2.43. Every weight in [−1.2, 1.2] now rounds to 0 — because one level is now 2.43 wide. The outlier hijacks the entire grid, and the rest of the layer collapses to near-ternary (just −1, 0, +1). Use the slider in Chapter 0 again and you will see this happen visually: push the outlier out and the grid ticks spread until the central spike falls between two ticks.

Chapter 3: PowerQuant — bending the grid

We want a grid whose steps are small near zero (where weights crowd) and large far out (where weights are rare). The trick: don't change the quantizer — change the input to it. Pass each weight through a warping function t first, quantize uniformly in the warped space, and the levels, mapped back, come out non-uniformly spaced.

Why this transform and not any other? The authors imposed one hard requirement: the quantizer must preserve multiplication, because that is what neural network layers do. Formally they wanted t(x)·t(y) = t(xy). The functions satisfying that on the positive reals are exactly the power functions t(x)=x^a — they are the automorphisms of (ℝ⁺, ×). The sign(x) wrapper just extends it to negatives. So the power function is not a guess; it is the only family that keeps the algebra intact.

Worked example — the same near-zero weight, now with a = 0.5

Reuse the tensor from Chapter 2: max|x| = 2.0, but now warp with a = 0.5 (square root) before quantizing. First warp the range: t(2.0) = 2.0^0.5 = 1.414, so the warped scale is s = 1.414 / 7 = 0.2020.

Compare to Chapter 2. The near-zero weight 0.050 that vanished under uniform quantization (error 0.050) now lands on level 1 and survives with error 0.009 — five times better. We spent a level we used to waste in the tail to rescue the dense spike. That is the whole win of bending the grid.

# PowerQuant: the entire operator in 6 lines
def power_quant(x, a=0.5, bits=4):
    qmax = 2**(bits-1) - 1                 # 7 for int4
    t    = torch.sign(x) * x.abs()**a       # warp toward zero
    s    = t.abs().max() / qmax             # scale in warped space
    q    = torch.clamp(torch.round(t / s), -qmax-1, qmax)   # int code
    deq  = torch.sign(q) * (s * q.abs())**(1/a)  # unwarp on the way back
    return q, deq

Chapter 4: The Grid — see the bend live (showcase)

This is the heart of the paper made interactive. Below is a real bell-shaped weight tensor (with an adjustable outlier). The vertical ticks are the quantization levels — the only 2^b values a weight can become. Drag exponent a from 1 (uniform) down toward 0.3 and watch the levels migrate toward zero, crowding where the weights actually live. Drag bits to see how many levels you get. The red bar is the total reconstruction error — your job is to make it small.

Play with the outlier slider while watching the error bar. With a = 1, pushing the outlier out makes the error explode — every other weight collapses to zero. With a ≈ 0.5, the bent grid keeps near-zero weights alive even as the outlier stretches the range. The bend is what tames outliers. This is precisely why the method shines on LLMs (Chapter 7).

Chapter 5: Learn the Weights — not just round them

So far the grid is good. But two weights that round to the same integer still have different true values — and we have only used the weights, never checked how the layer's output changes. Calibration-based methods (GPTQ / AdaRound) fix that: feed a tiny batch through, compare the quantized layer's output to the original's, and nudge the quantization to match.

AdaRound's nudge is tiny by design. For each weight it learns one number ε ∈ [0,1] that decides whether to round the floored weight up or down — a one-level wiggle, nothing more.

Here f_l is layer l; X_fp the original (full-precision) input to that layer; X_q the quantized input from upstream; ⌊W/s⌋ the floored integer code; and σ(ε)∈[0,1] the learnable "round up by this much." The loss says: pick rounding so the quantized layer reproduces the float layer's output.

NUPES's move: relearn the integer outright

Instead of learning a fractional rounding offset, NUPES learns the whole quantized integer as a single free parameter ε ∈ ℤ — initialized at the current code, but allowed to move by many levels, up or down. A near-zero weight can jump several steps to compensate; a tail weight can be held still. To keep that integer differentiable, NUPES borrows the differentiable soft quantization (DSQ) function:

Don't panic at the symbols — the picture is simple. dsq is a soft staircase: a smooth, differentiable curve that hugs the integer steps. β is a steepness knob. Small β → gentle ramps you can slide a value across between integers (gradients flow, weights move freely). Large β → sharp steps that snap ε firmly onto an integer (so the final result is a true int). NUPES schedules β: start soft to let weights migrate, end steep to lock them onto the grid.

Chapter 6: Learn the Exponent — making the quantizer trainable

PowerQuant froze one a for the whole network. But Chapter 4 showed every layer has a different distribution — the first layer is flat, deep layers are peaky. They want different bends. NUPES's second contribution: make a a per-layer trainable parameter, learned by gradient descent alongside the weights. To do that we need the gradient of x^a with respect to a:

Clean on paper. A nightmare in practice. The authors hit three landmines, and the fix for each is a single line in their Algorithm 1. This is a beautiful "concept → realization" moment — the math is one line, the engineering is the contribution.

# NUPES Algorithm 1: gradient for the exponent a (the three fixes)
# scales are recomputed but DETACHED — no gradient flows through them
scale_x = scale(sign(X) * X.abs()**a).detach()      # fix 2
scale_w = scale(sign(W) * W.abs()**a).detach()      # fix 2

Xc = X.abs().clamp(min=1e-6)                        # fix 1: no log(0)
Wc = W.abs().clamp(min=1e-6)                        # fix 1

g_inputs  = (Xc**a * Xc.log() / scale_x).mean()       # fix 3: average each
g_weights = (Wc**a * Wc.log() / scale_w).mean()       # fix 3
grad_a    = g_inputs + g_weights                       # then combine

The payoff (Figure 6 in the paper): the learned exponents land in a ∈ [0.213, 0.774] across layers — vindicating the old hand-picked 0.5 — but now each layer gets its own value, and the first layer reliably wants a higher a (its distribution is flatter, closer to uniform). The quantizer is no longer a hyperparameter; it is learned.

Chapter 7: Outliers & LLMs — where the method earns its name

Everything above was warm-up for the case that actually matters in 2023: large language models. Their weight and activation distributions have a vicious feature ConvNets don't — massive outliers. The paper measures how far the most extreme value sits from the mean, in standard deviations:

Now recall Chapter 4: bending the grid toward zero with a < 1 packs many fine levels into that central bulk while still reaching the outlier with the coarse outer levels. The outlier is represented; the bulk is preserved. The power transform is, almost by accident, an outlier-handling mechanism — which is why NUPES is so strong exactly where uniform methods are weakest.

Chapter 8: Experiments — does it actually work?

The headline result (paper's Figure 1, Table 3): in the brutal W4/A4 setting — 4-bit weights and 4-bit activations — NUPES reaches near full-precision accuracy on both ConvNets and Transformers, where every prior method is strong on one and weak on the other. Toggle the chart below between architectures and watch the methods line up.

The ablation that proves each piece matters

NUPES has two switches: learn the weights, learn the exponent. The paper turns them on one at a time (Table 3):

Both knobs help; together they help most. The biggest single jump is learning the weights (the Chapter 5 contribution), but learning the exponent on top recovers the last few points — and on the hardest cases (RetinaNet went from PowerQuant's 2.24 to 33.1 mAP) it is the difference between a broken model and a working one.

Chapter 9: Connections & cheat sheet

NUPES sits at the intersection of three lineages: non-uniform quantization (codebooks, log-quant), gradient-based PTQ (AdaRound, BrecQ), and LLM-specific outlier handling (LLM.int8, QLoRA, GPTQ/OPTQ). Its distinctive bet is that a single learnable scalar per layer — the power exponent — plus the freedom to relearn each integer code, beats both the rigidity of fixed grids and the fragility of full codebooks, while staying integer-only.

Cheat sheet — every symbol in one place

Where to go next

• Lloyd-Max quantization (EE269 lesson) — the classical optimal-non-uniform-quantizer theory NUPES's grid approximates.
• Pruning & LLM deployment (TinyML lesson) — quantization's sibling compression techniques and where they sit in the deployment stack.
• LLM inference internals (lesson) — why memory bandwidth, not FLOPs, is what quantization actually buys you.