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

이번에 다뤄볼 논문은 Low-Rank Adaptation of Large Language Models입니다.

Abstract

자연어 처리의 중요한 패러다임 중 하나는 일반 도메인 데이터에서 대규모 사전 학습을 진행한 후 특정 작업 또는 도메인에 대한 adaptation을 하는 것입니다. pre-training하는 모델이 커질수록, 전체를 fine-tuning하는 것은 비효율적입니다.

이 논문은 ‘Low-Rank Adaptation(LoRA)’를 제안합니다. LoRA는 사전 학습된 모델 가중치를 고정하고 Transformer 아키텍처의 각 레이어에 학습 가능한 랭크 분해 행렬(rank decomposition matrices)를 주입함으로써 downstream task에 대한 학습 가능한 파라미터 수를 크게 줄이는 방법입니다. 이를 통해 GPT-3 175B를 Adam으로 fine-tuning하는 것과 비교하여 LoRA는 학습 가능한 파라미터 수를 10,000배 줄이고 GPU 메모리 요구 사항을 3배로 줄일 수 있습니다.

LoRA는 RoBERTa, DeBERTa, GPT-2 및 GPT-3에서 모델 품질 면에서 fine-tuning과 동등하거나 더 나은 성능을 보입니다. 그럼에도 불구하고 LoRA는 더 적은 학습 가능한 파라미터를 가지며 높은 학습 처리량을 제공하며, adapter와는 달리 추가적인 inference latency를 갖지 않습니다.

Introduction

Abstract에서 소개했던 것처럼 자연어 처리는 하나의 대규모로 pre-train된 언어 모델을 downstreak task에 맞게 fine-tuning하는 방식으로 진행됩니다. fine-tuning을 통해 pre-trained model의 모든 파라미터를 업데이트하기 때문에, pre-trained model이 커질수록 업데이트해야 할 파라미터 역시 많아집니다.

많은 연구에서 이 문제를 완화하기 위해 일부 파라미터만 업데이트시키거나 새로운 task를 위해 외부 모듈을 학습하는 등의 방법을 사용했습니다. 이렇게 함으로써 각 task마다 사전 학습된 모델에 추가로 task 특정 파라미터를 저장하고 로드하기만 하면 되어, 배포 시 운영 효율을 크게 항상시킬 수 있습니다. 그러나 기존의 기술은 모델의 깊이를 확장하거나 모델의 사용 가능한 시퀀스를 줄이는 방식으로 model inference latency를 도입합니다. 이와 관련된 내용은 아래에서 다시 설명하겠습니다.

본 논문은 학습된 과도한 매개변수를 가진 모델이 사실 low intrinsic dimenstion에 위치한다는 연구에서 영감을 받았습니다. 본 연구는 model adaptation 동안의 가중치의 변화도 “low intrinsic rank”를 가지고 있다고 가정하여, 이로부터 ‘Low-Rank Adaptation(LoRA)’ 접근법을 제안합니다. LoRA를 사용하면, 사전 학습된 가중치를 그대로 유지하면서, adaptation 중 dense layer의 변화에 대한 랭크 분해 행렬(rank decompostion matrices)을 최적화함으로써 신경망 내의 일부 dense layer를 간접적으로 학습시킬 수 있습니다.

LoRA는 여러가지 주요 이점을 갖고 있습니다.

• pre-trained model은 다양한 task에 대해 공유할 수 있습니다. 공유된 모델을 고정시키고 그림 1의 행렬 A와 B를 각 task에 맞게 교체함으로써 task를 효과적으로 전환할 수 있어 storage requirement와 task-switching overhead를 크게 줄일 수 있습니다.

• LoRA는 대부분의 파라미터에 대한 그래디언트를 계산하거나 옵티마이저 상태를 유지할 필요가 없기 때문에 adaptive optimizers를 사용할 때 학습을 더 효율적으로 만들고 하드웨어 집입 장벽을 최대 3배 낮출 수 있습니다. 대신, 주입된 훨씬 작은 low-rank matrices만을 최적화합니다.

