PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Johann Birnick, Rayan Saab

arXiv:2609.14706v2cs.LGcs.AI

TL;DR

Post-training quantization raises the question of how to choose a Kronecker-factored Hessian for quantizing weights. The paper combines two-sided GPTQ with waterfilling and uses full-Hessian distortion to define a factor-selection criterion, then reports that its FlipFlop Hessian improves on prior Hessian choices empirically.

  • Problem

    The paper asks how to choose a Kronecker-factored Hessian approximation when the distortion of quantized weights should reflect the full Hessian.

    Source: §4 · p. 7

  • Method

    WaterKron applies two-sided GPTQ with row- and column-dependent waterfilling scales, while a mismatch factor guides Kronecker-factor selection and yields the FlipFlop Hessian through Gaussian covariance fitting.

    Source: §1 · p. 2, §4 · p. 8, §5 · p. 9

  • Results

    Under Gaussian and entropy-coding assumptions, WaterKron has an approximately 0.255-bit asymptotic gap to the information-theoretic limit for Kronecker distortion and empirically outperforms prior Hessian choices.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The mismatch factor provides a distortion-based criterion for selecting Kronecker factors, and its minimization justifies the FlipFlop Hessian for quantization.

    Source: §1 · p. 2, §4 · p. 8

    The FlipFlop derivation is developed for H expressible as an expectation over outer products of Kronecker-factored vectors, with Fisher information given as a natural instance.

    Source: §5 · p. 9

Abstract

from arXiv · show

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 $\Phi$. This factor quantifies the asymptotic distortion penalty due to the Kronecker Hessian approximation and provides a criterion for selecting the factors optimally. Minimizing $\Phi$ 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.

1 INTRODUCTION

The paper studies post-training weight quantization under quadratic-loss error propagation, where GPTQ uses input correlations to guide quantization. It develops WaterKron for two-sided geometry and uses full-Hessian distortion to motivate a mismatch criterion for choosing Kronecker factors and a FlipFlop approximation.

1 INTRODUCTION

  • Post-training quantization approximates trained high-precision weights with lower-precision weights to retain accuracy while reducing memory requirements.The paper treats each linear module W separately, seeking a low-precision approximation V.Source: §1 · p. 1
  • The work uses GPTQ's quadratic-loss error propagation, which exploits correlations among a module's input channels to guide adaptive rounding.The loss accounts for the effect of replacing W with V on the module output, rather than only measuring the weight-matrix difference.Source: §1 · p. 1
  • Two-sided GPTQ uses a Kronecker-factored Hessian A ⊗ B to include output geometry as well as the input geometry used by standard GPTQ.The paper positions its work within adaptive rounding under a quadratic loss and extends that setting with information-theoretic rate allocation.Source: §1 · pp. 1–2
  • WaterKron extends two-sided GPTQ with optimal row and column scales, while the FlipFlop Hessian minimizes a mismatch factor to select Kronecker factors.Under Gaussian and entropy-coding assumptions, WaterKron retains an approximately 0.255-bit asymptotic gap to the information-theoretic limit for Kronecker distortion; the paper reports empirical improvements over prior Hessian choices.Source: §1 · p. 2

2 SETTING AND BACKGROUND

The section sets up weight-only quantization and GPTQ, then extends GPTQ to two-sided Kronecker geometry. It motivates waterfilling scales by showing that equal-sized quantization cells minimize distortion at fixed volume.

2.1 ALPHABET

  • The paper studies weight-only quantization, representing the lower-precision matrix through integer entries and row- and column-specific scales.The integers are unbounded and encoded; the scale vectors are stored in high precision, with the paper noting their practical memory cost is small.Source: §2.1 · p. 3

2.2 GPTQ: ADAPTIVE ROUNDING

  • GPTQ uses reverse Cholesky factors of the input covariance to encode input-channel correlations, sequentially quantizing channels and modifying W to partially correct earlier errors.The expectation is taken over an empirical calibration distribution.Source: §2.2 · p. 3

2.3 TWO-SIDED EXTENSION AND KRONECKER-FACTORED HESSIANS

  • Two-sided GPTQ replaces the input-only Hessian with A ⊗ B, so the algorithm accounts for output geometry as well as input correlations.When B = I_m, the extension is equivalent to standard GPTQ.Source: §2.3 · p. 4

2.4 INFORMATION-THEORETIC VIEW AND WATERFILLING FORMULA

  • For standard GPTQ, the high-rate waterfilling scale is α_j ∝ 1/L_j,j, with L the reverse Cholesky factor of the input covariance.The section describes this allocation as a way to distribute precision non-uniformly across columns.Source: §2.4 · p. 5
  • At fixed cell volume, equal side lengths minimize distortion, yielding the waterfilling scales and cube-shaped quantization cells.The geometric explanation views GPTQ as Babai nearest-plane rounding on a lattice.Source: §2.4 · pp. 5–6

