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

이상민·2023년 12월 13일

논문리뷰

목록 보기

7/29

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

Abstract

Natural Language Processing 분야에서 중요한 패러다임은 일반적인 분야의 데이터로 large-scale pretraining하고 특정 task나 분야에 적응시키는 것이다. 큰 pre-trained model을 full fine-tuning 하는 것은 너무 비싸서 점점 하기 힘들다. 그래서 “LoRA(Low-Rank Adaptation)”을 제안한다. Pre-trained model의 가중치들을 얼려놓고 훈련 가능한 rank decomposition matrices를 transformer 구조의 각 layers에 주입하면서, downstream tasks를 위한 훈련 가능한 parameters의 수를 크게 줄였다. Adam을 이용해 GPT-3 175B를 fine-tuning한 것과 비교해서 LoRA는 훈련 가능한 parameters의 수를 10,000배 줄였고, GPU memory 요구량을 3배 줄였다. LoRA는 RoBERTa, DeBERTa, GPT-2, GPT-3에서 모델의 quality를 fine-tuning하는 것보다 비슷하거나 더 나은 성능을 보인다. 그렇지만 훈련 가능한 parameters를 더 적게 가지고, 훈련 처리량이 높고, adapters와 달리 추가적인 inference latency(추론 지연)가 없다. LoRA의 효율성을 보여줄 수 있는, LM adaptation에서 rank-deficiency에 대한 경험적인 조사를 제공한다.

Rank-deficiency
full-rank 가 아닌 것을 의미한다.

1. Introduction

요즘 NLP에서 applications는 하나의 큰 pre-trained Language model을 다수의 downstream applications에 적응(adapt)시킨다. 이러한 적응(adaptation)은 대게 모델의 모든 parameters를 업데이트하는 fine-tuning을 통해 이루어진다. Fine-tuning의 주요한 단점은 new model이 original model과 같은 수의 parameters를 갖는다는 것이다. 몇 달마다 더 큰 모델이 훈련됨에 따라, 이는 GPT-2(Radford et al., b) 또는 RoBERTa large(Liu et al., 2019)에서는 단순한 "불편"이었지만, 1,750억 개의 훈련 가능한 매개 변수를 가진 GPT-3(Brown et al., 2020)에서는 중요한 배포 과제로 바뀐다.

많은 사람들은 약간의 parameters만 적응시키거나, 외부 모듈을 학습시키는 것을 통해 이를 완화시키려고 했다. 이렇게 하면, 각 task에 대해 pre-trained model에 작은 수의 task-specific parameters만 저장하고 로드하면 되므로, 배치할 때 운영 효율성을 크게 향상된다. 하지만 기존의 방법들은 모델의 깊이를 늘리거나 모델의 사용 가능한 sequence 길이를 줄임으로써 inference latency를 도입한다. 더 중요한 것은 이러한 방법이 fine-tuning baselines과 일치하지 않는 경우가 많아, 효율성과 모델 퀄리티의 trade-off를 일으킨다.

학습된 over-parametrized models이 실제로 낮은 고유 차원에 존재한다는 것을 보인 논문들에서 우리는 영감을 받았다. 모델 적응 동안의 가중치의 변화 또한 낮은 고유 rank를 가진다는 가설을 세우고, 이는 우리가 제안할 LoRA approach로 이어진다. (’Adaptation 중 weight의 변화는 low intrinsic rank를 가진다’ 가설 → LoRA approach) LoRA는 neural network의 몇몇 dense layers를 간접적으로 훈련시킬 수 있게 해준다. Pre-trained 가중치는 얼려놓으면서 대신, adaptation 중에 dense layers’ change의 rank decomposition matrices를 최적화한다. Figure 1에서 볼 수 있음!

GPT-3 175B를 예로 들면, full rank( 즉, $d$ )가 12,288만큼 높아도 매우 낮은 rank(즉, Figure 1의 $r$ 은 1이나 2일 수 있음)도 충분하다는 것을 보였다. 이는 LoRA가 저장적으로도 계산적으로도 효율적이게 만들어준다.
LoRA는 몇몇 주요 이점을 갖는다.

