论文日记 | Tokenformer: 下一代Transformer架构-优快云博客

作者 | 大鲸鱼编辑 | 自动驾驶之心

原文链接：https://zhuanlan.zhihu.com/p/14266356312

点击下方卡片，关注“自动驾驶之心”公众号

>>点击进入→自动驾驶之心『Transformer』技术交流群

本文只做学术分享，如有侵权，联系删文

1. 导言

Transformer架构已经成为当今大模型的基石，不管是NLP还是CV领域，目前的SOTA模型基本都是基于Transformer架构的，比如NLP中目前的各种知名大模型，或者CV中的Vit等模型

本次介绍的论文标题为：Tokenformer: Rethinking Transformer Scaling with Tokenized Model Parameters,” 顾名思义，本文提出了Tokenformer架构，其优势在于增量学习能力：在增加模型尺寸时，无需从头开始重新训练模型，大大降低了成本。 本文代码已开源。

2. Transformer vs Tokenformer - 结构比较

首先我们从顶层设计的角度，对于传统 Transformer 架构和本文提出的 Tokenformer 架构进行比较，如下图所示：

2.1 Transformer 架构

自注意力机制是Transformer的核心，主要包括以下几个步骤：

计算注意力分数：通过查询和键的点积，再除以缩放因子，并通过Softmax函数得到归一化的注意力权重：

如上图所示，一个Transformer层主要由两个部分组成：

多头自注意力机制（Multi-Head Self-Attention） ：输入首先经过一个线性投影模块，以计算注意力模块的输入，即矩阵 Q、K 和 V。然后利用子注意力机制计算出Token之间的权重
前馈神经网络（Feed-Forward Network, FFN） ：对于注意力层的输出进行投影，计算出下一层的输入

2.2 Transformer 架构的缺陷

传统Transformer在处理token与参数的交互时，依赖于固定数量的线性投影，这限制了模型的扩展性，这句话本身较难理解，因此接下来详细论述架构的缺陷。

2.2.1 模型的拓展性是什么

模型的拓展性（Scalability）指的是模型在需要更强大性能时，能够有效地增加其规模（如参数数量、计算能力等）而不导致性能下降或计算成本过高的能力。

简而言之，拓展性好的模型可以在保持或提升性能的同时，灵活且高效地扩大其规模。

2.2.2 为什么说传统Transformer的固定线性投影限制了模型的扩展性

固定线性投影指的是，Transformer中用来生成查询、键、值的权重矩阵,,是预先定义且固定的。这带来了以下几个限制：

参数数量固定：传统Transformer的线性层,,的维度是固定的。例如，如果输入维度d增加，那么每一层transformer中线性层的维度、输出投影层的维度也必须进行修改。这意味着模型的整体参数数量会急剧增加。
需要重新训练：如果要增加模型的规模（如增加d或），必须从头开始训练整个模型。耗时而且需要大量计算资源。尤其是随着模型规模的增加，重新训练的成本（时间和计算资源）呈指数级增长，导致在实际应用中难以持续扩展模型。

3. TokenFormer的解决方案

为了解决模型维度固定导致的模型缺乏拓展性的问题，TokenFormer提出了一种创新的方法，通过将模型参数视为tokens，并利用注意力机制来处理token与参数之间的交互，从而实现更高效、更灵活的模型扩展。

3.1 模型参数Token化

传统Transformer将参数（如,,）作为固定的权重矩阵来处理。而TokenFormer将这些参数表示为一组可学习的tokens。具体来说：

参数Tokens：原本transformer模型的Q、K、V投影层不再是固定的矩阵，而是转化为一组向量（tokens），例如：

其中,n是参数tokens的数量，和分别是输入和输出的维度。

3.2. Token-Parameter Attention（Pattention）层

Pattention层是TokenFormer的核心创新，它通过注意力机制来处理token与参数之间的交互。从而替代原本的Q,K,V，具体过程如下：

查询与参数Tokens交互：输入tokens作为查询（Query） ，参数tokens作为键（Key）和值（Value） ，通过注意力机制进行交互，其中，X是输入tokens，和是参数tokens，是一种修改后的softmax函数，用于稳定优化：

注意力分数（部分）的计算如下：

4. 总体结构

为方便阅读再把图扔到这：

与传统transformer结构相同，其总体上也包括两层：多头自注意力层和前馈网络层。

4.1 多头自注意力（Single-Head Variant：

其中，,等是不同投影的参数tokens。即，首先计算出Q, K, V，然后同样计算自注意力，然后将计算结果放入Pattention层。

4.2 前馈网络（Feed-Forward Network, FFN）

其中，是经过Layer Normalization后的中间表示。

这里也可以看到，相对于Transformer，Tokenformer就是将所有的投影层从固定的全连接网络也变成了Pattention层。

4.3 与transformer的比较

下方公式左侧代表传统Transformer的自注意力机制，右侧代表tokenformer的自注意力机制：

从上边的图中可以清楚看到，相对于transformer，本论文只是将投影层与连接层替换成了新的层。

5. 可扩展性

之前说过，相对于transformer，tokenformer主要是解决可拓展性的问题，那么假设我们要增加参数数量，或者要增加输入维度，tokenformer如何进行增量学习？

如上图所示，若需要扩展模型，可以简单地追加新的和tokens：

这样，模型的参数量可以按需扩展。

初始化策略：新增的参数tokens初始化为零，类似于LoRA技术（Low-Rank Adaptation），确保模型能够在保持原有知识的基础上，快速适应新的参数扩展。

6. 实验部分

与从零重训练的 Transformer 相比，如上图所示，Y 轴代表模型性能，X 轴代表训练成本。蓝线代表使用 3000 亿个 token 从头开始训练的 Transformer 模型，不同的圆圈大小代表不同的模型大小。

其他线条代表 Tokenformer 模型，不同颜色代表不同的Token数量。例如，红线从 1.24 亿个参数开始，扩展到 14 亿个参数，其训练集为从300B token中抽样出的30B Token。最终版本模型的性能与相同规模的 Transformer 相当，但训练成本却大大降低。

黄线显示，使用 60B个Token来训练的增量版本在更低的训练成本下，性能已经比 Transformer 更优。

① 2025中国国际新能源技术展会

自动驾驶之心联合主办中国国际新能源汽车技术、零部件及服务展会。展会将于2025年2月21日至24日在北京新国展二期举行，展览面积达到2万平方米，预计吸引来自世界各地的400多家参展商和2万名专业观众。作为新能源汽车领域的专业展，它将全面展示新能源汽车行业的最新成果和发展趋势，同期围绕个各关键板块举办论坛，欢迎报名参加。

② 国内首个自动驾驶学习社区

『自动驾驶之心知识星球』近4000人的交流社区，已得到大多数自动驾驶公司的认可！涉及30+自动驾驶技术栈学习路线，从0到一带你入门自动驾驶感知（端到端自动驾驶、世界模型、仿真闭环、2D/3D检测、语义分割、车道线、BEV感知、Occupancy、多传感器融合、多传感器标定、目标跟踪）、自动驾驶定位建图（SLAM、高精地图、局部在线地图）、自动驾驶规划控制/轨迹预测等领域技术方案、大模型，更有行业动态和岗位发布！欢迎扫描加入

③全网独家视频课程

端到端自动驾驶、仿真测试、自动驾驶C++、BEV感知、BEV模型部署、BEV目标跟踪、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、车道线检测、轨迹预测、在线高精地图、世界模型、点云3D目标检测、目标跟踪、Occupancy、CUDA与TensorRT模型部署、大模型与自动驾驶、NeRF、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频（扫码即可学习）