摘要

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

1 引言

设计强大的神经网络架构是机器学习领域长期以来的目标。基础模型（FMs）的最新发展表明了Transformer（Vaswani等，2017）作为通用计算架构的潜力。凭借其灵活性和可扩展性，Transformer在自然语言处理（NLP）（Radford等，2018；Alec等，2019；Brown等，2020）、视觉建模（Dosovitskiy等，2021；Liu等，2021）、视觉-语言（Liu等，2023；Wang等，2024）、图表示（Ying等，2021）和3D视觉（Wang等，2023a；b）等各种领域都取得了最先进的性能。

Transformer通常将处理单个令牌所需的计算分为两个不同部分：与其他输入令牌的交互（令牌-令牌交互）和涉及模型参数的计算（令牌-参数交互）。注意力机制（Vaswani等，2017）促进了令牌-令牌交互，使现代通用基础模型能够将多模态数据编码为统一的令牌序列，并有效地捕获它们之间的复杂依赖关系（Liu等，2023；Zhu等，2023；Wang等，2023d）。相反，令牌-参数计算在很大程度上依赖于线性投影（Dunford＆Schwartz，1988），其中输入令牌乘以一组固定参数。这种规定的设计限制了可扩展性，因为增加模型规模需要更改核心架构组件，这通常需要对整个模型进行从头开始的重新训练。随着模型规模的扩大，这导致资源消耗过大，变得越来越不切实际。在本文中，我们引入了一种新架构，增强了令牌-参数交互的灵活性，允许逐步扩展模型参数并有效重用先前训练的模型，从而显著降低了训练负担。

为了实现这一目标，我们引入了Tokenformer，这是一种全新架构，它通过完全使用注意力机制来统一令牌-令牌和令牌-参数交互的计算。我们的令牌-参数注意力层的灵活性以及其处理可变数量参数的能力，从根本上增强了模型的可扩展性，促进了逐步高效扩展。

如图1所示，我们通过保留输入令牌之间的计算模式，并使用交叉注意力机制重新公式化所有线性投影，来扩展Transformer架构。具体来说，为了将具有输入和输出维度$D_{1}$和$D_{2}$的特征进行投影，我们使用两组参数，每组包含$N$个可学习令牌，其通道维度分别为$D_{1}$和$D_{2}$。在这种公式中，输入令牌作为查询，模型参数作为键和值。这种灵活性使我们的模型参数能够随着变量$N$的变化而自然扩展，从而可以通过不断添加新的键值参数对来实现高效扩展。图1显示，我们的模型可以从1.24亿参数逐步扩展到14亿参数，达到与从头开始训练相似的性能，同时节省了超过一半的训练成本。

本文的主要贡献概括如下：1) 如图1所示，我们提出了Tokenformer，这是一个完全基于注意力机制的神经网络，它将模型参数视为令牌（tokens），从而最大化了令牌-参数计算的灵活性，同时在语言和视觉领域的标准基准测试中均取得了具有竞争力的性能。2) 得益于这种设计，我们的模型可以通过逐步添加新的键值参数对来自然地扩展。与从头开始训练的方法（Biderman等人，2023；Kaplan等人，2020）相比，我们的方法在实现近乎相同性能的同时，极大地降低了训练成本。

2 相关工作

Transformer（Vaswani等人，2017）凭借其多功能的注意力机制，已成为深度学习领域的基础架构，能够处理任何令牌化的数据，并适应包括语言建模（Radford等人，2018；Touvron等人，2023）、图像处理（Dosovitskiy等人，2021）、多模态理解（Liu等人，2023；Wang等人，2024；2023b；2022）、决策制定（Chen等人，2021b）、图学习（Yun等人，2019）等众多领域。虽然Transformer有效地以灵活的方式处理输入令牌之间的交互，但这种特性并不扩展到涉及模型参数的计算，这些计算是通过规定的线性投影进行的。在本文中，我们通过开发一个完全基于注意力的网络来重构令牌-参数交互，该网络通过注意力机制统一了令牌-令牌和令牌-参数的计算，从而进一步扩展了网络的灵活性。

大规模训练已被证明是开发强大基础模型的有效方法。如GPT系列模型（Radford等人，2018；Alec等人，2019；Brown等人，2020）所示，当得到更大训练数据集和模型规模（以参数衡量）的支持时，简单的架构往往比更复杂的算法表现更好。扩大数据规模通常更具成本效益，因为它独立于模型的架构，并允许通过微调现有模型来持续集成新数据（Kaplan等人，2020）。相比之下，增加模型规模通常会产生极高的成本，因为它会改变架构细节，并且通常需要在每个扩展步骤中从头开始重新训练整个数据集（Biderman等人，2023）。这极大地提高了在行业中构建逐步扩大规模模型的费用。

