[ICLR 2022] LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS

Abstract

안녕하세요, 강희승입니다. 오늘 리뷰해볼 논문은 LoRA입니다.

LoRA는 해당 분야에서 획기적으로 학습할 파라미터의 수를 줄였고, 일부 파라미터만 fine-tuning함에도 불구하고, Full fine-tuning의 성능에 비해 drop이 크지 않고 일부 task에 대해서는 오히려 성능의 향상을 이루었습니다.

5년이 지나면서 지속적으로 모델의 크기가 커짐에 따라 PEFT의 중요도가 꽤나 커진 것으로 생각됩니다. 여담으로 ICML에 다녀왔을 때, LoRA라는 키워드를 기반으로 연구되는 PEFT 기법들이 꽤 많이 보여서, 체감하게 된 것 같습니다.

Introduction

Transformer의 등장 이후, NLP;Natural Language Processing 분야의 연구 성장 속도는 기하 급수적으로 빨라졌습니다. 그에 따라, LLM의 모델들의 규모도 커졌습니다. 하지만, 커져가는 규모에 특정 downstream, domain으로의 full fine-tuning 하는 것과 해당 tuned model들을 모두 배포하는 것은 효율/용량 관점에서 굉장한 부담이었습니다.

이를 해결하기 위해 많은 연구들은 일부 파라미터만 학습하거나, 새로운 task에 대해 추가 module들을 학습하는 방식을 사용했습니다. 덕분에, 적은 task-specific parameters를 배포할 수 있게 되었습니다.

하지만, 기존의 방식들은 추가적 모듈로 인한 inference latency를 발생시키거나, 모델이 처리할 수 있는 token sequence의 길이를 제한한다는 단점이 있었습니다.

저자는 over-paramterized model이 사실 Low Intrinsic Dimension에 존재한다는 것을 보인 ‘Intrinsic dimensionality explains the effectiveness of language model fine-tuning’ 논문에서 영감을 받아, fine-tuning 과정에서 모델의 parameter가 instrinsic rank가 존재 하며 그에 따라 가중치가 업데이트 된다는 가설을 세웠습니다.

LoRA를 통해, GPT3-175B 모델에서 full-rank가 최대 12,288에 달함에도 불구하고, 매우 낮은 rank에서도 충분히 학습할 수 있음을 보였습니다.

Existing Solutions

LoRA가 등장하기 이전부터, PEFT; Parameter Efficient Fine-tuning 연구는 계속 되어왔습니다. 전이 학습이 처음 등장하고, 모델을 downstream에 adaption하는데, 더욱 파라미터/연산 효율적인 방식들에 대해 연구 되어 왔습니다.

기존의 PEFT 연구 방식은 크게 두가지 방향으로 나뉘어져 있었습니다.

1. Adapter를 적용하는 방식입니다. 대표적으로 transformer block 당 2개의 adapter를 각 attention과 FFN 뒤에 적용하는 방식이 있습니다.
   ‘Parameter-Efficient Transfer Learning for NLP’
2. Prefix-tuning 방식입니다. Input embedding의 앞부분에 learnable parameter(soft prompt)를 추가하는 방식과 다르게, 각 transformer layer 내부 attention key, value의 앞부분에 learnable parameter(prefix)를 추가하여 학습하는 방식입니다. 이때, pre-trained model은 freeze하며, prefix를 tuning하는 방식을 사용한다고 합니다.
   ‘Prefix-Tuning: Optimizing Continuous Prompts for Generation’

LoRA에서는 위 두가지 방식의 대표적 방식에 대해서 비교 실험을 진행했고, 해당 방식들이 가지는 문제점을 제시하였습니다.

Adapter 방식은 추가적인 layer로 구성되어있기 때문에, Inference latency를 발생시킵니다. Transformer block 내부에 추가된 Adapter는 연산 과정에서 우회할 수 있는 직접적 방식이 없기 때문에, Inference latency가 불가피하다고 저자는 언급합니다.

(구체적으로, 대규모 attention 연산 이후, 소규모 adapter를 연산하기 위해 반드시 순차적으로 처리해야하기 때문에, 전체 모델의 1% 미만의 파라미터를 갖도록 구조적 설계를 했음에도 불구하고, inference latency가 나타난다고 합니다.)

Prefix tuning 방식은 최적화 하는 것이 어려우며, 성능이 일관되지 않는다는 단점을 언급하였습니다. 또한 구조적으로 모델이 처리할 수 있는 최대 sequence length를 제한한다는 단점이 있습니다.

