线性注意力背后的视角转换
我最近阅读了 Google Research 的两篇论文【1】、【2】,并回溯了最早提出线性注意力机制的论文【3】。我在阅读过程中再一次体会到视角转换如何为一个理论框架赋予新的想象空间。如果你没有经历过这种转换,你会低估它的深远意义。联想量子力学的发展历史:
矩阵力学 / 波动力学 => 量子场论 / 正则量子化 => 路径积分
每次变革不仅仅是用一套新的数学公式重写既有知识,而是用一套全新的视角和语言审视理论。在这个意义上,(数学)语言不仅是思维的边界,它内嵌了思维的空间结构。
线性注意力机制打开了什么空间?传统的注意力机制本质上做了两件事情:
1,将输入X映射到三个矩阵:Key,Value,Query
2,计算输出:Y=Softmax(Q*K)*V(为了表述简便,忽略归一化与掩码矩阵)
要实现全局注意力,就要一个包含位置信息的全局矩阵来编码注意力矩阵Q*K,与之伴随的是空间平方的复杂度,和超长上下文的工程难题。如果我们放松Softmax的计算要求,用线性核来解耦注意力矩阵:Softmax(Q*K) ~= phi(Q)*phi(K),那么我们可以改换输出公式里的乘法顺序(忽略核函数):Y=Q*(K*V)。此时我们只需要维护一个与上下文长度无关的固定矩阵K*V即可,计算复杂度降为空间线性。缺点也是显而易见的:我们无法精准地编码位置信息;随着文本增长,早期记忆更容易被灾难性遗忘。正如【3】所指出的,线性注意力机制本质上是一种RNN。
这种重述带来了新的视角。我们把线性注意力的输出公式拆分为两步:
1,M+=K*V
2,Y=Q*M
中间状态M描述了一个随着输入流更新的记忆状态,也就是 RNN 里的隐状态。这是一个典型的关联记忆模型(associative memory),它编码了从 Key 到 Value 的映射关系。面对一个新的 Query,输出行为就是从M中提取出与Q关联的记忆。关联记忆模型正是去年 Hopfield 和 Hinton 获得物理学诺贝尔奖的工作的基础。启发它的赫布理论以及衍生出的突触可塑性依然是今天主流神经生物学的基础理论。
中间状态M的引入,让我想起量子力学历史上矩阵力学与波动力学之争。深受马赫主义影响的海森堡认为物理理论应当避免引入缺乏操作定义的量,比如电子在轨道上的精确位置和经典动量等,只有轨道之间的跃迁信息是可观测量,应当用矩阵来描述。与此同时,德布罗意从光的波粒二相性出发,认为一切物质都具有波动性,需要用波函数描述这种物质波。注意:在波恩的概率诠释之前,德布罗意认为这是真实存在的实体波(就像电磁波一样),而不仅仅是便于计算的隐喻。虽然两个理论很快被证明在数学上是等价的——从本体论层面上,波函数从物质实体降格为量子态在位置表象下的表征(representation)——但是波函数相较抽象晦涩的矩阵拥有几何直觉与可视化的优势。它在帮助我们想象(而非理解)理论的可能性上仍有价值。
在我看来,线性注意力的中间状态M就是这个多出来的波函数。它不仅在方法论上帮助我们理解线性注意力机制,更是在本体论层面打开了一个新的探索空间,即对于“记忆”本身的建模与训练。
在此基础上,【1】提出了一个更深刻的视角转换。它将第一步,记忆更新,视作一个在线学习(on-line learning)的动力学过程,指导这个过程的,是一个全局损失函数L。于是,记忆更新行为被视作:为了最小化损失函数,参数M对数据流的动态响应。即:
M=argmin L(M(K),V)
我们对反过来的表述都很熟悉。传统的机器学习就是基于这样一套逻辑:预设一个目标函数,通过已标注训练集(在线或离线),训练出一个最佳模型。此处,目标函数是具象的(预测与标注的误差),模型参数是抽象的。而在“从关联记忆到在线学习”的图景转换里,模型参数是具象的(记忆),参数的动力学过程也是具象的(记忆更新),我们要反推出一个目标函数,是相对抽象的。不过【1】还是为这个目标函数赋予了一个诠释,称为“惊讶度”,直观理解为:人们对习以为常的现象不会加以过多关注,意外事件更容易被记住,即触发记忆更新机制。
任何一个接受过高等物理学训练的人都会对这个逆向思维非常敏感,因为它是分析力学的核心思想。基于速度、加速度、力等矢量概念的牛顿力学表述是动力学过程,它用“动力因”描绘了一整套机械宇宙图景;但是当代力学体系所依赖的是一系列标量:哈密顿量、拉格朗日量、作用量、路径积分、配分函数。这些标量不仅完美兼容动力学表述,还可以简洁地描述系统遵循的对称性,深刻揭示对称与守恒的关系,优雅地处理边界条件和自由度。它们极其晦涩抽象,但是在数学上如此优越,足以让人们抛弃一个更直观的动力学图景。
一旦为线性注意力机制赋予了关联记忆和在线学习的诠释,我们就可以泛化以上两个步骤。【4】将泛化自由度归为四类:
1,记忆框架,用神经网络取代矩阵编码K->V,获得更高的自由度同时降低过拟合风险。
2,优化目标:设计损失函数。
3,遗忘门:记忆与遗忘一体两面,学习因子和遗忘因子本身可以是可学习的参数。
4,优化器:从损失函数到记忆更新的优化机制,Gradient Descent, SGD, Adam…
【1】将在线学习称为 Test-time Training(TTT)。它打了个很形象的比方:今天的大模型是一个顺行性遗忘症(anterograde amnesia)患者。它拥有在某个时间节点前(预训练语料库截止时间)的所有记忆。在运行时,它也拥有执行该任务所需要的所有短期记忆(上下文)。但是,一旦任务结束,这些上下文被立刻抛弃,无法形成长期记忆。在线学习可以突破这个藩篱,因为它不必保留所有上下文,只需要通过流经的短期记忆形成长期记忆(也就是M)即可。与之相对的基础模型(提供了从输入X到QKV的映射矩阵)除了提供对世界的基本理解以外,还担负着“更有效地形成长期记忆”的职责。
对于这一点,【2】认为【1】的做法不够精致。从更新频率的视角来看,【1】本质上提供了二元更新频率,即:
外层循环,f=0,基础模型,静态,不更新。
内层循环,f=1,以token流为单位,每流经一个token就更新一次M。
而人的记忆更像是一个层级结构:从原初的感官信息(sensory registeration)开始,部分进入工作记忆(working memory)并被编码为具体的情境记忆(episodic memory)。随着系统整合的发生,这些记忆经历“去情境化”(decontextualization),最终提炼为稳固的语义知识(semantic knowledge),并内化为深层的认知图式(cognitive schemas)与世界观。(当然实际情况比这个单一链条复杂得多,比如我们今天的大部分抽象知识都在语义知识空间里直接生成,并没有对应的情境记忆。)这不仅仅是一个抽象的认识论模型,它更有着神经生物的基础,在大脑中与高频(gamma)到低频(delta)的神经振荡相对应。这也是当前神经科学研究的核心范式之一。
【2】设计了这样一套层级的记忆与学习模型框架。如同深度神经网络和stacked transformer在空间维度实现信息的层级抽象,记忆行为在时间维度实现层级抽象。当然它同时也是空间维度的,因为无论【2】还是大脑,不同频率的记忆与学习发生在不同的区域(从后脑到前额叶)。【2】提出了一个洋葱式的内嵌结构,最外层直面原初信息流,更新最快;最内层是基础模型,不更新;中间每一层都在某个特定的频率上对内层提炼与浓缩。【1】可以被看作只有两层的洋葱。
“层级结构”再一次触动了我的物理学神经。从某种程度上来说,物理学是研究“尺度”的学科。很幸运,我们并没有生活在一个缩放对称的物理世界(不然世界会变成一幅宏大、氤氲、永恒流动的分形抽象画)。指导整个物理世界运作的,是几个有量纲的物理常数,其中最基本的就是普朗克时间和普朗克空间。另一方面,物理学探索跨越时空的永恒性与普世性,那么时空尺度就必然成为任何理论都无法绕过的属性。重整化群(Renormalization Group,RG)便是探索物理理论与尺度关系的数学工具。这个令人困惑的名字来自量子场论。“重整化”最初是用来消除量子电动力学中臭名昭著的紫外发散问题的“奇技淫巧”,之后人们发现它背后蕴含某种深刻的不变性,被凝聚态物理借用后发扬光大,成为一门关于尺度的普世方法(Wilsonian RG)。“群”意味着它处理的尺度变换符合结合律。对这部分物理感兴趣的可以读我的科普书【6】。
重整化群是一套普世方法,无论对象是量子电动力学、永磁体还是沙堆,无论尺度变换是连续的还是离散的,它都用同一个框架去审视系统。在这套框架里最核心的概念是不动点(fixed point),或严格来说是非平庸的不动点(称之为“临界点”)。临界点针对参数空间而言,它指系统在一套特定的参数下(比如磁铁所处的温度和外磁场强度)不随时间或空间的尺度(比如磁铁原子的自旋块,block spin)的变化而变化。临界点并不是一个“点”,而是一切符合尺度不变性的参数子空间。在深度学习模型里,空间尺度坐标是神经网络的深度(从下到上);在【2】里,时间尺度坐标是洋葱的层数(从外到里)。(其实重整化流的图像对扩散模型最直观,因为它的时间和空间尺度变化是绑定的。)
一个好的模型一定在一个临界点上,反之未必。我们可以想象两类平庸的不动点。红外不动点——过拟合,模型仅仅是一本字典,没有泛化能力;紫外不动点——白噪音 / 灾难性遗忘,模型记不住任何东西,至少无法可靠地记住任何东西。而在非平庸的临界点上,模型收敛到了某个固定的参数流形上。它有效地冲刷掉了原初信息中的噪音,保留了记忆和认知中最深层的结构。注意,临界点性质不足以判断模型好坏,但它提供了一套归类方法(普世类,Universality Class),将模型的微扰自由度分为相关算符(relevant operator)和无关算符(irrelevant operator)。在同一个普世类下,所有无关算符都会被积分掉(integrated out),而相关算符则不断被强化。我想提出的是,对于记忆模型而言,特别是在个性化AI的语境下,我们真正关心的不是那些强相关算符,而是在两者边界上的边缘算符(marginal operator)。强相关算符对应的是基础模型中已经习得的通用知识(语法、基础知识、普世价值观),它对记忆模型而言并不重要;真正重要的是从“我”的记忆中提取的关于“我”的偏好、思维方式和价值观,这些信息最有可能编码在边缘算符中。令人期待的是,重整化群不仅为这种概念区分提供了可量化定义的蓝图,还提供了一个系统方案来计算模型在相关算符附近的状态方程(equation of state)。它不是第一性原理层面上的动力学方程(equation of dynamics),但它提供了丰富的模型库,来探索普世类流形的动态平衡关系。
以上是我最近的阅读体验。如果你看完一头雾水,这不怪你(当然也不怪我)。首先我自己并没有想清楚许多细节,特别是如何将【4】里的自由度映射到RG框架里;其次我也没有面向普通读者写作(和我的科普书不同)。我想象中的读者是兼具物理学和AI研究背景的科学家——这在今天的AI圈不在少数。物理学对于AI的滋养远未结束,尤其在关于“尺度”的研究上——这远不是一个“Scaling Law”可以涵盖的丰富宝藏。
参考:
【1】Titans: Learning to Memorize at Test Time (Dec 2024)
【2】Nested Learning: The Illusion of Deep Learning Architecture (Nov 2025)
【3】Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (Aug 2020)
【4】It’s All Connected: A Journey Through Test-Time Memorization,
Attentional Bias, Retention, and Online Optimization (Apr 2025)
【5】线性注意力简史:从模仿、创新到反哺 By 苏剑林,2025-06-20
【6】什么是物理?用物理学的视角看世界(下):近代物理篇,2024