[NLP] LoRA: Low-Rank Adaptation of Large Language Models

LoRA: Low-Rank Adaptation of Large Language Models

 

link : https://arxiv.org/abs/2106.09685

 

 

 

✲ Abstract & Introduction

 

LLM 학습에 있어 full fine-tuning (모델의 모든 파라미터를 학습하는 것)은 비효율적이다. 이를 해결하기 위해 Low-Rank Adaptation(LoRA)라는 효율적인 방법을 제안. LoRA는 pre-trained 모델의 가중치를 freeze하고, 각 transformer layer에 학습 가능한 rank decomposition matrices를 삽입해 다운스트림 task에서 필요한 학습할 파라미터 수를 획기적으로 줄인다. GPT-3 175B 기준으로 LoRA는 full-finetuning 보다 학습 파라미터수를 10,000배 줄이고 GPU 메모리를 3배 줄일 수 있다고 한다. 

 

 

 

이전에도 다운스트림 task를 위해 매개변수의 일부만을 학습하거나, 추가적인 모듈을 사용하려는 시도는 있었으나 효율성과 모델 품질 사이에 trade-off가 발생하였다. LoRA는 over-parametrized model이 실제로는 low intrinsic dimension에 있다는 것을 발견하고 이를 활용한 방법을 제안한다. model adaptation에서 가중치의 변화는 "intrinsic rank"를 가지며, 이는 제안된 LoRA 접근 방식으로 이어진다.

 

LoRA를 사용하면 Figure 1과 같이 사전 학습된 가중치는 고정한 상태로 유지하면서 adaptation 중 dense layer의 변화에 대한 rank decomposition matrices를 최적화하여 신경망의 일부 dense layers를 간접적으로 학습할 수 있다. GPT-3 175B 기준으로 full rank (figure 1의 d)가 12,288로 매우 크다고 하더라도 LoRA는 1또는 2와 같은 very low rank로 충분하다.

 

 

‣ LoRA의 주요 이점

 

1. (모델 재사용성) 큰 pre-trained model 하나를 공유하고 여러가지 task에 대해 여러 작은 LoRA 모듈을 사용할 수 있다. matrix A or B만 교체하여 작업을 전환할 수 있어 메모리와 작업 전환 오버헤드를 크게 줄임

 

2. (학습 효율성) 하드웨어 메모리의 3배 큰 모델을 학습할 수 있다. 대부분의 파라미터에 대한 gradient 계산이나 optimizer 상태 유지를 필요로 하지 않고 low rank matrices만 최적화 하면 된다.

 

3. (inference latency 없음) 간단한 linear design 덕분에 학습 가능한 matrices를 동결된 가중치와 함께 합쳐 배포할 수 있어, full fine-tuning 모델과 비교해 추론 지연이 없다

 

4. (호환성) LoRA는 기존 대부분의 방법들(ex. prefix-tuning ..)과 호환가능하고 결합가능하다.

 

 

 

✲ Problem Statement

 

GPT와 같은 generic multi-task learner인 pretained autoregressive language model $P_{\Phi}(y|x)$ 가 있다고 가정해보자. 그리고 이 모델을 다운스트림 conditional text generation task에 adapting한다고 할 때 이러한 다운스트림 task는 학습 데이터의 context-target pairs $Z = \{(x_i, y_i)\}_{i=1, .. N}$로 나타낼 수 있다.

 

full fine-tuning 중에 pre-trained weights $\Phi_0$ 로 모델은 초기화되고, conditional LM objective를 maximize 하기 위한 gradient를 따르는 $\Phi_{0}+\Delta\Phi$ 로 업데이트 된다:

 

 

 

full fine-tuning 주 단점 중 하나는 각각의 다운스트림 task를 위해 우리는 서로 다른 set of parametres $\Delta\Phi$를 학습해야 한다는 것이고 $|\Delta\Phi|$는 $|\Phi_0|$와 같다 즉, pre-trained 파라미터 dimension과 동일한 사이즈를 학습해야 한다는 것이다. 또한 pre-trained model이 매우 큰 경우 많은 독립적인 fine-tuned model들을 저장하고 deploying하는 것은 매우 어렵다.

 

 

