Contraction-Gauge Preconditioning for Quantized Matrix Multiplication

技術要約：量子化行列積における収縮ゲージ・プリコンディショニング（Contraction-Gauge Preconditioning）

問題定義

本論文は、両方の因子 $A \in \mathbb{R}^{m \times K}$ と $B \in \mathbb{R}^{K \times n}$ が量子化されている低精度行列積（$C = AB$）における課題に対処している。主要な困難は、量子化誤差が均一に伝播しない点にある。すなわち、特定の要素$A_{ik}$の誤差は、対応する$B$の行$B_{k,:}$のエネルギーに比例して出力行$C_{i,:}$ に影響を与え、その逆も同様である。チャネルグルーピングや直交変換などの既存手法は、これらの非対称な誤差伝播と、変換されたオペランドの複数のコピーを保存するコストを考慮した、統一的なプロダクトレベルの目的関数を欠いていることが多い。著者らは、確率的な仮定の下で厳密な誤差恒等式を導出し、変数の表現を量子化前に最適化するための「収縮ゲージ（contraction-gauge）」フレームワークを定式化すると同時に、誤差削減と変換されたコピーの保存/通信コストとのトレードオフを明示的に管理することを目指している。

手法

1. 理論的基礎：積誤差恒等式（Product-Error Identity）
著者らは、既知の分散場を持つ、互いに独立で平均ゼロの要素ごとの誤差の下での、期待平方積誤差に関する厳密な有限次元恒等式を導出した。量子化された因子 $\hat{A} = A + E_A$ および $\hat{B} = B + E_B$ に対して、期待平方フロベニウスノルム誤差は以下の通りとなる：
$$\mathbb{E}\|\hat{A}\hat{B} - AB\|_F^2 = \sum_{i,k} v^A_{ik} \|B_{k,:}\|_2^2 + \sum_{k,j} v^B_{kj} \|A_{:,k}\|_2^2 + \sum_{k} \left(\sum_i v^A_{ik}\right)\left(\sum_j v^B_{kj}\right)$$
この恒等式は、非オーバーロードの減算ディザリング（subtractive dither）および残差依存の分散を持つ独立した確率的丸めに対して厳密である。これは、反対側の因子のエネルギーによって誤差を重み付けする役割を果たす。

2. 収縮ゲージ・プリコンディショニング（Contraction-Gauge Preconditioning）
本手法は、任意の可逆行列 $T \in GL(K)$ に対する代数的等価性 $AB = (AT)(T^{-1}B)$ を利用する。著者らは、量子化前に共有の収縮次元 $K$ に適用される $T$ の選択を「収縮ゲージ」と定義する。

• ゲージ・ドメイン（Gauge Domains）： 単一の出力ブロック（例：行ブロック）に単一のゲージが割り当てられた固定領域。
• 再利用記述子（Reuse Descriptor, $n_{opp}$）： プリコンディショニングのコストは、必要とされる反対側の因子の、変換・量子化された異なるコピーの数によって測定される。単一のゲージがすべてのドメインで共有される場合、$n_{opp}=1$ であり、個別のゲージが使用される場合は $n_{opp}$ が増加する。

3. 構造的ファミリーと最適化規則
著者らは、異なるゲージのファミリーに対して具体的な最適化戦略を提案している：

• 対角フォールド（スケーリング）： 共有ドメインにおいて、$(AD, D^{-1}B)$ が誤差を最小化する最適な対角スケーリング $D$ を見つける問題は、**幾何計画法（Geometric Program, GP）**として定式化される。これにより、共有フォールドのグローバルな最適選択が可能となる。また、アイデンティティ・フォールドが既に最適であるかをテストするための線形計画法（LP）も導出されている。
• パーティショニング（分割）： 精度とコピー数のバランスを取るため、本論文はスカラーノルムによるソート（ランダムなタイにより性能が線形に低下する可能性がある）を避け、類似のプロファイルを持つ行をグループ化するための対数絶対値座標におけるクラスタリングを提案している。
• 回転（Rotations）： 直交ゲージについては、ブロック局所コヒーレンスを用いて潜在的な利得を限定している。ランダムな回転（例：アダマール変換）は、座標の最大値を対数因子内で減少させられることを示している。
• 部分的回転（ハウスホルダー変換）： 重み付きエネルギーが低ランクである場合、著者らは、重み付きグラム行列のトップ固有空間をフラットなターゲットに合わせるために、少数のハウスホルダー反射器（$t$）を使用することを提案しており、そのコストは $O(tK)$ である。
• 階層構造（Hierarchies）： ブロック対角直交ゲージにおいて、スライスエネルギー反相関基準が導出されている。スライスにわたる $A$ と $B$ のエネルギープロファイルが反相関である場合、階層的な精緻化が誤差を減少させるが、そうでなければフラットな回転が好ましい。テレスコーピング恒等式により、ボトムアップの深さ選択が可能となる。

4. 量子化器のデザイン

