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

arxiv 2021 paper code
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen
Microsoft Corporation
17 Jun 2021

들어가기 전에..

LoRA는 fine-tuning을 비교적 저차원에서 가능케하는 Parameter efficiency fine-tuning (PEFT)방법론으로써 LLM을 위해 처음 공개되었지만, 나는 Stable Diffusion 모델에 fine-tuning을 하며 처음 알게 되었다. 당시, 논문을 깊게보지 않고 실무에 필요한 정보만 얻었기에 이번 기회에 논문을 정독하며 리뷰 포스트를 남기고자 한다.

Abstract

저자는 LLM 학습 시 모든 파라미터를 업데이트 하기에는 비용이 너무 크다는 점을 지적하며 pre-trained weight를 freezing하고 학습가능한 저차원 매트릭스를 적용하는 LoRA(Low-Rank Adaptation)를 제안하였다. GPT-3 175B + Adam 기준, trainable parameter와 VRAM 각각 $ 1/10000 $, $ 1/3 $ 수준으로 줄일 수 있었다고 한다. 이러한 LoRA는 추가적인 추론 시간 지연 없이 RoBERTa, DeBERTa, GPT-2, GPT-3에서 Full Fine-Tuning(FT)과 동등하거나 더 좋은 성능을 달성했다.

1. Introduction

Large model에 대해 전체 파라미터를 tuning하는 것은 너무나 비효율적이므로, 작은 파라미터를 tuning하는 연구들이 진행되고 있다. 그 중, adapter 방식은 inference latency가 증가되는 문제가 있다. 저자는 작은 파라미터에서 tuning하며 inference latency를 증가시키지 않는 Low-Rank Adaptation (LoRA)를 제안한다. 위 그림과 같이 pre-trained weight는 freezing하고 dense layer를 저차원 rank(dimension)를 통해 연산하는 두 매트릭스로 분해하여 학습하는 방식이다. 이때, rank의 크기는 1 내지 2와 같은 작은 값에서도 잘 작동한다.

저자는 LoRA의 주요 장점을 다음과 같이 정리하였다.

• pre-trained는 고정하고 LoRA 모델만 학습 후 갈아끼우는 방식으로 사용할 수 있어 저장공간과 작업 전환 오버헤드를 크게 줄일 수 있다.
• 학습 효율을 높혀 3배 가량 낮은 컴퓨팅 환경에서 학습이 가능하다.
• 간단한 linear 디자인으로 설계되어 다양한 frozen weight에 merge가 가능하다. 때문에, 추가적인 inference latency가 없다.
• 이전의 다양한 방법이 똑같이 적용 가능하며 prefix-tuning에도 사용 가능하다.(Appendix E 참고)

본 논문에서는 기존 Transformer의 용어 및 규칙을 그대로 따른다고 한다.

• $ d_{model} $ : input and output dimension size of a Transformer layer.
• $ d_{q}, d_{k}, d_{v}, d_{o} $ : query, key, value, output projection matrics in the self-attention modules.
• $ W or W_{0} $ : pre-trained weight matrix
• $ \Delta W $ : accumulated gradient update during adaptation
• $ r $ : rank of a LoRA module
• Transformer MLP feedforward $ d_{ffn} = 4 \times d_{model} $

2. Problem Statement

$$ \max_{\Phi} \sum_{(x,y) \in \mathcal{Z}} \sum_{t=1}^{|y|} \log(P_{\Phi}(y_t | x, y_{<t})) $$

저자는 다양한 Downstream task에 대해 각 태스크마다 모든 Parameter를 학습해야하고 이것을 저장하며 관리하는 것은 Full Fine-tuning(FT)의 대표적인 단점이라고 지적한다. 위 수식은 기존 FT에 대한 수식이며 저자는 LoRA를 이용한 다음 수식을 제안한다.

$$ \max_{\Theta} \sum_{(x,y) \in \mathcal{Z}} \sum_{t=1}^{|y|} \log p_{\Phi_0 + \Delta\Phi(\Theta)}(y_t | x, y_{<t}) $$

