TokenFormer: Riconsiderare la scalabilità del Transformer con il modello tokenizzato Parametri

TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

October 30, 2024

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

cs.AI

Abstract

I Transformer sono diventati l'architettura predominante nei modelli di base grazie alle loro eccellenti prestazioni in vari domini. Tuttavia, il notevole costo di scalare questi modelli rimane una preoccupazione significativa. Questo problema deriva principalmente dalla dipendenza da un numero fisso di parametri all'interno delle proiezioni lineari. Quando vengono introdotte modifiche architettoniche (ad esempio, dimensioni dei canali), l'intero modello richiede tipicamente un nuovo addestramento da zero. Man mano che le dimensioni del modello continuano a crescere, questa strategia comporta costi computazionali sempre più elevati e diventa insostenibile. Per superare questo problema, introduciamo TokenFormer, un'architettura nativamente scalabile che sfrutta il meccanismo di attenzione non solo per i calcoli tra i token di input, ma anche per le interazioni tra i token e i parametri del modello, migliorando così la flessibilità architetturale. Trattando i parametri del modello come token, sostituiamo tutte le proiezioni lineari nei Transformer con il nostro strato di attenzione token-parametro, dove i token di input agiscono come query e i parametri del modello come chiavi e valori. Questa riformulazione consente una scalabilità progressiva ed efficiente senza la necessità di un nuovo addestramento da zero. Il nostro modello scala da 124M a 1.4B parametri aggiungendo incrementalmente nuove coppie di parametri chiave-valore, ottenendo prestazioni paragonabili ai Transformer addestrati da zero riducendo notevolmente i costi di addestramento. Il codice e i modelli sono disponibili su 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.