이러한 문제를 해결하기 위해 LoRA가 제시 되었고, 저자가 언급하는 LoRA의 keypoint는 다음과 같습니다.

1. 사전 학습된 모델을 공유하여, 다양한 task를 위한 LoRA module을 구축할 수 있다는 점입니다. 이를 통해 pre-trained model 즉 shared model을 freeze한 상태에서 LoRA module만 교체하는 것만으로도 task switch가 가능합니다. 따라서 저장 공간과 task switch로 인한 overhead를 획기적으로 줄일 수 있습니다.
2. LoRA는 대부분의 parameter에 대해 gradient를 계산하거나 optimizer의 상태를 유지할 필요가 없어, adaptive optimizer를 사용할 때 학습 효율을 높일 수 있습니다.
3. 단순한 linear 구조를 채택하여, 배포 시 learnable matrix를 fixed matrix와 병합이 가능하며 구조적으로 full fine-tuning 모델에 비해 inference latency가 발생하지 않습니다.
4. LoRA는 기존 여러 방법들과 orthogonal하며, prefix-tuning과 같은 기존 기법들과 결합하여 사용이 가능합니다.

저자는 굉장히 단순한 구조를 채택하면서도 위와 같은 contribution을 달성했습니다. Method를 처음 봤을 때, 구조는 굉장히 단순하지만, 논리와 접근은 단순하지 않다는 생각이 들었습니다.

Method

Low-Rank-Parameterized Updata Matrices

대부분의 인공지능 신경망은 dense layer로 구성되어 있고 matrix multiplication을 수행합니다. 각 layer들이 가지는 행렬들은 full-rank를 가집니다. 이전 연구에서, 특정 task로 tuning할 때, 사전학습 모델이 낮은 intrinsic dimension을 가지며, 더 작은 subspace로 임의 투영하여 학습하더라도 더욱 효율적인 학습이 가능하다라는 것을 보였습니다. 이에 영감을 받은 저자는 가중치 또한 intrinsic rank를 가질 것이라는 가설을 세웠습니다.

사전 학습된 가중치 W_0 \in \mathbb{R}^{d \times k} 에 대해, tuning 되며 업데이트 되는 weight를 다음과 같이 decomposition하였습니다.

W_0 + \Delta W = W_0 + BA, B \in \mathbb{R}^{d \times r}, A \in \mathbb{R}^{r \times k}, r \ll min(d,k)

사전학습된 W_0의 경우 freeze하여 gradient를 흐르지 않도록 하고, 오롯이 BA만을 학습합니다. 또한, 각각은 동일한 input으로 곱해지며 output vector를 coordinate-wise로 더합니다.

B는 0으로 초기화하며, A는 파라미터가 \mathcal{N}(0,\sigma^2)의 분포를 가지도록 초기화 해줍니다. 따라서 결과적으로 BA는 0이 되며, full fine-tuning과 같이 초기 pre-trained model의 parameter에서 학습하는 효과를 냅게 됩니다.

No Additional Inference Latency

만약, W_0x + BAx 와 같은 방식으로 별도로 계산을 수행한다면, 여전히 inference latency가 존재합니다. 하지만, LoRA는 별도의 비선형 함수가 없는 선형 연산으로 이루어져 있기 때문에, 가중치 간 병합이 가능합니다. 즉, 결과적으로 (W_0 + \Delta W)x = W_0x+\Delta Wx 이 성립하게 됩니다. 따라서 연산 과정에서 추가적인 Inference Latency가 발생하지 않는다고 합니다.

Deployment 관점에서, 다른 downstream task로의 전환이 필요할 때, 기존 task에 해당하는 BA를 단순히 빼줌으로써 W_0을 복원할 수 있고, 새로운 task에 해당하는 B’A’ 를 더하여 전환할 수 있습니다. 이러한 빠른 matrix operation을 통해 아주 작은 memory overhead만으로도 다양한 task로 전환이 가능합니다.

A Generalization of Full Fine-tuning

위와 같은 설계로, tuning 시에 모델의 학습할 weight \Delta W를 full-rank 행렬로 학습할 필요가 없다고 설명합니다. 이때, BA가 표현 가능한 weight update의 최대 rank (r)의 크기를 충분히 키운다면 full fine-tuning이 가지는 weight의 공간 크기 만큼 가질 수 있고, 그 결과를 표현할 수 있습니다. 다만, 이는 표현의 가능성에 관한 것으로, full fine-tuning의 동일 해로 수렴을 보장하지는 않는다고 합니다.