Pre-trained model은 공유될 수 있고, 다른 tasks에 대한 많은 작은 LoRA 모듈을 만드는데 사용될 수 있다. Figure 1의 matrix A와 B를 교체하여 tasks들을 효율적으로 전환할 수 있어, storage requirement와 task-switching overhead를 크게 줄인다.
Adaptive optimizers를 이용할때, LoRA는 훈련을 더 효율적으로 만들고 hardware 진입 장벽을 3배 낮춘다. 왜냐하면 대부분의 parameters에 대해 gradients를 계산하고, optimizer 상태를 유지할 필요가 없기 때문이다. 대신에 주입된 훨씬 작은 low-rank matrices를 최적화한다.
간단한 linear 디자인은 배치 시 훈련가능한 matrices와 얼린 가중치를 병합할 수 있게 하며, 구조에 의해 fully fine-tuned model에 비해 inference latency를 도입하지 않는다.
LoRA는 많은 이전의 방법들에 수직적이고, prefix-tuning과 같은 많은 방법들과 합쳐질 수 있다. Appendix E에서 예시를 볼 수 있다.

Terminologies and Conventions

Transformer 구조를 자주 참조하고, 그 차원에서 기존의 단어를 사용한다. Transformer layer의 input과 output dimension을 $d_{model}$ 이라고 부른다. Self-attention module에서 query/key/value/output projection matrices의 의미로 $W_q, W_k, W_v, W_o$ 을 사용한다. $W$ or $W_0$ 은 pretrained weight matrix와 adaptation 중 누적 gradient update $\Delta W$ 를 의미한다. LoRA module의 rank을 $r$ 로 쓴다. 우리는 Trasnformer 논문의 규칙들을 따르고, 모델 optimization으로 Adam을 이용하고, Transformer MLP feedforward dimension으로 $d_{ffn} = 4 \times d_{model}$ 을 이용한다.

2. PROBLEM STATEMENT

우리의 제안은 훈련 목표와는 무관하지만, 동기 부여 사용 목적으로써 언어 모델링의 초점을 맞춘다. 아래는 언어 모델링 문제의 간결한 설명이 있고, 특히 주어진 task-specific prompt의 조건부 확률 최대화이다.

$\Phi$ 로 parameterized된 autoregressive model $P_{\Phi}(y|x)$ 이 주어졌다고 하자. $P_{\Phi}(y|x)$ 는 Transformer 구조를 기반으로한 GPT와 같은 일반적인 multi-task learner가 될 수 있다. 이러한 pre-trained model이 요약, MRC, NL2SQL과 같은 downstream 조건부 text generation task에 적응하는 것을 생각해보자. 각 downstream task은 context-target 쌍의 훈련 데이터셋으로 표시된다:

\mathcal{Z} = \{(x_i, y_i)\}_{i=1,\dots,N}

여기서 $x_i$ 와 $y_i$ 는 tokens의 sequence이다.

예를 들어, N2SQL 경우 $x_i$ 는 자연어 query이고 $y_i$ 는 그에 대응하는 SQL 명령어이고, summarization 경우 $x_i$ 는 기사의 내용이고 $y_i$ 는 그의 요약이다.

Full fine-tuning 동안 모델은 $\Phi_0$ 의 가중치로 초기화되고, 조건부 언어 모델링의 목표를 최대화하기 위해 기울기를 반복적으로 따라 $\Phi_0 + \Delta \Phi$ 로 업데이트한다:

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

Full fine-tuning의 주요 단점 중 하나는 각 downstream task마다 dimension $|\Delta \Phi|$ 이 $|\Phi_0|$ 과 같은 parameters $\Delta \Phi$ 의 다른 집합들을 배운다. 그래서 pre-trained model이 크면( $|\Phi_0| \approx$ 175B인 GPT-3과 같은), fine-tuned 모델들의 많은 독립적인 예시들을 저장하고 배치하는 것이 어려울 수 있다. 저장하고 배치하는 것이 가능하다고 한다면!

이 논문에서는 더 parameter-efficient approach를 선택할 것이다. 이는 task-specific parameter increment $\Delta \Phi = \Delta \Phi (\Theta)$ 가 $|\Theta| \ll |\Phi_0|$ 인 parameter $\Theta$ 의 훨씬 더 작은 집합으로 인코딩된 parameter-efficient approach다. 그래서 $\Delta \Phi$ 를 찾는 문제는 $\Theta$ 에 대해 최적화하는 문제가 된다:

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

