改进Transformer，解读Tokenformer论文：基于参数分词化重新思考Transformer的扩展策略

Transformer 训练成本高昂的问题日益凸显，不仅需要耗费巨额的资金与大量的计算资源，还对环境产生了不可忽视的影响，最近由北京大学与谷歌联合发表的一篇论文，为这一棘手难题带来了全新的曙光。论文中提出的创新方案，有望打破 Transformer 训练成本居高不下的困境，

一、研究动机

如今，很难想象人工智能领域如果没有Transformer会怎样。这些模型是大型语言模型的基础架构，为人工智能带来了革命性变化。而且，它们的影响力不仅限于自然语言处理领域。在计算机视觉等其他领域，视觉Transformer（ViT）也发挥着重要作用。随着技术的发展，模型规模不断扩大，从头开始训练模型的成本越来越高，且这种做法难以持续，还引发了环境方面的担忧。

训练Transformer模型成本高昂

二、引入Tokenformer

我们今天要解读的论文题目是。该论文对Transformer架构提出了一项引人注目的改进，名为Tokenformer。Tokenformer能够在增加模型规模时，避免从头开始重新训练模型，从而大幅降低成本。

在这里插入图片描述

三、Transformer与Tokenformer架构对比

让我们从整体上对比一下传统Transformer架构和Tokenformer，相关对比图来自论文。
在这里插入图片描述

3.1 Transformer整体流程

左边展示的是原始Transformer架构的简化图。底部是输入的词元序列。

输入首先经过线性投影层，计算注意力模块的输入，即矩阵Q、K和V。这一阶段涉及模型参数和输入词元（token）之间的相互作用，通过线性投影计算得出。
然后，自注意力组件让输入词元之间相互作用，这一过程通过注意力模块计算。
最后，前馈神经网络（FFN），为下一层生成输出，这同样体现了词元和参数之间通过线性投影计算的相互作用。

3.2 Tokenformer的核心思路

通常，词元与参数之间的相互作用是通过固定参数大小的线性投影组件计算的。这就导致在增加模型规模时，必须从头开始训练。

Tokenformer的主要思路是构建一个完全基于注意力机制的模型，包括词元与参数之间的相互作用，从而拥有更灵活的架构，支持逐步增加参数数量。

3.3 Tokenformer整体流程

在上面图的右侧，我们可以看到Tokenformer架构的简化图。

为了计算自注意力模块的输入（Q、K和V矩阵），我们将输入词元输入到一个名为词元 - 参数注意力的新组件中。在这个组件中，除了输入词元，我们还会传入参数。输入词元代表词元 - 参数注意力模块中的查询部分，参数则代表键和值部分。
接着，我们使用与之前相同的自注意力组件。
为了为下一层准备输出，我们用另一个词元 - 参数注意力模块取代前馈神经网络。在这个模块中，我们将自注意力模块的输出作为查询，同时为键矩阵和值矩阵使用不同的参数。

3.4 Tokenformer架构详解

在这里插入图片描述

计算自注意力模块的输入：深入研究Tokenformer架构，输入词元在经过归一化处理后，会被送入三个Pattention模块（Pattention即词元 - 参数注意力模块的名称）。三个Pattention模块各自拥有权重，分别负责为自注意力模块创建输入，即每个Pattention模块生成一个Q、K或V矩阵。
Pattention与注意力模块的区别：
- 原始注意力模块计算公式为：（此处原文未给出原始注意力公式）。
- Pattention模块计算公式为：
  $K_p, V_p)=\Theta(X ⋅ K_p^T ) ⋅ V_p$ 。
- 输入词元作为查询，键矩阵和值矩阵由Pattention模块的参数给出。与注意力模块相比，Pattention模块使用了带有不同缩放机制的改进softmax函数（即 $\Theta$ 函数），以提高优化的稳定性。
- $X$ 与 $K$ 相乘后的改进softmax函数将输入词元与键参数进行匹配，得到注意力分数。然后，利用V参数矩阵对注意力分数进行加权，生成该模块的最终输出。
用Pattention取代FFN：在标准自注意力模块处理完三个Pattention模块的输出后，我们不再使用前馈神经网络，而是使用两个连续的Pattention模块。从上图左下角可以看到，在这两个模块之间，我们通过残差连接将其与输入词元合并。
模型的增量式扩展：在架构图的右下角，展示了通过添加新参数来逐步增加模型规模的示意图。我们在每个Pattention模块的键矩阵和值矩阵中追加更多的参数词元行。已训练的参数词元会被保留，无需从头开始训练。然后，我们训练规模更大的模型，并在这个过程中更新所有词元。从下面的结果可以看出，与从头开始训练相比，增量式训练的模型训练速度要快得多。

四、实验结果

在这里插入图片描述

上图展示了Tokenformer相对于从头开始训练的Transformer在性能和训练成本方面的对比。

蓝色线代表使用3000亿词元从头开始训练的Transformer模型。不同大小的圆圈代表不同的模型规模，中间灰色圆圈列出了具体规模。其他线条代表Tokenformer模型，每种颜色代表训练增量版本时使用的不同词元数量。例如，红线代表模型从1.24亿参数扩展到14亿参数，在用于从头训练的3000亿词元中，仅使用了300亿词元进行训练。最终版本在性能上与相同规模的Transformer相当，但训练成本却大幅降低！黄线表明，使用600亿词元训练增量版本的模型，性能优于Transformer，且训练成本更低。