기술 요약: 유계 비선형성을 가진 넓은 얕은 모델에 대한 경사 하강법의 전역 수렴

1. 문제 제기

본 논문은 넓은 얕은 신경망과 어텐션 메커니즘을 학습할 때 경사 하강법 (특히 경사 흐름) 의 전역 수렴성을 조사한다. 경사 하강법은 비볼록 지형에서 비최적 임계점에 수렴할 위험이 있는 1 차 알고리즘이지만, 경험적 증거는 과매개변수화 영역에서 전역 최소점을 찾는 경우가 많음을 시사한다.

이전 이론적 연구, 특히 [CB18] 은 특정 조건 하에서 넓은 네트워크의 전역 수렴성을 확립했다:

• ReLU 활성화: 매핑이 양의 2-동차 (positively 2-homogeneous) 인 경우.
• 스칼라 출력 가중치를 가진 시그모이드 활성화: 매핑이 부분적으로 1-동차 (partially 1-homogeneous) 인 경우.

그러나 이러한 결과는 다음을 다루지 못한다:

1. 벡터 출력 가중치: 다중 클래스 분류와 어텐션 메커니즘 (여기서 출력 차원 $d_w > 1$) 에서 일반적으로 사용됨.
2. 어텐션 레이어: 특히 소프트맥스와 같은 유계 비선형성과 벡터 값 파라미터에 의존하는 멀티헤드 어텐션 레이어.
3. 일반적인 초기화: 이전 증명들은 종종 초기 파라미터 분포가 컴팩트한 지지집합 (compact support) 을 가져야 한다고 요구하여, 널리 사용되는 가우스 초기화를 배제했다.

저자들은 멀티헤드 어텐션과 벡터 출력 가중치를 가진 2 층 네트워크를 포함한 광범위한 모델 클래스에 대해, 모든 비전역 최소점이 경사 하강법 역학 하에서 불안정함을 증명함으로써 이러한 격차를 해소하고자 한다. 결과적으로, 초기 분포가 완전한 지지집합 (예: 가우스 분포) 을 가진다면, 무한한 폭의 극한에서 연속 시간 경사 흐름은 전역 최소점으로 수렴한다.

2. 방법론

2.1 평균장 프레임워크

본 분석은 $m$ 개의 뉴런의 파라미터를 파라미터 공간 $\Omega$ 위의 확률 분포 $\mu_t$ 로 인코딩하는 평균장 극한 (mean-field limit) 에 의존한다. $m \to \infty$ 일 때, 이산적인 경사 하강법 역학은 확률 측도 공간 $\mathcal{P}(\Omega)$ 위에 정의된 함수 $F$ 의 와서슈타인 경사 흐름 (Wasserstein gradient flow) 으로 수렴한다.

목적 함수는 $F(\mu) = \mathcal{R}(\int_\Omega \Phi d\mu)$이며, 여기서 $\mathcal{R}$은 손실 함수이고 $\Phi(w, \theta) = \phi(\theta)w$는 모델을 나타낸다. 여기서 $w$는 출력 가중치 (또는 어텐션의 값 행렬) 에 해당하고, $\theta$는 숨겨진 파라미터 (또는 쿼리/키 행렬) 에 해당한다.

2.2 잘 정의됨과 안정성 (2 절)

수렴성을 분석하기 전에, 저자들은 평균장 역학의 잘 정의됨 (well-posedness) 을 확립한다.

• 가정: 초기 측도 $\mu_0$는 서브가우스 (sub-Gaussian) 이다 (이는 [CB18] 에서 요구된 컴팩트한 지지집합보다 약한 조건임). 여기에는 가우스 초기화가 포함된다.
• 결과: 그들은 모든 시간 $t \geq 0$에 대해 와서슈타인 경사 흐름의 존재성과 유일성을 증명한다. 또한, 안정성 추정을 확립한다: 경험적 측도 열 $\mu_{m,0}$이 $W_2$ 거리에서 $\mu_0$로 수렴하면, 해당 흐름은 유한 시간 구간에서 균일하게 수렴한다.
• 기법: [CB18] 이 거리 공간 경사 흐름 이론을 사용한 것과 달리, 본 증명은 상호작용 입자 시스템에 대한 Dobrushin 유형의 추정을 기반으로 하는 고정점 논증 (fixed-point argument) 을 활용하며, 속도장의 선형 성장을 처리하기 위해 서브가우스 꼬리 제어를 활용한다.

2.3 "활성 집합 탈출"을 통한 전역 수렴 (3 절)

핵심 이론적 기여는 비전역 최소점이 불안정하다는 증명이다.

• 메커니즘: 저자들은 "탈출 활성 집합" (escaping active set, $E(g)$) 개념을 도입한다. 후보 극한점 $\mu$가 전역 최소점이 아니라면, 파라미터 공간의 어떤 영역에서 출력 가중치 $w$가 양의 비율로 균일하게 증가해야 하는 영역이 존재한다.
• 정리 2: 초기 분포가 완전한 지지집합을 가진다면, 와서슈타인 경사 흐름은 연관된 벡터장 $g_\mu$가 탈출 활성 집합을 허용하는 어떤 측도 $\mu$로도 수렴할 수 없다.
• 귀결: 모든 비영 $g_\mu$가 탈출 활성 집합을 허용한다면, 흐름은 오직 전역 최소점 (여기서 $g_\mu \equiv 0$) 으로만 수렴한다.

