用闭式解优化网络最后一层，提升回归任务训练效率。
Closed-Form Last Layer Optimization
- 将最后一层权重设为主干参数的函数，仅优化主干参数。
- 在神经正切核框架下保证收敛，比SGD更快达到最优解。
- 适用于回归任务，尤其在神经算子与因果推断中表现优异。
神经网络通常使用随机梯度下降的变体进行优化。但在平方损失下，线性最后一层权重的最优解可解析求得。本文提出在优化过程中利用这一性质，将最后一层视为主干参数的函数，仅对主干参数进行优化。该方法等价于交替执行主干的梯度下降与最后一层的闭式更新。针对随机梯度场景，通过权衡当前批次损失与历史批次信息实现改进。理论分析表明，在神经正切核框架下，该方法能收敛至最优解，并在一步分析中量化其相对于标准SGD的增益。实验验证了该方法在多个回归任务（包括神经算子和因果推断）中优于SGD和Adam的效果。
原文摘要 · Abstract (English)
Neural networks are typically optimized with variants of stochastic gradient descent. Under a squared loss, however, the optimal solution to the linear last layer weights is known in closed-form. We propose to leverage this during optimization, treating the last layer as a function of the backbone parameters, and optimizing solely for these parameters. We show this is equivalent to alternating between gradient descent steps on the backbone and closed-form updates on the last layer. We adapt the method for the setting of stochastic gradient descent, by trading off the loss on the current batch against the accumulated information from previous batches. We provide theoretical analyses showing convergence of the method to an optimal solution in the neural tangent kernel regime, as well as quantifying the gains compared to standard SGD in a one-step analysis. Finally, we demonstrate the effectiveness of our approach compared with SGD and Adam on a squared loss in several regression tasks, including neural operators and causal inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。