[논문 리뷰] Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path

X. Y. Han, Vardan Papyan외 1명|arXiv (Cornell University)|2021. 06. 03.

Domain Adaptation and Few-Shot Learning참고 문헌 33인용 수 19

한 줄 요약

이 논문은 평균 제곱 오차(MSE) 손실로 훈련된 딥 네트워크에서 신경망 붕괴(Neural Collapse, NC)를 조사하며, NC가 교차 엔트로피(CE) 훈련과 유사하게 나타남을 보여준다. MSE 손실을 NC와 관련된 구성요소와 분리하여 분해하고, 분류기가 항상 MSE 최적임을 유지하는 이론적 궤적인 '중앙 경로(central path)'를 도입함으로써, 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.

연구 동기 및 목표

• 평균 제곱 오차(MSE) 손실로 훈련된 딥 네트워크에서 신경망 붕괴(Neural Collapse, NC)가 교차 엔트로피(CE) 손실과 유사하게 발생함을 입증하는 것.
• MSE 손실을 NC와 관련된 해석 가능한 구성요소로 분해하여, NC에 적합한 항과 이격된 항을 구분하는 것.
• 분류기가 훈련의 각 단계에서 특성 활성화에 대해 항상 MSE 최적임을 유지하는 이론적 궤적인 '중앙 경로(central path)'를 도입하여 정확한 동역학 분석을 가능하게 하는 것.
• 중앙 경로를 따라 정확한 기울기 흐름 역학을 유도하여 NC의 발생을 예측하고, NC에 대한 수학적으로 엄밀한 기초를 마련하는 것.

제안 방법

• 분류기가 최소 제곱 분류기일 때 NC 이론에 의해 해석 가능한 구성요소(A)와 이 이상적인 분류기에서의 이격 항(B)으로 MSE 손실을 분해하는 것.
• 실증적으로, 표준 네트워크와 데이터셋에서 훈련 중에 항(B)이 무시할 만큼 작음을 보여주어 중심 경로 근사가 타당함을 입증하는 것.
• 중앙 경로를 각 훈련 단계에서 특성 활성화에 대해 분류기가 항상 MSE 최적임을 유지하는 궤적으로 정의하는 것.
• 중앙 경로를 따라 재정규화된 기울기 흐름을 분석하여 NC 수렴를 지배하는 정확한 역학을 도출하는 것.
• 중앙 경로와 재정규화된 역학을 활용하여 NC의 최종 행동, 즉 특성 붕괴, 단체 ETF 수렴, 분류기 정렬을 분석적으로 예측하는 것.
• MSE 손실의 분석 가능성을 활용하여, 복잡한 CE 손실보다 더 접근하기 쉬운 이론적 프레임워크를 제공하여 NC의 발생을 이론적으로 설명하는 것.

실험 결과

연구 질문

• RQ1 평균 제곱 오차(MSE) 손실로 훈련된 딥 네트워크에서 신경망 붕괴(Neural Collapse)가 교차 엔트로피(CE) 손실과 유사하게 나타나는가?
• RQ2 MSE 손실을 NC의 메커니즘을 직접 반영하는 구성요소로 분해할 수 있는가?
• RQ3 분류기가 항상 MSE 최적임을 유지하는 이론적 궤적—예를 들어 '중앙 경로'—이 존재하는가? 이는 NC의 정확한 동역학 분석을 가능하게 하는가?
• RQ4 중앙 경로를 따라 정확한 기울기 흐름 역학이 NC 수렴을 예측할 수 있는가?
• RQ5 MSE 훈련 하에서 NC가 발생하는 수학적 조건는 무엇이며, CE 기반 훈련과 비교해보면 어떻게 다른가?

주요 결과

• 세 가지 표준 아키텍처와 다섯 개의 벤치마크 데이터셋에서 MSE 손실로 훈련된 딥 네트워크에서 신경망 붕괴(Neural Collapse)가 나타나며, 이는 CE 훈련 외부에서도 NC가 강건함을 확인한다.
• MSE 손실은 분류기가 최소 제곱 최적일 때 NC와 일치하는 구성요소와 훈련 중에 실증적으로 무시할 만큼 작은 이격 항으로 분해되며, 이는 NC 이론에 부합한다.
• 중앙 경로는 분류기가 훈련 전반에 걸쳐 항상 MSE 최적임을 유지하는 이론적 구성으로 도입되며, 정확한 동역학 분석을 가능하게 한다.
• 중앙 경로를 따라 재정규화된 기울기 흐름에서 도출된 정확한 역학은 특성 붕괴와 단체 ETF 수렴을 포함한 NC 수렴을 예측한다.
• MSE 손실 기반 이론적 프레임워크는 복잡한 CE 손실보다 더 분석 가능성이 높아, NC 발생의 정확한 유도를 가능하게 하며, NC의 발생을 이해하는 데 더 유리한 길을 제공한다.
• 결과는 NC가 CE 손실의 산물이 아니라, 표준 최적화 하에서 과다 매개변수화된 딥 러닝의 근본적인 성질임을 시사한다.