跳转至

Mechanistic Interpretability with Sparse Autoencoder Neural Operators

会议: NeurIPS 2026
arXiv: 2509.03738
领域: 可解释性 / Interpretability
关键词: 稀疏自编码器, 神经算子, 机制可解释性, 傅里叶神经算子, 泛函表征假设

一句话总结

提出稀疏自编码器神经算子(SAE-NO),将概念从欧氏空间标量激活拓展为连续函数空间中的算子参数化,通过概念稀疏性与域内稀疏性解耦表征表达位置与强度,实现了分辨率无关的可解释性表征与跨尺度泛化。

研究背景与动机

稀疏自编码器(SAE)作为机制可解释性的核心分析工具,广泛用于在大语言模型和视觉模型中分解重叠表征并提取单义性特征。现有的主流 SAE 体系高度依赖线性表征假设(Linear Representation Hypothesis, LRH),采用浅层 MLP 编码器与线性字典解码器,将每个潜在概念严格建模为欧氏空间中的固定维度向量,并通过单一标量激活来表征该概念的存在性及其强度。

然而,这种基于固定离散化和欧氏向量的参数化方式在处理具备连续空间、时间或场结构的数据(如图像、物理场模拟信号、时空动力学等)时暴露出根本局限。标量激活只能编码概念“是否存在”,无法表达概念在空间域中“如何表达”以及“在何处表达”;在面对图像平移或时间演化时,传统 SAE 无法复用相同的结构概念,而被迫为不同位置的相同模式分配完全独立的正交特征字典项。此外,固定网格维度的绑死使得传统 SAE 无法外推到未见过的分辨率或连续物理域,在遭遇高频频谱扰动时重建极易崩溃。

为了打破固定网格与欧氏向量表征的束缚,本文引入映射于无穷维函数空间之间的神经算子(Neural Operators),提出泛函表征假设(Functional Representation Hypothesis, FRH)。核心 idea:将稀疏概念参数化为傅里叶积分算子(SAE-FNO),通过“概念稀疏性”与“域内稀疏性”的联合稀疏约束,将概念建模从单纯的存在性检测拓展为函数空间中对概念表达位置与形态的结构化解耦建模。

方法详解

整体框架

SAE-NO 的整体目标是将连续函数空间 \(\mathcal{X}\) 中的输入数据 \(x\) 分解为无穷维函数空间 \(\mathcal{Z}\) 上的稀疏特征图集合。在具体实现中,本文基于傅里叶神经算子(FNO)实例化为 SAE-FNO,并可结合高维升维(Lifting)与降维投影(Projection)算子提升优化效率。

整个处理流程包括升维映射、频域卷积编码、联合稀疏化投影以及频域卷积解码四个核心步骤:首先通过线性升维算子将输入映射至通道升维空间;编码器在频域中执行截断模式下的核积分变换获得密集特征图;接着依次施加概念稀疏性(挑选激活概念)与域内稀疏性(约束激活区域);最后解码器通过傅里叶逆变换与局部投影算子将稀疏函数表征重构回原始物理域。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入数据 x"] --> L["1. 升维预条件算子<br/>Lifting 变换至高维通道空间"]
    L --> Enc["2. 傅里叶积分编码<br/>频域卷积变换获取稠密特征图"]
    Enc --> S1["3. 联合稀疏化约束<br/>概念群稀疏性: 筛选活跃特征图"]
    S1 --> S2["域内空间稀疏性: 约束局部表达位置"]
    S2 --> Dec["4. 连续函数空间解码<br/>傅里叶逆变换与逆投影重构"]
    Dec --> Out["重构数据 x_hat"]

关键设计

