将 MLP 转换为闭式多项式

如何在少假设数据特征的前提下，把已训练 MLP 与 GLU 变成可用线性代数检查的多项式。

• 问题场景：机制可解释性里，MLP 神经元常在多样、看似无关的上下文上激活，即多义性（Elhage et al., 2022）。作者称现行做法是字典学习，用宽且稀疏的 MLP 对原激活最小化均方误差；SAE 能抽出特征（Huben et al., 2023; Paulo et al., 2024），但不清楚有多少来自模型、多少来自数据。
• 作者的假设：数据分布取低阶矩约束下的最大熵分布。已知均值与协方差时，这就是高斯 N(µ, Σ)。逼近在该分布上全局最小化 MSE，从而少依赖训练集的具体样本。
• 与二次基座的差别：Pearce et al. (2024) 从双线性层出发，用 interaction matrices 的特征分解做解释、转向和对抗。作者称那要求不常规的基座；本文改为逼近已经训好的 MLP 与 GLU。
• 本文给出的对象：任意次数多项式的闭式最小二乘最优逼近；R² 高时可视化线性与二次系数的特征分解；用逼近随训练步的拟合检验 distributional simplicity bias；并用线性系数做 SVD 输入干预，看准确率是否随逼近一起变化。

相关讨论在引言、Background 和 Section 4，没有单独的 Related Work。

多义性与稀疏字典

• Elhage et al. (2022)：MLP 神经元常是多义的。作者用这一点提出要学模型本身的特征，而不是只学数据上的字典。
• Huben et al. (2023) 与 Paulo et al. (2024)：用 SAE 从 MLP 激活和输出抽取可解释特征，后者还做大规模自动解释。作者称尚不清楚 SAE 学到的是模型特征还是数据特征。

二次层与谱方法

• Pearce et al. (2024)：双线性层每个输出单元是输入的二次型，对 Hessian（interaction matrices）做特征分解。作者转述，最大幅度特征向量常可直接查看，并用于转向和对抗。作者称该路线要从非常规的双线性或二次基座做起；本文的二次逼近带有可做同样谱分解的张量参数。
• Shazeer (2020)、Mnih and Hinton (2007)、Dauphin et al. (2017)：注意力块后常用 GLU；恒等激活给出双线性层。作者称其性能与 SwiGLU、ReGLU 相近。论文未把这些架构当作实验对照。

复杂度随训练上升

• Nakkiran et al. (2019)、Refinetti et al. (2023)、Belrose et al. (2024)：网络随 SGD 表示更复杂的函数，并按阶利用统计矩。作者把后两篇称为 distributional simplicity bias，并用本文的闭式线性、二次逼近去测。
• Teney et al. (2024)：随机网络往往简单。作者用它解释训练最初几百步 FVU 下降。

基线与数据

论文没有外部解释器或攻击方法作为实验基线。对照是同一 ReLU MLP 的线性与二次最小二乘逼近，以及 SVD 投影下原网络和两种逼近的测试准确率。数据是 MNIST（LeCun et al., 1998）。作者称高斯混合拟合 MNIST 后的样本，对人很难与真样本区分（Belrose et al., 2023, Figure 2）。

作者把本文定位为：在最大熵高斯上给已训练的 MLP/GLU 做闭式多项式逼近，而不是更换基座或训练 SAE。

在高斯或高斯混合输入下，闭式求出单隐层 MLP 与 GLU 的最小二乘多项式系数，而不是在任意分布上用 SGD 估系数。作者给出的理由是：SGD 计算贵，且不利于理解归纳偏置。先处理单个高斯，Section 3.4 再推广到混合分布。

线性系数

