TokenFormer: Rethinking Transformer Scaling with Tokenized Model
  Parameters
Authors: name: Haiyang Wang; name: Yue Fan; name: Muhammad Ferjad Naeem; name: Yongqin Xian; name: Jan Eric Lenssen; name: Liwei Wang; name: Federico Tombari; name: Bernt Schiele
Published: 2024-10-30 16:19:00

Abstract: 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.