한마디로, 연산과 메모리 모두 효율적인 LoRA Adapter인 $\Theta$만 학습하여 전체 파라미터를 학습하는것과 같은 효과를 보이고자 한다. GPT-3 175B 기준, trainable parameters는 FT 대비 0.01%밖에 되지 않는다.

3. Aren't Existing Solutions Good Enough?

이전에도 연산과 메모리를 효율화하는 학습 방법론에 대한 연구는 계속되어 왔다. 하지만 Large scale이나 inference latency에 민감하다는 문제가 있었다.

Adapter Layers Introduce Inference Latency

트랜스포머 블록 당 Adapter Layer를 2개 추가하거나 1개와 LayerNorm을 추가한 선행 연구가 존재한다. 하지만 여기서의 Adapter Layer는 순차적으로 연산이 되기에 latency가 발생한다. 위 두 방식과 LoRA에 대해 Inference Latency를 비교한 결과, LoRA 대비 2~30% 차이가 난다고 한다. 특히, 짧은 Sequence length에서 그 차이가 더 크게 난다고 한다.

Directly Optimizing the Prompt is Hard

Prefix tuning이란 방법은 입력 token 앞에 약간의 token을 입력하고 해당 token에 대해서만 학습을 진행하는 것이다. 하지만 이 방식은 최적화가 잘 안된다는 기본적인 단점이 있다.

4. Our Method

4-1. Low-Rank-Parametrized Update Matrices

일반적인 dense layer를 full-rank라고 하고, LoRA는 작은 dimension(rank)에서 연산이 수행되므로 low-rank라고 한다.

$$ h = W_0 x + \Delta W x = W_0 x + BAx $$

• $ W_0 \in \mathbb{R}^{d \times k} $ : 크기가 $d \times k $인 pre-trained weight matrix
• $BA$ : $\Delta W$를 Decompose하여 각각의 matrix로 나타냄
• $B \in \mathbb{R}^{d \times r}$
• $A \in \mathbb{R}^{r \times k}$
• $r \leq \min(d, k)$

위 그림에서 A는 랜덤 Gaussian으로 초기화하고, B는 0값으로 초기화한다. 따라서 첫 step에서의 $\Delta W = BA$ 결과값은 0이 된다. 이후, $\Delta W$를 $\frac{\alpha}{r}$로 scailing 후 $W_0$ 결과값과 덧셈 연산을 하게된다. Adam Optimizer 사용 시, $\alpha$는 $learning rate$와 같은 역할을 수행한다.

A Generalization of Full Fine-tuning

LoRA의 Rank r을 full-rank만큼 크게 설정한다면 대략적으로 FT의 표현력을 복원할 수 있다. 하지만 기존의 Adapter 방식은 MLP로 수렴하게 되고, prefix tuning의 경우, 긴 입력 시퀀스를 처리할 수 없는 모델로 수렴한다.

No Additional Inference Latency

제품 배포 시, $W = W_0 + BA$를 계산하고 저장한 후 추론을 진행한다. 이때, $W_0$와 $BK$ 모두 $R^{d \times k}$에 속한다. 다른 downstream task에서는 BA를 빼고 새로운 BA를 더하여 $W_0$를 복구할 수 있다.

4-2. Applying LoRA To Transformer

저자는 4개의 self-attention module($W_q, W_k, W_v, W_o$)과 2개의 MLP Layer가 있는 다양한 Transformer architecture에 LoRA를 적용해보았다. 본 논문에서는 간결성과 효율성을 위해 오직 Attention 부분에서만 LoRA를 적용하였다. 나머지 MLP, LayerNorm, biases 부분은 추후 진행할 연구로 남겨둔다.

Practical Benefits and Limitations

