注意力机制的数学重构与参数重参数化

Transformer 的核心魅力在于其注意力层(Attention Layer),它通过 Query 和 Key 的内积计算、Softmax 加权,将前一层嵌入重新加权,从而实现对下一个 Token 的精准预测。然而,直接分析包含 $W_Q, W_K, W_U$ 等大量参数的原始梯度方程极其复杂且混乱。为了解决这一难题,研究团队引入了关键的重参数化(Reparameterization)技巧,将复杂的模型参数转化为两个独立变量:$y$ 代表解码器顶层权重,$z$ 代表自注意力矩阵的成对对数几率(pairwise logits)。

这种转化使得动态分析变得可行。正如演讲者所述:

"So we actually make two very important tricks. So first of all, was that you want to make sure that you reparameterize the problem with independent variable... So this is actually try to reparameterize these many variables into y and z."

通过这种数学上的简化,研究者能够更清晰地写出 $y$ 和 $z$ 的动态方程,为后续理解多层 Transformer 的行为奠定了坚实的数学基础。这一策略不仅适用于单层模型,也为后续更复杂的多层结构分析提供了通用框架。

单层注意力模型的理论假设与设定

为了深入理解注意力机制的本质,研究首先聚焦于单层设置(One-layer settings)下的简单输入输出结构。输入由上下文 Token ($x_1$ 到 $x_{t-1}$) 和一个查询 Token 组成,经过自注意力机制生成概率分布,再通过归一化、编码层和 Softmax 预测下一个 Token。在此框架下,研究基于几个关键假设展开:

  1. 无位置编码:暂时忽略位置编码,因为已有研究表明即使没有位置编码,模型仍能获得合理性能。
  2. 序列长度趋于无穷:假设序列长度无限长,以简化数学推导。
  3. 学习率差异:解码器 $Y$ 的学习率显著大于自注意力层 $Z$ 的学习率。

此外,数据分布被假定为包含 $k$ 个序列类别,其中 Token 分为两类:Distinct Tokens(独特 Token)仅出现在一个类别中,而 Common Tokens(常见 Token)则出现在多个类别中。

"We assume there's a bunch of sequence classes from one to k... There's one token called distinct tokens that only appears in one class but not the others. There's also common tokens. That appears in multiple sequence classes."

这些假设构建了一个理想化的数学环境,使得研究者能够剥离噪声,专注于注意力机制本身的动态演化规律。

注意力动态演化的核心结论:CL Title 分析

研究的核心发现围绕一个关键统计量 CL Title 展开,它定义为共现概率(来自数据集)与注意力矩阵成对对数几率的指数之积。通过分析 CL Title 的行为,研究者揭示了自注意力层在不同类型 Token 上的差异化动态:

对于 Common Tokens(常见 Token),由于其在多个序列类别中出现,梯度方向一致指向减小注意力分数。数学上表现为 $z$ 对时间的导数小于零($\frac{dz}{dt} < 0$),这意味着随着训练进行,常见 Token 的注意力权重会被不断压制,逐渐变小。

相反,对于 Distinct Tokens(独特 Token),由于其仅在一个特定类别中出现,对应的注意力对数几率 $z$ 会随着时间增长而显著上升。这种分化现象表明,Transformer 的注意力机制具有天然的类别区分能力:它会自动抑制跨类别的常见模式,同时强化特定类别的独特特征。

"For common tokens... the gradient will basically drag all the attention score of the common tokens error to be smaller and smaller... On the other hand, you can actually show for distinct tokens. Their corresponding z... will grow over time."

这一结论揭示了 Transformer 在单层层面如何通过动态调整注意力权重,实现对数据内在结构的高效提取和分类。

注意力机制作为可学习的TF-IDF模型

在多层Transformer的训练动力学分析中,注意力分数(Attention Score)的变化揭示了其内在的统计学习机制。随着训练进行,注意力分数会变得更加紧密(tighter),这种机制类似于自然语言处理早期广泛使用的TF-IDF(词频-逆文档频率)模型。TF-IDF的核心逻辑在于平衡两个标准:一是寻找在特定类别中高频出现的词(Term Frequency),二是抑制在多个文档中普遍出现但信息量低的词(Inverse Document Frequency)。

"So from these two, you can already see the attention mechanism is kind of like the learnable TF-IDF"

具体而言,第一类标准关注对每个分类至关重要的词汇,而第二类标准则要求模型过滤掉那些跨文档过于常见、缺乏信息量的Token。例如,在文档分类任务中,通用词汇应当被抑制,而特定类别的专属词汇(如体育类中的"Messi")则应获得更高的注意力权重。这种机制使得模型能够更精准地进行文档分类,通过强化特定Token与类别之间的连接来提升决策质量。

上下文稀疏性与注意力分数的增长速率

