IDScore: Task-Conditioned Likelihood Correction for Low-Bit QAT

Quantization-aware training (QAT) is a key approach for optimizing quantized neural networks. However, deterministic quantization introduces piecewise-constant mappings that make direct gradient-based optimization difficult, so QAT relies on surrogate gradients whose quality is critical to effective training. We propose Identification-Score (IDScore), a correction that improves surrogate gradient estimation in QAT with little additional training cost. This improvement comes from a new perspective: motivated by quantized system identification, IDScore exploits additional structure information from the quantizer itself, beyond the external signals used by existing methods. Concretely, IDScore constructs a local likelihood score from a task-selected neighboring target, the latent parameter, and quantization thresholds, and uses it to form a bounded residual correction. Our theory explains why IDScore provides a meaningful correction direction and how its correction strength should be allocated across parameters. Formally, we prove calibrated target-direction consistency and constrained-optimal likelihood allocation, and connect the resulting correction to threshold progress and a local sufficient condition for task-loss descent. Experiments on vision and language tasks show that IDScore consistently improves the tested LSQ and QuEST baselines at low additional cost. Evaluations span W1A1-W4A4 precision and models with up to 430M non-embedding parameters. It also yields further gains when combined with FOGZO or CAGE. Controlled ablations support the likelihood-based allocation mechanism characterized by our theory. Together, these results provide both an efficient approach with practical potential for low-bit training and a broader perspective on how quantization structure can inform backward optimization.