• 逼近取仿射形式。斜率的最小二乘解为 β=Cov[x]⁻¹Cov[f(x),x]。截距 α 等于 E[f(x)] 减去 β 的转置与 E[x] 的乘积，对应式 (5)。
• 网络写成 f(x)=W₂ ϕ(W₁x+b₁)+b₂，ϕ 逐元作用。x 为高斯时预激活 y 仍为高斯。ReLU、GELU 的 E[ϕ(y)] 有解析式（附录 A）；其他激活用 Gauss-Hermite 求积。再由期望的线性性得到 E[f(x)]。
• 交叉协方差化为 E[ϕ(y)x^T]。x 与 y 联合高斯，Stein 引理给出 𝔼[ϕ(y_i)x_j]=Cov(y_i,x_j)𝔼[ϕ′(y_i)]。ReLU 与 GELU 的导数在高斯下的期望也可闭式计算。

二次与任意次数

• 二次特征映射收集 x_i x_j（i≤j），与 x 拼成 z。x 的任意二次函数都是 z 的仿射函数，系数仍由 OLS 估计。同一替换可用于任意次数。
• 需要 Cov[f(x), z] 中的 E[ϕ(y_i) x_k x_l]，以及 z 的高阶矩。高阶矩用 Isserlis 定理对非中心高斯的推广：按单点集与配对划分求和。
• Theorem 3.1 把 E[g(X) 与若干 Y_i 的乘积] 写成一元积分 E[g(X) X^k] 的线性组合。系数由 Y_i 对 X 的 OLS 截距、斜率和残差协方差解析得到。n=2 时论文给出三项简化式。一元积分再用数值积分或附录 A 的闭式。

混合分布与 GLU

• 分量指标为类别分布时，交叉协方差等于条件协方差的期望，加上条件均值之间的交叉协方差。对输入自身使用该式得到输入协方差；把其中一项换成 f(X)，得到输入与 MLP 输出的交叉协方差。
• GLU 是 ϕ(xW+b) 与 (xV+c) 的逐元积。高斯输入下两路预激活联合高斯。均值的每个分量是 E[ϕ(Y_i) Z_i]，用 Stein 引理；与输入的交叉协方差用 Theorem 3.1。作者称更高次多项式可同样处理。

实验中的拟合取舍

• 实验网络是单隐层 ReLU MLP，形式同式 (1)。训练用 schedule-free AdamW，warmup 1K 步，batch size 64，weight decay 0.1，对数间隔存检查点。
• 线性逼近按与 MNIST 各类均值、协方差一致的高斯混合做闭式拟合。二次的混合分布解析解在 MNIST 上不可行：二次特征维数为 d(d+1)/2，d=784 时为 307,720，协方差有 9.47×10^10 项，常用线性求解为 O(d^6)。x 为 N(0,1) 时该协方差为对角阵，元素在 {1, 2}。
• 因此二次系数先在 N(0,1) 下闭式计算，再用 schedule-free SGD 在高斯混合样本上微调。作者称目标凸，故这些系数应很好地近似真最小二乘值，并称 Figure 2 中二次 FVU 的噪声来自这一步。

谱分解与输入干预

• 对线性系数做 SVD，和式写到 10 项，U 把输入维映射到 R^10。对测试集每个样本施加同一正交投影，减去前 k 个左奇异向量张成的成分。
• 线性与二次网络也可用 SVD 或特征分解拆出重要性项。作者称，若攻击迁到 ReLU 网络，就说明逼近抓住了原网络的分布外行为。效果用测试集准确率衡量，不用单独的危害性评审。

实验设置

• 模型与训练：一个单隐层 ReLU MLP，在 MNIST 上用 schedule-free AdamW 训练。warmup 1K 步，batch size 64，weight decay 0.1，对数间隔存检查点。Figure 1 另有一次更长训练。论文未报告隐层宽度、学习率、重复次数和测试集条数。
• 拟合分布：与 MNIST 各类均值、协方差一致的高斯混合。作者称该混合的样本对人很难与真样本区分（Belrose et al., 2023, Figure 2）。
• 两种逼近：线性逼近在该混合下闭式拟合。二次逼近先在 N(0,1) 下闭式求解，再在混合样本上用 schedule-free SGD 微调（脚注 3，附录 B）。
• 指标：FVU、R²，以及逼近与网络输出的 KL（nats）。输入干预用测试集准确率。论文未使用 LLM 评审，也未报告人工一致性。
• 图轴：正文与 Figure 2 图注说的是未解释方差比例；图中纵轴文字是「Fraction of variance explained (FVU)」。Figure 3 纵轴为 KL divergence (nats)。Figure 5 纵轴为 Accuracy (%)，横轴为 SVD Components Ablated，图例为 ReluMLP、Linear、Quad 三种 SVD Attack。
• 对照：没有外部基线。对照是同一网络的线性逼近、二次逼近，以及原来的 ReLU MLP。

