lucidrains/vector-quantize-pytorch

何を解決するか

このライブラリは、生成AIで連続データ（画像や音声など）を学習されたコードの離散集合にマッピングするために使用されるベクトル量子化（VQ）の包括的なツールセットを提供します。VQにおける一般的な課題、たとえば「死んだ」コードブックエントリ（一部のコードが一度も使われない状態）や、非微分可能な量子化層の訓練の難しさに対処します。

動作方法

このライブラリは、入力ベクトルを離散インデックスに変換し、再び量子化されたベクトルに戻すさまざまな量子化戦略を実装しています：

• 標準VQ： コードブックと指数移動平均（EMA）を使用して辞書を更新します。
• リジッドVQ： 元のベクトルと量子化されたベクトルの差分（残差）を再帰的に複数の量子化器で量子化することで、より高い精度を実現します。
• FSQ（有限スカラ量子化）： スカラを離散レベルに丸めるというシンプルなプロセスにより、学習されたコードブックやコミットメント損失の必要性を排除します。
• LFQ（ルックアップフリー量子化）： コードブックを完全に排除し、独立したバイナリラテンを使用します。
• FSP（有限スカラ摂動）： 離散化を構造化ノイズとして扱い、表現に注入することで訓練を安定化させます。
• 高度な勾配手法： 伝統的なストレートスルー推定器（STE）よりも優れた勾配フローを提供する「回転トリック」とDiVeQを実装しています。

対象ユーザー

VQ-VAE、VQ-GAN、EncodecやSoundStreamなどの音声コーデックなど、生成モデルを構築している研究者や開発者向けに設計されています。PyTorchで柔軟で高性能な量子化層が必要な方々に最適です。

特徴

• 多様な実装： 標準VQ、リジッドVQ、グループ化リジッドVQ、FSQ、LFQ、SimVQをサポート。
• コードブックの健全性： コードブックの崩壊を防ぐ技術として、低いコードブック次元、コサイン類似度、古くなったコードの置換を含む。
• 柔軟な初期化： 最初のバッチデータに基づいてk-meansによるコードブック初期化をサポート。
• マルチプロセス対応： 複数プロセス間でコードブックを自動同期します。