LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS

Abstract

NLP의 중요한 패러다임은 우선 제너럴 도메인 데이터에 대해 large scale 사전학습을 시킨 후, 특정 태스크나 도메인들에 대해 적응(adaptation)시키는 것이다. 우리가 더 큰 모델을 full finetuning으로 학습시킴에 따라 full finetuning의 실현 가능성은 점점 더 작아진다. (full fine tutning = 모델의 모든 파라미터를 포함) 175B 파라미터를 가진 GPT-3를 예로 봤을 때, 175B개 파라미터를 가진 독립적인 인스턴스를 배포하는 것은 비용이 매우 매우 커진다. 이에 연구진은 Low-Rank Adaptation, LoRA를 제안한다. 이는 pre-trained 모델의 weights를 freeze 시키고 trainable rank decomposition matrices를 트랜스포머 아키텍처의 각 레이어에 주입하는 것이다. 이렇게 함으로써 목표 태스크(downstream task)의 trainable parameters의 수를 대폭 줄일 수 있다. 즉, LLM의 모든 파라미터를 full finetuning은 현실적으로 비싸니까 그냥 freeze 시키는 대신 각 레이어에 rank decomposition matrices를 넣어줌으로써 trainable parameter를 줄이겠다는 것. 최대한 압축시켜 fine tuning을 가볍게 할 수 있는 방법이라고 볼 수 있겠다.

아담으로 파인튜닝된 GPT-3 175B와 비교하면, LoRA는 trainable parameters의 수를 10000배 줄일 수 있으며 GPU 메모리 필요량은 3배나 감축 가능하다고 한다. LoRA는 RoBERTa, DeBERTa, GPT-2, and GPT-3에 대해 파인 튜닝 퀄리티가 동등하거나 더 높은데, 심지어 trainable parameters도 더 적고, training 처리량(throughput)이 높으면서, adapters와 달리 추가적인 inference 지연이 없다. 연구진은 LoRA의 효능을 입증하는 언어 모델 adaptation에서 rank-deficiency에 대한 실험적인 조사 또한 제공한다.

Introduction

NLP에서 대부분의 application은 하나의 large scale, pre-trained 모델을 여러개의 다운스트림 application에 adapt 시키는 것에 의존한다. 그러한 adaptation은 보통 파인튜닝에 의해 수행되는데, 이는 사전 학습 모델의 전체 파라미터를 업데이트한다. 파인튜닝의 가장 주요한 단점은 새로운 모델의 파라미터 수가 원래 모델의 파라미터 수와 같다는 것이다. 더 큰 모델이 몇 달마다 훈련됨에 따라, GPT-2나 RoBERTa large 모델에서는 그저 ‘불편함’에 지나지 않았다면 GPT-3처럼 175B개 파라미터를 가진 큰 모델로 가게 되면서 ‘크리티컬한 배포 챌린지’가 되게 되었다.

그래서 많은 연구들이 이를 완화시키기 위해 일부 파라미터만 adapt 시키거나 새로운 태스크를 위해 외부 모듈을 학습하는 등의 방법을 찾았다. 이러한 방식은 각 태스크를 위한 pre-trained 모델과 더불어 적은 수의 task-specific 파라미터들만 저장하고 로드시키면 되기 때문에 배포 시 운영 효율을 크게 향상시킨다. 그러나, 현존하는 기술들은 모델 depth를 확장하거나 모델의 사용가능한 시퀀스 길이를 줄임으로 인해 대부분 inference latency를 동반한다. 더욱 중요한 것은, 이러한 방법들이 종종 파인튜닝 베이스라인과 매칭되는 데 실패하는 경우가 많으며 효율성과 모델 퀄리티간의 trade-off가 발생한다는 것이다.

연구진은 학습된 over-parameterized 모델들은 실제로 낮은 intrinsic 차원에 존재한다는 Li et al. (2018a); Aghajanyan et al. (2020)에서 영감을 얻었다. 연구진은 모델 adaptation 중 weight의 변화 또한 낮은 intrinsic rank를 가진다고 가정하며, LoRA에 대한 접근법을 제안한다. LoRA는 미리 훈련된 weights들을 고정한 채로 유지하면서, 밀집된 레이어들의 변화에 대한 rank decomposition 행렬들을 최적화시킴으로써 신경망의 일부 밀집된 레이어들을 간접적으로 훈련 가능하게 해준다. GPT-3의 175B 모델을 예로 들면, full rank가 12,288로 높더라도 아주 낮은 rank로 충분할 수 있음을 보여주며, LoRA가 저장과 계산에 모두 효율적이라는 것을 입증한다. 아래는 LoRA의 핵심 장점들이다.

  • pre-trained 모델은 share 될 수 있으며 각기 다른 태스트들을 위한 많은 작은 LoRA 모듈들로 build될 수 있다. shared model은 고정시켜놓고, 행렬 A와 B를 바꿔가며 효율적으로 작업을 스위치한다. Figure 1을 보면, 필요 저장공간과 작업 스위치의 오버헤드를 눈에 띄게 단축시키는 것을 볼 수 있다.
  • LoRA는 훈련을 더욱 효과적으로 만들고, 하드웨어 장벽을 최대 3배까지 낮춘다. adaptive optimizer를 사용할 때 대부분의 parameter들에 대해 그라디언트를 계산하거나 옵티마이저 state들을 유지할 필요가 없기 때문이다. 대신에, 보다 훨씬 작은 주입된 low-rank 행렬들만 최적화시킨다.
  • 이러한 단순한 선형 디자인은 배포 시 고정된 weight과 trainable 행렬들을 병합할 수 있기 때문에 fully fine-tuned 모델과 비교해 inference latency가 없다.
  • LoRA는 이전의 많은 방법들과 직교하며 대부분과 병합될 수 있다. 예를 들어 prefix-tuning. 이에 대한 예시는 Appendix E에 나와있음.

