TokenFormer: Repensando la Escalabilidad del Transformer con un Modelo Tokenizado de Parámetros

TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

October 30, 2024

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

cs.AI

Resumen

Los Transformers se han convertido en la arquitectura predominante en los modelos fundamentales debido a su excelente rendimiento en diversos dominios. Sin embargo, el costo sustancial de escalar estos modelos sigue siendo una preocupación significativa. Este problema surge principalmente de su dependencia de un número fijo de parámetros en proyecciones lineales. Cuando se introducen modificaciones arquitectónicas (por ejemplo, dimensiones de canal), todo el modelo generalmente requiere ser entrenado nuevamente desde cero. A medida que los tamaños de los modelos continúan creciendo, esta estrategia resulta en costos computacionales cada vez más altos y se vuelve insostenible. Para superar este problema, presentamos TokenFormer, una arquitectura nativamente escalable que aprovecha el mecanismo de atención no solo para cálculos entre tokens de entrada, sino también para interacciones entre tokens y parámetros del modelo, mejorando así la flexibilidad arquitectónica. Al tratar los parámetros del modelo como tokens, reemplazamos todas las proyecciones lineales en los Transformers con nuestra capa de atención token-parámetro, donde los tokens de entrada actúan como consultas y los parámetros del modelo como claves y valores. Esta reformulación permite un escalado progresivo y eficiente sin necesidad de un nuevo entrenamiento desde cero. Nuestro modelo escala desde 124M hasta 1.4B de parámetros mediante la adición incremental de nuevos pares de parámetros clave-valor, logrando un rendimiento comparable a los Transformers entrenados desde cero mientras se reducen significativamente los costos de entrenamiento. El código y los modelos están disponibles en 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.