进一步地,研究者量化了注意力分数的增长速率,发现其具有显著的上下文稀疏性(Contextual Sparsity)。在自注意力矩阵中,元素 $Z_{ML}$ 随时间变化,其增长速度与给定序列类别下的条件概率密切相关。理论定理表明,不同Token的增长比率取决于它们与Query的共现概率。

"this is called the contextual sparsity, because over time you're going to see this C title, which is attention score, will actually become sparser and sparser after normalization"

当某个Distinct Token与Query具有高共现概率时,其注意力分数会迅速增长,远高于那些仅偶然共现的Token。这种"赢家通吃"(Winner takes all)效应导致归一化后的注意力分布越来越稀疏:富者愈富,贫者愈贫。最终,只有少数几个Token脱颖而出。这种稀疏性是依赖Query的(Query-dependent),即不同的Query会触发完全不同的注意力增长曲线和模式,体现了动态的上下文选择机制。

两阶段注意力动力学:扫描与锁定

尽管直观上认为注意力会持续呈现"赢家通吃"的增长,但实际分析揭示了两阶段效应:首先是注意力扫描(Attention Scanning)阶段,此时"富者愈富"效应显著,注意力模式快速变化;其次是注意力锁定(Attention Snapping)阶段,增长放缓并趋于饱和。

"One is called attention scanning... On the other side, the second state called attention snapping. In which everything slows down a little bit, and finally it saturates"

在扫描阶段后,注意力模式会冻结,形成一种"少数派通吃"(Minority take all)的稳定结构,而非无限的赢家通吃。这一结论基于"Scan and Snap"论文的分析,并通过WikiText2等数据集的实验得到验证。在特定的学习率结构下,训练过程中确实观察到了注意力模式的稀疏化变化。为了证明这一理论,研究者采用了简化策略:利用无限序列长度将项转化为概率分布统计量,并假设解码器(Decoder)的学习率远快于自注意力矩阵(Self-attention matrix),从而将注意力矩阵视为常数来分析解码器动力学,再将梯度反向传播回注意力层进行分析。

从单层到多层:引入DRAMA框架

早期的"Scan and Snap"分析存在诸多理想化假设,如忽略嵌入向量、残差连接,且假设解码器学习率远快于注意力层,这在实践中并不成立。为克服这些局限,后续研究引入了DRAMA框架,旨在更深入地理解训练动力学。

DRAMA工作的核心创新在于发现MLP下层与自注意力层之间的动力学性质,并将它们结合到一个简化的修改版MLP(Modified MLP)中。这一模型不仅捕捉了两层结构的动态交互,还比单独分析自注意力和MLP层更为简洁。通过消除之前的严格假设,DRAMA为理解多层Transformer中更复杂的层间动态提供了更贴近实际的理论工具,标志着从单层简化模型向多层复杂系统分析的演进。

Drama框架的核心架构与残差机制

Drama(Dynamics of Residual Multi-layer Attention)框架旨在深入剖析多层Transformer内部的训练动态。与早期仅关注线性近似的研究不同,Drama明确引入了残差连接(Residual Connections)和非线性激活函数(Nonlinearity)。演讲者指出,之前的分析往往忽略了残差连接的作用,而Drama通过将这些元素重新纳入模型,展示了它们在Transformer中的实际角色。特别是在Decoder层中,模型引入了非线性函数 $\phi$,这与之前假设从输入到Softmax全程线性的模型形成了鲜明对比。这种架构上的改进使得研究者能够揭示更丰富的动态特性。正如演讲者所述:

"So for this one, we basically like put like a drag down like a drag one layer of the transformers, right? And then we put like a bunch of symbols to analyze what's going on."

通过这种更贴近实际Transformer结构的建模,Drama不仅保留了计算的可解释性,还增强了理论模型对真实训练过程的拟合能力。

嵌入矩阵的正交性假设与实证验证

Drama框架建立在一个关键假设之上:嵌入矩阵(Embedding Matrix U)中的向量在训练过程中保持相互正交。尽管审稿人曾对此假设提出质疑,但演讲者通过实证数据证明了这一假设在实践中的有效性。为了验证这一点,研究团队绘制了不同层之间嵌入向量的余弦相似度随训练动态变化的图表。余弦相似度的取值范围为0到1,其中1表示完美对齐,0表示完全正交。

实验结果显示,在训练过程中,不同层嵌入向量之间的余弦相似度并未显著增加,而是稳定在极低的水平。这意味着即使在训练的不同阶段,嵌入向量依然近似保持正交状态。具体数据如下表所示,展示了PyTear模型及其他测试模型在不同层间的相似度表现:

模型名称 测试维度 平均余弦相似度范围 正交性状态描述
PyTear 多层对比 ~0.03 几乎完全正交
其他测试模型 多层对比 极低值 近似正交

"The cosine similarity will vary between zero and one, right? So one means like perfect aligned, and zero means like uh orthogonal to each other."

这一发现支持了嵌入矩阵正交性的假设,尽管嵌入本身的动态变化仍需进一步研究,但这为当前的理论推导奠定了坚实基础。