해당 논문에서는 위에서 정의한 adapting을 parameter-efficient하게 만드는 방법을 제안한다. task-specific parameter increment $\Delta\Phi = \Delta\Phi (\Theta)$ 는 훨씬 작은 사이즈의 parameter set $\Theta$로 인코드되며 이는 원래 parameter dimension인 $| \Phi_{0} |$보다 더 작은 사이즈이다. 그리고 $\Delta\Phi$를 찾는 작업은 $\Theta$에 대해 최적화된다.

 

 

이후 섹션에서는 $\Delta\Phi$를 인코드 하기 위한 low-rank representation을 소개하며 이는 comute-, memory- efficient한 방법이다. 

 

 

✲ Aren't existing solutions good enough?

 

Adapting(fine-tuning) 단계에서 학습 파라미터 수를 효율적으로 줄여보려는 시도는 이전에도 존재하였다.

 

‣ Adapter Layers Introduce Inference Latency

 

Adapter 레이어를 사용하는 경우 중 두 가지 방법에 중점을 둔다.

transformer block당 2개의 adaptor layer를 포함한 경우 / block 당 하나의 adapter layer와 추가적인 layer norm이 있는 경우 (최신방법)

 

 

adapter layer의 latency를 줄이기 위해 layer를 pruning하거나 multi-task setting을 활용할 수 있지만 adapter layer의 계산 비용을 우회할 직접적인 방법은 없다. adapter layer는 parameter가 적고 bottleneck 차원이 작기 때문에 Flops를 제한할 수는 있지만 큰 신경망은 하드웨어 병렬성을 이용해 latency를 낮추는데, adapter layer는 순차적으로 처리되므로 latency가 증가하게 된다.

 

특히 배치가 1인 경우(챗봇, 음성인식 등에서는 배치를 1로 처리해야 함) online inference에서는 이러한 latency가 더 두드러진다. gpt2를 단일 gpu에서 실행한다면 adapter를 사용할 경우 latency가 눈에 띄게 증가한다. 모델을 sharding하는 경우, adapter layer의 추가 깊이로 인해 더 많은 synchronous GPU operations(GPU가 동시에 작업을 처리하는 방식, 나눠진 각 모델의 결과를 기다려야함)이 필요하여 문제가 심화된다.

 

 

‣ Directly Optimizing the Prompt is Hard

 

Prefix tuning은 최적화하기 어렵고, 학습 가능한 매개변수에 따라 성능이 비선형적으로 변한다는 관찰이 있었다.(참고) 더 근본적으로, 시퀀스 길이의 일부를 adaption을 위해 설정하면 다운스트림 task를 처리할 수 있는 시퀀스 길이가 줄어들어 prompt tuning 성능이 다른 방법에 비해 낮아질 가능성이 있다.

 

 

 

✲ Our Method

‣ Low-rank parametrized update matrices

 

뉴럴넷은 행렬곱을 수행하는 많은 dense layers를 포함하고 있다. 이러한 레이어들의 weight matrices는 일반적으로 full-rank이다. 이전 논문에서 특정 task에 adapting을 할 때 pre-trained LM은 더 낮은 "intrisic dimension"을 가지고 있고, 더 작은 subspace로 random projection 함에도 여전히 효율적으로 학습할 수 있음을 보였다. 이에 영감을 받아 저자들은 adapting에서 가중치 업데이트 또한 low "intrinsic rank"를 가진다고 가정하였다. 

 

사전학습 가중치 행렬 $W_0 \in \mathbb{R}^{d \times k}$의 업데이트 $\Delta W$에 대해

low-rank decomposition 로 표현한다. $W_0 + \Delta W = W_0 + BA$, where $B \in \mathbb{R}^{d \times r}, A \in \mathbb{R}^{r \times k}$, the rank $r << min(d,k)$.

 

학습 동안 $W_0$는 frozen되어 gradient updates를 하지 않고 A와 B는 학습 가능한 파라미터들이다. 두 벡터 $W_0, \Delta W = BA$는 동알한 input에 곱해지고, 이들의 각각 출력 벡터는 coordinate-wise로 합해진다. $h=W_0 x$의 경우 수정된 forward pass는 다음과 같다.

 

 

 