主结果

论文没有数值总表。摘要、结论和 Section 4.2 里的原数如下。

Swipe horizontally to read the table

量	条件	数值	出处
二次逼近解释的输出方差	训练结束；结论写明 MLP 在 MNIST 上训练	over 95%；结论写 well over 95%	摘要；Conclusion
线性 FVU 与 KL 上升的步数	模仿 MNIST 训练集的高斯混合	500 到 1K 步起，持续到约 4K 步	Section 4.2

• 作者把线性拟合变差、二次 FVU 近乎不动的这一段称为二次阶段，并认为网络从线性特征转到大体二次的非线性特征，与 DSB 大体一致。
• Section 4 的预期还包括：线性 R² 下降之后，二次 R² 也近单调下降。正文对该时段的报告是二次 FVU 近乎不变，没有给出二次 R² 下降的数字。最初几百步两种 FVU 都下降，因此从训练起点起也不是单调变差。
• 作者称早期 FVU 下降有些出乎意料，并认为随机初始化往往简单，SGD 可能先去掉噪声。Figure 3 的 KL 没有这一早期下降，其余趋势与 FVU 图相似。论文未把 over 95% 对应到 Figure 2 的某个读数。

特征可视化

• Figure 1 画类别「3」的二次顶特征向量和该类的线性系数，列为 2^4、2^6、2^8、2^10、2^12、2^14、2^16、2^18。各列上方标有 0.01、0.01、0.01、0.01、0.04、0.10、0.13、0.20；正文没有单独定义这些数。
• 图注称可辨的「3」从噪声中出现，与学习发生的区间一致；线性结构在 FVU 最小时最直观，该步编号在文本中未完整印出。图注称二次特征更晚成形，最末期出现视觉伪影。
• Section 4.3 写，线性与二次特征在 2^12 与 2^14 步之间最直观，之后模型开始过拟合。这是作者对图像的解读；文本转换看不到像素。

SVD 输入干预

• 对测试集全部样本施加同一投影，去掉线性系数 SVD 的前 k 个分量，k 从 1 到 10。同时记录 ReLU MLP、线性逼近、二次逼近的测试准确率。Figure 4 上行是对「3」从 1 到 10 个分量的干预，下行是十个分量都去掉后的各个 digit。
• 作者称原网络与两种逼近的准确率同步下降（图注用 perfect lockstep）。消融 4 个分量后准确率已低于 50%。去掉全部 10 个分量，即 784 维中的 10 维，作者称 MLP 不优于随机。图注称这些图像仍与原图高度相似、人仍能辨认，但 MLP 不能分类。
• 作者把这看作线性逼近抓住了对原模型有因果作用的输入特征，并捕获了分布外行为。结论里作者称这些对抗样本可被证明有效。论文没有给出每个 k 的准确率百分点。

其他消融与分析

• 二次闭式解：d=784 时 z 维数 307,720，Cov(z) 有 9.47×10^10 项，时间复杂度 O(d^6)；N(0,1) 时 Cov(z) 对角，元素属于 {1, 2}（附录 B）。
• 二次 FVU 曲线噪声：作者归因于先 N(0,1) 再 SGD 微调，而不是混合分布的全解析解（脚注 3）。
• FVU 与 KL：前几百步 FVU 下降，KL 不下降（Figure 2、Figure 3）。
• 论文未报告隐层宽度、学习率、重复次数，也未报告 GLU 或 GELU 网络上的 FVU 与准确率。

