🤖 machine learning

JacQuant: STE-Free Quantization-Aware Training via Learned Jacobian Surrogates

JacQuant introduces a novel Quantization-Aware Training framework that replaces the brittle Straight-Through Estimator with a lightweight, learned Jacobian surrogate to stabilize training and achieve higher accuracy in ultra-low-bit LLM quantization without modifying forward quantizers or incurring significant runtime costs.

Original authors: Kai Yi, Vignesh Vivekraja, Harshit Khaitan, Steven Li

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kai Yi, Vignesh Vivekraja, Harshit Khaitan, Steven Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Broken Compass" in AI Training

Imagine you are trying to teach a robot to navigate a maze. To save space and energy, you decide to give the robot a very low-resolution map (this is called quantization). Instead of knowing the exact distance to a wall, the robot only knows if it's in "Zone A," "Zone B," or "Zone C."

The problem arises when the robot is right on the line between Zone A and Zone B.

The Reality: If the robot moves a tiny bit, it might stay in Zone A. If it moves a tiny bit the other way, it might jump to Zone B. The map is "blocky" and doesn't change smoothly.
The Old Method (STE): For years, when the robot tried to learn from its mistakes, the training algorithm used a shortcut called the Straight-Through Estimator (STE). It pretended the map was smooth and continuous. It told the robot, "You moved 1 inch, so the map changed by 1 inch."
The Result: This is like giving the robot a broken compass. When the robot is near a boundary, the compass spins wildly or points in the wrong direction because the map isn't actually smooth. This makes training unstable, especially when the map is very low-resolution (like 1 or 2 bits).

The Solution: JacQuant (The "Smart Sensor")

The authors of this paper, working at Meta AI, introduced JacQuant. Instead of pretending the map is smooth, they teach the robot a lightweight sensor that understands exactly how the blocky map reacts to movement.

Here is how it works, step-by-step:

1. The "Surrogate" Sensor

Imagine the robot has a small, cheap sensor attached to its foot. This sensor doesn't change the map; the map is still blocky. But, this sensor measures: "If I wiggle my foot a tiny bit, does the zone actually change?"

If the robot is in the middle of a zone, the sensor says, "Yes, moving a little changes things." (Sensitivity = 1).
If the robot is stuck against a wall (saturation) or right on a boundary where moving doesn't change the zone, the sensor says, "No, moving a little does nothing." (Sensitivity = 0).

This sensor is called a Learned Jacobian Surrogate. It's a simple mathematical map that tells the training algorithm how "sensitive" the current location is.

2. Fixing the Compass

When the robot makes a mistake, the training algorithm looks at the sensor's reading before sending a correction signal.

Old Way (STE): "You made a mistake! Push hard in this direction!" (Even if the robot is stuck against a wall, it pushes hard, causing it to bounce around uselessly).
JacQuant Way: The sensor says, "Hey, you're stuck against a wall; moving won't help." So, the algorithm dampens the push. It says, "Okay, don't push so hard here; let's try a different direction."

This stops the robot from bouncing around wildly near the boundaries and helps it settle into the best path much faster.

