Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
下記の文章は論文のアブストラクトをChatGPTに解釈させたものであり、正確性は保証しません。
本研究は、トランスフォーマベースの言語モデルにおいて、FLOPs（計算量）をシーケンスの特定の位置に動的に割り当てることができることを示しています。通常、トランスフォーマモデルは、各トークンごとに同じ量の計算を行いますが、本研究では、異なるレイヤーごとにシーケンス上で最適な割り当てを学習することができます。具体的には、各レイヤーで自己注意とMLP計算に参加できるトークンの数（𝑘）を制限し、ネットワークがトップ-𝑘ルーティングメカニズムを使用して処理するトークンを決定します。この方法は、事前に定義された𝑘に基づいて静的な計算グラフを使用するため、他の条件付き計算技術とは異なり、計算グラフのサイズが既知です。しかし、𝑘トークンのアイデンティティが変動するため、この方法では計算量を時間とモデルの深さの次元で均等に分散することができます。したがって、計算量の総計は完全に予測可能ですが、トークンレベルで動的かつ文脈に応じたものです。このように学習されたモデルは、計算を動的に割り当てることができるだけでなく、効率的に割り当てることもできます。同等のFLOPsとトレーニング時間でベースラインの性能に達することができますが、1回のフォワードパスあたりのFLOPsは一部の場合でのみ必要であり、ポストトレーニングサンプリング時には50％以上高速です。
本研究の目的は、トランスフォーマベースの言語モデルにおいて、静的な計算予算を使用して計算を効率的に割り当てる方法を開発することです。通常、トランスフォーマモデルは各トークンに対して同じ量の計算を行いますが、本研究では、利用可能な計算予算を各レイヤーごとにトークンごとに動的に割り当てる方法を提案しています。これにより、ハードウェアの効率性を向上させることができます。具体的には、ユーザーが定義した総計算量をトレーニング前に設定し、ネットワークのオンザフライの決定に依存しないようにします。これにより、メモリの使用量や1回のフォワードパスあたりのFLOPsを削減するなど、ハードウェアの効率を最大化することができます。モデルの全体的な性能を犠牲にすることなく、これらの利点を得ることができます。
本研究で使用されたデータについては、文中では具体的な情報は提供されていません。したがって、データの名前や利用可否についてはわかりません。
本研究で使用された手法は、Mixture-of-Depths（MoD）と呼ばれます。MoDは、Mixture of Experts（MoE）トランスフォーマと似たアプローチを採用しており、ネットワークの深さにわたってトークンレベルのルーティングの決定を行います。MoEとは異なり、MoDでは、通常のトランスフォーマと同様にトークンに計算を適用するか、リジッドコネクションを介して計算をスキップするかを選択します。また、MoEとは異なり、このルーティングは前向きのMLPとマルチヘッドアテンションの両方に適用されます。したがって、ルーティングはトークンの更新だけでなく、アテンションの対象となるトークンの選択にも影響を与えます。この手法は、トランスフォーマの深さに沿ってトークンが異なる数のレイヤーまたはブロックを通過することを強調するために、Mixture-of-Depths（MoD）と呼ばれています。
本研究によって明らかになったことは、MoDトランスフォーマが効率的に計算を割り当てることができるということです。MoDトランスフォーマは、同等のトレーニングFLOPs（計算量）に対してベースラインのトランスフォーマと比較して、最終的な対数確率トレーニング目標で1.5％改善することができます。また、1回のフォワードパスあたりのFLOPsの一部を使用するため、ステップが高速化され、トレーニング時間が短縮されます。これにより、MoDトランスフォーマは、不必要な計算をスキップすることで（つまり、割り当てる必要のない計算を行わずに）インテリジェントにルーティングすることができることが示されています。
この記事が気に入ったらサポートをしてみませんか？