Boundary-Aware Quantization: Finite-Scale Decision Geometry of Neural Classifiers

Post-training quantization (PTQ) converts a trained full-precision model into low-bit weights without task-level retraining, while quantization-aware training (QAT) incorporates quantization into the training loop. Although PTQ is efficient and often accurate at moderate bitwidths, it can fail sharply at aggressive bitwidths; QAT is more expensive but can often recover the lost accuracy. We propose a unified geometric framework that explains both PTQ failure and QAT recovery. We model full-precision training as following a low-loss \emph{river} inside a wider \emph{valley}: a normal neighborhood of the river forms a nearly flat \emph{basin}, while leaving this basin incurs a sharp loss increase. When the quantization grid is comparable to the basin width, local PTQ objectives, including rounding and Hessian-based second-order reconstruction, can select a high-loss deployed quantized point outside the basin even when nearby low-loss quantized points exist. In this regime, straight-through-estimator-based QAT has a useful bias: it evaluates gradients at the deployed quantized weights while updating latent full-precision weights, causing the gradient to sense the valley wall and acquire an inward component that steers subsequent quantized iterates back into the basin. We formalize this mechanism through a local landscape model, construct a geometric PTQ failure mode, and prove finite-time QAT recovery under local quantizer-compatibility assumptions. Experiments across vision and language models under multiple neural-network quantization schemes corroborate the predicted basin-crossing failure of PTQ and the corresponding recovery mechanism of QAT.

Quantization is known to hurt below four bits, but nobody can say which of a model's decisions will change at a given bit-width. This matters most where a model acts rather than answers: a compressed agent stops calling its tools and, one bit lower, loses roughly half its safety refusals, while benchmark scores barely move. Prior work assumes the added noise has a roughly fixed size, which would make confident decisions safe. We measure the decision instead: the margin, the picked option's score minus its best alternative's, tracked before and after quantization across 16 models from 8 families under round-to-nearest, seven under AWQ, two under GPTQ and one under GGUF, at 8 down to 2 bits. The damage is proportional, not fixed in size: the margin is multiplied by a factor that collapses with bit-width (median 0.86 at 4 bits, 0.33 at 3, 0.00 at 2), which we call margin shrinkage. Contraction removes the protection a large margin affords; the model's own biases pick the direction: at 3 bits the decision to call a tool collapses toward inaction while the choice of which tool is untouched. No additive account, including one whose noise grows with the margin, wins a single damaged whether-to-call or safety cell (378 of 378). Given a condition's own constants the relation predicts held-out flip rates to a median 1.7 points, calibrated per decision (error 0.004 over 161,744 predictions), no flip used in the fit. Borrowed constants are wrong by 18-33 points at 3 bits, so the paired margin set has to be measured per model and bit-width: it locates breaking decisions without replacing measurement. At 4 bits the measurement is anchored to behaviour (the most likely token over the whole vocabulary is one of the item's two options in 85% of tool items); we treat the 2-bit floor as where the instrument stops measuring. No label-free repair we tested recovers more than one more bit does.

Quantization-aware training (QAT) is widely deployed but typically relies on the Straight-Through Estimator (STE), which passes gradients through non-differentiable quantizers by fiat. This often makes training brittle near bin boundaries and weakly aligned with the actual behavior of the low-precision model. We introduce JacQuant, a QAT framework that learns a lightweight surrogate of the model's local sensitivity to parameter changes and uses it to stabilize and accelerate training within standard variance-reduced optimizers. The surrogate is inexpensive (diagonal or block-diagonal), data-driven, and compatible with common weight and activation quantizers. On code-preserving training phases, we prove convergence for non-convex objectives and obtain linear rates under a PL condition, and we relate the learned sensitivity to end-to-end output fidelity via a simple calibration argument. Across LLM benchmarks at $\leq 2$ bits, JacQuant consistently reaches higher accuracy than STE-based QAT, and the runtime analyses on various models show that the added cost remains negligible under practical group sizes. The method is drop-in and requires no changes to the forward quantizers; our empirical claims are scoped to ultra-low-bit LLM QAT.