선형 분해 블록을 위한 기하학 기반 AdamW 최적화
Stiefel-AdamW: Geometry-Aware AdamW for Linear Factorization Blocks
arXiv최신 딥러닝 모델에서 흔히 사용되는 선형 분해 블록($W=BA$)은 구조적 비유일성 때문에 학습이 불안정해지는 문제가 있었습니다.
- Stiefel-AdamW는 두 인자 중 하나를 스티플 다양체에 제약하여 요인 폭발을 막고, 기존 AdamW의 실용성과 리만 기하학적 안정성을 결합한 최적화 알고리즘입니다.
- LoRA나 트랜스포머 모델 등 다양한 구조에 적용 가능하며, 추가 비용 없이 강력한 기준 성능 대비 일관된 개선 효과를 입증했습니다.
- 구현 오버헤드가 매우 적으면서도 리만 다양체의 안정성과 표준 최적화 알고리즘의 수렴 보장을 동시에 확보했다는 기술적 강점이 핵심입니다.
현대 딥러닝에서 흔히 사용되는 선형 분해 블록($W = BA$)은 어댑터, 저랭크 압축 레이어 등 다양한 곳에 나타나지만, 두 행렬의 곱으로 이루어져 있어 구조적 비유일성이라는 문제를 안고 있습니다. 이로 인해 학습이 불안정해지거나 사용 가능한 학습률이 제한되는 문제가 발생하며, 기존에는 이러한 블록을 최적화할 때 근본적인 기하학적 특성을 무시하는 표준 유클리드 방법을 사용했습니다.
Stiefel-AdamW는 이러한 문제를 해결하기 위해 도입된 옵티마이저입니다. 이 방법은 두 인자 중 하나를 스티플 다양체(Stiefel manifold)에 제약하고 다른 하나는 유클리드 공간에 남겨둠으로써, 전체 $\mathrm{GL}(\mathbb{R}^r)$ 게이지 대칭성을 컴팩트한 직교 대칭성으로 완화합니다. 이를 통해 요인 폭발을 방지하면서도 AdamW가 가진 좌표별 대각 사전 조건화(diagonal preconditioning)의 실용적인 강점을 유지할 수 있습니다.
Stiefel-AdamW는 구현 오버헤드가 최소화되었으며, 리만 다양체 방법론의 안정성 이점과 표준 최적화 알고리즘의 수렴 보장을 모두 계승합니다. GPT2, ViT, Mistral 7B에 대한 LoRA 스타일 및 OpenWebText를 이용한 GPT2 전체 사전 학습을 통해 검증되었으며, 추가 비용 없이 강력한 기준 성능 대비 일관된 개선 효과를 입증했습니다.
용어 풀이
- LoRA
- 모델 전체 대신 작은 추가 부분만 학습시켜 적은 비용으로 미세 조정하는 기법.
- 매개변수
- 모델이 학습하면서 조정한 내부 숫자. 개수가 많을수록 대체로 모델이 크고 무거워요.
- 파인튜닝
- 이미 학습된 모델을 특정 목적의 데이터로 조금 더 학습시키는 일.