Mechanism of feature learning in deep fully connected networks and kernel machines that recursively learn features

Contributions

这篇论文试图回答一个非常硬的问题：深度全连接网络/Deep Fully Connected Networks 到底是怎样学到特征的。作者提出 Deep Neural Feature Ansatz，核心断言是：网络在每一层学特征，本质上是在把那些对输出最敏感的方向持续放大，而这个“方向重要性”可以由 平均梯度外积/Average Gradient Outer Product/AGOP 来刻画。更具体地说，层内权重的 神经特征矩阵/Neural Feature Matrix/NFM，即 $W_i^{\top }{W}_i$，会与该层输入上的平均梯度外积成比例。论文用这个观点解释了全连接网络中的 simplicity bias、spurious feature、grokking、lottery ticket 等现象，并且把这一机制从神经网络里抽出来，构造出无需反向传播也能学特征的 Recursive Feature Machine/RFM。

真正让这篇论文进入 representation learning 主线的，不只是作者给了一个解释，而是他们把“好表征是怎样被选出来的”写成了一个可计算对象。与此同时，论文的边界也很明确：理论结果集中在 全连接网络，而且最强的证明依赖 零初始化、无限宽极限、Gradient Independence Ansatz/GIA 与 高斯输入 等条件；实践上最亮眼的性能提升主要出现在 tabular data，RFM 也主要对应输入层特征，而不是更深层的层级表征。

1. Introduction

神经网络为什么会比很多经典方法更强，一个常见回答是：因为它会自动学习并且利用任务相关的特征。但这句话往往只是在重复现象，而不是解释机制。机器学习研究的一个目标，就是识别这种神经特征学习/Neural Feature Learning 究竟通过什么机制发生，以及模型最终选择了哪些特征。理解了这一机制，就可以有机会设计出在可靠性和模型透明性上都更优的网络。

过去很多工作会把 Neural Feature Learning 描述为：网络内部的中间表征在训练过程中发生的变化，可这种说法并没有告诉我们：到底是哪一种数学对象在决定某些方向被放大、另一些方向被忽略。这篇论文的价值就在这里。作者不再把 feature learning 当成笼统的 representation drift，而是试图给出一个可验证、可计算、还能外移到其他模型族上的具体机制。

论文的直觉很直接。如果一个模型对某个输入方向非常敏感，那么沿着这个方向做微小扰动，就会显著改变预测结果。于是，真正值得被“学成特征”的，应该正是这些会强烈影响输出的方向。作者把这个想法形式化成 平均梯度外积：先看梯度告诉我们每个样本上哪些方向最影响输出，再把这些梯度外积在数据分布上求平均，于是得到一个全局的方向重要性矩阵。论文最核心的命题是，深度全连接网络在训练时，实际上就在逐层实现这个过程。

这一步一旦成立，feature learning 就不再只是神经网络内部的神秘副产品，而变成一个显式对象。你可以用它解释为什么网络会偏向简单但可能是伪相关的特征，也可以解释为什么某些稀疏子网络会更好，还可以反过来把这一机制移植到原本不会学特征的模型上。作者正是沿着这条路，把 kernel machine 改造成了会递归更新特征矩阵的 RFM，并在 tabular benchmark 上取得了非常强的结果。

从表征学习的角度看，这篇论文最重要的不是又发明了一个 pretext task，而是换了一个视角：表征的形成可以被看作模型对“输出敏感方向”的连续重加权。这会把 representation learning 从“设计一个看起来合理的目标函数”往前推进一步，进入“特征到底通过什么几何机制被选出来”的层面。

2. Problem Setup

作者考虑的是一个带有 $L$ 个隐藏层的全连接网络。记输入为 $x\in {\mathbb{R}}^d$，层间特征为 $h_i(x)$，其中 $h_1(x)=x$，而后续层通过权重矩阵 $W_i$ 和逐元素非线性 $\varphi$ 递推得到，也就是 $f(x)=W_{L+1}{h}_L(x)$，$h_l(x)=\varphi (W_{l-1}{h}_{l-1}(x))$。论文关心的不是某一层激活值本身，而是每一层权重矩阵诱导出的 神经特征矩阵/Neural Feature Matrix/NFM：

