21년 여름에 혜성처럼 등장해 최근 2년 동안 LLM 학습을 위해 사용되고 있는 LoRA 를 소개한다.
LoRA 를 소개할 때 수식에 대한 내용은 최대한 제거하고 원리는 무엇인지, 모델 구조에서 학습을 위해 LoRA 가 어떻게 활용되는지를 위주로 설명하겠다.
때문에 논문 리뷰같지 않은 리뷰라고 느껴질 수 있다.
수식에 대한 내용이 궁금한 독자는 다른 사람의 논문 리뷰를 참고하기 바란다.
Target 독자: Deep Learning 전문가
목차
1. Fully Fine-Tuning 이 힘든 이유 (간략)
2. LoRA 학습 방법
3. ( LoRA_B x LoRA_A ) 는 어디 Pretrain layer weight 에 더해지는가
4. 당연한 마무리 (이런게 좋아요~)
1. Fully Fine-Tuning 이 힘든 이유 (간략)
LLM 의 weight 는 최소 1.5~3B 이다.
Model 을 GPU 에 load 하는것만 해도 웬만큼 비싼 GPU 가 아닌 이상 불가능하다.
모델을 Fine-Tuning 학습하는 건 말할 것도 없다.
Forward & Backward, 이를 통한 Model weight update 는 gradient 를 전부 GPU 에 저장해야 된다.
gradient 뿐만 아니라 optimizer 를 위한 이전 기록(=tensor)들도 GPU 에 저장해야 된다.
결국 Fully Fine-Tuning 을 위해서는 모델의 weight 수 * 2~3배의 GPU vram 이 필요하다.
weight 의 수가 많은 LLM 을 Fully FIne-Tuning 하지 않는 이유가 이것이다.
그래서 탄생한 학습 방법이 바로 LoRA 이다.
2. LoRA 학습 방법
LoRA 의 사상은 다음과 같다.
- Fully FIne-Tuning 하지 않는다.
- Model weight 를 Freeze 한다.
- 학습하는 Layer 는 LoRA_A & LoRA_B 이다. (둘 다 nn.linear 형태)
- Transformer Layer 에 있는 Query, Key, Value, Output(=self attention) 중 선택하여 ( LoRA_B x LoRA_A ) 를 단순히 더해준다.
즉, Model weight 를 freeze 하지만 Inference 시 사용되는 weight 값은 update 가 된다.
Model weight 에 ( LoRA_B x LoRA_A ) 를 더해줬기 때문이다.
그림을 통해 쉽게 설명해보겠다.
- 왼쪽에 있는 파란색 box 는 Pretrained Model 의 Weight 이다.
- Transformer layer 의 Query, Key, Value, Output(=self attention) layer 의 차원은 (d x k) 이다.
- d 는 hidden_size, k 는 웬만하면 d 와 같다.
- Fully Freeze 되어 학습 시 weight update 가 되지 않는다.
- gradient, optimize 관련된 tensor 값이 gpu 에 load 되지 않는다.
- 즉, vram 을 save 할 수 있다.
- 오른쪽에 있는 주황색 A 가 LoRA_A layer 이다.
- LoRA_A 는 nn.linear 이다.
- {d x r} 의 차원을 가지고 있다.
- d 는 hidden_size , r 은 사용자가 설장한 낮은 차원(r << d)이다.
- 오른쪽에 있는 주황색 B 가 LoRA_B layer 이다.
- LoRA_B 는 nn.linear 이다.
- {r x k} 의 차원을 가지고 있다. (k 는 웬만하면 d 와 같다.)
- k 는 웬만하면 d 와 같고, r 은 사용자가 설장한 낮은 차원(r << d)이다.
그렇다면 ( LoRA_B x LoRA_A ) 에 의해 만들어진 행렬은 Pretrain Model 의 어디에 더해질까?
3. ( LoRA_B x LoRA_A ) 는 어디 Pretrain layer weight 에 더해지는가
논문에 있는 표를 통해 설명하겠다.
해당 표는 모델을 학습을 위해 18M 의 weight 만을 가지고 2개의 Downstream task 에 대한 LoRA 학습 결과를 보여준다.
18M 의 weight 는 LoRA_A, LoRA_B 에 의한 weight 수이다.
그렇다면 Weight Type 은 무엇일까?
Weight Type 이 바로 ( LoRA_B x LoRA_A ) 가 더해지는 Transformer layer 이다.
- W_q : query layer 를 뜻한다. (차원은 d x k 이다.)
- W_k : key layer 를 뜻한다. (차원은 d x k 이다.)
- W_v : value layer 를 뜻한다. (차원은 d x k 이다.)
- W_o : self-attention layer 를 뜻한다. (차원은 k x d 이다.)
모든 Weight Type 의 차원은 웬만하면 d x d 가 되기 때문에 weight update 가 되고 있는 ( LoRA_B x LoRA_A ) 행렬과 더할 수 있다. (d 와 k 의 차원은 웬만하면 동일하기 때문)
이를 통해 Model weight 를 직접적으로 update 하지는 않았지만 update 하는 효과를 볼 수 있다고 한다.
논문 저자는 ( LoRA_B x LoRA_A ) 행렬을 어디에 더해줘야할지 고민했다고 한다.
위의 표가 바로 고민을 통해 나온 결과이다.
어디 Weight Type 에 더해주는지에 따른 성능 비교를 한거다.
결과는 query 와 key layer 에 더해줬을 때 가장 좋은 성능을 보였다.
Rank r 이 의미하는건 LoRA_A, LoRA_B 행렬의 차원인 r 인 것 같다..
확실하지는 않아 더 알아볼 예정이다...
실제 코드를 통해 알아보니 r 은 LoRA_A, LoRA_B 행렬의 차원이 r 이 맞다.
4. 당연한 마무리 (이런게 좋아요~)
모든 논문의 끝은 거의 동일하다.
LoRA 를 사용하면 좋아요~ 이다.
무엇이 좋을까?
- Pretrain Model weight 를 update 하지 않고도 Fully Fine-Tuning 한 결과와 비슷하거나 더 좋은 성능을 보였다.
- weight update 에 필요한 weight 가 오직 LoRA_A, LoRA_B layer 에 있는 weight 뿐이기 때문에 vram 을 상당부분 save 할 수 있다.
- 학습된 모델로 inference 해도 Pretrain Model 로 inference 할 때와 동일한 연산량이다. (추론 시간이 거의 동일)
- Pretrain Model weight 를 Freeze 한 상태로 ( LoRA_b x LoRA_A ) 행렬을 단순히 더해주기 때문에 Pretrain Model weight 로 다시 원복하기 쉽다. (더해준 만큼 빼면 되기 때문)
마무리,,
필자가 근무하는 회사에서도 LLM 학습을 위해 LoRA 를 활용한다.
그동안 LoRA 는 Pretrain Model weight freeze 만 알고 있었기 때문에 학습 시 왜 성능이 좋아지는건지 의문이었다.
드디어 그 의문이 풀렸다 :)
다음 포스팅에서는 LoRA 학습을 위한 코드 리뷰를 해보겠다.
댓글