Terminologies and Conventions

  • Transformer 레이어의 input과 output 차원 ⇒ $d_{model}$
  • $W_q, W_k, W_v, W_o$ ⇒ self-attention 모듈의 쿼리/키/value/output projection 행렬들
  • $W or W_0$ ⇒ pre-trained weight 행렬
  • $\Delta W$ ⇒ accumulated gradient update during adaptation
  • $r$ ⇒ rank of LoRA module
  • Adam optimization 사용
  • Transformer MLP feedforward dimention $d_{ffn} = 4*d_{model}$ 사용

Problem Statement

LoRA는 목적함수에 관계 없이 적용 가능하지만, 언어 모델링에 집중한다. 아래는 언어 모델링 문제에 대한 간략한 설명이며, 특히 작업별 프롬프트가 주어진 조건부 확률의 최대화에 관한 것이다.

$\phi\(phi\)P_{\phi}(y x)$ pre-trained autoregressive language model이 주어졌다고 생각하자. 예를 들어, $P_{\phi}(y x)$ 는 Transformer 기반의 GPT와 같은 일반화된 멀티태스크 모델이 될 수 있다. 이를 summarizaiton, machine reading comprehension, natural language to SQL과 같은 조건부 텍스트 생성 태스크에 adapting한다고 생각해보자. 각 목적 태스크들은 context-target pairs $Z={(x_i, y_i)}_{i=1, …,N}$, ($x_i$와 $y_i$는 토큰들의 sequences)의 training dataset에 의해 나타내어진다. 예를 들어, NL2SQL 태스크에서는, $x_i$가 자연어 쿼리이고 $y_i$가 이에 대응하는 SQL 명령이 되고, summarization에서는 $x_i$가 article의 내용 그리고 $y_i$ 가 이의 summary가 된다. full fine-tuning동안에 모델은 조건부 언어 모델링 objective를 최대화하기 위해 pre-trained weights $\phi_0$ 으로 초기화되며 gradient를 따라 반복적으로 $\phi_0 + \Delta \phi$ 로 업데이트 된다.
\[max_\phi \sum_{(x,y)\in Z} \sum_{t=1}^{\vert y \vert} log(P_\phi (y_t\vert x,y_{<t}))\]

full fine-tuning의 가장 주요한 단점 중 하나는 각 목적 태스크에 대해, 차원 $\vert \Delta \phi \vert$가 $\vert\phi_0 \vert$인 각기 다른 세트의 $\Delta \phi$ 를 학습한다는 것이다. 따라서, pre-trained 모델이 큰 경우(GPT3과 같이) 많은 독립적인 인스턴스들을 저장하고 배포하는 것이 챌린징할 수 있다. (can be challenging, if at all feasible)

이 연구에서는, 파라미터적으로 효율적인 접근법을 채택했고, 이는 task-specific 파라미터 증분 $\Delta \phi = \Delta \phi (\Theta)$ 가 $\vert \Theta \vert \ll \vert \Phi_0 \vert$인 훨씬 작은 사이즈의 파라미터 세트 $\Theta$로 인코딩된다. 따라서 $\Delta \Phi$를 찾는 태스크는 $\Theta$에 대해 최적화된다.

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

후속 섹션에서, $\Delta \Phi$를 인코딩하기 위한 계산, 메모리에 효율적인 low-rank representation을 제안한다. pre-trained 모델이 GPT-3 175B 일 때, trainable parameters $\vert \Theta \vert$는 $\vert \Phi_0 \vert$의 0.01%만큼 축소될 수 있었다.

Our Method

Low-Rank-Parameterized Update Matrices

