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

arXiv: https://arxiv.org/abs/2605.25469 · HTML: https://mlsanity.com/paper/2605.25469/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 45 references that are in the map.
- First posted on arXiv 25 May 2026.

## Builds on (most cited first, 30 of 39)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2211.10438/) (2022)
- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://mlsanity.com/paper/2306.00978/) (2023)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Learned Step Size Quantization](https://mlsanity.com/paper/1902.08153/) (2019)
- [Up or Down? Adaptive Rounding for Post-Training Quantization](https://mlsanity.com/paper/2004.10568/) (2020)
- [DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients](https://mlsanity.com/paper/1606.06160/) (2016)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers](https://mlsanity.com/paper/2206.01861/) (2022)
- [PACT: Parameterized Clipping Activation for Quantized Neural Networks](https://mlsanity.com/paper/1805.06085/) (2018)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [QuIP: 2-Bit Quantization of Large Language Models With Guarantees](https://mlsanity.com/paper/2307.13304/) (2023)
- [QuIP#: Even Better LLM Quantization with Hadamard Incoherence and Lattice Codebooks](https://mlsanity.com/paper/2402.04396/) (2024)
- [Momentum-Based Variance Reduction in Non-Convex SGD](https://mlsanity.com/paper/1905.10018/) (2019)
- [Extreme Compression of Large Language Models via Additive Quantization](https://mlsanity.com/paper/2401.06118/) (2024)
- [Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets](https://mlsanity.com/paper/1903.05662/) (2019)
- [SARAH: A Novel Method for Machine Learning Problems Using Stochastic Recursive Gradient](https://mlsanity.com/paper/1703.00102/) (2017)
- [LSQ+: Improving low-bit quantization through learnable offsets and better initialization](https://mlsanity.com/paper/2004.09576/) (2020)
- [Soft-to-Hard Vector Quantization for End-to-End Learning Compressible Representations](https://mlsanity.com/paper/1704.00648/) (2017)
- [ZeroQuant-V2: Exploring Post-training Quantization in LLMs from Comprehensive Study to Low Rank Compensation](https://mlsanity.com/paper/2303.08302/) (2023)
- [PV-Tuning: Beyond Straight-Through Estimation for Extreme LLM Compression](https://mlsanity.com/paper/2405.14852/) (2024)
- [VPTQ: Extreme Low-bit Vector Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2409.17066/) (2024)

## Built on by (most cited first, 2 of 2)

- [Quantization-Aware Pre-Training with Constrained Empirical Weight Distribution](https://mlsanity.com/paper/2609.32659/) (2026)
- [Few Bits, One Law: Toward W2A4KV2](https://mlsanity.com/paper/2610.09202/) (2026)

## Abstract (arXiv)

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 ≤ 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.