BaKron、両側 Hessian 量子化を3次計算量に削減――8192次正方行列カーネルでYAQAの60倍高速

BaKronが解決しようとしているのは、GPTQ系の学習後量子化における情報のトレードオフだ。GPTQは、キャリブレーションデータから得た入力相関行列を用いて重みごとの丸めを調整する。これは `A⊗I` 形式の Hessianを使用することに相当し、異なる出力座標間の相関を無視する。BoAやYAQAなどの両側手法は、これを `A⊗B` に拡張することで入力と出力の幾何構造を同時に表現できるが、ベクトル化した重みに直接作用させると4次の計算量が必要になる。BaKronは、互いに独立した更新を反対角線に沿ってバッチ処理し、さらに分割統治によって行列乗算を再利用することで、m×nの重み行列に対するカーネルの計算量を `O(mn(m+n))` に削減する。逐次深度は `O(m+n)` に維持される。

単一のNVIDIA RTX PRO 6000上で、float32のカスタムTritonカーネルを使用したマイクロベンチマークでは、BaKronが8192×8192行列を1.839秒で処理したのに対し、YAQAの同等実装には110.379秒を要し、カーネルは60倍高速だった。4096×14336およびその転置形状でも約43倍高速だった。ただし、GPTQは同じ8192次正方行列をわずか0.059秒で処理する。これは、GPTQが依然としてより単純な片側の幾何構造を使用しているためだ。したがって、BaKronの価値は速度でGPTQを上回ることではなく、従来はスケールさせることが難しかった両側量子化を、GPTQと同じ3次計算量のクラスにまで削減した点にある。

品質評価では、Llama 3とQwen3を重み当たり約2.81 bitに量子化した。Meta-Llama-3-8Bでは、GPTQのWikiText2パープレキシティが53.47だったのに対し、逆伝播によるK-FAC Hessianを用いたBaKronは11.90まで低下した。非量子化モデルは7.44だった。一方、このBaKron処理には1,939秒を要し、GPTQは530秒だった。これは、Hessianの蓄積、キャリブレーションデータに対する順伝播・逆伝播、Cholesky分解が、依然としてエンドツーエンドのコストを支配していることを示す。論文の評価も、単一GPU、特定のスケーリング、2.81-bit設定に限られており、公開コードも提供されていない。エンジニアリングチームは、実装の公開に加え、より多様なビット幅と実際の推論スループットによる検証を待つべきだ。