Mixture-of-Modules: Reinventing Transformers as Dynamic Assemblies of Modules

要約

Transformers の浅い層から深い層までトークンを計算する必要は常にありますか?
バニラトランスフォーマーとそのバリエーションの継続的な成功は、疑いの余地のない「はい」を示唆しています。
しかし、この研究では、モジュール混合 (MoM) と呼ばれる新しいアーキテクチャを提案することで、深さ順の慣例を打破しようとします。これは、位置に関係なく、どの層も計算に使用できるという直観に動機付けられています。
必要な処理能力を備えている限り、トークンは必要ありません。
MoM の構築は、マルチヘッド アテンション ネットワークとフィードフォワード ネットワークによって定義されたモジュールの有限セットから始まり、各モジュールは独自のパラメータ化によって区別されます。
次に、2 つのルーターがセットからアテンション モジュールとフィードフォワード モジュールを繰り返し選択してトークンを処理します。
選択により、トークンのフォワード パスで計算グラフが動的に拡張され、最終的にモジュールのアセンブリが完成します。
MoM は、Transformer とその多数のバリアントに統合されたフレームワークを提供するだけでなく、Transformer のパラメータ化の冗長性を削減するための柔軟で学習可能なアプローチも提供することを示します。
OpenWebText を使用してさまざまな MoM を事前トレーニングします。
経験的な結果は、パラメータ数が異なる MoM が、GLUE ベンチマークと XSUM ベンチマークの両方で一貫してバニラ トランスフォーマーよりも優れていることを示しています。
さらに興味深いことに、固定パラメーター バジェットを使用すると、MoM-large では GPT-2-large と比較して計算グラフの深さが 38% 以上増加し、その結果、GLUE で 1.4、XSUM で 1 の絶対ゲインが得られます。
一方、MoM-large では、レイヤーあたりのモジュール数を増やしながら深さを 60% 以上削減することも可能で、GPT-2-large と比較して、同等の性能を維持しながら、TFLOP が 16% 削減され、メモリ使用量が 43% 削減されます。
パフォーマンス。

要約(オリジナル)

Is it always necessary to compute tokens from shallow to deep layers in Transformers? The continued success of vanilla Transformers and their variants suggests an undoubted ‘yes’. In this work, however, we attempt to break the depth-ordered convention by proposing a novel architecture dubbed mixture-of-modules (MoM), which is motivated by an intuition that any layer, regardless of its position, can be used to compute a token as long as it possesses the needed processing capabilities. The construction of MoM starts from a finite set of modules defined by multi-head attention and feed-forward networks, each distinguished by its unique parameterization. Two routers then iteratively select attention modules and feed-forward modules from the set to process a token. The selection dynamically expands the computation graph in the forward pass of the token, culminating in an assembly of modules. We show that MoM provides not only a unified framework for Transformers and their numerous variants but also a flexible and learnable approach for reducing redundancy in Transformer parameterization. We pre-train various MoMs using OpenWebText. Empirical results demonstrate that MoMs, of different parameter counts, consistently outperform vanilla transformers on both GLUE and XSUM benchmarks. More interestingly, with a fixed parameter budget, MoM-large enables an over 38% increase in depth for computation graphs compared to GPT-2-large, resulting in absolute gains of 1.4 on GLUE and 1 on XSUM. On the other hand, MoM-large also enables an over 60% reduction in depth while involving more modules per layer, yielding a 16% reduction in TFLOPs and a 43% decrease in memory usage compared to GPT-2-large, while maintaining comparable performance.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google