🤖 machine learning

Enjoy Your Layer Normalization with the Computational Efficiency of RMSNorm

本論文は、数学的に中心化操作を上位の線形層に折り込むことで、任意の深層ニューラルネットワークにおけるレイヤー正規化（LN）層をより効率的なRMSNormに特定し変換するフレームワークを提案し、これによりモデル性能を損なうことなく2%から12%の正確またはほぼ正確な推論加速を実現する。

原著者： Yuxin Guo, Yihao Yue, Yunhao Ni, Yizhou Ruan, Jie Luo, Wenjun Wu, Lei Huang

公開日 2026-05-15

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Yuxin Guo, Yihao Yue, Yunhao Ni, Yizhou Ruan, Jie Luo, Wenjun Wu, Lei Huang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたが非常に忙しいレストラン（ディープニューラルネットワーク）を運営していると想像してください。すべての料理（データサンプル）は、顧客に提供される前に完璧に準備される必要があります。

このレストランには、**レイヤー正規化（LN）**と呼ばれる特定の工程があります。これは、すべての料理を一口ずつ試食し、バッチ全体の平均的な味を確認してから、すべての材料を調整して「平均的な味」が正確にゼロになるようにする、シェフの役割と考えると良いでしょう。これにより、料理はバランスよく一貫性のあるものになります。しかし、この「試食と調整」の工程は遅いです。すべての料理の平均的な味を計算するのに時間がかかり、特に何千もの注文が殺到している場合、キッチン全体が遅延します。

スピードを上げるために、RMSNormと呼ばれるより高速な方法が考案されました。これは、「平均的な味の試食」という工程を完全に省略する、副料理長の役割のようなものです。彼らは材料の「量」や「強度」を確認し、それを拡大または縮小するだけです。平均値を見つけるための計算をしなくて済むため、はるかに高速です。しかし、注意点があります。「平均をゼロにする」という工程を省略するため、料理が時々塩辛すぎたり、味が薄すぎたり（不安定になったり）し、最終的な味は元のシェフのものほど良くない可能性があります。

大きな問い：
副料理長（RMSNorm）のスピードを維持しつつ、シェフ（LN）がもたらす完璧なバランスを失うことなく得ることはできるでしょうか？

この論文の解決策：「折りたたみ可能な」レシピ
この論文の著者たちは言います：はい、ただしキッチンが特定の仕様に設定されている場合に限ります。

彼らは**「折りたたみ（folding）」**と呼ばれる巧妙なトリックを提案しています。

問題点： 通常、シェフを副料理長に単純に差し替えることはできません。なぜなら、シェフはデータを中心化するという、キッチンの他の部分が依存している特定の役割を果たしているからです。
トリック： 著者たちは、シェフに入ってくる材料がすでに完璧にバランスが取れている（平均がゼロ）場合、シェフは「試食と調整」の部分を実行する必要がないことに気づきました。彼らはそれをスキップし、代わりにスケーリング部分（RMSNorm）を実行するだけで済みます。
実行方法： 彼らは、材料がシェフに到達する前に「事前にバランスを取る」方法を見つけました。これは、前の調理ステーション（線形層）のレシピをわずかに調整することで行われます。彼らはこれを**カラムベースの重み中心化（CBWC）**と呼びます。
- 前のステーションをミキサーだと想像してください。著者たちは、そのミキサーを調整して、すでに完璧にバランスの取れた材料を自動的に排出するようにします。
- 材料がすでにバランスが取れているため、シェフ（LN）を高速な副料理長（RMSNorm）に置き換えても、料理の最終的な味は変わりません。

「マップ」（ゼロ平均グラフ）
すべてのキッチンの設定が同じわけではありません。材料が複雑に混合され、分割され、奇妙な方法で再結合される経路を持つキッチンもあります。

著者たちは、キッチンのレイアウトを確認するためのマップ（グラフアルゴリズム）を作成しました。
マップが、シェフに到達する材料が、事前に調整可能なミキサーの「直線」から来ていることを示している場合、そのシェフは**「折りたたみ可能」**です。
経路があまりにも複雑な場合（バランスを崩すような方法で材料が接着されているなど）、折りたたむことはできません。

彼らが発見したこと：

速度： 人気のある AI モデル（GPT-2、BERT など）にこの手法を適用した結果、「提供」フェーズ（推論）中にキッチンが2% から 12% 高速化されることがわかりました。
味：決定的なことに、料理の味は完全に同じです。顧客（ユーザー）は品質の違いに気づきません。
トレーニング： キッチンが学習している間（トレーニング中）にもこれをテストしました。キッチンが混沌として学習している間は、完璧な「折りたたみ」の条件が常に満たされるわけではありませんが、彼らの手法は遅いシェフとほぼ同等に機能し、はるかに高速でした。

