技術的サマリー：SSA：より優れたスコアリング関数によるパフォーマンス向上

問題定義

トランスフォーマーモデルは強力なコンテキスト内学習（ICL）能力を示す一方で、分布シフト（テスト時の入力やアルゴリズム構造が訓練分布と異なる状況）への汎化においては体系的な失敗を示す。最近の文献は、大規模言語モデル（LLM）が「低確率設定」において苦労し、根本的なルールが変更されていない場合でも汎化できないことを特定している。

著者らは、制御された ICL タスクを用いて、これらの失敗のアーキテクチャ的起源を調査する。彼らは、アテンション機構内のSoftmaxスコアリング関数が、これらの汎化失敗の主要な要因であると特定する。具体的には、入力シーケンスに訓練中に観測されたものよりも著しく大きな絶対値を持つ値が含まれる場合、Softmax がアテンション飽和を引き起こすことを観察する。この現象により、アテンション重みが最大値を持つ単一のトークンにほぼ完全に収束し、タスクに不可欠な他のすべてのコンテキストトークンを事実上無視してしまう。この「アテンション収束」は、モデルがシーケンス全体にわたって情報を統合することを妨げ、論理的集約や機能的推論を必要とするタスクにおいて体系的な誤りを引き起こす。

手法

制御された実験設定

アーキテクチャの影響を事前学習のアーティファクトから分離するため、著者らは 2 つの合成 ICL タスクに対してゼロからトランスフォーマーモデルを訓練した。

1. 量化タスク: 数値のシーケンスに対する論理量化詞（「すべて」と「いくつか」）の解釈能力を評価する。
2. 線形関数タスク: コンテキスト内の例からアフィン関数 $f(x) = ax + b$ を推論する能力を評価する。

モデルは標準正規分布 $N(0, 1)$ から抽出されたデータで訓練され、より長いシーケンス長や分散が高い分布（$\sigma \ge 1$ である $N(0, \sigma)$）から抽出された入力など、分布シフト条件下でテストされた。

提案される解決策：スケーリング符号付き平均（SSA）

著者らは、Softmax の代替として学習可能な**スケーリング符号付き平均（SSA）**を提案する。SSA は、指数関数的スコアリング関数を要素ごとに適用されるパラメータ化関数に置き換える。
$$x \mapsto (1 + b|x|)^{\text{sgn}(x)n}$$
ここで、$b > 0$ および $n \ge 1$ は学習可能なスカラーパラメータである。

主要な理論的性質:

• 多項式成長対指数成長: 対数オッズの差が増加した際に「ハードマックス」分布へ急速に飽和する指数関数的に成長する Softmax と異なり、SSA は多項式的に成長する。
• 理論的保証: 著者らは定理 1 と定理 2 を用いて、グローバルな対数オッズの増幅または内部の不均衡（あるトークンのスコアが他を大幅に上回る状況）において、Softmax の重みは勝者に対して 1 に、他者に対して 0 に指数関数的に速く収束することを証明する。これに対し、SSA は多項式的にのみ収束し、二次的なトークンに対して非ゼロの「コンテキスト質量」を保持する。これにより、モデルは極端な値が存在する場合でも、コンテキストのバランスの取れた視点を維持できる。

評価範囲

本研究は、SSA を以下の 3 つの異なる設定で評価する。

1. 制御された ICL タスク: ゼロから訓練された小型トランスフォーマー（12 レイヤー、8 ヘッド）。
2. デコーダ専用ベンチマーク: FineWeb コーパスでゼロから訓練された 1 億 1400 万パラメータの Nemotron 風モデル。
3. エンコーダ専用アーキテクチャ: 子供向け発話（AO-CHILDES コーパス）で訓練された BabyBERTa の変種を、文法プロービングスイートで評価。

著者らはまた、Sparsemax、Entmax、Temperature-scaled Softmax、CosFormer、SA-Softmax などの既存の Softmax 代替案を体系的にテストし、特定のタスクにおいて標準的な Softmax よりも一貫した改善を提供するものは何もないことを発見した。

主要な結果

1. 制御された ICL タスク

• シフト下での汎化: 分布外（OOD）データでテストされた際、SSA は量化タスクおよび線形関数タスクの両方において Softmax を大幅に上回る。
• 絶対値に対する頑健性: 量化タスクにおいて、Softmax モデルはシーケンス内に単一の大きな値が現れると失敗し、その値のみに基づいて予測を行う。一方、SSA モデルはシーケンス全体にわたって情報を正しく集約し、極端な入力値が存在する場合でも低い誤り率を維持する。
• 線形関数の推論: 分布シフトされた分布から抽出された係数を持つ関数を推論する際、SSA モデルは Softmax ベースラインと比較して大幅に低い平均二乗誤差（MSE）を達成する。

2. デコーダ専用ベンチマーク

• NLP パフォーマンス: SSA で訓練された 1 億 1400 万パラメータのデコーダ専用モデルは、分布内（FineWeb）および分布外（Wikipedia）の両方のテキストにおいて、Softmax 対応モデルよりも低いパープレキシティを達成した。
• ベンチマークでの改善: 標準的な NLP ベンチマーク（ARC、BoolQ、HellaSwag など）において、SSA ベースのモデルはゼロショット評価で Softmax ベースラインを一貫して上回り、arc_easy や boolq などのタスクで顕著な改善が見られた。

3. エンコーダ専用モデル

• 文法プロービング: SSA ベースの BabyBERTa モデルは、Softmax ベースのモデルと比較して複数の文法プロービングタスク（主語 - 動詞の一致、項構造など）でパフォーマンスが向上し、その恩恵がマスク言語モデルアーキテクチャにも及ぶことを示した。

意義と主張

本論文は、Softmax 飽和が、それらのトークンがタスクに関連するかどうかに関わらず、大きな絶対値の入力に直面した際にアテンション機構が単一のトークンに収束させる根本的なアーキテクチャ的制限であると主張する。この失敗様式は、分布シフト下での ICL における不良な汎化の根本原因として特定されている。

この研究の意義は以下の点にある。

1. 具体的なメカニズムの特定: モデルの規模やデータの複雑さではなく、特定の ICL シナリオにおける汎化失敗の主要な駆動力として Softmax スコアリング関数を特定すること。
2. シンプルで効果的な代替案: 各アテンションヘッドあたりわずか 2 つのスカラーパラメータを追加するだけで、Softmax の「ドロップイン」代替として機能する SSA を導入すること。
3. 広範な適用性: SSA がデコーダ専用およびエンコーダ専用の両方のアーキテクチャ、ならびに合成および実世界の NLP タスクにおいてパフォーマンスを向上させることを実証すること。

著者らは謙虚な立場を維持し、SSA は Softmax の柔軟性の欠如に起因する失敗を軽減するが、入力と関数の分布の両方が訓練と著しく異なる極端な分布シフトを含むシナリオなど、すべての汎化課題を解決するものではないと認めている。彼らは、トークンの絶対値と関連性を混同するアテンションの基礎的な数学的構造は、将来の研究における未解決の領域であると結論付けている。