기술적 요약: 프로젝션 헤드의 기하학: 조건부, 불변성, 그리고 붕괴

문제 제기

자기지도학습 (SSL) 은 특정 아키텍처 구성 요소인 비선형 프로젝션 헤드에 크게 의존합니다. 훈련에 필수적이지만, 이 헤드는 하류 미세 조정 전에 일반적으로 폐기됩니다 ("기울 효과"). 이로 인해 이론적 난제가 발생합니다. 헤드가 사전 훈련 손실을 최소화하는 데 필수적이라면, 왜 그것이 생성한 표현은 추론에 최적화되지 않는 것일까요? 또한, 기존 설명들 (예: 정보 병목 현상이나 차원 붕괴 방지) 은 '무엇'이 일어나는지는 설명하지만, 헤드가 최적화 지형을 어떻게 변형시켜 정보를 필터링하고, 수렴을 가속화하며, 붕괴된 안장점을 탈출시키는지에 대한 기하학적 메커니즘은 설명하지 못합니다.

방법론

저자들은 프로젝션 헤드를 백본 표현 다양체 (manifold) 상의 학습 가능한 리만 계량으로 모델링하는 기하학적 이론을 개발했습니다. 핵심 이론적 프레임워크는 다음과 같습니다:

1. 유효 헤시안 분해: SSL 목적 함수에 대한 백본 표현의 유효 헤시안($H_{eff}$) 을 통해 최적화 지형을 분석합니다. 이 연산자는 풀백 계량($G(z) = J_h^\top \nabla^2_h L J_h$) 과 헤드의 고유한 곡률에 의해 주도되는 상호작용 항($\sum [\nabla_h L]_i \nabla^2_z h_i$) 으로 분해됩니다.
2. 기하학적 조건부: 선형 및 비선형 헤드가 손실 지형을 어떻게 조건화하는지 조사합니다. 선형 헤드는 암시적 부분공간 화이트닝 수행 능력을 분석하고, 비선형 헤드는 곡선 궤적을 따라 지역 계량을 적응시키는 능력을 연구합니다.
3. 안정성 분석: 표현이 상수 벡터로 매핑되는 차원 붕괴의 안정성을 $H_{eff}$의 고유값을 분석하여 붕괴 평형점에서 검토합니다. 저자들은 연속 시간 기울기 흐름과 이산 시간 동역학을 구분합니다.
4. 계량 특이성과 불변성: 불변성 목표와 계량 퇴화 사이의 관계를 공식화합니다. 저자들은 증강 궤도에 대한 불변성을 강제하는 것이 특이한 풀백 계량을 필요로 하여, 증강 접공간을 따라 정보를 파괴함을 증명합니다.
5. 실증적 검증: 이론은 헤시안 추적 (조건수와 음의 고유값 모니터링), 궤적 압축 분석 (증강 다양체 내 분산 감소 측정), 그리고 다양한 아키텍처 (ResNet, ViT) 와 데이터셋 (CIFAR-10/100) 에 대한 활성화 함수 (ReLU 대 Swish/GELU 와 같은 부드러운 활성화) 및 헤드 깊이에 대한 제거 연구 (ablation studies) 를 통해 검증됩니다.

주요 기여

1. 암시적 부분공간 화이트닝으로서의 선형 헤드

저자들은 선형 프로젝션 헤드가 암시적 부분공간 화이트닝을 수행함을 증명합니다. 가중치 행렬 $W$를 학습함으로써 헤드는 표현 공간의 관련 부분공간을 임베딩 공간에 등거리적으로 매핑하는 마할라노비스 계량을 유도합니다. 그러나 이 조건부 처리는 전역적이며, 국소 곡률이 변하는 곡선 다양체에는 불충분합니다.

2. 지역 계량 어댑터로서의 비선형 헤드

비선형 헤드 (MLP) 는 곡선 최적화 궤적을 따라 지역적으로 계량을 적응할 수 있는 능력을 가진 것으로 나타났습니다. 이 논문은 선형 헤드가 곡선 손실 다양체를 전역적으로 평평하게 할 수 없다는 것 (명제 3.3) 을 증명하는 반면, 충분히 깊은 비선형 헤드는 궤적을 따라 유효 헤시안을 $\epsilon$-등방성으로 만드는 이상적인 계량을 근사할 수 있음을 증명합니다 (정리 3.2).

3. 붕괴의 곡률 유도 불안정성

핵심 기여 중 하나는 명시적 음의 샘플 없이도 비대조적 방법들 (예: BYOL, SimSiam) 이 차원 붕괴를 피하는 이유에 대한 기하학적 해결책입니다.

