FastKron: Efficient Quantization with Kronecker-Factored Hessians

Efficient Quantization with Kronecker-Factored Hessians" by Johann Birnick and Rayan Saab.

The paper introduces BaKron, an efficient solver for post-training neural network quantization that uses Kronecker-factored Hessian approximations to inform the rounding of weights. The authors study quantization of a linear unit W in R m times n by replacing it with a low-precision matrix V, while preserving network accuracy. Unlike GPTQ, which uses a one-sided Hessian approximation of the form XX T I m, BaKron handles a general two-sided Kronecker-factored Hessian A B, capturing correlations across both input and output features.

The abstract states: "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 squared n 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."

The paper derives the algorithm step by step. The starting point, BaKron-naive, is GPTQ in the vectorized weight domain: running GPTQ on vec(W) with Hessian A B, using Proposition 3.1 to express operations unvectorized. This requires mn sequential steps and O(m squared n 2) total work. The next variant, BaKron-antidiagonal, processes anti-diagonals in parallel: the update caused by an entry W i,j only updates entries to the bottom-right of itself, i.e., it only updates entries W i',j' which satisfy i' i and j' j. This reduces sequential steps to O(m+n), but retains O(m squared n 2) total work. The final algorithm, BaKron, combines anti-diagonal parallelism with a recursive divide-and-conquer approach over the anti-diagonal dimension: "We process anti-diagonals in parallel, and we do a recursive divide-and-conquer approach over the 'anti-diagonal dimension'. The latter means that we recursively process the first half of anti-diagonals, then we propagate the error to the second half of anti-diagonals, and after that we recursively process the second half of anti-diagonals." This yields both O(m+n) sequential steps and O(mn(m+n)) total work.

The paper proves (Theorem 3.2) that all four algorithms—BaKron-naive, BaKron-antidiagonal, BaKron-recursive, and BaKron—are equivalent and produce the same outputs. It provides an error guarantee (Corollary 3.4): with L(A) = RevCholesky(A) and L(B) = RevCholesky(B), the output V satisfies

[

L(B)(V-W)(L(A)) T F squared 1 over 4 tr(A) times tr(B).

]

More generally, "If A and B encode the geometries that are relevant to the layer, then B 1/2(V-W)A 1/2 F squared is a more natural proxy for the loss induced by quantization. Moreover, the bound depends on tr(A) tr(B) rather than on tr(A)m."

The paper also considers possible Kronecker-factored Hessians H W about A B to use with BaKron. For the loss, it considers a global Fisher information matrix using backpropagation: H W = E[aa T gg T], where g = grad b L x,y, and local Hessians that do not require backward passes. For local attention Hessians, following BoA, it reports H W Q = XX T KK T, H W K = XX T QQ T. For MLP modules in Transformers with GLU activation, it derives:

[

H W up = E[xx T (gg T W down T W down)], g = sigma(W gatex),

]

[

H W gate = E[xx T (ff T W down T W down)], f = sigma'(W gatex) W upx.

]

For Kronecker-factorizing an expected Hessian, the paper considers a K-FAC-style independence assumption H about E[aa T] E[bb T], and a Shampoo-style power-iteration scheme to find the best Kronecker approximation, yielding e.g. A about E[tr(b T b) aa T], B about E[tr(a T a) bb T] after one iteration with identity initialization.

The paper additionally proposes an efficient technique to compute backpropagated (global) Hessians. Naively, one could do a backward pass per layer, costing O(2) compute, or accumulate all layer Hessians in a single backward pass as in YAQA, requiring O memory with a big constant hidden in the asymptotic notation. The authors propose a recursive halving technique: "Given a range of layers, initially the full list of layers, the recursive algorithm requires the input to the first layer and the gradient of the output of the last layer. Then a forward and backward pass is required to compute the activations as well as their gradient at exactly the middle of the layer list... This divide-and-conquer approach requires O compute and O memory. They explain that the practical improvement is reducing accelerator-resident memory: our flow reduces the accelerator-resident memory needed for the global Hessians from several times the model size to essentially the size of a single layer, at the price of a O factor of additional compute."

Benchmarks show practical speedups of BaKron over the previously best implementation (BaKron-antidiagonal, equivalent to YAQA). In Table 3, for a 8192 times 8192 matrix, BaKron runs in 1.839 seconds versus 110.379 seconds for BaKron-antidiagonal, a 60.0 times speedup. The speedup grows with matrix size, as the quartic versus cubic total cost predicts. BaKron-naive is impractical beyond small shapes, costing about 50 microseconds per weight.

Finally, the paper evaluates BaKron on Llama-3.2-1B, Llama-3.2-3B, Meta-Llama-3-8B, Qwen3-1.7B-Base, Qwen3-4B-Base, and Qwen3-8B-Base at 2.81 bits per weight, with calibration on 256 sequences of 2048 tokens from The Pile. The experiments compare GPTQ with three BaKron variants—BaKron-MlpLocal, BaKron-FullyLocal, and BaKron-Backprop—each with K-FAC-style and Shampoo-style factorizations. The Core share column shows that the core quantization algorithm accounts for at most 12.2% of total pipeline time, confirming that even with a two-sided Kronecker-factored Hessian, the pipeline remains dominated by accumulating the Hessians and by sending the calibration data through the model. Regarding quality, BaKron-Backprop-Shampoo usually attains a lower Wikitext2 perplexity than GPTQ, although not on every model. The local variants stay closer to GPTQ.

Improved AI systems can use BaKron to achieve the following specific improvements:

• Better low-bit LLM quantization

Quantize models such as Llama-3 and Qwen3 to 2.8 bits per weight with lower Wikitext-2 perplexity than GPTQ, by exploiting two-sided Kronecker-factored Hessian curvature that captures correlations across both input and output features.

• Faster post-training quantization

Reduce quantization runtime from O(m squared n 2) to O(mn(m+n)) total work while keeping only O(m+n) sequential steps. For an 8192 times 8192 weight matrix, this gives a 60× speedup over previous antidiagonal/YAQA-style implementations.

• Practical scalability to large models

Quantize very large weight matrices in seconds instead of minutes, making low-bit compression feasible for billion-parameter models on limited compute budgets.

• Memory-efficient Hessian estimation

Compute global backpropagated Hessians using a recursive halving technique that reduces accelerator-resident memory from several times the model size to roughly one layer’s size, at the cost of only O additional compute. This enables curvature-aware quantization on models that previously could not fit in GPU memory.

• Modular integration into existing pipelines

Plug BaKron into any quantization framework because it is independent of the base quantizer and Hessian estimator. The core quantization algorithm takes at most 12% of total pipeline time, so the system can spend its budget on better Hessian accumulation and calibration data.

• Flexible Hessian approximation choices

Use either K-FAC-style independence assumptions or Shampoo-style power iteration to Kronecker-factor Hessians, allowing a trade-off between computational cost and quantization quality depending on the model and task.

• Better compression of attention and MLP layers

Apply improved per-layer Hessians for Transformers, including GLU-based MLP modules, enabling the quantized system to preserve more task accuracy in attention projections (W Q, W K) and feed-forward layers.