WaterKron and FlipFlop Hessian: Information-Theoretically Grounded Quantization with Kronecker-factored Hessians

Johann Birnick; Rayan Saab

cs.LGSep 13, 2026

WaterKron and FlipFlop Hessian: Information-Theoretically Grounded Quantization with Kronecker-factored Hessians

Authors: Johann Birnick, Rayan Saab

Organizations: Department of Mathematics University of California San Diego La Jolla, CA 92093, USA · Department of Mathematics and HDSI University of California San Diego La Jolla, CA 92093, USA

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.

Figures & tables

Figure 1: Babai rounding for

L=\bigl(\begin{smallmatrix}2&0\\ 1&3\end{smallmatrix}\bigr)

. Lattice points

L\operatorname{diag}(\alpha)\mathbb{Z}^{2}

are represented by black dots. Target vectors

Lw

with

w\sim\mathcal{N}(0,I_{2})

are represented by blue dots. Each target is mapped to the lattice point at the center of its cell. The fundamental region

E

that contains the quantization errors is highlighted. On the left, the scaling factors are

\alpha=(1,1)

, giving rectangular cells. On the right, they are chosen optimally as

\alpha=(3/\sqrt{6},\,2/\sqrt{6})

, giving square cells while preserving the volume.

Name	Input factor $A$	Output factor $B$	Notes
Input	$\mathbb{E}[xx^{T}]$	$I_{m}$	Recovers WaterSIC.
Marginal	$\mathbb{E}[xx^{T}]$	$\mathbb{E}[gg^{T}]$	Treats $x,g$ as independent.
Frobenius	$\mathbb{E}[(g^{T}Bg)\,xx^{T}]$	$\mathbb{E}[(x^{T}Ax)\,gg^{T}]$	Targets $\lVert H-A\otimes B\rVert_{F}$ .
FlipFlop	$\mathbb{E}[(g^{T}B^{-1}g)\,xx^{T}]$	$\mathbb{E}[(x^{T}A^{-1}x)\,gg^{T}]$	Targets $\Phi(H,A\otimes B)$ .

Table 1: Different Hessian approximations

A\otimes B

H=\mathbb{E}[xx^{T}\otimes gg^{T}]

. Frobenius and FlipFlop approximations are computed using the alternating fixed-point iteration, starting from the identity matrix.

Figure 2: Llama-3.2-1B KL/PPL across rates for various Hessian approximations. (See Figure 3 .)

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: More detailed version of Figure 2 . Llama-3.2-1B quantized with all eight Hessian variants across various rates. “FlipFlop” outperforms all other Hessian approximations.

	Pile10k		WikiText-2
Hessian	KL $\downarrow$	PPL $\downarrow$	KL $\downarrow$	PPL $\downarrow$
Unquantized	0.00	7.3	0.00	9.8
Input (WaterSIC)	1.14	22.5	1.60	49.3
Marginal	0.70	14.6	0.98	26.3
Frobenius-1	0.62	13.5	0.91	24.4
Frobenius-2	0.74	15.2	1.09	29.4

Appendix

Table 2: Llama-3.2-1B at a target rate of 2 bits per weight.

	Pile10k		WikiText-2
Hessian	KL $\downarrow$	PPL $\downarrow$	KL $\downarrow$	PPL $\downarrow$
Unquantized	0.00	6.2	0.00	7.8
Input (WaterSIC)	0.46	9.5	0.72	15.8
Marginal	0.41	9.2	0.71	15.6
Frobenius-1	0.39	9.0	0.68	15.2
Frobenius-2	0.46	9.6	0.82	17.5

Appendix

Table 3: Llama-3.2-3B at a target rate of 2 bits per weight.

	Pile10k		WikiText-2
Hessian	KL $\downarrow$	PPL $\downarrow$	KL $\downarrow$	PPL $\downarrow$
Unquantized	0.00	5.3	0.00	6.2
Input (WaterSIC)	0.41	7.9	0.62	11.6
Marginal	0.40	7.9	0.64	11.9
Frobenius-1	0.42	8.0	0.67	12.2
Frobenius-2	0.47	8.4	0.69	12.5

Appendix

Table 4: Llama-3.1-8B at a target rate of 2 bits per weight.

	Pile10k		WikiText-2
Hessian	KL $\downarrow$	PPL $\downarrow$	KL $\downarrow$	PPL $\downarrow$
Unquantized	0.00	9.1	0.00	16.7
Input (WaterSIC)	0.78	17.2	1.11	37.2
Marginal	0.74	16.1	1.09	36.1
Frobenius-1	0.64	14.2	0.90	28.8
Frobenius-2	0.76	15.7	1.07	32.3

Appendix

Table 5: Qwen3-1.7B at a target rate of 2 bits per weight.

	Pile10k		WikiText-2
Hessian	KL $\downarrow$	PPL $\downarrow$	KL $\downarrow$	PPL $\downarrow$
Unquantized	0.00	7.8	0.00	13.7
Input (WaterSIC)	0.39	10.1	0.61	19.7
Marginal	0.36	9.9	0.58	19.2
Frobenius-1	0.42	10.2	0.61	18.9
Frobenius-2	0.47	10.6	0.66	19.9

Appendix

Table 6: Qwen3-4B at a target rate of 2 bits per weight.

	Pile10k		WikiText-2
Hessian	KL $\downarrow$	PPL $\downarrow$	KL $\downarrow$	PPL $\downarrow$
Unquantized	0.00	6.5	0.00	9.7
Input (WaterSIC)	0.28	7.9	0.45	13.0
Marginal	0.29	8.0	0.47	13.1
Frobenius-1	0.31	8.0	0.49	13.6
Frobenius-2	0.34	8.2	0.50	13.6

Appendix

Table 7: Qwen3-8B at a target rate of 2 bits per weight.

Neural Network Quantization

Neural network quantization represents weights, activations, and states with fewer bits to reduce memory use, computation, and deployment costs.

See all 394 papers

BaKron: Efficient Quantization with Kronecker-Factored Hessians
Aug 6, 2026Johann Birnick, Rayan SaabPost-Training Quantization Neural Network Quantization
Department of Mathematics · University of California San Diego · La Jolla, California, United States +1
KronQ: LLM Quantization via Kronecker-Factored Hessian
Jul 8, 2026Donghyun Lee, Yuhang Li, Ruokai Yin +1Mixed-Precision Quantization LLM Quantization
University of Southern California · Yale University
High-Rate Quantized Matrix Multiplication II
May 13, 2026Or Ordentlich, Yury PolyanskiyLLM Quantization Rate-Distortion Theory
the Hebrew University of Jerusalem, Israel · MIT, USA

Aug 6, 2026J/K move · Enter open · S save

BaKron: Efficient Quantization with Kronecker-Factored Hessians

Johann Birnick, Rayan Saab

Department of Mathematics · University of California San Diego · La Jolla, California, United States+1

Post-Training Quantization Neural Network Quantization Kronecker-Factored Approximation

Open paper AskPDF arXiv

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)

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.

Abstract

Neural Network Quantization

Explore similar work

BaKron: Efficient Quantization with Kronecker-Factored Hessians

KronQ: LLM Quantization via Kronecker-Factored Hessian

High-Rate Quantized Matrix Multiplication II