LoRA의 대표적인 이점은 메모리와 저장공간을 아낄 수 있다는 것이다. Adam을 이용한 Transformer 학습 시 VRAM을 기존 대비 2/3 수준으로 아낄 수 있다. GPT-3 175B 기준에서는 VRAM을 1.2TB에서 350GB로 줄였으며, r=4일때 query, value 부분만 LoRA를 적용한다고 할때에 checkpoint 사이즈를 10,000배 가량 줄일 수 있게된다(1개의 checkpoint당 350GB -> 35MB). 덕분에 I/O bottleneck을 피하면서 훨씬 적은 GPU 자원에서 학습이 가능하게 된다. 또다른 이점으로는 task별로 LoRA를 쉽게 바꿔서 적용할 수 있다는 것이다. 결론적으로 저자는 GPT-3 175B 모델을 학습할 때 25%의 속도 향상을 할 수 있었다고 한다.

물론 LoRA도 한계점이 존재한다. inference-latency를 해소하기 위해 $W_0$에 $BA$를 흡수한 경우, 두 개의 다른 태스크를 single forward pass에 입력하여 처리하는 것은 간단하지 않다. (당연하지 않나..? 이정도만 되도 대단한건데,,) 물론, merge하지 않고 배치의 각 요소 별 LoRA를 선택해서 pass하게 하면 가능하고 한다.

5. Empirical Experiments

저자는 다양한 모델의 다양한 Language Task에 대한 Downstream task 결과를 공유하였다.

5-1. Baselines

기본적으로 각 모델에서 공개한 설정값들을 그대로 사용하였다.

Fine-Tuning(FT) 는 가장 흔한 adaptation 접근법이다. 보통 전체 weight를 학습하거나 일부 layer에 대해 학습한다. 본 실험에서는 전체 weight를 Fine-tuning하는 $FT$와 GPT-2 모델에 대해 마지막 두 layer만 학습하는 $FT^{Top2}$를 다루었다.

Bias-only or BitFit 은 Bias vector만 학습하고 나머지는 모두 freezing 하는 것이다.

Prefix-embedding tuning (PreEmbed) 는 special token들을 입력하여 해당 token의 word embedding에 대해서만 학습하는 것이다. prefixing은 prompt의 앞부분에 추가하는 것이고 infixing은 prompt 중간에 삽입하는 것이다. 여기에서 $l_p$($l_i$)는 prefix(infix) token 수를 나타내며 trainable parameter 수는 $d_{model} \times (l_p + l_i)$ 이다.

Prefix-layer tuning (PreLAyer) 는 위 PreEmbed를 확장한 것이다. word embedding만 학습하는것이 아닌, 매 Transformer Layer의 activations를 학습한다. 이전 layer의 activations는 학습 가능한 layer로 대체된다. trainable parameters 수는 $L \times d_{model} \times (l_p + l_i)$이다. 여기에서 $L$은 Transformer layer 수를 가리킨다.

Adapter tuning 는 self-attention module 사이에 Adapter layer와 Residual connection을 삽입하는 것이다(MLP도 마찬가지). 두 개의 fully connected layer와 bias로 구성되어 있으며 nonlinearity하다. 이 구조를 $Adapter^H$ 로 부른겠다. 최근 제안된 효율적인 Adapter는 오직 MLP module과 LayerNorm 다음에 adapter layer를 삽입하였다. 이것을 $Adapter^L$ 로 부르겠다. 이와 매우 유사한 다른 Adapter는 $Adapter^P$ 라 부르겠다. 효율을 최대화하기 위해 몇 layer를 drop 한 AdapterDrop는 $Adapter^D$ 라 부르겠다. 학습 파라미터 수는 다음과 같다.
$$L^\text{Adpt} \times (2 \times \text{dmodel} \times r + r + \text{dmodel}) + 2 \times L^\text{LN} \times \text{dmodel}$$

• $L^\text{Adpt}$ : adapt layer 수
• $L^\text{LN}$ : LayerNorm 수

LoRA 는 실험 간소화를 위해 $W_q$와 $W_v$만 학습하였으며 Trainable Parameter 수는 다음과 같다. $2 \times L^\text{LoRA} \times \text{dmodel} \times r$