$$N_i\triangleq W_i^{\top }{W}_i.\text{\hspace{1em}(1)}$$

这个对象之所以重要，是因为它决定了该层如何对输入特征做缩放与旋转。从线性代数上看，$N_i$ 的特征值与特征向量刻画了哪些方向被放大、哪些方向被压缩，因此它正好提供了一个观察 feature selection 几何结构的窗口。

接着，作者引入另一个对象：梯度外积/Gradient Outer Product。对于标量输出预测器 $f$，在样本 $x$ 上的梯度 $\nabla f(x)$ 描述的是，沿不同输入方向做无穷小扰动时，预测会怎样变化。于是，单样本上的方向重要性可以写成

$$G(x)\triangleq \nabla f(x)\nabla f(x{)}^{\top }.\text{\hspace{1em}(2)}$$

如果输出是向量，论文则改用 Jacobian Gram 矩阵 $J_f(x)J_f(x{)}^{\top }$。关键在于，这个对象不是在问哪一维值大，而是在问哪一个方向对输出的局部影响大。

Deep Neural Feature Ansatz

对于第 $i$ 层，记 $h_i(x_p)$ 是样本 $x_p$ 进入该层时的表示，$f_i$ 是该层之后的子网络。作者提出：

$$W_i^{\top }{W}_i\propto \frac{1}{n}\sum _{p=1}^n\nabla f_i(h_i(x_p))\nabla f_i(h_i(x_p){)}^{\top }.\text{\hspace{1em}(3)}$$

也就是说，层内学到的特征矩阵，与该层输入上的平均梯度外积成比例。从直觉上说，网络会系统性地上调那些最能影响输出的方向。

这个设定把 feature learning 问题改写得很干净。过去我们常说网络“学到了更好的表示“，但这里作者给出的不是结果性描述，而是一个操作性命题：如果你能估计出跨样本平均的梯度外积，你就能近似恢复网络真正学到的特征方向。论文还特别指出，在第一层、标量输出的特例下，这个对象和统计学里的 Expected Gradient Outer Product 有直接关系，因此它并不是凭空出现的新构造，而是和经典监督降维工具相连的。

3. Algorithm / Methods / Model

3.1 Mechanism View: feature learning as reweighting influential directions

有了上面的设定以后，论文的核心机制就可以用一句话概括：网络通过训练逐步把那些更能改变输出的方向变成更强的特征方向。如果某个方向上的微小扰动会显著改变预测，那么它在梯度外积中的权重就会更大；而如果某个方向的平均梯度外积持续变大，对应的 NFM 也会在这个方向上变强。于是，feature learning 不再是抽象的“内部表示变化”，而是一个非常具体的几何过程。因此这个 ansatz 就可以将很多似乎不同的现象统一起来：无论是简单性偏置、伪特征，还是 lottery ticket、grokking，都可以看成是同一机制在不同数据分布和训练条件下的自然后果。

3.2 Recursive Feature Machines: exporting the mechanism out of neural nets

作者使用该 ansatz 揭示的表征学习的机制，提供了一个将特征学习整合进任意机器学习模型的算法框架。该 ansatz 的一个关键洞见在于：神经特征学习是通过平均梯度外积发生的，而平均梯度外积又是一个可以作用于任意函数的数学运算。正因为它具有这种普适性，我们就可以把它应用到任意机器学习模型上，从而赋予模型特征学习能力。

因此我们就可以使用一个迭代的两步策略：先训练一个任意的预测器，然后利用该预测器的平均梯度外积来直接学习特征。一个例子就是将经典的 Kernel Machines 改造成一个递归学习特征的模型。他们从带有可学习 Mahalanobis 度量的 Laplace kernel 出发：

$$K_M(x,z)=\exp \left(-\gamma \Vert x-z{\Vert }_M\right),\Vert x-z{\Vert }_M^2=(x-z{)}^{\top }M(x-z),\text{\hspace{1em}(4)}$$

