FHEを用いた暗号化された大規模言語モデルに向けて

TL;DR

Hugging FaceとZamaは、完全準同型暗号（FHE）を使用して、大規模言語モデル（LLM）の一部を暗号化されたデータ上で実行する方法を実証しました。このアプローチにより、LLMはサービスプロバイダーに生データを見られることなく機密性の高いユーザーのクエリを処理することができ、同時に、オンプレミス展開によるモデル所有者の知的財産の漏洩を防ぐことができます。

完全準同型暗号によるLLMのプライバシー問題の解決

完全準同型暗号（FHE）は、暗号化されたデータ上で直接関数を実行することを可能にし、処理前に情報を復号する必要性を排除します。Hugging Faceの transformers ライブラリのGPT-2実装を適応させ、Concrete-Pythonを使用することで、モデルの推論プロセスにおける特定のセクションをFHEの等価物に変換できます。

TFHE（Torus Fully Homomorphic Encryption）スキームでは、モデルの重みと活性化関数は整数として表されます。非線形関数は、暗号化されたデータに対してテーブルルックアップ（TLU）を実行し、任意の計算を可能にするために暗号文をリフレッシュするProgrammable Bootstrapping（PBS）を介して処理されます。PBSは線形演算よりも計算コストが高いものの、LLMの計算のあらゆる部分、あるいは全体をFHEで表現することを可能にします。

技術的実装：アテンションヘッドの暗号化

暗号化されたLLMレイヤーを実装するために、システムはハイブリッドアプローチを採用しています。クライアントが初期のローカル推論を実行し、中間操作を暗号化してサーバーに送信します。サーバーは、暗号化されたデータに対してアテンションメカニズムの一部を適用し、クライアントが復号してローカル推論を継続できるように結果を返します。

量子化の要件

FHEは整数で動作するため、モデルの重みと活性化関数は量子化される必要があります。再学習の必要がないポストトレーニング量子化を使用することで、Zamaは4ビット量子化が元のモデル精度の96%を維持することを見出しました（約80文のデータセットに基づく）。

GPT-2との統合

実装には、量子化された演算子を含めるために特定のモジュールのフォワードパスを書き換える作業が含まれます。例えば、GPT2LMHeadModel は、最初のマルチヘッドアテンションモジュールを QGPT2SingleHeadAttention モジュールに置き換えることで変更できます。

このセットアップでは、クエリ、キー、バリュー行列の投影を含むマルチヘッドアテンションメカニズムの最初のヘッドが、FHEフレンドリーな演算子を使用して実行されます。その他の計算は浮動小数点のまま、クライアント側で非暗号化状態で実行されます。

計算複雑性とパフォーマンス

アテンションメカニズムは、クエリ、キー、バリューの乗算により、トランスフォーマーモデルの中で最も計算負荷の高い部分です。FHEでは、暗号ドメインでの乗算の複雑さと、シーケンス長が長くなるにつれて演算が二次関数的に増加することにより、このコストがさらに増大します。

主なパフォーマンスの観察結果は以下の通りです：

• 計算負荷: 長さ6のシーケンスには11,622回のPBS演算が必要です。
• 現状: 現在の実装は最適化されていない最初の実験であり、CPU上で数秒で実行されますが、かなりの計算能力を必要とします。
• 将来の展望: Zamaは、将来のASICハードウェアによってレイテンシが1,000倍から10,000倍改善され、処理時間がCPUでの数分から100ミリ秒未満に短縮される可能性があると予測しています。

結論

FHEをLLMに統合することは、ユーザーのプライバシーが完全に尊重され、モデルの知的財産が保護されるクラウドベースのAIサービスへの道を開きます。transformers ライブラリとZamaの Concrete および Concrete-ML ライブラリを活用することで、開発者は機械学習モデルをFHEの等価物に変換し、暗号化されたデータに基づいて予測を開始することができます。