Riemannian Gradient Descent for Low-Rank Architectures

기술 요약: 저계수 아키텍처를 위한 리만 경사 하강법 (Riemannian Gradient Descent)

문제 정의
현대 딥러닝에서 모델 파라미터는 종종 계산 비용과 파라미터 수를 줄이기 위해 저계수(low rank)로 제약된 행렬 형태를 가집니다 (예: $W = AB^T$와 같은 멀티헤드 어텐션 메커니즘). 표준 학습 방식은 유클리드 기울기(Euclidean gradient)를 사용하여 인자 $A$와 $B$를 독립적으로 업데이트하며, 이는 계수 행렬의 기하학적 구조를 무시합니다. 본 논문은 계수 행렬 $W$를 매끄러운 부분 다양체(smooth submanifold) 위의 한 점으로 취급하고 리만 최적화(Riemannian optimization)를 적용하는 것이 더 나은 수렴성이나 해를 제공할 수 있는지 조사합니다. 핵심 과제는 $W=AB^T$라는 분해가 유일하지 않으며($A \to AS, B \to BS^{-T}$에 대해 불변), 계수 행렬의 집합은 벡터 공간이 아닌 다양체라는 점입니다. 저자들은 모델 아키텍처의 변경 없이 오직 옵티마이저만을 수정하면서, 매 반복당 복잡도를 표준 방식과 유사한 수준($O(mr^2 + nr^2)$)으로 유지하는 것을 목표로 합니다.

방법론
저자들은 저계수 분해된 파라미터를 위해 조정된 리만 경사 하강법 프레임워크를 제안합니다. 일반적인 반복 과정은 다음과 같습니다:

1. 인자 기울기 $G_A$와 $G_B$로부터 주변 유클리드 기울기 $G_W$를 복구합니다.
2. 이 기울기를 특정 다양체의 접공간(tangent space)에 투영하여 리만 기울기 $G_W^R$을 얻습니다.
3. 모멘텀을 업데이트하고 단계 크기(step size)를 정규화합니다.
4. 측지선(geodesic)을 따라 리트랙션(retraction)을 사용하여 다음 반복점으로 이동합니다.
5. 모멘텀 벡터를 새로운 접공간으로 평행 이동(parallel transport)합니다.

본 논문은 기하학 및 제약 조건에 따라 분류된 10가지 특정 알고리즘 설계에 대해 이 템플릿을 구현합니다:

1. 고정 계수 행렬을 위한 기하학 ($R_{r}^{m \times n}$):

   • 임베디드 기하학 (Embedded Geometry): 주변 프로베니우스 메트릭(Frobenius metric)을 사용합니다. 리트랙션은 $W + \Xi$의 계수-$r$ 절단 SVD를 통해 수행됩니다. 평행 이동은 메트릭 투영을 사용합니다.
   • 쿼션티에트 기하학 (Quotient Geometry): 다양체를 인자 $(A, B)$의 곱 공간의 쿼션티(quotient)로 취급합니다. 메트릭은 "위쪽(upstairs)"의 인자들에 대해 정의되며, 이는 쿼션티 상의 메트릭을 유도합니다. 리트랙션은 인자들에 대한 가법(additive) 방식으로 이루어지며, 트랜스포트는 수평 투영(horizontal projection)을 포함합니다.

2. 계수-$r$ 부분 등거리 변환(Partial Isometries, $P_{r}^{m \times n}$)을 위한 기하학:

   • 여기서 $W$는 0이 아닌 특잇값(singular value)이 모두 1이 되도록 제약됩니다 (이는 $A$와 $B$가 열 직교함을 의미함).
   • 임베디드 기하학: 프로베니우스 메트릭을 제한합니다. 리트랙션은 특정 직교 제약 조건을 가진 절단 SVD를 사용합니다.
   • 쿼션티에트 기하학: 스티펠 다양체(Stiefel manifolds, $St_{m,r} \times St_{n,r}$)의 곱을 "위쪽" 공간으로 사용합니다. 리트랙션은 QR 또는 폴라 분해(Polar decomposition)를 사용합니다.
   • 캐노니컬 기하학 (Canonical Geometry): 스티펠 인자들의 캐노니컬 메트릭을 사용하여 유도된 곱 메트릭을 사용합니다.

