기술 요약: LoRA 랭크에 따른 학습률 스케일링 및 풀 파인튜닝으로의 전이

문제 정의

저차원 적응(Low-Rank Adaptation, LoRA)은 사전 학습된 가중치에 저차원 인수분해 $\alpha BA$를 추가하는 매개변수 효율적 파인튜닝(PEFT)의 표준적인 방법입니다. 그러나 LoRA의 훈련 역학은 하이퍼파라미터 간의 상호작용, 특히 초기화 방식(어떤 인자를 무작위화하고 어떤 인자를 0으로 설정할 것인지), 랭크 의존적 스케일링 승수 $\alpha$, 그리고 학습률 $\eta$에 매우 민감합니다.

실무자들은 중요한 문제에 직면해 있습니다: 최적의 학습률은 LoRA 랭크 $r$에 대해 불변하지 않습니다. 랭크를 변경하거나 표준 설계 선택(예: $Init[A]$대$Init[B]$, 또는$\alpha=1$대$\alpha=r^{-1}$) 사이를 전환하면 이전에 최적이었던 학습률이 최적이 아니게 되어, 비용이 많이 드는 재튜닝이 필요하게 됩니다. 또한, LoRA 설정을 풀 파인튜닝(Full Finetuning, FFT)으로 어떻게 전이할 수 있는지에 대해서도 여전히 불분명하며, FFT는 계산 비용이 매우 높습니다. 본 연구가 다루는 핵심 질문은 다음과 같습니다: 최적의 학습률 $\eta$는 서로 다른 LoRA 구성 하에서 어댑터 랭크 $r$과 모델 너비 $n$에 따라 어떻게 의존하는가?

방법론

저자들은 사전 학습에 사용되는 Maximal-Update Parametrization($\mu P$)에서 영감을 얻은 이론적 프레임워크인 **Maximal-Update Adaptation ($\mu A$)**을 도입합니다. 분석은 모델 너비 $n$과 LoRA 랭크 $r$이 모두 무한대로 향하는 공동 극한($n, r \to \infty$)에서 수행됩니다.

이론적 프레임워크

1. 안정성 및 비소멸 업데이트: 이 프레임워크는 "안정적인 특징 학습(stable feature learning)"을 폭발을 방지하기 위해 업데이트가 유계($O(1)$)되면서도, 학습이 멈추는 것을 방지하기 위해 비소멸적($\Theta(1)$)인 상태로 정의합니다.
2. 업데이트의 분해: 특징 업데이트 $\Delta Z_B$를 한 인자는 고정하고 다른 하나는 선형적으로 업데이트하는 항, 그리고 두 인자 모두 곱하는 항을 나타내는 세 가지 항($\delta_1, \delta_2, \delta_3$)으로 분해합니다. 이론은 서로 다른 초기화 방식 하에서 이러한 항들의 점근적 스케일링을 분석합니다.
3. 가정: 분석은 유계된 순방향/역방향 신호를 가정하며, 폐쇄형 스케일링 규칙을 도출하기 위해 Adam의 다루기 쉬운 대리 모델인 SignSGD 옵티마이저 추상화를 활용합니다.
4. 초기화 방식:
   • Init[A]: $A$는 무작위로 초기화되고($A \sim \mathcal{N}(0, 1/n)$), $B=0$입니다.
   • Init[B]: $B$는 무작위로 초기화되며($B \sim \mathcal{N}(0, 1/r)$), $A=0$입니다.
5. 스케일링 승수: 상수 값($\alpha=1$)과 랭크 의존적 값($\alpha=r^{-1}$)을 포함하여 흔히 사용되는 승수 $\alpha$의 선택지를 조사합니다.

스케일링 규칙 도출

$\Delta Z_B = \Theta(1)$ 조건을 강제함으로써, 저자들은 $n$과 $r$의 함수로서 $\eta$에 대한 명시적인 스케일링 처방을 도출합니다:

• Init[A]의 경우: 최적의 학습률은 $n$과 $r$ 모두에 의존합니다. 구체적으로, $\alpha=1$이면 $\eta \propto n^{-1/2}r^{-1/2}$입니다. 만약 $\alpha=r^{-1}$이면, $\eta \propto n^{-1/2}$가 되어 랭크에 불변합니다.
• Init[B]의 경우: $\alpha=1$일 때, 최적의 학습률은 $\eta \propto n^{-1}$로 스케일링되며 랭크 $r$에 독립적입니다. 결정적으로, 이 스케일링은 동일한 가정 하에서 선형 레이어의 풀 파인튜닝에 대한 이론적 요구 사항과 일치합니다.