3. The "Amortized" Cost (Why it's cheap)

You might think, "Measuring this sensitivity sounds expensive!"
The authors solved this by being lazy (in a smart way). They don't check the sensor every single second.

They check the sensor occasionally (like once every 100 steps).
They use that reading for a long time until the robot moves far enough to need a new check.
This is called amortization. It's like checking the weather forecast once in the morning and using that info all day, rather than stepping outside every 5 minutes. The cost is so low (less than 2% extra time) that it feels like free.

The Results: Smoother, Faster, Better

The paper tested this on Large Language Models (LLMs) like LLaMA and Qwen, specifically when squeezing them down to ultra-low precision (1 or 2 bits).

The Analogy: Imagine trying to fit a giant, high-definition painting into a tiny, pixelated frame. The old method (STE) made the image look blurry and shaky. JacQuant acts like a smart filter that knows exactly which pixels to keep sharp and which to smooth out, resulting in a much clearer picture.
The Evidence: In their tests, models trained with JacQuant:
- Learned faster (converged more quickly).
- Made fewer mistakes (lower "perplexity," which is a measure of how confused the AI is).
- Answered questions better (higher accuracy on reasoning tasks).
- Stayed stable (didn't crash or oscillate near the "walls" of the quantization zones).

Summary

JacQuant is a new way to train AI models that are compressed to be very small. Instead of blindly guessing how the compressed model learns (using the old "Straight-Through Estimator"), it learns a simple, cheap "sensitivity map." This map tells the AI exactly how much it can trust its own movements near the edges of its compressed world. The result is a more stable, accurate, and efficient training process for the smallest, most memory-friendly AI models.

Technical Summary: JacQuant

Problem Statement

Quantization-aware training (QAT) is a standard approach for deploying Large Language Models (LLMs) under tight memory and latency constraints, particularly in ultra-low-bit regimes (e.g., $\le 2$ -bit weights). However, modern QAT pipelines almost universally rely on the Straight-Through Estimator (STE). The STE approximates the non-differentiable quantizer by passing gradients through as if the quantizer were the identity function.

The paper identifies a critical flaw in this approach: hard quantizers are piecewise constant with bin boundaries and saturation/clipping regions. In ultra-low-bit settings, many parameters lie near these boundaries or hit clipping limits. The STE's assumption of an identity Jacobian creates a persistent mismatch, causing gradients to be over-propagated in saturated regions where the quantizer is locally insensitive. This leads to brittle training, instability near bin boundaries, and solutions that are weakly aligned with the actual behavior of the low-precision model.

Methodology: JacQuant

JacQuant proposes a framework that replaces the fixed identity Jacobian of the STE with a learned, lightweight surrogate Jacobian, denoted as $B(W)$ . The core philosophy is to treat QAT as a "missing-Jacobian" problem: while the forward pass retains the deployed hard quantizer $Q_\Delta(\cdot)$ , the backward pass learns how perturbations in full-precision weights affect their quantized proxies.

Core Mechanism

Instead of the standard STE update $\nabla_W \ell \approx \nabla_{\hat{W}} \ell$ , JacQuant uses:
$\nabla_W \ell \approx B(W) \nabla_{\hat{W}} \ell$
where $\hat{W} = Q_\Delta(W)$ and $B(W)$ is a learned sensitivity map.

Key Design Choices

Parameterization: $B(W)$ is constrained to be diagonal or block-diagonal, with one scalar gain per quantization group (e.g., 128 weights). This keeps the memory overhead negligible (adding only ~1.5% to Adam optimizer states).
Estimation Strategies: The paper introduces two practical, data-driven estimators to update $B(W)$ $B (W)$ :
- JacQuant-Probe: Uses Gaussian random probes $\delta_g$ to measure the local slope $\Delta q_g = Q_\Delta(W_g + \delta_g) - Q_\Delta(W_g)$ . A least-squares fit estimates the local sensitivity.
- JacQuant-Dither: Uses subtractive dithering (adding uniform noise) to create a differentiable proxy for the quantizer, estimating sensitivity based on the expected response.
Amortized Updates: To maintain training throughput, $B(W)$ is not updated at every step. It is refreshed occasionally (e.g., every 50–200 steps) using an Exponential Moving Average (EMA). Between refreshes, the current $B(W)$ is reused.
Variance Reduction: The framework integrates with variance-reduced optimizers (e.g., loopless-SVRG, SAGA) to stabilize the stochasticity introduced by the probing/dithering process.

Theoretical Contributions

The paper provides a convergence analysis for JacQuant under non-convex objectives, structured around four key arguments:

Target Gradient Definition: The authors define a "mean-field" target gradient $g^\dagger(W) = J(W)\bar{v}$ , where $J(W)$ is the true sensitivity of the dither-averaged quantizer. They prove that the bias of JacQuant is proportional to $\|B(W) - J(W)\|$ , whereas STE's bias is proportional to $\|I - J(W)\|$ .
Consistency of Estimators: They demonstrate that both the Probe and Dither estimators are consistent with the target sensitivity $J(W)$ . Specifically, the Probe update is a local regression that converges to the mean-field sensitivity, and the Dither update has $J(W)$ as its unique population fixed point.
Stability: The learned $B(W)$ acts as an adaptive gain controller. It remains near identity in responsive bin interiors but dampens gradients (gain $< 1$ ) in saturated regions where the quantizer is insensitive. This prevents the oscillations common in STE-based training.
Convergence Rates: Under a Polyak-Łojasiewicz (PŁ) condition and assuming the "code-preserving" property (quantization codes do not change frequently within a local window), the paper proves that JacQuant achieves linear convergence rates. The convergence error is bounded by the Jacobian tracking error and the stochastic variance floor.

Experimental Results

The authors evaluate JacQuant on several LLM families (LLaMA3-1B/3B/8B, Qwen3-1.7B) in ultra-low-bit regimes ( $W \le 2$ , $A8/A16$ ).

Performance: JacQuant consistently outperforms STE-based baselines (including ParetoQ and WinQ) and the "Rotation Trick" (a fixed geometric transform method).
- On LLaMA3-1B (W2A16), JacQuant improved perplexity from 13.5 to 12.3 and accuracy from 55.7% to 56.6% over the best STE-based baseline.
- On LLaMA3-8B (W2A16), JacQuant showed monotonic improvements over ParetoQ across 30K, 60K, and 90K training steps.
Efficiency: The runtime overhead is negligible. On LLaMA3-1B, overhead ranges from 1–5% depending on group size; on LLaMA3-8B, it remains under 1.6% for practical group sizes.
Mechanism Validation: Finite-difference diagnostics confirm that JacQuant learns group gains $< 1$ in clipped/saturated regions, effectively damping gradients where STE would over-propagate them.

Significance and Claims

The paper positions JacQuant as a drop-in upgrade for existing QAT pipelines. Its significance lies in:

Ste-Free Training: It removes the reliance on the heuristic identity Jacobian without requiring changes to the forward quantizer or the model architecture.
Geometry Awareness: By learning the local sensitivity, it naturally adapts to the geometry of the quantization landscape, specifically addressing the saturation and clipping issues that plague ultra-low-bit training.
Practicality: It achieves these gains with minimal computational cost, making it viable for large-scale LLM training.

The authors explicitly scope their empirical claims to ultra-low-bit LLM QAT ( $\le 2$ bits). They acknowledge that while the method is a drop-in replacement for the backward rule, the forward quantizer remains a hard, deployed operator. The paper does not claim to solve all quantization challenges but specifically targets the gradient mismatch inherent in STE-based optimization. Future work directions mentioned include extending the method to mixed-precision policies, non-uniform quantizers, and refining the estimation scheduling.