2.4 [CB18] 의 수정 및 확장

• 스칼라 경우 ($d_w=1$): 저자들은 스칼라 출력과 유계 비선형성 (시그모이드) 에 대한 탈출 활성 집합의 존재에 대해 수정된 증명을 제공한다. [CB18, 부록 C] 의 원래 증명은 활성 집합의 경계에서 그래디언트의 유계성에 관한 오류가 있었다. 저자들은 매핑과 그 도함수에 대한 균일 수렴 가정을 사용하여 무한대에서의 거동을 분석함으로써 이를 수정한다.
• 벡터 경우 ($d_w > 1$): 저자들은 벡터 출력 가중치로 분석을 확장한다. 그들은 벡터 역학의 복잡성 (방향 변화) 으로 인해 스칼라 증명의 자연스러운 일반화가 실패함을 보여준다. 그들은 고정된 방향 벡터 $v$를 사용하여 활성 집합의 두 번째 구성을 제안한다. 그들은 이 구성이 $d_\theta = 1$ (스칼라 숨겨진 파라미터) 인 경우에 작동함을 증명하지만, 일반적인 벡터 경우에 대한 엄밀한 증명은 여전히 열려있으며, 약한 정규성 가정 하에서 결과가 성립할 것이라고 추측한다.

3. 주요 결과

1. 서브가우스 초기화에 대한 잘 정의됨: 평균장 학습 역학은 컴팩트한 지지집합을 요구했던 이전 결과를 확장하여, 서브가우스 초기 분포에 대해 잘 정의되고 안정적이다.
2. 유계 비선형성에 대한 전역 수렴: 유계 비선형성 (예: 시그모이드) 과 스칼라 출력 가중치를 가진 모델의 경우, 초기 분포가 완전한 지지집합을 가진다면 와서슈타인 경사 흐름은 전역 최소점으로 수렴한다.
3. 어텐션 및 벡터 가중치에 대한 적용:
   • 이 프레임워크는 멀티헤드 어텐션 레이어 ( $\Phi(A, V) = V \psi(A)$로 모델링되며, 여기서 $A=K^T Q$) 에 적용된다.
   • 저자들은 파라미터가 커짐에 따라 소프트맥스 어텐션 메커니즘이 "하드맥스 (hardmax)" 메커니즘으로 균일하게 수렴함을 증명하며, 이는 독립적인 관심 대상이 되는 결과이다.
   • 그들은 이러한 아키텍처에 대해 전역 수렴 조건이 충족됨을 확립하지만, 벡터 출력 경우 ($d_w > 1$) 에 대한 엄밀한 증명은 고차원에서의 탈출 활성 집합 존재에 관한 추측에 의존한다.
4. 이전 연구의 수정: 본 논문은 시그모이드 네트워크의 비전역 최소점의 불안정성에 관한 [CB18, Proposition C.4] 증명에서 결함을 식별하고 수정한다.

4. 중요성과 주장

본 논문은 과매개변수화 딥러닝에서 경사 하강법이 성공하는 이유에 대한 이론적 이해를 크게 확장한다고 주장한다:

• 격차 해소: 특징이 학습되지 않는 제한적인 NTK 영역과 ReLU 활성화의 특정 경우를 넘어선다. 현대 아키텍처 (Transformer, 시그모이드 네트워크) 에 만연한 유계 비선형성을 가진 모델에 대한 특징 학습 영역 (feature learning regime) 을 다룬다.
• 현실적인 아키텍처: 멀티헤드 어텐션 레이어와 벡터 출력 가중치를 포함함으로써, 단순화된 2 층 네트워크가 아닌 현재 최첨단 모델의 핵심이 되는 아키텍처에 결과가 적용된다.
• 강건한 초기화: 초기화 요구사항을 컴팩트한 지지집합에서 서브가우스로 완화함으로써, 이론이 실제 학습 관행 (예: 가우스 초기화) 과 더 잘 부합한다.
• 열린 문제에 대한 겸손: 저자들은 일반적인 숨겨진 차원 ($d_\theta > 1$) 을 가진 벡터 출력 경우 ($d_w > 1$) 에 대한 완전한 증명은 여전히 열린 문제임을 명시적으로 밝힌다. 그들은 충분 조건과 국소 분석을 제공하지만, 일반적인 벡터 경우에 대한 완전한 증명을 주장하지는 않는다. 또한, 그들의 분석은 실제에서 흔하지만 현재 이론적 범위를 벗어난 정규화 항 (예: $\ell_2$ 가중 감쇠) 과 GELU/SiLU 와 같은 비유계 활성화 함수를 제외한다고 지적한다.

요약하자면, 본 논문은 유계 비선형성을 가진 넓은 얕은 모델에 대해, 학습 손실의 지형이 초기화가 충분히 넓다면 경사 흐름이 비최적 임계점에서 밀려나게 됨을 시사하는 엄밀한 수학적 기초를 제공한다. 이는 이러한 설정에서 경사 하강법이 전역 해를 찾는다는 경험적 관찰을 지지한다.