TokenFormer : Repenser l'échelle du Transformer avec un modèle à tokens paramétrés

TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

October 30, 2024

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

cs.AI

Résumé

Les Transformers sont devenus l'architecture prédominante dans les modèles de base en raison de leurs excellentes performances dans divers domaines. Cependant, le coût substantiel de mise à l'échelle de ces modèles reste une préoccupation majeure. Ce problème découle principalement de leur dépendance à un nombre fixe de paramètres dans les projections linéaires. Lorsque des modifications architecturales (par exemple, les dimensions des canaux) sont introduites, l'ensemble du modèle nécessite généralement une nouvelle formation à partir de zéro. À mesure que les tailles des modèles continuent de croître, cette stratégie entraîne des coûts de calcul de plus en plus élevés et devient insoutenable. Pour surmonter ce problème, nous introduisons TokenFormer, une architecture nativement évolutive qui exploite le mécanisme d'attention non seulement pour les calculs entre les jetons d'entrée, mais aussi pour les interactions entre les jetons et les paramètres du modèle, améliorant ainsi la flexibilité architecturale. En traitant les paramètres du modèle comme des jetons, nous remplaçons toutes les projections linéaires dans les Transformers par notre couche d'attention jeton-paramètre, où les jetons d'entrée agissent en tant que requêtes et les paramètres du modèle en tant que clés et valeurs. Cette reformulation permet une mise à l'échelle progressive et efficace sans nécessiter une nouvelle formation à partir de zéro. Notre modèle passe de 124M à 1,4 milliard de paramètres en ajoutant progressivement de nouveaux paires de paramètres clé-valeur, atteignant des performances comparables à celles des Transformers formés à partir de zéro tout en réduisant considérablement les coûts de formation. Le code et les modèles sont disponibles sur 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.