기술적 요약: SGD 가 브라운 운동이 아닌 이유

문제 제기

확률적 경사 하강법 (SGD) 은 딥러닝의 기초 최적화 알고리즘이지만, 이에 대한 이론적 이해는 여전히 불완전합니다. 주된 연구 흐름은 SGD 를 랑베주 동역학의 이산화로 모델링하여, 미니배치 노이즈를 고정된 퍼텐셜 (평균 손실) 에 작용하는 브라운 운동으로 취급합니다. 이 프레임워크는 연속 시간 극한과 특정 $\sqrt{\eta}$ 노이즈 스케일링을 가정하여 정상 깁스 분포를 유도합니다.

그러나 본 논문은 표준 랑베주 근사가 유한 학습률 ($\eta$) 에서 이산 SGD 동역학을 충실히 재현하지 못한다고 주장합니다. 핵심적인 불일치는 확률성의 스케일링에 있습니다. 실제 이산 업데이트에서는 노이즈가 미니배치 손실을 통해 $\eta$ 차수로 진입하는 반면, 랑베주 공식화는 연속 시간에서의 확산을 보존하기 위해 $\sqrt{\eta}$ 스케일링을 부과합니다. 결과적으로 랑베주 근사는 정확한 이산 동역학에 존재하는 $\eta^2$ 차수의 항들을 생략하게 되어, 정상 분포의 존재 여부와 임계점 근처의 매개변수 분산 행동에 관해 질적으로 잘못된 예측을 초래합니다.

방법론

저자들은 SGD 를 고정된 퍼텐셜 내의 브라운 운동이 아닌, 요동치는 손실 지형에서의 결정론적 동역학으로 취급하는 새로운 프레임워크를 제안합니다. 무작위성은 미니배치 샘플링에서 내재적으로 발생하며, 이는 매 단계마다 목적 함수를 변화시킵니다.

방법론은 다음과 같이 진행됩니다:

1. 정확한 마스터 방정식: 정확한 이산 SGD 업데이트 규칙 $w_{n+1} = w_n - \eta \nabla L_n(w_n)$ 에서 시작하여, 매개변수 분포 $p_n(w)$ 의 진화에 대한 정확한 마스터 방정식을 유도합니다.
2. 이산 포커 - 플랑크 유도: 마스터 방정식을 학습률 $\eta$ 의 거듭제곱으로 전개하고 2 차 항에서 절단하여 이산 포커 - 플랑크 방정식을 유도합니다. 이 방정식은 연속 시간 근사를 도입하지 않고 매개변수 밀도의 진화를 포착합니다.
3. 국소 분석: 이 프레임워크를 평균 손실의 임계점 주변에 적용합니다. 손실 지형의 국소 2 차 근사를 사용하여, 평균 헤시안의 고유기저에서 매개변수 분포의 평균과 공분산의 동역학을 분석합니다.
4. 가정: 해석 가능한 폐쇄형 식을 얻기 위해 분석은 다음을 가정합니다:
   • 복원 추출에 의한 독립적인 미니배치 샘플링.
   • 평균 헤시안 고유기저에서 헤시안 요동의 근사적 대각성.
   • 기울기 노이즈 공분산이 헤시안 구조에 비례한다는 경험적 관계 ($E[G_i G_j] \approx \gamma E[H_{ij}]$).

주요 기여

1. 이산 포커 - 플랑크 프레임워크

본 논문은 표준 랑베주 기반 포커 - 플랑크 방정식과 다른 매개변수 밀도 진화의 2 차 전개식 (식 4) 을 유도합니다. 구체적으로, 이산 전개는 $\eta^2$ 차수에서 $(\nabla \bar{L})^2$ 에 비례하는 추가 항을 포함합니다. 표준 랑베주 공식화는 유한 $\eta$ 에서의 이산 업데이트와 일관되지 않은 $\sqrt{\eta}$ 노이즈 스케일링을 가정하므로 이 항을 생략합니다.

2. 랑베주 동역학과의 불일치