まとめ：
この論文は、AI モデルのどの部分を、遅く慎重な正規化ステップから、速く単純なステップに差し替えることで高速化できるかを示すルールブックとツールを提供しています。彼らは、遅い部分を数学的に前のステップに「隠す」ことで、AI の知能や精度を損なうことなく、より高速に実行できるようにします。これは、同じ目的地に到達するが、歩く時間が少ない迷路の抜け道を見つけるようなものです。

技術的概要：RMSNorm の計算効率を活かしてレイヤー正規化を享受する

問題定義

レイヤー正規化（LN）は、特にトランスフォーマーにおける現代の深層学習アーキテクチャの基盤的コンポーネントであり、トレーニングの安定化と最適化の改善に寄与しています。しかし、推論段階において、LN はサンプルごとの中心化（平均の減算）とスケーリングを必要とし、正規化を不要とする代替手段と比較して無視できない計算オーバーヘッドをもたらします。

RMSNorm は、中心化操作を除去しスケーリングのみを保持することで、より効率的な代替手段として提案されました。RMSNorm は大幅な高速化（先行研究では実行時間の 7%〜64% の削減が報告されている）を提供しますが、中心化操作を廃棄します。この除去により、出力範囲の安定性が低下し、活性化値の変動が大きくなる可能性があり、先行文献（LeCun ら、1990; Schraudolph、1998; Ba ら、2016）で確立された中心化に関連する恩恵が失われる可能性があります。LN と RMSNorm を同等にする以前の試み（例：Jiang ら、2023）は、特定のプリノーム・トランスフォーマー構造に限定されており、モデルの機能やトレーニングダイナミクスを変更することなく、任意の深層ニューラルネットワーク（DNN）において LN を RMSNorm に置換できる一般的なフレームワークが存在するかどうかという問いは未解決のままです。

手法

本論文は、任意の DNN における LN レイヤーを、推論およびトレーニングの両方で数学的同等性を維持しつつ RMSNorm に置換可能かどうかを判定する一般的なフレームワークを提案します。中核的な戦略は、LN の中心化操作を上位の線形レイヤーに「折りたたむ」ことです。

1. 列中心化制約（CCC）と列ベース重み中心化（CBWC）

著者はまず、LN が線形レイヤーの直後に続く場合を分析します。線形レイヤーの出力がニューロン間で平均ゼロである場合、LN における中心化ステップは冗長になることが観察されます。

列中心化制約（CCC）: 重み行列の各列の要素の和がゼロである場合、その線形レイヤーは CCC を満たします。この制約の下では、任意の入力に対して線形レイヤーの出力が平均ゼロになることが保証されます。
列ベース重み中心化（CBWC）: モデルの容量を制限することなくトレーニング中に CCC を強制するために、著者は再パラメータ化手法を導入します。代理パラメータ行列 $W$ を使用し、これを射影を通じて有効な重み行列 $V$ に変換します： $V = (I - \frac{1}{m}\mathbf{1}_m\mathbf{1}_m^\top)W$ 。この変換により、有効な重みがゼロ和列制約を満たしつつ、元のモデルの最適化ダイナミクスを保持することが保証されます。

命題: CBWC を用いた線形レイヤーの後に RMSNorm を用いるモデルは、再パラメータ化なしの同じ線形レイヤーの後に標準的な LN を用いるモデルと数学的に同等です。

2. 任意の DNN への一般化：折りたたみ可能な LN

このフレームワークは、グラフベースの分析を用いて任意のネットワーク構造に拡張されます。

折りたたみ可能な LN: 入力の上流レイヤーに対する制約を通じて、その入力が平均ゼロになることが保証される場合、その LN は「折りたたみ可能」と定義されます。
ゼロ平均グラフ: 著者は、LN へのデータフローを表す有向グラフを構築します。LN から逆方向に遡ることで、「リーフノード」（上流レイヤー）を特定します。
- すべてのリーフノードが一般的な線形レイヤー（共有重みを持つ畳み込みや RNN を含む）または本質的に平均ゼロの出力を保証するレイヤーに対応する場合、その LN は折りたたみ可能です。
- スカラー操作（推論モードにおけるドロップアウト、温度スケーリングなど）は、ゼロ平均の性質を保持し、折りたたみ可能性を破りません。
- 残差接続（加算）は、各ブランチが個別に折りたたみ可能であれば、折りたたみ可能性を保持します。
- 非線形操作（ReLU、Softmax など）または結合は、一般的にゼロ平均の性質を破り、補助的な中心化が挿入されない限り折りたたみを妨げます。

