【要約】TurboQuant を numpy で実装してみたら、論文どおりの QJL でハマった話 [Zenn_Python] | Summary by TechDistill
> Source: Zenn_Python
Execute Primary Source
// Problem
筆者がTurboQuantの論文に基づき、QJLを用いた残差補正を実装した際、内積推定の精度が著しく低下するという問題に直面した。理論上のバイアス除去が、実用環境では致命的なノイズとなることが判明した。
- ・QJLを再構成ベクトルに加算する素朴な実装では、内積推定の分散がMSE-only版の6000倍以上に増大する。
- ・Softmax attentionは分散の増幅に弱いため、理論的な無偏性よりも低分散なMSE-only版が実用的に安定する。
- ・KとVでノルムが大きく異なる場合、一律の量子化では誤差がノルムの2乗に比例して増大する。
// Approach
筆者はnumpyを用いてTurboQuantのコアアルゴリズムを最小実装し、4つの実験を通じて理論と現実の差分を検証した。数学的な挙動を可視化することで、実装上の注意点を浮き彫りにしている。
- ・Lloyd-Maxアルゴリズムを用いた、MSEを最小化するスカラー量子化器の構築。
- ・ランダム直交行列による回転を用い、座標ごとの分散の偏りを解消する手法の検証。
- ・QJL(Quantized Johnson-Lindenstrauss)による残差補正の実装と、内積推定への影響評価。
- ・ベクトルのノルム差が量子化誤差に与える影響の定量的測定。
// Result
実験の結果、理論的なバイアス除去よりも、分散の抑制が実用上重要であることが判明した。論文の理論をそのまま再構成プロセスに適用することの危険性が示された。
- ・3bit量子化でもコサイン類似度0.98以上を維持できることを確認した。
- ・QJLによる残差補正は、再構成プロセスに組み込むと分散を爆発させるため、実用にはMSE-only版が適している。
- ・KとVのノルム差を考慮した非対称なビット配分が、効率的な量子化に不可欠である。
Senior Engineer Insight
> 理論的な「無偏性」が、分散の増大によって実用性を損なう典型例である。大規模推論エンジンへの実装においては、QJLを再構成プロセスに組み込むのではなく、Attention計算のカーネル内で直接内積を推定する高度な設計が求められる。また、K/Vのノルム差に基づいた非対称なビット配分は、メモリ効率と精度のトレードオフを最適化する上で必須の検討事項である。