Non-Euclidean Gradient Descent Operates at the Edge of Stability

이 논문은 방향성 매끄러움 (Directional Smoothness) 개념을 비유클리드 노름으로 확장하여 일반화된 날카로움 (sharpness) 지표를 정의함으로써, 다양한 옵티마이저에서도 안정성 한계 (Edge of Stability) 현상이 보편적으로 관찰됨을 이론적으로 해석하고 실험적으로 입증합니다.

핵심

1. 배경: "안정의 가장자리 (Edge of Stability)"란 무엇인가?

가상 세계를 상상해 보세요. 여러분은 눈 덮인 언덕 (손실 함수, Loss Landscape) 위에서 눈썰매 (모델) 를 타고 내려가고 있습니다. 목표는 가장 아래쪽 (최소 오차) 에 있는 보물 상자에 도달하는 것입니다.

• 전통적인 생각: 눈썰매를 탈 때는 너무 빠르게 달리면 넘어지니까, 아주 천천히, 조심스럽게 내려가야 한다고 배웠습니다. (수학적으로 '매끄러운' 조건을 만족해야 함)
• 현실 (딥러닝): 그런데 실제로는 눈썰매를 아주 빠르게 (큰 학습률) 타고 내려가도, 모델은 넘어지지 않고 오히려 가장자리에서 위아래로 살짝 흔들리면서 (진동) 계속 내려갑니다.
• 이 현상의 이름: 이 상태를 '안정의 가장자리 (Edge of Stability, EoS)' 라고 부릅니다. 눈썰매가 넘어질 듯 말 듯 하는 그 위험한 경계선에서 가장 효율적으로 내려가는 것입니다.

기존 연구들은 이 현상이 일반적인 눈썰매 (유클리드 경사 하강법, GD) 에서만 일어난다고 생각했습니다. 하지만 이 논문은 "아니요, 다른 종류의 눈썰매 (비유클리드 방법) 를 타도 똑같은 일이 일어난다!" 고 주장합니다.

2. 새로운 눈썰매들: "비유클리드"란?

우리가 평소에 타는 눈썰매는 땅이 평평하다고 가정하고 미끄러집니다. 하지만 이 논문은 땅이 뾰족하거나, 구불구불하거나, 방향에 따라 마찰이 다른 다양한 지형에서도 눈썰매가 어떻게 작동하는지 연구했습니다.

• $\ell_\infty$-descent (무한 노름): 마치 벽을 타고 기어오르는 등반가처럼, 가장 높은 벽 (가장 큰 오차) 에 집중해서 움직이는 방식입니다.
• Block CD (블록 좌표 하강): 한 번에 한 층씩만 수리하는 건축가처럼, 모델의 특정 부분만 골라서 업데이트합니다.
• Spectral GD (스펙트럴): 회전하는 원반처럼, 데이터의 구조에 맞춰 회전하며 움직이는 방식입니다. (최근 화제의 'Muon' 옵티마이저의 기반이 됩니다.)

이 논문은 이 다양한 눈썰매들도 모두 안정의 가장자리에서 춤을 추며 내려간다는 것을 발견했습니다.

3. 핵심 발견: "날카로움 (Sharpness)"의 재정의

눈썰매가 넘어지지 않으려면 언덕이 너무 가파르면 안 됩니다. 수학적으로 이 가파른 정도를 '날카로움 (Sharpness)' 이라고 부릅니다.

• 기존의 날카로움: "이 언덕이 얼마나 급한가?" (가장 큰 경사도)
• 이 논문의 날카로움 (일반화된 날카로움): "이 눈썰매가 이 지형에서 얼마나 위험한가?"

저자들은 눈썰매의 종류 (노름, Norm) 에 따라 '날카로움'을 다시 정의했습니다.

• 일반적인 눈썰매는 일반적인 날카로움을 봅니다.
• 벽타기 눈썰매 ($\ell_\infty$) 는 벽타기에 적합한 날카로움을 봅니다.
• 회전 눈썰매 (Spectral) 는 회전하는 데 적합한 날카로움을 봅니다.

