2024. 5. 29. 14:37ㆍ카테고리 없음
Abstact
LoRA는 pre-trained 모델로 downstream task를 학습할 때 pre-trained weight 전체를 학습시키는 대신 각 transformer layer마다 rank decomposition matrix를 추가하여 학습시키는 기법으로 학습 파라미터 개수를 크게 줄여준다.
LoRA는 기존 fine-tuning 기법에 비해 학습 파라미터 개수가 적고 학습 처리량이 높음에도 불구하고 동등하거나 더 나은 성능을 보인다. 그 뿐만 아니라, 타 adaptation 기법과 다르게 추론 지연이 없다.
Introduction
기존에는 pre-trained 모델을 fine-tuning할 때 pre-trained weight를 모두 학습하는 비효율적인 방식을 사용했다.
효율을 높이기 위해 일부 파라미터만 학습하거나 external module를 학습시키는 방식을 적용했다.
하지만, 이는 모델의 최대 입력 길이가 줄어지거나 모델이 깊어져 추론이 지연되는 문제를 유발시킬 뿐만 아니라,
fine-tuning 모델에 비해 성능이 낮게 나와, 성능과 효율의 trade-off를 야기한다.
저자는 학습된 over-parameterized model은 (고유) 차원이 낮기 때문에, 해당 모델을 fine-tuning해도 (고유) 차원이 낮을 것으로 가정했으며, 해당 아이디어는 LoRA 기법을 창안하는데 영감을 주었다.
LoRA는 dense layer를 간접적으로 학습시키는 방식을 사용한다.
구체적으로 말하자면, dense layer의 pre-trained weight를 학습시키는 대신 해당 layer 변화(=학습)량의 rank decomposition matrix를 최적화하는 방식을 사용한다. 이는 여러 이점을 제공한다.
1. 오직 pre-trained 모델 instance 하나만으로 task 별 LoRA 모듈을 여러 개 만들 수 있으며, LoRA 모듈로 task 전환을 효율적으로 할 수 있다.
2. LoRA는 효율적인 학습을 가능케하고, GPU 진입 장벽을 낮춰준다.
3. pre-trained weight와 rank decomposition matrix를 합쳐 추론 지연을 방지할 수 있다.
4. LoRA는 타 adapation 기법과 독립적으로 작동하기 때문에 결합하여 활용될 수 있다.
Problem Statement
파라미터 $\Phi$로 구성된 pre-trainedautoregressive language model $P_\Phi(y|x)$가 주어졌을 때, 이를 downstream conditional text generation task에 adaptation 해보는 문제를 풀어보자.
각 downstream task의 학습 데이터셋이 $Z = \{(x_i, y_i)\}_{i=1, 2, \cdots, N}$일 때, 목표 함수는 다음과 같다.
$$\underset{\Phi} \max \sum_{(x, y) \in Z} \sum_{t=1}^{|y|} \log(P_\Phi(y_t|x, y_{<t}))$$
full fine-tuning은 pre-trained weight $\Phi_0$로 모델을 초기화한 후, conditional language modeling 목표 함수값을 극대화 하기 위해 경사 하강법으로 $\Delta \Phi$를 구해 $\Phi_0$를 $\Phi_0 + \Delta \Phi$로 업데이트한다.
이때, full fine-tuning은 각 downstream task마다 $\Phi$와 동일한 차원(=크기)를 가진 $\Delta \Phi$를 새로 구해야 한다는 단점이 존재한다.
그렇기 때문에, pre-trained 모델이 클 경우, 많은 fine-tuned 모델을 저장하고 배포하는 과정이 매우 어렵다.
이번 논문에서는 새로운 parameter-efficient 접근법을 제시했다.
직관적으로, task 별 $\Delta \Phi$를 보다 작은 양의 파라미터 $\Theta$로 인코딩하는 방식인 $\Delta \Phi(\Theta)$을 제시했다.
이는 $\Delta \Phi$를 찾는 문제를 $\Theta$를 찾는 문제로 전환시켜줘 목표 함수가 다음과 같이 수정된다.
$$\underset{\Theta} \max \sum_{(x, y) \in Z} \sum_{t=1}^{|y|} \log(P_{\Phi_0 + \Delta \Phi(\Theta)}(y_t|x, y_{<t}))$$
저차원 표현을 사용해 $\Delta \Phi$을 인코딩하여 compute-, memory-efficient를 챙길 수 있다.
Aren't existing solutions good enough?
paramter-, compute-efficient하게 모델 adaptation하기 위해 이전부터 여러 기법이 제안되었다.
그 중 대표적으로 두 가지 기법(1. 추가된 adapter layer 최적화 2. 입력 layer의 activations(=prefix) 최적화)이 존재한다.
하지만, adapter layer를 추가하면 추론 시간이 지연되는 문제가 발생한다.
prefix tuning은 최적화하기가 어렵고, 파라미터 개수 증가에 따라 성능이 단조롭게 증가하지 않는다. 그뿐만 아니라, 최대 입력 토큰 길이가 줄어드는 단점도 존재한다.
Our Method
Low-rank-parameterized update matrices
Aghajanyen은 pre-trained 언어 모델은 낮은 내제 차원을 가지고 있음을 보일 뿐만 아니라, 이를 보다 좁은 부분공간으로 무작위 사영을 하여도 효율적인 fine-tuning 학습이 가능함을 보였다.
이에 영감을 받아 저자는 weight의 업데이트도 낮은 내제 차원을 갖을거라 가정해 weight의 업데이트를 다음과 같이 low-rank decomposition matrix로 표현하는 제악을 두었다.
$$W_0 + \Delta W = W_0 + BA, \ \text{where } B \in \mathbb R^{d \times r}, A \in \mathbb R^{r \times k} \text{ and } r \ll \min(d, k)$$
위 그림과 같이, $A$와 $B$는 각각 가우시안 분포와 0으로 초기화한다. $\Delta W = BA=0$에서 학습이 사작되기 위함이다.
LoRA는 full fine-tuning의 일반화 형태이다.
fine-tuning의 일반화 형태는 일부 pre-trained weight를 학습하는 것이다.
한 걸음 더 나아가, LoRA는 weight의 업데이트의 rank가 무조건 full rank일 필요가 없다.
이는 모든 weight에 LoRA를 적용하고 각 weight의 업데이트 rank를 full rank로 하면 full fine-tuning과 동일함을 의미한다.
LoRA는 추론 지연 시간 없다.
모델을 배포할 때, $W = W_0 + BA$를 저장하고 계산하기 때문에 추론 지연 시간이 없다.
그리고, 다른, task로의 전환이 필요할 경우, $W - BA$로 $W_0$을 찾고 다른 $B^\prime A^\prime$를 더한다. 해당 과정은 계산 및 메모리 비용이 적게 든다.
Applying LoRA to Transformer
위 연구에서는 오직 attention layer에만 LoRA를 적용해 실험해볼 것이다.
LoRA의 가장 큰 이점은 메모리 및 저장 공간 사용량 감소다.
LoRA를 활용해 GPT-3 175B를 학습시켰더니, VRAM 사용량이 1.2TB에서 350GB로 줄어들었다. frozen 파라미터의 optimizer states를 저장할 필요가 없어졌기 때문이다.
그 뿐만 아니라, 체크포인트 크기도 10,000x 줄어들었다.
이는 "LoRA를 활용할 경우, 학습에 필요로 GPU 개수가 줄어들고 동시에 GPU의 I/O 병목현상을 피할 수 있음"을 시사한다.
또 다른 이점은 배포된 상태에서 다른 task로 전환하는 비용 감소다.
이는 "원하는 모델로 즉시 교체해줄 수 있음"을 의미힌다.
학습 당시 학습 속도가 25% 향상되었음을 확인했다. 모든 파라미터의 기울기를 계산할 필요가 없어졌기 때문이다.
물론 단점도 존재한다. 서로 다른 task를 가진 입력 A, B가 같은 입력 배치에 들어올 경우, 추론 지연 없이 일괄 처리할 수 없다.
pre-trained weight와 task 별 rank decomposition matrix를 동적으로 합쳐 일괄 처리할 수 있지만, 추론 지연은 불가피하다.
Understanding the Low-Rank updates
What is the Optimal rank $r$ for LoRA?
rank $r$에 따라 모델 성능이 어떻게 변하는지 실험해봤다.
그 결과, $r$이 작아도 충분히 downstream task dataset을 adapting할 수 있음을 확인했다.
이는 weight 업데이트인 $\Delta W$의 내제 차원이 매우 낮다는 것을 시사한다.
저자는 rank $r$을 키워도 의미 있는 부분 공간을 찾아주지 못한다고 생각한다. 그렇기 때문에 low-rank adapation matrix를 사용하는 것이 적절하다고 판단하는 것이다.
Subspace similarity between different $r$
다른 크기의 $r$로 학습시킨 $\Delta W$가 있을 때, column space가 얼마나 겹치는지 파악하고자 한다.
$A_{r=8}$과 $A_{r=64}$이 각각 $r=8$, $r=64$에서 학습된 adaptation matrix고, $U_{r=8}$과 $U_{r=64}$은 SVD로 구한 right-singular unitary matrix다.
이때, $i$ singular vector가 span하는 공간과 $j$ singular vector가 span하는 공간이 얼마나 겹치는지 파악하기 위해 다음과 같은 수식을 사용했다.
$$\phi(A_{r=8}, A_{r=64}, i, j) = {\lVert{U_{A_r=8}^i}^\top U_{A_r=64}^j\rVert_F^2 \over \min(i,j)} \in [0, 1]$$
계산 결과, $A_{r=8}$의 top singular vector의 반향만 $A_{r=64}$의 column space에 겹친다 (>0.5). 그 반대도 마찬가지다.
이는 top singualr vector만 유용하고 나머지는 대부분 무작위 노이즈를 포함하고 있다. 그렇기 때문에, adaptation matrix는 매우 낮은 내제 차원을 가진다고 볼 수 있다.
Subspace similarity between different random seeds
random seed를 바꿔도 column space가 어떻게 변할지 파악하고자 한다.
$\phi(A_{r=64}, A_{r=64}^\prime, i, j)$로 계산해본 결과, $\Delta W_q$가 $\Delta W_v$보다 공통된 singular value direction을 더 많이 가지고 있었다.
이는 $\Delta W_q$가 $\Delta W_v$보다 더 높은 내제 차원을 갖고 있다고 볼 수 있다.
참고로, 다른 random seed로 gaussian matrix를 만들어 column space 유사도를 계산해본 결과, 공통된 부분 공간이 하나도 없었다.
How does the adaptation matrix $\Delta W$ compare to $W$?
나중에... 이해가 안되서...
Conclusion
거대한 언어 모델을 fine-tuning 하는 것은 너무 비쌀 뿐만 아니라, 다양한 task를 처리하기 위해 여러 모델 instances를 호스팅하기 위해 저장/전환하는 비용도 매우 비싸다.
LoRA는 매우 효율적인 adaptation 기법으로 기존 fine-tuning 기법과 거의 동일한 성능을 보임과 동시에 추론 시간 지연과 입력 토큰 개수 감소도 없다. 그리고, pre-trained 모델 instance 하나만으로 다양한 task를 처리하기 할 수 있기 때문에 저장/전환하는 비용도 매우 저렴해졌다.