技術的概要：PowerSoftmax：暗号化データ上での安全な LLM 推論に向けた取り組み

問題定義

現代のプライバシー保護型大規模言語モデル（LLM）は、しばしば復号化なしに暗号化データ上で計算を可能にする同型暗号（HE）、特に CKKS 方式に依存しています。しかし、CKKS は多項式演算のみをサポートします。標準的なトランスフォーマー・アーキテクチャは、非多項式コンポーネント、特に自己注意（self-attention）における Softmax 関数に大きく依存しており、これには指数関数、除算、対数演算（数値的安定性のための log-sum-exp 法など）が含まれます。

このギャップを埋める既存のアプローチは、重大な限界に直面しています：

1. 事後近似（PTA）： 事前学習済みモデルにおいて、高次多項式を用いて非多項式関数を直接近似することです。これは、高い乗算深度と蓄積されたノイズにより、非効率的な推論をもたらすことがよくあります。
2. 近似考慮型トレーニング（AAT）： トレーニング前に、非多項式プリミティブを近似しやすい代替手段（例えば、ポイントワイズ注意やスケーリングされた ReLU）に置き換えることです。小規模モデルには効果的ですが、これらの手法は安定性とスケーラビリティにおいて困難に直面しており、数十億パラメータを持つ LLM や深いアーキテクチャ（例：32 レイヤー）のサポートに失敗しています。

核心的な課題は、標準的な Softmax 型トランスフォーマーの表現力、推論能力、およびコンテキスト内学習（ICL）を、数十億パラメータ規模で維持しつつ、HE の制約（低次多項式、安定した数値範囲）と本質的に互換性のあるトランスフォーマー・アーキテクチャを設計することです。

手法

著者は、Softmax の核心原則を再現しつつ多項式近似に適した、新しい HE 友好型の自己注意バリアントであるPowerSoftmaxを提案します。この手法は、3 つの明確な段階を含みます：

1. HE 友好型注意設計

著者は、Softmax の本質的な特性（[0, 1] への正規化、スコア差の増幅、順序の保存）を抽出します。そして、Softmax における指数関数 $\exp(x)$ を、べき関数 $x^p$（ここで $p$ は正の偶数）に置き換えます。
提案された注意メカニズムは以下のように定義されます：
$$\text{PowerSoftmax}(x)_j = \frac{x_j^p}{\sum_i x_i^p}$$
この定式化は、Softmax の直接的多項式近似に内在する指数関数やゼロ除算の問題を回避します。

2. トレーニングの安定性と数値的ロバスト性

べき関数によるトレーニングの不安定性（入力 $>1$ で発散し、入力 $<1$ で消失する）に対処するため、著者は 2 つの主要な修正を導入します：

• 安定型バリアント： log-sum-exp 法に触発され、べき関数を適用する前に定数 $c = \|x\|_\infty + \delta$ によって入力を正規化します。これにより、入力が管理可能な範囲 $(0, 1)$ に留まり、オーバーフロー/アンダーフローを防ぎ、浮動小数点精度を維持します。
• $\epsilon$-有界除算： 非多項式であり、ゼロ付近で敏感な除算演算を処理するため、分母を $\epsilon + \sum x_i^p$ に修正します。これにより $1/\epsilon^2$-リプシッツ連続関数が生成され、その後の除算の多項式近似が大幅に扱いやすくなり、必要な乗算深度が減少します。

3. 長さ非依存近似

標準的な除算近似（例えば Goldschmidt アルゴリズム）は、系列長さ $L$ に比例して線形にスケールする注意スコアの和の無制限性に対して困難をきたします。著者は長さ非依存 PowerSoftmaxを提案します：
$$\text{Length-Agnostic}(x)_j = \frac{x_j^p}{L \cdot \text{Mean}(x^p)}$$
$L$（系列長さ）は公開情報であるため、項 $1/L$ は近似なしに直接計算できます。モデルは和の除算ではなく、平均の除算を近似するため、近似領域は系列長さに関係なく有界に保たれます。

4. 構築パイプライン（アルゴリズム 1）

多項式 LLM を作成するための完全なパイプラインは以下の通りです：

