📊 statistics

Closed-Form Last Layer Optimization

本文提出了一种闭式最后一层优化方法，将最后一层权重视为骨干网络参数的函数，从而构建一种交替优化方案，该方案在神经正切核机制下能够高效收敛，并在平方损失回归任务上优于标准的随机梯度下降和Adam算法。

原作者： Alexandre Galashov, Nathaël Da Costa, Liyuan Xu, Philipp Hennig, Arthur Gretton

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

原作者： Alexandre Galashov, Nathaël Da Costa, Liyuan Xu, Philipp Hennig, Arthur Gretton

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

想象一下，你正在教一个机器人画画。这个机器人主要有两个部分：

艺术家（骨干网络）：这部分负责学习观察世界，识别形状、颜色和纹理。它复杂、富有创造力，且需要很长时间来学习。
画家（最后一层）：这是最后一步，机器人决定具体要画哪些笔触，以匹配特定的目标图像。

旧方法：“逐步进行”

通常，在训练这些机器人时，我们使用一种称为**随机梯度下降（SGD）**的方法。这就像一名徒步者在雾中试图找到山谷的底部。

徒步者（计算机）迈出一小步，检查是否更低，然后再迈出一步。
他们同时对“艺术家”和“画家”都这样做。
问题在于？“画家”实际上非常简单。如果你确切知道“艺术家”到目前为止画了什么，你就可以通过数学计算瞬间得出完美的笔触。但旧方法迫使“画家”像“艺术家”一样迈出一小步、缓慢前行，尽管答案就在那里等待被求解。

新想法：“即时修正”

这篇论文提出了一种更聪明的训练机器人方法。它意识到，对于**“画家”**来说，我们不需要猜测和检查。我们可以瞬间解决这个数学问题（即“闭式解”）。

因此，新方法的工作原理如下：

艺术家移动：机器人迈出一小步，以改善其视觉（骨干网络）。
画家瞬间就位：机器人不再迈出一小步，而是瞬间计算出针对该特定视觉的完美画作。这就像“画家”瞬间调整手部动作，以完美匹配画作。
重复：艺术家再次移动，画家瞬间重新调整。

“小批量”带来的问题

在现实世界中，我们无法一次性向机器人展示整个世界；我们只能向它展示小片段（小批量）。

如果你要求“画家”仅基于一个微小片段瞬间解出完美图像，他们可能会感到困惑并反应过度。他们可能会过度记忆那个单一的小片段，从而在下一个片段上失败。这被称为过拟合。
想象一下，如果你要求一位厨师仅通过品尝一粒米就烹饪出一顿完美的饭菜。因为那一粒米很咸，他们可能会加入过多的盐。

解决方案：“近端正则化器”

为了解决这个问题，作者添加了一个“温和的推动”或记忆项。

当“画家”为新的小片段计算完美解时，他们会被告知：“针对这个小片段做到完美，但不要改变你的风格，使其与之前针对上一个片段所做的风格相差太大。”
这使“画家”保持稳定。他们仍然能为当前数据找到最佳答案，但不会剧烈地来回摆动。这就像一位舞者，为了配合音乐调整姿势，但通过扶着栏杆（之前的状态）来保持平衡。

为何这很重要（结果）

该论文在多项任务上测试了这种方法，例如预测分子的化学性质（量子化学）和求解复杂的物理方程（流体动力学）。

速度与稳定性：新方法更加稳定，尤其是在机器人每次只能看到少量数据（小批量）时。如果没有“温和的推动”，旧的“即时修正”方法在少量数据上会崩溃并失败。
优于标准训练：在许多情况下，这种新方法比标准的“逐步进行”方法学习得更快，犯错更少。
因果推断：它在一种棘手的问题类型——“工具变量回归”（用于经济学和科学以确定因果关系）中表现尤为出色，因为它消除了在末尾重新计算所有内容的缓慢且昂贵的第二步需求。

局限性

论文指出，这种魔法技巧仅在目标是最小化平方误差（基本上就是“我的预测偏离了多少？”）时才有效。它非常适用于回归（预测数值）。

当他们尝试将其用于分类（猜测类别，例如“这是猫还是狗？”）时，它在较小的数据集（如 CIFAR-100）上表现令人惊讶地好，但在拥有数千个类别的庞大数据集（如 ImageNet）上却显得吃力。作者建议，对于那些巨大的类别列表，标准的“交叉熵”方法仍然是王者，而将这种新技巧改编以适应那里的工作，则是未来的任务。

总结

可以将这篇论文理解为：通过让“艺术家”缓慢而谨慎地学习，同时让“画家”针对每一个新草图瞬间调整到完美位置（前提是他们不要摆动得太剧烈），来教机器人画画。这是一种使神经网络训练更快、更稳定、更智能的方法，特别适用于涉及数值和物理的任务。

技术摘要：闭式最后一层优化

问题陈述

深度神经网络通常使用随机梯度下降（SGD）或其变体进行优化，对所有层一视同仁。然而，在平方损失函数下，给定骨干网络产生的特征，神经网络线性最后一层的最优权重存在闭式解（具体而言，是岭回归解）。

标准 SGD 将最后一层视为通过梯度步长更新的参数集，忽略了这种解析最优性。这造成了一种不匹配：虽然骨干参数（ $\theta$ ）需要迭代优化，但最后一层权重（ $W$ ）理论上可以直接计算。此外，在具有小批量（mini-batches）的随机设置中，天真地为每个批次计算闭式解会导致对该特定批次的严重过拟合，从而引发不稳定性。现有试图利用这一特性的方法（例如在因果推断中）通常需要对闭式逆矩阵进行反向传播，计算代价高昂，或者将最后一层与骨干网络解耦，导致优化不稳定。