이후의 sections에는 계산도 메모리도 효율적인 $\Delta \Phi$ 를 인코딩하는데 low-rank representation를 사용하는 것을 제안한다. Pre-trained model이 GPT-3 175B일 때, 훈련 가능한 parameters $|\Theta|$ 의 수는 $|\Phi_0|$ 의 0.01%만큼 작아질 수 있다.

3. AREN’T EXISTING SOLUTIONS GOOD ENOUGH?

우리가 해결하기 위해 시작한 문제는 새로운 것이 아니다. Transfer learning이 시작된 이후, 수십 개의 작업들이 보다 계산 효율적, 메모리 효율적으로 모델 적응을 만들기 위해 노력해왔다. Language modeling을 예시로 이용하면, 효율적인 적응에 대한 2개의 주요한 전략들이 있다:

Adapter layers를 추가하는 것
Input layer activations의 몇몇 형식을 최적화하는 것

그러나 두 전략 모두 한계가 있었고, 특히 큰 규모와 latency-sensitive production 시나리오에 한계가 있었다.

Adapter Layers Introduce Inference Latency

많고 다양한 adapters가 있지만, transformer block당 두 개의 adapter layers를 가지는 original design(2019)과 transformer block 당 additional LayerNorm(2016)과 한 개의 adapter layer 가지고 있는 보다 최근 것(2020)에 초점을 둘 것이다. 하나는 계층을 가지치기 하거나 multi-task setting을 활용하여 전체 latency를 줄일 수 있지만, adapter layer에서 추가 계산을 우회할 직접적인 방법은 없다. Adapter layer는 그들이 추가할 수 있는 FLOPs을 제한할 정도의 작은 bottleneck dimension을 가지면서 적은 parameter(때때로 <1% original model)로 설계되었기 때문에 위는 문제가 되지 않는 것 같다. 그러나 큰 neural networks는 latency를 낮게 유지하기 위해 hardware 병렬에 의존하고 있고, adapter layers는 순차적으로 처리되어야 한다. 이는 일반적으로 배치 크기가 1개만큼 작은 online inference setting에서 차이를 만든다. 단일 GPU에서 GPT-2 medium에서 추론을 실행하는 것과 같이 모델 병렬성이 없는 일반적인 시나리오에서는 매우 작은 병목 차원에서도 adapter를 사용할 때 latency이 눈에 띄게 증가하는 것을 볼 수 있다. (표 1)

(Adapter layers는 pretrained model에 순차적 방식으로 더해지는 외부 모듈, 반면 LoRA는 병렬적 방식으로 pretrained model에 더해짐. Adapter layers는 base model 외에도 추가적으로 계산되어야 하기에 추가 latency가 발생할 수 밖에 없다.)

이 문제는 모델을 공유해야 할 때(shard the model; “Megatron-lm: Training multi-billion parameter language models using model par- allelism, 2020”) 더 안좋아진다. 왜냐하면 adapter parameters를 여러 번 중복 저장하지 않는 한 추가 깊이에는 AllReduce 및 Broadcast와 같은 더 많은 동시의(synchronous) GPU 작업이 필요하기 때문입니다.

Directly Optimizing the Prompt is Hard

예를 들어 prefix tuning과 같은 다른 방법은 다른 문제에 직면한다. Prefix-tuning은 최적화하기가 힘들고, 훈련 가능한 parameters에서 성능이 비단조적으로(non-monotonically) 변화하는 것을 관찰했다. 더 근본적으로, adaptation을 위해 sequence 길이의 일부를 예약하는 것은 downstream task를 처리하는데 이용가능한 sequence 길이를 반드시 줄인다. 이로 인해 prompt를 조정하는 것이 다른 방법에 비해 성능이 떨어지는 것으로 의심된다. Task 수행에 대한 연구를 Section 5로 연기한다.

4. OUR METHOD

LoRA의 간단한 디자인과 실용적인 이점을 설명할 것이다. 여기에 적용된 원칙들은 딥러닝의 어느 dense layer에나 적용되지만, motivating use case로써 Transformer LM의 특정 가중치에 초점에만 초점을 맞춘다.

4.1 LOW-RANK-PARAMETRIZED UPDATE MATRICES

