LoRA(Low-Rank Adaptation)를 파악해보자아앗!!

LoRA(Low-Rank Adaptation)

LoRA(Low-Rank Adaptation)는 Large Language Model(LLM)이나 신경망을 효율적으로 미세 조정(Fine-tuning) 하기 위한 방법 중 하나이다.

LoRA를 한 줄로 요약하면, Pre-trained 모델 가중치를 동결시키고 Transformer 아키텍처의 각 층에 학습 가능한 rank decomposition matrices을 넣는 방법이다.

We propose Low-Rank Adaptation, or LoRA, which freezes the pretrained model weights and injects trainable rank decomposition matrices into each layer of the Transformer architecture, greatly reducing the number of trainable parameters for downstream tasks

 

이제 풀어서 구체적으로 들어가 보자.

 

LoRA의 배경

Large Language Model의 시대가 도래하면서 이 LLM을 개발은 충분한 리소스가 있는 대기업들에서만 다룰 수 있었다.

자신의 task에 맞게 fine-tuning을 하고 싶더라도 pre-trained 모델의 파라미터 수만큼 업데이트시켜야 하고, 몇 달에 거쳐 fine-tuning을 모두 시킨다고 해도 배포를 하는 데에도 리소스에 한계가 있다. GPT-3 175B로 예를 들면, 각각 1750억 개의 파라미터를 가진 독립된 fine-tuned 인스턴스를 배포해야 하는데 이를 위해 정말 많은 자원이 필요하다

 

그러다 “단순히 정말 내가 하고자 하는 Task(예를 들어 summary task)에 이렇게 많은 파라미터가 필요할까?” 라는 의문이 제기되었다. 이전 연구에서 over-parametrized model은 실제로 low intrinsic dimension에 핵심이 존재한다는 것에서 영감 받아 이 논문에서는 모델이 새롭게 학습하는 과정에서 변경되는 가중치들도 사실 low intrinsic rank에 있을 가능성이 높을 것이라는 가정을 세운다.

We take inspiration from Li et al. (2018a); Aghajanyan et al. (2020) which show that the learned over-parametrized models in fact reside on a low intrinsic dimension. We hypothesize that the change in weights during model adaptation also has a low “intrinsic rank”, leading to our proposed Low-Rank Adaptation (LoRA) approach

 

 

*Low Rank

선형대수학에서 Rank는 행렬에서 선형적으로 독립인 column의 최대 갯수를 뜻하는데 그것은 해당 행렬이 얼마나 많은 독립된 정보를 담고 있는지를 의미한다.

그래서 독립된 정보를 담고 있는 것을 Rank라고 하고, 구체적으로 모델이 많은 수의 매개변수를 가지고 있어 고차원 공간에 위치하더라도 모델을 새로운 작업에 맞게 조정하는 데 필요한 실제 변화는 훨씬 더 작고 저차원적인 부분 공간(lower-dimensional subspace)에서 일어날 수 있다는 것이다.

 

 

LoRA의 구조

 

LoRA의 구조를 보면 핵심은 오른쪽의 주황색 블럭이다. Pre-trained weights(왼쪽 파란색 블록) 옆에 붙어서 인코더 디코더 구조같이 r 만큼 줄어들었다가 다시 커지는 형태를 갖고 있다. 결국 Low Rank는 dxd 크기의 W 전체를 학습시키기보다는 더 작은 행렬을 두 개를 사용하여 파라미터를 줄이면서 학습시킨다.

이때 Pre-trained weight를 froze 한 상태(왼쪽 파란색 블록)에서 pre-trained weights의 정보를 간접적으로 Low-rank를 통해 훈련시킬 수 있다는 것이 메인 아이디어이다.

LoRA allows us to train some dense layers in a neural network indirectly by optimizing rank decomposition matrices of the dense layers’ change during adaptation instead, while keeping the pre-trained weights frozen, as shown in Figure 1.

 

관련 실험으로는 GPT3의 12,888개의 원본 정보를 LoRA의 r을 1, 2로도 표현이 가능하다고 한다.

