TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

要約

トランスフォーマーは、さまざまなドメインにわたって優れたパフォーマンスを発揮するため、基礎モデルの主要なアーキテクチャとなっています。
ただし、これらのモデルを拡張するための多大なコストが依然として大きな懸念事項です。
この問題は主に、線形投影内の固定数のパラメータに依存することから発生します。
アーキテクチャの変更 (チャネルの寸法など) が導入されると、通常、モデル全体を最初から再トレーニングする必要があります。
モデルのサイズが増大し続けると、この戦略では計算コストがますます高くなり、持続不可能になります。
この問題を解決するために、入力トークン間の計算だけでなく、トークンとモデル パラメーター間の相互作用にもアテンション メカニズムを活用する、ネイティブにスケーラブルなアーキテクチャである TokenFormer を導入します。これにより、アーキテクチャの柔軟性が向上します。
モデル パラメーターをトークンとして扱うことで、Transformers のすべての線形投影をトークン パラメーター アテンション レイヤーに置き換えます。入力トークンはクエリとして機能し、モデル パラメーターはキーと値として機能します。
この再定式化により、最初から再トレーニングすることなく、段階的かつ効率的なスケーリングが可能になります。
私たちのモデルは、新しいキーと値のパラメーターのペアを段階的に追加することで 1 億 2,400 万から 1.4 億のパラメーターまで拡張し、トレーニング コストを大幅に削減しながら、ゼロからトレーニングした Transformers に匹敵するパフォーマンスを実現します。
コードとモデルは \url{https://github.com/Haiyang-W/TokenFormer} で入手できます。

要約(オリジナル)

Transformers have become the predominant architecture in foundation models due to their excellent performance across various domains. However, the substantial cost of scaling these models remains a significant concern. This problem arises primarily from their dependence on a fixed number of parameters within linear projections. When architectural modifications (e.g., channel dimensions) are introduced, the entire model typically requires retraining from scratch. As model sizes continue growing, this strategy results in increasingly high computational costs and becomes unsustainable. To overcome this problem, we introduce TokenFormer, a natively scalable architecture that leverages the attention mechanism not only for computations among input tokens but also for interactions between tokens and model parameters, thereby enhancing architectural flexibility. By treating model parameters as tokens, we replace all the linear projections in Transformers with our token-parameter attention layer, where input tokens act as queries and model parameters as keys and values. This reformulation allows for progressive and efficient scaling without necessitating retraining from scratch. Our model scales from 124M to 1.4B parameters by incrementally adding new key-value parameter pairs, achieving performance comparable to Transformers trained from scratch while greatly reducing training costs. Code and models are available at \url{https://github.com/Haiyang-W/TokenFormer}.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google