[논문]LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS

논문

https://arxiv.org/abs/2106.09685

LoRA: Low-Rank Adaptation of Large Language Models

 

1. Introduction

  Language 모델을 활용하는 다양한 분야에서 특정 분야에 특화된 모델을 만들기 위해서 사전에 학습된 pre-trained model을 활용해서 fine-tuning을 진행하게 된다. 하지만 파라미터가 굉장히 많은 GPT 모델 같은 경우에는 이런 fine-tuning 작업에도 많은 리소스가 활용되고 있다. 이런 방식을 해결하기 위해서 많은 방식을 활용을 했었다. 하지만 이런 방식들도 모델의 효율성과 모델의 퀄리티에 대한 trade-off는 해결하지 못했다. 

 

Low-Rank Adaptation 방식은 "Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning"이라는 논문에서 등장했던 내용에서 영감을 받았다. 이 논문에서는 over-parametrized 모델을 더 작은 의미 있는 차원으로 조정하여서 모델을 변경하는 그런 아이디어를 제시하였다. 그래서 여기서 영감을 받아서 LoRA는 원래 네트워크보다 더 작은 dense 레이어를 학습을 함으로써 fine-tuning을 진행한다.

 

본 논문에서 제시하는 4가지 장점이 존대한다.

 

- fine-tuning 과정에서 여러 가지 LoRA 모듈을 만드는 게 가능하고 필요한 스토리지를 줄이는 게 가능하다.

- 훈련중 기존 방식보다 더욱 효과적이며  하드웨어에서 계산하는 그래디언트 같은 부분이 3배 정도의 줄어든다.

- 더 작은 차원의 모델을 학습하고 학습을 하지 않았던 기존 모델의 파라미터와 합치는 구조로 되어있기 때문에 추가적인 latency가 발생하지 않는다.

- 이전의 많은 방식들과 적용이 가능하다.

 

 Terminologies and Conventions

 

 

• $d_{model}$: Transformer 레이어의 입력 및 출력 차원 크기
• $W_q$, $W_k$, $W_k$, $W_v$, $W_o$: self-attention 모듈에서 query, key, value, output projection 행렬
• $W_0$, $W$: 사전 학습된 가중치 행렬
• $\Delta W$: Adaptation 중에 누적된 기울기 업데이트
• $r$: LoRA 모듈의 rank

2. Problem Statement

$P_\Phi (y|x) $ 는 $\Phi $를 기준으로 pre-trained 모델의 파라미터라고 한다면 이런 모델들은 multi-task로 GPT와 같이 적용될 수 있다. 그래서 이런 테스크들에 대한 데이터는 $z = \left\{(x_i,y_i) \right\}_{i=1,...,N}$ 로 표현하는 게 가능하다.

 

 

기존 모델

기존의 사전에 학습된 언어 모델의 weight를 $\Phi_0$ 로 표현하며 업데이트는 $\Phi_0 + \Delta \Phi$ 로 업데이트 한다.

$\Delta \Phi$  는 다음 수식을 통해서 구하게 된다.

해석

$\max_\Phi$  : 학습하는 파라미터를 나타낸다.

$\sum_{(x,y) \in \mathcal{Z}}$ : 학습하는 데이터 세트를 이야기한다.

$\sum_{t=1}^{|y|}$ : 학습하는 데이터를 기준으로 label sequence의 길이를 기준으로 토큰의 타임스텝을 지정한다.

$\log \left( P_{\Phi}(y_t \mid x, y_{<t}) \right)$ : 입력 시퀀스 $x$ 와 타겟 시퀀스의 토큰을 타임스텝 $t$(현재) 를 기준으로 $t$ 이전까지의 데이터를 활용하여 현재에 나올 토큰이 $y_t$ 를 예측하도록 하는 최댓값 $\Phi$ 가 되도록 하는 값을 구하는 것이 최종적인 목표이다.

 

