Training for quantization: QAT, 1-bit BitNet and 1.58-bit weights

In the last post we quantized a model that was already trained (post-training quantization, PTQ). The weights were fixed, and we looked for the best way to squeeze them onto an integer grid.

But what if the model knew it was going to be quantized while it was learning? That is quantization-aware training (QAT). QAT is what eventually gives you BitNet, where every weight is -1, 0 or +1. Let's start with the basics of QAT.

Part 1: QAT

The idea

PTQ: you train and retain the weights in full precision, then quantize them and deliver the quantized model for inference.

QAT: you put the quantization of weights inside the training loop. On every forward pass the weights are quantized, and the loss is computed with the quantized values. The optimizer then adjusts the weights so that the model does well despite the rounding.

How does the training change to account for quantization?

Normal training finds a point where the loss is low, and it stops caring about the shape of the loss around that point. What I mean by shape of the loss, is that if you change the weights from that minima, how quickly does the loss functions value rise again. That is fine with standard training practices, because once you publish the model the weights never move.

Quantization breaks that assumption. Each weight now lands somewhere within a small error range δ of where training left it. This is the diff in weights that can come due to Quantization. So what matters is not just how low the loss is at the minimum, but how fast the loss climbs as you move away from it.

Near a minimum the slope is zero, so the loss increase from a nudge of size δ is approximately

ΔL ≈ ½ · H · δ²

where H is the curvature (second derivative, in simple terms how quickly the loss changes with a change in weight). Two minima with the same loss can behave very differently:

• Sharp minimum (large H): a small rounding error costs a lot of loss.
• Wide minimum (small H): the same rounding error barely matters.

A note on terms, because I kept mixing them up: slope is the first derivative, and it is zero at both minima. Curvature is the second derivative, and it is what differs.

Try it yourself to build up your intuition for this, focus on how the loss changes with the change in weight from the point of minima for both wide and sharp minimas: slope and curvature playground. Drag the quantization error and watch the loss rise 8x faster on the sharp minimum.

So QAT optimizes two things at once: a low loss, and a gentle curvature around it so quantization error stays in a controlled range.

How the gradient flows: fake quantization and the STE

There is a catch. Rounding has a gradient of zero almost everywhere, so backprop would learn nothing through it.

The fix is the straight-through estimator (STE): in the forward pass, use the quantized value. In the backward pass, pretend the rounding was the identity function and pass the gradient straight through.

import torch

def fake_quant(w, bits=4):
    qmax = 2 ** (bits - 1) - 1
    scale = w.abs().max() / qmax
    w_q = torch.round(w / scale).clamp(-qmax, qmax) * scale
    # forward: w_q. backward: gradient of w (identity)
    return w + (w_q - w).detach()

The model computes with rounded weights, but the optimizer updates the underlying full-precision weights. Hold on to that idea, because BitNet pushes it to the limit.

In practice QAT tends to beat PTQ, most visibly at low bit-widths such as 4-bit.

Part 2: BitNet and 1-bit weights

Paper: BitNet: Scaling 1-bit Transformers for Large Language Models.

First, what is a linear layer?

A linear layer is the most basic building block in a neural network: it takes a vector of numbers in, multiplies it by a matrix of learned numbers (the weights), and gives a new vector out.

