Post

[논문 리뷰] LoRA: Low-Rank Adaptation of Large Language Models

Microsoft Corporation arXiv 2021

📖 핵심 훑어보기 !!

  • Low-Rank Parametrized Update Matrices: LoRA는 down-stream task에 fine-tuning(가중치 업데이트)을 수행하는 대신, dense layer의 가중치 업데이트를 low-rank 행렬 $A$와 $B$로 분해하여 표현하고, 원래 가중치 행렬은 고정한 상태로 $A$와 $B$만 학습함으로써 효율성을 높임.
  • Intrinsic Rank: LoRA는 모델의 가중치 변화가 실제로는 low intrinsic rank를 가지고 있다고 가정하며, 이를 통해 적은 수의 파라미터만 학습해도 모델의 성능을 유지할 수 있게 함. 이는 large-scale 모델의 매개변수 효율을 극대화하는 핵심 원리임.
  • 효과 및 결과: GPT-3 175B 모델에 LoRA를 적용하면 VRAM 사용량이 1.2TB에서 350GB로 줄고, fine-tuning 중 메모리 요구량과 checkpoint 크기를 크게 줄이면서도 학습속도는 25% 향상됨.




Introduction


자연어 처리(NLP)의 많은 애플리케이션은 대규모 pre-train된 언어 모델을 다양한 task에 맞게 fine-tuning 하여 사용한다. 하지만, 이러한 fine-tuning 방식은 원래 모델의 매개변수 수와 동일한 수의 매개변수를 학습해야하므로, 특히 GPT 같은 초대형 모델에서는 배포와 운영에 큰 부담을 준다. 이를 해결하기 위해 task-specific한 일부 매개변수만 학습하거나 외부 모듈을 추가하는 방식이 제안되었지만, 이러한 방법은 종종 모델의 depth를 증가시키고, 성능을 저하시키거나 추론 지연(inference latency)을 발생시킨다.

저자는 학습된 over-parametrized 모델이 실제로는 low intrinsic 차원(고유 차원)에 있다는 것에 영감을 받아, 모델 adaptation 중 가중치의 변화 또한 낮은 “intrinsic rank(고유 랭크)"를 가지고 있다고 가정한다. 여기서 LoRA(Low-Rank Adaptation)라는 새로운 방식이 제안된다. 이 방식은 pre-train된 모델의 가중치를 고정하고, 모델 adaptation 중에 row-rank matric을 학습함으로써 효율성을 높인다. 이를 통해 저장 요구 사항과 학습 효율성을 개선하고, GPT-3과 같은 초대형 모델에서도 성능을 유지하면서도 하드웨어 요구 사항을 크게 줄일 수 있음을 보여준다.

LoRA의 주요 장점은 다음과 같다.

  • Pre-train된 모델을 공유하여, 다양한 task를 위한 여러개의 LoRA 모듈을 만드는데 사용할 수 있다.
  • 모델 adaptation 과정에서 대부분의 파라미터에 대해 gradient를 계산하지 않고 optimizer 상태를 유지할 필요가 없으므로, 효율적인 학습이 가능하고 하드웨어 장벽을 최대 3배까지 낮춘다.
  • 간단한 linear 설계로 인해, fine-tuning된 모델과 비교하여 inference latency가 존재하지 않는다.
  • 이전에 있던 많은 방법들과 결합해서 사용할 수 있다(e.g. prefix-tuning).






Problem Statement


논문의 방법은 training object와 관계 없이 사용가능하지만, 본문에서는 LLM 학습을 예로 들어 설명했다.

LoRA_1.png

  • 예를 들어, pre-train된 LLM의 확률함수는 $P_{\Phi} (y∣x)$이다(매개변수 $\Phi$). 이 때, LLM fine-tuning 과정을 살펴보자.
  • 위의 (1)번 Objective function과 같이 Maximum Likelihood를 이용하여 전체 파라미터 $\Phi$를 업데이트한다. Objective를 최대화하기 위해 gradient를 반복적으로 업데이트한다: $\Phi_0 + \Delta \Phi$
  • 여기서 fine-tuning의 문제점은, 각 down-stream task에 대해 전체 파라미터 수만큼 업데이트를 진행해야한다는 점이다.
  • 따라서 본 논문에서는 parameter-efficient한 접근을 위해, task-specific 파라미터 증가량 $\Delta \Phi = \Delta \Phi(\Theta)$를 훨씬 더 적은 파라미터 $\Theta$로 인코딩한다. 따라서 $\Delta \Phi$를 찾는 작업은 $\Theta$에 대한 최적화가 된다(식 (2)).

Pre-train된 모델이 GPT-3 175B인 경우, 학습 가능한 매개변수 수 $|\Theta|$는 $|\Phi_0|$의 0.01%만큼 줄어드는 효과가 있다고 한다.






Method


이제 설명할 LoRA 방법은 모든 딥러닝 구조에 적용할 수 있지만, 논문에서는 Transformer 언어 모델에 초점을 맞춰 설명한다.


1. LOW-RANK-PARAMETRIZED UPDATE MATRICES

신경망에는 행렬곱을 수행하는 많은 dense layer가 있다. Dense layer의 가중치 행렬(weight matrices)은 일반적으로 full-rank를 갖는다. 이전의 연구 결과에서 영감을 얻어, 특정 task에 모델을 adaptation 할 때 가중치 업데이트가 “intrinsic rank”를 갖는다고 가정한다.

