TIDE：每一层都知晓上下文中的令牌

# TIDE: 每一层都知晓上下文之下的 Token 身份

来源: https://arxiv.org/html/2605.06216
Lauren Hannah, Han-Byul Kim, Duc Hoang, Mehrdad Farajtabar, Minsik Cho
Apple

*(2026年5月7日)*

###### 摘要

我们重新审视了现代大语言模型（LLM）中普遍接受但研究不足的一个设计选择：Token索引仅在输入嵌入层查找一次，随后被永久丢弃。这种*单次注入假设*导致了两个结构性缺陷：(i) **稀有Token问题**，即词汇量的Zipf型分布导致稀有Token的嵌入由于接收到的累积梯度信号远少于常见Token而长期训练不足；以及 (ii) **上下文坍塌问题**，即参数有限的模型将分布相似的Token映射到无法区分的隐藏状态。为了应对这两个问题，我们提出了**TIDE**，它在标准Transformer基础上引入了**嵌入记忆（Embedding Memory）**：由$K$个独立的**记忆块（Memory Blocks）**组成的集成，将Token索引映射到与上下文无关的语义向量，该向量计算一次后，通过带有可学习空位库（learnable null bank）的深度条件Softmax路由器注入到每一层。我们从*理论上*和*实证上*确立了TIDE在解决与单次Token身份注入相关问题方面的优势，并在多种语言建模及下游任务中*提升了性能*。

## 1 引言

扩展现代大语言模型（LLMs）涉及投入大量的表示能力，通过创新注意力机制、扩大前馈模块以及堆叠深层Transformer层来实现对Token的*上下文语境化*。相比之下，在最近的进展中，一个关键的LLM组件——**Token索引**（唯一明确标识Token身份的信息）——被广泛忽视。Token索引仅在输入嵌入层查找一次，随后被永久丢弃。所有$L$层Transformer中的后续计算都作用于语境化的隐藏状态，不再直接查询正在处理的词汇表条目。这种单次注入假设导致了两种不同的失效模式：

❶ **稀有Token问题**：自然语言词汇表遵循幂律缩放，具体为Zipf定律（zipf1949human; pilgrim2021bias）：频率最高的1%的Token占据了语料库中约~80%的出现次数。在随机梯度下降（SGD）下，每个Token嵌入的累积梯度信号与其频率成正比（第2.1节），导致稀有Token的嵌入（如稀有命名实体、技术术语、低频形态形式）持续训练不足（图1）。

❷ **上下文隐藏状态坍塌**：在训练期间，前馈网络（FFN）被迫进行*表示过载*，同时执行残差流的结构变换并作为Token特定事实知识的主要存储库（meng2022rome; dai2022knowledgeneurons）。Token索引在中间层从未被重新咨询，FFN区分深层中两个Token的唯一机制依赖于残差和注意力输出的上下文混合。然而，当两个语义不同的Token出现在几乎相同的句法环境中时，上下文提供的区分信号有限，导致它们在整个网络中的隐藏状态变得几乎无法区分（图2）。

受这些挑战的启发，我们提出一个关键问题：**如何为每个Transformer层提供持久的、依赖于Token身份的知识，使其独立于上下文残差流并与其互补？** 与侧重于事后分析FFN（geva2022vocabspace; meng2022rome; meng2023memit）或在推理时改装外部检索（lewis2020rag; borgeaud2022retro; izacard2023atlas）的先前方法不同，我们采用了另一种方法：从头设计和训练一种新的Transformer架构，该架构维护一个由静态Token身份信息直接索引的专用语义记忆。

在这项工作中，我们提出了**TIDE**（Token Identity Delivered Everywhere），这是对标准Transformer的一种架构修改，它维护一个由Token身份直接索引的专用语义记忆（图3）。TIDE引入了**嵌入记忆**，即由$K$个独立的**记忆块**组成的集成，每个块将Token索引映射到静态且与上下文无关的学习语义向量，这些向量可以通过持久化、Token条件化的信号并行注入到每个Transformer层中，与上下文残差流并行。我们的主要贡献总结如下：

- **架构方面**：TIDE引入了*Token级统一嵌入记忆*，为Token级梯度累积提供了$K$个不相交的路径。记忆嵌入张量在*每次前向传播中计算一次*，并通过基于注意力后隐藏状态的逐层Softmax路由机制注入到*每一个*Transformer层。
- **理论方面**：我们形式化了标准Transformer中的两种失效模式，并证明TIDE (i) 渐近地泛化了标准Transformer；(ii) 将每个Token的累积梯度信号放大了$K$倍；以及 (iii) 通过暴露离散的、Token索引的输入（无需义务绑定到隐藏状态），绕过了FFN的Lipschitz约束。
- **实证方面**：我们实证验证了TIDE显著有利于*稀有*Token并缓解了上下文坍塌问题。在从350M到11B参数的模型规模范围内，TIDE在各种语言建模数据集（*如* Wikitext, PubMed, DCLM）以及下游任务（*如* HellaSwag, ARC, PIQA）上， consistently 提供优于标准Transformer的性能提升。