作者指出的局限与后续方向

• 分布假设：解析推导假定输入服从某种高斯混合；作者称不清楚能否有用地推广到其他分布（Limitations）。
• 适用性：作者认为高斯性可以避免逼近过拟合训练集的具体性质，但不一定适合所有应用（Limitations）。
• 非高斯时的用法：若高斯性过严，作者建议用解析多项式初始化梯度优化器，再在目标分布样本上微调（Limitations）。
• 次数与规模：n 次逼近要物化含 d k^n 个参数的张量（k 为输入维，d 为神经元数）。作者称中间激活维数到数万时，能解析计算的最高次数是 n=2；若最好的二次逼近较差，单靠本方法难以解释模型（Limitations）。
• 更深的模块：作者写本文聚焦简单图像分类，并认为同一逼近或可用于深层网络里的前馈模块，例如 transformer（Future Work）。
• 过完备基：作者认为对 transformer FFN 的二次逼近做特征分解，可给 d 维残差流 d² 个特征向量；该过完备基或许能以类似 SAE 的方式处理多义性，更复杂的张量分解或许更好。基向量或可用 unembedding 投到下一 token 预测空间查看（Future Work）。

实验覆盖范围

• 实验模型是在 MNIST 上训练的单隐层 ReLU MLP；附录 B 将输入维写作 784。论文未报告隐层宽度、学习率和重复次数。
• 线性与二次逼近的 FVU、KL 画在模仿 MNIST 训练集的高斯混合上（Figure 2、Figure 3），检查点为对数间隔。二次系数在 N(0,1) 闭式解上初始化后，用 schedule-free SGD 微调。
• 可视化是 Figure 1 中类别「3」的线性系数和二次逼近的顶特征向量，步数从 2^4 到 2^18。Section 4.3 写最可解释的是 2^12 到 2^14 步。
• 输入干预是对测试集统一施加线性系数的 SVD 投影，消融分量数从 1 到 10，并同时报告 ReLU MLP、线性逼近和二次逼近的准确率（Figure 4、Figure 5）。
• 附录 A 给出 ReLU 与 GELU 的高斯积分，Section 3.5 给出 GLU 的论证。数值结果只报告了该 ReLU MLP 的线性与二次逼近；论文未报告 GLU、GELU 或 transformer 上的 FVU、KL 或准确率。

闭式推导给出单隐层 MLP 与 GLU 在高斯或高斯混合下的最小二乘多项式，用来检查网络函数，而不是训练 SAE。

• 问题：作者称 SAE 分不清特征来自模型还是数据。他们把输入设为低阶矩约束下的最大熵分布，即已知均值与协方差时的高斯，并在该分布上最小化 MSE。
• 做法：线性斜率由输入协方差和网络–输入交叉协方差确定，交叉项用 Stein 引理。更高次通过单项式特征映射变成 OLS；Master Theorem 把所需多元积分拆成一元积分。混合分布用全协方差定律。MNIST 上二次的混合解析解不可行，故先在 N(0,1) 下求解，再用 SGD 微调。
• 拟合轨迹：被测对象是 MNIST 上的单隐层 ReLU MLP。摘要称训练结束时二次逼近解释超过 95% 的输出方差，结论写 well over 95%。500 到 1K 步线性 FVU 与 KL 上升，并持续到约 4K 步；同期二次 FVU 近乎不变。作者称之为二次阶段，并认为与 DSB 大体一致。最初几百步 FVU 下降，KL 没有这一下降。
• 输入干预：去掉线性系数的 4 个 SVD 分量后，测试准确率低于 50%；去掉 10/784 维后，作者称 MLP 不优于随机，且与两种逼近同步下降。作者称这表示线性逼近抓住了对原模型有因果作用的特征。
• 作者的收束：拟合方差高时，谱分解上的观察可以迁到原 ReLU 网络。作者同时指出，推导依赖高斯混合，高维时解析次数止于二次。