3. 検出アルゴリズム

本論文では、計算グラフの前方トラバースを実行して、折りたたみ可能な LN を自動的に検出し、CBWC が必要となる特定の上流の一般的な線形レイヤーを特定するアルゴリズム（アルゴリズム 1）を提示します。このプロセスは計算的に軽量であり、単一の前方パスのみを必要とします。

主要な貢献

一般的なフレームワーク: 先行研究がプリノーム・トランスフォーマーに限定されていた制約を超え、任意の DNN における LN の RMSNorm への置換可能性を判定する理論的フレームワーク。
数学的同等性: LN の中心化操作を上位の重みに正確に折りたたむことを可能にする CBWC の導入により、トレーニングおよび推論の両方で、CBWC+RMSNorm 方式が元の LN ベースのモデルと同一の出力および勾配を生成することを保証。
折りたたみ可能な LN の検出: 「折りたたみ可能な LN」の定義と、与えられたアーキテクチャ内のどの LN が変換可能かを特定するグラフベースの検出アルゴリズムの開発。
実用的な拡張: 正確な折りたたみ条件が満たされない場合（例：トレーニング中のドロップアウトなど）の処理に関する議論。正確な同等性が部分的に破れた場合でも、CBWC+RMSNorm がバニラ LN と競合する性能を維持することを示唆。

実験結果

著者は、複数のアーキテクチャ（GPT-2、BERT、OPT、Phi、ViT、Swin-Tiny）およびタスク（テキスト分類、翻訳、画像分類、ファインチューニング）において、手法を検証しました。

推論の高速化:
- 本手法は、モデルの予測を変更することなく、代表的なモデル（GPT-2、BERT、OPT など）において、エンドツーエンドの推論速度を**2% から 12%**向上させました。
- 高速化は、より長いシーケンス長で顕著に現れます（BERT において 16,384 トークンの場合、最大 24.95% の高速化）。
- スループットは大幅に増加します（A100 GPU 上で +1,119 から +3,385 トークン/秒）。
- 折りたたみアルゴリズムのオーバーヘッドは最小限であり（数回の推論ステップ後に償却）、メモリフットプリントを増加させません。
トレーニングおよびファインチューニングのパフォーマンス:
- 正確な折りたたみ条件が成立するシナリオ（例：推論）では、本手法はベースラインと同一です。
- ドロップアウトなどの条件が正確な同等性を破る実用的なトレーニング設定において、CBWC+RMSNormは、標準的な LN とバニラ RMSNorm の間で一貫して性能を発揮します。
- 画像分類（Swin-Tiny を用いた ImageNet-100）において、本手法は 3 つの変種（LN、RMSNorm、CBWC+RMS）の中で最高のテスト精度と損失を達成しました。
- テキスト生成およびファインチューニングタスク（Alpaca 上の GPT-3）において、本手法は効率性の向上を提供しつつ、最高のテスト損失およびパープレキシティを達成しました。

意義と主張

本論文は、レイヤー正規化の計算コストを削減しつつ、その理論的および実用的な恩恵を可能な限り維持する原理的な手法を提供すると主張しています。

正確な同等性: 「折りたたみ可能な」LN に対して、本手法は数学的に正確な置換を提供し、再トレーニングや性能の低下なしに、元々 LN でトレーニングされたモデルへの RMSNorm の展開を可能にします。
効率性: 広く使用されているアーキテクチャにおいて測定可能なエンドツーエンドの推論高速化（2%〜12%）を実現し、大規模言語モデル（LLM）やビジョン・ランゲージモデル（VLM）の大規模展開において重要です。
堅牢性: 正確な折りたたみ条件が完全に満たされない場合（例：トレーニング中のドロップアウト時）でも、本手法は標準的な LN と競合する性能を維持し、忠実度と効率性の間で有利なトレードオフを提供します。
一般性: 先行のアプローチとは異なり、このフレームワークは特定のトランスフォーマー変種だけでなく、任意の DNN に適用可能であり、プロセスを自動化する検出アルゴリズムを含んでいます。

著者は、分析が代表的なアーキテクチャを網羅しているものの、すべての現代の大規模モデルを網羅しているわけではないことを認めています。トレーニング中のドロップアウトなどのモジュールを持つアーキテクチャでは正確な同等性が失敗する可能性があることに言及し、今後の研究でそのようなモジュールに対する処理を精緻化できる可能性を示唆しています。