기술 요약: 곡선 가중치 공간에서의 학습 (SymExpLin)

1. 문제 정의

표준 신경망 훈련은 선형 파라미터 공간에서의 가산적 업데이트(additive updates)에 의존합니다. 그러나 많은 내부 네트워크 연산은 본질적으로 곱셈적(multiplicative)입니다 (예: 특징 진폭의 스케일링, 정규화 게인 또는 특이값). 선형 공간에서는 동일한 크기의 상대적 변화(예: 가중치를 2배로 늘리는 것과 절반으로 줄이는 것)가 불균등한 최적화 거리를 요구합니다. 예를 들어, 가중치를 1.0에서 2.0으로 이동하려면 1.0의 스텝 크기가 필요하지만, 1.0에서 0.5로 이동하려면 0.5의 스텝이 필요합니다.

이러한 불일치는 다음과 같은 최적화 문제를 야기합니다:

• 헤비 테일 분포 (Heavy-Tailed Distributions): 사전 훈련된 모델들은 종종 헤비 테일 가중치 분포를 보입니다. 즉, 소수의 파라미터가 초기화 스케일보다 훨씬 더 크게 성장하는 반면(예: 30배), 대부분은 초기값 근처에 머물러 있습니다.
• 옵티마이저 불변성 (Optimizer Invariance): Adam과 같은 적응형 옵티마이저는 좌표별로 업데이트를 정규화하여 스텝 크기를 파라미터 크기에 무관하게 만듭니다. 결과적으로, 작은 파라미터(예: 0.004)와 큰 파라미터(예: 0.4)는 유사한 절대적 업데이트를 받게 됩니다. 이는 100배의 상대적 섭동(perturbation) 차이를 발생시켜, 작은 파라미터는 과도하게 섭동되는 반면, 크고 중요한 파라미터는 과소 섭동될 가능성을 높입니다.
• 정적 vs 동적 스케일링 (Static vs Dynamic Scaling): Maximal Update Parameterization (µP)와 같은 기존 솔루션은 초기화 시점의 텐서 너비에 따라 정적으로 스케일 불일치를 해결하지만, 훈련 중 개별 파라미터의 동적인 성장을 반영하지는 못합니다.

2. 방법론: SymExpLin (SEL)

저자들은 대칭-지수 경로(symmetric-exponential pathway)와 선형 경로(linear pathway)로 유효 가중치를 분해하는 비선형 가중치 재매개변수화 기법인 **SymExpLin (SEL)**을 제안합니다. 이를 통해 원시 공간(raw space)에서의 가산적 업데이트가 유효 가중치 공간에서는 크기에 비례하는 변화로 매핑되는 곡선 파라미터 공간 기하학을 생성합니다.

2.1 변환 (The Transform)

유효 가중치 $w_{eff}$는 다음 구성 요소들을 사용하여 원시 가중치 $w$로부터 계산됩니다:

1. 대칭-지수 경로 ($E_\theta$): 대칭 지수 함수(symexp)를 기반으로 한 학습 가능한 부호 보존 함수입니다. $|w|$가 작을 때는 선형적으로 동작하며, $|w|$가 클 때는 지수적으로 성장합니다. 이를 통해 큰 원시 스텝이 더 큰 유효 변위(displacement)를 생성할 수 있게 합니다.
   • $E_\theta(|w|) = e_w (\exp(\beta |w| \cdot m - n) - \exp(-n))$
   • 파라미터: $\beta$ (곡률), $m$ (내부 승수), $n$ (원점에서의 기울기 오프셋), $e_w$ (출력 스케일).
2. 선형 경로 ($L_\theta$): 직접적인 그래디언트 경로를 제공하여 최적화를 안정화하는 것으로 가설 세워진 항등 함수 형태의 경로 $L_\theta(w) = w \cdot l_w$입니다.
3. 부호 래핑 (Sign Wrapping): 원시 가중치의 부호가 경로들의 합에 적용됩니다: $w_{eff} = \text{sign}(w) \cdot (E_\theta(|w|) + L_\theta(w)) + \Delta_\phi$.
4. 잔차 ($\Delta_\phi$): 선택적인 저계수(low-rank, LoRA 스타일) 가산 보정항입니다.

2.2 주요 설계 선택 사항

• 불일치 초기화 (Mismatched Initialization): 저자들은 "합치된(congruent)" 조합 규칙(선형 및 지수 부분이 부호 래핑 외부에서 더해지는 방식)을 역전시켜 원시 가중치를 초기화하는 방식과, 순전파 시 부호 래핑 내부에 두 경로를 포함하는 방식 사이의 유익한 비대칭성을 발견했습니다. 이는 음수 가중치가 양수 가중치보다 작은 크기로 초기화되도록 하여, 조기 대칭성 깨기(symmetry-breaking) 메커니즘으로 작용합니다.
• 구조화된 스케일 파라미터: 학습 가능한 스케일($e_w, l_w, m, n$)은 표현력과 계산 비용의 균형을 맞추기 위해 행-열 곱셈 패턴(rank-1 벡터)을 사용하여 매개변수화됩니다.