Neural Network는 matrix multiplication을 하는 많은 dense layers를 포함하고 있다. 이런 layers에서 가중치 matrices는 일반적으로 full-rank이다. 특정 task에 적응할 때는 Aghajanyan et al. (2020)은 pre-trained LM이 낮은 “intrinsic dimension”을 가지고, random projection을 더 작은 subspace에 하더라도 여전히 효율적으로 학습할 수 있음을 보였다. 이것에 영감을 받아, 적응동안 가중치 업데이트도 낮은 “intrinsic rank”를 가진다고 가설을 세웠다. Pre-trained 가중치 matrix $W_0 \in \R^{d \times k}$ 의 경우, $W_0 + \Delta W = W_0 + BA$ 으로 $\Delta W$ 를 $BA$ 인 low-rank decomposition으로 표현하여 업데이트를 제한한다. 여기서 $B \in \R^{d \times r}, A \in \R^{r \times k} \text{rank} r << \min (d,k)$ . 훈련동안 $A,B$ 는 훈련가능한 parameters를 포함하고 있지만, $W_0$ 는 얼리고 기울기 업데이트를 받지 않는다. $W_0$ 와 $\Delta W = BA$ 는 같은 input과 곱해지고, 그들 각각의 output 벡터들은 좌표별로 더해진다. $h=W_0x$ 의 경우, 우리가 변화시킨 forward pass는 다음과 같다:

h=W_0x + \Delta Wx = W_0x+BAx

Figure (1)에서 reparametrization을 설명한다. $A$ 는 random Gaussian 초기화으로, $B$ 는 0을 이용했기에, $\Delta W = BA$ 는 훈련 시작 시 0이다. 다음, $\Delta W$ 를 $\frac{\alpha}{r}$ 로 scaling한다. 여기서 $\alpha$ 는 $r$ 에서 상수이다. Adam을 최적화할 때 우리가 적절히 초기화 스케일링을 한다면, $\alpha$ 를 튜닝하는 것은 학습률을 튜닝하는 것과 비슷하게 본다. 결과적으로, 단순히 $\alpha$ 를 우리가 시도하는 첫 번째 $r$ 로 설정하고 조정하지 않는다. 이러한 스케일링은 다양한 $r$ 에 대해 hyperparameter를 재튜닝할 필요를 줄이는데 도움이 된다. ( $\alpha$ 가 첫 번째 $r$ 이니까, $r$ 이 다양해져도 $\alpha$ 는 1개 고정되어 있으니까!)

A Generalization of Full Fine-tuning

보다 일반적인 형태의 fine-tuning을 통해 pre-trained parameters의 subset을 훈련할 수 있습니다. LoRA는 한 단계 나아가, 적응 중에 축적된 기울기 업데이트가 full-rank 가중치 matrices를 가질 필요가 없다. 이것은 모든 가중치 matrices에 LoRA를 적용하고 모든 편향을 훈련할 때, LoRA rank $r$ 을 pre-trained 가중치 matrices의 rank로 설정함으로써 full fine-tuning의 표현력을 대략 회복할 수 있다. 다르게 표현하면, 훈련 가능한 parameters의 수를 증가시키면, LoRA 훈련은 본래의 모델을 훈련하는 것에 대략적으로 수렴한다. 반면 Adapter 기반의 방법은 MLP로 수렴하고, prefix 기반 방법은 긴 input sequences를 받을 수 없는 모델로 수렴한다.

No Additional Inference Latency

Production을 배치할 때, $W = W_0 + BA$ 를 명시적으로 계산하고 저장할 수 있고, 평소와 같이 inference를 수행할 수 있다. $W_0$ 와 $BA$ 는 $\R^{d \times k}$ 에 있다. 다른 downstream task로 변경이 필요할 때, 우리는 $BA$ 를 빼고, 다른 $B'A'$ 를 추가함으로써 $W_0$ 를 회복할 수 있다. 이는 매우 작은 메모리 overhead를 가진 빠른 작업이다. 결정적으로, 이는 구조에 의해 fine-tuned model과 비교했을 때, inference 시에 아무 추가적인 latency를 도입하지 않는 것을 보장한다.

4.2 APPLYING LORA TO TRANSFORMER

