BaKron: Efficient Quantization with Kronecker-Factored Hessians

Neural network quantization aims to reduce the precision of weight matrices while maintaining model accuracy. Traditional methods like GPTQ optimize quantization by considering the geometry of input activations, typically using a one-sided Hessian approximation. This paper addresses the challenge of incorporating two-sided Hessian information—which captures correlations across both input and output features—without incurring the prohibitive computational costs associated with naive implementations of such methods.

The authors introduce BaKron, an algorithm designed to solve the quantization problem using a Kronecker-factored Hessian approximation (A ⊗ B). While previous attempts to use two-sided approximations (such as BoA and YAQA) were computationally expensive, BaKron optimizes the process by integrating two key techniques:

1. Anti-diagonal Parallelism: By analyzing the dependency graph of weight updates, the algorithm processes entries along anti-diagonals in parallel, significantly reducing the number of sequential steps.
2. Recursive Divide-and-Conquer: The authors implement a recursive strategy that batches updates, effectively lowering the total work from O(m²n²) to O(mn(m+n)).

This combination allows BaKron to match the cubic scaling of GPTQ while utilizing richer geometric information about the network's loss landscape.

BaKron provides a more robust way to quantize neural networks by accounting for output-side correlations, which are ignored by standard one-sided methods. By reducing the computational complexity to match that of GPTQ, the authors make it practical to apply these more sophisticated, geometry-aware quantization techniques to large-scale models. The paper also provides practical guidance on Hessian estimation and efficient computation, offering a modular framework that can be adapted to different quantizers and Hessian estimators.