🤖 machine learning

Semantic Optimal Transport for Sparse Autoencoder Feature Matching and Circuit Compression

本文提出了一种基于 Wasserstein 距离的统一分布框架，将稀疏自编码器特征表示为激活加权的分布，以实现跨层的鲁棒语义匹配，并将特征电路自动压缩为可解释的超节点。

原作者： Tue M. Cao, Nguyen Do, My T. Thai

发布于 2026-05-28

📖 1 分钟阅读☕ 轻松阅读

原作者： Tue M. Cao, Nguyen Do, My T. Thai

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是论文《用于稀疏自编码器特征匹配与电路压缩的语义最优传输》的解释，已用通俗易懂的语言并辅以生动的类比进行翻译。

宏观图景：解码 AI 的“大脑”

想象一个大型语言模型（就像你正在聊天的 AI）是一座巨大的多层图书馆。在这座图书馆里，信息并非整齐地存放在书籍中，而是散落在数百万个微小的、发光的尘埃颗粒中，这些颗粒被称为特征。这些特征是 AI 思维的基本构建块。

为了理解 AI 如何思考，研究人员使用一种称为**稀疏自编码器（SAE）**的工具。可以将 SAE 想象成一台高倍显微镜，让我们能够看见这些尘埃颗粒。然而，两个大问题使得这项工作变得困难：

“同一概念，不同楼层”问题：同一个概念（如“正义”或“数字相加”）可能由 1 楼的一个尘埃颗粒表示，而由 50 楼一个外观完全不同的尘埃颗粒表示。目前的工具很难意识到这两个是同一个概念，因为它们各自所在的“房间”里看起来截然不同。
“房间太多”问题：当我们追踪一个特定的思维（即一个“电路”）时，往往会得到一个由数百个尘埃颗粒组成的纠缠网络。人类不可能阅读每一个颗粒。我们需要将它们分组为“超节点”（就像将所有“厨房”物品归为一类），但目前这需要人工标记每一个物品，无法扩展。

旧方法：比较静态快照

以前，研究人员试图通过拍摄每个特征的单一“快照”来解决这个问题。想象给一个尘埃颗粒拍张照片，测量它的颜色和亮度。如果 1 楼的照片与 50 楼的照片看起来相似，他们就假设这是同一个特征。

缺陷：这就像试图仅通过观察一个人模糊的阴影照片来识别他。它忽略了上下文。特征不仅仅是一个静态的点；它是成千上万句不同的句子中，该特征何时以及以多强的强度被激活的模式。旧方法丢弃了这种丰富的上下文，导致错误，特别是在比较图书馆中相距甚远的楼层时。

新方案：“人群地图”与“搬家卡车”

作者提出了一种利用最优传输（Optimal Transport）来观察这些特征的新方法。这是一个听起来很花哨的数学概念，但实际上非常直观。

1. 从单点到人群地图

新方法不再拍摄特征的单一照片，而是创建一张**“人群地图”**。

想象一个特征是一位名人。
旧方法：你只测量这位名人的身高。
新方法：你绘制一张每一位前来观看的粉丝的地图，并标记他们确切站的位置以及他们的兴奋程度（即他们的“激活权重”）。
这张地图捕捉了特征的上下文。它知道这个特征在人们谈论“猫”时会亮起，但在谈论“狗”时不会。

2. 共享参考空间（中立地带）

由于 1 楼和 50 楼的布局不同（不同的“流形”），你无法直接比较它们的地图。

作者将所有这些粉丝地图投影到一个共享参考空间中。想象这是一个巨大的、中立的市公园。
他们将来自 1 楼的粉丝和来自 50 楼的粉丝都放置在这张相同的公园地图上。现在，他们处于同一个街区，变得可比较了。

3. 最优传输（搬家卡车）

那么，如何衡量两个特征的相似度呢？

