核心机制与重参数化技巧
Transformer 的核心魅力在于其注意力层,它通过查询(Query)和键(Key)的内积计算,利用 Softmax 值对前一层嵌入进行加权,从而预测下一个 Token。这种机制看似简单却蕴含“魔法”,但直接分析其梯度极其复杂。为了深入理解这一机制,研究引入了 Scan and Snap 论文(ECCV 2023)及其后续工作 Drama(ICLR),旨在从单层设置出发,逐步构建多层 Transformer 的数学理解框架。在单层模型中,输入为上下文 Token 序列 $x_1$ 至 $x_{t-1}$,最后一个 Token 作为 Query 进行自注意力计算,生成概率分布后经过归一化和编码层,最终通过 Softmax 预测 $x_{t+1}$。
"The majority of this lecture is going to contribute and to talk about like how we should understand this mechanism."
面对包含 $W_Q, W_K, U$ 等大量参数的复杂公式,研究者采用了两个关键技巧来简化分析。首先是 重参数化(Reparameterization),将众多变量转化为独立变量 $y$ 和 $z$。其中,$y$ 代表顶层解码器权重,$z$ 代表自注意力矩阵的成对 Logits。这一转换使得动态分析变得可行,尽管初始方程看似复杂,但结合后续引入的 Drama 框架后,整体结构将显著简化,从而揭示 Transformer 的本质属性。
关键假设与数据分布设定
为了推导注意力层的动态行为,研究建立了一系列严格的数学假设。首先,模型 不包含位置编码(Positional Encoding),因为处理位置编码会增加难度,且已有研究表明即使没有位置编码,模型仍能获得合理性能。其次,假设 序列长度趋于无穷大,以简化长序列下的统计特性分析。第三,假设解码器 $Y$ 的学习率大于自注意力层 $Z$ 的学习率,这一时间尺度分离是分析动态过程的关键。此外,还包含若干技术性假设以支撑理论推导。
在数据分布方面,假设存在 $k$ 个序列类别,每个类别包含两类 Token:Distinct Tokens(独特 Token) 和 Common Tokens(常见 Token)。Distinct Tokens 仅在一个特定类别中出现,其条件概率在该类别外为零;而 Common Tokens 则出现在多个序列类别中。这种分布设定允许研究者区分不同 Token 在注意力机制下的不同演化路径,为后续分析注意力分数的动态变化奠定基础。
注意力动态演化结论
基于上述假设,研究揭示了自注意力层的核心动态行为,关键在于分析 CL 标题(CL title),即数据集共现概率与注意力矩阵成对指数项的乘积。对于 Distinct Tokens,由于其共现概率仅在单一类别非零,其对应的 CL 标题呈现出截然不同的模式。而对于 Common Tokens,动态分析显示其梯度随时间变化率为负($\frac{dz}{dt} < 0$),这意味着 Common Tokens 的注意力分数会随时间逐渐减小。这一结论表明,模型倾向于抑制在多个类别中通用的 Token 的注意力权重,以减少混淆。
"For common tokens... the gradient will basically drag all the attention score of the common tokens error to be smaller and smaller."
相反,对于 Distinct Tokens,其对应的注意力 Logits $z$ 会随时间 持续增长。这种分化机制使得模型能够逐渐聚焦于具有类别特异性的 Token,从而提升预测准确性。这种动态演化过程展示了 Transformer 如何通过自注意力机制自动区分通用信息与特异性信息,无需显式的位置编码或复杂的预处理,仅凭学习动态即可实现有效的模式识别。这一发现为理解多层 Transformer 中的信息流动和特征提取提供了重要的理论依据。
注意力机制作为可学习的 TF-IDF 模型
在 Transformer 架构中,注意力分数的增长趋势揭示了其内在的统计学习机制。随着训练进行,注意力分数不仅会变得更紧密,还会呈现出类似传统自然语言处理中 TF-IDF(词频-逆文档频率)的特征。注意力机制本质上是一种可学习的 TF-IDF 模型,它旨在平衡两个关键因素:一是识别对特定类别具有高频率的重要词汇,二是抑制那些在多个文档中普遍出现但信息量低的通用词汇。这种机制类似于文档分类任务中的特征选择过程,通过过滤掉跨文档过于常见的非信息性令牌,模型能够更精准地捕捉关键语义。
"So from these two, you can already see the attention mechanism is kind of like the learnable TF-IDF, right?"
具体而言,如果某些令牌(如体育新闻中的 "Messi")与特定类别高度相关,模型会赋予其更高的注意力权重。反之,通用词汇会被抑制。这种基于条件概率的令牌筛选机制,使得模型能够在复杂的文本序列中,动态地调整对不同信息的关注程度,从而实现更高效的特征提取和分类决策。这种类比不仅解释了注意力的功能,也为理解其训练动态提供了直观的统计视角。
上下文稀疏性与注意力分数的增长定理
训练过程中,注意力分数的增长速率并非均匀分布,而是表现出显著的上下文稀疏性(Contextual Sparsity)。通过量化自注意力矩阵中元素 $Z_{ML}$ 随时间 $t$ 的变化,研究发现,与查询令牌具有高共现概率的令牌,其注意力分数增长更快。这种增长差异可以通过一个定理来描述,该定理量化了不同令牌之间增长率的比例关系。随着训练的进行,注意力模式会逐渐变得稀疏,呈现出"富者愈富,穷者愈穷"的马太效应。
| 令牌类型 | 与查询的共现概率 | 注意力分数增长速率 | 最终稀疏化表现 |
|---|---|---|---|
| 高共现令牌 | 高 | 快 | 显著突出,占据主导 |
| 低共现令牌 | 低 | 慢 | 被抑制,趋于零 |
这种稀疏性是依赖于查询的(Query-dependent)。对于不同的查询令牌,模型会展现出截然不同的注意力曲线模式。例如,查询令牌 1 可能与令牌 A 有高共现,而查询令牌 2 可能与令牌 B 有高共现。这种动态的、查询依赖的稀疏化过程,使得模型能够根据当前语境灵活地聚焦于最相关的信息,而非静态地分配注意力权重。这种机制确保了模型在处理长序列时,能够有效地过滤噪声,保留关键信号。
两阶段动态:扫描与冻结(Scan and Snap)
尽管直觉上认为注意力会无限增长直至"赢家通吃",但实际训练动态呈现出更复杂的两阶段效应。第一阶段被称为注意力扫描(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)的稳定结构,而非绝对的赢家通吃。这种动态在 WikiText2 等数据集的实验中得到验证,显示出明显的稀疏模式。这种两阶段模型解释了为什么 Transformer 在训练后期能够形成稳定且高效的注意力分布,同时也为理解模型收敛行为提供了理论依据。通过这种扫描-冻结机制,模型在保持灵活性的同时,确保了最终表示的稳定性。
理论分析与实验验证策略
为了证明上述动态,研究采用了一套简化的分析策略。首先,利用无限序列长度的假设,将复杂项转化为包含概率分布的统计项。其次,假设解码器(Decoder)的学习率远快于自注意力矩阵(Self-attention),从而在分析解码器动态时将注意力矩阵视为常数。这种分层分析策略简化了多层 Transformer 的复杂性,使得推导梯度反向传播至注意力层的机制成为可能。
"We want to use this power of infinite sequence length to basically turn all the terms into statistic terms that include the probability distribution."
在实验层面,通过在 WikiText2 上应用特定的学习率结构,观察到了注意力模式的稀疏化变化。然而,初步分析存在简化假设,如忽略嵌入向量和残差连接,且假设解码器学习率更快,这在实践中并不总是成立。为克服这些限制,后续研究引入了Drama 论文,通过结合 MLP 层和自注意力层的动态,构建了修改后的 MLP 层,以更准确地捕捉多层 Transformer 的训练动态。这种从简化模型到复杂模型的演进,体现了研究团队对 Transformer 内部机制深入探索的科学路径。
DRAMA 框架引入残差连接与非线性机制
在之前的分析中,Transformer 模型往往忽略了残差连接(Residual Connections)的作用,而 DRAMA(Dynamics of Residual Multi-layer Attention)框架的核心突破在于将这些连接重新纳入考量。演讲者指出,通过在 Decoder 层引入非线性激活函数(记为 $\phi$),能够更准确地描述模型内部的动态变化。这与之前仅考虑线性映射至 Softmax 的简化模型形成了鲜明对比,从而揭示了更丰富的网络行为特性。
"So for this one, we basically like put like a drag down like a drag one layer of the transformers... we can add them back and show how it actually play the role here, and we also can use these nonlinearities."
这种改进使得分析不再依赖于“Decoder 层训练速度远快于自注意力层”这一强假设。相反,DRAMA 通过建立下层 MLP 与自注意力层之间的联合动态方程,实现了更统一的理论描述。该方程不仅涵盖了注意力机制,还整合了下层 MLP 的动态,使得研究者无需单独处理注意力层的训练过程,只需将相关方程代入动态系统中即可推导出整体行为。
嵌入矩阵正交性的实证检验
为了验证理论假设,研究团队对嵌入矩阵(Embedding Matrix $U$)的性质进行了实证分析。理论假设指出,嵌入向量之间保持正交性且不随时间变化。尽管这一假设曾受到审稿人的质疑,但通过绘制不同层嵌入向量之间的余弦相似度随训练批次变化的曲线,研究提供了有力的数据支持。实验选取了 PyTear 及其他模型进行测试,结果显示嵌入向量在训练过程中确实保持了高度的正交性。
| 指标 | 数值/范围 | 说明 |
|---|---|---|
| 余弦相似度范围 | 0 - 1 | 1 表示完美对齐,0 表示正交 |
| 平均余弦相似度 | ~0.03 | 训练过程中基本保持稳定,未显著增加 |
| 曲线对应关系 | 不同颜色代表不同层 | 各层间的嵌入向量保持近似正交 |
从实验数据来看,随着训练的进行,不同层嵌入向量间的余弦相似度并未显著上升,而是稳定在 0.03 左右。这意味着即使在不同的网络层中,嵌入向量也几乎保持正交状态。这一发现证实了正交性假设在实践中的近似真实性,尽管嵌入动态本身的详细分析可能成为未来研究的独立课题。
DRAMA 与 SCAN/SNAP 的理论对比与统一
为了回应学术界关于不同理论框架适用性的疑问,演讲者对比了 DRAMA 与 SCAN/SNAP 两种框架的假设与设置。两者共同的基础假设包括:下层权重 $W$ 和自注意力权重 $Z$ 均为可训练参数;嵌入矩阵 $U$ 固定且列正交;输入 Token 在第一层为 One-hot 编码。然而,两者的关键差异在于对激活函数和归一化的处理。
| 特性 | SCAN / SNAP 框架 | DRAMA 框架 |
|---|---|---|
| 激活函数假设 | 线性激活 | 线性或非线性(Identity 或 $\phi$) |
| 归一化层 | 包含 L2 归一化 | 无归一化或不同设置 |
| 训练速度假设 | 下层权重 $W$ 学习快于 $Z$ | 无需此假设,依赖联合动态 |
| 注意力稀疏性 | 预测注意力变稀疏 | 预测注意力变稀疏 |
DRAMA 框架通过引入非线性激活,扩展了理论的适用范围。验证结果显示,当将 DRAMA 应用于线性激活场景时,其预测结果与 SCAN/SNAP 高度一致。具体而言,两种框架均预测注意力机制将变得更加稀疏。此外,具有最高“Token 判别力乘以相干性”(token discriminacy times coherence)的组件,其权重曲线增长最快,并率先达到饱和,而其他组件则增长较慢。这一结论不仅验证了 DRAMA 的理论一致性,也展示了其在非线性环境下的扩展能力。
"The prediction basically say that attention will become sparser... consistent with scanner step... if the term is highest... you're going to see it's the curve will go the fastest."
通过这种统一视角,DRAMA 不仅简化了对 Transformer 训练动态的分析,还为理解非线性激活下的模型行为提供了新的理论工具。
非线性激活下的注意力动态演变
在探讨Transformer内部可解释性时,线性激活与非线性激活展现出截然不同的注意力演化轨迹。理论推导指出,非线性激活会导致注意力机制先变得稀疏,随后在后期激活中重新变得密集。这一现象与线性激活形成鲜明对比:在线性激活中,注意力仅会持续变稀疏直至饱和,而不会在后期回升。这种差异源于非线性函数对特征选择的动态调节作用,使得模型在训练初期聚焦于高显著性特征,而在后期扩展至更广泛的潜在关联。
"Actually, tell you that the attention will first become sparser. Which is what the scatter plot paper will predict, but at the same time it will become denser later in the later activations."
为了验证这一理论预测,研究团队在真实数据集上训练了多层Transformer模型。实验选取了WikiText 2和WikiText 103两个真实世界文本数据集,旨在观察模型是否复现理论所预测的曲线形态。实验结果证实,在包含非线性激活的模型中,注意力熵确实呈现出先下降后上升的趋势,这与理论推导高度一致。特别是在单层Transformer中,这种"先降后升"的模式尤为明显,表明模型并非一次性学习所有特征,而是经历了一个动态的筛选与扩展过程。
单层与多层模型中的注意力熵变化
在单层Transformer的训练过程中,注意力熵的变化揭示了特征学习的阶段性。起初,由于共现频率高的token被优先选中,注意力分布迅速集中,导致熵值急剧下降至较低水平。然而,随着训练的进行,模型开始捕捉那些虽不显著但具有时间鲁棒性的特征,导致注意力分布重新分散,熵值随之回升。这一过程并非在训练的第一步发生,而是需要经历数千个mini-batch的迭代,表明特征学习是一个渐进的、分阶段的过程。
"The reason why it go up is that it also pick up features and tokens that is not salient. but still reasonable resilient over time."
在多层Transformer中,这种"先稀疏后密集"的模式在顶层层中同样显著。尽管不同学习率可能导致数值尺度的变化,但曲线的整体形态保持稳定。这一发现挑战了直觉:对于单层模型而言,一次性学习所有特征似乎更高效,无需经历这种复杂的动态变化。然而,实验数据表明,这种动态特性在多层架构中可能具有更深层的功能意义,特别是在处理复杂层级结构的数据时。
层级表征学习与潜在变量假设
为了解释多层Transformer中这种看似非直觉的行为,研究提出了层级数据表征(Hierarchical Data Representation)的假设。在深层神经网络中,模型仅能观察到顶层标签和底层token,中间过程由潜在变量(Latent Variables)控制。Transformer通过自监督学习,自动推断这些潜在的结构化表示。理论证明,在层级生成模型中,层级距离相近的token共现概率极高(接近1),而层级距离远的token共现概率极低(接近0)。这一性质由共同潜在祖先的高度(Height)决定,与Transformer架构本身无关,但决定了模型的学习策略。
"If you have two tokens that is very close in hierarchy. Then their probability that they co-occur will be very high, almost equal to one."
基于此,Transformer的学习机制被解释为:首先学习高共现的局部关联(如m与r,m'与r'),随后在顶层组合这些局部结构。这种策略避免了直接学习低共现但间接相关的特征(如m与r'),而是通过层级结构间接捕获。因此,注意力先稀疏后密集的现象,实质上是模型从局部关联构建到全局层级表征的必经之路。这种机制不仅适用于浅层潜在分布,也适用于深层层级分布,解释了为何多层Transformer需要这种动态的注意力演化来实现高效的层级特征学习。
浅层Transformer的自适应学习机制
Transformer模型在训练过程中展现出一种无需预先知晓数据内部结构的自适应能力。起初,模型会优先学习高频出现的词元(tokens),例如字母 'm' 和 'l'。随着训练的深入,即使这些词元之间的对应关系非常微弱,模型也能逐渐建立起 'm' 与 'l prime' 之间的关联。这种机制表明,注意力机制能够同时捕捉深层和浅层的潜在分布,而无需显式地理解数据的底层结构。
"The transformer will first learn the tokens that occur a lot... But eventually, it will learn a correspondence between m and l prime... the same mechanism... can learn both the deep latent distribution and also shallow latent distribution at the same time without knowing the internal structure of the data."
这种自适应特性解释了为什么Transformer能够执行如此多样的任务。在单层的Transformer模型中,我们观察到非常明显的“先下降后反弹”(bounce back)的学习曲线行为。这种反弹现象反映了模型在初步学习强关联后,转而探索较弱关联的过程。由于单层模型缺乏其他层级的辅助,它必须独自完成从强特征到弱特征的完整学习过程,因此这种波动在浅层结构中尤为显著。
多层Transformer的层级分工与解释消除
当Transformer的层数增加时,学习动力学发生了显著变化。在深层Transformer中,底层网络的学习曲线在下降后不再反弹。这是因为多层结构引入了层级分工:底层首先学习高频共现的强关联特征。一旦这些强关联被底层捕获,上层网络便无需再重复学习这些已被“解释掉”(explained away)的弱共现关系。
"If you have multiple layer of transformers... you first learn these co-occurrences that co-occur a lot. And once you learn this, you don't need to learn the weak co-occurrence... because this bouncing back property has already been explained away by all the by the previous mechanisms."
因此,在深层模型中,底层网络一旦完成了对强特征的提取,就不会再花费精力去拟合那些已经被上层处理过的弱特征,导致其损失曲线呈现单调下降而非反弹。相反,顶层网络仍然保留了这种反弹特性,因为它们需要处理更复杂的、尚未被底层完全解释的语义信息。这种理论解释了为何深层模型的底层表现更为平滑,而顶层则保留了探索弱关联的动力学特征。
合成数据实验与相关性验证
为了验证这一理论,研究者通过生成合成分层数据进行了实验,以观察Transformer不同隐藏层与潜在变量之间的对应关系。实验结果显示,Transformer的隐藏层与数据的潜在变量之间存在极高的相关性,尤其是在前几层中,相关性接近100%。这证实了模型确实在不同层级中学习了不同的潜在变量。
| 实验配置 | 观察到的现象 | 理论解释 |
|---|---|---|
| 单层Transformer训练 | 底层损失曲线出现明显“先降后升”反弹 | 模型需独自学习从强关联到弱关联的所有特征 |
| 多层Transformer训练(底层) | 损失曲线下降后不再反弹 | 强关联已被底层捕获,弱关联被上层解释,底层无需重复学习 |
| 多层Transformer训练(顶层) | 保留“反弹”特性 | 顶层需处理未被底层完全解释的复杂/弱关联特征 |
| 合成分层数据实验 | 隐藏层与潜在变量相关性接近100% | 验证了不同层级确实对应学习不同的潜在变量 |
需要注意的是,目前的分析主要基于相关性而非因果性。虽然强共现可能暗示因果关系,但实验并未直接证明因果作用。如果未来引入因果性约束,可能需要设计特定的属性来强制注意力矩阵具备因果方向性,例如确保如果第一个事件被阻断,第二个事件也不会发生,反之则不一定成立。
参数化方案与学习率的影响
在参数化方案方面,研究者尝试了YZ参数化的Transformer(即YZ Transformer)。实验表明,虽然三层深度的YZ Transformer可以训练并工作,但在多层结构中,其性能明显不如原始Transformer。这主要归因于嵌入表示(embedding)的缺失:在YZ Transformer中,嵌入信息全部被压缩进一个巨大的词汇表大小的Z矩阵中,导致在实际训练中难以有效优化。
"In the YD transformers, there's no embedding. The embedding is all put into this Z matrix... This huge matrix may not be trained well during the practical training."
此外,学习率的选择对训练曲线形态有决定性影响。实验发现,展现出“反弹”特性的训练曲线往往对应着最高的验证准确率。这意味着,允许模型经历从强特征到弱特征的探索过程(即反弹现象),可能是更正确的数据学习方式。相比之下,其他学习率设置下的训练曲线缺乏这种反弹特性,且最终性能较低。如果冻结底层参数后再解冻,可能会观察到底层再次反弹随后又下降的复杂动态,这为未来的实验提供了有趣的方向。
注意力矩阵动态与Token谱系的连续化分析
在多层Transformer的理论分析中,注意力矩阵的动态演化是一个尚未被完全解释的复杂领域。现有的理论框架往往需要简化假设,例如将Token严格划分为“多类发生”或“单类发生”的极端情况。然而,现实数据中存在大量中间状态,例如一个Token在90%的情况下属于某一类,但在其他类中也有少量分布。为了解决这一理论局限,研究者引入了DRA(Dynamic Representation Analysis)框架,通过构建特定的数学公式,自然衍生出一种连续谱系(Continuous Spectrum)。在这种谱系中,高频通用Token与低频独特Token不再是非黑即白的分类,而是通过一个连续的量值来区分。当该量值幅度较大时,代表独特Token;当量值接近零时,代表通用Token。这种处理方式消除了对极端假设的依赖,使得相干性(Coherence)和判别性(Discriminancy)成为公式中的自然项,无需额外特殊处理。
“The great thing about drama is that once you have the formulation. Then naturally you will have this kind of property... Now we have a continuous spectrum covering all the tokens.”
尽管这一进展解决了注意力机制中部分动态问题,但理论分析仍需进一步分解。目前的研究主要聚焦于注意力机制本身,而嵌入矩阵(Embedding Matrix)的动态变化尚未得到充分解释。未来的工作将致力于分析嵌入矩阵在训练过程中的演变,以完善对Transformer内部动力学的整体理解。
理论局限与未来研究方向:嵌入与位置编码
当前的理论研究存在一个显著的理想化假设:即假设嵌入向量是固定的且彼此正交。虽然实证结果支持这一假设在某种程度上成立,但这并非理想状态,因为在实际训练中,嵌入层也是被学习的参数。研究者承认,目前尚未深入分析嵌入层在训练过程中的具体动态,这是一个遗憾之处。此外,位置编码(Positional Encoding)的动态变化也是未来研究的重要方向。
在多層Transformer的设置下,潜在概念(Latent Concept)的学习过程缺乏具体的定理支持。目前我们仅知道概念层面的学习逻辑,但缺乏关于潜在向量如何随时间被学习的严格数学描述。这一过程非常棘手,难以处理。因此,后续论文计划专门解决嵌入层和位置编码的动态分析问题,以填补当前理论框架的空白。
MLP中的模式叠加与稀疏性加速推理
除了注意力机制,MLP层中的模式叠加(Pattern Superposition)也是一个重要发现。实证研究表明,MLP层中的单个向量可以编码多种不同的模式,这一现象在各向异性(Anisotropic)的研究中被强调,而DRA框架则从训练动力学的角度证明了这种叠加在从头训练时是可行的。基于这些理论洞察,研究转向了基于注意力稀疏性的应用,旨在加速大模型的推理过程。
核心思路是仅保留高激活的Token,以加速推理。然而,实证发现许多注意力头和MLP层在每次前向传播中并未被激活,导致交叉注意力权重被无效加载。为了解决这一问题,研究者提出了DJW方法,通过预测哪些头在推理中不会被激活来优化性能。例如,某些头(如头42、43)给出均匀分布,而另一些头(如头42、44)则具有清晰的聚焦模式。通过为每层、每头甚至每个激活学习一个简单的预测器,系统可以预测上下文稀疏性,从而只加载必要的权重行和列。
DJW加速性能与H2O稀疏注意力机制
在GPU架构中,内存加载到缓存的速度极慢,通常需要数百个GPU周期,而实际计算仅需几个周期。因此,减少需要加载的权重数量可以显著提升前向传播速度。通过工程优化,DJW方法在多种基准测试中实现了比Hugging Face和Fast Transformers更低的延迟。该方法还与量化技术兼容,特别适用于小批量(Low Batch Size)场景。
此外,另一项名为H2O的研究利用稀疏注意力机制加速推理。由于标准注意力计算的复杂度为$O(n^2)$,每个新Token都需要计算与所有先前Token的注意力。通过利用稀疏性,H2O旨在减少这种计算负担。虽然具体性能数据在片段中未完全展开,但核心逻辑是通过预测和跳过非激活部分来降低计算复杂度。
| 对比项 | Hugging Face | Fast Transformers | DJW (本研究) | H2O (本研究) |
|---|---|---|---|---|
| 核心机制 | 标准实现 | 稀疏注意力优化 | 动态预测稀疏性+权重加载优化 | 稀疏注意力机制 |
| 延迟表现 | 基准 | 较低 | 更低 | 信息不足 |
| 兼容性 | 标准 | 标准 | 兼容量化 | 信息不足 |
| 适用场景 | 通用 | 通用 | 小批量场景 | 长序列处理 |
“When you do computation. It's actually a few cycles usually, so much much slower. So if you can load like fewer weights, you actually will have a much faster feedforward speed.”
注意力分布的非均匀性与Heavy Hitters概念
在Transformer模型的自注意力机制中,一个核心疑问是:是否真的需要保留所有历史token的信息?实证分析表明,注意力分数的分布呈现出极度的非均匀性,这与词共现分布高度相似。通过绘制激活值分布可以发现,绝大多数token的注意力得分极低,而少数token占据了绝大部分的注意力权重。这些高注意力分数的token被称为"Heavy Hitters"(重击者)。实验证实,如果移除这些Heavy Hitters,下游任务的性能会出现断崖式下跌;反之,仅保留它们进行推理,模型能维持极高的性能水平。这一发现为简化KV Cache提供了理论依据,即通过筛选关键token来优化计算资源。
"If we get rid of this token that has a lot of attention, then all of a sudden the downstream task performance drops substantially."
H2O算法与KV Cache压缩策略
基于上述发现,研究者提出了H2O(Heavy Hitter Oracle)算法,这是一种局部贪心策略。该算法在每一个解码步骤中,累加先前所有token的注意力分数,并丢弃累积注意力分数最低的token,仅保留一个固定预算内的KV Cache token。实验数据显示,仅保留20%的KV Cache token(即剔除80%的低注意力token),模型在多个基准测试中的表现几乎与完整上下文保持一致。相比之下,如果仅保留局部token(Local tokens),性能则会大幅下降。这种策略不仅显著节省了内存空间,还大幅提升了推理速度。此外,移除部分限制性的token甚至可能增加模型的创造性,因为它减少了模型对冗余历史信息的过度依赖。
Attention Sink现象与Streaming LM原理
进一步分析揭示了一个反直觉的现象:前几个token(通常是特殊标记)占据了约50%的注意力分数,这种现象被称为"Attention Sink"(注意力 sinks)。由于Softmax函数的归一化约束,注意力总和必须为1,模型被迫将大量注意力分配给初始token以平衡其他token的权重。研究发现,这些初始token的位置重要性远高于其内容重要性,甚至可以用随机token替换它们而保持困惑度(Perplexity)基本不变;但若直接移除它们,模型性能会崩溃,困惑度飙升至约5000,输出变为无意义的乱码。基于此,Streaming LM技术应运而生,它通过固定长度的KV Cache预算,始终保留初始token和最新生成的token,从而实现了无限上下文窗口的聊天能力。
"The first few tokens are important, but only their position is important rather than the content."
性能验证与可学习同步Token
Streaming LM在多种场景下进行了验证,结果显示其能保持稳定的困惑度,意味着输出始终具有语义意义,且内存占用随时间保持恒定,不受上下文长度影响。实验表明,该技术可将模型上下文扩展至400万token以上,且性能不衰减。为了进一步优化,研究者将初始token处理为可学习的同步Token(Sync Tokens),包括零嵌入(Zero Sync)和可学习嵌入(Learnable Sync)。实验对比显示,使用可学习同步Token的效果略优于固定嵌入,进一步验证了初始位置信息的关键作用。这种机制使得大型模型能够在硬件资源受限的情况下,实现长期、稳定的对话交互,解决了传统Transformer因KV Cache无限增长而导致的内存溢出问题。
| 模型/策略 | 上下文保留方式 | 性能表现 (困惑度/任务得分) | 内存/推理效率 | 备注 |
|---|---|---|---|---|
| Full Attention | 保留所有Token | 基准性能 | 高内存,慢推理 | 标准Transformer |
| H2O (20% Budget) | 仅保留Heavy Hitters | 几乎无下降 | 内存减少80%,推理加速 | 性能与全量接近 |
| Local Only | 仅保留近期Token | 性能大幅下降 | 低内存 | 丢失长期依赖 |
| Remove Initial | 移除前几个Token | 崩溃 (Perplexity ~5000) | 低内存 | 输出为垃圾信息 |
| Streaming LM | 保留Initial + Recent | 稳定,可扩展至4M+ tokens | 恒定内存,无限上下文 | 支持长期对话 |
| Learnable Sync | 初始Token可学习 | 略优于Zero Sync | 恒定内存 | 进一步优化初始位置编码 |
同步令牌与注意力机制的局限性
在Transformer架构中,同步令牌(sync token)的出现是一个关键转折点。一旦该令牌被引入,模型的困惑度(perplexity)会显著降低,使得用户能够顺畅地与模型进行对话,而不再受限于上下文窗口的长度问题。这一发现并非孤立存在,Meta等机构的研究团队也独立发现了类似的结构,特别是在视觉Transformer中,这种机制被称为Registers,其本质对应于注意力汇(attention sinks)。
然而,这种基于保留概率而非完整信息的机制存在根本性局限。它通过丢弃部分token来维持整体概率分布,虽然能给出“合理”的回答,但在需要精确检索特定信息(如密码提取)时,如果对应的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."
这种基本局限性意味着仅靠概率保留无法真正解决长上下文中的信息丢失问题,未来需要更深入的解决方案。
流式处理与HTO的结合
为了进一步提升性能,近期研究尝试将流式推理模型(Streaming RMs)与HTO(Hierarchical Token Optimization? 或类似架构)方法相结合。流式推理模型可以被视为HTO的一种特殊情况,因为Transformer通常只保留前几个token,而HTO则根据注意力分数动态选择token。将两者结合不仅能带来更优越的性能表现,还能保持无限聊天(infinite chat)的特性,即在极长序列中维持稳定的推理能力。
这一更新为快速推理提供了新的改进方向,通过动态选择而非固定截断,更好地平衡了上下文信息的保留与计算效率。尽管具体实现细节仍在迭代中,但这种混合架构展示了在处理超长上下文时的巨大潜力。
位置编码插值扩展上下文窗口
针对长上下文窗口的扩展,传统方法通常依赖预训练后的大量微调,但这需要消耗巨大的计算资源。一项被撤回的论文(因竞争过于激烈而未提交)提出了一种通过微调(fine-tuning)而非重新预训练来扩展上下文窗口的方法。该方法的核心在于改变Llama模型中的位置嵌入(positional embeddings)缩放比例。
具体而言,若要将模型从2K上下文扩展至4K,只需将每个位置编码索引除以2,从而将4K窗口映射回原有的2K空间。这种插值(interpolation)策略利用了注意力函数在插值时表现良好、而在外推(extrapolation)时表现极差的特性。通过这种方式,仅需极少量的样本和步骤即可完成高效微调。
"We actually discover that attention function is actually ill behaved... if you do extrapolation, and all of a sudden it gives you super bad scores... So that's actually the reason why we should do interpolation rather than extrapolations."
性能数据与微调效率对比
实验结果表明,通过改变位置编码进行插值,模型可以在极少的训练步骤内实现上下文窗口的显著扩展。以下是不同方法在扩展上下文窗口时的效率对比数据:
| 扩展目标 | 传统微调所需步骤 | 插值微调所需步骤 | 性能保持情况 |
|---|---|---|---|
| 2K -> 4K | >1000 steps | <1000 steps (约200 steps) | 基准测试性能基本不变 |
| 2K -> 8K | 资源消耗巨大,效果不佳 | 200 steps | 基准测试性能基本不变 |
| 2K -> 16K | 资源消耗巨大,效果不佳 | 200 steps | 基准测试性能基本不变 |
| 2K -> 32K | 资源消耗巨大,效果不佳 | 200 steps | 基准测试性能基本不变 |
数据显示,仅需200个步骤即可将上下文窗口从2K扩展至8K、16K甚至32K,即原始长度的8倍。相比之下,传统微调即使使用1000个步骤,也难以有效扩展到2.5K以上,且资源消耗巨大。这种方法不仅大幅降低了计算成本,还保持了模型在各项基准测试中的原有性能,证明了位置编码插值在长上下文扩展中的高效性与可行性。
经验规律与第一性原理的理论分野
当前大语言模型(LLM)领域存在两种截然不同的研究范式:一种是基于经验观察的缩放定律(Scaling Laws),另一种是试图从第一性原理(First Principles)出发的理论推导。对于像 Siri 这样深层堆叠的模型,其展现出的涌现能力(Emergent Capabilities)往往通过经验曲线来预测。例如,根据缩放定律,当模型参数量和数据量同时扩大十倍时,性能变化遵循特定的曲线趋势,这种经验性描述在实际应用中非常有用,因为它能直接指导工程实践。
然而,从第一性原理出发的理论则更为艰难。它要求从基础公式出发,推导在特定设定下必然产生的后果。虽然这种路径难以直接应用于即时场景,但如果能成功打通从基础理论到宏观现象的全链路,将能覆盖多种适用场景。正如对话中所言:"> 对于我讨论的理论,更多的是从第一性原理出发,从公式化开始,思考这种设定会达成什么样的后果。"
从微观推导至宏观缩放定律的路径
将第一性原理与缩放定律连接的关键,在于模拟与拟合的过程。理论上,我们可以基于基础方程模拟拥有更多数据时的系统行为,从而推导出边界条件,进而形成类似经验观察到的缩放定律。这一过程类似于物理学中的研究路径:从基础方程出发,模拟一千个原子或 $10^{32}$ 个原子的行为,最终总结出全局性的经验定律,并与真实实验数据进行拟合。
这种从微观到宏观的推导并非一蹴而就,而是需要经历漫长的探索。早期的理论研究往往充满痛苦,研究者可能在某个无法证明的问题上耗费数月,最终发现该证明并无实际益处。但随着经验的积累,研究者能够建立起对系统演化的粗略图景(Rough Picture),从而明确需要证明的定理方向。即使初始定理不完全正确,只要其大致指向正确的方向,便可通过添加约束条件和实证检验逐步完善,直至获得既逻辑严密又具实证价值的结论。
理论研究的迭代方法论:符号与现实的博弈
理论研究的精进过程被描述为一种“艺术”,需要在符号结构(Symbolic Structures)与现实世界结构(Real World Structures)之间反复迭代。研究者首先从现实世界出发提出假设,将其转化为符号层面的推导,检验这些假设在符号系统中是否简化了问题或具有理论意义,然后再回到实践中验证其相关性。通过这种“来回穿梭”的过程,最终收敛于既具有符号趣味性又与现实世界紧密相关的假设体系。
以“Drama”论文为例,其核心思想——关于嵌入绘制(Drawing Embeddings)的概念,可能在几天内即可确立。一旦推导出关键方程,后续工作便围绕该核心展开,极大地提高了研究效率。在数学处理上,为了捕捉图像的本质感觉,研究者常采用$t \to \infty$的极限假设,忽略有限的数学细节,从而快速构建理论框架。然而,有一个关键约束是绝不假设特征固定(Feature is fixed),因为这将违背模型动态演化的本质。这种从直觉到形式化,再回归实证的循环,是提升理论推导效率的核心方法论。
假设与物理图景的平衡
在构建模型解释性理论时,数学上的便利性往往与物理图景的真实性存在冲突。如果为了简化数学推导而做出过于理想化的假设,可能会彻底破坏对模型实际运行机制的理解。因此,研究者必须在简化假设与保持物理真实性之间找到平衡点。即使假设能让数学变得简单,也必须确保该假设不会扼杀核心的物理图景。这种平衡是理论构建中的关键难点,需要反复验证假设是否偏离了实际现象。
"Even if this will make the mathematics much easier, this basically kills the entire physics picture. That's not a good thing to do."
层次化表示的长期探索
作者回顾了其在层次化表示(Hierarchical Representation)领域的长期研究历程。早在硕士阶段,他就致力于视觉模型中的层次化表示研究,尽管当时发表的论文因符号过于复杂而被拒绝,但这成为了他后续研究的重要动机。对于卷积神经网络(CNN),层次化结构是其固有特性,但分析起来极具挑战性。作者指出,符号复杂性是阻碍该领域理论突破的主要障碍,许多优秀的层次化分析工作因此未能发表。这一经历促使他转向更易于形式化分析的Transformer架构,试图在保持层次化洞察的同时,解决数学表达上的难题。
高维空间与嵌入表示的效率优势
Transformer与CNN在表示机制上存在本质差异,核心在于嵌入(Embedding)的强大能力。在高维空间中,嵌入向量具有独特的几何性质:在$d$维空间中,可以容纳$2^d$数量的向量,且它们之间几乎保持相互正交(Orthogonal)。这意味着在有限的样本空间中,可以编码指数级数量的对象或概念,而不会发生严重的重叠或干扰。这种高维正交性使得Transformer能够高效地存储和区分大量信息,这是传统方法难以企及的效率优势。
| 特性维度 | CNN表示机制 | Transformer表示机制 |
|---|---|---|
| 表示类型 | 近似One-hot表示 | 密集嵌入向量(Dense Embedding) |
| 正交性要求 | 概念间需严格正交,维度需求极高 | 高维空间中自然保持近似正交 |
| 容量效率 | 线性或低效扩展 | 指数级容量($2^d$) |
| 概念组合 | 难以解耦不同属性 | 支持向量的线性组合与解耦 |
推理能力与属性的解耦表示
Transformer的推理能力与其嵌入表示的结构密切相关。作者推测,复杂的概念(如单词"a")可能由多个嵌入向量的线性求和构成,例如一个向量代表字母属性,另一个代表大小写属性。这种表示方式实现了属性的解耦(Decoupling),使得Transformer可以在计算过程中独立处理不同的属性特征。这种解耦机制不仅提高了表示的灵活性,也为理解模型内部的推理过程提供了新的视角。未来的研究重点将在于如何进一步分析这种解耦表示在多层Transformer中的动态演化过程。
"Another point that makes things efficient is that for transformers, it can also do lots of reasoning tasks. And the reasoning, I would believe, is definitely related to representation of the embeddings."