원칙적으로, 훈련 가능한 parameters의 수를 줄이기 위해 neural network에서 가중치 matrices의 어느 subset에도 LoRA를 적용할 수 있다. Transformer 구조에서 self-attention 모듈 안에 4개의 가중치 matrices( $W_q, W_k, W_v, W_o$ )가 있고, MLP 모듈에는 2개가 있다.

MLP 에서 왜 2개의 가중치가 있을까?
나는 Feed Forward Neural Network에서 하나의 가중치만 있는 줄 알았는데 그게 아니었다. 애초에 식이
$FFN(x) =ReLU(xW_1 + b_1)W_2 + b_2$
였다! 그래서 2개의 weight matrices가 쓰인 것이었따.. 하하

출력 차원이 보통 attention heads로 잘리지만, 우리는 $W_q(\text{or } W_k, W_v)$ 를 $d_{model} \times d_{model}$ 차원의 한 개의 matrix로 취급한다. 우리는 연구를 downstream task에 대해 오직 attention 가중치들을 적응하는 것으로 제한하고, 간단함과 parameter 효율성 모두를 위해 MLP 모듈을 얼린다(그래야 downstream tasks에 대해 훈련하지 않는다). 우리는 나아가 section 7.1에서 Transformer 안 다른 타입의 attention 가중치 matrices를 적응시키는 효과를 공부할 것이다. MLP layers, LayerNorm layers, biases를 적응시키는 것의 경험적 조사는 향후 작업으로 남겨놓을 것이다.

Practical Benefits and Limitations.

가장 중요한 이점은 메모리와 저장소 사용량의 감소이다. Adam으로 훈련된 큰 trasnformer의 경우, 만약 $r \ll d_{model}$ 이면 frozen parameters에 대해 optimizer 상태를 저장할 필요가 없기에 VRAM 사용량을 2/3까지 줄였다. GPT-3 175B에서는 훈련 중 VRAM 소비량을 1.2TB에서 350GB로 줄인다. $r=4$ 이고 오직 query, value projection matrices만 적응되면, checkpoint 크기는 대략 10,000배 줄어든다(350GB에서 35MB로). 이는 상당히 적은 GPUs를 이용하여 훈련을 할 수 있게 했고, I/O bottleneck을 피할 수 있게 했다. 또 다른 이점은 모든 parameters가 아닌 LoRA 가중치만 바꿔 훨씬 저렴한 비용으로 배포하면서 작업들을 전환할 수 있다는 것이다. 이는 pre-trained weights를 VRAM에 저장하는 기계에서 필요에 따라 빠르게 교체할 수 있는 많은 사용자 정의 모델을 생성하게 해준다. 또한 대부분의 parameters에 대한 기울기를 필요가 없기 때문에 GPT-3 175B를 훈련할 때, full fine-tuning 과 비교하여 25% 속도 향상을 관찰했다.

LoRA는 또한 한계들이 있다. 예를 들어, 추가적인 inference latency를 제거하기 위해 $A$ 와 $B$ 를 $W$ 로 흡수하기로 선택한 경우, 다른 A,B을 가진 다른 tasks들의 inputs을 하나의 forward pass로 일괄처리하는 것은 간단하지 않다. 그렇지만 Latency가 중요하지 않은 시나리오에 대해서는 배치의 샘플을 사용하기 위해 가중치들을 병합하지 않고 동적으로 LoRA 모듈을 선택할 수 있다.

느낀점

기존의 pretrained된 weight matrices들은 얼려놓고, LoRA의 adapter layers를 이용해서 훈련을 시키는데 LoRA의 adapter layers인 $\Delta W$ 의 parameters를 줄이기 위해 low-rank decomposition을 이용했다. 그래서 $\Delta W =BA$ 를 이용했고, 훈련을 시키는 parameters 수를 극적으로 줄였다.

예전부터 pretrained 된 정보는 저장하고, fine-tuning하는 부분을 따로 학습해줄 수 있는 방법이 어떤 것이 있을까? 하고 궁금했었다. 프로젝트를 할 때는 embedding layers를 얼리고 fine-tuning을 한 적도 있었다. 이런 LoRA를 이용하면, 훈련 시키는 parameter 수를 줄여줄 수도 있고, pretrained 된 정보도 온전히 가지고 있을 수 있어서 좋은 방법인 것 같고, 흥미로운 방법이다!