学习大语言模型激活的生成式元模型 Learning a Generative Meta-Model of LLM Activations
用扩散模型学习LLM激活分布,作为先验提升引导干预和可解释性
前置知识
残差流激活(Residual Stream Activations)
在Transformer架构中,残差流是各层输出累加形成的向量通道。每一层的注意力或MLP块将信息写入残差流,后续层从中读取。对于一个d维模型(如Llama1B的d=2048),每个token位置对应一个d维向量,这些向量编码了模型对当前上下文的内部表征。本文所建模的正是这些中间层的激活向量,而非网络权重。
本文的核心目标就是对这些激活向量的分布进行建模,理解其几何结构是理解全文的基础。
流匹配扩散模型(Flow Matching Diffusion Model)
流匹配是扩散模型的一种变体,其前向过程将数据点 $z_0$ 通过线性插值添加噪声:$z_t = (1-t)z_0 + t\epsilon$,其中 $t \in [0,1]$,$\epsilon$ 为标准高斯噪声。反向过程则训练神经网络预测速度场 $\hat{u}_\theta(z_t, t)$ 来近似目标速度 $u = \epsilon - z_0$,从纯噪声 $z_1 \sim \mathcal{N}(0,I)$ 迭代采样回到干净数据。这种方法避免了传统扩散模型复杂的噪声调度设计,训练更简洁稳定。
GLP的训练目标直接采用流匹配框架,理解其前向/反向过程是理解模型如何学习激活分布的关键。
Fréchet距离(Fréchet Distance, FD)
Fréchet距离(也叫Fréchet Inception Distance, FID在图像领域的变体)衡量两个多元高斯分布之间的距离。它计算生成分布和真实分布在特征空间中的均值和协方差差异:$FD = \|\mu_r - \mu_g\|^2 + \text{Tr}(\Sigma_r + \Sigma_g - 2(\Sigma_r \Sigma_g)^{1/2})$。值越低表示生成样本越接近真实分布。由于激活向量无法直接用肉眼评估质量,FD成为衡量GLP生成质量的核心指标。
论文使用FD作为评估GLP生成质量的主要指标,也是验证模型是否学到有意义分布的关键证据。
稀疏自编码器(Sparse Autoencoder, SAE)
SAE是当前LLM可解释性的主流方法之一。它将激活向量通过一个稀疏瓶颈层进行编码和解码:输入激活 $a \in \mathbb{R}^d$,编码为稀疏隐变量 $h = \text{ReLU}(W_e a + b_e)$(隐变量维度远大于输入维度,如16384维),再解码重构 $\hat{a} = W_d h + b_d$。训练目标是重构误差加稀疏正则化 $\|a - \hat{a}\|^2 + \lambda \|h\|_1$。SAE假设概念是激活空间中的线性方向,这一强假设是本文要突破的核心限制。
SAE是本文最重要的基线方法,论文的多个实验都与SAE进行对比,理解SAE的线性假设及其局限是理解本文创新点的前提。
激活引导(Activation Steering)
激活引导是一种在推理时控制LLM行为的技术。它在模型的中间层激活上加一个方向向量来注入某种概念,如正面情感或特定人格。例如DiffMean方法通过对比两组数据的均值差来提取概念向量 $w$,然后在推理时修改激活 $a' = a + \alpha w$。核心挑战是:引导系数 $\alpha$ 越大,概念注入越强,但激活越容易偏离自然流形,导致生成文本质量退化(不流畅、重复等)。
激活引导是本文GLP最重要的应用场景之一,GLP通过将偏离流形的激活投影回自然分布来解决引导强度与质量的矛盾。
研究动机
现有分析LLM内部激活的方法——如主成分分析PCA和稀疏自编码器SAE——都依赖于较强的结构性假设。PCA假设最重要的结构是数据的主方向(线性子空间),SAE假设概念以稀疏线性方向编码在激活空间中。这些假设在很多场景下是合理的,但它们的致命弱点在于:当用这些方法对激活进行干预(如引导)时,重构或修改后的激活很可能会偏离激活的真实分布流形,产生'离流形'(off-manifold)的激活向量。具体来说,当引导系数增大时,SAE解码器产生的激活和DiffMean方向向量加到原始激活后的结果,都会使LLM生成质量严重退化——表现为文本重复、语义混乱、不连贯等问题。Templeton等人(2024)和Vu & Nguyen(2025)的工作也观察到了类似现象。这些结构性假设本质上是用简化模型去逼近一个复杂的高维分布,当干预强度超过假设的适用范围时,简化模型的局限就会暴露。
本文的目标是本文的具体目标是:训练一个无结构性约束的生成模型来学习LLM激活向量的分布,使其能够(1)生成与真实激活不可区分的样本,验证模型确实学到了激活流形的结构;(2)作为下游任务的先验,将偏离流形的激活投影回来,改善引导干预的输出质量;(3)作为特征编码器,通过模型内部的'元神经元'自动发现可解释的概念单元。此外,作者希望验证这个生成式元模型是否遵循可预测的缩放定律(scaling laws),即更多的计算投入是否能持续改善模型质量和下游性能。
与已有工作不同的是,本文的独特切入点在于:它不把激活视为需要分解的'信号',而是视为需要建模的'数据分布'。这一视角转换的深远意义在于——就像图像扩散模型学习了自然图像的流形结构,从而能将不真实的图像投影回自然图像流形(如SDEdit所做的),一个学习了激活分布的扩散模型也可以将偏离流形的激活投影回自然激活流形。已有工作如Li等人(2024)在DINO激活上探索了扩散模型,但将其用作生成条件信号或潜在空间,而本文的核心创新是直接利用元模型本身——其学到的先验知识和内部表征——而非其生成的样本。此外,论文系统性地验证了生成式元模型的缩放行为,证明扩散损失可以作为下游任务性能的可靠预测指标,这在已有元模型文献中尚属首次。
核心方法
想象一下:我们有数以亿计的LLM内部激活向量,它们构成了一个高维空间中的复杂流形。本文的核心思路是——用扩散模型去学习这个流形的形状。就像图像扩散模型学习了'什么样的图像是自然的',GLP(Generative Latent Prior,生成式潜在先验)学习了'什么样的激活是自然的'。技术路线是:首先从LLM(如Llama1B或Llama8B)的中间残差流层提取十亿级激活向量,然后训练一个基于MLP的扩散模型(采用流匹配目标),使其能从纯噪声生成与真实激活不可区分的向量。训练完成后,这个模型有两重用途:一是作为先验——将干预后的激活通过扩散采样投影回自然流形;二是作为编码器——其内部神经元(元神经元)自动隔离出可解释的概念,用于探测任务。
本文与已有方法最本质的区别在于:它放弃了所有关于激活结构的先验假设。SAE假设激活是稀疏线性组合,PCA假设主要结构是线性子空间,而GLP不对激活分布做任何参数化假设——它让数据自己说话,通过强大的扩散模型自然地学到激活流形的真实几何结构。这意味着GLP可以捕获非线性、多模态、高阶相关性等线性方法无法表达的结构。这一自由度带来了两个关键优势:第一,作为先验时,GLP能更精确地将激活投影回流形(而非投影到某个线性子空间),从而在保持语义内容的同时不牺牲生成质量;第二,作为编码器时,GLP的深层非线性MLP中的神经元比SAE的线性隐变量更擅长将概念隔离到单个单元中。此外,作者发现扩散损失(diffusion loss)可以作为一个统一的缩放指标:它随计算量呈幂律下降,并且直接预测下游任务(引导、探测)的性能。
方法步骤详情
GLP的完整流程分为训练和应用两个阶段。训练阶段:(1) 数据收集:使用vLLM和nnsight库高效提取LLM中间层的残差流激活,每个token位置一个向量,从FineWeb数据集采样10亿个token,排除句首特殊token;(2) 架构设计:采用类似Llama3的SwiGLU MLP块堆叠(而非注意力层),每层通过乘法调制注入扩散时间步 $t$ 的条件信息;(3) 训练目标:使用流匹配,前向过程 $z_t = (1-t)z_0 + t\epsilon$,训练去噪网络 $\hat{u}_\theta(z_t, t)$ 预测目标速度 $u = \epsilon - z_0$,损失函数为均方误差 $\mathcal{L} = \|\hat{u}_\theta - u\|^2$;(4) 数据管道:采用生产者-消费者模式,生产者将激活缓存到固定大小的缓冲区,消费者从缓冲区读取训练,解决十亿级数据的内存-速度权衡。应用阶段分为两种:先验应用(引导):给定一个被干预的激活,向其添加噪声($t_{\text{start}} = 0.5$),然后通过20步扩散采样将其去噪,相当于SDEdit的激活空间类比;编码器应用(探测):对输入激活加噪到特定时间步(如 $t = 0.1$),然后提取各MLP层SwiGLU门控的内部表征作为'元神经元'特征。
技术新颖性
GLP的技术新颖性体现在三个层面。第一,建模范式的新颖:已有元模型工作(如Schmidhuber 1992, Peebles 2022, Wang 2024)聚焦于生成网络权重,而本文选择建模激活——这避免了权重建模中排列不变性(permutation symmetry)的难题,以及数据生成需要昂贵优化的问题。第二,评估视角的新颖:已有元模型的价值在于生成的样本,而GLP的价值在于训练后的模型本身——它编码了训练分布的结构,可以作为先验或特征提取器。第三,技术贡献的新颖:扩散损失与下游任务性能之间存在可预测的幂律关系这一发现是全新的。在图像领域,扩散损失已被用作图像典型性度量(Li et al., 2023a; Siglidis et al., 2024),本文将其引入激活域,证明GLP的损失可以作为异常激活检测的指标。此外,与SAE不同,扩散目标不依赖于稀疏性假设,可以应用于任意架构,且训练过程更稳定。
实验结果
本文的实验结果可以从四个维度来分析。首先是生成质量:在Fréchet距离指标上,Llama1B上3.3B参数的GLP达到FD=0.53(下界为0.22),远优于SAE重构的1.99;Llama8B上3.4B参数的GLP达到FD=5.93,同样优于SAE的6.91。更令人印象深刻的是,在Delta LM Loss(LLM注入重构激活后的困惑度增加量)上,GLP以0.0513大幅超越SAE的0.1976(Llama8B-Base),这说明GLP生成的激活比SAE重构的激活更接近自然分布。第二是缩放行为:扩散损失随计算量呈平滑幂律下降 $L(C) = 0.52 + 435.1 \cdot C^{-0.169}$,不可约误差 $E=0.52$,每增加60倍计算量损失减半。关键发现是:这一缩放行为直接转移到下游任务——更好的GLP产生更好的引导和探测性能。第三是引导改进:在SAE特征引导任务上,GLP后处理将概念-流畅度的Pareto前沿向外推(图5);在人格激发(Persona elicitation)任务上,GLP在邪恶、谄媚、幻觉三种人格上均扩大了Pareto前沿(图6);在情感控制任务上,Llama1B上不同规模GLP的引导性能随计算量持续提升(图2b)。第四是探测能力:在113个二分类探测任务上,GLP元神经元在Llama1B上达到0.84 AUC,在Llama8B上达到0.87 AUC,显著超越SAE(0.70/0.76)、原始层输出(0.77/0.77)甚至原始MLP神经元(0.79/0.82)。特别值得注意的是,Llama1B的GLP(0.84)甚至超越了Llama8B的所有原始激活基线(最高0.82),表明GLP是一种超越单纯模型缩放的可解释性途径。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Fréchet距离(生成质量) | FD (↓) | GLP 3.3B: 0.53 (Llama1B) / GLP 3.4B: 5.93 (Llama8B) | SAE: 1.99 (Llama1B) / SAE: 6.91 (Llama8B) | Llama1B降低73.4%,Llama8B降低14.2% |
| Delta LM Loss(重构质量) | Delta LM Loss (↓) | GLP: 0.0513 (Llama8B-Base) / 0.0860 (Instruct) | SAE: 0.1976 (Base) / 0.2224 (Instruct) | Base降低74.0%,Instruct降低61.3% |
| 1-D探测(113个二分类任务) | Probe AUC (↑) | GLP: 0.84 (Llama1B) / 0.87 (Llama8B) | SAE: 0.70/0.76, Raw Layer: 0.77/0.77, MLP Neuron: 0.79/0.82 | 相对最佳基线提升6.3%/6.1% |
| 情感控制引导 | Concept & Fluency Mean (↑) | GLP+DiffMean: Pareto前沿外推 | DiffMean alone | 相同概念分数下流畅度提升,或相同流畅度下概念分数提升 |
局限与改进
作者明确承认了三个局限。第一,当前GLP独立建模单个token的激活,忽略了跨位置的序列结构——多token建模可能捕获位置间依赖关系,开启新的应用场景(如序列级别的引导)。第二,GLP是无条件模型,不以干净激活为条件输入——有条件建模(条件于噪声版本的干净激活)可能减少信息损失,改善引导效果。第三,本文只关注单层残差流激活,扩展到其他类型(如注意力头输出)或多层联合建模可能获得更丰富的表征。从我自己的观察来看,还有几个值得注意的问题:GLP的推理成本显著高于SAE——20步扩散采样的计算量远超SAE的一次前向传播,这在实时应用中可能是瓶颈;此外,论文的探测实验使用的是特定的113个二分类任务集,这些任务偏向于事实性知识和语言理解,对于更细粒度的语义概念(如情感强度、隐喻识别)的探测效果尚未验证;最后,GLP与SAE的Delta LM Loss比较中,GLP使用了插值加噪的'重构'方式(t_start=0.5),这与SAE的端到端重构在机制上不完全对等,可能影响结论的公平性。
独立分析的弱点
第一个弱点是推理效率:GLP的引导后处理需要20步扩散采样,每步都涉及一次完整的MLP前向传播(对于3.3B参数的GLP来说计算量不小)。在实际部署中,如果需要对每个token的激活都做后处理,这可能成为延迟瓶颈。改进方向包括:蒸馏到更小的模型、减少采样步数(论文显示4步已能覆盖大部分模式,见图3b)、或设计单步预测器。第二个弱点是条件建模的缺失:当前GLP是无条件生成模型,不以原始激活为条件。在引导场景中,我们知道原始激活是什么,但GLP无法利用这个信息,导致信息损失。改进方向是训练条件扩散模型,在引导时以加噪的干预激活为条件,直接预测去噪版本。第三个弱点是单一激活类型的局限:论文只建模了残差流输出,而Transformer中还有注意力输出、MLP输出等丰富的激活类型。不同激活类型的统计特性不同,联合建模或多类型建模可能揭示更丰富的内部结构。第四个弱点是缺乏与更多非线性可解释性方法的对比,如SelfIE、LatentQA等自然语言描述激活的方法——这些方法与GLP互补但竞争关系尚不清楚。
未来方向
作者提出了几个明确的未来方向,我也基于论文成果延伸一些思考。作者提出的方向:(1) 多token建模——捕获跨位置结构,可能用于序列级引导或异常检测;(2) 条件生成——以干净激活为条件减少信息损失;(3) 扩展到其他激活类型和多层建模。我的延伸思考:(1) 异常检测——图像领域已用扩散损失衡量图像典型性,GLP的损失可能类似地标识'异常'或'越狱'激活,这对安全应用很有价值;(2) 激活空间插值——扩散模型天然支持插值,可以在两个概念的激活之间平滑过渡,这可能用于渐进式引导或概念融合;(3) 与机制可解释性结合——GLP学到的元神经元可以作为电路分析的新工具,帮助理解信息如何在层间流动;(4) 跨模型迁移——探索在一个模型上训练的GLP能否迁移到其他模型(如从Llama-Base到Llama-Instruct的初步迁移已显示可行性);(5) 激活编辑——类似于图像扩散中的img2img技术,在激活空间中进行更精细的概念编辑。
复现评估
从复现角度来看,本文的条件较为友好。数据方面,作者使用FineWeb数据集(公开可用)采样10亿token,激活提取依赖vLLM和nnsight两个开源库,数据管道代码将开源。模型架构方面,GLP采用标准的Llama风格SwiGLU MLP块,没有复杂的定制组件。训练方面,最大规模的实验(3.3B参数模型)在单张A100 80GB GPU上训练5.6天,这个算力需求对学术实验室来说是可接受的(远低于训练同等规模的LLM)。超参数设置简单:单epoch训练,batch size 4096,学习率5e-5,余弦调度。主要的复现挑战在于:(1) 需要先有源LLM(Llama1B或Llama8B)来提取激活,虽然这些模型已公开但需要额外算力;(2) 十亿级激活的缓存和训练管道需要仔细的工程优化;(3) 评估涉及LLM-as-a-judge范式,可能引入评估模型的偏差。总体而言,论文的复现难度中等偏下,算力需求合理,代码将开源。
论文图表