TokenFormer: Repensando a Escala do Transformer com Modelo Tokenizado 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

Resumo

Os Transformers tornaram-se a arquitetura predominante em modelos fundamentais devido ao seu excelente desempenho em vários domínios. No entanto, o custo substancial de escalar esses modelos continua sendo uma preocupação significativa. Esse problema surge principalmente da dependência de um número fixo de parâmetros em projeções lineares. Quando modificações arquiteturais (por exemplo, dimensões de canais) são introduzidas, o modelo inteiro geralmente requer um novo treinamento do zero. À medida que os tamanhos dos modelos continuam crescendo, essa estratégia resulta em custos computacionais cada vez mais altos e se torna insustentável. Para superar esse problema, introduzimos o TokenFormer, uma arquitetura nativamente escalável que aproveita o mecanismo de atenção não apenas para cálculos entre tokens de entrada, mas também para interações entre tokens e parâmetros do modelo, aumentando assim a flexibilidade arquitetural. Ao tratar os parâmetros do modelo como tokens, substituímos todas as projeções lineares nos Transformers pelo nosso camada de atenção token-parâmetro, onde os tokens de entrada atuam como consultas e os parâmetros do modelo como chaves e valores. Essa reformulação permite escalonamento progressivo e eficiente sem a necessidade de um novo treinamento do zero. Nosso modelo escala de 124M para 1.4B parâmetros, adicionando incrementalmente novos pares de parâmetros chave-valor, alcançando desempenho comparável aos Transformers treinados do zero, enquanto reduz significativamente os custos de treinamento. Código e modelos estão disponíveis em 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.