[논문리뷰] LoRA+ : Efficient Low Rank Adaptation of Large Models

arXiv paper code
Soufiane Hayou, Nikhil Ghosh, Bin Yu
UC Berkeley
19 Feb 2024

들어가기에 앞서..

본 논문은 LoRA 논문의 확장 연구로서 기본적인 LoRA 개념은 다루지 않을 예정입니다. 혹시 LoRA를 아직 보지 않으신 분은 이 곳을 통해 먼저 개념을 익히고 들어오시기를 추천드립니다.

Abstract

본 논문에서는 LoRA에서 Matrix B, A에 같은 Learning rate를 사용하는 것은 효율적이 않다고 주장하였다. 그러면서 간단한 방법으로 learning rate를 다르게 설정하는 LoRA+를 제안하였다. 이는 같은 컴퓨팅 자원에서 LoRA 대비 1~2%pt의 성능 향상과 약 2배까지의 속도 향상이 가능하다고 주장한다.

1. Introduction

이전의 LoRA 논문의 경우, rank의 하한선만 언급할 뿐 학습 속도(lr)는 1e-4로 고정된 값을 사용했다고 지적하였다.

2. Related work

QLoRA는 LoRA를 weight를 최소 4bit까지 절감하는 Quantization 버전이며 단일 GPU에서 LLaMA-65B까지도 학습할 있다.
LoftQ는 양자화를 통한 LoRA 학습의 연산을 효율화하기 위해 제안 된 것으로 더 나은 초기값을 설정하는 방법이다.
VeRA는 $N(0, \sigma^2)$ 분포에서 샘플링되는 어댑터를 freezing하여 사용하고 더 적은 Trainable parameter를 갖는 LoRA이다.

하지만 위 연구들 모두에서 Learning rate에 대한 언급은 없었다.

Contributions

본 논문에서는 신경망의 scaling 이론을 통해 learning rate를 설정하는 가이드라인을 제세한다고 한다. 구체적으로, 신경망의 폭/깊이를 무한대로 가정하고 하이퍼 파라미터(learning rate, initialization variance 등)에 따른 극한의 동작이 어떻게 변화하는지 이해한다. 이후, 원하는 목표(예 : feature learning)를 달성하기 위해 하이퍼파라미터의 원칙적인 선택을 도출한다.

구체적으로, LoRA+는 위 그림과 같이 A모듈과 B모듈의 학습 속도(lr)를 다르게 설정하는 방법이다. 매우 간단한 방법이지만 이에 대한 성능은 다양한 언어 모델을 통해서 결과로서 검증이 되었다.

Setup and definitions

Scaling of Neural Networks
일반적으로, 모델 폭(embedding dim)이 커진다면 초기값 설정 방식과 학습을 조정하여 효율적인 학습을 해야하는것으로 알려져 있다. 예를들어, initilization weights의 분산은 모델 폭 'n'이 커지면 large pre-activations을 방지하고자 $\frac{1}{n}$으로 스케일링 되어야 한다(He init). 저자는 모델 폭이 무한할 때의 LoRA를 사용한 학습 역학에 대해 연구하였다.

3. An intuitive analysis of LoRA

LoRA의 행렬 A와 B는 Transpose 모양을 하고있으며(위 그림을 보면 이해에 도움됨) 두 행렬에 대해 learning rate를 다르게 설정해야하는지에 대해 질문을 남길 수 있다. 실제로 대부분의 SOTA 모델은 모델의 폭(embedding dim)이 너무 크다. 따라서 저자는 폭이 $\infty$이 될 때의 학습 역학을 연구하는 것이 합리적이라고 한다. 정확히 말하면 저자가 이끌어낸 질문은 다음과 같다.

How should we set the learning rates of LoRA modules A and B as model width n → ∞ ?

간단한 Linear model을 통해 무한한 폭의 fine-tuning 역학에 대해 분석하고 폭에 따라 '최적'의 학습 속도에 대해 추론한다.

3-1. LoRA with a toy model

위 수식은 입력크기 n을 받으면 출력크기 1로 출력하는 간단한 Linear model이며 이때 LoRA의 Rank r=1이다.

