TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

基本信息

• 📝 原文链接: https://arxiv.org/abs/2410.23168
• 👥 作者: Haiyang Wang, Yue Fan, Muhammad Ferjad Naeem, Yongqin Xian, Jan Eric Lenssen, Liwei Wang, Federico Tombari, Bernt Schiele
• 🏷️ 关键词: Progressive Scaling, Attention mechanism
• 📚 分类: 机器学习

摘要

中文摘要

由于在各个领域都表现出色，Transformer已成为基础模型中主要的架构。然而，扩展这些模型的巨大成本仍然是一个重要的问题。这个问题主要源于它们依赖于线性投影中固定数量的参数。当引入架构修改（例如，通道维度）时，整个模型通常需要从头开始重新训练。随着模型规模的不断增长，这种策略导致计算成本越来越高，并变得不可持续。为了克服这个问题，我们引入了TokenFormer，这是一种原生可扩展的架构，它不仅利用注意力机制进行输入token之间的计算，还利用注意力机制进行token与模型参数之间的交互，从而增强了架构的灵活性。通过将模型参数视为token，我们用我们的token-参数注意力层替换了Transformer中的所有线性投影，其中输入token充当查询，模型参数充当键和值。这种重新表述允许渐进式和高效的扩展，而无需从头开始重新训练。我们的模型通过增量添加新的键值参数对，从1.24亿个参数扩展到14亿个参数，实现了与从头开始训练的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.

论文解读

一句话总结

Tokenformer提出了一种基于完全注意力机制的可扩展Transformer架构，通过将模型参数视为token并使用token-parameter attention层替换线性投影，实现了在不重新训练的情况下高效扩展模型规模，并显著降低了训练成本。

问题1：这篇论文想要解决什么具体问题？

• 问题背景：Transformer模型在各种领域表现出色，但扩展模型规模的成本很高，主要是因为线性投影中参数数量固定。修改架构（例如改变通道维度）通常需要从头开始重新训练，随着模型规模的增长，这种策略的计算成本越来越高，难以承受。
• 现有方案不足：现有的模型重用方法通常涉及复制、堆叠或组合模型权重，这可能会破坏预训练模型的分布，导致预训练知识丢失并减慢收敛速度。
• 研究目标：提出一种能够渐进式且高效扩展的Transformer架构，在不从头开始重新训练的情况下实现与从头训练的Transformer模型相当的性能，并显著降低训练成本。

问题2：论文的核心创新点是什么？

• 技术创新：提出了一种名为Tokenformer的全新架构，该架构完全基于注意力机制，将模型参数作为token处理。引入了token-parameter attention层，用以取代Transformer中的所有线性投影。
• 方法改进：通过将模型参数视为token，并利用交叉注意力机制，实现了模型参数的渐进式扩展。通过添加新的键值对参数，可以有效地扩展模型规模，而无需从头开始重新训练。
• 优势：与从头开始训练的Transformer相比，Tokenformer在扩展模型规模时显著降低了训练成本，同时实现了可比的性能。

问题3：实验结果如何验证了方法的有效性？

• 关键实验：在语言建模和图像分类任务上进行了实验，将Tokenformer与从头开始训练的Transformer模型进行了对比，并评估了渐进式模型扩展的有效性。使用了OpenWebText语料库和ImageNet数据集。
• 性能提升：在语言建模任务中，Tokenformer在1.4B参数规模下实现了与从头训练的Transformer模型相当的困惑度（perplexity），同时训练成本降低了三分之二以上。 图像分类任务中，Tokenformer达到了与标准ViT模型相当的性能。
• 对比结果：图3和图4展示了Tokenformer与从头开始训练的Transformer模型在训练成本方面的对比，Tokenformer显著降低了训练成本。表1和表2分别展示了Tokenformer在语言建模和图像分类任务上的性能与其他基线模型的对比，展现出其在各个任务上的竞争力。

问题4：这个研究的实际应用价值是什么？

• 应用场景：Tokenformer架构适用于各种需要处理大型模型的应用场景，例如自然语言处理、计算机视觉等领域的大型基础模型的训练和部署。
• 实施建议：论文提供了Tokenformer的代码和模型，方便研究人员进行实验和应用。建议根据具体应用场景选择合适的模型规模和训练策略。
• 局限与展望：论文中未明确指出Tokenformer的局限性。未来的研究方向可能包括：扩展混合专家（MoE）范式，改进参数高效微调技术，整合视觉和语言模型，以及增强模型的可解释性。