Tokenformer: 下一代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的核心，主要包括以下几个步骤：

• 输入：假设有一个长度为TTT的序列，每个token用一个维度为ddd的向量表示，记为矩阵 X∈RT×dX \in \mathbb{R}^{T \times d}X∈RT×d。

• 线性投影： 将输入 XXX 通过三个不同的线性层，分别得到查询（Query）、键（Key）和值（Value）矩阵，其中，WQ,WK,WV∈Rd×dkW_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}WQ​,WK​,WV​∈Rd×dk​ 是可学习的权重矩阵，dkd_kdk​ 是查询和键的维度。：

Q=X⋅WQ,K=X⋅WK,V=X⋅WVQ = X \cdot W_Q, \quad K = X \cdot W_K, \quad V = X \cdot W_VQ=X⋅WQ​,K=X⋅WK​,V=X⋅WV​

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

Attention(Q,K,V)=softmax(Q⋅K⊤dk)⋅V\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q \cdot K^\top}{\sqrt{d_k}} \right) \cdot VAttention(Q,K,V)=softmax(dk​​Q⋅K⊤​)⋅V

• 输出投影： 将注意力输出通过一个线性层进行投影，其中，WO∈Rdv×dW_O \in \mathbb{R}^{d_v \times d}WO​∈Rdv​×d 是输出投影矩阵。：

O=Attention(Q,K,V)⋅WOO = \text{Attention}(Q, K, V) \cdot W_O O=Attention(Q,K,V)⋅WO​

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

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

2.2 Transformer 架构的缺陷

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

2.2.1 模型的拓展性是什么

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

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

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

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

1. 参数数量固定： 传统Transformer的线性层 WQ,WK,WVW_Q, W_K, W_VWQ​,WK​,WV​ 的维度是固定的。例如，如果输入维度 ddd 增加，那么每一层transformer中线性层的维度、输出投影层的维度也必须进行修改。这意味着模型的整体参数数量会急剧增加。

2. 需要重新训练： 如果要增加模型的规模（如增加 ddd 或 dkd_kdk​），必须从头开始训练整个模型。耗时而且需要大量计算资源。尤其是随着模型规模的增加，重新训练的成本（时间和计算资源）呈指数级增长，导致在实际应用中难以持续扩展模型。

3. TokenFormer的解决方案

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

3.1 模型参数Token化

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

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

KP∈Rn×d1,VP∈Rn×d2K_P \in \mathbb{R}^{n \times d_1}, \quad V_P \in \mathbb{R}^{n \times d_2}KP​∈Rn×d1​,VP​∈Rn×d2​

其中，nnn 是参数tokens的数量，d1d_1d1​ 和 d2d_2d2​ 分别是输入和输出的维度。

3.2. Token-Parameter Attention（Pattention）层

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

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

Pattention(X,KP,VP)=Θ(X⋅KP⊤)⋅VP\text{Pattention}(X, K_P, V_P) = \Theta\left(X \cdot K_P^\top\right) \cdot V_PPattention(X,KP​,VP​)=Θ(X⋅KP⊤​)⋅VP​

1. 注意力分数（Θ\ThetaΘ部分）的计算如下：

Sij=f(Aij×τ∑k=1n∣Aik∣2),∀i,j∈{1,…,n}S_{ij} = f\left( \frac{A_{ij} \times \tau}{\sum_{k=1}^n |A_{ik}|^2} \right), \quad \forall i, j \in \{1, \dots, n\}Sij​=f(∑k=1n​∣Aik​∣2Aij​×τ​),∀i,j∈{1,…,n}

其中，A=X⋅KP⊤A = X \cdot K_P^\topA=X⋅KP⊤​，τ\tauτ 是缩放因子（默认设置为nnn），fff 是非线性函数（如GeLU）。

总结一下，Pattention层的详细计算流程如上图所示，首先输入的tokens：X∈RT×d1X∈R^{T×d_1}X∈RT×d1​与KP∈Rn×d1K_P \in \mathbb{R}^{n \times d_1}KP​∈Rn×d1​计算出注意力分数，而后计算结果和VP∈Rn×d2V_P \in \mathbb{R}^{n \times d_2}VP​∈Rn×d2​相乘，即：

Pattention(X,KP,VP)=Θ(A)⋅VPA=X⋅KP⊤\text{Pattention}(X, K_P, V_P) = \Theta(A) \cdot V_P \\ A = X \cdot K_P^\topPattention(X,KP​,VP​)=Θ(A)⋅VP​A=X⋅KP⊤​

4. 总体结构

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

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

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

Q=Pattention(X,KPQ​​,VPQ​​

K=Pattention(X,KPK​​,VPK​​)

V=Pattention(X,KPV​​,VPV​​)

Xatt​=softmax(d​Q⋅K⊤​)⋅V

​Oatt​​=Pattention(Xatt​,KPO​​,VPO​​)​

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

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

Offn=Pattention(Xffn,KPffn,VPffn)O_{\text{ffn}} = \text{Pattention}(X_{\text{ffn}}, K_{P_{\text{ffn}}}, V_{P_{\text{ffn}}})Offn​=Pattention(Xffn​,KPffn​​,VPffn​​)

其中，XffnX_{\text{ffn}}Xffn​ 是经过Layer Normalization后的中间表示。

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

4.3 与transformer的比较

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

 

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

5. 可扩展性

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

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

KPscale=[KPold,KPnew]VPscale=[VPold,VPnew]K_{P_{\text{scale}}} = [K_{P_{\text{old}}}, K_{P_{\text{new}}}]\\ V_{P_{\text{scale}}} = [V_{P_{\text{old}}}, V_{P_{\text{new}}}]KPscale​​=[KPold​​,KPnew​​]VPscale​​=[VPold​​,VPnew​​]

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

初始化策略：新增的参数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 更优。