Initialization
일반적으로라면 다음과 같이 Gaussian initialization으로 weight를 초기화한다.

$$ a_i \sim N(0, \sigma_a^2), \quad b \sim N(0, \sigma_b^2) $$
다만 LoRA에서는 첫 iteration에서 $ ba^\top $의 값이 0이 되도록 하고 있다. 이때, $a$와 $b$ 모두 0으로 하게된다면 saddle point로서 학습이 일어나지 않는다. 따라서 $a$와 $b$ 중 하나만 0이 되게 해야한다. 만약 $a$ 값을 0으로 초기화하지 않는 것을 선택한다면, $a^\top x$가 모델 폭에 따라 explode하지 않도록 표준 초기화 방식(He Init, LeCun Init)을 이용해서 $\sigma_a^2 = \Theta (n^{-1})$를 따라야 한다. 이는 중심극한정리(Central Limit Theorem; CLT)에 의해 정리되는 개념이다. 만약 $b$를 0이 아닌 값으로 초기화 한다면, $\sigma_b^2 = \Theta(1)$을 따라야 한다.

결과적으로 다음 두 가지 방법을 선택할 수 있게된다.

• $\sigma_b^2 = 0, \quad \sigma_a^2 = \Theta(n^{-1})$
• $\sigma_b^2 = \Theta(1), \quad \sigma_a^2 = 0$

Learning rate
$$\frac{\partial L}{\partial b} = a^\top x (f(x) - y), \quad \frac{\partial L}{\partial a} = b (f(x) - y) x$$
위 수식은 $a$와 $b$ weight가 학습될 때 각각의 Gradient 값을 구하는 식을 나타낸다.($b$에 대한 값을 구할 때에는 $a$가, $a$에 대한 값을 구할 때에는 $b$가 수식에 등장함 -> 이해가 안간다면 backpropagation 방법에 대한 자료가 구글에 많으니 그것을 참고하길 바람)

위 수식의 우항인 세 term은 각각 $a$, $b$, $ab$의 업데이트에 의한 변화를 가리킨다.

n이 커짐에 따라 바랍직한 값은 $\Delta f_t = \Theta(1)$가 되는것이다. 즉, 모델이 확장할 때 업데이트가 영향을 받지 않는다는 것을 의미한다. 이러한 모델 확장에 영향을 받는 예로는 scaling에 영향을 받는 lazy training가 있다. 이는 n이 무한에 가까울 때, $\Theta(n^{-1/2})$ 에 따라 학습이 발생하지 않는다는 것을 의미한다. 다시 한 번 정리하자면 $\Delta f_t = \Theta(1)$ 조건은 업데이트가 모델 폭에 따라 폭발하지 않는다는 것을 의미하여 바람직한 속성이다.

$\Delta = \Theta(1)$을 만족하는 것은 세 가지 항 중 적어도 하나가 $\Theta(1)$인 것을 의미한다. 또한, $\delta_1^t = o(1)$인 경우에는 $n \rightarrow \infty$일 때, 모델이 $a$가 고정되어 있고 $b$만 학습되는 것과 같다. (반대의 경우에도 마찬가지)

결론적으로, n이 무한에 가까울 때(클 때) $a$의 불안정한 학습을 막기위해 학습률을 낮추면 $b$의 학습이 비효율적으로 될 수 밖에없다. 이를 쉽게 해결하기 위해 저자는 $a$와 $b$에 다른 learning rate를 적용하고 $b$의 learning rate를 좀 더 설정하는 것을 제안하였다.

실험 결과, LoRA와 같은 컴퓨팅 자원에서 1~2%의 성능 향상을 기록하였으며 학습 속도는 2배가량 높아졌다고 한다.

---

본 논문에서 제안하는 학습 트릭이 너무 간단하여 논문 또한 쉽기 편할 줄 알았는데 생각보다 분석한 내용이 많아서 의외였다. 다음에 관련 대회에 참가할 기회가 생기면 QLoRA와 DoRA, LoRA+의 이점을 합쳐 Fine-tuning을 해볼 계획이다.