模型重用。以前重用模型的方法通常涉及通过复制（Chen等人，2015；2021a）、堆叠（Gong等人，2019）或组合（Wang等人，2023c）模型权重，用预训练的小型模型初始化更大的模型。虽然这些方法可能有效，但它们通常会干扰小型模型预先建立的分布，增加了丢失预训练知识和收敛速度减慢的风险。相比之下，我们的模型能够以自然且无缝的方式实现参数扩展，并保留了现有模型的完整性。

3 方法论

在本节中，我们首先在第3.1节中回顾了传统的注意力机制。然后，第3.2节介绍了Tokenformer，这是一种以灵活的令牌-参数注意力层为中心的、可原生扩展的架构。最后，第3.3节详细介绍了Tokenformer的增量模型扩展。

3.1 预备知识

Transformer模型（Vaswani等人，2017）已在深度学习中确立了其基础架构的地位，并在广泛的任务中表现出色。它们成功的基石是自注意力机制，该机制允许模型动态评估每个令牌的重要性，并有效地建模它们之间的复杂依赖关系。

给定一组包含T个输入令牌 $X \in \mathbb{R}^{T \times d}$ ，其中通道维度为d，自注意力模块首先通过三个不同的线性投影推导出依赖于输入的查询Q、键K和值V，分别为

$Q=X \cdot W^{Q}, \quad K=X \cdot W^{K}, \quad V=X \cdot W^{V},$

其中 $W^{Q}, W^{K} \in \mathbb{R}^{d \times d_{k}}$ 和 $W^{V} \in \mathbb{R}^{d \times d_{v}}$ 是可学习的权重矩阵。通过测量查询和键向量之间的相似性，然后应用softmax函数获得归一化权重来计算注意力分数。这些分数随后用于计算缩放点积注意力的输出，即

$\operatorname{Attention}(Q, K, V)=\operatorname{softmax}\left[\frac{Q \cdot K^{\top}}{\sqrt{d}}\right] \cdot V,$

其中 $\sqrt{d}$ 是用于缓解softmax引起的小梯度的缩放因子。最后，输出为

$O=X_{\mathrm{att}} \cdot W^{O},$

其中 $X_{\text {att }}$ 是注意力输出，$W^{O} \in \mathbb{R}^{d_{v} \times d}$ 是输出投影矩阵。

上述架构设计使模型能够灵活地管理不同长度令牌之间的交互，从而使现代通用模型能够同时处理任何形式和数量的令牌化多模态数据。这种能力显著推动了当前人工智能领域的发展，并且是基于Transformer的系统成功的关键。

3.2 Tokenformer

尽管Transformer在各个领域都表现出色，但其可扩展性受到由规定的标记-参数交互（即线性投影）引起的高计算开销的限制。因此，调整架构组件（例如通道维度）的扩展策略通常需要从头开始重新训练整个模型，导致计算资源利用效率低下。

为了克服这一挑战，我们提出了Tokenformer，一种完全基于注意力机制的架构。Tokenformer的核心创新是标记-参数注意力（Pattention）层，它引入了一组可训练的标记作为模型参数，然后采用交叉注意力来管理输入标记和这些参数标记之间的交互。通过这种方式，Pattention层引入了一个额外的维度——参数标记的数量——该维度独立于输入和输出通道维度。这种解耦使得输入数据能够与可变数量的参数动态交互，为通过重用预训练模型实现渐进式模型扩展提供了所需的灵活性。因此，在训练更大模型时，速度大大加快，同时实现了与从头训练的Transformer相当的性能。

Pattention层。设输入标记和输出标记分别为$\mathcal{I} \in \mathbb{R}^{T \times d_{1}}$和$\mathcal{O} \in \mathbb{R}^{T \times d_{2}}$，其中$T$是序列长度，$d_{1}$和$d_{2}$分别是输入和输出维度。为了实现我们的Pattention机制，我们引入了两组$n$个可学习的参数标记：$K_{P} \in \mathbb{R}^{n \times d_{1}}$表示键，$V_{P} \in \mathbb{R}^{n \times d_{2}}$表示值。来自缩放点积Pattention层的输出$\mathcal{O}$计算如下：

$\operatorname{Pattention}\left(X, K_{P}, V_{P}\right)=\Theta\left(X \cdot K_{P}^{\top}\right) \cdot V_{P}$

其中，$\Theta$是用于Pattention层稳定优化的修改后的softmax操作。输出Pattention分数$S \in \mathbb{R}^{n \times n}$的计算公式为：

$S_{i j}=f\left(\frac{A_{i j} \times \tau}{\sqrt{\sum_{k=1}^{n}\left|A_{i k}\right|^{2}}}\right), \quad \forall i, j \in 1 \ldots n$

