一个乱码提示当钥匙，10个注意力头定任务：拆开GPT-J的黑箱

10万条句子训练出来的一个"软提示"，只有4万个参数，解码成人类语言是一串乱码。就这个东西，能从GPT-J的一个中间状态里，读出模型三个词之后要说什么。参数量是它五千倍、约2亿的线性解码器，干同样的活干不过它。

这是Cognitive Revolution播客2024年4月5日一期讲的事。主持人Nathan Labenz请来美国东北大学Bau实验室的三个人：教授David Bau，Future Lens论文一作Koyena Pal，Function Vectors论文作者Eric Todd。两篇论文，同一件事：把语言模型的黑箱撬开一条缝，看里面到底在算什么。

先从一个番茄说起

Bau在Google工作多年，2015年前后看到领域要变天，辞职出来做研究。他的转折点是2012年的AlexNet。神经网络的想法来自上世纪50年代，训练方法来自80年代，学界长期认定这么简单的流程不可能走到最后。AlexNet证明了一件事：把神经网络堆大，它就真的能行。Bau起初研究生成式视觉模型GAN，本打算长期做下去，结果GAN之后来了扩散模型，之后又是语言模型，一直是移动靶。

实验室的动机说起来很简单。机器学习是一种新型软件，传统计算机科学的调试工具对它全部失效，手里剩下的都是统计性、黑箱式的手段。Bau想回答两个问题：能不能恢复调试和理解软件的能力？能不能像改传统程序那样，改模型里的一行"代码"？

他打了个比方：这就像问能不能完全控制一个番茄的内部机制。番茄从进化里长出来，模型从训练里长出来，都属于"长"出来的东西。生物学家想改良番茄，得先逆向工程搞清楚它怎么运作。Bau实验室做的就是给神经网络建立这样一门科学。Nathan在节目里顺带指出，这个领域最早一批做内部机制研究的人，比如Chris Olah，也是从视觉模型起步的，Bau的路径正好踩在这条暗线上。

这条线上的前期代表作是ROME和MEMIT。ROME定位了事实知识在模型里的位置：把"迈克尔·乔丹打篮球"改成"打棒球"，改完还得经得起换着花样的提问。后续的MEMIT一次能改1万个事实。即将发在ICLR的LRE论文再往深拆一层，研究"某人住在某城""某人打某运动"这类关系本身的表示。结果发现很多关系确实是个线性映射，但也有相当一部分案例呈现非线性。模型为什么要把某些关系建得那么复杂，那些复杂性换来了什么，还是开放问题。

Bau说这一切背后有个根本问题："知识"在神经网络里有没有物理实体？有一派观点认为一切都弥散在所有神经元里，谈不上定位。但无论生物大脑还是人工网络，研究者这些年一次次看到局部性。ROME回答的是单个关联存在哪，LRE开始回答关系存在哪。

Future Lens：模型在提前想后面的词

Koyena的项目起点是研究模型的记忆力。她注意到一个现象：给ChatGPT某段许可证代码的开头几个token，它能一字不差背出训练数据里接下来的11到20个token。研究过程中她撞上了更大的问题：一个隐藏状态里，除了当前token，还存了多少未来的信息？

先交代工具。logit lens是博主nostalgebraist搞的，把transformer每一层的中间状态直接投影到解码层，把数字翻译成词。模型在浅层"想"的词很粗糙，越往深层越精确，最后收敛到正式输出。Bau推荐没见过的人去搜一下，那种一层层叠起来的词网格非常好看。

logit lens的短板是只能看模型"现在"在想什么。模型处理"New"这个token时，心里想的是New York还是New Jersey？看不到。

Nathan给这个问题画了个坐标系：层与层之间的加工是纵向的，token之间的信息流动是横向的。logit lens看的是纵向，Koyena要探的是横向。英文里早有旁证，模型选冠词a还是an，取决于后面名词是不是元音开头，冠词选对了，等于提前看了后面的词。

Koyena用的是GPT-J，60亿参数，隐藏状态4096维，28层。她从Pile语料抽了10万条文本，筛选条件是模型必须把下一个token预测对，另留1000条做测试。对每条文本取最后一个token在每一层的隐藏状态，然后问：从这个4096维向量里，能读出未来几个token？

五种方法赛马：

