プロンプトによるプライバシー漏洩が心配ですか?このフレームワークにより、LLaMA-7Bは安全性の推論を実行できる。

現在、ディープラーニングサービスを提供する事業者は数多く存在します。これらのサービスを利用する際には、利用者は自分の情報をプロンプトに記入してサービス提供者に送信する必要があり、プライバシー漏洩などの問題が生じます。一方、サービスプロバイダーは、苦労してトレーニングしたモデルパラメータを公開したがらないのが一般的です。

この問題に対処するために、Ant Group の研究チームは、モデルのパフォーマンスに影響を与えずに安全な推論を実現できる PUMA フレームワークを提案しました。それだけでなく、関連コードもオープンソース化しました。

• 論文: https://arxiv.org/abs/2307.12533
• コード: https://github.com/secretflow/spu/blob/main/examples/python/ml/flax_llama7b/flax_llama7b.py

事前トレーニング済みの Transformer モデルは、多くの実用的なタスクで優れたパフォーマンスを発揮し、多くの注目を集めています。現在、このテクノロジーをベースにしたツールが数多く存在し、多くの場合、「Deep Learning as a Service (DLaaS)」パラダイムでサービスを提供しています。ただし、これらのサービスにはプライバシーの問題が伴うことがよくあります。たとえば、ユーザーが ChatGPT を使用する場合、ユーザーはサービス プロバイダーに独自のプライベート プロンプトを提供するか、サービス プロバイダーがユーザーに独自のトレーニング済み重み構成を提供する必要があります。

Transformer モデル サービングのプライバシー問題に対処するための 1 つの解決策は、推論中のデータとモデルの重みのセキュリティを確保できるセキュア マルチパーティ コンピューティングです。しかし、マルチパーティコンピューティング (MPC) における単純な基本 Transformer 推論の時間と通信コストは非常に高く、実際のアプリケーションで使用するのは困難です。より高い効率を実現するために、いくつかの研究結果では、Transformer モデルのセキュリティ推論を高速化するさまざまな方法が提案されていますが、これらの方法には依然として次の 1 つ以上の欠点があります。

• 交換は困難です。最近のいくつかの研究では、コストを削減するために、GeLU や softmax などの高コスト関数の代わりに、2 次関数や ReLU 関数などの高速近似法を使用できると提案されています。ただし、これらの関数を単純に置き換えると、Transformer モデルのパフォーマンスが大幅に低下し (追加のモデルトレーニング、つまり微調整が必​​要になる場合があります)、デプロイメントの問題が発生する可能性があります。
• 推論コストは高いです。いくつかの研究では、より正確な多項式関数を使用して高コストの非線形関数を近似することが提案されていますが、それらの近似方法では、GeLU と Softmax の特殊な特性が考慮されていません。したがって、近似値を使用した後でも、この方法のコストは依然として高くなります。
• 展開が簡単ではありません。最近のいくつかの研究では、埋め込みプロセスを分解し、線形層を再編成するなど、Transformer モデル アーキテクチャを変更することで、安全な推論を加速することが提案されています。さらに悪いことに、Crypten フレームワークは安全な LayerNorm をサポートしていないため、BatchNorm のみを使用してコストをシミュレートすると、セキュリティ推論の結果が不正確になります。これらの変更は、既存のプレーンテキスト Transformer システムと競合します。

まとめると、MPC Transformer 推論の分野では、モデルのパフォーマンスと効率の両方を実現することは難しく、次のような疑問が生じる可能性があります。

事前トレーニング済みの大規模なトランスフォーマー モデルを、さらに再トレーニングすることなく、プレーンテキスト モデルと同様の精度を達成しながら、安全かつ効率的に評価できますか?

Ant Group が提案した PUMA フレームワークは、この問題を解決するために作成されました。このフレームワークは、エンドツーエンドのセキュアな Transformer 推論を安全かつ正確に実行できます。この論文の主な貢献は次のとおりです。

• 非線形関数を近似する新しい方法。この論文では、Transformer モデルで高価な非線形関数 (GeLU や Softmax など) を近似するために使用できる、より正確で高速な近似方法を提案します。従来の方法とは異なり、新しく提案された近似法は、これらの非線形関数の特殊な特性に基づいており、精度と効率のバランスをとることができます。
• より高速かつ正確なセキュリティ推論。研究者らは、6 つのトランスフォーマー モデルと 4 つのデータセットを使用して広範な実験を実施し、その結果、PUMA フレームワークを使用すると、精度はプレーンテキスト モデルに近くなり、速度と通信効率が約 2 倍向上することが示されました (ただし、MPCFORMER は PUMA ほど正確ではないことに注意してください)。 PUMA は、LLaMA-7B の評価を 5 分以内に完了して 1 つの単語を生成することもできます。著者らは、このような大規模な言語モデルの評価に MPC が使用されたのは今回が初めてであると述べています。
• オープンソースのエンドツーエンドのフレームワーク。 Ant Group の研究者たちは、MPC 形式で安全な埋め込みおよび LayerNorm プログラムを設計し、実装することに成功しました。その結果、PUMA のワークフローは、モデル アーキテクチャを変更することなくプレーンテキスト Transformer モデルに従い、事前トレーニング済みのプレーンテキスト Transformer モデル (Huggingface からダウンロードしたものなど) を簡単に読み込んで評価できるようになります。著者らは、これは再トレーニングなどのさらなる変更を必要とせずに、事前トレーニング済みの Transformer モデルの正確な推論をサポートする最初のオープンソース MPC ソリューションであると述べています。

