Back Home

模型推論與量化

BaKron Reduces Two-Sided Hessian Quantization to Cubic Complexity, With Its 8192-Square-Matrix Kernel 60× Faster Than YAQA

Using anti-diagonal parallelization and recursive divide-and-conquer, BaKron reduces total work from quartic to cubic complexity while preserving curvature information on both the input and output sides. In 2.81-bit experiments on Llama 3 8B, it improves perplexity over GPTQ, but total quantization time remains longer, and no public implementation is currently available.

Gabor Eszes (UED77) · CC BY-SA 3.0 · Image source
zh-Hant

BaKron addresses the information trade-off in GPTQ-style post-training quantization. GPTQ uses an input correlation matrix derived from calibration data to adjust weight-by-weight rounding, which is equivalent to using a Hessian of the form `A⊗I` and ignores correlations among different output coordinates. Two-sided methods such as BoA and YAQA extend this to `A⊗B`, capturing both input and output geometry, but operating directly on vectorized weights requires quartic work. BaKron processes mutually independent updates in batches along anti-diagonals, then uses divide-and-conquer to reuse matrix multiplications, reducing the core cost for an m×n weight matrix to `O(mn(m+n))` while maintaining a sequential depth of `O(m+n)`.

In microbenchmarks on a single NVIDIA RTX PRO 6000 using a custom float32 Triton kernel, BaKron processes an 8192×8192 matrix in 1.839 seconds, versus 110.379 seconds for an equivalent YAQA implementation—a 60× kernel speedup. It is also about 43× faster on 4096×14336 matrices and the transposed shape. However, GPTQ processes the same 8192-square matrix in just 0.059 seconds because it still uses simpler one-sided geometry. BaKron’s value, therefore, is not in beating GPTQ on speed, but in reducing previously hard-to-scale two-sided quantization to the same cubic complexity class as GPTQ.

Quality tests quantized Llama 3 and Qwen3 to approximately 2.81 bits per weight. On Meta-Llama-3-8B, GPTQ achieved a WikiText2 perplexity of 53.47, while BaKron using a backpropagation-based K-FAC Hessian reduced it to 11.90; the unquantized model scored 7.44. However, the BaKron pipeline took 1,939 seconds, compared with 530 seconds for GPTQ, showing that Hessian accumulation, forward and backward passes over calibration data, and Cholesky decomposition still dominate end-to-end cost. The paper also evaluates only a single GPU, specific scaling choices, and a 2.81-bit configuration, and provides no public code. Engineering teams should wait for an implementation and validation across more bit widths and real-world inference throughput.

Sources

  1. BaKron: Efficient Quantization with Kronecker-Factored Hessians
  2. GPTQ Official Implementation
  3. Meta-Llama-3-8B Model Card