想象你有一堆沙子（代表粉丝）表示特征 A，另一堆沙子表示特征 B。
最优传输就像雇佣一辆搬家卡车，将沙子从 A 堆移动到 B 堆。
“成本”是沙子需要移动的距离。
如果 A 堆中的粉丝站在与 B 堆中的粉丝完全相同的位置，卡车就不需要把它们移动很远。成本很低。低成本 = 高相似度。
如果粉丝位于公园完全不同的区域，卡车就必须行驶很长的距离。高成本 = 含义不同。

这种方法之所以强大，是因为它观察的是活动的整体分布，而不仅仅是单一点。它能够区分乍一看相似但具有不同“粉丝模式”的两个特征。

他们的成就

利用这种“人群地图”和“搬家卡车”的方法，该论文声称取得了三大胜利：

更精准的匹配：他们现在可以准确地匹配 AI 不同层之间的特征，即使这些层相距甚远。这就像识别出地面层的一种特定树木与屋顶上的树木属于同一物种，尽管由于风和光线的原因，它们看起来不同。
自动压缩：他们可以将数百个纠缠的特征自动分组为清晰、易懂的“超节点”。算法不再需要人工整理 500 个物品，而是根据“粉丝地图”的相似度将它们分组。
发现细微差异：他们成功识别出执行非常具体任务的特征，例如“将两个数字相加”。其他方法无法区分那些仅仅是在“做数学”的特征，但这种方法发现了具体的“数字相加”模式。

“为何有效”的保证

该论文还包含了数学证明（即“收据”），以展示为何此方法有效：

尺度不变性：无论特征是“响亮”（非常活跃）还是“安静”（活跃度较低），只要谁在听的模式相同，结果就不受影响。该方法忽略音量，专注于人群的形状。
稳定性：如果你添加一点点噪声（比如几个额外的粉丝随机出现），“搬家卡车”的成本不会剧烈变化。该方法具有鲁棒性。
可恢复性：如果两个特征之间的差异足够大，数学上保证该方法能找到正确的匹配，即使数据不完美。

总结

简而言之，这篇论文指出：“停止将 AI 特征视为单一的静态点。将它们视为动态的活动人群。通过在一张中立地图上，利用数学搬家系统来比较这些人群，我们可以自动理解 AI 思维如何在各层之间演变，并将复杂的电路简化为可读的摘要。”

技术摘要：用于稀疏自编码器特征匹配与电路压缩的语义最优传输

1. 问题定义

稀疏自编码器（SAEs）已成为解释大型语言模型（LLM）的主要工具，通过将潜在空间分解为可解释的特征来实现。然而，在扩展 SAE 分析时仍面临两个关键挑战：

跨层特征匹配：识别 LLM 不同层之间语义相似的特征。现有方法通常依赖于比较解码器向量。这种方法在匹配非连续层之间的特征时失效，因为解码器向量被训练用于重建不同的激活分布，导致直接向量比较不准确。
特征电路压缩：特征电路（连接跨层特征的因果图）通常包含数百个节点，使得人工检查不切实际。当前的压缩方法依赖于手动将语义相似的节点分组为“超节点”，该过程无法扩展。

作者指出的核心问题是估计位于不同激活流形上的特征之间的语义距离。现有方法将特征视为单个向量（例如解码器权重），这丢弃了特征激活的“位置”和“强度”的丰富上下文，导致信息丢失，并无法处理多流形比较。

2. 方法论

作者提出了一种分布框架，将特征比较重新框架化为**最优传输（OT）**问题。

2.1 作为分布的特征表示

该方法不再用单个解码器向量表示特征，而是将第 $\ell$ 层的每个特征 $i$ 表示为在其激活的隐藏状态上的激活加权经验分布。

给定一个令牌位置语料库 $T$ ，该方法选择特征 $i$ 激活最强的前 $K$ 个令牌位置。
每个选定的令牌 $t$ 被分配一个权重 $w_{i,t}$ ，与其激活幅度成正比。
特征被表示为分布 $\hat{\mu}_{i,T} = \sum w_{i,t} \delta_{x_t}$ ，其中 $x_t$ 是令牌 $t$ 处的隐藏状态。

2.2 共享参考空间

