Kronecker-factored Hessians improve post-training quantization accuracy
WaterKron and FlipFlop Hessian: Information-Theoretically Grounded Quantization with Kronecker-factored Hessians
Machine LearningArtificial Intelligence
Summary
Post-training quantization is a way to shrink large AI models so they run faster or use less memory, but it’s tricky to keep their accuracy. This paper introduces WaterKron, a method that smartly adjusts how these models are compressed by analyzing how quantization errors impact the model’s behavior. The authors also developed FlipFlop Hessian, a mathematical approach to better approximate the errors introduced during compression, making the quantization process more reliable. They tested their ideas and found that FlipFlop Hessian consistently leads to better results compared to other common ways of approximating error.
What this means in practice
- •For machine learning engineers: Use FlipFlop Hessian to improve accuracy when compressing pretrained models for deployment without retraining.
- •For cloud infrastructure teams: Optimize storage and bandwidth by applying WaterKron quantization to reduce model size with minimal impact on inference quality.
Authors
Johann Birnick, Rayan Saab
Abstract
How should a Kronecker-factored Hessian approximation be chosen for post-training quantization? We address this question through WaterKron, which combines two-sided GPTQ with row- and column-dependent waterfilling scales and entropy coding. We derive its high-rate distortion with respect to the full Hessian using an explicit Kronecker-Hessian mismatch factor $Φ$. This factor quantifies the asymptotic distortion penalty due to the Kronecker Hessian approximation and provides a criterion for selecting the factors optimally. Minimizing $Φ$ leads to a Gaussian covariance-fitting problem with classical ``flip-flop'' updates. We thus give a rate-distortion justification for using the resulting FlipFlop Hessian in quantization. We evaluate it empirically, finding that the FlipFlop Hessian consistently improves KL divergence and perplexity over input-only, marginal, and Frobenius-based Hessian choices.