BaKron: Efficient Quantization with Kronecker-Factored Hessians

2026-08-06

AI Summary

This work addresses the high computational cost of modeling the Hessian matrix via bidirectional Kronecker factorization in neural network quantization, which, while effective at capturing output dimension correlations, incurs prohibitive overhead. The authors propose BaKron, a novel solver that integrates anti-diagonal parallelism with a recursive divide-and-conquer strategy to achieve a serial step complexity of O(m+n) and an overall computational complexity of O(mn(m+n)). BaKron preserves the cubic complexity of GPTQ while enabling modular integration of base quantizers and Hessian estimators. Experimental results demonstrate that BaKron substantially reduces computational costs across diverse Hessian configurations and consistently improves quantization accuracy.

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.

Problem

• neural network quantization
• Kronecker-factored Hessian
• adaptive rounding
• computational efficiency
• weight correlation

Innovation

• Kronecker-factored Hessian
• neural network quantization
• adaptive rounding
• divide-and-conquer
• computational efficiency