Search papers, labs, and topics across Lattice.
This paper introduces BaKron, an efficient algorithm for neural network quantization that leverages two-sided Kronecker-factored Hessian approximations to capture correlations across output coordinates. By employing a recursive divide-and-conquer approach combined with anti-diagonal parallelism, BaKron reduces computational complexity from \(O(m^2n^2)\) to \(O(mn(m+n))\), while maintaining the cubic scaling of existing methods like GPTQ. Experimental benchmarks demonstrate BaKron's effectiveness in quantization tasks, showcasing its modularity with respect to quantizers and Hessian estimators.
BaKron accelerates neural network quantization by harnessing richer curvature information, achieving significant computational efficiency without sacrificing performance.
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.