output = input × weights   (+ a bias, which I'll ignore here)

Almost all of a model's parameters are entries in weight matrices like this. Two words I'll keep using:

• Weights are the stored, learned numbers. They are fixed once training is done.
• Activations are the numbers flowing through the model for a particular input. They change with every token.

In a normal model, both are FP16.

Where linear layers sit in a transformer

A transformer is a stack of identical blocks. Each block has two parts, and almost everything with weights in it is a linear layer:

• Attention has four: the Q, K and V projections, and an output projection. The step in the middle, softmax(QKᵀ)V, multiplies two sets of activations together. It has no stored weights, so it isn't a linear layer.
• The feed-forward network has two: an up projection and a down projection, with a non-linearity (GELU) between them.

Watch this if you want a quick crash course on Transformer Internals!

That's why BitNet goes after the linear layers: shrink those, and you've shrunk the model.

So what does BitLinear change?

BitNet replaces most of the Linear with a BitLinear. The job is the same (output = input × weights), but the numbers going in are much smaller:

	Standard Linear	BitLinear
Weights	FP16 (16 bits each)	±1 (1 bit each)
Activations	FP16	INT8
The matrix multiply	FP16 multiply-adds	INT8 inputs against ±1 weights

Multiplying by +1 or -1 just keeps a number or flips its sign, so the matrix multiply collapses into additions and subtractions. Multiplication is one of the most expensive things a chip does, so this saves time and energy as well as memory.

Quantizing both sides loses scale information, though: a ±1 weight no longer says how big the original was. So BitLinear keeps two scale factors, one for the weights and one for the activations, and uses them to stretch the result back to the right size at the end. The recipe below shows how.

The BitLinear recipe

1. Quantize the input from FP16 to INT8 using absmax scaling (after a LayerNorm, so the activations have a steady scale to quantize).
2. Binarize the weights. Take the mean of the weights, α. Anything above it becomes +1, anything below becomes -1. (Centring first keeps the +1s and -1s roughly balanced.)
3. Matrix multiply with the integer inputs and ±1 weights. This is mostly additions, which is much cheaper than multiplications.
4. Dequantize the output using the activation scale and a weight scale β (the mean absolute weight), so the result lands back at the right magnitude.

It is two quantizers layered together: one on the activations, one on the weights.

My question: why not just store 1 bit?

If training is learning which weights should be "on" and which "off", why keep anything bigger than a bit?

Because of the gradient. The binarizing function is sign(). Its gradient is zero everywhere, except at zero where it is undefined. If the weights were just bits, gradient descent would have nothing to update.

So, as with QAT, the model keeps latent weights in high precision (the paper trains with high-precision weights, FP16/BF16 in practice). The training loop looks like this:

1. Keep a full-precision latent weight.
2. Forward pass: binarize it with sign, use the ±1 version.
3. Backward pass: use the STE, so the gradient lands on the latent weight.
4. Update the latent weight.

I first assumed the latent weights were INT8. That doesn't work, and it is a good illustration of why. Learning rates are small. A tiny gradient step on an INT8 weight would round back to the same integer, and the update would simply vanish. The latent weights need enough resolution to accumulate many small nudges until one finally flips the sign.

Only the ±1 version is needed at inference. The latent weights are a training-time scaffold.

Try it: latent weights playground. Press "Train 10 steps": the 1-bit-only weight never moves, while the latent weight quietly accumulates the updates until it crosses zero and flips.

"So is 1-bit a sham?"

I was suspicious of a claim that BitNet gets closer to FP16 quality as models grow. Doesn't that mean you just need a bigger model, and the savings vanish?

No. The saving is per parameter. A 1-bit model may need more parameters to match quality, but each parameter is up to 16x smaller than FP16. It only fails if you need more than 16x as many parameters. The trend in the paper is that the quality gap shrinks as the model scales.

Part 3: 1.58 bits

A follow-up paper, The Era of 1-bit LLMs (BitNet b1.58), found that pure 1-bit needs a lot of scale to match full-precision quality. The authors added one more state: zero.

Weights are now -1, 0 or +1.

Why zero helps

A zero lets a weight say "this connection doesn't matter", instead of being forced to push the output up or down. The weight matrix becomes more expressive at almost no extra cost.

The recipe

1. Compute γ, the mean of the absolute weights.
2. Divide each weight by γ and round to the nearest of -1, 0, +1.

Equivalently:

• |w| <= γ/2 becomes 0
• w > γ/2 becomes +1
• w < -γ/2 becomes -1

This is "absmean" quantization. The activations switch to symmetric absmax quantization (scaled around zero), rather than the 0-to-max range used before. Everything else about the training setup, including latent weights and the STE, stays the same.

def ternary_weights(w, eps=1e-5):
    gamma = w.abs().mean()
    return (w / (gamma + eps)).round().clamp(-1, 1), gamma

The five-weight example

Take w = [0.9, -0.3, 0.05, -1.2, 0.4].

	w0	w1	w2	w3	w4
original	0.90	-0.30	0.05	-1.20	0.40
1-bit (β = 0.57)	+0.57	-0.57	+0.57	-0.57	+0.57
ternary (γ = 0.57)	+0.57	-0.57	0	-0.57	+0.57

The tiny weight 0.05 is forced to ±0.57 in the 1-bit version, an error of 0.52. With a zero available, its error drops to 0.05. The other four are unchanged.

A bit of etymology : Why "1.58" bits?

Three states need log2(3) ≈ 1.58 bits of information. But hardware cannot store 1.58 bits, so real implementations pack the values. The simplest packing uses 2 bits per weight. A tighter one fits 5 ternary values into one byte (3⁵ = 243 ≤ 256), which is 1.6 bits per weight. So 1.58 is the information-theoretic figure, and practical storage lands slightly above it.

Versus FP16, that is roughly 10x smaller per parameter.

What to take away

• PTQ rounds a finished model. QAT trains with the rounding in the loop, so the model finds weights that survive it.
• The reason it works is geometry: prefer wide minima (low curvature) so quantization error costs little.
• Rounding has no gradient, so we use the STE and keep high-precision latent weights that accumulate small updates.
• BitNet applies this with sign(); b1.58 adds a zero, using absmean scaling, for a better quality/size trade-off.

Interactive pages: slope and curvature, latent weights.