3 TWO-SIDED RATE ALLOCATION

The section extends information-theoretic rate allocation to two-sided GPTQ by choosing both row and column scales. It establishes corresponding two-sided rate-distortion results and notes that nontrivial scales on both axes call for entropy coding that handles entry-specific scales.

3 TWO-SIDED RATE ALLOCATION

  • The two-sided allocation chooses scales proportional to inverse diagonal entries of the reverse Cholesky factors for A and B.The section establishes two-sided information-theoretic and achievability results by relating the two-sided algorithm to standard GPTQ on a transformed input.Source: §3 · p. 6
  • The two-sided rate and distortion definitions follow from the standard results by treating A ⊗ B and vec(W) as the transformed input-side problem.The derivation relates the determinant term for A ⊗ B to the separate determinants of A and B.Source: §3 · pp. 6–7
  • With nontrivial row and column scales, each matrix entry has a unique scale, making per-column Huffman coding inefficient; arithmetic-coding variants can handle this setting.The remark contrasts this with the input-only case, where entries in a column share a scale.Source: §3 · p. 7

4 KRONECKER-HESSIAN MISMATCH FACTOR

Because A⊗B only approximates the full Hessian H, the section relates two-sided distortion to full-Hessian distortion and introduces Φ as a criterion for choosing the Kronecker factors.

4 KRONECKER-HESSIAN MISMATCH FACTOR

  • Asymptotically, full-Hessian distortion equals two-sided distortion multiplied by tr(H(A⊗B)^−1)/mn.The correction term is generally present but becomes negligible as quantization errors become asymptotically uniform, which occurs in the high-rate limit under the Gaussian setting.Source: §4 · pp. 7–8
  • Minimizing Φ(H, A⊗B) selects the Kronecker factors that minimize the A,B-dependent term in the full-distortion result.Φ normalizes that term by |H|^(1/mn).Source: §4 · p. 8
  • Φ is the arithmetic-to-geometric mean ratio of generalized eigenvalues and is at least 1, with equality only when H is proportional to A⊗B.Source: §4 · p. 8
  • Up to rescaling A⊗B by a positive scalar, Φ is essentially the exponential of the Gaussian KL divergence between covariances H and A⊗B.Source: §4 · p. 8

5 THE FLIPFLOP HESSIAN APPROXIMATION

The paper chooses the Kronecker factors by minimizing the mismatch factor Φ, using alternating updates for A and B. This optimization is framed for Hessians that are expectations of outer products of Kronecker-factored vectors.

5 THE FLIPFLOP HESSIAN APPROXIMATION

  • The optimization minimizes Φ(H, A ⊗ B) for Hessians H expressed as expectations of outer products of Kronecker-factored vectors.In the Fisher setting described, x is a calibration input and g is the gradient of the negative log-likelihood with respect to the network output.Source: §5 · p. 9
  • For fixed A, B is proportional to E[GA^-1G^T], and for fixed B, A is proportional to E[G^TB^-1G], yielding alternating factor updates.The objective has no general closed-form solution; the paper presents an efficient iteration that converges to a solution.Source: §5 · p. 9
  • The alternating updates are called the flip-flop algorithm in matrix-normal maximum-likelihood estimation, and yield the FlipFlop Hessian.For positive-definite H, the iteration converges to a unique fixed point up to scaling.Source: §5 · p. 9
  • The Frobenius variant instead minimizes ||H − A ⊗ B||F and is described as an additive approximation, while FlipFlop is multiplicative.The paper reports that FlipFlop outperforms Frobenius in practice.Source: §5 · p. 9

6 EXPERIMENT FINDINGS

Across the reported experiments, FlipFlop is consistently the best Hessian variant, while Frobenius generally worsens with more iterations. The reported Figure 2 comparison covers KL divergence and perplexity across rates for Llama-3.2-1B.

6 EXPERIMENT FINDINGS

  • WaterKron generally outperforms input-sided WaterSIC when the Hessian approximation is chosen carefully; Frobenius generally worsens with more iterations.For larger 8B models, the Frobenius variant is clearly worse than the input-sided Hessian.Source: §6 · pp. 9–10
  • FlipFlop is consistently the best of the evaluated variants, and generally improves with its second iteration.The authors say this supports minimizing the mismatch factor Φ rather than using the Frobenius-based criterion.Source: Figure 2 · p. 10, §6 · p. 10
  • Figure 2 covers KL divergence and perplexity across rates for Llama-3.2-1B, comparing different Hessian approximations.The caption refers readers to Figure 3 as well.Source: Figure 2 · p. 10

