TokenFormer:用Token化模型参数重新思考Transformer扩展
TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
AI总结:
针对Transformer扩展成本高、架构修改需从头重训的问题,提出将参数视为token的TokenFormer架构,用token-参数注意力层替换线性投影,实现渐进式高效扩展,124M到1.4B规模下性能媲美从头训练的Transformer且成本更低。
摘要:
Transformer凭借在多领域的优异性能,已成为基础模型的主流架构。然而,扩展这类模型的高昂成本仍是一个重大问题。该问题主要源于Transformer依赖线性投影中固定数量的参数。当引入架构修改(如通道维度调整)时,通常需要从头重新训练整个模型。随着模型规模持续增长,这种策略导致计算成本越来越高,变得难以为继。为解决这一问题,我们提出TokenFormer——一种原生可扩展架构,它利用注意力机制不仅处理输入token间的计算,还实现token与模型参数的交互,从而提升架构灵活性。通过将模型参数视为token,我们用token-参数注意力层替换Transformer中的所有线性投影,其中输入token作为查询,模型参数作为键和值。这种重构使得无需从头重训练即可进行渐进式高效扩展。我们的模型通过逐步添加新的键-值参数对,参数规模从124M扩展到1.4B,性能与从头训练的Transformer相当,同时大幅降低了训练成本。代码和模型可在https://github.com/Haiyang-W/TokenFormer获取。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。