5-2. RoBERTa Base(125M)/Large(355M)

RoBERTa는 BERT의 학습 트릭들을 최적화하여 별도의 파라미터 증가 없이 성능을 향상시킨 모델이다. LoRA와의 공정한 비교를 위해 기존의 설정에서 다음 두 가지를 변경하였다.

1. Sequence length를 128로 설정하고 같은 배치 사이즈 사용
2. MNLI와 같이 fine-tuning이 되지 않은 pre-trained weight를 사용한다. 표에서 †로 표기

5-3. DeBERTa XXL(1.5B)

DeBERTa는 BERT보다 훨씬 큰 스케일로 학습된 모델로, GLUE와 superGLUE 벤치마크에서 매우 큰 성능 향상을 이룬 모델이다.

5-4. GPT-2 MEDIUM/LARGE

저자는 NLU 태스크에서 LoRA가 FT 대비 우세햇으므로 중형급 모델인 GPT-2에서의 성능도 기대하며 본 모델에 대한 실험을 진행하였다고 한다.

5-5. Scaling Up To GPT-3 175B

마지막으로 대형 언어 모델(LLM)인 GPT-3에 대한 실험을 진행하였다. 실험 결과, 학습 파라미터 수가 많다고 무조건적으로 성능이 올라가지는 않은 것을 확인할 수 있었다. 또한, 256개의 special token을 갖는 prefix-embedding tuning이나 32개의 special token을 갖는 prefix-layer tuning에서 성능이 크게 감소한 것을 볼 수 있었다. (Li & Liang (2021) 논문에 같은 내용이 있는 것으로 추정됨) 저자는 많은 special token을 사용하면 사전학습 데이터 분포로 부터 distribution shift가 많이 일어나서 발생한 것으로 추정한다고 한다. (본 논문의 범위 밖이라 각설함)

6. Related Works 생략

7. Understanding The Low-Rank Updates

저자는 LoRA에 대한 몇 가지 궁금증에 대해 자문자답하는 실험을 진행하였다.

1. 예산 문제로 Parameter 수에 제약이 발생할 때, Transformer의 어떤 부분에 LoRA를 붙이는 것이 좋은가?
2. 최적의 Rank r은 몇인가?
3. $W$와 $\Delta W$의 상관관계가 어떻게 되는가?

7-1. Which Weight Matrices In Transformer Should We Apply LoRA To?

저자는 어느 attention의 어느 weight에 LoRA를 적용하는 것이 좋을지를 Rank r과 함께 비교하였다. Trainable Parameters가 18M으로 동등한 조건에서의 실험 결과는 기본적으로 $W_q$와 $W_v$를 같이 사용할 때 성능이 좋았으며, 위 표와 같이 q,k,v,o 모두 쓰는것이 두 태스크 모두에서 좋은 것으로 나타났다. 이는 더 큰 rank를 갖는것 보다 다양한 weight matrix에 LoRA를 적용하는 것이 더 의미있음을 나타낸다.

7-2. What Is The Optimal Rank r For LoRA?

위 표에서 볼 수 있듯이 r이 낮아도 충분히 좋은 성능을 나타냄을 알 수 있다. 다른 말로, 굳이 더 높은 r을 사용하지 않아도 됨을 이야기한다. 그러나 저자는 pre-training과 fine-tuning 시 언어가 다를 경우, Full fine-tuning이 LoRA를 사용했을 때 보다 성능이 더 좋을 것이라고 추정한다고 한다.

8. Conclusion And Feature Work

저자는 본 논문에서 LLM의 Fine-tuning 시 겪게되는 storage/switching 문제를 inference-latency와 성능 저하없이 해결하는 LoRA를 제시하였다. 또한, 다양한 task에 대해 빠르게 switching하며 service할 수 있는 구조가 특징이며 Transformer 언어모델 외에 Dense Layer에도 일반적으로 사용할 수 있는 기술이다.