QAT and BitNet: A Deep Dive into 1-bit and 1.58-bit Quantization-Aware Training

QAT: Putting Quantization into the Training Loop

From PTQ to QAT

The previous article discussed post-training quantization (PTQ): the weights are already fixed, and we just need to find the best integer grid to compress them. But what if we make the model aware of quantization during training? That's quantization-aware training (QAT). The core idea of QAT is simple: quantize the weights in the forward pass, compute the loss using the quantized values, and then let the optimizer adjust the weights so that the model still performs well "despite the rounding."

Why QAT Works: The Geometry of the Loss Surface

Standard training finds a point with the lowest loss and then stops caring about the shape of the loss surface around that point. But quantization breaks this assumption: each weight now falls within a small error range δ around its position at the end of training. What really matters is not only how low the loss is at the minimum, but also how fast the loss rises as you move away from the minimum.

Near the minimum the slope (first derivative) is zero, and the increase in loss is approximately ΔL ≈ ½ · H · δ², where H is the curvature (second derivative). Two minima with the same loss can behave very differently:

• Sharp minimum (large H): a small rounding error causes a large increase in loss.
• Flat minimum (small H): the same rounding error barely matters.

So QAT optimizes two things at once: low loss, and a gentle curvature around the low-loss point, keeping quantization error manageable. The original article provides an interactive page where you can drag a quantization-error slider and see that the loss rises 8 times faster on a sharp minimum than on a flat one.

Curvature and generalization: The curvature of the loss surface (the eigenvalues of the Hessian matrix) not only determines quantization robustness but is also closely related to the model's generalization ability. A sharp minimum usually means the model is more sensitive to both input and parameter perturbations, whereas a flat minimum corresponds to a wider region of optimal solutions. By explicitly introducing quantization noise, QAT is effectively performing a kind of regularization during training, forcing the model to converge to flatter regions. This is similar in spirit to methods like Sharpness-Aware Minimization (SAM), except that QAT's perturbation is structured and determined by the quantization grid rather than being random noise.

How Gradients Flow: Fake Quantization and STE

The rounding operation has a gradient of zero almost everywhere, so backpropagation would learn nothing. The solution is the straight-through estimator (STE): the forward pass uses the quantized values, while the backward pass pretends that rounding is the identity function and passes 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 the rounded weights, but the optimizer updates the underlying full-precision weights. This idea is exactly the cornerstone of BitNet. In practice, QAT usually outperforms PTQ, especially at low bit widths such as 4-bit.

The mathematical nature of STE and its variants: The straight-through estimator (STE) essentially replaces the gradient of the rounding function with the gradient of the identity function, i.e., ∂L/∂w ≈ ∂L/∂w_q. This approximation is not mathematically exact, but it works remarkably well in practice. Researchers have proposed various improved variants, such as attenuating the gradient when |w| is near a quantization boundary, using soft rounding, or using learnable step functions to mitigate the gradient mismatch. The success of STE shows that deep learning optimization has some tolerance for gradient inaccuracy, as long as the gradient direction is roughly correct.

BitNet and 1-bit Weights

Where Linear Layers Sit in a Transformer

Linear layers are the most basic building block of neural networks: an input vector is multiplied by a weight matrix to produce an output vector. Almost all parameters in a model are elements of such weight matrices. In a Transformer, each block contains two parts:

• The attention mechanism has four linear layers: the Q, K, and V projections plus the output projection. The softmax(QKᵀ)V in the middle multiplies activations and stores no weights, so it is not a linear layer.
• The feed-forward network has two linear layers: an up-projection and a down-projection, with a GELU nonlinearity in between.

BitNet targets exactly these linear layers: compress them and you compress the entire model.

What Does BitLinear Change?

BitNet replaces most Linear layers with BitLinear. The task is the same (output = input × weights), but the numbers involved are much smaller:

	Standard Linear	BitLinear
Weights	FP16 (16 bits per weight)	±1 (1 bit per weight)
Activations	FP16	INT8
Matrix multiplication	FP16 multiply-add	INT8 input against ±1 weights

Multiplying by +1 or -1 just keeps the number or flips its sign, so the matrix multiplication collapses into addition and subtraction. Multiplication is one of the most expensive operations on a chip, so this saves both time and energy as well as memory.

But quantizing both sides loses scale information: ±1 weights no longer say anything about the original magnitude. So BitLinear keeps two scaling factors, one for the weights and one for the activations, and stretches the result back to the correct size at the end.

The BitLinear Recipe