其中 $M⪰0$ 是待学习的特征矩阵。然后在每一轮迭代里，先用当前 kernel 做 kernel regression，再用训练后 predictor 的平均梯度外积来更新 $M$：

$$M^{(t+1)}=\frac{1}{n}\sum _{x\in X}\nabla f_t(x)\nabla f_t(x{)}^{\top }.\text{\hspace{1em}(5)}$$

这就是 Recursive Feature Machine/RFM。初始化时 $M=I$，也就是普通的各向同性 kernel；迭代若干轮之后，kernel 中的距离不再是固定欧氏几何，而是逐步变成与任务相关的、数据自适应的特征几何。这个设计很关键，因为它说明作者并不是在说“神经网络神奇地做到了某件事”，而是在说：平均梯度外积本身就是一个独立的 feature learning primitive。

Tip

RFM 的本质不是“又一个 kernel trick”，而是把 representation learning 写成了 predictor learning 与 geometry update 的交替过程。传统 kernel 方法只学习函数；RFM 同时让 kernel 的几何也随着任务不断重写。

3.3 Theoretical evidence and its conditions

理论部分不是为了证明所有真实神经网络都严格满足 ansatz，而是为了说明：在若干重要极限和结构条件下，这个命题确实可以被推出来，而不是纯经验猜想。

首先，作者在一个简化 setting 中给出 Proposition 1。考虑 $f(z)=g(Bz)$，只训练第一层线性映射 $B$，并且在单样本上做梯度下降。如果初始化满足 $B^{(0)}=0$，那么任意训练步 $t$ 都有

$$\nabla f_t(z)\nabla f_t(z{)}^{\top }\propto B^{(t)\top }{B}^{(t)}.\text{\hspace{1em}(6)}$$

这个结论非常关键，因为它点明了为什么零初始化会让 ansatz 更干净。只要初始权重不为零，$B^{(t)\top }{B}^{(t)}$ 就会混入初始几何结构，因而破坏和梯度外积之间的精确比例关系。作者据此把“feature learning 强不强”与“ansatz 成不成立”直接联系起来。

更一般的理论结果是 Theorem 1。在一个带 ReLU 非线性的深度全连接网络里，如果第一层权重固定、其余层随机独立初始化并在无限宽极限下取期望，同时输入满足高斯分布，并使用 Gradient Independence Ansatz/GIA，那么第一层的 NFM 与输入梯度外积的期望相等到一个尺度因子。论文进一步指出，这个结论可以按层递归使用，从而为深层网络的 layer-wise 版本提供理论支持。

Theorem 1 的实际含义

这一定理并不是在说“任何实际训练出的有限宽网络都被完全刻画了”，而是在说：把第一层学到的特征矩阵解释成输入梯度外积，并不是拍脑袋的经验拟合，它在深度 ReLU 全连接网络的一个经典分析极限下有严格依据。

Why zero initialization matters

论文的理论线索反复强调零初始化，并不是偶然。对于单层可训练映射，训练后 $B^{(t)\top }{B}^{(t)}$ 可以被分成“由训练诱导出的部分”和“由初始化遗留的部分”。当 $B^{(0)}=0$ 时，后者直接消失，于是特征矩阵更纯粹地反映了模型在任务上真正放大的方向；当初始化不为零时，NFM 会混入与任务无关的初始几何，因此与梯度外积的对应变弱。论文把这一点同时用于解释为什么近零初始化更有利于 feature learning，也用于区分真正的 feature learner 与固定随机特征模型。

4. Experiments

4.1 Empirical verification of the ansatz

论文首先验证最核心的经验命题：训练后的 NFM 是否真的和平均梯度外积对齐。作者在 121 个 tabular classification 任务上训练五层隐藏层、每层宽度 1024 的 ReLU 全连接网络，并额外在 CelebA 与 SVHN 的 6 个图像分类任务上做验证。结果相当强：在 121 个 tabular 任务上，训练后 NFM 与平均梯度外积的 Pearson 相关系数中位数超过 0.85，而且系统性地高于“训练后 NFM 与初始化 NFM”的相关性；在 CelebA 上，第一层 NFM 与平均梯度外积的相关系数甚至能超过 0.97。这说明作者抓到的不是一个松散趋势，而是一个相当稳定的对应关系。