• LoRA의 간단한 선형 설계는 배포할 때 가중치를 동결한 상태로 학습 가능한 행렬과 병합할 수 있어, 구성적으로 완전히 fine-tuning된 모델과 비교하여 inference latency를 도입하지 않습니다.

• LoRA는 이전의 다양한 방법과 별개로 적용 가능하며, 예를 들어 prefix-tuning과 같은 다양한 방법과 결합할 수 있습니다.

Terminologies and Conventions

더욱 자세한 설명 이전에 간단하게 용어를 정리하고 가겠습니다.

본 논문은 Transformer 아키텍처를 참조하며, 해당 아키텍처의 차원에 대한 전통적인 용어를 사용합니다.

Transformer 레이어의 입력 및 출력 차원 크기를 $ d_{model} $ 로 사용합니다.
Self-attention 모듈에서 query/key/value/output projection matrices를 나타내기 위해 각각 $ W_q, W_k, W_v, W_o $ 를 사용합니다.
$ W $ 또는 $ W_0 $는 pre-trained weight matrices를 가리킵니다.
$ ∆W $는 fine-tuning 동안 누적된 그래디언트 업데이트를 나타냅니다.
LoRA 모듈의 랭크를 $r$로 표현합니다.
model optimizer로 Adam을 사용하고, Transformer MLP feedforward 차원은 $ d_{model} $의 4배인 $ d_{f fn} = 4 \times d_{model} $을 사용합니다.

Problem Statement

$\Phi$로 매개변수화된 사전 학습된 autoregressive language model $P_{\Phi}(y\vert x)$ 가 주어진다고 가정해봅시다. 이 pretrained model을 downstream conditional text generation task에 적응시키는 경우를 생각해보겠습니다. 이러한 downstream task에는 요약(summarization), 기계 독해(MRC),자연어에서 SQL로의 변환(NL2SQL) 등이 포함됩니다.

각 downstream task는 context-target 쌍의 학습 데이터 집합 $ Z = {(x_i, y_i)}_{i=1,…,N} $로 표현됩니다. 여기서 $ x_i $와 $ y_i $는 모두 token sequence입니다. 예를 들어, NL2SQL에서 $ x_i $는 자연어 $z$ 쿼리이고 $ y_i $는 해당하는 SQL 명령문입니다. 요약 task의 경우, $ x_i $는 기사의 내용이고 $ y_i $는 해당 기사의 요약일 것입니다.

전체 fine-tuning 과정에서 모델은 사전 학습된 가중치 $\Phi_0$로 초기화되고 조건부 언어 모델링 목적을 최대화하기 위해 반복적으로 기울기를 따라가면서 가중치를 $ \Phi_0 + ∆ \Phi $로 업데이트합니다.

전체 fine-tuning의 주요 단점 중 하나는 각 downstream task마다 전체 파라미터 차원인 $ ∆ \Phi $를 학습한다는 것입니다. 따라서 pre-trained 모델이 큰 경우, 많은 독립적인 fine-tuning 모델을 저장하고 배포하는 것은 매우 어려운 일입니다.

이 논문에서는 parameter-efficient approach 방식을 채택합니다. task-specific 파라미터 $∆ \phi (\theta) $는 훨씬작은 크기의 매개변수 집합으로 더 적절하게 인코딩되었습니다. 따라서 $ ∆ \Phi $를 찾는 작업은 $ \theta$에 대한 최적화를 의미합니다.

pre-trained model이 GPT-3 175B인 경우, 학습 가능한 파라미터 수 $ \vert \theta \vert $는 $ \vert \Phi_0 \vert $ 의 0.01% 정도로 작습니다.

Aren’t Existing Solutions Good Enough?

이전 연구들도 model adaptation을 보다 파라미터 및 계산 효율적으로 만들려고 노력했습니다.

Adapter Layers Introduce Inference Latency