其中，$A$是从$\left(X \cdot K_{P}^{\top}\right)$得出的分数，$\tau$是缩放因子，默认设置为$\sqrt{n}$，$f$是非线性函数，在我们的公式中设置为GeLU函数（Hendrycks & Gimpel, 2016）。这种设计提高了我们架构中的梯度稳定性，并且与标准softmax操作相比，性能更好（详见附录A和表4）。

我们的Pattention层采用交叉注意力机制来管理标记和参数之间的交互，从而完全保留了注意力机制的适应性特征。类似于Transformer模型中的自注意力如何处理可变长度的序列，我们的Pattention层被设计为独立于特征投影中使用的输入和输出通道维度来处理可变数量的参数。这使得网络参数可以沿着参数标记轴无缝扩展，从而能够有效地重用预训练权重，并为模型扩展提供了一种自然的渐进式方式。

总体架构。图2展示了Tokenformer的架构。给定输入标记$X_{\text {in }} \in \mathbb{R}^{T \times d}$，我们遵循预归一化Transformer的设计，Tokenformer层的输出计算如下：

$X_{\text {inter }}=X_{\text {in }}+\operatorname{MHA}\left(\operatorname{LN}\left(X_{\text {in }}\right)\right)$

$X_{\text {out }}=X_{\text {inter }}+\operatorname{FFN}\left(\operatorname{LN}\left(X_{\text {inter }}\right)\right)$

其中，LN表示层归一化（Ba, 2016; Zhang & Sennrich, 2019），MHA和FFN分别指我们修改后的多头自注意力和前馈层。

在多头自注意力块中，为了简化，我们考虑单头变体，并将$d_{k}$和$d_{v}$都设置为$d$。然后，我们用我们的Pattention层替换所有线性投影。设$\mathrm{LN}\left(X_{\text {in }}\right)$为$X$，该块公式如下：

$\begin{aligned} Q=\operatorname{Pattention}\left(X, K_{P}^{Q}, V_{P}^{Q}\right), K & =\operatorname{Pattention}\left(X, K_{P}^{K}, V_{P}^{K}\right), V=\operatorname{Pattention}\left(X, K_{P}^{V}, V_{P}^{V}\right), \\ X_{\mathrm{att}} & =\operatorname{softmax}\left[\frac{Q \cdot K^{\top}}{\sqrt{d}}\right] \cdot V, \\ O_{\text {att }} & =\operatorname{Pattention}\left(X_{\text {att }}, K_{P}^{O}, V_{P}^{O}\right), \end{aligned}$

其中，等式8和10表示标记-参数注意力，而等式9表示标记-标记注意力。QKV投影的键值参数标记为$\left(K_{P}^{Q}, V_{P}^{Q}\right) \in \mathbb{R}^{n_{q} \times d},\left(K_{P}^{K}, V_{P}^{K}\right) \in \mathbb{R}^{n_{k} \times d},\left(K_{P}^{V}, V_{P}^{V}\right) \in \mathbb{R}^{n_{v} \times d}$，而$\left(K_{P}^{O}, V_{P}^{O}\right) \in \mathbb{R}^{n_{o} \times d}$用于输出投影层。

为了保持一致性和简洁性，Tokenformer中的前馈块使用单个Pattention层。设$\mathrm{LN}\left(X_{\text {inter }}\right)$为$X_{\mathrm{ffn}}$，FFN计算由下式给出：

$O_{\mathrm{ffn}}=\operatorname{Pattention}\left(X_{\mathrm{ffn}}, K_{P}^{\mathrm{ffn}}, V_{P}^{\mathrm{ffn}}\right)$

其中，$\left(K_{P}^{\mathrm{ffn}}, V_{P}^{\mathrm{ffn}}\right) \in \mathbb{R}^{n_{\mathrm{ffn}} \times d}$是FFN块的可学习键值对。

通过这种方式设计架构，我们将包括输入数据和模型参数在内的所有基本组件表示为计算框架内的标记。这种以标记为中心的视角允许利用成功的注意力机制来统一Transformer中的两种主要计算，即标记-标记和标记-参数交互，从而建立了一个以完全注意力为基础、具有非凡灵活性的神经网络。

架构配置。我们的模型精确地反映了标准Transformer架构的超参数配置。以GPT-2（Radford等人，2018）为例，其包含12层Transformer和一个768的隐藏维度，我们的模型复制了这一配置，具有相同的层数和维度。查询-键-值投影和输出投影中的键值参数对的数量直接对应于隐藏维度。相比之下，前馈神经网络（FFN）模块使用的参数对数量是隐藏大小的四倍。这种架构对齐使得我们的模型能够使用预训练的Transformer来初始化参数，从而确保无缝集成到Transformer预训练生态系统中。

3.3 渐进式模型扩展