这组实验最有说服力的地方在于，它不只是报一个最终精度，而是直接比较了训练后真正学到的特征几何与一个外部可计算对象之间的吻合程度。对表征学习来说，这类证据比“某方法精度更高”更有解释力，因为它真的在回答“网络学到了什么结构”。

4.2 What the ansatz explains: simplicity bias, spurious features, lottery tickets, grokking

接下来，作者用同一套机制去解释几类经常被分开讨论的现象。

关于 Simplicity Bias/简单性偏置，其指的是神经网络倾向于使用最简单的可用特征来进行预测，即使有多个特征对类别标签同样具有指示作用。简单性偏置的一个结果，就是 Spurious Feature/伪特征 的出现。所谓伪特征，是指那些与预测目标相关，但并不一定与目标存在因果关系的模式。很多时候，这些伪特征比那些我们认为真正具有因果预测意义的模式更加简单。由于它们与标签高度相关，扰动这些简单特征或伪特征，往往会比扰动其他可用特征，包括那些与预测目标有因果关系的特征，更强烈地改变训练后模型的输出。

他们在 CIFAR10 与 MNIST 拼接图像上训练分类器，结果模型优先依赖 digit 区域而不是更复杂的自然图像部分，扰动 digit 区域的像素，会导致预测发生最大的变化；关于伪特征，他们在特别地在角落人为加入星形图案的图像上训练网络，对第一层 NFM 的对角线以及平均梯度外积进行可视化后可以看到，发现模型几乎完全依赖这个伪相关 pattern 做判断。

更有意思的是，在 CelebA 的 lipstick 分类任务里，NFM 与平均梯度外积都显示模型异常依赖眼睛区域而不是嘴唇，随后通过 mask test 进一步证实：遮住嘴唇只会轻微降低精度，而遮住被 ansatz 标出的眼部区域会显著打击性能。

Lottery Ticket 指的是，一个随机初始化的神经网络中，存在一个子网络；如果将这个子网络单独拿出来训练，它的性能可以达到甚至超过完整训练后的原网络。通常，人们通过剪去那些绝对值最小的权重来找到这样的子网络。本文所识别出的特征矩阵的稀疏性，为这一假说提供了直接证据。作者观察到第一层特征矩阵往往非常稀疏，于是直接按 NFM 中最重要的坐标阈值化输入，再重新训练相同架构的网络。以 CelebA 为例，他们根据 NFM 只保留约 2% 的关键像素，剪掉 98% 的输入维度，重训练后的性能反而更好。这提供了一种很不一样的解释：所谓 Lottery Ticket 的子网络不一定只是在 weight space 里稀疏，也可能意味着模型已经在 feature space 里识别出少量关键方向。

Grokking 也被放进同一机制下解释，其原始定义是：当深层网络在训练准确率已经达到 100% 之后继续训练时，测试准确率却会出现戏剧性的上升。作者构造了一个修改版 STL-10 任务：训练集很小、类别极度不平衡，并在人为添加的星形像素中埋入标签信号。网络可以很快达到接近 100% 的训练准确率，但测试准确率先停留在约 80%，随后在继续训练时突然升到 99.38%。作者给出的解释是：早期模型已经能记住训练集，但还没把真正泛化所需的简单特征方向稳定写进 NFM；随着训练继续，平均梯度外积开始持续放大星形像素对应方向，于是测试性能才突然跃迁。

4.3 Recursive Feature Machines on tabular benchmarks

