TokenFormer: トークン化されたモデルパラメータを用いたTransformerスケーリングの再考

TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

October 30, 2024

著者: Haiyang Wang, Yue Fan, Muhammad Ferjad Naeem, Yongqin Xian, Jan Eric Lenssen, Liwei Wang, Federico Tombari, Bernt Schiele

cs.AI

要旨

トランスフォーマーは、様々な領域で優れたパフォーマンスを発揮することから、基礎モデルにおいて主要なアーキテクチャとなっています。ただし、これらのモデルをスケーリングする際の膨大なコストは依然として重大な懸念事項です。この問題は、線形射影内のパラメータの固定数に依存していることから主に発生します。アーキテクチャの変更（例：チャネルの次元）が導入されると、通常、モデル全体をゼロから再学習する必要があります。モデルのサイズが拡大するにつれ、この戦略はますます高コンピュテーションコストをもたらし、持続不可能となります。この問題を克服するために、私たちはTokenFormerを導入します。これは、入力トークン間の計算だけでなく、トークンとモデルパラメータとの相互作用にも注意メカニズムを活用する、ネイティブにスケーラブルなアーキテクチャです。モデルパラメータをトークンとして扱うことで、トランスフォーマー内のすべての線形射影をトークン-パラメータ注意層で置き換え、入力トークンがクエリとして機能し、モデルパラメータがキーおよび値として機能するようにします。この再定式化により、ゼロからの再学習を必要とせずに段階的かつ効率的なスケーリングが可能となります。私たちのモデルは、新しいキー-値パラメータペアを段階的に追加することで124Mから1.4Bのパラメータにスケールし、トランスフォーマーと同等のパフォーマンスを達成しながら、トレーニングコストを大幅に削減します。コードとモデルはhttps://github.com/Haiyang-W/TokenFormer で入手可能です。

English

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 https://github.com/Haiyang-W/TokenFormer.