[论文解读] On the Optimization Landscape of Neural Collapse under MSE Loss: Global Optimality with Unconstrained Features

Jian Zhou, Li, Xiao+4 more|arXiv (Cornell University)|Mar 2, 2022

Stochastic Gradient Optimization Techniques被引用 6

一句话总结

本论文首次对在无约束特征模型下使用均方误差（MSE）损失训练的深度神经网络进行了全局优化景观分析。证明了所有全局最小值点均表现出神经坍缩（NC），且不存在虚假局部最小值——仅有严格鞍点。同时，通过重标定MSE损失可改善优化景观与泛化性能，该结论在ResNet18上得到实证验证。

ABSTRACT

When training deep neural networks for classification tasks, an intriguing empirical phenomenon has been widely observed in the last-layer classifiers and features, where (i) the class means and the last-layer classifiers all collapse to the vertices of a Simplex Equiangular Tight Frame (ETF) up to scaling, and (ii) cross-example within-class variability of last-layer activations collapses to zero. This phenomenon is called Neural Collapse (NC), which seems to take place regardless of the choice of loss functions. In this work, we justify NC under the mean squared error (MSE) loss, where recent empirical evidence shows that it performs comparably or even better than the de-facto cross-entropy loss. Under a simplified unconstrained feature model, we provide the first global landscape analysis for vanilla nonconvex MSE loss and show that the (only!) global minimizers are neural collapse solutions, while all other critical points are strict saddles whose Hessian exhibit negative curvature directions. Furthermore, we justify the usage of rescaled MSE loss by probing the optimization landscape around the NC solutions, showing that the landscape can be improved by tuning the rescaling hyperparameters. Finally, our theoretical findings are experimentally verified on practical network architectures.

研究动机与目标

• 理论证明在实践中表现与交叉熵（CE）相当或更优的均方误差（MSE）损失下，神经坍缩（NC）的合理性。
• 分析在无约束特征模型下，原始非凸MSE损失的全局优化景观，其中特征为自由优化变量。
• 证明NC解是唯一的全局最小值点，且所有其他临界点均为具有负曲率的严格鞍点，从而支持高效优化。
• 研究重标定超参数在改善MSE损失的优化景观与泛化性能方面的作用。
• 在miniImageNet等数据集上，通过ResNet18等实际架构的实证训练，验证理论发现。

提出的方法

• 形式化定义无约束特征模型，其中最后层的特征与分类器为自由变量，与网络权重解耦。
• 通过临界点分类分析MSE损失景观：证明全局最小值点恰好对应于NC解。
• 通过证明其Hessian矩阵具有负曲率方向，将非全局临界点表征为严格鞍点。
• 引入带有两个超参数（α 和 M）的重标定MSE损失，以改善NC解附近的景观。
• 利用景观可视化与ResNet18上的实证训练，研究重标定对NC程度与测试准确率的影响。
• 在miniImageNet上进行实验，通过调整α与M，关联重标定与NC增强及泛化性能提升的关系。

实验结果

研究问题

• RQ1 在无约束特征模型下，MSE损失是否唯一地导致神经坍缩作为全局解？
• RQ2 非凸MSE损失的全局优化景观如何？具体而言，是否存在虚假局部最小值，或仅存在严格鞍点？
• RQ3 重标定MSE损失如何影响优化景观与泛化性能？
• RQ4 重标定的理论优势是否可在深度网络的实际训练中被实证观察到？
• RQ5 基于MSE的NC与基于CE的NC在训练动态与鲁棒性方面有何比较？

主要发现

• 在无约束特征模型下，MSE损失的所有全局最小值点均为神经坍缩解，其中类别均值形成单纯形等角紧框架（ETF），且类内差异坍缩为零。
• 所有非全局临界点均为严格鞍点，因其Hessian矩阵至少有一个负特征值，可通过一阶方法高效逃离。
• 通过超参数α与M对MSE损失进行重标定，可改善优化景观，减少靠近决策边界的点数，并提升泛化性能。
• 在ResNet18上的实证结果表明，增大α或M均与更强的神经坍缩及更高的训练与测试准确率相关。
• NC程度与模型性能与重标定参数正相关，表明重标定对于MSE损失实现最优性能至关重要。
• 研究结果将神经坍缩的普遍性扩展至MSE训练的网络，表明其在任意训练算法（只要能逃离严格鞍点）与足够表达力的架构下均会出现。