如果说前两组实验是在证明 ansatz 有解释力，那么 RFM 实验是在证明这个机制有工程价值。作者在 Fernández-Delgado 等人的 121 个 tabular classification 数据集 benchmark 上，把 RFM 与 179 种方法做比较，包括传统树模型、kernel 方法、神经网络，以及使用 Neural Tangent Kernel/NTK 的 kernel ridge regression。RFM 在主要指标上都排得非常靠前：平均准确率 85.37%，P90 92.56%，P95 85.96%，PMA 97.36 ± 4.04，Friedman rank 17.79。作为直接对照，不带 feature learning 的 Laplace ridge regression 平均准确率是 83.76%，Friedman rank 是 28.48，说明收益确实来自递归特征更新，而不是仅仅来自某个幸运 kernel。

更值得注意的是训练代价。论文报告 RFM 在这 121 个任务上跑完只需要大约 40 分钟，而全连接神经网络基线需要大约 5 小时，两者都在同样的双 Titan Xp GPU 服务器上测量。作者还在第二个 tabular benchmark 上把 RFM 和 tabular transformer、ResNet、gradient boosting tree 一起比较，结论仍然是：RFM 通常以更低计算成本达到或超过这些方法的性能。

4.4 Grounded critical reading

这篇论文的实验设计整体是扎实的，因为它没有把“解释机制”与“做出更强模型”拆开说，而是两边都给了证据。不过边界也同样清楚。第一，最强的经验验证和最亮眼的实用收益都集中在 全连接网络与 tabular 数据；这并不能自动推出同样的机制在卷积网络、Transformer 或多模态 foundation model 里同样充分成立。第二，关于 spurious feature、grokking、lottery ticket 的展示虽然很形象，但很多任务带有明显的构造性，它们更像是机制演示而不是最终生态有效性证明。第三，RFM 成功的核心仍然是输入层几何重写，而不是对深层层级表征的直接递归建模，所以它更像是在说明“第一层 feature learning 已经很值钱”，而不是已经完整取代深网络的全部表征优势。

理论和实验之间也还有一条必须记住的缝。理论部分依赖无限宽、GIA、高斯输入等条件；实验部分则是有限宽、真实优化器、真实数据。这并不推翻作者结论，但意味着我们最好把论文理解成：他们发现了一个解释力很强、可迁移性很高的机制性近似，而不是已经给出一个对所有深度网络一锤定音的闭式理论。

5. Related Work & Future Work

从相关工作的组织来看，这篇论文其实把几条过去分散的线重新绑到了一起。第一条线是 监督降维/Supervised Dimension Reduction 与 Expected Gradient Outer Product。统计学里早就有人用目标函数的梯度来识别相关坐标或低维子空间，特别是在 multi-index model 与 sufficient dimension reduction 语境里。论文的一个重要价值，就是把这类经典对象接到了现代深度网络的 feature learning 机制上。

第二条线是 metric learning。一旦把 NFM 或 RFM 里的特征矩阵看成一个 Mahalanobis 几何，整个问题就和“学习什么距离度量最适合任务”直接连上了。作者也明确指出，这件事还和 manifold learning、非线性降维、早期 radial basis network 的度量构造有关系。换句话说，这篇论文虽然在讲 neural feature learning，但它真正动到的是更普遍的“任务相关几何如何被学习”。

第三条线是 NTK 与 after kernel。在很宽的网络里，模型会趋向线性化并表现得像固定 kernel regression；但论文恰恰强调，真正的 feature learning 出现在没有完全退化到固定 kernel 的 regime 中。作者因此提出一个很自然的未来方向：既然 RFM 学到的特征和神经网络第一层特征高度相似，它是否能作为一种更便宜的方式，逼近训练后经验 NTK 或所谓 after kernel。

未来工作基本沿着三个方向展开。其一，是把这种“先学 predictor，再直接用平均梯度外积更新几何”的思路推广到 CNN、GNN、Transformer 等更现代架构。其二，是认真回答 深层 feature learning 何时真正必要，因为本文的 RFM 主要抓的是输入层特征，但已经在 tabular 上很强。其三，是继续把 feature learning 的解释从经验相关推进到更强理论，把有限宽网络、真实初始化与真实数据分布下的误差项讲清楚。只要这些问题继续往前推进，这篇论文提供的视角就不会只是一个漂亮的解释，而会变成一种更通用的 representation learning 设计原则。