기술 요약: 에지 플로우 (Edge Flow)

문제 정의

딥러닝에서의 경사 하강법(Gradient Descent, GD)은 흔히 **에지 오브 스태빌리티(Edge of Stability, EoS)**라고 알려진 영역에서 작동한다. 이 영역에서 손실 함수의 헤시안(Hessian)의 최대 고유값(샤프니스, $S(w_t)$)은 안정성 임계값인 $2/\eta$ 근처에서 머문다 (여기서 $\eta$는 학습률이다). 이 과정에서 훈련 손실은 일반적으로 감소하지만, 반복적인 궤적은 준주기적 진동(quasi-periodic oscillations)을 보이며, 샤프니스는 발산하는 대신 자기 조절 메커니즘을 통해 안정화된다.

경사 흐름(gradient flow, $\dot{w} = -\nabla L(w)$)이나 데센트 레마(descent lemma)와 같은 고전적인 연속 시간 분석 도구들은 이 영역에서 실패한다. 왜냐한 이들은 학습률에 무관하며, $\eta \cdot S(w)$의 비율에 의존하는 현상을 포착할 수 없기 때문이다. 기존의 EoS용 연속 시간 모델들은 다음과 같은 한계를 갖는다:

• 센트럴 플로우 (Central Flow, Cohen et al., 2025): 진동의 중심을 투영된 경사 흐름(projected gradient flow)을 통해 모델링한다. 이는 안정화되기는 하지만, 투영 제약 조건으로 인해 시스템을 시뮬레이션하거나 이론적으로 분석하기 어렵게 만든다.
• 로드 플로우 (Rod Flow, Regis and Chewi, 2026): 반복 궤적을 단단한 막대(rigid rod)의 끝점으로 모델링한다. 개념적으로는 매력적이지만, 진동 크기에 대한 업데이트 규칙이 EoS 시작 단계에서의 자기 안정화 역학을 포착하지 못한다.

본 논문이 다루는 핵심적인 미해결 질문은 다음과 같다: EoS에서의 GD 역학을 설명할 수 있는, 다루기 쉽고(tractable), 충실하며(faithful), 예측 가능한(predictive) 연속 시간 모델은 무엇인가?

방법론: 에지 플로 (Edge Flow)

저자들은 GD 역학을 세 가지 천천히 변화하는 양으로 분해하는 세 개의 결합된 상미분 방정식(ODE) 시스템인 **에지 플로우(Edge Flow)**를 제안한다:

1. 중심 (The Center, $\bar{w}_t$): 진동의 서서히 이동하는 중점.
2. 진동 방향 (The Oscillation Direction, $u_t$): 헤시안의 최상위 고유벡터를 추적하는 단위 벡터.
3. 진동 크기 (The Oscillation Magnitude, $x_t$): 진동의 반-길이(half-length).

GD 반복 궤적은 $w_t \approx \bar{w}_t + (-1)^t x_t u_t$로 모델링된다. 에지 플로우 시스템은 다음과 같이 정의된다:

1. 중심 역학 (Center Dynamics):
   $$\frac{d\bar{w}_t}{dt} = -\frac{1}{2} \left( \nabla L(\bar{w}_t + x_t u_t) + \nabla L(\bar{w}_t - x_t u_t) \right)$$
   이는 대칭화된 손실(symmetrized loss)(두 끝점에서의 그래디언트 평균)에 대한 경사 흐름을 나타낸다. 이는 $x_t$에 대한 3차 미분 보정 항을 포함하여 표준 경사 흐름과 다르며, 이 항은 중심을 더 낮은 곡률을 가진 영역으로 편향시킨다.

2. 방향 역학 (Direction Dynamics):
   $$\frac{du_t}{dt} = 2 \left( \nabla^2 L(\bar{w}_t) - (u_t^\top \nabla^2 L(\bar{w}_t) u_t) I \right) u_t$$
   이는 단위 구(unit sphere) 상의 **레일리 몫 경사 흐름(Rayleigh quotient gradient flow)**이다. 이는 파워 메서드(power method)의 연속 시간 극한과 유사하게, $u_t$를 $\bar{w}_t$에서의 헤시안의 최상위 고유벡터로 유도한다.

3. 크기 역학 (Magnitude Dynamics):
   $$\frac{dx_t}{dt} = \left( S(\bar{w}_t) - \frac{2}{\eta} \right) x_t$$
   이 스칼라 선형 ODE는 진동 크기의 성장 또는 감소를 지배한다. 만약 샤프니스 $S(\bar{w}_t) > 2/\eta$이면 $x_t$는 지수적으로 성장하고, $S(\bar{w}_t) < 2/\eta$이면 $x_t$는 감소한다.

이산화: 에지 경사 하강법 (Edge Gradient Descent, EGD)

저자들은 **에지 경사 하강법(EGD)**이라는 자연스러운 이산화 방식을 제안한다. 각 반복 단계에서는 다음이 필요하다:

• 두 번의 그래디언트 계산: $\nabla L(\bar{w}_k \pm x_k u_k)$.
• 한 번의 헤시안-벡터 곱: $\nabla^2 L(\bar{w}_k) u_k$.
• 고유값 분해나 최적화 서브루틴은 필요하지 않다.

이 알고리즘은 두 개의 서로 다른 스텝 사이즈를 사용한다: $\eta$ (샤프니스 임계값을 정의하는 원래의 학습률)와 $\rho$ (ODE의 적분 단계를 위한 스텝 사이즈, 여기서 $\rho \le \eta$이다).

주요 기여

1. 단순하고 다루기 쉬운 모델: 에지 플로우는 GD 구조에 대한 제1원리와 기초적인 관찰로부터 유도되었다. 이는 시뮬레이션과 분석이 용이한 세 개의 결합된 ODE 시스템이다.
2. 창발적 자기 안정화 (Emergent Self-Stabilization): 안정성을 강제하기 위해 제약 조건을 부과하는 기존 모델들과 달리, 에지 플로우는 결합된 역학의 자연스러운 결과로서 자기 안정화 메커니즘(Damian et al., 2023이 식별함)을 재현한다. 샤프니스가 임계값을 초과하면 진동 크기가 증가하고, 이는 대칭화된 그래디언트를 통해 중심을 더 낮은 곡률 영역으로 밀어내어 결과적으로 샤프니스를 감소시킨다.
3. 충실한 추적 및 예측: 경험적으로 에지 플로우는 다양한 아키텍처(MLP, CNN, ResNet, ViT)와 손실 함수에 대해 풀 배치(full-batch) GD 궤적, 손실, 그리고 샤프니스를 추적한다. 결정적으로, 이 모델은 센트럴 플로우의 즉각적인 투영이 놓치는 특징인 EoS 시작 단계에서의 **오버슈트(overshoot)와 감쇠 진동(damped oscillation)**을 해결한다.
4. 불안정성에 대한 통찰: 이 모델은 EoS 시작 시 발생하는 손실 스파이크(loss spikes)의 원인이 경직된(stiff) 중심 역학의 이산화 아티팩트임을 식별한다. 저자들은 두 가지 원칙적인 해결책을 제시한다:
   • 이산화 정밀도 개선: 더 작은 적분 단계 $\rho$ (더 큰 $K = \eta/\rho$)를 사용하면 손실 스파이크를 억제할 수 있다.
   • 기초 진동 증가: 작은 초기 크기 $\epsilon$을 주입하면 자기 안정화 피드백 루프를 가속화하여 오버슈트를 줄일 수 있다.

결과

• 추적 성능: CIFAR-10 서브셋 실험에서 에지 플로우는 GD 및 센트럴 플로우의 벌크 역학(bulk dynamics)과 일치한다. 그러나 에지 플로우는 유일하게 샤프니스의 과도적 진동 동작과 EoS 시작 시 관찰되는 특유의 "링잉(ringing)" 현상을 재현한다(그림 4).
• 손실 예측: 모델은 특유의 스파이크를 포함한 훈련 손실 궤적을 정확하게 예측한다.
• 안정화 메커니즘: 결합된 역학은 명시적인 투영 제약 없이 어떻게 시스템이 자연스럽게 $2/\eta$ 임계값 근처에 안착하는지를 성공적으로 보여준다.
• 이산화 효과: 실험 결과, 이산화 해상도($K$)를 높이면 자기 안정화 메커니즘을 변경하지 않고도 손실 스파이크를 제거할 수 있음을 보여주며, 이는 스파이크가 EoS 현상 자체의 본질이 아니라 수치적 아티팩트임을 확인시켜 준다.

의의 및 주장

본 논문은 에지 플로우가 큰 스텝 사이즈 영역에서의 불안정성을 이해하고 완화하기 위한 원칙적인 프레임워크를 제공한다고 주장한다. 중심, 방향, 크기로 역학을 분해함으로써, 복잡한 투영 흐름이나 단단한 막대 모델에 대한 다루기 쉬운 대안을 제시한다.

저자들은 에지 플로우가 충실하며(과도 상태를 포함한 개별 궤적을 추적함), 예측력이 있고(노이즈 주입이나 적응형 스텝 사이즈와 같은 개입을 제안함), 다루기 쉽다(시뮬레이션 비용이 저렴함)는 점을 강조한다. 이 모델이 다음과 같은 현실적인 경로를 열어줄 것이라고 상정한다:

• 큰 학습률에 대한 암묵적 정규화(implicit regularization) 결과의 증명.
• 손실 곡률, 아키텍처, 일반화 사이의 상호작용 이해.
• EoS 역학의 경직된 특성에 맞춤화된 더 효율적인 이산화 방식 설계.

논문은 유도 과정에서 고차 결합 보정을 생략했다는 점을 언급하며 범위에 대해 겸허한 태도를 유지하며, 엄격한 오차 분석 및 확률적 설정(SGD)과 적응형 방법론으로의 확장은 향후 연구 과제로 남겨두었다.