Adapter 방법은 Transformer 블록에 adapter 레이어를 추가하는 방법입니다. 이 방법은 레이어를 제거하거나 멀티태스크 설정을 활용하여 latency를 줄일 수 있지만, adapter 레이어에서 발생하는 추가 계산을 직접 우회할 방법은 없습니다. adapter 레이어가 작은 병목 차원을 가져 파라미터가 적게 설계되었기 때문에 문제가 없는 것처럼 보입니다. 그러나 대규모 신경망은 low latency를 유지하기 위해 하드웨어 병렬 처리에 의존하며, adapter 레이어는 순차적으로 처리해야 합니다. 표 1에서 확인할 수 있듯이, 단일 GPU에서 GPT-2 medium에서 inference하는 경우와 같은 모델 병렬 처리(model parallelism)가 없는 일반적인 시나리오에서 adapter를 사용하면 매우 작은 병목 차원이더라도 latency가 늘어나는 것을 관찰합니다.

Directly Optimizing the Prompt is Hard

Prefix tuning 방법은 학습 가능한 파라미터에 대해 성능이 비단조적으로 변하기 때문에 최적화하기 어렵다는 단점을 가지고 있습니다. 또한 필연적으로 downstream task를 처리할 수 있는 시퀀스 길이를 줄이게 되며, 이로 인해 다른 방법에 비해 성능이 떨어집니다.

Our Method

LoRA의 간단한 디자인과 실제적인 이점을 설명합니다. LoRA는 딥 러닝 모델의 dense layer에 적용 가능하지만, Transformer language model에 중점을 두고 연구를 진행했습니다.

Low-Rank-Parameterized Update Matrices

특정 task에 adaptation할 때 Aghajanyan et al.(2020)은 pretrained language model이 “low intrinsic dimension”을 가지며, 작은 부분 공간으로의 임의의 투영에도 효율적으로 학습할 수 있음을 보여주었습니다.

이를 영감으로 받아, 가중치 업데이트도 adaptation 동안에 “low intrinsic rank”를 가질 것이라고 가설을 세웠습니다. 사전 학습된 가중치 행렬 $ W_0 \in \mathbb R^{d \times k} $ 의 경우, 이를 low-rank decomposition인 $ W_0 + ∆W = W_0 + BA $ 로 대체합니다. 여기서 $ B \in \mathbb R^{d \times r}, A \in \mathbb R^{r \times k} $, 그리고 랭크 $ r $은 $ min(d, k) $보다 훨씬 작습니다. 학습 중에 $W_0$는 동결되어 기울기 업데이트를 받지 않으며, A와 B의 파라미터만 학습합니다.

따라서 $ h = W_0 x $의 경우, 수정된 순전파 과정은 다음과 같습니다.

그림 1에서 확인할 수 있듯이, A에 대해서는 random Gaussian initialization을 사용하고 B에 대해서는 zero initialization을 사용하여, 학습 시작 시 $ ∆W = BA $는 0입니다. 그 후, $ ∆Wx $를 $ α \over r $로 스케일링합니다. 여기서 $ α $는 $r$의 상수입니다. 이 스케일링은 $r$을 변화시킬 때 하이퍼파라미터를 재조정하는 필요를 줄이는 데 도움이 됩니다.

A Generalization of Full Fine-tuning

LoRA를 모든 가중치 행렬에 적용하고 모든 bias를 학습할 때, LoRA 랭크 r을 사전 학습된 가중치 행렬의 랭크로 설정하면, full fine-tuning의 표현력을 대략적으로 회복합니다. 다시 말해, 학습 가능한 파라미터 수를 늘릴수록 LoRA 학습은 원래 모델을 학습하는 것으로 대략 수렴한다는 것을 의미합니다. 이에 반해 adapter 방법은 MLP로, prefix tuning 방법은 긴 입력 시퀀스를 처리할 수 없는 모델로 수렴합니다.

No additional Inference Latency

production 환경에서 배포할 때, 명시적으로 $ W = W_0 + BA $ 를 계산하고 저장한 다음 inference를 수행할 수 있습니다. 이 때 $ W_0 $와 $ BA $ 모두 $ \mathbb R^{d \times k} $ 입니다. 다른 downstream task로 전환해야 할 때 $ BA $를 빼서 $W_0$을 복구할 수 있고, 그 후 다른 $ B’A’ $를 더하면, 매우 적은 메모리로 빠르게 전환할 수 있습니다. 이로써 inference 시 fine-tuned 모델과 비교하여 추가적인 latency를 도입하지 않습니다.

