TokenFormer: Het heroverwegen van Transformer-schaling met getokeniseerd model Parameters

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

Samenvatting

Transformers zijn de dominante architectuur geworden in basismodellen vanwege hun uitstekende prestaties in verschillende domeinen. Echter blijft de aanzienlijke kosten van het schalen van deze modellen een belangrijke zorg. Dit probleem ontstaat voornamelijk door hun afhankelijkheid van een vast aantal parameters binnen lineaire projecties. Wanneer architecturale aanpassingen (bijv. kanaaldimensies) worden geïntroduceerd, vereist het gehele model doorgaans opnieuw getraind te worden vanaf het begin. Naarmate de modelgroottes blijven groeien, leidt deze strategie tot steeds hogere computationele kosten en wordt onhoudbaar. Om dit probleem te overwinnen, introduceren we TokenFormer, een van nature schaalbare architectuur die het aandachtsmechanisme benut niet alleen voor berekeningen tussen invoertokens, maar ook voor interacties tussen tokens en modelparameters, waardoor de architecturale flexibiliteit wordt verbeterd. Door modelparameters te behandelen als tokens, vervangen we alle lineaire projecties in Transformers door onze token-parameter aandachtslaag, waarbij invoertokens fungeren als vragen en modelparameters als sleutels en waarden. Deze herformulering maakt progressief en efficiënt schalen mogelijk zonder de noodzaak van opnieuw trainen vanaf het begin. Ons model schaalt van 124M naar 1.4B parameters door geleidelijk nieuwe sleutel-waarde parameterparen toe te voegen, met prestaties vergelijkbaar met Transformers die vanaf het begin zijn getraind, terwijl de trainingskosten aanzienlijk worden verlaagd. Code en modellen zijn beschikbaar op 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.