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

• URL: http://arxiv.org/abs/2605.25469v1
• Date: Mon, 25 May 2026 06:19:49 GMT
• Title: JacQuant: STE-Free Quantization-Aware Training via Learned Jacobian Surrogates
• Abstract summary: Quantization-aware training (QAT) is widely deployed but typically relies on the Straight-Through Estimator (STE), which passes through non-differentiable quantizers by fiat.<n>This often makes training aligned near bin and weakly with actual low-precision model.<n>We introduce JacQuant, a framework that learns a lightweight surrogate model's local sensitivity to parameter changes.
• Score: 9.03449876943648
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Sensitivity-Aware Thresholding and Token Routing for Activation Sparsification in Large Language Models [2.593510481959727]
  Efficient inference in Large Language Models (LLMs) requires deciding where computation can be reduced while preserving model quality.<n>We study this problem through multilayer perceptron (MLP) activation sparsification and token-level conditional routing.
  arXiv  Detail & Related papers  (2026-07-09T23:40:36Z)
• LC-QAT: Data-Efficient 2-Bit QAT for LLMs via Linear-Constrained Vector Quantization [9.37026660976194]
  Quantization-aware training (QAT) is essential for extremely low-bit large language models (LLMs)<n>We propose LC-QAT, a 2-bit weight-only VQ-QAT framework that represents quantized weights via a learned affine mapping over discrete vectors.<n>We show that LC-QAT consistently outperforms state-of-the-art QAT methods while using only 0.1%--10% of the training data.
  arXiv  Detail & Related papers  (2026-06-09T08:02:56Z)
• LSGQuant: Layer-Sensitivity Guided Quantization for One-Step Diffusion Real-World Video Super-Resolution [52.627063566555194]
  We introduce LSGQuant, a layer-sensitivity guided quantizing approach for one-step diffusion-based real-world VSR.<n>Our method incorporates a Dynamic Range Adaptive Quantizer (DRAQ) to fit video token activations.<n>Our method has nearly performance to origin model with full-precision and significantly exceeds existing quantization techniques.
  arXiv  Detail & Related papers  (2026-02-03T06:53:19Z)
• StableQAT: Stable Quantization-Aware Training at Ultra-Low Bitwidths [49.94623294999562]
  Quantization-aware training (QAT) is essential for deploying large models under strict memory and latency constraints.<n>Common approaches based on the straight-through estimator (STE) or soft quantizers often suffer from mismatch, instability, or high computational overhead.<n>We propose StableQAT, a unified and efficient QAT framework that stabilizes training in ultra low-bit settings.
  arXiv  Detail & Related papers  (2026-01-27T08:00:57Z)
• What Makes Low-Bit Quantization-Aware Training Work for Reasoning LLMs? A Systematic Study [59.44848132298657]
  Post-training quantization (PTQ) usually comes with the cost of large accuracy drops, especially for reasoning tasks under low-bit settings.<n>In this study, we present a systematic empirical study of quantization-aware training (QAT) for reasoning models.
  arXiv  Detail & Related papers  (2026-01-21T11:22:29Z)
• CAGE: Curvature-Aware Gradient Estimation For Accurate Quantization-Aware Training [73.46600457802693]
  We introduce a new method that counteracts the loss induced by quantization.<n>CAGE significantly improves upon the state-of-theart methods in terms of accuracy, for similar computational cost.<n>For QAT pre-training of Llama models, CAGE matches the accuracy achieved at 4-bits (W4A4) with the prior best method.
  arXiv  Detail & Related papers  (2025-10-21T16:33:57Z)
• End-to-End On-Device Quantization-Aware Training for LLMs at Inference Cost [53.25965863436039]
  Quantization-aware training (QAT) provides a more principled solution, but its reliance on backpropagation incurs prohibitive memory costs.<n>We propose ZeroQAT, a zeroth-order optimization-based QAT framework that supports both weight and activation quantization.<n>Experiments show that ZeroQAT consistently outperforms representative PTQ and QAT baselines while requiring significantly less memory.
  arXiv  Detail & Related papers  (2025-08-21T01:18:27Z)
• RoSTE: An Efficient Quantization-Aware Supervised Fine-Tuning Approach for Large Language Models [53.571195477043496]
  We propose an algorithm named Rotated Straight-Through-Estimator (RoSTE)<n>RoSTE combines quantization-aware supervised fine-tuning (QA-SFT) with an adaptive rotation strategy to reduce activation outliers.<n>Our findings reveal that the prediction error is directly proportional to the quantization error of the converged weights, which can be effectively managed through an optimized rotation configuration.
  arXiv  Detail & Related papers  (2025-02-13T06:44:33Z)
• LLMC: Benchmarking Large Language Model Quantization with a Versatile Compression Toolkit [55.73370804397226]
  Quantization, a key compression technique, can effectively mitigate these demands by compressing and accelerating large language models. We present LLMC, a plug-and-play compression toolkit, to fairly and systematically explore the impact of quantization. Powered by this versatile toolkit, our benchmark covers three key aspects: calibration data, algorithms (three strategies), and data formats.
  arXiv  Detail & Related papers  (2024-05-09T11:49:05Z)
• Norm Tweaking: High-performance Low-bit Quantization of Large Language Models [21.855106896725598]
  We introduce a technique called norm tweaking, which can be used as a plugin in current PTQ methods to achieve high precision. Our method demonstrates significant improvements in both weight-only quantization and joint quantization of weights and activations. Our simple and effective approach makes it more practical for real-world applications.
  arXiv  Detail & Related papers  (2023-09-06T06:51:15Z)
• Gradient $\ell_1$ Regularization for Quantization Robustness [70.39776106458858]
  We derive a simple regularization scheme that improves robustness against post-training quantization. By training quantization-ready networks, our approach enables storing a single set of weights that can be quantized on-demand to different bit-widths.
  arXiv  Detail & Related papers  (2020-02-18T12:31:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.