我们的模型在参数轴上表现出对大规模模型训练的强大适应性，这归功于Pattention层的通用设计，它允许通过重用较小、预训练模型的参数来逐步开发更大的模型。
为了便于理解而不牺牲一般性，我们采用单个Pattention层来举例说明模型扩展的复杂性。考虑一个现有的Tokenformer模型，其配备了一组预训练的键值参数标记，记为 $K_{P}^{\text {old }}, V_{P}^{\text {old }} \in \mathbb{R}^{n \times d}$。如图2所示，为了扩展模型，我们通过添加新的键值参数标记 $K_{P}^{\text {new }}, V_{P}^{\text {new }} \in \mathbb{R}^{m \times d}$ 来增强这个集合，如下所示：

$K_{P}^{\text {scale }}=\left[K_{P}^{\text {old }}, K_{P}^{\text {new }}\right], \quad V_{P}^{\text {scale }}=\left[V_{P}^{\text {old }}, V_{P}^{\text {new }}\right]$

其中，$\left[{ }^{\text {old }},{ }^{\text {new }}\right]$ 表示沿标记维度的拼接操作，$K_{P}^{\text {scale }}, V_{P}^{\text {scale }} \in \mathbb{R}^{(m+n) \times d}$ 是扩展后的参数集。扩展模型的前向传递定义为：

$O=\operatorname{Pattention}\left(X, K_{P}^{\text {scale }}, V_{P}^{\text {scale }}\right)$

这种扩展方案允许集成任意数量的参数，而无需改变输入或输出维度。如图3所示，这种方法显著提高了更大规模模型的训练效率，而不会降低性能。重要的是，通过将 $K_{P}^{\text {new }}$ 初始化为零（类似于LoRA技术（Hu等人，2022）），我们的模型可以完美地恢复预训练阶段的模型状态，而不会丢失已学习的知识，从而加快了收敛速度并加速了整体扩展过程。

4 实验

在本节中，我们展示了上述技术的实验结果。第4.1节验证了我们的模型持续扩展的能力。第4.2节突出了我们的模型在处理语言和视觉领域任务方面的有效性。第4.3节进行了深入比较，突出了我们的模型相对于标准Transformer模型的优势。最后，第4.4节详细描述了为评估Tokenformer中每个模块的重要性而进行的消融实验。

4.1 渐进式模型扩展

数据集。我们的模型使用Gokaslan和Cohen（2019）中描述的OpenWebText语料库进行训练。该语料库被广泛认为是OpenAI专有WebText数据集的开源近似版本，该数据集用于GPT-2（Alec等，2019）的开发。该数据集包含从8,013,769个Reddit共享文档中提取的文本数据。在训练过程中，我们随机从这些文档中抽取片段。

从头开始训练基线Transformer。为了评估我们渐进式模型扩展策略的有效性，我们通过从头开始训练一个Transformer模型来建立基线。遵循Karpathy（2022）和Kaplan等（2020）概述的训练程序，我们使用了AdamW优化器（Loshchilov和Hutter，2019），批处理大小为512个序列，每个序列包含1024个标记。为了与我们逐步扩展的方法进行公平比较，我们根据训练标记的总数配置了两种训练变体。第一种变体进行了$6 \times 10^{5}$步训练（大约3000亿个标记），与Karpathy（2022）为复制GPT-2性能所采用的训练步骤一致。第二种变体被限制为$6 \times 10^{4}$步训练（大约300亿个标记），以确保与我们的渐进式扩展的每个阶段具有可比性。在我们分析中包括的所有训练中，除非另有说明，否则采用$6 \times 10^{-4}$的学习率，先进行2000步预热，然后采用余弦衰减至零。

带渐进式模型扩展的Tokenformer。在上述训练协议的基础上，我们测试了模型扩展的性能，参数大小从1.24亿到14亿不等。与上述从头开始训练的方法不同，每次扩展迭代都利用预先训练好的较小Tokenformer来部分初始化3.3节中描述的较大模型的权重。扩展过程首先从大约3000亿个标记上从头开始训练初始源模型开始，这与Transformer基线类似。对于扩展，我们选择参数数量与目标大小最接近的预训练模型进行权重初始化。例如，为了训练一个具有3.54亿参数的模型，我们使用1.24亿参数的模型作为部分初始化器，并使用减少的计算预算（例如，150亿、300亿或600亿个标记）对整个模型进行重新训练。此迭代过程继续进行，扩展到7.57亿参数，然后是14亿参数。值得注意的是，为了简化扩展过程，在整个过程中，新参数和现有参数都使用相同的训练超参数进行等效训练。

我们的训练优化了1024个标记上下文上的平均自回归对数似然（即交叉熵损失），并将测试集上评估的对数困惑度作为测试分数。