为了比较位于不同流形上的不同层的特征，该方法将这些特定于层的分布投影到一个共享参考空间 $\mathcal{R}$ 中。

对于跨层匹配，参考空间通常是目标层的隐藏状态空间。源层的激活被投影到目标层中相应的令牌位置。
对于电路压缩，参考空间是所有层隐藏状态的连接。

2.3 基于 Wasserstein 距离的语义距离

一旦投影到共享空间，两个特征之间的语义距离就使用Wasserstein 距离（地球移动距离）来计算。

距离 $W_c(\hat{\mu}, \hat{\nu})$ 衡量将一个特征的激活分布的概率质量传输到另一个特征的最小成本。
基础成本 $c(z, z')$ 衡量在共享空间中匹配两个上下文表示的代价。
这种方法捕捉了激活模式的几何结构，而不仅仅是点对点向量相似性。

2.4 下游任务

特征匹配：对于目标特征，该方法找到具有最小 Wasserstein 距离（最近邻）的源特征。
电路压缩：该方法将电路节点视为分布空间中的点，并应用聚类（例如凝聚聚类）将具有小成对 Wasserstein 距离的节点分组为超节点，以最小化簇内的语义距离。

3. 主要贡献与理论保证

本文为提出的框架提供了三个主要理论保证：

激活重缩放不变性：投影后的特征分布对激活值的全局缩放具有不变性。由于表示依赖于激活的相对排名和比例分布（Top-K 选择和归一化权重），激活幅度的任意标量偏移不会改变语义距离。
扰动下的稳定性：特征分布之间的 Wasserstein 距离在分布的平滑扰动下是稳定的。距离的变化受基础成本的 Lipschitz 常数以及扰动分布与原始分布之间的 Wasserstein 距离所限制。这确保了令牌选择或 SAE 重新训练中的噪声具有鲁棒性。
基于边界的恢复：作者建立了经验最近邻匹配和语义分组完美恢复真实（oracle）匹配的条件。如果“边界”（最佳匹配与次佳匹配之间的差距）超过估计误差的两倍，则经验决策保证是正确的。

4. 实验结果

该方法在基于 GPT-2 和 Gemma-2-2B 训练的 SAE 上进行了评估，与包括 SAE Match（解码器向量 L2）、Feature Flow（解码器向量余弦）、Attribution Patching 和基于 LLM 的选择在内的基线进行了比较。

特征匹配：所提出的方法在LLM 评估（语义相似度评分）以及连续层、中间到最后一层、以及第一层到最后一层设置中的匹配准确率方面，均优于所有基线。至关重要的是，它在下游指标（交叉熵损失和解释方差）上保持了有竞争力的性能，表明它识别的是具有相似功能角色的特征，而不仅仅是相似的激活令牌。它特别成功地在其他方法失败的微妙“数字加法”特征匹配中取得了成功。
电路压缩：与手动分组和其他基线相比，该方法在将特征电路压缩为直观的超节点方面实现了更高的准确率。它成功识别了模块化特征组（例如“归纳头”和“令牌检测”），这些组在不同的电路中始终执行相同的功能角色。
消融研究：实验证实，将特征表示为分布比将它们折叠为单个质心向量保留了更多的语义信息，即使它们被投影到同一流形上。

5. 意义与主张

本文声称提供了一种统一的语义度量，将跨层特征匹配和电路压缩作为同一基本问题（即估计跨流形的语义距离）的实例来解决。

可扩展性：通过将特征分组为超节点的过程自动化，该方法消除了对人类标注的依赖，使得大规模特征电路的分析成为可能。
鲁棒性：分布方法捕捉了特征激活的“上下文”，允许在特征跨层发生显著演变（分布不相似）时进行准确匹配，而基于向量的方法在此类场景中会失效。
可解释性：自动将复杂电路压缩为可解释的超节点的能力增强了对模型机制的理解，正如在“达拉斯”（地理）、“押韵”和“相反”电路的案例研究中所展示的那样。

作者将这项工作定位为迈向可扩展、自动化的机械可解释性的基础步骤，超越了以向量为中心的特征比较的局限性。