Token化一切，甚至网络！北大&谷歌&马普所提出TokenFormer

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。

本论文第一作者是汪海洋，北京大学20级博士生，目前主要关注是通用模型的架构设计和学习算法。指导教授主要包括王立威，北京大学智能学院教授；BerntSchiele，德国马普计算所教授；FedericoTombari谷歌人工智能科学家等。

新一代通用灵活的网络结构TokenFormer:RethinkingTransformerScalingwithTokenizedModelParameters来啦！

TokenFormer不仅像原始Transformer一样Token化了inputdata，并且Token化了网络参数，将attention机制拓展到Token和parameters的交互中，最大化了Transformer的灵活性，真正得到了一个Fullyattention-based的网络结构。

这种方式打破了原有人们区别看待data和model的观念，即所有的计算都归纳为不同类型的Token（e.g.,data,paramtoken）通过灵活的attention来交互。得益于这一灵活的性质，TokenFormer允许incrementalscalingmodelsize，基于训好的模型上增量的拓展新的更大的模型，大大节省了计算的开销：

这项名为TokenFormer的新工作，由谷歌，马普计算所和北大的研究者提出，在Twitter，HackerNews,Reddit上得到广泛的讨论和关注(Twitter上有150K+的浏览量)。

目前代码、模型和项目主页均已放出：

论文链接：https://arxiv.org/pdf/2410.23168开源代码：https://github.com/Haiyang-W/TokenFormer开源模型：https://huggingface.co/Haiyang-W

背景介绍

得益于其处理各种数据的灵活性，Transformer网络结构在各个AI领域都取得了巨大的成功。

Transformer模型通常将处理单个Token所需的计算分为两个部分：与其他Token的交互（Token-TokenInteraction）和涉及模型参数的计算（Token-ParameterInteraction）。

Attention促进了Token-Token之间的交互，使现代通用基础模型能够将多模态数据编码成统一的Token序列，并有效捕捉它们之间的复杂依赖关系。

相反，Token-Parameter计算主要依赖于固定的linearprojection，大大限制modelsize的scaling。Scalingmodel是通常改变模型结构，往往需要从头训练整个模型，带来了过多的资源消耗，使其越来越不切实际。

在本文中，研究团队使用token这一概念建模所有的计算，即将modelparameters也视为一种token，网络的计算统一为各种不同的token(e.g.,datatokensandparametertokens)之间通过attention来进行交互，大大增强了Token-Parameter交互的灵活性，从而能够增量式的扩展模型参数，有效地重用先前训练的模型，从而显著降低了训练负担。

为实现这一目标，研究团队引入了TokenFormer。统一Token-Token和Token-ParametersInteraction的计算。其Token-Parameterattention具有灵活性，并能够处理可变数量的参数，从而本质上最大化了Transformer的灵活性，增强了模型的可扩展性。

TokenFormer提供一种新的看待模型的视角，即网络的计算就是一些Tokens相互任意交互。基于这些Tokens（e.g.,datatoken,parametertoken,memorytoken）和attention机制可以灵活地构造任意的网络结构。

该团队希望TokenFormer作为一种通用的网络结构，不仅在incrementalmodelscaling上有贡献，还在SparseInference,Parameter-EfficientTuning,VisionandLanguageModels,Device-CloudCollaboration和ModelInterpretability等领域有更多的贡献。

方法

Tokenformer的核心创新是Token-ParameterAttention（Pattention）Layer，它结合了一组TrainableTokens作为modelparameters，并通过cross-attention来管理InputToken与这些ParameterTokens之间的交互。

通过这种方式，Pattention层引入了一个额外的维度—ParameterToken的数量，这一维度独立于输入和输出维度。此解耦方式使得输入数据可以与variablenumberofparameters进行交互，提供了增量模型扩展所需的灵活性。

应用：天生的增量式ModelScaling

有了TokenFormer这一灵活的性质，可以延伸出很多应用。这里以增量式modelscaling为例。

实验结果

增量式modelscaling：如下右图所示，模型在已经训好的124M的模型的基础上，采用增量式训练，只用十分之一的数据就可以达到从头训练策略相近的性能，让模型可以不断迭代，真正地活起来了。

LanguageModeling：如下表所示，研究团队比较了Transformer-based的模型和TokenFormer在语言建模上的能力。

在相同规模、相同模型尺寸下，TokenFormer在大大增加灵活性的前提下达到了比Transformer更好的zero-shot性能。这里研究团队follow了pythia标准的训练代码以及数据集：Pile(300B)。上述结果展现了TokenFormer在语言模型建模上的能力。

VisualModeling:为了进一步验证TokenFormer的表达能力，研究团队还和标准的visiontransformer进行了对比。

在ImageNet-1K的监督训练的setting上，使用相同的训练策略，TokenFormer的性能超过了vision-transformer，验证了其在visualmodeling上的能力。

未来研究方向

极致的专家混合（Mixture-of-Experts）范式

研究团队认为Tokenformer是专家混合（MoE）框架的极致实例化，其中每一组键-值参数对都充当一个独立的专家。这种创新的类MoE架构有可能显著减少与Token-Parameter交互相关的计算成本。

新的参数高效微调范式

Tokenformer的扩展方法通过集成额外的key-valueparameterpairs，展现了一种参数高效的微调策略。当面对新任务或数据集时，该模型可以通过加入新的TokenParameters来扩展其预训练参数，从而快速适应特定任务需求。

整合视觉和语言模型

利用Tokenformer的参数高效微调能力，可以实现视觉和语言模态的无缝集成。具体方法是将预训练的VisualTokenformer和LanguageTokenformer的key-valueparameterTokens统一为一个参数集，然后引入新的TrainableTokens来执行视觉-语言对齐和指令微调。

端云协同

Tokenformer可以在设备-云协作中充当云端知识库，为设备端的大语言模型（LLM）提供支持，其中每组key-valueparametertokens代表一个可学习模式，通过设备进行实时处理，并利用云端执行密集任务。

增强模型的可解释性

由于Tokenformer完全基于注意力机制，它自然受益于在Token-Parameter交互中与注意力相关的可解释性特性。这一特点增强了模型的可解释性，为AI社区开发更透明、易理解的模型贡献力量。