BaKron: Efficient Quantization with Kronecker-Factored Hessians
2026-08-06 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors improve methods for simplifying neural networks by turning their weights into smaller, faster numbers (quantization). They focus on using mathematical information about the network's curvature (Hessian) in a way that looks at relationships from both input and output sides. Their new method, called BaKron, reduces the time and work needed to perform this process compared to previous approaches while keeping the benefits of richer curvature information. They also show how BaKron works with different setups and test it in practice.
neural network quantizationHessian matrixKronecker factorizationadaptive roundingcurvature informationGPTQdivide-and-conquer algorithmcomputational complexityweight matrix
Authors
Johann Birnick, Rayan Saab
Abstract
We accelerate a family of algorithms for neural network quantization whose geometry is informed by any Kronecker-factored approximation of the Hessian. GPTQ-style adaptive rounding typically uses one-sided information derived from input activations. Two-sided Kronecker-factored Hessian approximations can additionally capture correlations across output coordinates, but applying GPTQ directly in the vectorized weight domain is computationally expensive. Building on the two-sided adaptive-rounding formulation used by BoA and YAQA, we introduce BaKron, an efficient solver that combines anti-diagonal parallelism with a recursive divide-and-conquer construction. For an $m\times n$ weight matrix, BaKron uses $O(m+n)$ sequential steps while reducing the total work from $O(m^2n^2)$ to $O(mn(m+n))$. Thus, it matches the cubic scaling of GPTQ while exploiting richer curvature information. Moreover, BaKron is modular with respect to both the base quantizer and the Hessian estimator. We also provide practical benchmarks, consider a range of Hessians that BaKron can be called with, find an efficient technique to compute these Hessians, and evaluate the algorithm experimentally.