LoRA_2.png

  • Pre-train된 가중치 행렬 $W_0 \in ℝ^{d \times k}$ 의 업데이트를 생각해보자: $W_0 + \Delta W$
  • 이 경우, $\Delta W$는 low-rank decomposition을 통해, $B \in ℝ^{d \times r}$ 와 $A \in ℝ^{r \times k}$의 곱으로 표현할 수 있다. 이 때, rank $r \ll \min(d, k)$ 이다: $W_0 + \Delta W = W_0 + BA$
  • 학습이 진행되는 동안, $W_0$는 frozen되고 gradient update를 하지 않는다. 대신 $A$와 $B$의 trainable 매개변수가 학습된다.
  • forward pass 과정을 수식으로 나타내면 다음과 같다.

    \[h = W_0x + \Delta Wx = W_0x + BAx\]



  • $A$를 random Gaussian 분포로 초기화하고, $B$를 0으로 초기화하여 $\Delta W = BA$는 초기에 0이 된다.
  • 이후 $\Delta W x$를 $\alpha / r$ 만큼 scaling한다. 여기서 $\alpha$는 $r$에 대해 일종의 constant 역할을 하며, Adam을 사용해서 optimization 할 때 초기화를 적절히 했다면 $\alpha$는 일종의 learning rate 역할을 한다. 이 방법은 r을 변경할 때 hyperparameter를 재조정할 필요성을 줄여준다고 한다.




2. APPLYING LORA TO TRANSFORMER

  • LoRA는 모든 dense layer에 적용할 수 있다(Transformer의 self-attention 모듈에는 4개의 가중치 행렬(Wq, Wk, Wv, Wo)이 있고, MLP 모듈에는 2개의 가중치 행렬이 있다).
  • 본 논문에서는 Transformer 아키텍처의 self-attention 모듈의 가중치 행렬(Wq, Wk, Wv, Wo)에만 적용하고, MLP 모듈은 학습하지 않고 frozen했다.
  • 이를 통해 학습 시 파라미터 효율성을 극대화하고 메모리 사용량을 줄일 수 있다. 실험 결과, GPT-3 175B 모델에서 LoRA를 적용하면 학습 중 VRAM 사용량을 1.2TB에서 350GB로, 모델 checkpoint 크기를 350GB에서 35MB로 감소시킬 수 있다. 또한 전체 fine-tuning과 비교하여 25%의 학습속도가 향상되었다.






Experiments


1. Baseline

LoRA_3.png

LoRA의 성능을 비교하기 위해 여러 기존 방식과의 성능을 비교했다.

  • Fine-Tuning (FT): 모델 전체를 미세 조정하는 방식. 일부 실험에서는 상위 두 레이어만 미세 조정하는 방식(FTTop2)도 포함되었다.
  • Bias-Only Tuning (BitFit): 편향 벡터만 학습시키고 나머지 파라미터는 고정한다.
  • Prefix-embedding Tuning (PreEmbed): 입력 토큰에 특별한 토큰을 추가해 학습하는 방식으로, 토큰을 앞에 추가하는 “prefixing”과 토큰을 뒤에 추가하는 “infixing” 으로 나뉜다.
  • Prefix-layer Tuning (PreLayer): Prefix-embedding Tuning의 확장으로, 일부 특수 토큰에 대해 단어 임베딩만 학습하는 대신, 모든 Transformer 계층 이후의 활성화를 학습한다.
  • Adapter Tuning: 어댑터 레이어를 추가하여 일부 파라미터만 학습하는 방식. 다양한 어댑터 구조(AdapterH, AdapterL, AdapterP, AdapterD)가 실험되었다.




2. RoBERTa base/large

LoRA는 RoBERTa base(125M)와 large(355M) 모델에서 GLUE 벤치마크로 평가되었다. 기존 연구의 어댑터 방식과 동일한 설정을 사용해 공정한 비교를 진행했다. LoRA는 어댑터 기반 방식과 유사한 성능을 보이면서도 메모리와 연산 효율성을 크게 개선했으며, 특히 파라미터 수가 적은 환경에서도 효율적인 성능을 입증했다. 실험 결과는 Table 2에 요약되어 있다.



3. DeBERTa XXL

DeBERTa XXL(1.5B) 모델에서 LoRA의 성능을 GLUE 벤치마크로 평가했다. LoRA는 완전 미세 조정과 거의 동일한 성능을 보여주었으며, 대규모 모델에서도 LoRA가 효율적인 파라미터 학습을 통해 성능을 유지할 수 있음을 확인했다. Table 2 하단에서해당 결과를 확인해볼 수 있다.



4. GPT-2 medium/large

LoRA_4.png

GPT-2 medium과 large 모델에서 LoRA의 NLG(Natural Language Generation) 성능을 평가했다. E2E NLG Challenge 데이터셋에서 LoRA는 기존의 미세 조정 방식과 성능 차이가 거의 없었다. 추가로 WebNLG와 DART 데이터셋에서도 유사한 성능이 나타났으며, LoRA가 NLG 작업에서도 효과적임을 입증했다(Table 3 참고).




5. Scaling up to GPT-3 175B

LoRA_5.png

LoRA는 GPT-3 175B 모델에서 마지막으로 확장된 실험을 수행했다. 높은 training cost 때문에 표준 편차만을 보고했으며(Table 4 참고), LoRA는 세 가지 데이터셋에서 기존 미세 조정 방식과 비슷하거나 더 나은 성능을 보였다. 특이하게도, Prefix-embedding과 Prefix-layer 방식은 더 많은 학습 가능한 토큰을 사용할 때 성능이 떨어지는 경향이 있었다(Figure 2 참고). 이는 입력 분포가 사전 학습 데이터 분포와 멀어지기 때문으로 추측된다.

LoRA_6.png



This post is licensed under CC BY 4.0 by the author.