결론: 어떤 눈썰매를 타든, 모델이 훈련되는 동안 이 '날카로움'이 항상 '2/학습률'이라는 한계선 근처에 머물러 있습니다. 마치 눈썰매가 넘어지지 않으려고 스스로 속도를 조절하듯, 모델이 스스로를 안정화시키는 것입니다.

4. 왜 중요한가요? (일상적인 비유)

이 발견은 **AI 개발자에게 하나의 '만능 나침반'**을 제공합니다.

• 과거: "어떤 옵티마이저 (학습 방법) 를 쓰든, 이 방법이 잘 작동하는지 알기 어렵다. 각각의 규칙을 따로 외워야 했다."
• 이제: "어떤 눈썰매를 타든, 날카로움이 2/학습률 근처에 있으면 그 모델은 잘 훈련되고 있다는 뜻이다!"

이는 마치 다양한 차종 (트럭, 스포츠카, 오토바이) 을 운전할 때, 모두 '엔진 RPM'이 특정 구간을 유지하면 최적의 성능을 낸다는 것을 발견한 것과 같습니다. 이제 개발자는 복잡한 수식을 외울 필요 없이, 이 '날카로움' 지표를 보면 어떤 알고리즘이 안정적으로 작동하는지 한눈에 알 수 있습니다.

5. 요약: 이 논문의 메시지

1. 현상: 딥러닝 모델은 훈련 중 넘어질 듯 말 듯 하는 '위험한 경계선'에서 가장 잘 작동합니다.
2. 확장: 이 현상은 우리가 생각했던 일반적인 방법뿐만 아니라, **기존에 연구되지 않았던 다양한 최신 방법들 (Muon, SignGD 등)**에서도 똑같이 일어납니다.
3. 해결책: 각 방법마다 맞는 '날카로움' (Sharpness) 측정 기준을 만들면, 이 모든 방법이 같은 원리 (안정의 가장자리) 로 작동함을 증명할 수 있습니다.
4. 의미: 이제 우리는 다양한 AI 학습 알고리즘을 하나의 통일된 시선으로 이해하고, 더 나은 모델을 설계할 수 있는 길을 열었습니다.

한 줄 요약:

"AI 모델은 넘어질 듯 말 듯 하는 위험한 경계선에서 가장 잘 달립니다. 이 논문은 어떤 종류의 눈썰매 (학습 방법) 를 타든, 그 모델이 스스로 위험한 경계선을 유지하며 최적의 속도를 찾는다는 것을 증명했습니다."

기술 요약

이 논문은 딥러닝 훈련 중 관찰되는 '안정성의 가장자리 (Edge of Stability, EoS)' 현상을 비유클리드 (Non-Euclidean) 경사 하강법 (Gradient Descent, GD) 으로 확장하여 분석한 연구입니다. 기존 연구가 유클리드 노름 ($\ell_2$) 기반의 표준 GD 나 사전 조건부 (preconditioned) GD 에 국한되어 있었다면, 본 논문은 임의의 노름을 사용하는 일반화된 GD 에서도 EoS 가 발생함을 증명하고 이를 설명할 수 있는 새로운 이론적 프레임워크를 제시합니다.

다음은 논문의 상세한 기술적 요약입니다.

---

1. 문제 제기 (Problem)

• 안정성의 가장자리 (EoS) 현상: 딥러닝 모델 훈련 시, 학습률 ($\eta$) 이 일정 수준 이상일 때 손실 함수가 단조 감소하지 않고 진동하면서도 전체적으로 감소하는 현상이 관찰됩니다. 이때 Hessian 행렬의 최대 고유값 (Sharpness, $\lambda_{max}$) 은 안정성 임계값인 $2/\eta$ 근처로 수렴하거나 이를 약간 상회하며 진동합니다.
• 기존 연구의 한계: EoS 현상은 주로 표준 GD ($\ell_2$ 노름) 나 Adam, Adagrad 같은 적응형 사전 조건부 방법들에 대해 연구되었습니다. 그러나 $\ell_\infty$-descent, Block Coordinate Descent (Block CD), Spectral GD (Muon) 등 비유클리드 노름을 기반으로 하거나 다른 기하학적 구조를 가진 최적화 알고리즘들에서는 EoS 현상이 어떻게 작동하는지, 그리고 이를 설명할 수 있는 '일반화된 Sharpness'가 무엇인지에 대한 이론적 이해가 부족했습니다.
• 핵심 질문: 다양한 노름을 사용하는 비유클리드 GD 에서도 EoS 현상이 발생하는가? 만약 그렇다면, 이를 설명하기 위한 적절한 'Sharpness'의 정의와 이론적 근거는 무엇인가?