1. Quantize the input: from FP16 to INT8, using absmax scaling (after LayerNorm, activations have a stable scale).
2. Binarize the weights: take the mean α of the weights; values above α become +1 and values below become -1 (centering first so that +1 and -1 are roughly balanced).
3. Matrix multiplication: integer inputs multiplied by ±1 weights, mostly additions, far cheaper than multiplications.
4. Dequantize the output: use the activation scaling factor and the weight scaling factor β (the mean absolute value of the weights) to restore the result to the correct magnitude.

This is a stack of two quantizers: one for the activations and one for the weights.

Why Not Just Store 1 Bit?

If training is just learning which weights are "on" and which are "off," why keep anything larger than 1 bit?

Because of gradients. The binarization function is sign(), whose gradient is zero everywhere and undefined at zero. If the weights really were just bits, gradient descent would have nothing to update.

So the model keeps high-precision latent weights (the paper trains with FP16/BF16). The training loop is as follows:

1. Maintain full-precision latent weights.
2. Forward pass: binarize with sign and use the ±1 version.
3. Backward pass: use STE so that gradients land on the latent weights.
4. Update the latent weights.

The latent weights cannot be INT8: the learning rate is small, and a tiny gradient step on INT8 would round back to the same integer, making the update vanish entirely. The latent weights need enough resolution to accumulate many tiny pushes until the sign flips at some point. At inference time only the ±1 version is needed; the latent weights are training-time scaffolding.

The resolution requirement for latent weights: The latent weights need to retain enough precision, and the core reason is that the magnitude of a gradient update is usually far smaller than the quantization step. Taking a learning rate of 1e-4 and a weight scale of 0.01 as an example, a single update is about 1e-6, while the smallest positive normal number in FP16 is about 6e-8, so it can still accumulate; with INT8 (step size about 4e-3), every update is rounded away to zero. This is why the BitNet paper chooses to maintain the latent weights in BF16/FP16, ensuring that tiny updates can accumulate gradually and eventually cross the threshold for a sign flip. This also explains why the cost of QAT training is comparable to full-precision training—the computation graph still keeps high-precision weights.

"Is 1-bit a Scam?"

Some people are skeptical of the claim that BitNet gets closer to FP16 quality as models grow larger: doesn't that mean you just need a bigger model, and the savings disappear?

Not at all. The savings are per parameter. A 1-bit model may need more parameters to match quality, but each parameter is up to 16 times smaller than FP16. It only fails if you need more than 16 times as many parameters. The paper's trend is that the quality gap shrinks as model size grows.

1.58 Bits: The Rise of Ternary Weights

A follow-up paper, The Era of 1-bit LLMs (BitNet b1.58), found that pure 1-bit requires a very large scale to match full-precision quality. The authors added one more state: zero. Weights become either -1, 0, or +1.

Why Zero Helps

Zero lets a weight say "this connection is not important" instead of being forced to push the output up or down. The weight matrix becomes more expressive at almost no additional cost.

The Ternary Weight Recipe

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

Equivalently: |w| <= γ/2 becomes 0; w > γ/2 becomes +1; w < -γ/2 becomes -1. This is "absmean" quantization. Activations switch to symmetric absmax quantization (scaled around zero) instead of the previous 0-to-maximum range. The training setup (including latent weights and STE) remains unchanged.

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

A Five-Weight Example

Take w = [0.9, -0.3, 0.05, -1.2, 0.4]. The 1-bit version (β=0.57) forces the tiny weight 0.05 to +0.57, an error of 0.52; the ternary version (γ=0.57) turns it into 0, an error of only 0.05. The other four weights are unchanged.

Why Is It Called "1.58" Bits?

Three states require log2(3) ≈ 1.58 bits of information. But hardware cannot store 1.58 bits, so practical implementations pack the values. The simplest packing uses 2 bits per weight; a more compact scheme fits five ternary values into one byte (3⁵=243 ≤ 256), i.e., 1.6 bits per weight. So 1.58 is the information-theoretic number, and actual storage is slightly above it. Compared with FP16, each parameter is roughly 10 times smaller.

Key Takeaways

• PTQ rounds a finished model; QAT includes rounding in the training loop so that the model finds weights that can withstand rounding.
• The principle is geometric: prefer flat minima (low curvature) so that quantization error is cheap.
• Rounding has no gradient, so we use STE and keep high-precision latent weights to accumulate tiny updates.
• BitNet does this with sign(); b1.58 adds a zero state and uses absmean scaling for a better quality/size trade-off.