Applying LoRA To Transformer

신경망 내의 가중치 행렬의 하위 집합에 LoRA를 적용하여 학습 가능한 파라미터의 수를 줄일 수 있습니다. Transformer 아키텍처에서는 self-attention 모듈에 네 개의 가중치 행렬($ W_q, W_k, W_v, W_o $)과 MLP 모듈에 두 개의 가중치 행렬이 있습니다. 일반적으로 출력 차원이 어텐션 헤드로 분할되지만, $ W_q, W_k, W_v $를 $ d_{model} \times d_{model} $ 차원의 단일 행렬로 취급합니다. 본 논문은 단순성과 파라미터 효율성을 위해 MLP 모듈의 가중치를 고정하고 attention 가중치를 downstream task에 맞게 조정하여 연구를 진행합니다. 또한 Transformer의 다양한 유형의 attention 가중치 행렬을 조정하는 연구를 추가로 진행했습니다.

Practical Benefits and Limitations

• 메모리 및 저장 공간 사용량의 감소

Adam으로 학습된 큰 Transformer의 경우, $ r $이 $ d_{model} $보다 훨씬 작기 때문에 동결된 파라미터 수의 옵티마이저 상태를 저장할 필요가 없으므로 VRAM 사용량을 최대 2/3까지 줄일 수 있습니다. GPT-3 175B의 경우, 학습 동안의 VRAM 소비량을 1.2TB에서 350GB로 줄일 수 있습니다. r=4로 설정하고 query와 value projection 행렬만 조정하는 경우, 체크포인트 크기가 대략 10,000배 감소하여(350GB -> 35MB) GPU 수를 크게 줄이고 I/O 병목 현상을 피할 수 있습니다.

• 낮은 비용으로 task 간 전환

배포 중에 LoRA 가중치만 교체하면 훨씬 낮은 비용으로 task 간에 전환이 가능합니다. 이로써 VRAM에 pre-trained weight를 저장하는 기계에서 동적으로 많은 customized model을 생성하고 교체할 수 있습니다. 또한, 대부분의 파라미터에 대한 기울기를 계산할 필요가 없기 때문에 GPT-3 175B에서 학습 중에 full fine-tuning과 비교하여 25%의 속도 향상을 얻습니다.

• Limitation

추가적인 inference latency를 제거하기 위해 A와 B를 W에 흡수하려는 경우, 서로 다른 A와 B를 가진 서로 다른 task의 입력을 하나의 forward pass로 처리하는 것은 간단하지 않습니다. 그러나 latency가 중요하지 않은 상황에서는 가중치를 병합하지 않고 배치의 각 샘플에 대해 동적으로 사용할 LoRA 모듈을 선택하는 것이 가능합니다.

Empirical Experiments

LoRA의 downstream task 성능을 RoBERTa, DeBERTa 및 GPT-2에서 평가한 후 GPT-3 175B로 확장합니다. 본 논문의 실험은 자연어 처리(NLU)부터 생성(NLG)까지 다양한 task를 다룹니다. RoBERTa와 DeBERTa의 경우 GLUE 벤치마크를 baseline으로 평가합니다. GPT-2의 경우 직접적인 비교를 위해 Li&Liang의 설정을 따르고 GPT-3에서 대규모 실험을 위해 WikiSQL(NL to SQL queries) 및 SAMsum(대화 요약)을 추가합니다.

Baselines

• Fine-Tuning(FT)
  Fine-tuning 중에 모델은 사전 학습된 가중치와 편향으로 초기화되고 모든 모델 파라미터는 기울기 업데이트를 받습니다. 레이어 중 일부만 업데이트하고 다른 레이어를 동결하는 간단한 변형도 있습니다. GPT-2의 경우 이전 연구에서 보고된 FTTop2라는 baseline을 포함시켰는데, 이는 마지막 두 레이어만을 조정합니다.