2. 방법론 (Methodology)

저자들은 방향성 매끄러움 (Directional Smoothness) 개념을 비유클리드 노름으로 확장하여 EoS 를 해석합니다.

• 비유클리드 GD 정의:
  임의의 노름 $\|\cdot\|$에 대해, 다음 식으로 정의되는 일반화된 GD 를 다룹니다.
  $$w_{t+1} = \arg \min_y \langle \nabla L(w_t), y - w_t \rangle + \frac{1}{2\eta}\|y - w_t\|^2$$
  이는 $\ell_2$ 노름일 때 표준 GD 가 되며, $\ell_\infty$ 노름일 때 SignGD, 스펙트럴 노름일 때 Muon 등으로 일반화됩니다.

• 방향성 매끄러움 (Directional Smoothness, $D_{\|\cdot\|}$):
  Mishkin et al. [2024] 의 개념을 차용하여, 두 연속된 반복점 $w_t, w_{t+1}$ 사이의 평균 곡률을 나타내는 양을 정의합니다.
  $$D_{\|\cdot\|}(w, y) = \frac{L(y) - L(w) - \langle \nabla L(w), y - w \rangle}{\frac{1}{2}\|y - w\|^2}$$
  이론적으로 손실이 감소하려면 $D_{\|\cdot\|} \leq 2/\eta$여야 하며, 손실이 진동하면 $D_{\|\cdot\|}$는 $2/\eta$ 주변에서 진동하게 됩니다.

• 일반화된 Sharpness ($S_{\|\cdot\|}$) 정의:
  방향성 매끄러움과 Hessian 의 관계를 통해 임의의 노름 $\|\cdot\|$에 대한 일반화된 Sharpness 를 정의합니다.
  $$S_{\|\cdot\|}(w) := \max_{d \neq 0} \frac{d^\top \nabla^2 L(w) d}{\|d\|^2} = \max_{\|d\| \leq 1} d^\top \nabla^2 L(w) d$$

  • 이 정의는 $\ell_2$ 노름에서는 기존 Hessian 최대 고유값과 일치합니다.
  • $\ell_\infty$ 노름에서는 Ising 스핀 글래스의 최대 에너지 문제와 유사한 NP-hard 문제를 포함하며, Frank-Wolfe 알고리즘을 사용하여 근사합니다.
  • Block $\ell_{1,2}$ 노름이나 Spectral 노름 ($\|\cdot\|_{2\to2}$) 에 대해서도 각각의 기하학적 구조에 맞는 Sharpness 를 정의합니다.

• 정규화된 비유클리드 GD:
  SignGD 나 Muon 과 같이 노름으로 정규화된 업데이트 규칙 ($w_{t+1} = w_t - \eta (\nabla L(w_t))^*$) 에 대해서도 유효한 step-size $\tilde{\eta} = \eta / \|\nabla L(w_t)\|_*$를 도입하여 동일한 EoS 현상이 발생함을 보였습니다.

3. 주요 기여 (Key Contributions)

