让模型参数像令牌一样参与计算,实现无需重训的渐进式扩容。
TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
- 把模型参数当作令牌,用注意力机制动态关联输入与参数
- 从124M到1.4B参数逐步扩展,性能媲美从头训练的Transformer
- 适合需要频繁调整规模但不想重训的模型部署场景
Transformer因其在多领域中的优异表现,已成为基础模型的主流架构。然而,模型扩展带来的高昂成本仍是主要挑战,根源在于线性投影中固定参数数量的设计。当调整通道维度等结构时,通常需从头训练,导致计算开销随模型增大而急剧上升,难以为继。为此,我们提出TokenFormer,一种原生可扩展架构,将注意力机制用于输入令牌与模型参数间的交互,提升结构灵活性。通过将模型参数视为令牌,用令牌-参数注意力层取代所有线性投影,其中输入令牌作为查询,参数作为键值对。该设计支持无须重训的渐进式扩展。我们实现了从124M到1.4B参数的逐步扩容,在性能上达到与从头训练的Transformer相当水平,显著降低训练成本。代码与模型已公开于https://github.com/Haiyang-W/TokenFormer。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。