신경망은 matrix multiplication을 수행하는 많은 밀집 레이어들을 포함하고 있다. weight matrices들은 흔히 full-rank를 가진다. 특정 태스크에 adapt 시킬 때, Aghajanyan et al.(2020)는 pre-trained 언어 모델들이 낮은 ‘intrinsic dimention’을 가지며 더 작은 subspacec에 대한 랜덤 프로젝션에도 불구하고 여전히 효율적으로 학습될 수 있다는 것을 보여주었다. 이것에 영감을 받아, 연구진은 adaptation중의 가중치의 업데이트도 또한 intrinsic rank를 가질 것이라 가설을 세웠다. pre-trained 가중치 행렬 $W_0 \in R^{d\times k}$에 대해, 후자를 low-rank decomposition $W_0 + \Delta W = W_0 + BA,$ 이 때 $B \in R^{d \times r}, A \in R^{r \times k}$으로 나타내면서 이의 업데이트를 제한한다. 이 때 rank $r \ll min(d,k)$ 이다. 훈련동안, $W_0$는 고정되며 그라디언트 업데이트를 받지 않지만 $A$ 와 $B$는 trainable parameters를 포함한다. $W_0$과 $\Delta W = BA$는 동일한 input에 곱해지며, 그들 각각의 output 벡터는 좌표 방식(coordinate-wise)으로 합산된다는 것을 기억하자. $h = W_0x$ 의 경우, forward pass는 다음과 같다.

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

reparametrization를 figure 1에 표현했다. $A$에는 무작위 가우시안 initialization을, $B$에는 0을 사용했으며, $\Delta W = BA$는 훈련 초기에 0이 된다. 그러고나서 $\Delta Wx$를 $\frac{\alpha}{r}$로 스케일한다. $\alpha$는 $r$의 상수이다. Adam을 이요하여 옵티마이징 할 때, 만일 initialization을 적절히 스케일했다면 $\alpha$를 튜닝하는 것은 대강 learning rate를 튜닝하는 것과 동일하다. 결과적으로, 간단하게 $\alpha$를 시도하는 첫번째 $r$로 세팅하고 조정하지 않는다. 이러한 스케일링은 $r$을 변경할 때, 하이퍼 파라미터를 재조정할 필요성을 줄이는 데 도움이 된다.

Applying LoRA to Transformer

원칙적으로 LoRA는 trainable parameter의 수를 줄이기 위해 가중치 행렬의 어느 subset에도 적용될 수 있다. 트랜스포머 아키텍처의 self-attention 모듈에는 총 네 개의 행렬($W_q, W_k, W_v, W_o)$들이 존재하고, MLP 모듈에 두 개의 행렬이 존재한다. 출력 차원이 일반적으로 attention heads로 분할되더라도 $W_q$ (or $W_k, W_v)$를 차원 $d_{model} \times d_{model}$의 단일 행렬로 취급한다. 이 연구는 attention weights를 adapt하는 데 제한하며 단순성과 파라미터 효율성을 위해 MLP 모듈들은 고정되고 목적 태스크에 대해 학습되지 않는다. 섹션 7.1에서 트랜스포머의 다양한 유형의 어텐션 가중치 행렬을 adapting하는데에 대한 효과를 더 연구한다. MLP 레이어, LayerNorm 레이어, biase에 대한 adapting의 실험적 조사는 future work로 남겨둔다.

Practical Benefits and Limitations

가장 주요한 이점은 메모리와 저장공간 사용량의 감축이다. Adam으로 학습되는 대규모 트랜스포머의 경우, 고정된 파라미터에 대해 옵티마이저 states를 저장할 필요가 없으므로 VRAM 사용량을 최대 2/3까지 줄였다.

Conclusion and Future Work

거대한 언어모델을 파인튜닝하는 것은 요구되는 하드웨어와 각기 다른 태스크의 독립적인 인스턴스들에 대한 저장공간/변환 비용이 엄청나게 크다. 이로써 LoRA를 제안했으며, 이는 inference latency가 없고, input sequance 길이를 줄이지 않으면서 높은 모델 퀄리티를 유지하는 효율적인 adaptation 전략이다. 중요한 것은, 대부분의 모델 파라미터를 공유하여 실제 서비스에서 배포 시 빠른 task 전환을 허용한다는 것이다. 해당 연구에서는 트랜스포머 언어 모델들에 집중했으나, 제안된 방법들은 밀집된 레이어를 가진 어느 신경망에도 일반적으로 적용 가능하다.

future works에 대해서는 여러 방향성이 있다. 1) LoRA는 다른 효과적인 adaptation 방법들과 결합되어 잠재적으로 직교 개선(orthogonal improvement)을 제공할 수 있다. 2) 파인튜닝 혹은 LoRA의 매커니즘은 명확하지 않다. pre-training 동안 학습된 피처들이 어떻게 목적 태스크에 잘 작동하도록 변환되는지에 대해.. full fine-tuning에 비해 LoRA는 이를 더욱 tractable 하게 만든다고 믿는다. 3) 연구진은 LoRA를 적용할 weight matrices들을 선택하는 데 대부분 휴리스틱에 의존했다. 보다 더 원칙적인 방법이 있을지? 4) 마지막으로, $\Delta W$의 rank-deficiency는 $W$ 또한 rank-deficient 할 수 있음을 시사하며, 이는 future work의 또다른 영감 소스가 될 수 있다.

참고문헌

https://www.youtube.com/watch?v=BJqwmDpa0wM