联合动态方程与理论简化

基于上述假设,Drama框架推导出了底层MLP权重与自注意力层权重的联合动态方程。这一突破性的方程允许研究者通过底层权重的幅度来直接计算自注意力的行为,从而将复杂的自注意力层训练动态简化为可解析的形式。这意味着,一旦获得该方程,研究者无需单独追踪自注意力层的复杂训练过程,只需将其代入动态方程即可预测整体行为。

这种简化极大地降低了分析难度,并且不再需要假设Decoder层的训练速度远快于自注意力层,这是相较于SCAN和SNAP框架的重大改进。演讲者强调:

"So this equation basically tell you the story of both attention. Plus the MLP, a lower layer, right? So then you can use that equation to analyze the dynamics of the of the of the transformers."

通过这种联合动态视角,Drama提供了一个更统一且简洁的分析工具,能够同时捕捉注意力机制和多层感知机(MLP)的演化故事。

与SCAN/SNAP框架的对比与一致性验证

为了回应学术界关于不同理论框架适用性的疑问,演讲者详细对比了Drama与SCAN、SNAP框架的异同。尽管设定不同,但Drama在特定条件下能够复现SCAN和SNAP的预测结果,证明了其理论的兼容性与稳健性。

特性/假设 SCAN / SNAP 框架 Drama 框架
激活函数 假设为线性 支持线性或非线性($\phi$)
归一化层 包含L2归一化等额外归一化 可能无归一化或不同设置
权重学习速度 假设 $W$ 学习远快于 $Z$ 无需此假设,通过联合动态处理
嵌入矩阵 U 固定且列正交 固定且列正交
输入表示 One-hot 向量 One-hot 向量(针对第一层)

在激活函数为线性的特定设置下,Drama的预测与SCAN/SNAP高度一致:注意力机制会变得更加稀疏,且具有最高“token判别力乘以相干性”的组件将最快饱和。这表明Drama不仅是一个独立的理论框架,更是一个能够涵盖并扩展现有理论的通用模型。

"The answer is yes. We can actually have the same very similar predictions there. So the prediction basically say that attention will become sparser."

这种一致性验证了Drama框架在更广泛场景下的有效性,同时也展示了其在处理非线性激活时的扩展潜力。

非线性激活下的注意力稀疏性动态变化

在多层Transformer的训练动力学中,非线性激活函数对注意力机制的影响与线性情况存在显著差异。理论推导表明,在线性激活假设下,注意力分布会持续变稀疏直至饱和;然而,在非线性激活环境下,注意力机制会呈现出一种先稀疏后密集的动态演变过程。具体而言,注意力首先会变得更加稀疏,这符合散点图论文(scatter plot paper)的预测;但随着训练的进行,在后续的激活层中,注意力又会重新变得密集。这种“先降后升”的曲线形态是理解非线性Transformer行为的关键特征,其背后的数学机制较为复杂,但实验结果验证了这一理论预测的一致性。

"Actually, tell you that the attention will first become sparser... but at the same time it will become denser later in the later activations."

真实数据集上的实验验证与现象复现

为了验证上述理论预测,研究团队在两个真实的文本数据集 WikiText 2 和 WikiText 103 上训练了多层Transformer模型,并观察到了与理论完全一致的曲线形态。实验数据显示,无论是在单层还是多层Transformer中,注意力熵(Attention Entropy)都表现出先下降后回升的趋势。特别是在顶层(top layers),这种先下降再回升的现象尤为明显。对于单层Transformer而言,这种看似反直觉的行为——即在初始阶段学习显著特征,随后又拾取非显著特征——需要通过更深层的机制来解释,因为单层模型本应一次性学习所有特征。

注意力熵变化的机制解释

注意力熵先降后升的原因在于模型学习特征的顺序和策略。在训练初期(如前1000个mini-batch),注意力熵急剧下降,这是因为模型优先拾取与查询(query)共现频率极高的token,导致注意力分布高度集中。随后,注意力熵开始回升,这是因为模型开始拾取那些虽然不显著(not salient)但在时间尺度上具有鲁棒性的特征和token。这种机制并非单层模型的必要特性,但在多层结构中,它有助于构建分层数据表示(hierarchical data representation)。模型通过先学习高频共现特征,再逐步引入低频但稳健的特征,从而在深层网络中隐式地构建出层级结构,避免了直接学习所有可能共现但非关键的特征组合。

层级生成模型中的潜在变量学习

