기술 요약: 안정성의 가장자리(Edge of Stability)가 데이터 분포 전반에 걸쳐 학습을 선택적으로 형성하는 방식

1. 문제 제기

기존의 안정성의 가장자리(Edge of Stability, 이하 EoS)에 관한 문헌들은 이를 주로 최적화의 전역적(global) 특성으로 다루어 왔으며, 여기서 최상위 헤시안 고유값(top Hessian eigenvalue)은 학습률 $\eta$의 임계값인 $2/\eta$ 근처에서 자기 안정화(self-stabilize)된다. 부호가 교차하는 진동(sign-alternating oscillations)과 느린 교정 드리프트(slow corrective drift)로 특징지어지는 이 체제(regime)의 역학은 잘 기록되어 있으나, 이 임계값에서 작동하는 것이 기능적으로 어떤 결과를 초래하는지는 불분명하다. 구체적으로, 안정성 제약이 균일한 이득을 제공하는지, 아니면 훈련 데이터의 서로 다른 하위 집합들에 대해 차별적으로 학습을 형성하는지는 알려져 있지 않다. 핵심 질문은 다음과 같다: EoS 체제가 선택적 학습 편향(selective learning bias)을 유도하는가? 만약 그렇다면, 어떤 기하학적 및 역학적 특성이 특정 데이터 하위 집합의 이득을 결정하는가?

2. 방법론

2.1 실험 설정 및 프로토타입 분류 체계

저자들은 CIFAR-10(자동차 vs 트럭)에 대한 이진 분류 태스크와 풀 배치 경사 하강법(full-batch gradient descent)으로 훈련된 두 개의 은닉층 MLP를 사용하여 통제된 실험 프레임워크를 구축한다. 데이터 기하학의 효과를 격리하기 위해, 저자들은 훈련 모델과는 독립적으로 클래스 중심점($\mu_c$)에 대한 위치와 라벨 일관성을 기준으로 훈련 분포를 네 가지 그룹으로 나누는 **프로토타입 분류 체계(prototype taxonomy)**를 정의한다:

1. 인라이어(Inliers): 클래스 중심점 근처의 고밀도 지점으로, 올바른 라벨을 가짐.
2. 경계(Boundary): 클래스 간 경계 근처의 분포 내 지점(높은 라벨 모호성).
3. 입력 아웃라이어(Input-outliers): 기하학적으로 이례적인 입력($\mu_c$에서 멀리 떨어짐)이며 올바른 라벨을 가짐.
4. 출력 아웃라이어(Output-outliers): 잘못된 라벨이 할당된 고밀도 입력(합성적으로 생성됨).

2.2 분기 개입 (인과적 식별)

EoS 체제에 머무는 것과 이를 벗어나는 것의 효과를 인과적으로 식별하기 위해, 저자들은 **분기 개입(branching intervention)**을 사용한다:

• 베이스라인 분기(Baseline Branch): 원래의 학습률 $\eta$로 훈련을 계속하여 EoS 체제에 머무름.
• 탈출 분기(Exit Branch): EoS 발생 시점(최상위 헤시안 고유값 $\lambda_1$이 $2/\eta$에 도달했을 때)에 학습률을 절반으로 줄임($\eta \to \eta/2$). 이는 안정성 임계값을 $4/\eta$로 높여 최적화 도구가 EoS 체제를 벗어나게 한다.
• 비교: 개입 후 두 분기 사이의 하위 그룹별 손실($\ell_k$)의 발산을 비교함으로써, 저자들은 EoS 안정성 제약의 인과적 영향을 격리한다.

2.3 메커니즘 격리

저자들은 특정 하위 집합이 EoS로부터 이득을 얻기 위한 두 가지 가설적 필요 조건을 테스트한다:

1. 방향 정렬(Directional Alignment): 하위 집합의 집계된 그래디언트(aggregate gradient)가 최상위 헤시안 고유벡터($v_1$)와 정렬되어야 한다. 이는 변위 거리는 유지하되 방향을 무작위화하여 일관된 정렬을 파괴한 "무작위 방향" 입력 아웃라이어를 구성하여 테스트된다.
2. 그래디언트 지속성(Gradient Persistence): 하위 집합이 소멸하지 않는 그래디언트 크기를 유지해야 한다. 이는 MSE 손실(그래디언트가 지속됨)과 CE 손실(확신이 높은 지점에서 그래디언트가 소멸함)을 비교함으로써 테스트된다.

3. 주요 결과

3.1 선택적 학습 트레이드오프

분기 개입 결과, EoS는 전역적이지 않고 선택적임이 드러났다. 안정성 제약은 최적화 노력을 불균등하게 재분배한다:

• 수혜자(Beneficiaries): 그래디언트가 $v_1$과 정렬되고 지속되는 하위 집합(예: MSE 하의 입력 아웃라이어, CE 하의 출력 아웃라이어)은 탈출 분기에 비해 가속화된 손실 감소를 보인다.
• 억제 대상(Suppressors): 직교하는 그래디언트나 소멸하는 크기를 가진 하위 집합(예: 특정 조건 하의 인라이어 및 경계 지점)은 탈출 분기에 비해 EoS 분기에서 진행이 억제된다.
• 이러한 트레이드오프는 다양한 아키텍처(CNN, ResNet), 옵티마이저(SGD, Momentum), 클래스 쌍에 걸쳐 재현된다.

3.2 메커니즘: 정렬 $\times$ 지속성

두 조건이 EoS 이득을 포착하기 위한 공동의 필요 조건이다:

1. 정렬(Alignment): 그룹의 집계된 그래디언트는 최상위 헤시안 고유벡터($v_1$)와 정렬되어야 한다.
   • 증거: 입력 아웃라이어를 (거리 대신 정렬을 파괴하도록) 무작위 방향으로 이동시켰을 때, 곡률 영향 $(\nabla \ell_k \cdot v_1)^2$이 한 자릿수(order of magnitude)만큼 붕괴하며 EoS 이득이 사라진다.
2. 지속성(Persistence): 그룹은 소멸하지 않는 그래디언트 크기를 유지해야 한다.
   • 증거: Cross-Entropy 손실 하에서, 확신이 높은 입력 아웃라이어의 그래디언트는 포화(소멸)된다. 결과적으로 EoS 이득은 라벨 불일치로 인해 그래디언트가 활성화되어 있는 출력 아웃라이어로 이동한다.

3.3 기하학이 수혜자를 결정한다

이득을 얻는 하위 집합은 데이터 분포의 기하학적 구성에 의해 결정된다:

• 표준 CIFAR-10 설정에서는 입력 아웃라이어(중심점에서 가장 먼 지점)가 $v_1$을 지배하고 이득을 얻는다.
• 더 어려운 클래스 쌍(고양이 vs 개)이나 다른 아웃라이어 생성 파라미터를 사용하는 경우, 경계 지점이 가장 먼 그룹이 되어 $v de 1$을 지배하고 주요 수혜자가 될 수 있다.
• 자연 데이터(Natural Data): 자연스러운 CIFAR-10 분포에서, EoS 발생 시점에 중심점 거리와 $v_1$과의 정렬 사이에 양의 상관관계가 나타나며, 이는 기하학적 아웃라이어가 자연스럽게 불안정한 모드(unstable mode)를 주도함을 확인시켜 준다.

3.4 일반화 및 강건성 함의

본 논문은 훈련 시의 선택성이 테스트 시의 행동으로 전이된다는 예비 증거를 제시한다:

• 적대적 강건성(Adversarial Robustness): 경계 지점이 $v_1$을 지배할 때만 EoS가 경계 지점에 대한 강건성을 개선한다. 만약 입력 아웃라이어가 지배한다면, 최적화 예산이 전용되기 때문에 EoS 분기는 경계 강건성 측면에서 더 나쁜 성능을 보인다.
• OOD 일반화: 외삽된 입력(extrapolated inputs)이 훈련 중에 $v_1$을 지배할 때만 EOS가 OOD 성능을 개선한다.

4. 이론적 프레임워크

저자들은 셀프 스테빌라이제이션 프레임워크(Damian et al., 2023)를 전역 손실에서 하위 집합별 손실로 확장한다. 그들은 EoS 체제에서의 주기 평균 드리프트(cycle-averaged drift)에 $-\nabla S$ (여기서 $S$는 sharpness)에 비례하는 추가 항이 포함됨을 유도한다. 하위 집합 $k$에 대한 차등적 효과는 선택자(selector)에 의해 지배된다:
$$Q_k = \langle \nabla \ell_k, \nabla S \rangle$$
단일 모드 근사(single-mode approximation) 하에서, 이 선택자는 최상위 고유벡터에 대한 하위 집합 그래디언트의 제곱 투영에 비례한다:
$$C_k = (\nabla \ell_k \cdot v_1)^2 = \|\nabla \ell_k\|^2 \cos^2 \theta_k$$
이 분해는 방향 정렬($\cos^2 \theta_k$)과 그래디언트 크기 지속성($\|\nabla \ell_k\|^2$)이 모두 하위 집합이 EoS 드리프트를 경험하는 데 필요함을 확인해 준다.

5. 의의 및 주장

본 논문은 EoS에 대한 이해를 단순한 스칼라적 안정성 경계에서 데이터 분포 전반에 걸친 학습 할당을 조절하는 메커니즘으로 전환한다고 주장한다.

• 평탄함(Flatness)의 재정의: 저자들은 "평탄함"이 솔루션의 스칼라적 속성이 아니라 데이터 기하학에 의해 결정되는 방향적 속성이라고 주장한다. 낮은 곡률의 기능적 이득은 어떤 하위 집합이 최상위 헤시안 고유 방향을 지배하느냐에 전적으로 달려 있다.
• 귀납적 편향(Inductive Bias): EoS는 단순히 동등한 최소화 지점들 사이를 선택하는 것이 아니라, 데이터 기하학적 관계를 기반으로 특정 하위 집합을 최적화하도록 하는 귀납적 편향으로 작용한다.
• 관점의 통합: 이 연구는 파라미터 공간에서의 암묵적 정규화(곡률 제약)와 데이터 분포에 대한 귀납적 편향(어떤 예시가 학습되는가)을 연결한다.
• 범위의 제한: 저자들은 헤시안-벡터 곱(Hessian-vector product)의 계산 비용으로 인해 실험이 작은 모델과 풀 배치 훈련에 국한되었다는 점을 인정하며, 이러한 선택적 효과가 더 큰 데이터셋, 미니 배치 옵티마이저, 또는 표현 공간(representation-space)의 프로토타입에서도 지속될지는 미결 과제로 남겨두었다.

결론적으로, 본 논문은 안정성의 가장자리(EoS)가 최상위 헤시안 고유 방향과 정렬되고 지속되는 하위 집합의 학습을 증폭하는 동시에 다른 집합은 억제하는 선택적 체제임을 보여줌으로써, 모델의 기능적 특성을 근본적으로 형성한다는 것을 입증한다.