• 정리 4.1: 저자들은 매끄러운 비선형 헤드의 경우, 유효 헤시안의 상호작용 항이 붕괴 평형점에서 음의 고유값을 주입함을 증명합니다. 이는 붕괴된 상태를 안정적인 최소점에서 불안정한 안장점으로 변환시켜, 기울기 하강을 위한 자연스러운 탈출 메커니즘을 제공합니다.
• ReLU 간극: 조각별 선형 활성화 (ReLU 등) 는 고유한 2 차 곡률 ($\nabla^2 \text{ReLU} = 0$) 이 부족합니다. 결과적으로 연속 시간 흐름 하에서 이러한 음의 곡률을 생성할 수 없습니다. 이 논문은 ReLU 헤드가 이산 시간 노이즈 (유한 단계 크기) 와 배치 정규화 (BatchNorm) 를 통해서만 붕괴를 피하는 반면, 부드러운 활성화 (Swish, GELU) 는 불안정성에 대한 본질적인 기하학적 보장을 제공한다고 주장합니다.

4. 계량 퇴화와 정보 - 불변성 트레이드오프

이 논문은 계량 특이성을 통해 "훈련 시 사용, 배포 시 폐기" 전략을 공식화합니다.

• 명제 5.2: 불변성 목표를 충족시키기 위해 프로젝션 헤드는 증강 접공간이 영공간 ($v^\top G(z) v = 0$) 에 놓이도록 특이한 풀백 계량을 유도해야 합니다.
• 정리 5.3: 이 특이성은 백본에 대한 피셔 정보 행렬의 랭크 결핍을 초래합니다. 헤드는 증강 방향을 따라 정보를 소멸시키는 기하학적 저역 통과 필터 역할을 합니다.
• 결론: 헤드를 폐기하는 것은 이론적 필수 사항입니다. 백본은 데이터 다양체의 전체 차원성을 보존하는 반면, 헤드는 불변성을 강제하는 데 필요한 계량 퇴화를 흡수합니다. 추론에 헤드를 사용하면 파괴된 정보 (또는 정보 부재) 가 유지되지만, 원시 백본은 선형적으로 분리 가능한 시맨틱을 유지합니다.

결과

• 헤시안 추적: 실증적 추적은 부드러운 헤드가 유효 헤시안에 음의 고유값을 적극적으로 주입하여 붕괴된 상태를 탈출시키는 "위상 전이"를 유발함을 확인했습니다. 반면, ReLU 와 선형 헤드는 그러한 음의 곡률을 보이지 않으며, 이산 동역학의 도움이 없는 한 연속적인 흐름 하에서 붕괴된 분지에 갇혀 있습니다.
• 궤적 압축: 실험은 프로젝션 헤드가 증강 궤적을 ResNet 의 경우 약 20 배, ViT 의 경우 약 4000 배까지 압축하면서도 클래스 간 시맨틱 분리는 보존함을 보여줍니다. 이는 헤드가 선택적 기하학적 압축기 역할을 함을 확인시켜 줍니다.
• 기반 모델: 사전 훈련된 모델 (DINO, VICReg, Barlow Twins) 에 대한 분석은 기하학적 이분법을 드러냅니다. DINO 헤드는 불변성을 강제하기 위해 궤적을 붕괴시키는 기하학적 압축기로 작용하는 반면, 중복 감소 헤드 (VICReg) 는 상관 제거 제약을 충족시키기 위해 분산을 보존하는 차원 확장기로 작용합니다. 두 경우 모두 헤드는 백본을 목적 함수의 경직된 제약으로부터 분리합니다.
• 깊이 제거 연구: 더 깊은 헤드 (3 층 MLP) 는 다양체를 왜곡하고 계량 특이성을 격리하는 능력이 뛰어나며, 얕거나 선형인 헤드에 비해 하류 백본 정확도가 더 좋습니다. 얕거나 선형인 헤드는 백본을 계량 퇴화로부터 적절히 보호하지 못합니다.

의의

이 논문은 SSL 에서 프로젝션 헤드의 실증적 필요성에 대한 기계적 기하학적 설명을 제공한다고 주장합니다. 두 가지 핵심 난제를 해결합니다:

1. 붕괴 회피: 매끄러운 비선형 헤드의 고유 곡률만으로도 사소한 붕괴된 해를 불안정하게 만들 수 있음을 보여주며, 이는 오직 배치 정규화나 이산 시간 노이즈에 의존하는 설명에 대한 기하학적 대안을 제공합니다.
2. 기울 효과: 헤드가 하류 작업에 필요한 정보를 파괴하는 계량 특이성을 유도하므로 폐기되어야 함을 증명합니다. 헤드는 사전 훈련 목적의 파괴적 제약을 흡수하여 시맨틱 백본을 보호하는 "보편적 기하학적 버퍼" 역할을 합니다.

이 연구는 ReLU 를 부드러운 활성화로 대체하는 것이 단순한 휴리스틱이 아니라 SSL 에서 견고한 최적화를 위한 위상학적 요구사항임을 시사하며, 표준 2~3 층 MLP 아키텍처가 이러한 기하학적 조건을 충족하는 데 필요한 최소 용량임이 실증적으로 검증되었음을 보여줍니다.