TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

Transformer 已成为基础模型中的主要架构，因其在各个领域的出色表现而备受青睐。然而，这些模型的扩展成本仍然是一大问题。这一问题主要源于它们在线性投影中依赖于固定数量的参数。当引入架构修改（例如通道维度）时，通常需要从头开始重新训练整个模型。随着模型规模的不断增长，这种策略导致计算成本越来越高，变得不可持续。为了解决这个问题，我们提出了 TokenFormer，这是一种原生可扩展的架构，它不仅利用注意力机制进行输入令牌之间的计算，还用于令牌与模型参数之间的交互，从而增强了架构的灵活性。通过将模型参数视为令牌，我们将 Transformer 中的所有线性投影替换为我们的令牌-参数注意力层，在该层中，输入令牌作为查询，模型参数作为键和值。这种重新表述允许逐步且高效地扩展，而无需从头开始重新训练。我们的模型通过逐步添加新的键-值参数对，从 1.24 亿参数扩展到 14 亿参数，实现了与从头开始训练的 Transformer 相当的性能，同时大幅降低了训练成本。代码和模型可在以下网址获取：[https://github.com/Haiyang-W/TokenFormer](https://github.com/Haiyang-W/TokenFormer)。

论文试图解决的问题是Transformer模型在扩展时面临的高成本问题。由于Transformer依赖于固定数量的参数进行线性投影，当模型架构发生变化时，通常需要从头开始重新训练整个模型，这导致随着模型规模的增长，计算成本变得越来越高，变得不可持续。

论文的关键思路是引入TokenFormer，这是一种新的可扩展架构，通过将模型参数视为tokens，利用注意力机制不仅处理输入tokens之间的计算，还处理tokens与模型参数之间的交互。这种方法通过用token-parameter注意力层替换Transformer中的所有线性投影，实现了模型的渐进和高效扩展，而无需从头开始重新训练。

论文的其他亮点包括：1) 实验展示了TokenFormer可以从124M参数扩展到1.4B参数，性能与从头训练的Transformer相当，但大大减少了训练成本；2) 论文提供了开源代码和模型，便于复现和进一步研究；3) 研究提出了一个新颖的视角，即将模型参数作为tokens处理，为未来的研究提供了新的方向。

近期在这个领域中，还有一些相关的研究，例如：1)《Adaptive Attention Span in Transformers》提出了一种自适应调整注意力范围的方法，以提高Transformer的效率；2)《Performer: Generalized Attention with RFF Kernels for Scalable Transformer Models》通过使用随机傅里叶特征核来实现高效的Transformer模型；3)《Reformer: The Efficient Transformer》通过局部敏感哈希技术减少Transformer的内存和计算需求。