PUMAの安全設計

PUMAの概要

PUMA の設計目標は、Transformer ベースのモデルが安全に計算を実行できるようにすることです。これを実現するために、システムはモデル所有者、クライアント、計算パーティという 3 つのエンティティを定義します。モデル所有者はトレーニング済みの Transformer モデルを提供し、クライアントはシステムにデータを提供して推論結果を受け取る責任を負い、コンピューティング パーティ (つまり、P_0、P_1、および P_2) はセキュア コンピューティング プロトコルを実行します。モデルの所有者とクライアントも計算パーティになることができますが、説明を簡単にするためにここでは区別しています。

安全な推論中は、重要な不変条件を維持する必要があります。計算は常に、クライアント入力の 3 分の 2 とモデルのレイヤー重みの 3 分の 2 の機密シェアで開始され、それらのレイヤーの出力の 3 分の 2 の機密シェアで終了します。これらの共有は当事者に情報を漏らさないため、これらのプロトコル モジュールを任意の深さまで順番に組み合わせることができ、Transformer ベースのモデルに対して安全な計算が可能になります。 PUMA が重点を置く主な課題は、必要なセキュリティ レベルを維持しながら、さまざまな計算パーティ間の実行コストと通信コストを削減することです。 PUMA は、複製された秘密共有と新しく提案された 3PC プロトコルを活用することで、3 者間設定での Transformer ベースのモデルの安全な推論を可能にします。

セキュアな組み込みプロトコル

現在の安全な埋め込みプロセスでは、クライアントがトークン ID を使用してワンホット ベクトルを作成する必要がありますが、これはプレーンテキスト ワークフローから逸脱し、Transformer 構造を破壊します。したがって、この方法は、実際の Transformer モデル提供アプリケーションに展開するのは簡単ではありません。

この問題を解決するために、研究者らは新しい安全な埋め込み設計を提案しました。トークンid∈[n]とすべての埋め込みベクトルを と表記すると、埋め込みは と表すことができます。 (id, E) は共有秘密なので、提案された安全な埋め込みプロトコルは次のように動作します。

• クライアントから ID ベクトルを受信した後、コンピューティング パーティはワンホット ベクトルを安全に計算します。具体的には、i∈[n]です。
• 各計算パーティは、安全な切り捨てを必要としない によって埋め込みベクトルを計算できます。

このように、ここでの Π_Embed は、Transformer モデルのワークフローを明示的に変更する必要はありません。

セキュアなGeLUプロトコル

現在のアプローチのほとんどは、GeLU 関数をより小さな関数で構成されているものとして扱い、その各部分を最適化しようとするため、プライベート GeLU 全体を最適化する機会を逃してしまいます。 GeLU関数が与えられた場合:

過去のアプローチでは、tanh 関数の効率的なプロトコルを設計するか、指数関数と逆関数の既存の MPC プロトコルをシグモイドに適合させることに重点が置かれてきました。

しかし、これらの方法のいずれも、GeLU関数が両辺とも本質的に線形であるという事実、すなわちx < −4の場合にはGeLU (x) ≈ 0であり、x > 3の場合にはGeLU (x) ≈ xであるという事実を利用していません。研究者らは、低次多項式の区分近似は、GeLUの短い区間[−4,3]において、より効率的で実装が容易なセキュリティプロトコルオプションであると提案した。具体的には、この区分低次多項式は式（２）で示される。

多項式F_0とF_1の計算は、式(3)に示すようにソフトウェアライブラリnumpy.ployfitによって実装されている。研究者らは、この多項式フィッティングは単純であるにもかかわらず、予想外に良好なパフォーマンスを発揮することを発見しました。実験結果の最大誤差は 0.01403 未満、中央値誤差は 4.41e−05 未満、平均誤差は 0.00168 未満です。

数学的に言えば、秘密の入力が与えられた場合、新しく提案された安全な GeLU プロトコルの構築は、以下のアルゴリズム 1 に示されています。

セキュア ソフトマックス プロトコル

関数 において、重要な課題は Softmax 関数を計算することです (ここで、M はバイアス行列として考えることができます)。数値安定性については、Softmax は次のように計算できます。