1. 泛函表征假设与傅里叶积分算子参数化:跳出固定离散化的向量表征约束 传统 SAE 将概念建模为离散字典矩阵 \(D \in \mathbb{R}^{m \times p}\),其概念基底与固定的网格离散化分辨率 \(m\) 深度绑定。本文提出的泛函表征假设(FRH)将生成过程推广为函数空间之间的线性积分算子 \(x = \mathcal{G}_{D^*}(z)\)。SAE-FNO 将解码器实现为傅里叶卷积积分算子: $$ \mathcal{G}D(z) = \mathcal{F}^{-1}\left(\sum z_c)\right) $$ 其中 }^p W_c \cdot (\mathcal{F\(W_c = \mathcal{F} D_c\) 为频域卷积核参数,\(z_c\) 为第 \(c\) 个概念对应的连续特征图。由于积分运算在频域中解耦为截断傅里叶模态上的逐点复数乘法,概念不再依赖于特定的像素采样网格,使得模型能够直接在任意不同分辨率下进行查询和推理,天然具备离散化不变性与抗高频噪声干扰能力。

2. 联合稀疏化机制:概念稀疏性与域内稀疏性解耦 针对标量激活无法表达空间结构的痛点,SAE-FNO 将隐表征扩展为特征图 \(z \in \mathbb{R}^{p \times H \times W}\),并引入联合稀疏化结构。第一层为概念稀疏性(Concept Sparsity),充当群稀疏化角色:根据各概念预激活特征图的全局 \(\ell_2\) 能量范数进行 Top-K 筛选,强制仅有少量核心概念处于激活状态: $$ v = \mathbf{1}_a \cdot \tilde{v}, \quad \text{其中} \quad a = \Pi\left(\left[|\tilde{v}_1|_2, |\tilde{v}_2|_2, \dots, |\tilde{v}_p|_2\right]^\top\right) $$ 第二层为域内稀疏性(Domain Sparsity):对被选中的每个活跃特征图 \(v_c\) 施加元素级空间稀疏投影 \(\mathcal{P}(v_c)\),限定该概念在空间或时间域内的具体表达范围。这种双层解耦使特征的几何结构(由域内稀疏性控制)与模型的整体表达容量(由概念稀疏度 \(k\) 控制)彼此分离,彻底解决了传统 SAE 在改变稀疏度 \(k\) 时特征结构发生剧烈畸变的问题。

3. 升维算子充当优化预条件子:加速概念恢复并平滑表征流形 受神经算子三阶段架构启发,论文引入了线性升维算子 \(L\)(映射到高维通道)与投影算子 \(P\)(重构回原始维度)。理论分析证明(Proposition 3.1),在梯度动力学下,升维和逆投影在原始参数空间中的有效更新方向等价于引入了对称正定阵 \(L^\top L\) 作为预条件子: $$ D_c^{(k+1)} = D_c^{(k)} + \frac{\eta_L}{M} \left((L^\top L)(x - \hat{x}^{(k)})\right) \star z_{c,t} $$ 当 \(P = L^\top\) 且 \(L^\top L = I\) 时,模型在推理与动力学上完全退化为标准 SAE-FNO;而在实际优化中,学习可自适应调节的升维算子 \(L\) 显著降低了字典原子间的有害互相关,平滑了损失曲面,使概念恢复的收敛速度得到大幅提升。

损失函数 / 训练策略

SAE-FNO 采用均方误差(MSE)作为数据重构目标: $$ \min_{W \in \mathcal{W}} \frac{1}{2} |x - \mathcal{F}^{-1}(W \cdot \mathcal{F} z)|_2^2 + \lambda \mathcal{R}(z) $$ 训练过程中使用 AdamW 优化器,学习率设为 \(10^{-3}\),并在通道维度执行单位归一化。为保持梯度和频域参数尺度的稳定,SAE-FNO 的编解码器权重除以空间特征尺寸的根号外,还按 \(\sqrt{\text{modes}}\) 进行重缩放。在特定扩展实验中,该架构还可无缝迁移至 BatchTopK 和 Matryoshka 多层级表征训练策略。

实验关键数据

主实验

在频域概念恢复任务(合成 1D 单频概念恢复实验)中,在无噪声(Clean)与高噪声(Noisy, \(\sigma_{\text{noise}}=0.5\))设定下,对比 SAE-MLP 与不同傅里叶截断模态下的 SAE-FNO 的互相关恢复分数(Cross-Correlation Recovery Score,5 次运行的均值与标准差):

初始化 / 噪声环境 SAE-MLP SAE-FNO (全频域模态) SAE-FNO (截断模态 m=151)
Clean (\(\sigma=0\)) (0.329, 0.210) (0.868, 0.303) (0.863, 0.324)
Noisy (\(\sigma=0.5\)) (0.329, 0.210) (0.798, 0.534) (0.853, 0.627)

在计算开销与吞吐量方面,评估在 L40S GPU 上不同数据分辨率(\(k=100, p=1000\))下的训练与推理吞吐量:

架构 数据分辨率 训练吞吐量 (img/s) 推理吞吐量 (img/s)
SAE-MLP \(16 \times 16\) 13,893 9,680
SAE-MLP \(32 \times 32\) 10,551 8,648
SAE-FNO \(16 \times 16\) 3,838 12,843
SAE-FNO \(32 \times 32\) 1,024 3,960

消融实验

在模型尺寸 \(p\)(字典概念数)扩展性测试中(输入尺寸 \(16 \times 16, k=200\)),考察系统吞吐量变化:

架构 字典规模 \(p\) 训练吞吐量 (img/s) 推理吞吐量 (img/s)
SAE-MLP 1,000 17,847 12,987
SAE-FNO 1,000 3,777 12,254
SAE-MLP 5,000 17,744 14,947
SAE-FNO 5,000 744 3,111
SAE-MLP 10,000 11,308 15,082
SAE-FNO 10,000 373 1,598

关于模型结构消融与泛化表现(综合正文定性与定量分析): - 域内稀疏性消融(Ablation of Domain Sparsity):去除域内稀疏性(仅保留概念级 Top-K)后,模型退化为在全域弥散的稠密特征,彻底丧失局部空间特征提取能力;完全去除稀疏性(普通 AE-FNO)则彻底退化为无结构高频噪声。 - 跨分辨率超分辨率泛化:在 \(28 \times 28\) 训练的模型直接推断到 \(2\times, 4\times, 8\times\) 上采样网格时,SAE-FNO 的重建均方误差与稀疏支撑集恢复精度保持极其平稳,而标准 SAE-MLP 与 SAE-CNN 随着分辨率提升误差发生灾难性爆炸。

关键发现

  • 平移不变性与时空概念复用:在移动数字(Moving MNIST)测试中,SAE-FNO 即使在完全无时序约束的 i.i.d. 训练下,依然能自发通过平移特征图内部的空间激活来追踪目标,帧间激活概念交并比(IoU)接近 1.0;而 SAE-MLP 频繁在完全不相关的字典原子间切换,IoU 迅速归零。
  • 概念表达效率(Capacity Utilization):SAE-FNO 的概念使用率曲线呈现极陡峭的幂律衰减,即仅使用极少量的核心功能算子即可覆盖大部分数据分布,在低字典规模下达成远高于 SAE-MLP 的 \(R^2\) 拟合优度。
  • 结构稳定性与解耦:通过 Hoyer 分数测定表明,SAE-FNO 学到的概念空间与频域局部性在不同的稀疏度 \(k\) 取值下保持完全一致,消除了传统方法中因调整 \(k\) 导致字典概念基底彻底改变的顽疾。

亮点与洞察

  • 将神经算子思想注入机制可解释性:打破了经典可解释性将神经元激活锁定为标量的思维定势,利用函数空间积分算子将“概念身份”与“概念的几何表达”优雅解耦。
  • 理论推导 Lifting 预条件机制:严格证明了高维升维算子 \(L\) 的梯度动力学效应相当于在字典空间施加各向同性预条件更新,从数学上解释了神经算子升维模块平滑损失曲面的原因。
  • 零额外时序代价实现时序稳定性:在完全不对时序关联建模的前提下,利用频域算子的卷积位移同变性,自发涌现出跨时间帧平移复用概念的能力,对动态系统的可解释性分析提供了重要启示。

局限与展望

  • 计算与内存开销明显增大:由于引入了 2D/3D 快速傅里叶变换及多通道特征图激活,SAE-FNO 在超大字典(\(p \ge 10\text{K}\))及高分辨率场景下的训练速度明显落后于矩阵乘法高度优化的 SAE-MLP。
  • 强烈的平滑谱归纳偏置:傅里叶参数化高度偏向于低频连续变化信号,在处理包含尖锐高频跳变、离散文本符号或表格数据等非光滑结构时难以适用。
  • 未来方向:将 SAE-FNO 从原始数据像素域推进至神经算子大模型(FNO/U-NO/DiffFNO 等科学基础模型)以及多模态基础模型的深层隐表征中,挖掘连续物理定律和时空演化的机制可解释性。

相关工作与启发

  • vs SAE-MLP (Anthropic / OpenAI 等传统方案):传统方法依赖线性表征假设,概念是欧氏离散向量,标量激活无法区分概念“在何处、如何”表达;SAE-FNO 扩展至函数空间,使用特征图与联合稀疏性统一建模局部性与分辨率无关性。
  • vs SAE-CNN:卷积 SAE 虽然引入了局部感受野,但仍然绑定在固定的像素网格与离散核尺寸上;SAE-FNO 具有无穷维连续算子表达能力,能在跨分辨率插值场景下自由外推。
  • vs 时序 SAE (Temporal SAEs):现有时间序列 SAE 通常需要人为设计时序平滑损失或循环训练结构;SAE-FNO 依靠连续积分卷积的天然物理性质,无需时序正则即可自发涌现时序概念一致性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出泛函表征假设,将神经算子与机制可解释性稀疏自编码器有机融合。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖从已知生成模型的理论恢复,到 MNIST/CIFAR/ImageNet 多数据集验证,并包含完备的超分辨率和噪声鲁棒性分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽,概念动机明确,图表表达极具说服力。
  • 价值: ⭐⭐⭐⭐⭐ 为科学机器学习(SciML)以及连续物理场模型的机制可解释性开辟了全新的研究范式。