实验分析。如图3所示，采用Tokenformer的渐进式扩展方法实现了与从头开始训练的Transformer模型相当的性能，同时大幅减少了训练预算。具体来说，从在3000亿个标记上训练的1.24亿参数模型开始，我们逐步扩展到3.54亿、7.57亿和14亿参数，仅需额外的300亿个标记，即与从头开始训练的Transformer相比，计算预算仅为十分之一。此扩展过程在14亿参数水平上实现了11.77的测试困惑度。相比之下，同样大小的从头开始训练的Transformer模型实现了11.63的相似困惑度，但训练成本是其三倍。重要的是，我们的方法报告了包括所有扩展阶段的累计训练成本，而Transformer基线仅考虑单个阶段。即使在这种比较下，我们的方法也显示出比从头开始训练Transformer更低的计算成本，从而验证了我们的方法的有效性。

图4展示了我们的模型和标准Transformer在每个扩展阶段的训练成本。与图3相比，成本节省更为显著。具体来说，我们的模型所需的训练成本仅为Transformer基线成本的十分之一。为了减轻不同训练数据的影响，我们还包含了使用等效计算预算（300亿个标记）从头开始训练的Transformer的性能曲线。在相同的计算约束下，我们逐步扩展的模型实现了11.77的更低困惑度，而Transformer的困惑度为13.34，从而突出了我们方法在效率和可扩展性方面的优越性。

4.2 模型表达能力基准测试

语言建模。我们通过标准的自回归语言建模任务评估了我们提出的架构的有效性，并与现有的基于Transformer的模型进行了基准测试。评估使用预训练指标（特别是困惑度）和零样本性能指标进行。训练在Pile数据集（Gao等，2020）上进行，遵循Biderman等（2023）中描述的训练协议。详细的训练程序和模型大小（深度和宽度）见附录F。

表1展示了Tokenformer在各种广泛认可的零样本下游任务中的性能。我们将其与同等规模的主要开源Transformer模型进行了比较，特别是Pythia（Biderman等，2023），它使用了与我们模型相同的分词器、数据集和训练时长（3000亿个标记）。如表所示，与标准Transformer相比，我们的模型实现了具有竞争力的性能，证明了我们的架构在作为基础模型时的表达能力潜力。

视觉建模。表2验证了我们的模型在视觉任务中的表达能力。我们将我们的方法与在ImageNet-1K数据集（Deng等，2009）上通过监督学习训练的标准视觉Transformer（ViT）（Dosovitskiy等，2021）进行了比较。为了公平比较，我们使用了MMDetection代码库（MMDetection贡献者，2018），并遵循了He等（2022）中使用的超参数和训练策略。如表所示，我们的模型在视觉建模中达到了与ViT相同的性能，证实了其在视觉任务中的表达能力。

4.3 与标准Transformer的比较

Transformer也可以在一定程度上实现模型复用。Net2Net（Chen等，2015）是一种经典的模型增长方法，它提出了一种通过复制神经元来扩展神经网络宽度的技术。在这种方法中，使用小型模型中Transformer层的预训练权重矩阵$W_{s}^{\text {old }} \in \mathbb{R}^{d_{s} \times d_{s}}$，来创建一个更大的权重矩阵$W_{l}^{\text {new }} \in \mathbb{R}^{d_{l} \times d_{l}}\left(d_{l}>d_{s}\right)$，以填充更大的模型。这种扩展的公式如下，

$W_{l}^{\text {new }}=\left[\begin{array}{cc}W_{s}^{\text {old }} & W_{l(12)}^{\text {new }} \\W_{l(21)}^{\text {new }} & W_{l(22)}^{\text {new }}\end{array}\right]$，

其中，$W_{l(12)}^{\text {new }} \in \mathbb{R}^{\left(d_{l}-d_{s}\right) \times d_{s}}$，$W_{l(21)}^{\text {new }} \in \mathbb{R}^{d_{s} \times\left(d_{l}-d_{s}\right)}$，和$W_{l(22)}^{\text {new }} \in \mathbb{R}^{\left(d_{l}-d_{s}\right) \times\left(d_{l}-d_{s}\right)}$是扩展所需的新参数。缩放过程与第4.1节中介绍的方案相同。

用于长上下文建模的符号-符号交互的可控成本。思维链（Chain-of-Thought，CoT）建模（Wei等，2022）的最新进展强调了在大型语言模型（Large Language Models，LLMs）中高效处理长文本序列的关键重要性（Tay等，2020）。如第11节所述，Transformer架构的训练成本主要分为两部分：涉及模型参数的交互和输入序列之间的交互。表3表明，基于Transformer的模型的计算复杂度在文本长度上呈二次依赖关系，与符号-参数交互呈线性关系，与符号-符号交互呈二次关系。因此，在扩展模型参数的同时，控制符号-符号交互部分的计算负担是至关重要的。

传统上，缩放Transformer模型涉及增加通道维度。对于固定的文本长度，这会导致更高的计算成本，主要是因为主导的符号-符号交互变得更加密集，这阻碍了模型在长文本上的性能。我们提出的模型采用了一种不同的方法，将符号-符号交互的计算成本从模型缩放中解耦出来。我们增加了参数大小，而不改变符号通道维度，从而保持了与符号-符号交互相关的计算成本。如图5所示，随着参数数量的增加，我们的模型在Transformer上展现出了越来越显著的计算优势，尤其是在处理较长序列时。