• bigram基线：只统计Pile里词对的共现频率来猜下一个词，猜对20%。Nathan说没想到这么高，Bau补了一句：n-gram在几年前还是最先进方法。
• 直接词表预测：训练一个线性层，从4096维直接映射到50400维词表，取概率最高的词。
• 线性近似：不直接猜词，猜未来的隐藏状态，再让模型自己的解码层去解。
• 固定提示移植：把隐藏状态移植进"Hello, could you please tell me more about"这类通用句子，继续前向传播。
• 软提示：放弃人类语言，直接学10个4096维的嵌入向量，训练它们把隐藏状态里的未来信息"钓"出来。

结果一层层揭开。两个线性方法的表现几乎一模一样，曲线全程贴着走。Nathan觉得这事有点诡异：两个线性分类器，两个不比一个强。Koyena的解读是，线性方法已经把这部分信息榨干了。它们预测当前token时比bigram基线强一点，预测后两个token时甚至不如基线。

固定提示移植更惨，某些设置下比bigram还差，而且层数越深表现越差。Nathan追问了好几轮，为什么这么自然的想法会失败。Bau给了一个很形象的类比：这就像往人脑子里随机插根探针，然后指望这能让人闻到香蕉味。往一个运行中的前向传播里硬塞外来状态，大概率只是造成某种脑损伤，模型根本接不住。他们试了很多提示，论文里这个已经是最好的。

软提示一骑绝尘，而且它只有4万个参数，比线性方法的约2亿小了五千倍。Nathan算完账感叹：表现最好的方法，参数最少。Bau总结成一句话：这是用transformer解释transformer自己的力量。模型权重全程冻结，软提示只是学会了怎么"问"。

顺带一提，Nathan联想到此前的通用越狱提示研究，那边优化出来的万能越狱串也经常是乱码。优化器找到的"钥匙"，常常长着人类读不懂的脸。

钥匙找到了，但它开的是模型自带的字典

论文的图4值得盯着看。四张图分别对应预测当前token、下一个、下两个、下三个。预测当前token那张复现了logit lens，越深越准。预测未来token的三张图里，软提示那条线远超其他方法，而且在中层达到峰值，再往后开始下降。测后两个、后三个token时，实验设定会把前面正确的token喂给模型，再看它能不能接上更远的词。

Bau对下降的解释：transformer的后期层在主动抹掉一些信息，把空间腾出来干别的。模型自己擦除自己的信息，这个现象本身就值得单独研究。Koyena的想象是：开头在建立"我要说什么"的理解，中间握着一批想法，末尾只想着立刻吐出下一个词。

Bau对结果的解读指向一个更深的结论：信息就在那里，但他们没法直接解码，transformer自己的电路是解码这些信息的必要环节。他的心智模型是，transformer的权重里存着一本自己的字典，记录着怎么把概念展开成词序列。隐藏状态里存的更像一个指针，一个紧凑的向量，要靠模型内部的机制去查表，最后滚出"Back to the Future"整个短语。论文里用的例子正是这个词组。

如果解码过程简单，直接解码方法早就找到了。编码复杂到必须借用模型本身，这解释了软提示为什么有效：它像一把钥匙，触发了模型内置的查表机制。钥匙本身不是说明书。

Nathan在节目里修正了自己中途下的一个判断。他原本想，两个线性变换不比一个强，那上整个transformer大概也没用。现在看，可能恰恰需要非线性。Koyena同意这个方向应该更强，但他们没试，10万条数据也不够训练一个transformer，真要训练得百万级甚至上亿样本。

Bau强调另一个数字的含义：10万条样本就解锁了所有这些短语，而10万条远远不够让一个小模型重新学会GPT-J那60亿参数里的东西。所以这更像解锁模型里已经存在的东西，模型没学到任何新能力。Nathan补了一个很妙的角度：再加一个transformer去解释transformer，就成了先有鸡还是先有蛋，某种程度上反而背离了可解释性。

10个注意力头，决定模型在做什么任务

Eric的问题更大：in-context learning到底怎么工作的？看几个例子就会做任务，GPT-3当年的惊艳全靠这个能力。

他找到的答案是一小撮注意力头。实验设计：给模型看几对反义词，big和small，short和tall，然后给出查询词bright。把模型处理这些示例时某几个注意力头的输出抽出来，加总成一个向量，加到另一段完全干净的提示的某一层上，比如"the word fast means"后面。模型没有像平常那样解释fast的定义，直接说出了slow。任务本身被作为一个向量注入了。