수식을 통해서 알 수 있듯이 $\Delta \Phi$ 의 파라미터를 여러 가지 태스크에 대하여 학습을 진행해야 합니다. 하지만 GPT-3에서는 $|\Phi|$는 약 175 Billion 정도 되기 때문에 모델을 fine-tuning 한다는 건 굉장히 어려운 문제이다.

 

LoRA

기존의 모델의 fine-tuning에 대한 문제를 해결하기 위해서. 본논문에서는 $\Delta \Phi = \Delta \Phi(\Theta )$ 를 통해서 수행하게 되는데 여기서 중요한 LoRA의 방식은 $|\Theta| \ll |Phi_0|$ 인 $\Theta$ 를 학습하는 방식으로 진행된다.

여기서 주목해서 볼 부분은 $\max_\Theta$ 와 $p_{\Phi_0+\Delta \Phi(\Theta)}$ 이 부분이다. 이전에 방식과는 다르게 $\Theta$ 는 기존에 존재하는 모델들보다. 더 작은 차원의 weight. 를이야기하고 $p_{\Phi_0+\Delta \Phi(\Theta)}$  이 부분에서는 기존 모델에서의 pre-trained $\Phi_0$ weight를 이용하고 학습한 weight와 합쳐주는 방식이다. 

 

때문에 기존에 fine-tuning 방식보다 메모리 효율적이다. 그리고 $|\Theta|$ 는 기존에 weight의 $|\Phi_0|$ 의 0.01 % 정도이다.

 

3.  Aren't Existing Solutions Good Enough?

Adapter Layers Introduce Inference Latency

논문: Parameter-Efficient Transfer Learning for NLP

 

위의 논문에서 등장한 방식이지만 추가적인 Adapter layer를 추가해서 연상량을 줄이는 방식으로 적용을 하였지만 이런 방식은 추가적인 latency를 만들기 때문에 효과적이지 않다.

Directly Optimizing the Prompt is Hard

논문: Prefix-Tuning: Optimizing Continuous Prompts for Generation

 

프리픽스 튜닝은 모델의 입력 시퀀스 앞에 프리픽스를 추가하여 모델을 튜닝하는 방법이지만, 이 방법은 최적화가 어렵고 성능이 불규칙적으로 변화할 수 있습니다. 또한, 프리픽스를 위해 시퀀스 길이의 일부를 할애하면 실제 작업을 처리하는 데 사용되는 시퀀스 길이가 줄어들어 성능이 저하될 수 있습니다. 

 

4. OUR METHOD

4.1 LOW-RANK-PARAMETRIZED UPDATE MATRICES

 

neural network에서는 여러 가지의 dense layer가 존재한다. 그러는 과정에서 이런 모델들은 작은 차원으로 줄여 나가는 과정에서 모델에 대한 matrix는 의미 있는 작은 rank로 줄어든다는 가정을 하게 되었다. 그래서 이런 방식을 이용하여서 $W_0 \in \mathbb{R}^{d \times k}$ pre-trained 모델의 차원이 d, k 차원이라고 한다면 low-rank로 분해하여서 $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)$ 즉 원래 구성하는 pre-trained weight에서도 작은 차원을 기준으로 더 작은 차원으로 줄이게 됩니다.

 

여기서 중요한 점은 BA를 matrix 곱 연산을 하게 되면서 결국에는 pre-trained 모델의 차원과 같아지는 결과가 됩니다.

 

아래의 그림을 예시로 든다면 학습하지 않는 pretrained weights는 $d \times d$ 이라고 한다면 $A = r \times d , B = d \times r$ 일 때 BA를 matrix multiplication을 하게 된다면 결국에는  $d \times d$ 차원으로 맞춰지기 때문에 pretrained weights 차원이 같아지는 효과가 납니다.

 

 

 