在不丢失已学习分布的情况下进行缩放。当新的关键参数初始化为零时，我们的Tokenformer可以保持现有的输出分布。这一特性有利于连续缩放模型以纳入更多数据，因为它可以在不中断当前训练过程的情况下增加模型容量，从而促进快速收敛。

为了评估Tokenformer的缩放效率，我们将基于Net2Net的Transformer缩放与Tokenformer缩放的损失曲线进行了比较。两个模型最初都有3.54亿个参数，并在OpenWebText数据集上进行了预训练。然后，我们引入了EnWik8数据集，并继续训练一个周期，将模型扩展到7.57亿个参数，以适应新数据。图6表明，Tokenformer不仅收敛速度更快，而且达到了更低的最终损失，这归因于其在恢复训练期间能够保持输出分布的能力。

增量缩放性能基准测试。在本研究中，我们使用前面详细介绍的Net2Net方法逐步扩展标准Transformer。为了公平比较，我们调整了所有超参数，包括参数大小、学习率、数据集等。如图7所示，与标准Transformer相比，我们的模型在缩放方面表现更好。

4.4 消融研究

注意力层中的优化Softmax函数。在符号-参数注意力层中，我们解决了与传统softmax函数相关的梯度减小所导致的训练不稳定性问题。传统的softmax操作主要包括两个步骤：计算注意力分数的指数值，然后进行$L_{1}$归一化。如表4所示，为了缓解小梯度问题，我们将指数非线性函数替换为GeLU函数（Hendrycks & Gimpel, 2016），从而在ImageNet分类基准上实现了+2.1个百分点的性能提升。随后，我们将$L_{1}$归一化替换为$L_{2}$归一化，又带来了+0.8个百分点的额外提升。这些修改共同使我们的模型达到了与标准Vision Transformer相当的性能。

非参数层归一化。为了实现模型扩展并为后续研究合并两个单独训练的参数符号集，我们通过移除可训练权重和偏置，将Transformer的层归一化修改为非参数变体。这种调整确保了模型中只有键值参数需要学习。表5中的实证结果表明，在丢弃可学习权重和偏置后，模型保持了相当的性能。

5 未来工作

扩展混合专家（Mixture-of-Experts）范式。我们将Tokenformer解释为混合专家（Mixture of Experts，MoE）框架的一个极端实例，其中每个键值参数对都作为一个单独的专家。这种创新的类似MoE的架构有可能显著降低与符号-参数交互相关的计算成本。此外，Tokenformer在符号-符号交互方面的可调节计算负载补充了MoE特性，有助于开发更具资源效率的基础模型。

推进参数高效调优。Tokenformer的缩放方法涉及整合额外的键值参数对，这体现了参数高效调优的一种策略。当面临新任务或数据集时，模型可以通过纳入这些新的参数符号来扩展其预训练参数，从而快速适应特定任务的要求。

整合视觉和语言模型。利用Tokenformer的参数高效调优能力，我们可以实现视觉和语言模态的无缝整合。这可以通过将来自预训练的视觉Tokenformer和语言Tokenformer的键值参数符号统一为一个参数集来实现。然后，引入新的可学习符号来执行视觉-语言对齐和指令调优。
设备-云端协同。在设备端大型语言模型（LLMs）的设备-云端协同中，Tokenformer可以作为云端知识库，其中每一对键值参数令牌代表一个可学习的模式，利用设备进行实时处理，利用云端处理密集型任务。

增强模型可解释性。由于Tokenformer完全基于注意力机制，它天然受益于令牌参数交互中注意力相关的可解释性。这一特性增强了模型的可解释性，有助于人工智能领域开发更透明、更易理解的模型。

6 结论

本文介绍了Tokenformer，这是一种自然可扩展的架构，利用注意力机制不仅促进了令牌间的计算，还促进了令牌与模型参数间的交互，从而增强了架构的灵活性。通过将模型参数表示为令牌，我们用我们的Pattention层替换了Transformer中的所有线性投影层，实现了无缝且高效的增量扩展，而无需从头开始重新训练。我们认为，这种架构比传统的Transformer提供了更大的灵活性，将进一步推动基础模型的发展。

附录

Pattention层的梯度

我们的令牌参数注意力机制采用$L_{2}$归一化，随后是GeLU（Hendrycks & Gimpel, 2016）激活函数，这与标准令牌-令牌注意力层中传统的SoftMax函数不同，后者采用指数变换后接$L_{1}$归一化。这一设计选择是基于我们的实验观察，即SoftMax倾向于增加输出的幅度，经常将它们推入梯度变得极小的区域，导致整体性能较差（见表4）。
具体来说，给定一个查询令牌和$n$个键值对，维度为$d$，设查询令牌和键令牌之间的相似度分数为$A \in \mathbb{R}^{1 \times n}$。在标准SoftMax注意力中，注意力分数$S \in \mathbb{R}^{1 \times n}$计算如下：