1. 방향성 매끄러움을 통한 EoS 해석: EoS 현상을 Hessian 의 최대 고유값이 아닌, 방향성 매끄러움의 관점에서 해석하고, 이것이 비유클리드 노름으로 자연스럽게 확장됨을 보였습니다.
2. 일반화된 Sharpness 의 정의: 임의의 노름에 대해 Sharpness 를 정의하고, 이것이 손실 감소 조건 ($D_{\|\cdot\|} \leq 2/\eta$) 과 직접적으로 연결됨을 증명했습니다.
3. 다양한 최적화 알고리즘에 대한 EoS 확장: MLP, CNN, Transformer 아키텍처를 대상으로 실험하여, $\ell_\infty$-descent, Block CD, Spectral GD (Muon) 등 기존 연구에서 다루지 않았던 방법들에서도 일반화된 Sharpness 가 $2/\eta$ 임계값 주변에서 진동하며 EoS 현상을 보임을 실증했습니다.
4. 이차 함수 (Quadratic) 에 대한 이론적 분석: 비유클리드 GD 가 이차 함수에서 $2/S$ (S 는 일반화된 Sharpness) 보다 큰 step-size 를 가질 때 발산할 수 있음을 증명하여, EoS 가 국소적인 이차 근사에서의 불안정성과 관련 있음을 이론적으로 뒷받침했습니다.

4. 실험 결과 (Results)

• 진행적 Sharpening 및 진동: 모든 실험 (Vanilla GD, $\ell_\infty$-descent, Block CD, Spectral GD) 에서 초기에는 Sharpness 가 증가하다가 (Progressive Sharpening), 이후 $2/\eta$ 임계값 근처에서 진동하거나 약간 상회하는 EoS 단계를 거쳤습니다.
• 일반화된 Sharpness vs 표준 Sharpness:
  • $\ell_\infty$-descent 나 Spectral GD 의 경우, 표준 $\ell_2$ Sharpness는 $2/\eta$보다 훨씬 낮게 유지되는 반면, **일반화된 Sharpness**는 명확하게$2/\eta$ 근처에 수렴했습니다. 이는 EoS 현상이 알고리즘이 사용하는 기하학적 구조 (노름) 에 맞춰 정의된 Sharpness 로 설명되어야 함을 시사합니다.
• Frank-Wolfe 알고리즘의 역할: $\ell_\infty$ 및 Block $\ell_{1,2}$ 노름에서의 Sharpness 계산은 NP-hard 문제이므로 Frank-Wolfe 알고리즘을 사용했습니다. 실험 결과, 충분한 수의 재시작 (restarts) 을 거치면 Sharpness 추정이 안정화되어 $2/\eta$ 근처 값을 잘 포착함을 확인했습니다.
• 예외적인 진동 구간: 일부 비유클리드 방법 (특히 $\ell_\infty$) 의 경우, Sharpness 가 $2/\eta$에 도달하기 전에도 방향성 매끄러움이 증가하며 진동이 시작되는 'EoS 이전의 진동 구간'이 관찰되었습니다. 이는 유클리드 GD 에서는 나타나지 않는 새로운 현상입니다.

5. 의의 및 결론 (Significance)

• 통일된 프레임워크 제시: 본 논문은 다양한 최적화 알고리즘 (GD, Adagrad, Adam, SignGD, Muon 등) 에 대한 EoS 현상을 단일한 기하학적 프레임워크 (일반화된 Sharpness) 로 통합하여 설명합니다.
• 새로운 최적화 알고리즘에 대한 통찰: Muon 과 같은 최신 최적화 알고리즘이 왜 잘 작동하는지에 대한 이론적 근거를 제공하며, EoS 현상이 특정 알고리즘에 국한된 것이 아니라 비유클리드 기하학을 사용하는 광범위한 최적화 방법의 보편적 특성임을 보여줍니다.
• 이론적 한계와 향후 과제:
  • 유클리드 GD 에서는 초기값이 거의 모든 경우에서 발산하지만, 비유클리드 GD 에서는 특정 초기값에서만 발산이 증명되어 이론적 격차가 존재합니다.
  • EoS 이전의 진동 구간 (Pre-EoS oscillatory regime) 의 정확한 메커니즘은 아직 완전히 규명되지 않았습니다.

결론적으로, 이 연구는 딥러닝 최적화의 핵심 현상인 EoS 가 단순히 Hessian 의 최대 고유값이 아니라, 최적화 알고리즘이 사용하는 노름 공간의 기하학적 구조에 정의된 일반화된 Sharpness에 의해 결정됨을 밝혔습니다. 이는 향후 더 효율적이고 안정적인 비유클리드 최적화 알고리즘을 설계하는 데 중요한 이론적 토대를 제공합니다.