> **图1说明**：稀有Token嵌入保持训练不足的实证证据：(a) LLaMa-Base-1B预训练检查点的平均嵌入$l_2$-范数，显示从稀有到常见bin的范数*单调增加*；(b) 稀有和常见Token的嵌入范数分布：*宽*的稀有分布与*窄*的常见峰值的存在证实了稀有嵌入仍然以噪声为主导且训练不足；(c) 每500亿Token在中间训练检查点上的分bin范数增长率：随着训练的继续，稀有Token的范数表现出*单调下降*，而常见Token的范数持续增长。

## 2 当上下文不够时：诊断标准Transformer

### 2.1 稀有Token问题。

##### 梯度饥饿瓶颈：

在批量大小为$B$、序列长度为$T$、每个Token的平方梯度范数有界为$G^2$的小批量SGD下，Token $v$的嵌入$e_v \in \mathbb{R}^d$仅在$v$出现在当前批次中时接收非零梯度。在此设置中，经过$\tau$步训练后，期望累积平方梯度范数满足：

$$ \mathbb{E}\left[\sum_{s=1}^{\tau}\|\nabla_{e_v}\mathcal{L}_s\|^2\right] \leq \tau \cdot f_v \cdot B \cdot T \cdot G^2, \quad (2.1) $$

其中$f_v := \Pr[\text{均匀抽取的Token位置等于 } v] \in (0,1)$是$v$的单字概率，且$\sum_{v \in \mathcal{V}}f_v=1$。如果对于某些$\epsilon \ll 1/(BT)$有$f_v=\epsilon$，则Token $v$是*稀有*的；如果对于某些不依赖于$|\mathcal{V}|$的常数$c>0$有$f_u \geq c$，则Token $u$是*常见*的。方程(2.1)的完整推导见附录C。

| 层级 | Bin $(s)$ | 数量 | $f_v$ | $\mathbb{E}[N_v]$ |
| :--- | :--- | :--- | :--- | :--- |
| 单现词 (最稀有) | 0 | 1 | $8.3 \times 10^{-9}$ | $\approx 1,660$ |
| 近单现词 | 1 | $\sim 4$ | $3.3 \times 10^{-8}$ | $\approx 6,640$ |
| 不常见 | 2 | $\sim 10$ | $8.3 \times 10^{-8}$ | $\approx 16,600$ |
| 中频 | 3–6 | $\sim 10^{2-3}$ | $\sim 10^{-6}$ | $\approx 10^{5-6}$ |
| 常见 (最高频) | 7–9 | $\sim 10^6$ | $8.3 \times 10^{-3}$ | $\approx 1.66 \times 10^9$ |

**表1**：使用2000亿训练Token时，Token分bin的期望非零梯度更新$\mathbb{E}[N_v]$。

在一个使用LLaMA-3分词器（$|\mathcal{V}|=128,256$）分词的Wikitext-103（merity2016pointer）示例语料库中生成频率bin（附录B），*稀有*和*常见*Token之间的梯度差异变得严重。在2000亿Token的训练预算下，设$B=8, T=2048$，Token $v$的嵌入的期望非零梯度更新次数为：

$$ \mathbb{E}[N_v] = \tau\left(1-(1-f_v)^{BT}\right) \approx \tau \cdot f_v \cdot BT \quad \text{for small } f_v. \quad (2.2) $$

参考附录B中定义的分bin，表1实例化了训练数据集中2000亿Token的这一情况，说明了稀有和常见Token之间存在巨大的梯度更新差异。此外，从图1(c)中可以实证推断出，这种差异不仅仅局限于冷启动伪影，而是随着训练的推进*单调*增长。稀有Token的范数下降，而常见Token的范数持续增加。

**稀有和常见Token的梯度信号比率**：对于稀有$v$（$f_v=\varepsilon$）和常见$u$（$f_u \geq c > 0$），设$G_{\min}^2 > 0$为条件于Token $u$出现在批次中时，每步平方梯度范数的下界。累积梯度信号的比率满足：

$$ \frac{\mathbb{E}\left[\sum_s \|\nabla_{e_v}\mathcal{L}_s\|^2\right]}{\mathbb{E}\left[\sum_s \|\nabla_{e_u}\mathcal{L}_s\|^2\right]} \leq \frac{\varepsilon BT G^2}{\kappa G_{\min}^2} = O\left(\varepsilon/c\right) \quad (2.3) $$

