Conservation Laws for Modern Neural Architectures

기술 요약: 현대적 신경망 구조를 위한 보존 법칙 (Conservation Laws)

문제 정의
과잉 매개변수화된 신경망(over-parameterized neural networks)의 훈련 역학을 이해하는 것은 이들의 성공을 설명하는 데 핵심적이다. 이러한 이해를 위한 주요 메커니즘은 보존 법칙(conservation laws) 개념이다. 이는 최적화 궤적(특히 경사 흐름, gradient flow)을 따라 불변하는 양을 의미한다. 이러한 법칙들은 선형 및 ReLU 네트워크에 대해서는 잘 규명되어 있으나, 현대적 아키텍처에 대해서는 여전히 미개척 영역으로 남아 있다. 기존 프레임워크는 부드러운 활성화 함수(예: GELU, SiLU), 복잡한 어텐션 메커니즘(위치 인코딩이 포함된 Multihead Attention), 그리고 Mixture-of-Experts(MoE) 설계를 사용하는 현대적 모델들에 대해 완전한 특성화를 제공하는 데 어려움을 겪고 있다. 또한, Multihead Attention(MHA)에 관한 기존 연구는 단일 헤드 설정에 국한되어 있어, 멀티 헤드 케이스는 미해결 문제로 남아 있었다.

방법론
본 논문은 현대적 딥러닝 모델을 위한 보존 법칙을 규명하기 위해 통합된 이론적 프레임워크를 개발한다. 이 접근 방식은 Marcotte 등(2023; 2025)의 연구를 기반으로 하지만, 분석 전략 측면에서 차별화된다:

1. 문제 정식화: 저자들은 보존 법칙 $h(\theta)$를 유클리드 경사 흐름 $\dot{\theta} = -\nabla L_D(\theta)$ 하에서 불변인 함수로 정의한다. 문제를 다룰 수 있는 수준으로 만들기 위해, 저자들은 보존 법칙이 모든 데이터셋에 대해 성립해야 한다는 가정을 채택하여, 아키텍처와 최적화 역학에 의해 유도되는 암묵적 편향(implicit bias)만을 분리해 낸다.
2. PDE로의 환원: 저자들은 리 군(Lie-theoretic) 기법에 의존하는 대신, 보존 조건을 편미분 방정식(PDE) 시스템으로 환원한다. 구체적으로, 보존 법칙의 기울기 $\nabla_\theta h$가 모든 입력에 대해 모델의 자코비안 벡터(Jacobian vectors)가 생성하는 스팬(span)과 직교할 것을 요구한다.
3. 복소해석적 기법: 증명 전략의 핵심은 직교 조건을 복소수 값 함수로 간주하는 것이다. 저자들은 이러한 함수들을 특정 부분 공간으로 제한하고, 메로모르픽 함수(meromorphic functions)의 항등 정리(identity theorem)를 활용함으로써, $h$의 기울기에 대한 입력 독립적인 제약 조건을 도출한다.
4. 불변량을 통한 특성화: 도출된 PDE 제약 조건은 "특성 불변량(characteristic invariants)"을 식별하기 위해 풀린다. 보존 법칙은 이 불변량들에 의해 정의된 레벨 셋(level sets)의 연결된 성분(connected components) 상에서 상수임을 보인다.

주요 기여 및 결과
본 논문은 세 가지 주요 현대적 아키텍처 클래스에 대한 보존 법칙의 완전한 특성화를 제공한다:

1. 피드포워드 네트워크 (FFNs):

   • GELU 및 SiLU: 저자들은 GELU 또는 SiLU 활성화 함수를 사용하는 FFN의 경우, 유일한 보존 법칙은 상수 함수뿐임을 증명한다. 비자명한 불변량은 존재하지 않는다.
   • SwiGLU: 이와 대조적으로, SwiGLU 아키텍처(게이팅 메커니즘 내에서 행렬 간의 곱셈적 상호작용을 포함함)는 비자명한 보존 법칙을 나타낸다. 불변량은 각 은닉 단위 $i$에 대해 $\|A_{:,i}\|^2 - \|C_{i,:}\|^2$인 양에 의해 결정된다 (여기서 $A$와 $C$는 게이팅 메커니즘에 관여하는 가중치 행렬이다).