1. アーキテクチャの修正： Softmax を安定型 PowerSoftmax バリアントに置き換え、範囲最小化の目的関数を用いてゼロから（または継続的事前学習を通じて）トレーニングします。
2. 範囲最小化： 非多項式層（PowerSoftmax、GELU、LayerNorm）の入力範囲を制約する専用の損失関数を適用し、低次多項式近似が正確であることを保証します。
3. 多項式置換： トレーニング後、残りの非多項式演算（除算、逆平方根）を低次多項式近似（例：Goldschmidt 法）に置き換え、完全な多項式モデルを生成します。

主要な貢献

1. PowerSoftmax メカニズム： 指数関数スケーリングを多項式スケーリングに置き換える新しい自己注意バリアント。HE 環境に特化しつつ、注意の核心設計原則を維持します。
2. スケーラブルな多項式 LLM： 10 億パラメータを超える（具体的には 32 レイヤー、14 億パラメータ）最初の多項式 LLM の開発。これらのモデルは、同規模の標準トランスフォーマーと同等の推論能力とコンテキスト内学習（ICL）能力を示します。
3. 安定したトレーニングおよび推論戦略： 以前の多項式トランスフォーマーアプローチの不安定性とスケーラビリティの問題を克服する、数値的に安定したトレーニングバリアントと長さ非依存の推論戦略の導入。
4. 包括的な評価： 暗号化計算の遅延の詳細な内訳と、設計選択（べきパラメータ $p$ や $\epsilon$-有界除算の影響など）を検証するアブレーション研究。

実験結果

著者は、因果的（GPT/Pythia）および双方向的（RoBERTa）アーキテクチャ上で手法を評価しました：

• パフォーマンス： 14 億パラメータの多項式モデルは、ゼロショットおよび少数ショットベンチマーク（Lambada、PIQA、WinoGrande、ARC など）において、元の Softmax 型 Pythia モデルと同等のパフォーマンスを達成しました。特に、以前多項式 LLM では観測されていなかった能力である、ARC-Challenge などの推論ベンチマークで競争力のあるパフォーマンスを示しました。
• 双方向モデル： 多項式 RoBERTa-Base バリアントは、GLUE タスク（SST-2、QNLI、MNLI）において、元の Softmax 型 RoBERTa と比較して約 1% の低下のみを示し、Nexus などの以前の多項式アプローチを上回りました。
• HE 推論： HElayers フレームワークを使用した Pythia 70M モデルの実験により、暗号化推論は平文推論とほぼ同一の結果（最大 MSE 損失 0.005）をもたらすことが確認されました。
• 効率性： 提案手法は、最先端の代替手段（Nexus、Cho ら）と比較して乗算深度を大幅に削減しました。128 トークンの系列の場合、PowerSoftmax は Taylor/Goldschmidt アプローチと比較して約 9.7 倍、平方戦略アプローチと比較して約 20 倍の遅延削減を実現しました。

意義と主張

この論文は、完全な多項式 LLM が、推論能力やコンテキスト内学習能力を犠牲にすることなく、10 億パラメータ規模まで拡張可能であることを実証することで、プライバシー保護機械学習の分野における画期的な進展を主張しています。

• スケーラビリティ： 小規模モデルや単純な分類タスクに限定されていた以前の研究とは異なり、このアプローチは 14 億パラメータを持つ深い（32 レイヤー）トランスフォーマーのトレーニングに成功しました。
• 非対話型セキュリティ： 完全多項式モデルを用いた完全準同型暗号（FHE）を利用することで、マルチパーティ計算（MPC）や対話型プロトコルに関連する通信オーバーヘッドや潜在的な脆弱性を回避する、非対話型推論を可能にします。
• 実用性： 詳細な遅延分析と 14 億パラメータモデルの成功した展開は、安全でプライバシーが保護された LLM 推論が実用的に可能になりつつあることを示唆しており、医療や金融などのデータに敏感な分野への応用への道を開いています。

著者は、モデルが良好に機能している一方で、暗号化データ上での逐次デコーディングにおける自己回帰的生成能力の完全な評価は、今後の研究の方向性であると認めています。