• Bias-only or BitFit
  Bias-only 또는 BitFit은 다른 모든 것을 동결시키고 bias 벡터만 학습시키는 baseline입니다.

• Prefix-embedding tuning(PreEmbed)
  입력 토큰들 사이에 특수 토큰을 삽입하는 방법입니다. 이러한 특수 토큰들은 학습 가능한 워드 임베딩을 가지고 있으며 일반적으로 모델의 어휘에는 포함되지 않습니다. “prefixing”과 “infixing”에 중점을 두고, “prefixing”은 이러한 특수 토큰을 프롬프트 앞에 덧붙이고, “infixing”은 프롬프트에 덧붙입니다. $ l_p $(접두사 토큰 수)와 $ l_i $(중간 삽입 토큰 수)를 사용하여, 학습 가능한 파라미터의 수를 $ |Θ| = d_{model} × (l_p + l_i) $로 표현합니다.

• Prefix-layer tuning(PreLayer)
  Prefix-layer tuning(PreLayer)은 prefix-embedding tuning의 확장입니다. 특수 토큰에 대한 워드 임베딩(또는 embedding 레이어 이후의 activation)만 학습하는 대신, 각 Transformer 레이어 이후의 activation을 학습합니다. 이전 레이어에서 계산된 activation은 학습 가능한 activation으로 간단히 대체됩니다. 결과적으로 학습 가능한 파라미터의 수는 $ |Θ| = L × d_{model} × (l_p + l_i) $이며, 여기서 L은 Transformer 레이어의 수입니다.

• Adpater tuning
  Houlsby et al.(2019)에서 제안한 Adapter tuning은 self-attention 모듈(및 MLP 모듈)과 다음 residual connection 사이에 adapter layer를 삽입합니다. adapter layer에는 중간에 비선형성의 bias가 포함된 2개의 fully connected layer가 있습니다. 이를 $ Adapter^H $라고 합니다. 최근에 Lin et al.(2020)는 더 효율적인 디자인을 제안했는데, adapter layer는 MLP 모듈 이후 및 LayerNorm 이후에만 적용됩니다. 이를 $ Adapter^L $이라고 합니다. 이는 Pfeiffer et al.(2021)에서 제안된 또 다른 디자인과 매우 유사한데, 이를 $ Adapter^P $라고 합니다. 또한 효율성을 높이기 위해 adapter layer 일부를 삭제하는 또 다른 baseline인 $ Adapter^D $를 포함시킵니다.(Ruckl “ e et al., 2020) 모든 경우에 $ |Θ| = \hat L_{Adpt} \times (2 \times d_{model} \times r+r+d_{model})+ 2 \times \hat L_{LN} \times d_{model} $이며, 여기서 $ \hat L_{Adpt} $는 adapter 레이어의 수이고 $ \hat L_{LN} $은 학습 가능한 LayerNorm의 수입니다.($Adapter^L$의 경우)

• LoRA
  LoRA는 기존의 가중치 행렬에 병렬로 학습 가능한 랭크 분해 행렬 쌍(pairs of rank decomposition matrices)을 추가합니다. 단순화를 위해 대부분의 실험에서 $ W_q $와 $ W_v $에만 LoRA를 적용합니다. 학습 가능한 파라미터의 수는 rank $ r $과 원래 가중치의 형태에 따라 결정됩니다. 따라서 $ |Θ| = 2 × \hat L_{LoRA} × d_{model} × r $이며, 여기서 $ \hat L_{LoRA} $는 LoRA를 적용하는 가중치 행렬의 수를 나타냅니다.

Experiments

Understanding The Low-Rank Updates

downstream task에서 학습한 low-rank adaptation의 특성에 대해 설명하고자 업데이트 가중치가 사전 학습 가중치와 어떤 상관이 있는지에 대한 연구를 진행했습니다. GPT-3 175B에 중점을 두고 연구를 진행하며, 이로써 학습 가능한 파라미터를 크게 줄이면서(최대 10,000배까지) downstream task의 성능에 불이익을 미치지 않는 결과를 얻었습니다.

