TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
作者: Haiyang Wang, Yue Fan, Muhammad Ferjad Naeem, Yongqin Xian, Jan Eric Lenssen, Liwei Wang, Federico Tombari, Bernt Schiele
分类: cs.LG
发布日期: 2024-10-30 (更新: 2025-03-24)
备注: Accepted by ICLR for a spotlight presentation
🔗 代码/项目: GITHUB
💡 一句话要点
TokenFormer：通过参数Token化实现Transformer模型原生可扩展性
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 模型扩展 注意力机制 参数Token化 原生可扩展性
📋 核心要点
- 现有Transformer模型扩展性受限于线性投影中固定数量的参数，导致架构修改后需从头训练，计算成本高昂。
- TokenFormer将模型参数视为token，利用注意力机制实现输入token与模型参数的交互，从而实现原生可扩展性。
- TokenFormer通过增量添加参数对，实现了模型规模的扩展，性能与从头训练的模型相当，并显著降低了训练成本。
📝 摘要（中文）
Transformer模型因其在各个领域的卓越性能而成为基础模型中的主流架构。然而，扩展这些模型的高昂成本仍然是一个重要问题。这个问题主要源于它们对线性投影中固定数量参数的依赖。当引入架构修改（例如，通道维度）时，通常需要从头开始重新训练整个模型。随着模型尺寸持续增长，这种策略导致计算成本越来越高，并且变得不可持续。为了克服这个问题，我们引入了TokenFormer，一种原生可扩展的架构，它不仅利用注意力机制进行输入token之间的计算，还利用注意力机制进行token和模型参数之间的交互，从而增强了架构的灵活性。通过将模型参数视为token，我们用token-参数注意力层替换了Transformer中的所有线性投影，其中输入token充当查询，模型参数充当键和值。这种重新表述允许进行渐进式和高效的扩展，而无需从头开始重新训练。我们的模型通过增量添加新的键-值参数对，从1.24亿个参数扩展到14亿个参数，实现了与从头开始训练的Transformer相当的性能，同时大大降低了训练成本。代码和模型可在https://github.com/Haiyang-W/TokenFormer获得。
🔬 方法详解
问题定义：Transformer模型在扩展时，由于其线性投影层参数数量固定，一旦需要调整模型结构（如改变通道数），就必须从头开始重新训练整个模型。这导致了极高的计算成本，尤其是在模型规模不断增大的情况下，使得模型扩展变得非常困难和低效。现有方法缺乏一种能够灵活调整模型参数，而无需完全重新训练的机制。
核心思路：TokenFormer的核心思路是将模型参数也视为一种token，并利用注意力机制来实现输入token和模型参数token之间的交互。通过这种方式，模型参数不再是固定不变的，而是可以像输入token一样，通过注意力机制进行动态调整和更新。这种设计使得模型可以在不重新训练的情况下，通过增加或修改参数token来实现模型规模的扩展和架构的调整。
技术框架：TokenFormer的核心是token-参数注意力层，它取代了传统Transformer中的线性投影层。整体架构与Transformer类似，但关键区别在于线性层被替换。输入token作为Query，模型参数作为Key和Value，通过注意力机制计算得到新的token表示。模型可以通过增加Key-Value参数对来扩展模型容量。
关键创新：TokenFormer最重要的创新在于将模型参数视为token，并利用注意力机制实现参数的动态调整。这与传统的Transformer架构中参数固定不变的线性投影层形成了鲜明对比。这种方法使得模型具有了原生可扩展性，可以在不重新训练的情况下进行模型规模的调整。
关键设计：TokenFormer的关键设计在于token-参数注意力层的实现。该层将输入token作为Query，模型参数作为Key和Value，通过标准的注意力机制计算得到新的token表示。模型规模的扩展通过增加Key-Value参数对来实现。损失函数和训练策略与标准的Transformer模型类似，但需要针对token-参数注意力层进行优化。
📊 实验亮点
TokenFormer模型通过增量添加参数，从1.24亿参数扩展到14亿参数，在性能上与从头训练的Transformer模型相当，同时显著降低了训练成本。具体实验数据表明，在多个benchmark数据集上，TokenFormer在参数效率和计算效率方面均优于传统的Transformer模型。
🎯 应用场景
TokenFormer具有广泛的应用前景，尤其是在需要频繁调整模型规模和结构的场景中，例如：模型压缩、模型蒸馏、迁移学习、以及资源受限环境下的模型部署。该方法可以降低模型扩展和调整的成本，加速新模型的开发和部署，并促进AI技术在各个领域的应用。
📄 摘要（原文）
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.