[论文解读] Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path

X. Y. Han, Vardan Papyan+1 more|arXiv (Cornell University)|Jun 3, 2021

Domain Adaptation and Few-Shot Learning参考文献 33被引用 19

一句话总结

本文研究了使用均方误差（MSE）损失训练的深度网络中的神经坍缩（NC），表明NC的出现方式与交叉熵（CE）训练相似。通过将MSE损失分解为与NC相关的分量，并引入‘中心路径’——一种理论轨迹，其中分类器在整个训练过程中保持MSE最优——推导出可精确预测NC的梯度流动力学，从而为理解NC的出现提供了数学上可处理的框架。

ABSTRACT

The recently discovered Neural Collapse (NC) phenomenon occurs pervasively in today's deep net training paradigm of driving cross-entropy (CE) loss towards zero. During NC, last-layer features collapse to their class-means, both classifiers and class-means collapse to the same Simplex Equiangular Tight Frame, and classifier behavior collapses to the nearest-class-mean decision rule. Recent works demonstrated that deep nets trained with mean squared error (MSE) loss perform comparably to those trained with CE. As a preliminary, we empirically establish that NC emerges in such MSE-trained deep nets as well through experiments on three canonical networks and five benchmark datasets. We provide, in a Google Colab notebook, PyTorch code for reproducing MSE-NC and CE-NC: at https://colab.research.google.com/github/neuralcollapse/neuralcollapse/blob/main/neuralcollapse.ipynb. The analytically-tractable MSE loss offers more mathematical opportunities than the hard-to-analyze CE loss, inspiring us to leverage MSE loss towards the theoretical investigation of NC. We develop three main contributions: (I) We show a new decomposition of the MSE loss into (A) terms directly interpretable through the lens of NC and which assume the last-layer classifier is exactly the least-squares classifier; and (B) a term capturing the deviation from this least-squares classifier. (II) We exhibit experiments on canonical datasets and networks demonstrating that term-(B) is negligible during training. This motivates us to introduce a new theoretical construct: the central path, where the linear classifier stays MSE-optimal for feature activations throughout the dynamics. (III) By studying renormalized gradient flow along the central path, we derive exact dynamics that predict NC.

研究动机与目标

• 建立神经坍缩（NC）在使用均方误差（MSE）损失训练的深度网络中出现，其行为与使用交叉熵（CE）损失时相似。
• 将MSE损失分解为与NC相关的可解释分量，区分NC兼容项与偏离项。
• 引入‘中心路径’——一种理论构造，其中分类器在整个训练过程中保持MSE最优——以实现NC的精确动态分析。
• 推导沿中心路径的精确梯度流动态，预测NC的出现，为NC提供数学上严谨的基础。

提出的方法

• 将MSE损失分解为（A）当分类器为最小二乘分类器时可通过NC理论解释的分量，以及（B）与该理想分类器的偏离项。
• 通过实证表明，在典型网络和数据集上训练过程中，项（B）可忽略不计，从而支持中心路径近似。
• 将中心路径定义为分类器在每个训练步骤中对特征激活保持MSE最优的轨迹。
• 沿中心路径分析归一化后的梯度流，推导控制NC收敛的精确动力学。
• 利用中心路径与归一化动力学，解析预测NC的极限行为：特征坍缩、单纯形ETF对齐与分类器对齐。
• 利用MSE损失的解析可处理性，提供一种比复杂CE损失更易理解NC的理论框架。

实验结果

研究问题

• RQ1 在使用均方误差（MSE）损失训练的深度网络中，神经坍缩是否如其在交叉熵（CE）损失下的行为一样出现？
• RQ2 MSE损失能否被分解为直接反映神经坍缩机制的分量？
• RQ3 是否存在一种理论轨迹（如‘中心路径’），使得分类器在整个训练过程中保持MSE最优，从而实现NC的精确动态分析？
• RQ4 沿中心路径的精确梯度流动力学能否预测NC的收敛？
• RQ5 在MSE训练下NC出现的数学条件是什么？与基于CE的训练相比有何异同？

主要发现

• 在三种典型架构和五个基准数据集上，使用MSE损失训练的深度网络中均观察到神经坍缩，证实其在CE训练之外也具有鲁棒性。
• MSE损失可分解为与NC对齐的分量（当分类器为最小二乘最优时）和一个在训练过程中实证上可忽略的偏离项。
• 引入中心路径作为理论构造，其中分类器在整个训练过程中保持MSE最优，从而实现精确的动态分析。
• 从中心路径上归一化梯度流推导出的精确动力学可预测NC的收敛，包括特征坍缩与单纯形ETF对齐。
• 基于MSE损失的理论框架为理解NC提供了比复杂CE损失更易解析处理的路径，可实现NC出现的精确推导。
• 结果表明，NC并非CE损失的产物，而是过参数化深度学习在标准优化下的基本属性，即使使用MSE损失也成立。