2. 멀티헤드 어텐션 (MHA):

   • Vanilla MHA: 저자들은 멀티 헤드 어텐션의 보존 법칙을 규명하는 미해결 문제를 해결하며, 보존 법칙이 각 헤드 $i$에 대한 불변량 $Q_i^\top Q_i - K_i^\top K_i$ 및 $V_i^\top V_i - O_i^\top O_i$에 의해 특징지어짐을 확인한다.
   • 사인형 위치 인코딩 (Sinusoidal PE): 사인형 PE는 입력에 대한 고정된 생체적(bijective) 이동으로 작용하므로, 보존 구조는 Vanilla MHA 케이스와 동일하게 유지된다.
   • 회전 위치 인코딩 (RoPE): RoPE는 어텐션의 내부 구조를 근본적으로 변화시킨다. 저자들은 별개의 불변량 집합인 $\|Q_i^{(j)}\|_F^2 - \|K_i^{(j)}\|_F^2$ (여기로 $Q$와 $K$가 $2 \times 2$ 블록으로 분해됨) 및 $V_i^\top V_i - O_i^\top O_i$를 도출한다. 이는 RoPE가 절대적 위치 인코딩과 다른 클래스의 보존 법칙을 도입함을 보여준다.

3. Mixture-of-Experts (MoE):

   • Dense 및 Sparse MoE: 본 연구는 밀집 게이팅(dense gating)과 Top-$k$ 라우팅을 사용하는 희소 MoE(Sparse MoE, SMoE)를 모두 다룬다. 결과에 따르면 SMoE의 보존 법칙은 그에 대응하는 Dense MoE의 보존 법칙과 일치한다.
   • 게이팅 메커니즘: 불변량은 전문가(expert) 수준에서 국소화되며 (각 전문가에 대한 SwiGLU 불변량을 따름), 게이팅 가중치 벡터의 합 $\sum W_i$에 대한 추가적인 불변량을 포함한다.
   • Sigmoid 게이팅: 저자들은 Softmax 게이팅을 정규화된 Sigmoid 게이팅으로 대체해도 보존 법칙의 집합이 변하지 않음을 입증한다. 즉, 불변량은 Softmax 게이팅 케이스와 동일하다.

실험적 검증
이론적 결과는 언어 모델링(WikiText-103, Penn Treebank) 및 컴퓨터 비전(CIFAR-10, ImageNet-1K) 태스크를 통해 경험적으로 검증되었다.

• 설정: 실험에는 연속 시간 흐름(continuous-time flow)을 근사하기 위한 풀 배치 경사 하강법(full-batch gradient descent)과 소규모 및 대규모 데이터셋에 대한 미니 배치 SGD가 사용되었다.
• 지표: 저자들은 훈련 단계에 따른 "보존 오차"(불변량이 초기값으로부터 벗어나는 정도)를 추적한다.
• 결과: 실험은 보존 법칙이 훈련 중에도 유계(bounded) 상태를 유지함을 확인한다. 구체적으로, 오차는 이론이 예측하는 바와 같이 학습률 $\tau$와 단계 수 $k$에 따라 $O(\tau^2 k)$ (일정한 스텝 사이즈의 경우)로 스케일링된다. 결과는 데이터셋이나 훈련 방식에 관계없이 MHA, SwiGLU FFN, RoPE, MoE 게이팅 메커니즘 전반에서 일관되게 나타났다.

의의 및 주장
본 논문은 광범위한 현대적 신경망 아키텍처에 대한 보존 법칙의 첫 번째 완전한 특성화를 제공함으로써 훈련 역학에 대한 이론적 이해를 진전시켰다고 주장한다.

• 이론적 통합: 선형/ReLU 네트워크에서 부드러운 활성화 함수 및 복잡한 어텐션 메커니즘으로 확장된 통합된 프레임워크를 구축한다.
• 미해결 문제 해결: 멀티 헤드 어텐션의 보존 법칙을 규명하는 미해결 문제를 명시적으로 해결하였으며, 위치 인코딩 및 MoE 게이팅으로 분석을 확장하였다.
• 암묵적 편향: 이러한 불변량을 식별함으로써, 본 연구는 현대적 아키텍처에 내재된 기하학적 제약을 밝혀내며, 초기화부터 수렴까지 지속되는 속성을 드러낸다.
• 한계점: 저자들은 복소해석적 기법(메로모르픽 확장)에 대한 의존이 현재 레이어 정규화(layer normalization)나 $\sqrt{x}$, $e^{1/x}$와 같은 함수를 포함하는 깊은 합성(deep compositions)을 포함하는 비-메로모르픽 구성 요소를 가진 아키텍처에 대한 프레임워크 적용을 제한한다는 점을 겸허히 언급한다. 또한, 아키텍처의 편향을 분리하기 위한 원칙적인 접근이지만, 더 넓은 맥락에서는 제한적일 수 있는 데이터셋 독립성 가정에 대해서도 인정한다.