그러니까 12,888 * 12,888에서 12,888 * 2 + 2 * 12,888 이렇게 바뀌는 것이기 때문에 파라미터가 엄청 줄어든다.

Using GPT-3 175B as an example, we show that a very low rank (i.e., r in Figure 1 can be one or two) suffices even when the full rank (i.e., d) is as high as 12,288, making LoRA both storage- and compute-efficient

 

언어 모델을 full fine-tuning 할 경우

일반적으로 Language model을 full fine-tuning 하고자 하는 경우 아래와 같은 확률 모델을 따른다.

 

full fine-tuning (from pre-trained model)

 

Pre-trained autoregressive 언어 모델 PΦ(y|x)가 Φ를 파라미터로 갖고 있을 때, 예를 들어 우리가 fine-tuning 하기 원하는 task가 summary라고 한다면 x는 기사 내용, y는 요약본으로 Z는 (x, y) 쌍으로 데이터가 구성될 수 있다. 그리고 그 데이터들은 모두 token sequense로 이루어져 있을 것이다.

 

y는 t 보다 작고(t는 시간을 의미) x가 주어질 때 y_t의 확률을 maximum 시키는 확률 모델이다. 쉽게 이야기하면 y가 생성되는 시점보다 이전의 데이터를 가지고 x가 주어졌을 때 y를 예측하는 것으로 이해하면 된다.

 

full fine-tuning 과정에서는 모델이 pre-trained 가중치 Φ0로 초기화되고, 위의 조건부 언어 모델링 식을 최대화하기 위해 gradient를 반복적으로 따라가면서 Φ0 + ∆Φ로 업데이트한다. 여기서 단점은 각 task 마다 |Φ0|와 동일한 차원의 다른 매개변수 집합 ∆Φ를 학습해야 한다는 점이다 (GPT-3의 경우 |Φ0| ≈ 1750억 개인데 어떻게 학습시키고 서빙할 거야..)

 

During full fine-tuning, the model is initialized to pre-trained weights Φ0 and updated to Φ0 + ∆Φ by repeatedly following the gradient to maximize the conditional language modeling objective.

One of the main drawbacks for full fine-tuning is that for each downstream task, we learn a different set of parameters ∆Φ whose dimension |∆Φ| equals |Φ0|.

 

 

그럼 LoRA는 어떻게 다른 걸까?

LoRA approach

pΦ0는 pre-trained의 weight를 그대로 가져온 것이고 기존에 있는 weight에다 ∆Φ만큼의 정보를 추가해 주는 것이다.

LoRA에서는 fine-tuning 할 때 업데이트는 Θ만 하면 된다. 그러니 Θ를 maximize 하는 파라미터를 찾는 것이 목적이다. ∆Φ(Θ)은 훨씬 작은 크기의 파라미터 집합 Θ로 인코딩 되기 때문에 Φ0에 비하면 굉장히 작은 사이즈로 효율적으로 최적화시킬 수 있다.

In this paper, we adopt a more parameter-efficient approach, where the task-specific parameter increment ∆Φ = ∆Φ(Θ) is further encoded by a much smaller-sized set of parameters Θ with |Θ| |Φ0|.

In the subsequent sections, we propose to use a low-rank representation to encode ∆Φ that is both compute- and memory-efficient. When the pre-trained model is GPT-3 175B, the number of trainable parameters |Θ| can be as small as 0.01% of |Φ0|

 

Low-rank-parametrized update matrices

 

원래 full fine-tuning의 forward pass는 W0x 일 테지만 W 변화량을 추가해 위와 같은 식으로 표현할 수 있다.

위의 수식은 hidden layer로 가는 수식이고, W0는 Pre-trained model의 weight이다. W0는 frozen 시켜놓고 우리는 그 옆에 변화량 ∆W만 업데이트할 것이다.

W0 ∈ R^d×k에 대해, 업데이트를 Low Rank Decomposition로 표현하여 W0 + ∆W = W0 + BA와 같은 형태로 바꿀 수 있다. 여기서 B ∈ R^d×r, A ∈ R^r×k이며, rank r은 min(d, k) 보다 작거나 같은 범위에 있다.