ここで、入力ベクトル x の最大要素です。通常の平文ソフトマックスの場合、ε = 0 です。 2次元行列の場合、各行ベクトルに式(4)が適用されます。

アルゴリズム 2 は、新しく提案されたセキュリティ プロトコル Π_Softmax の詳細な数学的説明を提供します。ここでは、2 つの最適化方法が提案されています。

• 最初の最適化は、式(4)のεをε=10^-6のような非常に小さい正の値に設定し、式(4)の累乗演算の入力がすべて負の値になるようにすることです。研究者たちは、これらの負のオペランドを利用して速度を上げました。具体的には、テイラー級数を使用して、単純なクリッピングによって累乗を計算します。

• 研究者らが提案する2番目の最適化は、分割数を減らすことで、最終的には計算コストと通信コストを削減します。この目的のために、研究者らは、サイズnのベクトルxに対して、Div(x, Broadcast(y))演算をx・Broadcast(1/y)に置き換えました。この置換により、n 回の除算が 1 回の逆数演算と n 回の乗算に実質的に削減されます。この最適化は、Softmax 操作に特に有益です。固定小数点値の設定では、Softmax 演算の 1/y は十分な精度を維持するのに十分な大きさです。したがって、このような最適化により、精度を確保しながら計算コストと通信コストを大幅に削減できます。

セキュア レイヤー標準プロトコル

サイズ n のベクトル x が与えられると、 (γ, β) はトレーニングされたパラメータであり、 であることを思い出してください。 MPC では、平方根による除算の式を評価することが重要な課題です。この式を安全に評価するために、CrypTen のアプローチでは、平方根、逆数、乗算の順序で MPC プロトコルを実行します。しかし、アント・グループの研究者たちは観察した。 MPC 側では、逆平方根 σ^{-1/2} を計算するコストは、平方根演算のコストに近くなります。さらに、前のセクションの2番目の最適化に触発されて、彼らは最初にσ^{-1/2}を計算し、次に(σ^{-1/2})をブロードキャストして、高速で安全なLayerNorm(x)をサポートすることも提案しました。アルゴリズム 3 は、Π_LayerNorm プロトコルの数学的形式を示します。

実験的評価

図 1: GLUE および Wikitext-103 V1 ベンチマークでのパフォーマンス。モデルに関しては、a は Bert-Base、b は Roberta-Base、c は Bert-Large、d には GPT2-Base、GPT2-Medium、GPT2-Large が含まれます。

表 1: 長さ 128 の入力文に対する Bert-Base、Roberta-Base、Bert-Large のコスト。時間コストは秒単位で計算され、通信コストは GB 単位で計算されます。

表 2: GPT2-Base、GPT2-Medium、GPT2-Large のコスト。入力文の長さは 32 で、これらは 1 つのトークンを生成するコストです。

表3: {2、4、8、16}文のバッチに対するBert-BaseとGPT2-Baseのコスト。 Bert-BaseとGPT2-Baseの入力長はそれぞれ128と32に設定され、GPT2のデータは1トークンを生成するコストになります。

表 4: 異なる入力長での Bert-Base と GPT2-Base のコスト (#Input)。 Bert-BaseとGPT2-Baseの入力長はそれぞれ{64, 128, 256, 512}と{16, 32, 64, 128}に設定されています。GPT2のデータは1トークンを生成するコストです。

図 2: 入力長が 32 の場合に GPT2-Base によって異なる出力トークンを生成するコスト。 a は実行時間コスト、b は通信コストです。

表 5: LLaMA-7B でセキュア推論を実行するコスト。#Input は入力文の長さを表し、#Output は生成されたトークンの数を表します。

LLaMA-7B の拡張にはわずか 5 分しかかかりません。研究者らは、それぞれ 128 スレッド、1 TB RAM、20 GB 帯域幅、0.06 ms ラウンドトリップ時間を備えた 3 台の Alibaba Cloud ecs.r7.32xlarge サーバーで PUMA を使用して大規模言語モデル LLaMA-7B を評価しました。表 5 に示すように、PUMA は合理的なコストで安全な推論のための大規模言語モデル LLaMA-7B をサポートできます。たとえば、8 つのトークンで構成される入力文が与えられた場合、PUMA は 1.865 GB の通信コストで、約 346.126 秒で 1 つのトークンを出力できます。研究者らによると、MPCアプローチがLLaMA-7Bに対して評価されたのは今回が初めてだという。

PUMA は一連のブレークスルーを達成しましたが、それはまだ学術的な成果であり、その理論が実装されるまでにはまだまだ時間がかかります。研究者たちは、将来、機械学習分野の最新の定量技術とハードウェア分野の最新のハードウェア加速技術を組み合わせることで、プライバシーを真に保護する大規模モデルサービスが実現するのもそう遠くないと考えています。