Transformer之所以采用这种先稀疏后密集的学习策略,是因为数据生成过程往往涉及潜在变量(latent variables)。理论证明,在层级生成模型中,两个token在层级结构中越接近,它们的共现概率越高(接近1);反之,若两个token在层级结构中相距甚远,其共现概率则极低(接近0)。这一特性由公式 $H$(共同潜在祖先的高度)所刻画。Transformer通过这种学习机制,隐式地构建了数据的层级结构:首先学习紧密关联的特征对(如m和r,m'和r'),然后在顶层将这些局部结构组合起来。这种方法有效地避免了学习那些共现概率低且非关键的特征对(如m和r'),从而提高了学习效率并增强了模型的泛化能力。

浅层 Transformer 的注意力反弹现象

在分析 Transformer 如何学习数据分布时,我们发现模型具有极强的自适应性,能够在不了解数据内部结构的情况下,同时捕捉深层和浅层的潜在分布。这种机制解释了为什么 Transformer 能处理如此多样的任务。在训练初期,模型会优先学习高频出现的 token(如 'm' 和 'l'),随后逐渐学习到它们之间较弱的对应关系(如 'm' 和 'l prime')。这一过程导致注意力机制的曲线呈现出先下降后上升的“反弹”特性。

"This basically provides you with kind of insights about why transformer can do so much interesting things without actually knowing the underlying distributions. It will be auto adaptive to the structure of the data."

当我们使用单层 Transformer 在相同数据集上进行训练时,可以观察到非常清晰且复杂的结构变化。这种浅层模型中的注意力反弹现象,反映了模型在初期忽略弱相关性后,重新拾起并学习这些弱相关特征的过程。这种动态行为是理解 Transformer 学习机制的关键线索,表明模型并非一次性完成所有特征提取,而是分阶段、分层次地进行。

深层 Transformer 的层级分工与解释消除

随着 Transformer 层数的增加,学习机制发生了显著变化。在多层 Transformer 中,底层网络首先学习高频共现关系。一旦这些强相关性被上层网络学习并“解释消除”(explained away),底层网络就不再需要重复学习那些弱共现关系(如 A 和 NPR 之间的弱关联)。因此,在深层 Transformer 的底层,我们观察不到浅层模型中那种“反弹”现象,曲线在下降后保持平稳,不再回升。

"Once you learn this, you don't need to learn the weak co-occurrence between A and NPR again... Because this bouncing back property has already been explained away by all the previous mechanisms."

这种层级间的功能分工意味着,上层网络承担了处理复杂、弱相关特征的任务,而底层网络则专注于基础的高频模式。如果底层试图重新学习已被上层解释的特征,不仅效率低下,而且缺乏必要性。这种理论解释了为何深层模型的训练曲线在底层表现为单调下降,而在顶层则保留了某种反弹或波动特性,体现了模型内部的信息传递与特征抽象的高效性。

合成数据实验与因果性探讨

为了验证这一理论,研究者使用合成分层数据进行了实验。结果显示,Transformer 的隐藏层与数据的潜在变量之间存在极高的相关性,尤其是在前几层,相关性接近 100%。然而,目前的分析仅基于相关性,尚未确立因果关系。虽然强注意力可能暗示因果关系(即两个 token 频繁共现可能意味着真实的因果联系),但目前的参数化方法(如 YZ Transformer)在处理多层结构时表现不佳。

"You can see almost 100% correlations between the latent layers, between the hidden layers of transformers and the latent variables, especially in the first few layers."

YZ Transformer 的性能瓶颈主要源于其巨大的嵌入矩阵(词汇表大小乘以词汇表大小),这在实践中难以有效训练。相比之下,标准 Transformer 通过嵌入层和权重分离的方式,更有效地模拟了这种注意力机制。未来研究需要探索如何设计具有因果属性的参数化方法,例如强制注意力矩阵具备非对称性,以更好地捕捉因果逻辑,而不仅仅是统计相关性。

学习率与解冻策略的影响

实验还发现,不同的学习率设置会导致截然不同的训练曲线。那些展现出“反弹”特性的训练曲线,往往对应着最高的验证准确率,这可能暗示了学习数据的正确方式。此外,如果在一到三个 epoch 后解冻底层参数,底层网络可能会再次出现反弹,随后又因上层网络的覆盖而再次下降。这种动态变化进一步证实了层级间特征的竞争与协作关系。

"The training curve that show this bounce back properties will have the highest validation accuracy. Which means that it might that might mean like the correct way of learning data."

这一发现提示我们,训练动态中的反弹现象并非噪声,而是模型有效学习潜在结构的标志。通过调整学习率和解冻策略,可以优化这种动态过程,使模型更好地平衡底层基础特征与高层抽象特征的学习。尽管目前尚无其他理论能完全解释这一现象,但这一视角为理解 Transformer 的内部工作机制提供了新的思路。

注意力机制与嵌入矩阵的动力学分析

在探讨多层Transformer的内部可解释性时,理论框架往往需要将复杂的系统拆解为独立部分进行逐一分析。目前的理论尚未完全解释注意力机制(Attention)和嵌入矩阵(Embedding Matrix)的动态变化过程,这是未来研究的重要方向。为了简化分析,早期研究(如“snapshot”方法)通常假设嵌入矩阵是固定且正交的,但这与训练过程中嵌入向量实际被学习的现实不符。

"We basically assume. The embedding is fixed and it's orthogonal to each other. While this is true if you check the empirical results, but it's still not ideal because the embedding should also be learned during the training."

针对早期理论中关于极端情况(如token仅属于单一类别或属于多个类别)的假设限制,新的理论框架(如Drama)引入了连续谱系(Continuous Spectrum)的概念。通过量化token的独特性(Distinctness)与普遍性(Commonness),模型不再需要人为区分极端案例,而是自然覆盖所有token。当量化值幅度较大时,代表独特token;当值接近零时,代表常见token。这种连续谱系保留了一致性(Coherence)和判别性(Discriminancy),使得理论分析更加自然且无需特殊处理假设条件。

MLP层的模式叠加与潜在概念学习

在多层Transformer的训练动力学中,多层感知机(MLP)层展现出一种称为模式叠加(Pattern Superposition)的特性。实证研究已观察到,单个MLP向量可以编码多种不同的模式,这一现象在Anisotropic等论文中被重点强调。Drama理论进一步从训练动力学的角度证明,从头开始训练(Train from scratch)时,这种模式叠加是自然发生的。

然而,对于潜在概念(Latent Concept)的学习过程,目前仍缺乏具体的定理来描述潜在向量(Latent Vector)随时间变化的动态规律。这是一个非常棘手且难以处理的难题。未来的工作不仅包括分析嵌入矩阵的动态变化,还将深入探讨位置编码(Positional Encoding)的影响,以及揭示多层设置下潜在概念学习的具体机制。这些理论空白构成了当前研究的核心挑战,也是后续论文试图解决的关键问题。

基于注意力稀疏性的推理加速实证

为了提升大模型的推理速度,研究团队提出利用注意力稀疏性(Attention Sparsity)和MLP激活稀疏性进行加速。在典型的自回归生成中,许多注意力头和MLP层的前向路径中并未被激活,导致交叉注意力权重(Cross-attention weights)被无效加载。为了解决这一问题,团队提出了DJW方法,旨在预测哪些头或层在推理过程中不会被激活。

例如,某些头(如第42、43头)可能输出均匀分布,表明其未激活;而其他头(如第42、44头)则显示出清晰的注意力分布模式。通过训练一个轻量级的预测器(Predictor),可以预先判断哪些头和MLP中间层需要激活。由于GPU中内存加载到缓存的速度极慢(数百个周期,而计算仅需几个周期),减少加载的权重数量能显著提升前向传播速度。

对比基准 延迟表现 (Latency) 备注
Hugging Face 较高 标准实现,未优化稀疏性
Fast Transformers 较高 标准稀疏实现
DJW (本研究) 较低 基于预测的上下文稀疏性加速

该方法与量化技术(Quantization)兼容,特别适用于小批量(Lower batch size)场景。此外,团队还开发了H2O项目,利用稀疏注意力机制加速推理。由于注意力计算复杂度为O(n²),每个新token需遍历所有先前token,利用稀疏性可大幅减少计算量,从而提升整体效率。

稀疏注意力与推理效率的权衡

在H2O项目中,核心挑战在于如何有效利用稀疏注意力来应对O(n²)的计算复杂度。在自回归生成中,随着序列长度n的增加,计算所有token间注意力的成本急剧上升。通过识别并仅保留高激活的token,可以显著减少需要计算的注意力对数量。

"We want to only keep the high activation tokens in attention MLP, and use that to accelerate the inference of the larger models."

实证结果显示,许多注意力头在前向传播中并未被激活,且MLP层的激活也呈现稀疏性。DJW方法通过预测这些稀疏模式,仅加载必要的权重行和列,避免了无效内存访问。这种上下文稀疏性(Contextual Sparsity)的预测机制,使得推理速度在保持模型性能的同时得到显著提升。特别是在低批量大小场景下,这种优化效果更为明显,为大规模Transformer模型的高效部署提供了新的技术路径。

优化维度 传统方法 DJW/H2O 方法 性能提升点
注意力计算 全连接 O(n²) 稀疏注意力 减少计算量
权重加载 全量加载 预测性加载 减少内存带宽压力
适用场景 通用 小批量/低延迟 优化特定负载

通过结合理论分析与实证优化,研究团队展示了从训练动力学理解到推理加速应用的完整闭环。尽管嵌入矩阵和位置编码的动态分析仍是未解之谜,但基于当前理论框架的稀疏性优化已展现出显著的工程价值。

注意力分布与“重击者”机制

在Transformer模型中,并非所有历史token都需要被同等对待。通过分析注意力分数分布可以发现,其形态与词共现分布高度相似,这为两者间的关联性提供了实证支持。进一步的研究揭示,那些获得高注意力分数的token对下游任务至关重要。如果移除这些高注意力token,模型性能会出现显著下降;反之,仅保留这些关键token进行推理,则能维持良好的表现。这些高频且高注意力的token被称为“重击者”(Heavy Hitters)。

基于这一发现,研究者提出了一种名为H2O(Heavy Hitter Oracle)的注意力算法。这是一种局部贪心算法,其核心逻辑是在每个解码步骤中累加先前token的注意力分数,并剔除累积注意力分数最低的token,仅保留一个固定预算内的KV缓存token。实验数据显示,即使仅保留20%的KV缓存token,模型在多个基准测试中仍能保持近乎相同的性能。相比之下,若仅保留局部token,性能会大幅跌落;而保留高注意力token则能在剔除超过80%的KV缓存后依然保持优异表现。这种方法不仅节省了内存,还显著提升了推理速度,甚至因移除限制创造力的token而提升了模型的创造性。

“如果你移除那些拥有大量注意力的token,下游任务的性能就会大幅下降。”

注意力Sink现象与无限上下文

H2O算法引出了一个更深层的发现:许多“重击者”实际上是Transformer中的前几个token。这些初始token占据了超过50%的注意力分数。研究者假设,由于Softmax规则要求注意力分数之和为1,模型被迫将多个强注意力分数过载到初始的几个token上,以平衡其他位置的注意力。有趣的是,这些初始token的具体内容并不重要,重要的是它们的位置。实验表明,可以用随机token替换前几个token,模型的困惑度(Perplexity)变化不大;但如果直接移除前几个token,模型性能会崩溃,困惑度飙升至约5000,输出变为无意义的乱码。

这一发现催生了Streaming LM(流式语言模型)的研究,并发表于ICLR。Streaming LM的核心思想是维护一个固定长度的KV缓存预算。当生成的token超过预算时,丢弃中间的token,但始终保留初始token和最近生成的token,并保留它们在缓存中的注意力位置。这种机制使得模型能够进行无限长度的对话,而无需受限于上下文窗口的大小。在部署层面,这种方法避免了因长序列KV缓存导致的内存溢出,同时也防止了因截断初始token而导致的模型失效。

“前几个token很重要,但重要的是它们的位置,而不是内容。”

性能基准与可扩展性实验

为了验证Streaming LM的有效性,研究者在多种场景下进行了广泛测试。实验结果显示,Streaming LM能够保持稳定的困惑度,意味着输出始终具有意义,同时内存占用随时间保持恒定,不受上下文窗口长度的影响。研究者进一步将模型上下文扩展至400万token,甚至更多,且困惑度保持一致,证明了该方法的可扩展性。

此外,研究者还对比了其他基线方法,并探索了将前几个token作为可学习参数的可能性。实验引入了专门的Sync Token(同步token),分别测试了零嵌入(zero sync)和可学习嵌入(learnable sync)的效果。结果显示,引入Sync Token后,模型的困惑度进一步降低,性能略优于固定初始token的方案。这些变体工作已被部署在硬件设备和库中,显示出巨大的应用潜力。

以下是Streaming LM与其他基线方法在关键指标上的对比数据(基于文中描述的趋势整理):

方法/配置 KV缓存保留策略 上下文长度 困惑度表现 内存占用 备注
标准Transformer 全部历史Token 有限 基准 随长度线性增长 长序列易OOM
仅保留局部Token 最近N个Token 有限 大幅下降 低 性能损失严重
H2O (保留20% KV) 高注意力Token 有限 接近基准 低 性能优异,速度快
Streaming LM 初始+最近Token 无限 稳定 恒定 支持无限对话
Streaming LM + Learnable Sync 初始Sync+最近Token 无限 略优于标准Streaming 恒定 性能微调提升

“我们可以将模型扩展到400万token,困惑度保持不变,这是非常好的结果。”

注意力机制中的“寄存器”与长上下文局限

在探讨长上下文窗口(Long Context Window)的解决方案时,Yuandong Tian 指出,虽然某些技术能显著降低显存占用,使模型能够进行无长度限制的对话,但这并不意味着解决了根本性的信息丢失问题。Meta 的并发研究也发现了类似的结构,即在 Vision Transformers 中存在被称为“Registers”(寄存器)的机制,它们对应于注意力机制中的“Attention Sinks”(注意力汇)。这些机制的核心逻辑是通过丢弃部分 Token 来维持概率分布的稳定性,从而节省计算资源。

然而,这种策略存在明显的局限性。> “It basically keep the keep the probability. It means that it will give you reasonable answers. But it may not necessarily retrieve the information that is being discarded.”

这意味着,虽然模型能给出看似合理的回答,但如果关键信息(如密码检索中的特定 Token)被丢弃,模型将无法提供正确答案。因此,仅靠丢弃 Token 无法真正解决长上下文中的信息检索难题,未来需要更深入的优化。

流式注意力与 HTO 的结合

为了克服上述局限,研究者提出了将流式注意力机制(Streaming RMs)与 HTO(Hierarchical Transformer Optimization 或类似层级优化方法)相结合的策略。Streaming RMs 是一种特殊的 HTO 方法,其特点是 Transformer 仅保留前几个 Token,而 HTO 则根据注意力分数动态选择 Token。将两者结合可以带来更优越的性能,同时保持无限对话(Infinite Chat)的特性。

这种组合方法在最近的更新中被证明能有效提升推理速度。通过动态筛选高注意力分数的 Token,而非简单地截断或保留前缀,模型能够在保持长上下文能力的同时,更精准地捕捉关键信息,从而在性能上优于单一的流式或层级优化方案。

通过插值扩展上下文窗口的预训练微调

在扩展上下文窗口方面,Yuandong Tian 分享了一篇因竞争过于激烈而未正式提交的论文。该研究旨在通过预训练阶段的参数调整来高效扩展上下文窗口,而非传统的微调。以 Llama 模型为例,其原生上下文窗口为 2K,若强行使用超过 2K 的输入,模型输出将变为乱码。传统方法是通过长文本微调来扩展,但这需要大量资源和批次。

研究团队发现,只需修改位置编码(Positional Embeddings)的缩放比例,即可大幅提高效率。具体做法是:当将窗口从 2K 扩展到 4K 时,将每个位置编码索引除以 2,使其映射回原来的 2K 范围。这种方法本质上是利用插值(Interpolation)而非外推(Extrapolation)。因为注意力函数在插值时表现良好,而在外推时会产生极差的分数。

插值微调的效率与性能数据对比

通过这种位置编码插值的方法,模型仅需极少的微调步骤即可显著扩展上下文窗口。研究表明,使用该方法可以在少于 1,000 步的微调中将上下文窗口扩展至原始长度的 8 倍。相比之下,不进行参数修改的直接微调需要更多步骤且效果不佳。

以下是不同方法在扩展上下文窗口时的效率对比数据:

扩展目标 (从 2K 开始) 插值微调所需步骤 直接微调所需步骤 性能保持情况
8K (4倍) 200 步 > 1000 步 基准测试性能基本持平
16K (8倍) < 1000 步 效果不佳/未完全扩展 基准测试性能基本持平
32K (16倍) < 1000 步 效果不佳/未完全扩展 基准测试性能基本持平

该研究指出,位置编码的计算依赖于 RoPE(Rotary Positional Embeddings),其公式涉及位置索引的差值。通过插值,模型在训练窗口内(如 2K)的行为得到保留,而在扩展窗口(如 4K-8K)内,模型通过插值而非外推来生成注意力分数,从而避免了外推带来的性能崩溃。尽管该论文因发布后两天内 Reddit 上出现更优方案而未正式提交,但其揭示的插值优于外推的原理为长上下文扩展提供了重要思路。

经验规律与第一性原理的理论分野

当前大模型领域的理论研究呈现出两条截然不同的路径:一条是基于经验现象总结的实证路线,另一条则是基于第一性原理的推导路线。以 Siri 为代表的深层堆叠模型,其发展很大程度上依赖于对缩放定律(Scaling Laws)和涌现行为(Emergent Behavior)等经验现象的描述。这些经验规律能够直观地预测模型规模与数据量变化对性能的影响,例如通过曲线预测将模型和数据同时扩大十倍后的性能表现,这在工程应用中具有极高的实用价值。

"Scaling laws and also emergent behavior... The curve will tell you predict what's going on."

相比之下,基于第一性原理的理论研究则试图从公式出发,推导模型在特定设定下必然产生的后果。这种路径虽然难以直接应用于实际场景,且过程极为艰难,但如果能够成功打通从基础理论到应用的全流程,将能实现对模型行为的全面掌控。这种从底层逻辑出发的探索,虽然起步艰难,但被认为具有更深层的解释力和未来的应用潜力。

从第一性原理推导缩放定律的可能性

将第一性原理理论转化为实际可用的缩放定律,被视为理论研究的下一个重要步骤。这一过程类似于物理学中的研究范式:从基本原理出发建立方程,然后模拟在拥有海量数据(如 $10^{32}$ 个原子)时的系统行为,进而推导出全局性的经验定律,最后将这些理论预测与真实实验数据进行拟合。

在当前的语境下,这意味着通过理论推演来模拟增加数据量后的模型表现,从而得出理论上应有的边界和缩放规律。这种由理论推导出的缩放定律,将比目前纯经验观察到的规律更进一步,因为它提供了更深层的因果解释。尽管这一过程充满挑战,但它为理解大模型在极端规模下的行为提供了理论框架,使得研究者能够从数学上预测而非仅仅观察模型的表现。

理论构建的迭代:从错误定理到收敛

理论构建的过程往往充满痛苦与反复。回顾早期研究,研究者可能在某个问题上耗费数月,最终证明的结论却毫无实际意义。然而,经过长期的探索,研究者逐渐掌握了构建理论的节奏:先形成一个粗略的理论图景,即使初步提出的定理是错误的,只要其方向正确且近似合理,便可以通过添加约束条件和进行经验验证来逐步修正。

"Wrong theorems, but they are approximately right, in the right direction. Then by adding constraint... you might have like a better idea about what this should look like."

以 Dream 论文为例,核心思想(如绘制嵌入)可能在几天内即可提出,一旦推导出关键方程,后续工作便有了明确方向。这种理论构建并非孤立进行,而是需要在符号结构与现实世界结构之间不断往返。研究者从现实假设出发,进入符号推导,再回到实践验证,直到找到既在符号上有趣又在现实中相关的假设。这种迭代过程使得理论既能保持逻辑正确,又具备实证效用。

理论假设的简化与收敛策略

在具体理论推导中,合理的假设简化是推进研究的关键。例如,在涉及时间 $t$ 的收敛性分析中,研究者往往先假设 $t$ 趋于无穷大,以捕捉核心图景,而暂时忽略复杂的数学细节。这种简化并非随意,而是基于对问题本质的深刻理解。

然而,某些假设是绝对不可妥协的。例如,研究者明确表示绝不会假设特征固定,因为这在理论上是错误的。理论构建需要在简化与严谨之间找到平衡:通过初步的简化假设(如 $t \to \infty$)快速建立直观理解,同时坚守核心逻辑的正确性(如特征动态变化)。这种策略使得研究者能够在短时间内推导出定理并寻找良好解,从而兼顾理论深度与实验效率,形成一种兼具艺术性与科学性的研究范式。

数学简化与物理直觉的平衡

在构建模型解释性理论时,数学处理的便利性往往与保留物理直觉之间存在张力。过度简化数学模型可能会抹杀对系统底层物理机制的理解,这是一种需要避免的倾向。因此,研究者必须在假设的有效性与模型的完整性之间寻找平衡点。如果假设过于理想化,虽然能让数学推导变得容易,但可能会破坏原本想要解释的核心现象。正如演讲者所言:

"this will make the mathematics much easier, this basically kill the entire physics picture. That's not a good thing to do. So there's always a balance."

这意味着,任何理论构建都必须确保其基本假设不会导致目标解释对象的失真。这是一种长期的探索过程,需要反复验证假设是否能在简化计算的同时,依然保留对复杂系统行为的准确描述。

层级表示的历史与CNN的挑战

演讲者回顾了自己多年来对层级表示(Hierarchical Representation)的研究历程。早在硕士阶段,他就开始研究视觉模型中的层级表示,尽管当时的论文鲜有人问津,但这奠定了他后续研究的基础。对于卷积神经网络(CNN),层级结构是其核心特征,也是他最初从视觉领域转向该研究的动机。

然而,在CNN中进行层级分析面临巨大挑战。早期的尝试往往因为符号系统过于复杂而被期刊拒绝。直到后来,通过改进数学表述,才成功发表了第一篇关于CNN层级分析的论文。这一经历表明,复杂的符号系统是阻碍理论深入的主要障碍。对于CNN而言,其层级结构是显而易见的,但如何将其形式化并纳入统一的解释框架,需要极大的耐心和创新。演讲者提到:

"All the previous hierarchical papers that we try to do the analysis, they all get rejected because the notation is too complicated."

这反映了理论研究中形式化表达的重要性,简洁而有力的符号系统是推动领域发展的关键。

Transformer嵌入向量的高维优势

Transformer与CNN在表示学习上的根本差异在于嵌入向量(Embeddings)的使用。Transformer利用高维空间中的嵌入向量,展现出强大的表示能力。在高维空间中,存在一个有趣的性质:可以在有限的维度 $d$ 中容纳 $2^d$ 个向量,且这些向量几乎彼此正交。这种指数级的内容存储能力使得Transformer能够在有限的样本空间中编码海量的对象或概念。

相比之下,CNN倾向于使用类似独热编码(One-hot representation)的方式,即每个概念对应一个特定的token,且要求这些token之间相互正交。为了保持正交性,独热编码需要极高的维度,这在效率上远不如Transformer的密集嵌入。这种高维正交性使得Transformer能够高效地处理复杂的信息组合,为后续的推理任务提供了坚实的基础。

组合式表示与推理机制

Transformer的嵌入向量不仅用于存储信息,还支持复杂的推理任务。演讲者推测,这种推理能力与嵌入向量的组合式表示(Compositional Representation)密切相关。例如,一个概念 $A$ 可能由多个嵌入向量相加而成,分别代表其不同的属性,如字母本身、首字母、大小写状态等。这种表示方式将不同属性解耦(Decoupled),使得Transformer可以在计算过程中独立处理这些属性。

这种解耦特性使得模型能够灵活地组合和变换信息,从而执行复杂的逻辑推理。未来的研究重点可能在于深入分析这种属性独立计算的机制,揭示Transformer如何通过嵌入向量的线性组合来实现高级认知功能。演讲者指出:

"This kind of representation is very useful because it actually is decoupled into different properties, and this property may be computed individually in the transformer computation."

这种从表示到推理的连贯性,是理解Transformer内部工作机制的关键线索。