换首都任务也一样：Paris对France，Moscow对Russia，念到Ottawa的时候读出那几个头的输出，里面编码的就是"报首都"这件事。Eric还试了把向量塞进自然文本、不同模板，长短句通吃，移植的任务和原任务长度不需要匹配。

找这些头的方法叫因果中介分析，也叫激活移植：一个干净句子正常展示任务，另一个句子把标签全部打乱，看不出在做什么任务。把干净设置里某个注意力头的输出移植到打乱的设置里，测量模型输出朝预期方向变化多少。单个头不够翻转预测，攒到10个左右，效果就明显了。

他们测了40多个任务，首都、翻译、反义词，全都由同一小撮头中介。Bau说这是他最大的意外：他本来以为会存在"反义词注意力头"，结果发现的是通用的任务理解头。从反义词任务里挑出的头，换到翻译任务上照样注入翻译任务。模型内部有一条专门的通路在传递"我现在在做哪类任务"这个信息，尽管提示里从没人明说过示例和标签之间的关系。

这批头有两个特征。位置上大部分埋在中层，Nathan数了数论文图3里前十名的分布：有的层挤两个，有的层一个，再往后就安静了。行为上注意力模式高度一致，都盯着前面示例的标签token，也就是每一对词里的"答案"。这和其他研究的发现对得上：模型做in-context learning时，标签token是存放关键信息的位置。

Eric对transformer有个两段论：前半段搞清楚"我要从上下文里拿什么"，后半段拿这个去预测下一个token。任务识别恰好发生在中段，就在模型开始输出之前。

更大的模型会更好懂吗

Nathan抛了个大胆的猜想：更大的模型，中间状态的语义结构会不会更清晰，反而更好读？Bau的态度很坦率：这是个反直觉但值得认真对待的猜想。传统预期是模型越大越难解释，如果规模带来的结构更多而不是更少，那会是个相当强的主张。他直接把这标记为实验室的知识缺口，需要系统实验。

Mamba这类无注意力架构那边有初步观察。Bau说他们做到一半，不想误导听众，但可以透露：这些没有传统注意力机制的模型，似乎也能完成一些transformer靠注意力才完成的复杂计算。具体怎么做到的，还在拆。Nathan顺带提到用StripedHyena混合架构做DNA基础模型的新工作，混合架构赢了纯transformer，他押注未来不少场景的赢家都是混合体。

聊到混合专家模型时，Bau把话题拉回到了整个领域最根本的问题上。Nathan的背景是：GPT-4据传是MoE架构，Mixtral开源带火了这条路，Gemini 1.5也被认为走的大规模MoE路线。Bau的回应是，注意力头里有结构，MLP里有结构，但十几个头协同工作时，正确的抽象可能比单个头更高一层。Function vectors是一个提案：把一堆头压成一个向量，丢失细节，但保留了强因果效应，还能做组合。

我们还不清楚理解这些系统该用什么抽象层级。我们知道它们是用什么造的，矩阵乘法、非线性、各种模块，但怎么看这些计算才是对的，还不确定。可解释性领域所有人都在追的，就是这个问题：正确的抽象层级到底是什么？

明天就能动手的三件事

第一，Future Lens的代码、笔记本和脚本全开在future.baulab.info。测试部分Colab Pro就能跑。Bau明说了，固定提示那条路他们没走通，但没准有更聪明的设置能走通，谁找到他们会很高兴。

第二，想做激活移植类实验，看nnsight这个库（nnsight.net），Baukit的继任者，几行代码就能完成移植干预。项目还在早期，入口是官网角落的Discord链接，社区不大但很活跃。库支持远程后端NDIF，本地没GPU也能跑Llama 7B这种自己不好搭的模型，改个flag就行。它甚至不限于transformer，已经有人在上面研究Mamba。

第三，Bau实验室论文页面上的可视化值得专门点开看。他还替学生Nikhil的论文打了广告：《fine-tuning enhances existing mechanisms》，研究微调和模型内部电路的关系，也要发ICLR，电路可视化被Bau称为他见过做得最好的。

这期节目给外行的最大启发可能就一条：黑箱内部有结构，有字典，有钥匙孔，有专门负责"识别任务"的岗位。拆解它的工具已经开源，门槛低到一台Colab Pro就能进门。