[翻訳] MoD(Mixture-of-Depths): Transformerベース言語モデルの計算最適化に向けたアプローチ、そしてMoDE(MoD+MoE)

(discuss.pytorch.kr)

7 ポイント投稿者 ninebow 2024-04-13 | まだコメントはありません。 | WhatsAppで共有

PyTorchKR

Jamba や Qwen1.5-MoE、あるいは DBRX のように、複数の「専門家」ブロックの中から1つ以上を選んで入力 token の処理を進める MoE(Mixture-of-Experts) 手法が近年広く使われています。こうした MoE 方式は、巨大な LLM の一部を動的に調整することで計算資源を効率よく割り当て、より高い性能を示しています。
このように MoE が複数の専門家の中から一部を選ぶことで「モデルの 幅(width)」を動的に調整していたのに対し、今回紹介する MoD(Mixture-of-Depths) は「モデルの 深さ(depth)」を動的に調整するアプローチを取ります。つまり、特定のトークンに対してモデルの全レイヤーで計算を行うのではなく、一部のレイヤーをスキップする形で動的に計算を行うというものです。では一緒に見ていきましょう。

MoD(Mixture-of-Depths): Transformerベース言語モデルの計算最適化に向けたアプローチ (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

紹介

従来、Transformer モデルは入力シーケンスのすべての部分に同じ計算資源を割り当ててきました。しかし今回紹介する MoD(Mixture-of-Depth) Transformer 論文では、Transformer の特定のシーケンス位置に計算資源を動的に割り当て、レイヤーごとにその割り当てを変化させることを試みています。この Mixture-of-Depths(MoD) 戦略を実現するために、著者らはどのトークンを処理するかを決める top-k ルーティング機構を用いて全体の計算量を管理します。この技術は、既存モデルに匹敵する、あるいはそれを上回る性能を示しながら、必要な FLOPs を大幅に削減し、学習後のサンプリング時には最大 50% の速度向上を提供します。

Attention is All You Need 論文で紹介された Transformer アーキテクチャ

近年の Transformer モデルや LLM(Large Language Model) は、その規模と複雑さが継続的に増大しています。これらのモデルは、人間レベルの言語理解と生成能力に近づくために数十億規模のパラメータを活用し、その結果として驚くべき性能を示しています。しかし、その進歩には大きな計算コストとエネルギー消費が伴い、モデル規模の拡大は学習時間と推論時間を急激に増加させます。

これは、Transformer が入力シーケンス中のすべてのトークンに対して同じ量の計算を行っているためです。しかし、あるトークンは文脈理解において中核的な役割を果たす一方で、別のトークンは相対的に重要度が低いなど、すべてのトークンが同じ量の情報を提供しているとは限らない、という発想から出発しています。

MoD Transformer は各トークンに対して動的に計算を割り当てることで、モデルが必要な場所に資源を集中できるようにします。これにより、モデル全体の計算コストを下げながら、重要なトークンにより多くの注意を向けることで情報処理の効率を高められます。特に LLM のような大規模モデルにおいて、MoD のアプローチは学習と推論の両面で時間と資源を節約できる重要な手法になり得ます。

MoD(Mixture-of-Depths) 手法の紹介

MoD 手法は、Transformer の各レイヤーでどのトークンを処理するかを動的に選ぶ top-k ルーティング機構を統合します。これにより、選ばれた一部のトークンだけが完全な計算過程を通過し、それ以外のトークンはレイヤーをスキップします。

図1 / 左: MoD(Mixture-of-Depths) Transformer の概要。MoE(Mixture-of-Experts) 構造と同様に、ルーターを通じて計算経路を選択します。ただし MoE と異なり、標準ブロック(Self-Attention および MLP) を計算するかどうかを決定します。このとき入力トークン($X_i$) がルーターの判断により完全な計算を有効化しない場合、そのレイヤーの計算を通過しません。
図1 / 右: 64 トークンに短縮されたシーケンスに対するルーティング決定結果です。(X軸: シーケンス、Y軸: レイヤー) 紫色はそのレイヤーの計算をすべて実行したことを示し、オレンジ色はそのレイヤーの計算をスキップしたことを示します。

上の図の説明の通り、MoD は Transformer アーキテクチャ内に予測ルーターを実装し、各トークンごとにレイヤー計算を進めるかどうかを決定します。これにより、計算資源を相対的に「重要な」トークンへ集中させることができます。

MoD Transformer の学習

MoD Transformer の学習過程は、基本的には標準的な Transformer と似ています。入力シーケンスはモデルに与えられ、各レイヤーを通過しながら必要な計算が実行されます。しかし MoD Transformer では、ルーターを通じて各トークンがどの計算経路をたどるかを決める追加ステップが含まれます。

ルーター重みの計算: ルーターはモデルの各レイヤーで各トークンに対する重みを計算します。この重みは、そのトークンが受けるべき計算量を表し、トークンの重要度に基づいています。
上位トークンの選択: ルーターは計算された重みに基づいて、各レイヤーごとに上位のトークンを選択します。これらのトークンが完全な計算を行う経路へ送られます。
計算経路の割り当て: 選択されたトークンは標準計算(Self-Attention および MLP) を実行する経路に割り当てられます。残りのトークンは残差接続を通じて渡され、追加計算なしに次のレイヤーへ進みます。
損失計算と逆伝播: モデルの出力は最終損失関数で評価され、勾配がモデル全体に逆伝播して重みが更新されます。この過程でルーターも学習され、時間の経過とともにより効率的なルーティング判断を下せるようになります。

MoD Transformer のルーティング方式 (Routing Schemes)

ルーティングの実装は MoD Transformer の中核です。著者らは 3 つのルーティング方式を比較し、Expert-Choice MoD 方式のトークンルーティングが最も優れた性能を示すことを確認しました。ここでは、紹介された他のルーティング方式とも比較してみます。

トークン選択(Token-choice) ルーティング

トークン選択ルーティングでは、各トークンが自分自身で通る経路を選びます。モデルは各トークンに対して、さまざまな計算経路(例: 異なるレイヤーや処理ブロック) に関する確率分布を生成し、その分布に従ってトークンを該当経路へ送ります。

トークン選択ルーティングでは、各トークンが最も適した計算経路を選べるため、非常に柔軟なルーティングが可能です。また、モデルが各トークンの重要度や文脈をより精密に考慮できるため、個々のトークンに最適化できます。

エキスパート選択(Expert-choice) ルーティング

エキスパート選択ルーティングでは、モデルが各経路(または「エキスパート」) で処理するトークンを直接選択します。これは、各トークンのルーター重みに基づいて上位のトークンを選別する top-k メカニズムによって行われます。

このようなエキスパート選択ルーティングを行うことで、トークンを各経路に均等に分配しやすくなり、計算負荷がバランスよく分散されるため、モデルの計算資源をより効率的に活用できます。また、上位トークンが明確に決定されるため、モデルの予測可能性と安定性も向上します。

エキスパート選択 MoD(Expert-choice MoD)

エキスパート選択 MoD は、エキスパート選択ルーティングの概念を MoD Transformer の構造に適用したものです。この方式では、モデルが各レイヤーまたは処理段階ごとに上位のトークンを選び、選ばれたトークンだけが計算を行う経路に沿って処理されます。残りのトークンは残差接続を通じて、より低コストな経路をたどります。

この方式では、選択されたトークンだけが計算を行うため、モデル全体の FLOP 使用量を大幅に削減できます。また、重要なトークンに対しては全レイヤーでの計算を保証しつつ、不要な計算を最小化してモデル性能を維持または向上できます。さらに、ルーターによる動的ルーティングを通じて、入力シーケンスの特性に応じて計算資源を柔軟に割り当てられる利点があります。

ルーティング実装方式

ルーティングの実装は、次の 3 つの主要段階で可能です。

トークン埋め込みとルーター重みの決定: モデルの各レイヤーに対して、与えられたシーケンスのトークン埋め込み集合 $X^i$ があります。各トークン埋め込み $x^l_i$ に対して、ルーターは線形変換を通じて、そのトークンが計算に参加する可能性を表すスカラー重み $r^l_i$ を計算します。
上位重みの決定: トークンのルーター重み集合 $R^l$ から、$\beta$ 番目の重みの百分位確率 $P_\beta(R^l)$ を求めます。このとき $\beta$ は、モデルが各レイヤーで処理するトークン比率を決定します。したがって、ルーター重みが $P_\beta(R^l)$ より大きいトークンが計算に参加します。
計算経路の選択: 各トークンはルーター重み $r^l_i$ と比較され、$P_\beta(R^l)$ 以上であればそのレイヤーの計算を実行し、小さければ次のレイヤーへ直接渡されます。これにより、選択されたトークンに対してのみ計算を行い、残りのトークンは計算コストを発生させない経路に沿って処理できます。

このようなルーティング実装では、次の点を特に重視する必要があります。

重みに基づく動的決定: ルーティング過程は各トークンの重みに基づいて動的に決まり、これはトークンの重要度を反映します。重要なトークンはより多くの計算を受けてモデル性能の最適化に寄与し、重要度の低いトークンは残差経路を通ってコストを節約します。
計算効率と性能最適化: このルーティング機構により、MoD Transformer は計算資源を効率よく使いながら、モデル性能を維持または向上できます。特に、計算コストの高いレイヤーで重要なトークンだけを処理することで、全体の FLOP 使用量を削減できます。
実装の単純さと効率性: ルーティングは線形変換と百分位数計算によって実装できるため、比較的シンプルで効率的です。これにより、モデル学習や推論に追加の大きな計算負荷を与えることなく、計算資源の動的割り当てを可能にします。

MoD Transformer の動作例

ここまで見てきた内容に沿って、MoD Transformer の動作を例で説明します。

たとえば、入力シーケンスに 100 個のトークンがあり、ルーターがを 10 に設定したと仮定してみましょう。この場合、各レイヤーでルーターは最も重要な 10 個のトークンを選び、それらのトークンだけが標準計算経路に沿って処理されます。

残りの 90 個のトークンは計算をスキップし、残差接続を通じて次のレイヤーへ直接渡されます。この過程を通じて、MoD Transformer は全体の FLOP 使用量を抑えつつ、重要なトークンに必要な計算を保証します。

MoD Transformer の性能比較

学習性能

ハイパーパラメータ最適化: 12.5% 容量のルーティングブロックを全面的な attention ブロックと交互に配置した MoD Transformer 変種を学習させる実験を通じて、最適なハイパーパラメータ設定を見つけました。この設定は isoFLOP 分析において最適ベースラインよりも優れた性能を提供しました。

計算効率: 特定の MoD 変種は、従来モデルと比べてより少ない FLOP で同等の性能を達成しつつ、66% 高速な処理速度を示しました。

ルーティング分析

ルーティング判断の疎性: 12.5% capacity ルーティングを用いることで、ほとんどのトークンはブロックを迂回し、ごく少数の重要なトークンだけがブロックを通過しました。これはルーター重み分布からも確認でき、ネットワークは深さに応じて特定トークンを優先してルーティングする傾向も見られました。

推論性能

自己回帰推論: 学習時の非因果的(non-causal) top-k ルーティング方式から、推論時には因果的(causal) な予測ベース手法へ切り替えても、性能低下はほとんどありませんでした。これは、その予測問題の学習が容易であり、学習初期から 97% 以上の精度を達成していたためと考えられます。

Mixture-of-Depths-and-Experts(MoDE) = MoD + MoE

MoDE(Mixture-of-Depths-and-Experts) は、MoD(Mixture-of-Depths) 手法と MoE(Mixture of Experts) 手法の統合を意味します。この方式は、2 つのアプローチの利点を組み合わせて、Transformer モデルの計算効率と性能を同時に最適化することを目的としています。MoD 方式が各トークンに対してより深く、あるいは浅く計算するかを選ぶ一方で、MoE 方式は複数の「エキスパート」の中から最適なものを通じてトークンを処理します。MoDE はこの 2 つの方式を統合することで、トークンが特定ブロックをスキップしたり選択的により深い計算を行えたりするようにすると同時に、複数のエキスパートの中から最適な処理方式を選べるようにします。

MoDE の実装方式

MoDE には 2 つの主要な実装方式があります。段階的(Staged) MoDE と 統合型(Integrated) MoDE です。

Staged MoDE(段階的 MoDE): 段階的 MoDE では、まず MoD 手法を適用して各トークンの計算経路を決定し、その後に MoE 手法を適用して、選択された経路で処理されるトークンをそれぞれのエキスパートブロックへルーティングします。こうすることで、モデルはまず計算の深さを調整し、その後で各段階における最適なエキスパートを選択できます。
Integrated MoDE(統合型 MoDE): 統合型 MoDE は、MoD と MoE の手法を 1 つのルーティング段階に統合して実装します。つまり、トークンをルーティングする際に 1 回の判断で計算経路の深さと処理するエキスパートを同時に選択します。この方式は実装がよりシンプルで効率的になり得るほか、各トークンに最も適した計算経路とエキスパートを同時に決定できる利点があります。

MoDE の利点

計算効率: MoDE は、各トークンごとに必要な計算量と担当するエキスパートを動的に決定することで、モデル全体の計算コストを効果的に管理します。
性能最適化: 各トークンを最適な計算経路とエキスパートに割り当てることで、MoDE はモデル性能を最大化します。これは特に、複雑な問題を解いたり多様な種類のデータを扱ったりする際に有用です。
柔軟性と拡張性: MoDE は多様なモデル構造やタスクに適用できる柔軟性を提供し、計算効率と性能を同時に高められる拡張可能なフレームワークを提供します。

MoDE は Transformer モデルの今後の発展方向を示すものであり、計算コストを最小化しながらモデル性能を最大化する方法を探る研究に重要な貢献をします。

さらに読む

https://arxiv.org/abs/2404.02258

この記事は GPT モデルで整理した内容をもとにしたものであり、原文の内容や意図と異なる形でまとめられている可能性があります。関心のある内容であれば、原文もあわせて参照してください。お読みの中で不自然な点や誤りを見つけた場合は、コメントでお知らせいただけると幸いです。

⚠️広告⚠️: PyTorch 韓国ユーザーコミュニティがまとめたこの記事は役に立ちましたか？会員登録すると主要な記事をメールでお届けします！ (基本は Weekly ですが、Daily に変更も可能です。)