A LIMITATIONS

The practical use of WaterKron is limited by its reliance on entropy coding, and the theory is developed in the high-rate regime while practical quantization and experiments use low rates.

A LIMITATIONS

  • WaterKron’s practical limitation is its use of entropy coding, whose decoders may be difficult to implement efficiently as GPU kernels.The paper notes this limitation is shared with Huffman-GPTQ and WaterSIC.Source: A LIMITATIONS · p. 12
  • Suggested workarounds are accelerator-native compression schemes or splitting Z into a dense fixed-datatype matrix and a sparse outlier matrix.Source: A LIMITATIONS · p. 12
  • The theory applies in the high-rate regime, whereas practical neural-network quantization and the experiments operate in the low-rate regime.Source: A LIMITATIONS · p. 12

B EXPERIMENT RESULTS

The experiments isolate the contributions of WaterKron and the FlipFlop Hessian in a basic setting. On Llama-3.2-1B across rates, the figure reports that FlipFlop outperforms the other Hessian approximations.

B EXPERIMENT RESULTS

  • The comparison isolates WaterKron and FlipFlop against basic WaterSIC without other GPTQ refinements.The authors omit drift correction, improved base quantizers, scale fine-tuning, and random or learned rotations to isolate the Hessian approximation contribution.Source: B EXPERIMENT RESULTS · p. 12
  • The experiments evaluate Llama and Qwen models up to 8B parameters at 2 bits per weight, plus Llama-3.2-1B at various rates.Tables 2–7 cover the 2-bit runs; Figure 3 covers the varying-rate run.Source: B EXPERIMENT RESULTS · p. 12, Table 2 · p. 13, Table 3 · p. 14
  • FlipFlop outperforms all other Hessian approximations for Llama-3.2-1B across the rates shown.Figure 3 compares all eight Hessian variants.Source: Figure 3 · p. 13
  • Evaluation uses KL divergence and perplexity on Pile10k and WikiText-2, averaged over tokens.The setup quantizes transformer-block linear modules, calibrates on 256 sequences of 2048 Pile10k tokens, and uses unquantized-model inputs.Source: B EXPERIMENT RESULTS · pp. 12–13

C TECHNICAL DETAILS

This section clarifies two technicalities that the authors left unaddressed while focusing on the paper’s core theorem statements.

C TECHNICAL DETAILS

  • The section concerns two technicalities omitted from the theorem statements.The authors say they deliberately left them without comment to focus on the core statements.Source: C TECHNICAL DETAILS · p. 15
  • The authors characterize the omission as deliberate rather than as part of the core theorem statements.Source: C TECHNICAL DETAILS · p. 15
  • The purpose of this section is to clarify those technicalities.Source: C TECHNICAL DETAILS · p. 15

C.1 ENCODING OF THE SCALING FACTORS α, β

The theorem statements omit coding the scaling factors, which has negligible practical memory cost; the lower bound also holds when the decoder receives the Hessian factors.

C.1 ENCODING OF THE SCALING FACTORS α, β

  • Theorem 2.2, Corollary 3.2, and Theorem 4.1 analyze encoding Z but omit encoding the scaling factors α and β.The omission keeps the theorem statements clean; in practice, the scaling factors’ memory cost is negligible.Source: §C.1 · p. 15
  • For the coding argument, the decoder may receive α and β directly, or even full access to A and B.With A and B, the decoder can recover α and β and reconstruct V from the entropy-coded Z.Source: §C.1 · p. 15
  • The lower bound of Corollary 3.1 still holds when the decoder has access to A and B.The cited standard proof already assumes that decoder access.Source: §C.1 · p. 15

C.2 LIMIT γ →0 AND mn →∞

The theorems take a high-rate limit and a growing-matrix limit, understood in an iterated order: first the rate tends to zero for fixed Hessian factors, then matrix size grows.

C.2 LIMIT γ →0 AND mn →∞

  • Theorem 2.2, Corollary 3.2, and Theorem 4.1 consider γ → 0 and mn → ∞.Source: §C.2 · p. 16
  • Taking γ → 0 places the analysis in the high-rate regime, where transformed quantization errors become uniformly distributed.Source: §C.2 · p. 16
  • The growing-matrix limit ensures the coding scheme can approach the entropy of Z; the joint limit is interpreted iteratively, first γ → 0 for fixed A and B, then mn → ∞.This order ensures the high-rate limit for each concrete pair of Hessian factors.Source: §C.2 · p. 16
Loading 2609.14706v2…