3. 학습률 어닐링 (Learning Rate Annealing): 스케일 파라미터를 위한 타겟 스케줄이 사용됩니다. 외부 스케일 $e_w$는 부호를 빠르게 결정하기 위해 높은 학습률로 시작하며, 내부 승수 $m$은 낮게 시작하여 스케일이 안정화된 후 곡률을 정교화하기 위해 점진적으로 상승합니다.
4. 곡률 교정 (Curvature Calibration): 곡률 파라미터 $\beta$는 모델 너비에 따라 휴리스틱하게 설정되어, 지수 영역이 전형적인 초기화 크기(너비가 증가함에 따라 작아짐) 근처에서 활성화되도록 보장합니다.

3. 주요 기여

• SymExpLin (SEL): 가중치를 독립적으로 학습 가능한 구조화된 스케일을 가진 대칭-지수 및 선형 성분으로 분해하는, 모든 학습 가능한 행렬 또는 바이어스에 적용 가능한 가중치 재매개변수화 기법입니다.
• 불일치 초기화: 초기화 역전 과정과 순전파 과정에서 서로 다른 조합 규칙을 사용하는 것이 조기 최적화를 개선하는 유익한 비대칭 편향을 만든다는 것을 발견했습니다.
• 구조화된 스케일 및 어닐링: 중복된 지수 제어를 관리하기 위해 행-열 곱셈 스케일 파라미터와 특정 학습률 어닐링 스케줄을 도입했습니다.
• 실용성: 이 방법은 구현이 용이합니다. 훈련 변환은 요소별(element-wise)로 수행되며, 훈련 후에는 유효 가중치가 계산되어 표준 선형 레이어에 통합되므로 추론 비용이 발생하지 않습니다.
• 상보성: SEL은 $\mu$P와 같은 정적, 텐서 단위 메서드를 보완하며, 동적인 파라미터별 축에서 작동합니다.

4. 실험 결과

본 방법은 OpenWebText를 사용하여 너비 $\{1024, 2048, 3072\}$ 및 깊이 $\{12, 24, 36\}$의 9가지 구성에 대해 자기회귀 언어 모델링에서 평가되었습니다.

• 손실 하강 (Loss Descent): SEL은 모든 구성에서 베이스라인(표준 파라미터화를 사용한 Adam)보다 일관되게 낮은 검증 손실에 도달했습니다.
• 스텝 효율성: SEL은 베이스라인 대비 1.32배에서 1.49배 적은 훈련 스텝만으로 동일한 검증 손실을 달고 달성했습니다. 이 이득은 더 큰 너비의 모델(예: 3072×24에서 1.49배 가속)에서 가장 두드러졌습니다.
• 실제 시간 가속 (Wall-Clock Speedup): 변환으로 인한 스텝당 오버헤드(작은 모델 1.7% ~ 큰 모델 5.5%)에도 불구하고, 전체 스텝 수의 감소 덕분에 추정된 실제 시간 가속도는 1.29배에서 1.37배였습니다.
• 절제 연구 (Ablations):
  • 선형 경로: 선형 경로를 제거하면 "학습 실패" 사례가 발생하여, 안정성을 위한 선형 경로의 필요성을 확인했습니다.
  • 불일치 vs 합치: 불일치 초기화가 합치 초기화보다 특히 작은 모델에서 일관되게 우수한 성능을 보였습니다.
  • 어닐링: 스케일 파라미터를 위한 어닐링된 학습률 스케줄은 고정된 승수 대비 일관되지만 작은 개선을 제공했습니다.

5. 의의 및 주장

본 논문은 SEL이 많은 신경망 연산의 곱셈적 성격과 표준 최적화의 가산적 성격 사이의 근본적인 불일치를 해결한다고 주장합니다. 파라미터 공간을 곡선화함으로써, SEL은 옵티마이저가 가중치 크기에 자연스럽게 비례하는 스텝을 밟을 수 있게 하며, 이는 별도의 파라미터별 학습률 튜닝 없이도 상대적 스케일링을 효과적으로 제공합니다.

저자들은 SEL이 다음과 같은 특징을 가진 표준 선형 레이어의 실용적인 드롭인(drop-in) 대체제임을 강조합니다:

1. 특히 헤비 테일 가중치 분포를 가진 모델의 수렴을 가속화합니다.
2. 추론 시 아키텍처 변경이 필요 없습니다.
3. 역전이 가능하여, 표준 사전 훈련된 가중치를 SEL 공간으로 변환하여 지속적인 훈련이나 미세 조정을 할 수 있습니다.

이 연구는 "소프트 로그 공간(soft log-space)" 파라미터화가 선형 가산 업데이트로는 도달하기 어려운, 네트워크 계산에 필요한 큰 크기의 목표치에 가중치가 도달하도록 돕는다는 점을 시사합니다.