훈련 동안 W0는 frozen 되어 gradiant 업데이트를 받지 않으며, A와 B 파라미터만 optimization 시킨다.

 

*LoRA에서 weight update가 일어나는 부분은 Adaptation이라고 하는데 이 weight metrix에 대해서 gradient update가 full rank를 가질 필요가 없어 메모리 요구량이 적다. 그래서 그 adaptation을 갈아 끼우는 것만으로 여러 개의 task를 수행할 수 있는 것이 LoRA의 특징이다.

 

 

Weight 초기화

A는 Gaussian 랜덤 초기화 시키고 B는 전부 0으로 초기화해서 ∆W = BA가 0이 되도록 한다.

처음엔 B가 0이니까 전부 0으로 초기화되어서 pre-trained weight만 FFN으로 넘어가고, 그리고 backpropagation 할 때 업데이트 될 테니 이렇게 한 것 같다.

 

 

Transformer에 LoRA 적용하기

 

Transformer 아키텍처에서는 셀프 어텐션 모듈에 네 개의 가중치 행렬(Wq, Wk, Wv, Wo)과 MLP 모듈에 두 개의 가중치 행렬이 있는데 일단 이 논문에서는 self-attention 모듈에만 LoRA를 적용해봤다고 한다. (대부분의 LLM 매개변수가 self-attention 층에 집중되어 있어 이 부분에 적용하는 것이 가장 큰 매개변수 절감 효과가 있다)

 

이를 통해 메모리와 저장 공간이 크게 감소하였고, 매개변수 r ≪ model_d로 설정하면 frozen 매개변수에 대해 optimizer 상태를 저장할 필요가 없기 때문에 VRAM 사용량을 최대 2/3까지 줄일 수 있었다고 한다. 또한, GPT-3 175B에서 full fine-tuning과 비교했을 때, 대부분의 매개변수에 대해 gradient를 계산할 필요가 없기 때문에 훈련 속도가 25% 향상되었다고 한다.

그러나 단일 forward pass에서 여러 개의 task에 대해 다른 A, B를 일괄로 한 번에 처리하는 것은 간단하지 않다.

 

it is not straightforward to batch inputs to different tasks with different A and B in a single forward pass, if one chooses to absorb A and B into W to eliminate additional inference latency.

 

 

LoRA를 사용하면 뭐가 좋길래?

1. 다양한 작업에 대해 여러 개의 작은 LoRA 모듈을 구축하는 데 사용 가능하다. 

pre-trained 모델을 공유하고 다양한 task에 대해 여러 개의 작은 LoRA 모듈을 구축할 수 있다. 공유된 모델을 동결하고, Matrix A와 B를 교체하는 방식으로 저장 공간과 작업을 효율적으로 전환할 수 있다. 예를 들어 llama 3 모델을 pre-train model로 공유하고, 특정 task에 맞는 adaptor(새로 학습된 A, B Matrix)만 바꿔서 사용할 수 있다.

 

2. LoRA는 대부분의 매개변수에 대해 gradient를 계산하거나 optimizer states를 유지할 필요가 없기 때문에, 삽입된 훨씬 작은 Low rank 행렬만을 최적화하여 학습을 더 효율적으로 만든다.

 

3. 배포 시에 학습 가능한 행렬을 동결된 가중치와 결합할 수 있게 하여, full fine-tuning 모델과 비교했을 때 더 효율적으로 사용이 가능하다.

 

LoRA로 fine-tuning을 마친 후

Pre-trained model과 LoRA adapter를 별도로 로드할 수도 있지만, 추가적인 latency issue가 발생할 수 있기 때문에 adapter 가중치를 pre-trained 모델과 병합하면 된다. 이렇게 하면 새로 병합된 모델을 독립형 모델로 효과적으로 사용할 수 있다.

 

 

Reference

https://www.youtube.com/watch?v=66GD0Bj5Whk&list=WL&index=1&t=603s

https://huggingface.co/docs/peft/main/en/conceptual_guides/lora

https://arxiv.org/pdf/2106.09685