[yongggg's] LoRA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS Review

안녕하세요 이번 장에서는 MicroSoft 사의 LoRA 논문을 소개하겠습니다.

2022년 하반기, 2023년 초반기 시점에서 GPT3, GPT4와 같은 LLM이 나와서 큰 이슈가 되고 있습니다.

이 model들은 강력한 성능을 갖고 있지만, 이 pre-trained model의 전체 parameter를 가지고 fine-tuning 하기엔 time, cost 적으로 너무 비효율적이라는 단점을 갖고있습니다.

이를 해결하기 위해 LoRA라는 연구가 나왔으며, 실제로 fine-tuning 성능도 좋은 것으로 알려져 있습니다.

그럼, 지금부터 해당 논문을 리뷰해 보겠습니다!

Abstract

LoRA(Low Rank Adaption)라는 기술은 pre-trained LLM의 모든 weights를 fine-tuning하지 않고, pre-trained LLM weights는 모두 freeze 하고 down-stream task를 수행하기 위해 훈련 가능한 rank decomposition matrice를 추가 함으로써 parameter를 효율적으로 훈련하는 방법을 제안한다.

sequential한 방식으로 외부 module에 추가되었던 adapter 방식과 다르게, parallel한 방식으로 module에 추가되기 때문에 추가적인 inference latency가 없으며, input sequence length를 줄이지 않고 훈련이 가능하다.

이로써, Adam으로 fine-tuning한 GPT-3와 비교하여, 10,000배 가량의 parameter 수를 줄였고 3배에 가까운 메모리를 줄였으며, RoBERTam DeBERTa, GPT-2, GPT-3에 fine-tuning한 model 보다 훨씬 더 적은 수의 parameter로 더 높은 throughput과 추가적인 inference latency 없이 훨씬 더 좋은 성능을 내고 있다.

Introduction

PLM은 model의 모든 parameter를 업데이트 하는 fine-tuning 방식을 이용하여 다양한 task에 적용했다. 그러나 가중치를 update하는 데에 사용되는 parameter는 극히 일부이기 때문에 자원을 비효율적으로 사용하게 된다.

down-stream task를 위해 external module을 학습하거나 몇몇의 parameter만 adapt하는 방식이 제안되어 왔으며, 다른 방법으로는 각각의 task를 수행하기 위한 pre-trained model에 task-specific parameter를 저장하고 load하는 방법이 제안되어 왔다.

하지만, 기존의 방식으로는 model depth를 늘림으로써 inference latency가 생겨나거나,  model의 입력 가능한 문장 길이를 줄여야 학습이 가능한 경우가 있어, 효율성과 model quality간의 trade off를 고려해야하며, fine-tuning base line에 맞지 않다.

 

따라서, 본 논문의 저자는 over-parametrized model이 실제로 낮은 고유 차원에 존재함을 보여주는 Li et al. (2018a); Aghajanyan et al. (2020) 연구에 영감을 받아 model adaptation을 하는 데, 가중치의 변화 또한 low "intrinsic rank"(고유한 랭크)를 가지고 있다고 가정하고 LoRA를 제안한다. 

 

LoRA는 pre-trained weights를 고정시키는 대신, adaptation 동안 dense layer 변화의 rank 분해 행렬을 최적화 하여 neural network에서 dense layer를 간접적으로 학습시킨다. GPT-3 175B를 예시로 들자면, 전체 rank(d=12888)가 매우 낮은 rank(r=1 or 2) 로도 충분하며, 이로 인해 storage와 computing 자원의 효율성 또한 높일 수 있음을 보여준다.

Benefits

1. replace LoRA for each down-stream task, not fine-tuning all parameters

pre-trained model은 LoRA 모듈이 서로 다른 task에 build하도록 공유되어 사용할 수 있다. pre-trained model의 weights를 freeze하고, 이를 행렬 A와 B로 대체함으로써 storage 사양과 작업 전환 overhead를 크게 줄일 수 있다.

2. no store gradient and just optimize low rank decomposition matrice

LoRA는 adaptation optimizer를 사용해서 대부분의 매개변수에 대한 optimizer state를 유지하거나 gradient를 계산할 필요가 없기 때문에, 학습을 보다 효율적으로 만드며, 진입에 대한 하드웨어 장벽을 최대 3배까지 낮춘다. 대신 low-rank 행렬을 최적화한다.

3. no inference latency

간단한 선형 구조로 추가적인 inference latency (input이 들어가 model이 예측하기까지 걸리는 시간)없이 trainable 행렬을 freeze한 weights와 병합할 수 있다.

4. combining LoRA with prefix tuning

LoRA는 기존에 존재하는 많은 방식들과 직교하기 때문에, prefix 기반 접근과 함께 사용할 수 있다.

Method

$$ \max\limits_{\Phi} \sum\limits_{(x,y) \in Z} \sum\limits_{t=1}^{|y|} log(P_{\Phi}(y_{t}|x,y_{ < t})) \\ \max\limits_{\Theta} \sum\limits_{x,y \in Z} \sum\limits_{t=1}^{|y|} log(p_{\Phi_{0} + \vartriangle \Phi (\Theta)} (y_{t} | x, y_{<t}))$$

• $ \vartriangle \Phi = \vartriangle \Phi (\Theta) ( | \Theta | << | \Phi_{0} | ) $
• $ \vartriangle \Phi $를 찾는 task는 결국 $ \Theta $에 대한 최적화가 된다.
• $ \vartriangle \Phi $를 encoding하기 위해, compute, memory에 효율적인 low-rank representation을 제안한다.

기존에 효율적인 adaptation 방법은 1. adapter layer를 external module에 추가하는 것과 2. 몇 개의 input layer activation form을 최적화 시키는 것이 있었다. 이 두 가지 전략에서 large scale language model에 각각 1. inference latency를 유도한다. 2. 활성화 함수 최적화는 prompt를 직접적으로 최적화하기 어렵다. 라는 문제가 있다.

 

따라서 본 연구의 저자는 다음과 같은 기법들을 제안한다.

1. Low-Rank Parametrized update matrice

학습시 low-rank 행렬과 A와 B를 업데이트하고, pre-trained model의 initiallized weight인 $W_{0}$는 freeze 하여 gradient update를 하지 않는다.

• $ W_{0} + \vartriangle W = W_{0} + BA$ 처럼 $\vartriangle$를 rank 분해 행렬곱으로 표현 하면, 아래와 같은 전진패스를 산출 할 수 있다.
  $ h = W_{0}x + \vartriangle Wx = W_{0}x + BAx $
  - $ W_{0}$는 고정하고 $A$는 무작위 가우스, $B$는 0으로 초기화한다.
  - Projection matrix $W_{q}, W_{k}, W_{v}, W_{o}$가 multi-head에 의해 나누어지더라도, 단일 $ d \times d $행렬로 취급하여 분해한다.
• $A$는 random Gausian initialization, $ B $는 0으로 초기화 되어, $ \vartriangle W $는 학습 초기에 0이다. $ \vartriangle W$는 $ {a \over r} $로 scaling 한다. ($a$는 $r$에서 constant) Adam으로 최적화할 때, 초기화를 적절히 scailing 했다면, $a$를 튜닝하는 것은 learning rate를 튜닝하는 것과 같다. 결과적으로 단순히 $a$를 첫 번째 $r$로 설정하고 조정하지 않는다.

2. Applying LoRA to Transformer

LoRA는 학습 가능한 parameter의 수를 줄이기 위해, 신경망에서 어떤 weights 행렬의 부분집합에 적용이 가능할지 고민했다. 연구 저자는 parameter-efficiency를 위해 down-stream task에 대해 attention weight만 adpating 했고, MLP에선 freeze 하였다. 

Limitation

추가 추론 대기 시간을 제거하기 위해 A와 B를 W로 흡수하기로 선택한 경우, 단일 순방향 패스에서 A와 B를 다른 여러 작업에 대한 입력을 배치 처리하는 것은 간단하지 않다. 가중치를 병합하지 않고 대기 시간이 중요하지 않은 시나리오의 경우 배치에서 샘플에 사용할 LoRA 모듈을 선택할 수 있다.

Experiment

Baseline

1. Prefix-Embedding tuning (PreEmbed)
   input token에 special token lp(prefix token의 개수)와 li(infix token의 개수)를 주입한다. 학습 가능한 word embedding을 가지고 있고, model vocab에는 포함되어 있지 않는다.
   "prefixing"에 초점을 맞춰, prompt에 그런 토큰을 배치하고 infixing(문장 중간에 들어가는 조사 등을 교체하는 작업)을 하여 prompt에 추가한다. 학습가능한 매개변수의 개수는 다음과 같이 계산한다. 
   $|\Theta| = d_{model} \times (lp + li)$
2. Prefix-layer tuning(PreLayer)
   prefix-embedding tuning의 확장 버전으로 word embedding이나 embedding layer 다음 activation만 학습하는 것이 아닌, 모든 transformer layer 이후의 activation을 학습한다. 학습가능한 매개변수의 개수는 다음과 같이 계산한다.
   $|\Theta| = L \times d_{model} \times (lp + li) $
3. Adapter tuning
   self attention module과 residual connection사이에 adapter layer를 추가하는 것이다.
   AdapterH: 중간에 비선형성이 있는 adapter layer에서는 bias를 가진 2개의 fully connected layer가 있다. (original)
   AdapterL: MLP module과 LayerNorm 후에만 적용되는 adapter layer로 되어 조금 더 효율적이다.
   AdapterP: 위의 Adapter와 비슷하지만 Pfeiffer에 의해 제안되었다.
   AdapterD: AdapterDrop으로 불린다.
4. LoRA
   LoRA는 기존 가중치 행렬과 parallel하게 훈련 가능한 rank decomposition 행렬 Pair를 추가하는 기법이다. 단순하게 LoRA를 $W_{q}$ 및 $W_{v}$에만 적용한다. 학습 가능한 매개변수의 수는 rank r과 원래 가중치의 모양에 따라 결정된다.
   $$|\Theta | =2 \times \hat{L}_{LoRA} \times d_{model} \times r \\ where \; \hat{L}_{LoRA} : LoRA를 \; 적용하는 \; 가중치 \; 행렬의 \; 개수 $$
   학습된 weight는 추론 중 main weight로 merge함으로써 어떠한 latency를 일으키지 않는다.
5. LoRA + PrefixEmbed (LoRA+PE)
   prefix-embedding tuning과 결합한 형태로, special token $l_{p} + l_{i}$를 넣어준다. 이러한 embedding은 훈련 가능한 매개변수로 취급한다.
6. LoRA + PrefixLayer (LoRA+PL)
   prefix-layer tuning과 결합한 형태로, special token $l_{p} + l_{i}를 넣어준다. 하지만, token의 hidden representation을 모든 transformer 블록 이후에 input-agnostic vector로 대체한다. 따라서, 임베딩과 활성화함수는 훈련 가능한 매개변수로 취급한다.

• 훈련 가능한 매개변수의 수를 늘리면 PrefixEmbed와 PrefixLayer 기법의 성능이 안좋아지는 반면, LoRA 기법은 안정적인 것을 볼 수 있다.
• LoRA + PrefixEmbed(LoRA + PE)가 LoRA와 LoRA + PrefixLayer (LoRA + PL)보다 성능이 좋은데, 이는 LoRA가 prefix-embedding tuning에 orthgonal하기 때문이다.
• LoRA + PrefixLayer(LoRA + PL)는 LoRA보다 더 많은 매개변수로 훈련함에도 불구하고 성능이 좋지않다. 이는 prefix-layer tuning은 learning rate 선택에 민감하여 LoRA+PL에서 LoRA 가중치를 최적화 하기 어렵기 때문이다. 

Low-Rank Update 는 무엇인가?

Downstream task를 위해 학습한 Low-rank adaptation은 parallel하게 훈련 가능하기에 하드웨어 장벽이 낮아지며, 업데이트된 가중치와 사전학습되었던 가중치가 어떤 상관관계가 있는지 더욱 잘 설명할 수 있다.

어떤 가중치 행렬에 LoRA를 적용해야 하는가?

• 모든 파라미터에 LoRA를 적용하기 보다, $W_{q}$와 $W_{v}$ㅇ[만 적용하는 것이 가장 좋은 성능을 내고있다.
• rank가 4인 경우에도 $\vartriangle W$에서 충분한 정보를 포착하기 때문에, rank가 더 큰 단일 유형의 가중치를 적용하는 것보다 더 많은 가중치를 갖는 matrix를 적용하는 것이 바람직하다.

LoRA를 위한 최적의 rank는 무엇인가?

• Update matrix rank는 작은 intrinsic rank를 갖고 있으며, r이 증가한다고 더 의미있는 subspace를 커버하지는 않는다. 따라서 low-rank adaptation 행렬로도 충분하다고 말한다.

r의 크기에 따라 학습된 subspace의 유사성

rank 8과 64로 학습한 adaptation 행렬을 singular value decomposition해서 얻은 right-singular unitary 행렬은 top i 개의 singular vector에 의해 얼마나 subspace가 span 되는지 Grassmann 거리를 기반으로 normalized subspace의 유사도를 정량화 했다.
$$ \phi (A_{r=8}, A_{r=64}, i, j) = {{||{U^{i}_{A_{r=8}}}^{T} U^{j}_{A_{r=64}}||}^{2}_{F}\over{min(i,j)}} \in [0,1] $$

• $A_{r=8}$과 $A_{r=64}$의 최상위 singular vector에서 overlap이 일어나며, 이 둘의 가중치 변화량은 dimension 1의 subspace의 상관관계가 0.5보다 크다. (본 연구에서 GPT-3의 Downstream 가 r=1로도 잘 수행되는지에 대한 이유가 됨)

여기서 singular vector란 고유값 분해로 분해된 직교행렬의 열 벡터를 말한다. $(A=U \sum V^{T}$ 에서 얻어진 $U$(left singular)와 $V$(right singular)의 열벡터 $)$

adaptation 행렬과 기존 가중치 행렬의 관계

low rank adaptation 행렬은 학습된 특정 downstream task에 대한 중요 feature를 강화하지만, 일반적인 사전학습 모델의 feature는 강화되지 않는다.