다음과 같은 질문에 답하기 위해 연구를 수행했습니다.

1) 주어진 파라미터 예산 제약 조건 하에서 사전 학습된 Transformer의 어떤 가중치 행렬 하위 집합을 adaptation해야 downstream task 성능을 극대화할 수 있을까요?
2) “Optimal” adaptation matrix ∆W가 실제로 low-rank인가요?
3) ∆W와 W 간의 어떤 연결이 있나요?

Which Weight Matrices In Transformer Should We Apply LoRA To?

앞서 언급했듯이, 단순성을 위해 self-attention 모듈의 가중치 행렬만을 고려합니다. GPT-3 175B에서 대략 35MB에 해당하는 18M의 파라미터 예산을 설정하며, 하나의 유형의 어텐션 가중치를 적응할 경우 r=8, 두 가지 유형의 어텐션 가중치를 적응할 경우 r=4로 설정합니다.

$∆W_q$ 또는 $∆W_k$에 모든 매개변수를 넣은 것은 성능이 크게 저하되는 반면, $W_q$와 $W_v$를 모두 적응하는 것이 최상의 결과를 얻었습니다. 이것은 랭크가 4인 경우에도 ∆W에서 충분한 정보를 포착하므로, 더 많은 가중치 행렬을 적응하는 것이 단일 유형의 가중치를 큰 랭크로 적응하는 것보다 선호된다는 것을 보여줍니다.

What Is The Optimal Rank r For LoRA?

모델 성능에 대한 rank r의 영향에 주목합니다. 비교를 위해 $ {W_q, W_v}, {W_q, W_k, W_v, W_c} $, 그리고 단순히 $W_q$를 적응시킵니다.

표 6은 놀랍게도 LoRA가 매우 작은 r에 대해 이미 경쟁력 있는 성과를 내고 있다는 것을 보여줍니다. 이는 업데이트 행렬 $ ∆W $가 매우 작은 “intrinsic rank”를 가질 수 있다는 가능성을 시사합니다. r을 늘린다고 해서 성능이 향상되지 않으며, 이는 low-rank adaptation matrix가 충분하다는 것을 보여줍니다.

How Does The Adaptation Matrix ∆W Compare To W?

∆W와 W 간의 관계를 더 연구했습니다. W를 ∆W의 r차원 부분 공간에 투영하려고 합니다. 이를 위해 $ U^\top W V^ \top $를 계산합니다. 여기서 $ U/V $는 ∆W의 왼쪽/오른쪽 특이 벡터 행렬입니다. 그런 다음 $ \Vert U^\top W V^\top \Vert_F $ 와 $ \Vert W \Vert_F $ 간의 Frobenius norm을 비교합니다.

표 7로부터 여러가지 결론을 얻습니다. 첫째, ∆W는 random matrix와 비교했을 때 W와 더 강한 상관 관계를 가지며, ∆W는 이미 W에 있는 어떤 특성을 강화시킨다는 것을 나타냅니다. 둘째, W의 상위 특이 방향을 반복하는 대신, ∆W는 이미 W에서 강조되지 않은 방향만을 강화합니다. 셋째, 이러한 증폭 요인은 상당히 큽니다. 이는 low-rank adaptation matrix가 일반 pre-training model에서 학습되었지만 강조되지 않았던 구체적인 downstream task에 중요한 기능을 증폭할 수 있다는 가능성을 나타냅니다.

Conclusion And Future Work

LoRA를 통해 추가적인 latency를 도입하지 않고 입력 시퀀스 길이를 줄이지 않으면서 높은 모델 품질을 유지할 수 있습니다. 많은 모델 파라미터를 공유함으로써 서비스로 배포될 때 빠른 작업 전환을 가능하게 하여 막대한 하드웨어 비용을 줄일 수 있습니다. 본 논문에서는 Transformer language model에 중점을 뒀지만, dense layer를 가진 모든 신경망에 일반적으로 적용될 수 있습니다.