其中$\kappa := 1-(1-c)^{BT} > 0$，且$BT, G^2, G_{\min}^2$为固定的正常数。完整推导见附录C.1。对于表1中的实证实例，稀有Token（Bin 1）与常见Token（Bin 9）之间的比率为$\varepsilon/c \approx 10^{-6}$，即在相同训练预算下，稀有和常见Token之间的梯度信号存在六个数量级的差异。

> **图2说明**：上下文坍塌的实证证据：热力图展示了来自三个上下文坍塌示例类别的250个模板句子中，Token对隐藏状态（LLaMa-Base-1B）之间的平均$l_2$-距离$\|h^{(\ell)}_u - h^{(\ell)}_v\|$。对于所有采样的对，在大多数层（除了末尾）距离保持*接近零*，证实了上下文坍塌的存在。

### 2.2 上下文坍塌与FFN的盲点。

如前所述，梯度饥饿问题导致稀有Token嵌入收敛到低范数、嘈杂的表示。更严重的是，当两个不同的Token携带相似幅度的训练不良嵌入时，会出现更深层的结构性故障：这些Token在所有Transformer层中产生的隐藏状态可能变得无法区分，这在共享相似上下文时尤为 problematic。我们将这种失效模式形式化，并表明这是由FFN连续域施加的Lipschitz连续性带来的固有后果。

##### 上下文坍塌现象：

在每一层$\ell$，Token $v$的隐藏状态$h^{(\ell)}_v \in \mathbb{R}^d$由作用于周围上下文的注意力机制产生。当两个Token $u \neq v$出现在几乎相同的句法环境中时，例如语法同音词（*their* 或 *there*）、数字身份Token（*1847*, *1851*, 或 *1849*），或稀有领域特定同义词（*ibuprofen* 或 *acetaminophen*），上下文不提供区分信号，因此注意力对两者产生相似的输出。

我们将其正式定义为：

###### 定义 2.1（上下文坍塌集）.

对于容忍度 $\delta > 0$，层 $\ell$ 的*上下文坍塌集*可以正式定义为：

$$ \mathcal{C}_\delta^{(\ell)} := \{(u,v) \in \mathcal{V}^2 : u \neq v, \ \|h^{(\ell)}_u - h^{(\ell)}_v\| \leq \delta\}, $$

其中隐藏状态是在代表性上下文语料库上平均的。

图2提供了LLaMa-Base-1B标准模型中上下文坍塌的直接实证证据，该证据是使用150个仅在被考虑的单Token对上有差异的模板句子估计的。对于三个示例经典类别中的每一个，平均$l_2$距离$\|h^{(\ell)}_u - h^{(\ell)}_v\|$在整个深度轴上保持持续较小，除了最后几层，证实了坍塌的普遍存在。注意，这种现象在数值Token类别中更为严重，即使在最后一层的隐藏状态中也存在显著的坍塌（小$\delta$）。

###### 命题 2.2（坍塌Token的FFN近似下界）.

设 $(u,v) \in \mathcal{C}_\delta^{(\ell)}$ 为坍塌Token对，并设 $g: \mathcal{V} \to \mathbb{R}^d$ 为满足 $\|g(u)-g(v)\|=C>0$ 的任何目标函数。那么对于*任何*权重选择 $W_1, W_2$：

$$ \max\left\{\|\mathrm{FFN}(h_u)-g(u)\|, \ \|\mathrm{FFN}(h_v)-g(v)\|\right\} \geq \frac{C - L_{\mathrm{FFN}}\delta}{2}. $$

当 $C > L_{\mathrm{FFN}}\delta$ 时，右侧严格为正：无论拥有多少参数，FFN都不能以任意精度逼近坍塌对 $(u,v)$ 上的 $g$。

**证明概要**。由于 $(u,v) \in \mathcal{C}_\delta^{(\ell)}$，Lipschitz界限强制 $\|\mathrm{FFN}(h_u)-\mathrm{FFN}(h_v)\| \leq L_{\mathrm{FFN}}\delta$。对目标分离 $C=\|g(u)-g(v)\|$ 应用三角不等式并代入此界限，得到 $\|\mathrm{FFN}(h_u)-g(u)\| + \|\mathrm{FFN}(h_v)-g(v)\| \geq C - L_{\mathrm{FFN}}\delta$。由于两个非负项的最大值至少是其和的一半，结果得证。详见附录D。

在此界限中，$\delta$ 由嵌入和注意力层决定；在FFN作用之前它是固定的。分离目标 $C$ 由下游任务决定。**Lipschitz常数** $L_{\mathrm{FFN}}$ 是FFN控制的唯一项，但在实践中它是受限的，因为大的 $L_{\mathrm{FFN}}$ 会放大*所有*输入扰动，降低大多数非坍塌Token的性能。该界限暴露了一个结构性限制：给定固定的上游表示，无论宽度如何，没有任何FFN能在不 destabilize 其他输入的情况下解决坍塌Token对。