Limitation

하지만, inference latency를 줄이기 위해 W에 A,B를 merge하게 되면, 서로 다른 task에 대한 batch input이 들어오게 된다면, 각 task에 대해 일괄적으로 처리하기 어렵습니다. 따라서, Inference Latency가 크게 중요하지 않다면, W와 A,B를 분리하여 각 task에 맞는 모듈을 유동적으로 선택하여 연산하는 방식도 가능하다고 언급하였습니다.

Experiments

Trainable Parameter, Performance

저자는 해당 table 들에 대해 깊게 해석하지는 않았습니다. 다만, table의 결과가 저자의 주장을 명확히 뒷받침하기 때문에 독자 입장에서 큰 의문 없이 쉽게 납득을 할 수 있었습니다. Trainable parameter를 비교했을 때, PEFT 방식 특성상 모두 FT에 비해 굉장히 적은 수의 parameter를 사용하는 것을 알 수 있었습니다. 이때, 적은 수의 파라미터를 사용했음에도 불구하고, FT의 성능을 보존하는 것을 확인할 수 있습니다.

Which Weight Matrices In Transformer Should We Apply LoRA To?

저자는 LoRA의 단순성과 parameter-efficiency를 위해 Transformer의 attention weight만 대상으로 제한하여 적용하였습니다. 다만, attention에만 적용하는 것이 최적이라는 의미가 아닌 실험적 setting 인 것이고, MLP, LayerNorm 등에 적용하는 것은 후속 연구로 남겼습니다.

해당 table은 동일한 파라미터에서 weight와 rank의 분배에 따른 결과입니다. 이를 통해, 개별 weight에 높은 rank로 설계하는 것보다, 여러 weight에 낮은 rank로 구성하여 학습하는 것이 더 효과적인 것을 알 수 있습니다. 이는 Attention의 각 weight의 역할과도 연계되어, 적은 rank로 구성되더라도, Q,K,V,O 모두 학습하여 downstream에 적응시키는 것이 더욱 효과적임을 보였습니다.

What Is The Optimal Rank For LoRA

또한 저자는 rank에 따른 성능 변화도 리포트 했습니다. 예상했던 경향과 달리 오히려, 낮은 rank일 경우 특히 WikiSQL에서는 rank가 1일 때 가장 높은 성능을 보이는 것을 보입니다. 이를 update matrix \Delta W가 굉장히 낮은 intrinsic rank를 가질수 있으며, rank를 키우더라도, 의미있는 subspace를 포함하지 않는다는 것으로 해석하였습니다. 따라서, 저자는 low-rank adaption matrix로도 충분하다고 주장하였습니다.

Conclusion

기존의 방식들처럼, 구조적인 접근이 아닌, parameter update의 대수적 관점에서 PEFT 연구를 진행했다는 점이 굉장히 인상이 깊었습니다.

Intrinsic dimension의 개념을 가져와, 별도의 구조적 조정 없이도, 효과적인 Fine-tuning을 가능케 했고, 기존 연구의 문제였던 Inference Latency와 Sequence 구조적 제약을 해결하였습니다. 또한, downstream task에 해당하는 무거운 tuned 모델들은 모두 배포할 필요 없이 Pre-trained Model과 여러 task에 대한 LoRA module만 배포하는 것으로 deployment의 부담을 획기적으로 줄였다는 점에서 지금까지도 LoRA를 많이 사용하게 된 주요 contribution이지 않을까 생각됩니다.

Future Works

저자는 아직 fine-tuning mechanism이 사전학습된 feature들을 downstream에 맞게 어떻게 변환하는지에 대해 명확하지 않다고 합니다. 이에 저자는 LoRA가 full fine-tuning보다 이에 대한 답변을 찾는데 더욱 용이할 것이라고 하였습니다.

또한, LoRA의 weight matrix 선택 방식이 heuristic한 방식에 의존할 수 밖에 없다는 점을 언급하였습니다. 이에 대해 더욱 원칙적인 방식이 있을까라는 의문을 남겼습니다.

---

이러한 저자의 의문점을 통해 지금까지도 LoRA의 연구가 지속되는 것 같습니다. 모델들의 scale은 점점 커지고, 이러한 대규모 모델들이 계속해서 나오고 있기 때문에 PEFT 연구 또한 멈추지 않을 것으로 생각됩니다!

이상으로 LoRA 리뷰를 마치겠습니다. 읽어주셔서 감사합니다.