方法论

作者提出了一种训练过程，将最后一层视为骨干参数的隐式确定性函数，仅针对骨干网络进行优化，同时确保最后一层在每一步都保持最优。

1. 理论基础（确定性情况）

对于固定的骨干参数 $\theta$ ，最优最后一层 $W^\star(\theta)$ 最小化正则化平方损失：
$W^\star(\theta) = \arg\min_W \sum \|y_i - W\phi_\theta(x_i)\|^2 + \beta\|W\|_F^2$
作者利用包络定理证明，复合损失 $L^\star(\theta) = L(W^\star(\theta), \theta)$ 关于 $\theta$ 的梯度与原始损失（其中 $W$ 固定为 $W^\star(\theta)$ ）关于 $\theta$ 的梯度相同：
$\nabla_\theta L^\star(\theta) = \nabla_\theta L(W^\star(\theta), \theta)$
这一关键结果表明，无需对计算 $W^\star(\theta)$ 所需的矩阵求逆进行反向传播。相反，算法在以下步骤之间交替进行：

计算 $W_{t+1} = W^\star(\theta_t)$ （闭式更新）。
更新 $\theta_{t+1} = \theta_t - \alpha \nabla_\theta L(W_{t+1}, \theta_t)$ （通过骨干网络的标准反向传播）。

2. 实际适应（随机小批量）

在随机设置中，在单个小批量上计算 $W^\star$ 会导致过拟合。为了解决这个问题，作者引入了一个近端正则化项，惩罚当前闭式解与先前估计 $W_t$ 之间的距离。目标函数变为：
$L_{B_t, W_t}^{\text{prox}}(W, \theta) = \sum_{(x,y) \in B_t} \|y - W\phi_\theta(x)\|^2 + \lambda \|W - W_t\|_F^2$
该近端损失的闭式解为：
$W^\star_{B_t, W_t}(\theta) = (Y_t \phi_\theta(X_t)^\top + \lambda W_t)(\phi_\theta(X_t)\phi_\theta(X_t)^\top + \lambda I)^{-1}$
该公式有效地在最后一层上实现了近似卡尔曼滤波，在适应当前批次的同时累积来自先前批次的信息。算法通过 SGD 更新骨干网络，然后使用此近端闭式解重新拟合最后一层。

3. 分类扩展

虽然该方法主要设计用于回归，但通过将独热标签视为平方损失的目标，将其扩展到了分类任务。输出通过 $\arg\max$ 操作进行解释。作者指出，虽然这种启发式方法适用于小到中等数量的类别，但它不能保证概率归一化。

主要贡献

优化框架：一种实用算法，将闭式最后一层更新集成到标准 SGD 流程中，无需对逆矩阵运算进行反向传播，利用了包络定理。
随机正则化：一种近端公式，通过将最后一层向其先前估计正则化，防止了小批量设置中的过拟合，被解释为近似卡尔曼滤波。
理论分析：
- NTK 区域：证明了在神经切线核（NTK）无限宽极限下，尽管函数空间中的损失是非凸的，仍能收敛到全局最小值。
- 单步分析：一个理论界限量化了闭式步骤相对于梯度下降步骤的优越性，表明增益与特征矩阵的条件数（奇异值的方差）有关。
实证验证：证明了该方法在多种任务上优于标准 SGD 和天真的闭式岭回归。

实验结果

该方法（" $\ell_2$ c.f. proximal ( $\lambda$ )"）与标准平方损失 SGD/Adam 以及天真的闭式岭回归基线（" $\ell_2$ c.f. ridge ( $\beta$ )"）进行了评估。

量子化学（QM9）：近端方法始终优于标准 SGD，特别是在小批量大小下，此时天真的岭回归基线无法收敛。它表现出对骨干优化器选择（SGD 与 Adam）的鲁棒性。
傅里叶神经算子（FNO）：在粘性 Burgers 方程的算子学习中，近端方法在不同批量大小下均保持有效，而天真的岭回归基线在小批量下失效。该方法有效地处理了高维空间自由度。
深度特征工具变量（DFIV）：在因果推断中，该方法消除了训练后在完整数据集上重新估计最后一层的昂贵需求，在计算高效的同时实现了与重新估计策略相当的性能。
分类（CIFAR-100, ImageNet）：在 CIFAR-100 上，该方法在小到中等批量大小下达到了与交叉熵相当或更好的准确率。然而，在 ImageNet（大量类别）上，其表现不如交叉熵，突显了平方损失在高基数分类中的局限性。

意义与主张

论文声称，利用最后一层的闭式解是一种可行且有效的策略，可加速优化并提高稳定性，特别是在回归任务和小批量区域中。

理论意义：该工作提供了对函数空间中修改后损失临界点的首次分析，并证明了在 NTK 区域下收敛到全局最小值，解决了问题固有的非凸性。
实际意义：近端公式解决了小批量闭式更新中固有的过拟合问题，使该方法具有鲁棒性并与标准训练流程兼容。
局限性：作者明确指出，该方法目前仅限于平方损失。虽然它在类别较少的分类任务中表现良好，但在大规模分类（ImageNet）中表现不如交叉熵，这表明将该方法适应交叉熵损失是未来工作的必要方向。该方法还会产生计算开销（矩阵求逆），其规模随最后一层维度增加而增加，使得在小批量区域下，对于非常大的最后一层，其速度慢于标准 SGD，尽管随着骨干网络规模的增加，这一成本会减小。