3. 가중치 공유 (그리드 다양체):

   • 위 내용을 블록 행 및 블록 열에 걸쳐 인자가 공유되는 아키텍처(예: Group-Query Attention)로 확장합니다.
   • **고정 계수(Fixed-Rank)**의 경우, 인자들을 더 큰 행렬으로 쌓음으로써 계수 보존이 블록 단위로 확인되는 한 싱글톤 알고리즘을 사용할 수 있음을 보여줍니다.
   • **부분 등거리 변환(Partial Isometries)**의 경우, 직교 제약 조건 때문에 단순한 적층(stacking)은 실패합니다. 저자들은 임베디드 케이스를 위해 쿼션티에-로 리프트하여 리트랙션과 트랜스포트를 수행한 후 다시 아래로 투영하는 "하이브리드" 접근 방식을 포함하여 더 복잡한 알고리즘을 개발했습니다.

주요 기여

• 체계적 탐색: 본 논문은 고정 계수 및 부분 등거리 변환 제약, 그리고 가중치 공유 시나리오를 모두 아우르는 저계수 행렬 파라미터를 위한 10가지 독특한 리만 최적화 변형에 대한 포괄적인 조사와 구현을 제공합니다.
• 구현 세부 사항: 계수의 비유일성과 다양체 제약(계수 및 직교성)의 보존 문제를 해결하면서, 구체적인 리트랙션(SVD, QR, Polar)과 벡터 트랜스포트의 상세 내용을 설명합니다.
• 실험적 평가: 소규모 언어 모델(6개 레이어, 512 임베딩 차원)의 Q, K, V, O 행렬에 대해 FineWeb 데이터셋을 사용하여 방법론을 적용했습니다.

실험 결과

• 성능: 학습률을 튜닝한 후에도, 리만 방법들이 AdamW 베이스라인을 결정적으로 능가하지 못했습니다. 멀티헤드 어텐션(MHA) 및 그룹 쿼리 어텐션(GQA) 설정에서 손실 값의 차이는 종종 랜덤 시드의 분산 범위 내에 있었습니다.
• 하이퍼파라미터: 메트릭 정규화는 약간의 개선(0.01–0.05 nats)을 제공했습니다. 강한 모멘텀($\nu=7/8$)은 일반적으로 성능을 저하시키거나 베이스라인 대비 미미한 이득만을 제공했습니다.
• 안정성: 일부 실행에서 손실 스파이크(예: 스텝 6660 부근)가 나타났으나 회복되었으며, 이는 종종 훈련 세트의 구조화된 데이터와 상관관계가 있었습니다.
• 복잡도: 저자들은 자신들의 방법이 수학적으로 타당하지만, AdamW보다 매 반복당 계산 복잡도가 높다는 점을 인정하며, 이로 인해 수행된 소규모 실험에서는 매력이 떨어질 수 있다고 언급했습니다.

의의 및 주장
본 논문은 자신의 독창성과 즉각적인 영향력에 대해 겸허한 입장을 유지합니다:

• 재발견: 저자들은 명백한 독창성이 있다면 그것은 "거의 확실히 재발견"일 것이라고 명시하며, 확립된 교과서(Absil et al., Boumal)와 기존 문헌에 크게 의존하고 있음을 밝힙니다.
• 휴리스틱 동기: 본 연구는 파라미터 간의 관계(다양체 기하학을 통해)를 더 잘 고려하는 것이 더 빠른 수렴이나 더 나은 해로 이어질 수 있다는 휴리스틱에 의해 동기 부여되었으나, 실험을 통해 이것이 결정적으로 증명되지는 않았습니다.
• 미래 잠재력: 소규모 모델에서의 성능 미달이 대규모 모델에서의 가치를 부정하는 것은 아니라고 주장합니다. 저자들은 접근 방식이 수학적으로 견고하며, "비용-정확도 트레이드오프"가 더 큰 모델이나 다른 아키텍처(예: MLA, MoLAE와 같은 텐서 기반 모델)에서 유리하게 바뀔 수 있다고 제안합니다.
• 한계점: 연구는 작은 모델 크기, 학습률 감쇠(learning rate annealing)의 부재, 그리고 하이퍼파라미터 튜닝 없이 단일 베이스라인(AdamW)만을 사용한 점 등 상당한 한계를 인정합니다.

결론적으로, 이 논문은 저계수 딥러닝 구성 요소에 리만 최적화를 적용하는 엄격한 기술적 시연으로서, 알고리즘 라이브러리와 베이스라인을 제공하는 동시에, 현재의 실험적 결과가 아직 AdamW와 같은 표준 옵티마이저를 대체할 만큼 충분하지 않음을 인정하며 마무리됩니다.