[論文レビュー] MatFormer: Nested Transformer for Elastic Inference

Devvrit Devvrit, Sneha Kudugunta+9 more|arXiv (Cornell University)|Oct 11, 2023

Opportunistic and Delay-Tolerant NetworksComputer Science被引用数 3

ひとこと要約

MatFormerは、各フィードフォワードネットワークブロック内にサイズの異なる複数のサブネットワークを統合的に学習する、ネストされたTransformerアーキテクチャを導入することで、追加の学習なしに数百の正確で行動的に一貫性のあるサブモデルを自由に抽出できるエラスティック推論を実現する。この手法により、言語および視覚タスクの両方で独立に学習されたモデルと同等の性能を達成し、予測的デコードやリtrieval分野において最先端の結果も得られている。

ABSTRACT

Foundation models are applied in a broad spectrum of settings with different inference constraints, from massive multi-accelerator clusters to resource-constrained standalone mobile devices. However, the substantial costs associated with training these models often limit the number of unique model sizes that can be offered. Consequently, practitioners are compelled to select a model that may not be optimally aligned with their specific latency and cost requirements. We present MatFormer, a novel Transformer architecture designed to provide elastic inference across diverse deployment constraints. MatFormer achieves this by incorporating a nested Feed Forward Network (FFN) block structure within a standard Transformer model. During training, we optimize the parameters of multiple nested FFN blocks with varying sizes, enabling the extraction of hundreds of accurate smaller models without incurring additional computational costs. We empirically validate the efficacy of MatFormer across different model classes (decoders and encoders) and modalities (language and vision), demonstrating its potential for real-world deployment. We show that a 850M decoder-only MatFormer language model (MatLM) allows us to extract multiple smaller models spanning from 582M to 850M parameters, each exhibiting better validation loss and one-shot downstream evaluations than independently trained counterparts. Furthermore, we observe that smaller encoders extracted from a universal MatFormer-based ViT (MatViT) encoder preserve the metric-space structure for adaptive large-scale retrieval. Finally, we showcase that speculative decoding with the accurate and consistent submodels extracted from MatFormer can lead to significant reduction in inference latency. Project website: https://devvrit.github.io/matformer/

研究の動機と目的

• モバイルデバイスや大規模クラスタなど、多様な推論環境に展開する際の基礎的モデルの高コストな学習と限られたモデルの粒度の問題に対処すること。
• 異なるサイズの独立して学習されたモデル同士における行動的一致性の欠如と推論最適化の互換性の欠如を克服すること。
• 1つの学習済みユニバーサルモデルから数百の正確なサブモデルを抽出することで、遅延・コスト・精度のトレードオフを細かく制御できること。
• 再学習なしに、多様なモデルサイズ、モダリティ（言語および視覚）およびタスクにおいて高いパフォーマンスを維持すること。
• サブモデル間のメトリクス空間構造とモデルの一貫性を保つことで、予測的デコードおよび適応的リtrievalのサポートを可能にすること。

提案手法

• フィードフォワードネットワーク（FFN）ブロックにネスト構造を導入し、小さなサブモデルがより大きなモデルに埋め込まれ、T₁ ⊂ T₂ ⊂ … ⊂ T₉ のような階層を形成する。
• すべてのネストされたサブモデルを1回の学習ランで、個々の損失を組み合わせることで同時に学習し、パラメータ共有と一貫性を保証する。
• ニューロンおよびアテンションヘッドを最も重要度の高いものから順に整理するマトリョーシカ表現学習の原則を適用し、階層的抽出を可能にする。
• Mix'n'Match抽出を可能にし、各層のサブブロックを任意に組み合わせることで、gˡ 個の異なるモデルが得られ、モデルの粒度が指数関数的に増加する。
• 推論効率の向上を最大化するため、非エンベッディングパラメータと遅延の60％以上を占めるFFNブロックに焦点を当てる。
• ビジョンエンコーダー（MatViT）においてもメトリクス空間構造を保持することで、適応的リtrievalシステムへの有効な応用が可能になる。

実験結果

リサーチクエスチョン

• RQ1 1回の学習で、追加の学習なしに、さまざまなサイズの正確で行動的に一貫性のあるサブモデルを数百個抽出できるユニバーサルなTransformerモデルを訓練可能か？
• RQ2 MatFormerから抽出されたサブモデルは、下流のNLPおよびビジョンタスクにおいて、独立に学習されたモデルと比べてどの程度の性能を示すか？
• RQ3 ネスト構造は、リtrieval応用を想定したビジョンエンコーダーにおけるメトリクス空間構造をどの程度保持するか？
• RQ4 抽出されたサブモデルを用いた予測的デコードは、精度を損なわせずに推論遅延を低減できるか？
• RQ5 ネストされたサブモデルを同時に学習させることで、さまざまな粒度における汎化能力にどのような影響を与えるか？

主な発見

• 2.6BパラメータのMatLMデコーダー・モデルは、1.5Bから2.6Bパラメータのサブモデルを抽出可能であり、それぞれが独立に学習されたモデルと同等の検証損失およびワンショット下流タスク性能を達成している。
• MatViTから抽出されたサブモデルは、メトリクス空間構造を保持しており、大規模な適応的リtrievalに有効に利用できる。
• MatFormerから抽出された正確で一貫性のあるサブモデルを用いた予測的デコードは、精度を損なわせずに推論遅延を低減する。
• TriviaQAベンチマークでは、1.55BのMatFormerサブモデルが26.15％のEMを達成し、1.65Bの独立学習モデルと同等の性能を示した。
• 14の下流タスクにおける1.55Bサブモデルの平均EMスコアは48.39％に達し、2.6Bのフルモデルの49.18％に非常に近い水準に達している。
• 抽出された1.55Bサブモデルの開発セットのログパープレキシティは2.663であり、フルモデルの2.65に非常に近い値であり、優れた一般化能力を示している。