• ビット割り当て： 固定ビット幅和制約（$b_A + b_B = b_{sum}$）の下での最適な分割は非対称である：$b_A - b_B = \frac{1}{2} \log_2(P_A/P_B)$。ここで $P_A$ および $P_B$ は、積で重み付けされた誤差係数である。
• クリッピングと診断： 本論文は、決定論的な最近接丸め（RTN）における格子整合を検出するための、切断された特性関数診断（$\Xi$）を導入しており、これはノイズモデルの独立性の仮定を損なう可能性がある。

主な貢献

1. 厳密な誤差計上： ディザリングおよび確率的丸め下での量子化行列積に関する、期待平方積誤差の厳密な恒等式の導出。これは、主要な誤差項と双線形交差項を明示的に分離している。
2. 収縮ゲージ・フレームワーク： 変換 $T$ の選択、その共有パターン（ゲージ・ドメイン）、および結果としての表現の再利用コストを結合した、プリコンディショニングのための統一された定式化。
3. グローバルに最適なフォールド選択： 境界付き共有フォールド問題が幾何計画法であり、グローバルな最適性に解けることの証明、およびアイデンティティ・フォールドの最適性をテストする厳密なLPの導出。
4. 決定統計量： すべてのファミリーの厳密な最適解を計算することなく、変換の選択を導くための計算可能な統計量の開発：
   • スケーリングのためのテイル指数（Tail index）。
   • パーティショニングのためのプロファイル拡散（Profile spread）。
   • 回転のためのブロック・コヒーレンスおよび重み付きグラム・スペクトル。
   • 階層の深さのためのスライスエネルギー共分散。
5. 再利用を考慮した設計： ゲージ共有パターンのストレージおよび通信コストを定量化するための指標として $n_{opp}$ を導入し、誤差削減と実装コストの間のトレードオフ分析を可能にした。

結果

制御された合成実験：

• スケーリング： パーティショニングにおける対数絶対値クラスタリングは、スカラーノルム・ソート（2.116倍）と比較して、プロファイル純粋参照に対してほぼ最適な広がり（1.015倍）を達成した。
• フォールド： ドメイン共有GPフォールドは、アイデンティティ・フォールドに対して約161倍、個別に最適化された片側境界に対して4.86倍、モデル化された誤差を減少させた。
• ビット割り当て： 非対称なビット分割（例：8/8ではなく9/7）は、理論的予測通り、対称分割よりも1.7倍低い誤差を達成した。
• 回転： ランダムなアダマール回転は、外れ値が注入されたヘビーテイルデータに対して最大60倍の誤差削減を提供し、ハウスホルダー反射器（$t=3$）はランク3のデータに対して完全な回転の性能に匹次した。
• 階層構造： 反相関基準は、階層的な精緻化がフラットな回転を上回るクロスオーバーポイントを正確に予測した。

学習済み分類器による検証：
手書き数字画像を用いた3ブロック画像分類器（ViT風）による検証：

• 相関： ディザリングモデルの予測は、決定論的なRTN誤差と高い中央値スペアマン相関を示した（8ビットで0.937、4ビットで0.918）。
• 性能： GPフォールドを使用することで、アイデンティティ・フォールドと比較して、保持されたプロダクト誤差が8ビットで18.0%、4ビットで20.5%（幾何平均）減少した。
• 比較： GPフォールドは、SmoothQuantスタイルのグリッドサーチ（キャリブレーションにより選択）に対し、12個のプロダクトのうち10個でより優れた結果を示し、テスト選択されたオラクルよりも低い幾何平均誤差を達成した。
• エンドツーエンド： 12個のプロダクトをすべてGPフォールドと共に量子化すると、アイデンティティ・フォールドモデルと比較して、ロジットMSEが8ビットで15.4%、**4ビットで26.4%**減少した。

意義と主張

本論文は、ヒューリスティックなスケーリングやアドホックなグルーピングを超えた、量子化行列積のための厳密かつ確率的な基礎を提供すると主張している。誤差の厳密な恒等式を確立し、それを「収縮ゲージ」フレームワークに結び付けることで、著者らは以下を実現している：

• 認定された選択： 幾何計画法による、境界付き対角ファミリー内でのグローバルな最適選択。
• 統一された目的関数： 多様な変換ファミリー（フォールド、回転、階層）とその共有パターンを評価するための共通の目的関数（積で重み付けされた誤差）。
• 実用的なトレードオフ： 変換されたコピーの保存コスト（$n_{opp}$）と誤差削減を明示的にバランスさせるためのメカニズム。これはデプロイメントにおいて極めて重要である。

著者らは、厳密な恒等式はディザリングおよび確率的丸めに対して成立するものの、導出された統計量および選択規則は、決定論的な最近接丸め（RTN）にも効果的に転移すること（学習済み分類器の実験により検証済み）を強調している。本研究は、全 $GL(K)$ 最適化問題を解決することを目的としているのではなく、制限された実用的なファミリー内において、高性能な候補を選択するための構造化された実行可能な手順を提供することを目指している。