본 논문에서는 A는 Gaussian initialization을 통해서 초기화를 하고 B는 전부 0으로 초기화를 합니다. 그래서 훈련 초기에는 BA는 zero가 됩니다.

 

그리고 $\Theta Wx$ 의 값을 $\alpha /r $을 기준으로 스케일링한다. 그래서 $\alpha$ 는 초기설정으로 두고 $r$의 변화에 따른 하이퍼파라미터 튜닝에 대한 재설정의 필요를 줄여준다. 이것의 효과는 $r$ 의 차원이 늘어난다고 하더라도 안정적인 학습 효과를 할 수 있는 이유입니다.

 

따라서 일반적인 Full Fine-tuning을 하는 것과 비교한다면 $r$ 즉 훈련 가능한 파라미터의 수를 증가시키면 기존의 모델을 학습하는 효과와 비슷해지며 prefix-based 방식의 긴 시퀀스를 처리하지 못하는 문제도 해결합니다.

 

그리고 이전에 설명했듯이 기존의 pre-trained 파라미터인 $W_0$ 에서 학습한 모델의 파라미터를 더하기 때문에 새로운 태스크를 적용한다면 기존의 파라미터를 빼주고 다른 $B'A'$ 를 더해주며 쉽게 다른 태스크 적용이 가능하다. 그래서 추가적인 latency가 발생하지 않습니다.

 

 

 

작동 예시

 

 

공식 코드에서 가지고온 LoRA의 코드입니다.

코드를 확인하면 LoRA의 차원을 설정하는 부분과 파라미터 초기화 하는 과정 등을 확인이 가능합니다.

 

여기서 확인이 가능한 부분은 self.scaling 값을 수식과 다르게 확인하는게 가능합니다. 

 

 

GPT-2 모델에서 LoRA를 적용한 코드를 보게 된다면 차원의 변화가 많은 Attention Layer의 Linear layer에 LoRA를 적요하는 모습을 알 수 있습니다.

4.2 APPLYING LORA TO TRANSFORMER

본 논문에서는 모델의 학습 과정에서 Transformer layer에서 MLP 구조의 모델의 가중치는 그대로 두고 q, k, v, o 값에 대한 어텐션 가중치만을 차원을 줄여서 효율성을 위해서 학습을 진행하였습니다. 

 

LoRA의 주요 이점은 메모리와 저장 공간 사용량을 크게 줄여 훈련 시 VRAM 사용량과 체크포인트 크기를 줄이고, 작업 전환 비용을 낮추며, 훈련 속도를 25% 향상하는 것입니다.

 

LoRA에도 한계가 있습니다. 서로 다른 태스크를 배치고 학습한다면 제대로 작동하지 않으며 다른 A, B 가중치로 선택한다면 추가적인 latency가 발생합니다. 

 

4.3 LoRA Advantages and Limitations

기존의 Fine-Tuning 방식보다 더 적은 파라미터를 자기고 downstream task 에 대하여 학습이 가능하고 더 좋은 성능을 낸다. 그리고 각각의 task마다 모듈 학습이 가능하며 기존의 pretrained 에 적용했던 파라미터를 더한걸 단순하게 빼고 새롭게 추가한 task에 대하여 추가하는 방식을 적용하여 쉽게 task별로 모델을 활용하는게 가능하다.

 

하지만 이런 방식에 대한 한계도 문제가 존재하는데 여기서는 각가 다른 task에 대한 데이터의 학습을 진행을 하는것에 어려움이 있으며 다른 task의 모델을 활용하기 위해서는 따로 따로 학습하는 시간이 필요로 한다. 

 

 

5. Evaluation

Transformer

 

transformer layer에 LoRA를 적용을 했을때 각 차원 rank 값이 줄어들어도 tine-tuning에서 학습 성능이 좋음을 알 수 있다.

 

 

이 그래프 에서도 기존의 Fine-Tuning 방식보다 적은 파라미터로 정확도 높은 성능을 내는 것을 보여준다.