$S_{i}=\frac{\exp \left(A_{i} / \sqrt{d}\right)}{\sum_{j=1}^{n} \exp \left(A_{j} / \sqrt{d}\right)}, \quad \forall i \in 1 \ldots n$

SoftMax函数关于$A_{i}$的导数为：

$\frac{\partial S_{i}}{\partial A_{j}}=\frac{1}{\sqrt{d}} S_{i}\left(\mathbb{1}_{i=j}-S_{j}\right)=\left\{\begin{array}{ll}\frac{1}{\sqrt{d}} S_{i}\left(1-S_{j}\right) & i=j \\-\frac{1}{\sqrt{d}} S_{i} S_{j} & i \neq j\end{array}\right.$

相比之下，我们的激活函数使用$L_{2}$归一化后接GeLU函数。设GeLU为$f$，注意力分数$Z \in \mathbb{R}^{1 \times n}$计算如下：

$\hat{S}_{i}=f\left(Z_{i}\right)=f\left(\frac{A_{i} \times \sqrt{n}}{\sqrt{\sum_{j=1}^{n}\left|A_{j}\right|^{2}}}\right), \quad \forall i \in 1 \ldots m$

对于我们的注意力函数当$i=j$时的导数，我们有：

$\begin{aligned}\frac{\partial \hat{S}_{i}}{\partial A_{i}} & =f^{\prime} \sqrt{m} \frac{\|A\|_{2}-A_{i} \frac{A_{i}}{\|A\|_{2}}}{\|A\|_{2}^{2}} \\& =f^{\prime} \sqrt{n} \frac{\|A\|_{2}^{2}-A_{i}^{2} \frac{1}{\|A\|_{2}}}{\|A\|_{2}^{2}} \\& =f^{\prime} \sqrt{n} \frac{1}{\|A\|_{2}} \frac{\|A\|_{2}^{2}-A_{i}^{2}}{\|A\|_{2}^{2}} \\& =f^{\prime} \sqrt{n} \frac{1}{\|A\|_{2}}\left(1-Z_{i}^{2} d\right) \\& =f^{\prime} \frac{1}{\sqrt{n}} \frac{1}{\|A\|_{2}}\left(d-Z_{i}^{2}\right)\end{aligned}$

当$i \neq j$时，导数为：

$\begin{aligned}\frac{\partial \hat{S}_{i}}{\partial A_{i}} & =-f^{\prime} \sqrt{n} \frac{A_{i} \frac{A_{j}}{\|A\|_{2}}}{\|A\|_{2}^{2}} \\& =-f^{\prime} \sqrt{n} \frac{1}{\|A\|_{2}} \frac{A_{i} A_{j}}{\|A\|_{2}^{2}} \\& =-f^{\prime} \frac{1}{\sqrt{n}} \frac{1}{\|A\|_{2}} Z_{i} Z_{j}\end{aligned}$

因此，我们的注意力函数的导数为：

$\frac{\partial \hat{S}_{i}}{\partial a_{j}}=\left\{\begin{array}{ll}f^{\prime} \frac{1}{\sqrt{n}} \frac{1}{\|A\|_{2}}\left(n-Z_{i} Z_{j}\right) & i=j \\-f^{\prime} \frac{1}{\sqrt{n}} \frac{1}{\|A\|_{2}} Z_{i} Z_{j} & i \neq j\end{array}\right.$

比较SoftMax（等式16）和我们方法（等式26）的梯度，关键区别在于我们的梯度依赖于乘积$Z_{i} Z_{j}$，而SoftMax依赖于$S_{i} S_{j}$。由于SoftMax的指数性质，$S_{i}$的分布往往更尖锐且更集中，这通常会使梯度趋于零。相反，我们的激活函数产生了更平滑的$Z$分布，缓解了梯度消失问题，并实现了更稳定的训练动态。

Tokenformer中的零初始化

如3.3节所示，在模型扩展过程中，将新键参数初始化为零可以使模型在最小干扰下继续训练。这是因为零初始化保持了模型的原始输出分布，防止了对已学习表征的重大干扰。
在本节中，我们证明了当新添加的参数被零初始化时，Pattention层对这些参数是不变的。设$X \in \mathbb{R}^{d}$为输入向量，并设Pattention层有$n$个键值对，表示为$K_{P}, V_{P} \in \mathbb{R}^{n \times d}$。Pattention层的输出计算如下：
$\begin{array}{c}A=K_{P} \cdot X, \\S=f\left(\frac{A}{\sqrt{\sum_{j} A_{j}^{2}}}\right), \\O=V_{P}^{\top} \cdot S.\end{array}$
当通过添加$m$个零初始化的新键值对来扩展模型时，输出变为：
$\begin{array}{l}\hat{A}=\left[\begin{array}{c}K_{P} \\0, \ldots, 0 \\\vdots \\0, \ldots, 0\end{array}\right] \cdot X=\left[\begin{array}{c}A \\0 \\\vdots \\0\end{array}\right], \\\hat{S}=f\left(\frac{\hat{A}}{\sqrt{\sum_{j} \hat{A}_{j}^{2}}}\right)=\left[\begin{array}{c}S \\0 \\\vdots \\0\end{array}\right], \\\hat{O}=\left[V_{P}^{\top}, V_{P}^{\text {new } T}\right] \cdot \hat{S}=O.\end{array}$
由于新添加的键参数被初始化为零，注意力机制不会修改原始输出。因此，输出$\hat{O}$与$O$保持一致。这一特性对于扩展模型是有利的，因为它在不中断已学习分布和正在进行的训练过程的情况下增加了模型容量，从而加快了收敛速度。

C 主要结果的表格

在这里，我们提供了与主论文中图3相对应的表格结果。表7展示了在OpenWebText验证集上的困惑度。Transformer模型是从头开始训练的，而Tokenformer模型则利用较小模型的参数重用（除了第一个具有1.24亿参数的Tokenformer模型，它也是从头开始训练的）。我们观察到，Tokenformer实现了与从头开始训练的Transformer相当的性能，但由于参数重用，训练成本显著降低。值得注意的是，最大的Tokenformer（14亿）甚至实现了比其对应的Transformer（11.60 vs 11.63）略好的困惑度。

表6比较了从头开始训练的Transformer模型和通过参数重用训练的Tokenformer模型在训练期间看到的令牌数量不同的情况下的性能。显然，使用相同数量看到令牌的Transformer无法达到使用参数重用的Tokenformer的性能水平。这表明参数重用成功地将知识从小模型迁移到大模型，减少了训练时间而没有牺牲性能。

D 渐进式模型扩展的实验细节

在本实验中，我们使用OpenWebText数据集（Gokaslan＆Cohen，2019）来评估模型扩展能力。该数据集包含8,013,769个文档，我们从中随机选择5％作为验证集，并报告该子集的困惑度。我们研究了四种模型大小：1.24亿、3.54亿、7.57亿和14亿参数。模型规格请参见表8。在Tokenformer的参数重用过程中，我们部分恢复旧模型的参数，并向Pattention层添加新的键值对，而不改变层数或特征维度。
Transformer和Tokenformer的训练方案相同：我们没有在两个模型中实现dropout，并且在最终层使用嵌入层计算logits。分词器来自GPT-NeoX-20B（Black等，2022）。我们使用AdamW优化器（Loshchilov，2019），其中$\beta_{1}=0.9$，$\beta_{2}=0.95$。学习率设置为$6 \times 10^{-4}$，在2000步内进行线性预热，然后采用余弦衰减至零。训练使用的批大小为512，序列长度为1024。

E 在不丢失已学习分布的情况下进行扩展的实验细节

在本实验中，我们使用EnWik8（Mahoney，2011）数据集来评估模型持续适应新数据的能力。EnWik8数据集包含以XML格式存在的英语维基百科的前1亿个字节。
我们从在OpenWebText（Gokaslan＆Cohen，2019）上预训练的3.54亿参数模型开始，然后将其扩展到Transformer和Tokenformer模型的7.57亿参数。对于Transformer，3.54亿和7.57亿模型仅在特征维度和多头注意力机制中的头数上有所不同，我们遵循Net2Net（Chen等，2015）的方法进行参数扩展。对于Tokenformer，我们将键值对的数量从2140增加到4850。
我们采用恒定学习率$6 \times 10^{-4}$，对数据集进行一次遍历，总共进行2204步训练。我们使用了AdamW优化器（Loshchilov，2019），并且没有从任何先前的运行中恢复优化器的内部状态。批大小设置为512，与预训练期间使用的批大小一致。

F 语言建模基准测试的实验

我们使用标准的自回归语言建模基准测试来评估我们提出架构的表达能力和性能，并将其结果与包括基于RNN的方法（Peng等，2023；Gu＆Dao，2023）在内的现有开源LLM的结果进行比较，如表9所示。评估使用了预训练指标（特别是困惑度）和零样本性能指标。对于训练，我们使用了单个周期的Pile数据集（Gao等，2020），并遵循Biderman等（2023）使用的训练策略。应用Adam优化器，学习率设置为$6 \times 10^{-4}$，每个批次包含1024个样本，序列长度为2048个令牌，这与Mamba的设置相同。训练过程包括14,300步，其中1430步为预热，相当于总训练步骤的1％。详细的模型规格请参见表8。