Figure 1에 보이는 것과 같이 제안된 reparametrization 방법은 먼저 A matrix은 random Gaussian initialization을 하고 B는 zero matrix로 초기화하며 따라서, $\Delta W = BA$ 는 학습 초기에 0 행렬이다. 그 다음 $\Delta W x$를 $\frac{\alpha}{r}$ 로 스케일하며 $\alpha$는 r의 상수이다. Adam으로 최적화할 때, $\alpha$ 초기화를 적절하게 조정하는 경우 learning rate를 조정하는 것과 거의 동일하다(updated weight는 W + alpha * A * B이므로 alpha크기에 따라 업데이트 변화정도가 바뀐다). 결과적으로 우리는 단순히 $\alpha$를 처음 r에 대해 설정하고 조정하지 않는다. 이 스케일링은 r을 변경할 때 하이퍼파라미터를 다시 조정할 필요성을 줄이는데 도움이 된다.

 

 A Generalization of Full Fine-tuning  파인튜닝의 more general form은 사전 학습된 매개변수의 subset을 학습할 수 있도록 한다. LoRA는 한 걸음 나아가 adapting 중 누적된 gradient update가 가중치 행렬 rank일 필요는 없다. 이는 LoRA를 모든 가중치 행렬에 적용할 때 rank r을 사전 학습된 가중치 행렬의 rank로 설정하면 전체 fine-tuning 표현력을 대략 회복할 수 있다. 즉 adapter 기반 방법은 MLP로 수렴하고 Prefix tuning 기반 방법은 긴 시퀀스를 처리할 수 없는 모델로 수렴하는 반면, LoRA를 학습하면 원래 모델을 학습하는 것과 거의 비슷하게 수렴한다.

 

No Additional Inference Latency production deploy를 할 때 우리는 직접적으로 $W = W_0 + BA$를 계산하고 저장하며 일반적인 추론을 수행한다. $W_0, \ BA$의 dimension은 d x k 이므로 우리가 다른 작업으로의 전환이 필요할 때 $W_0$에 $BA$를 빼거나 차이 $B'A'$를 더함으로써 recover할 수 있어 적은 메모리 오버헤드로 빠르게 연산이 가능하다. 따라서 추가적인 latency 없이 일반적인 fine-tuning 방법의 추론 시간과 동일하게 추론이 가능하다.

 

 

‣ Applying LoRA to transformer

 

이론적으로 LoRA를 뉴럴넷의 weight matrices의 어떤 부분이든 적용이 가능하다. 하지만 해당 연구에서는 Transformer 아키텍처에서 셀프어텐션 모듈의 4개의 weight matrices $(W_q, W_k, W_v, W_o)$와 2개의 MLP module이 있고 그 중 attention weights에 적용하는 것으로 제한한다. MLP 모듈은 freeze하며 이는 simplicity와 parameter-efficiency를 위함이다. 그리고 출력 차원이 어텐션 헤드 크기만큼 나뉘지만 $W_q (\text{or } W_k, W_v)$를 $d_{model} \cdot d_{model}$ 차원을 갖는 a single matrix로 다룬다.

 

No Additional Inference Latency LoRA의 가장 큰 이점은 메모리와 저장용량을 줄인다는 것이다. Adam을 사용해 Transformer를 학습할 경우 VRAM 사용량을 2/3까지 줄일 수 있다(if $r << d_{model}$). frozen 파라미터를 위한 optimizer states를 저장하지 않아도 되기 때문이다(전체 모델을 메모리에 올려야 하긴 하고 모델 파라미터 업데이트에 필요한 메모리만 주는 것임). r=4이고 query, value projection atrices만 어댑팅한다고 하면 checkpoint 사이즈가 대략 만배 줄어든다. 또한 GPT-3 175B 기준, 방대한 파라미터의 gradient를 계산할 필요가 없으므로 학습스피드도 25%감소하였다.

더 적은 GPU를 사용하게 하며 I/O 병목을 피할 수 있고 훨씬 적은 비용으로 여러 task를 switch하여 다양한 맞춤형 모델을 학습할 수 있다. 하지만 배치에 여러 task가 있고 서로 다른 A, B를 사용해야 하는 경우에는 문제가 있다. 배치처리를 하려면 배치 내 각 샘플에 대해 같은 A, B를 사용해야 한다. 이때 A,B를 W에 합치면 배치 처리가 쉬워질 수 있지만 배치 내에서 동적으로 A,B를 바꿔줘야 하므로 추론 속도가 느려질 수 있다. A, B를 W에 합치지 않고 동적으로 선택한다면 추가적 계산 비용이 발생한다

 

 

✲ Empirical experments

 

RoBERTa - base, large / DeBERTa - XXL

 

GPT-2 - medium, large

 

GPT-3 175B