주요 기여

1. $\mu A$ 프레임워크: 최대(안정적) 적응 조건을 정의하며, $n, r \to \infty$ 공동 극한에서 LoRA 파인튜닝 역학을 특징짓는 이론적 프레임워크를 도입했습니다.
2. 두 가지 영역 식별:
   • 랭크 의존적 영역(Rank-Dependent Regime): 상수 $\alpha$를 사용하는 $Init[A]$하에서, 최적의 학습률은 랭크에 반비례하여 스케일링됩니다 (예:$\eta \propto r^{-1/2}$).
   • 랭크 불변 영역(Rank-Invariant Regime): $\alpha=1$인 $Init[B]$ 하에서, 최적의 학습률은 랭크 변화에 불변합니다.
3. LoRA-to-FFT 전이: 최적의 학습률 스케일링($\eta \propto n^{-1}$)이 풀 파인튜닝과 일치하는 특정 구성($Init[B]$와$\alpha=1$)을 식별했습니다. 이는 재튜닝 없이 LoRA 설정을 풀 파인튜닝으로 전이하기 위한 필요 조건을 제공합니다.
4. 경험적 검증: 다양한 태스크(언어, 비전, 비전-언어, 이미지 생성, 강화 학습)와 아키텍처(Llama, Qwen, RoBERTa, ViT, Stable Diffusion)에 걸친 광범위한 실험을 통해 도출된 스케일링 규칙을 검증했습니다.

결과

실험 결과는 이론적 예측과 일관되게 일치합니다:

• $\alpha=1$인 Init[A]: 랭크 $r$이 증가함에 따라 최적의 학습률은 감소합니다. 랭크가 4배 증가할 때 최적의 학습률은 약 2배 감소합니다 ($\eta \propto r^{-1/2}$와 일치).
• $\alpha=r^{-1}$인 Init[A]: 최적의 학습률은 랭크에 불변하지만, 안정적인 범위가 좁습니다.
• $\alpha=1$인 Init[B]: 최적의 학습률은 랭크에 불변하며, 대부분의 모델에서 풀 파인튜닝의 최적 학습률과 밀접하게 일치합니다.
• 전이 가능성: $Init[B], \alpha=1$ 구성을 사용하여 LoRA에서 튜닝된 학습률은 풀 파인튜닝으로 안정적으로 전이되며, 이는 FFT를 위한 하이퍼파라미터 탐색 비용을 획기적으로 줄여줍니다.
• RLVR 민감도: 검증 가능한 보상이 있는 강화 학습(RLVR)에서, 최적이 아닌 학습률은 성능을 저하시킬 뿐만 아니라 모델이 탐색 단계에 갇히게 만들어 훈련 시간을 2~4배 증가시킵니다. $\mu A$ 규칙은 이를 피하기 위한 안정적인 영역을 식별하는 데 도움을 줍니다.

의의 및 주장

본 논문은 휴리스틱한 규칙을 넘어, LoRA에서의 학습률 선택에 대한 원칙적인 접근 방식을 제공한다고 주장합니다. 너비와 랭크의 공동 극한을 규명함으로써, 저자들은 왜 특정 구성이 랭크 의존적 학습률을 요구하는 반면 다른 구성은 그렇지 않은지를 설명합니다.

주된 의의는 하이퍼파라미터의 전이 가능성에 있습니다. 저자들은 올바른 초기화($Init[B]$)와 스케일링($\alpha=1$)을 선택함으로써 실무자들이 다음을 수행할 수 있음을 보여줍니다:

1. LoRA 랭크를 변경할 때 학습률을 다시 튜닝하는 것을 방지합니다.
2. 저차원 LoRA 실험에서 최적의 학습률을 풀 파인튜닝 설정으로 직접 전이하여, 상당한 계산 자원을 절약할 수 있습니다.

저자들은 이론이 점근적이며 유한한 $(n, r)$ 환경의 노이즈에 민감할 수 있고, 도출된 일치 조건이 모든 전이 시나리오에 대해 충분(sufficient)하다는 것이 증명된 것은 아니라는 점을 들어 제한 사항을 조심스럽게 언급했습니다. 그러나 경험적 증거는 다양한 도메인에서 이론과 실제 사이의 강력한 일치를 보여줍니다.