저자들은 이 생략이 단순한 정량적 보정이 아니라 질적 차이를 초래할 수 있음을 입증합니다. 1 차원 토이 예시에서 표준 랑베주 근사는 유한 분산을 갖는 정상 상태를 예측하는 반면, 정확한 이산 SGD 동역학은 분산의 지수적 성장 (불안정성) 을 예측하는 조건에서 이를 예측합니다. 이 불일치는 랑베주 절단이 유지된 확산 항과 동일한 형식적 차수의 항들을 폐기하기 때문에 발생합니다.

3. SGD 동역학의 통합적 그림

분석은 SGD 동역학이 평균 헤시안의 고유값 ($\lambda_i$) 에 따라 뚜렷한 영역으로 분해됨을 보여줍니다:

• 구속된 (날카로운) 방향: 충분히 큰 양의 곡률을 가진 방향의 경우, 매개변수 분산이 유한한 정상 극한으로 수렴합니다. 극한 분산은 학습률 $\eta$ 에 대략 비례하며 곡률 크기와 무관합니다 (역분산 - 평탄도 관계와 일치).
• 확산적 (평평한) 방향: 거의 평평한 방향 ($\lambda_i \approx 0$) 의 경우, 분산은 관측된 시간 범위 내에서 정상 분포로 수렴하지 않습니다. 대신 시간과 선형적으로 비례하여 ($\propto n$) 증가하며, 이는 손실 지형의 계곡을 따라 효과적인 확산에 해당합니다.

4. 실증적 검증

이론적 예측은 컴퓨터 비전 (MNIST 의 MLP) 과 자연어 처리 (셰익스피어 및 WikiText-2 의 NanoGPT) 의 신경망 모델에서 검증되었습니다.

• 소규모 실험: 정확한 밀집 헤시안을 사용하여, 매개변수 공분산이 평균 헤시안 고유기저에서 대략 대각선임을 확인했으며, 분산 프로파일이 이론적 예측과 일치하여 포화되는 방향과 포화되지 않는 방향 사이의 명확한 분리를 보여주었습니다.
• 대규모 실험: 660 만 개 파라미터를 가진 NanoGPT 모델에서 스토캐스틱 란초스를 사용하여 헤시안 스펙트럼을 추정한 결과, 실험적 분산 포화 수준은 이산 예측 (식 19) 과 일치하며 학습률이 증가함에 따라 표준 랑베주 예측과 크게 벗어났습니다.

결과 및 의의

본 논문은 SGD 를 고정된 퍼텐셜 내의 브라운 운동이 아닌, 요동치는 손실 지형에서의 결정론적 동역학으로 이해하는 것이 더 정확하다고 주장합니다. 이 관점은 유한 학습률에서 관찰된 이론과 실무 간의 불일치를 해결합니다.

이 연구의 의의는 다음과 같습니다:

• 이론적 기반의 교정: 표준 랑베주 근사가 실용적인 학습률에서 이산 SGD 에 대한 나쁜 근사임을 강조하며, 안정성과 정상성을 결정하는 중요한 $\eta^2$ 항들을 누락하고 있음을 지적합니다.
• 동역학 영역의 설명: 날카로운 방향에서는 구속된 행동을 보이고 평평한 방향에서는 지속적인 확산을 보이는 SGD 의 현상에 대한 통합적 설명을 제공합니다. 이는 종종 경험적으로 관찰되지만 평형 통계 물리학 모델로는 완전히 포착되지 않았던 현상입니다.
• 유한 단계 효과: 최적화 동역학에서 이산 시간 효과의 중요성을 강조하며, SGD 의 '노이즈'는 외부의 확률적 강제 항이 아니라 알고리즘의 구조와 손실 지형 기하학의 내재적 결과임을 시사합니다.

저자들은 국소 2 차 근사에 대한 의존성, 헤시안 요동의 독립성에 대한 가정, 그리고 복원 추출 샘플링으로의 제한 등 한계를 지적합니다. 그러나 다양한 아키텍처와 작업에 걸친 실증적 증거는 유도된 동역학 영역이 고차원 최적화 지형의 일반적인 속성임을 시사합니다.