본문 바로가기

DEVELOP_NOTE/LLM

[PEFT] LoRA(Low-Rank Adaptation of Large Language Models) Fine-Tuning 학습 방식 뜯어보기! (내부 학습 코드 해석)

오늘은 최근 Fine-Tuning 시 가장 많이 사용하는 매서드중 하나인 LoRA에 대해서 정리해보도록 하자.

 

먼저 LoRA는 PEFT로 대표되는 Fine-Tuning 학습 방법이고, 

최근 LoRA에서 파생된, QLoRA, MoRA 등 매서드들이 소개되면서, Fine-Tuning의 주류로 자리잡고있다. 

LoRA를 비롯한 최근의 Fine-Tuning 방법들에 대해서 이해하기 위해서 모태가 된 PEFT 방식이

왜 각광받게 되었는지 먼저 알아볼 필요가 있다.

 

1. DonwStream Task에 Fine-Tuning 하기에 모델이 너무 커지고 있다...

몇년전부터, 대부분의 언어 모델이 규모의 경쟁으로 돌입했고,

현재도 여전히 규모를 통해 성능을 끌어올린다 라는 공식에는 변함이 없다.

물론 엄청난 규모의 자원과 데이터를 통한 학습이 SOTA 모델의 성능을 끌어올리고 있는것은 맞지만,

실제 다양한 도메인에서 요구하는 퀄리티를 만족시키기 위해서는 모델의 Downstream Task Fine-Tuning은 필수적이다.

 

하지만, 쓸만한 모델은 모두 최소 8B이상이 되어야하는 현실에서,

8B모델을 Load하는데만, float32기준으로 32GB, Full Fine-Tuning 진행 시, 보통 2~3배정도 소요되기때문에,

현재 2~3000만원 정도하는 A100 GPU 한장으로도 학습이 쉽지 않다.(Quantization은 일단 배제하고 얘기하자.)

 

2. Catastrophic Forgetting (파괴적 망각)

Fine-Tuning을 진행할때, weight update가 이루지면서, 어쩌면 기존 학습 내용중 일부를 망각하는것은 자연스러운 현상이다.

다만, 이것을 최소화 하기위해서는 일반화된 기존의 학습데이터와 새롭게 학습해야하는 Downstream task data가

균형있게 구성되어야하고, 잘 만들어진 데이터로 학습해야, 이 현상을 최소화 할 수 있는데,

데이터 생성 비용자체가 매우 많이 들기때문에, 쉽지 않다. 이때문에 Fine-Tuning 진행 시, 기존 성능을 잘 유지하느것이 

가장 큰 노하우 이자 관건이 된다.

 

 

이러한 이유들로 인해, LLM의 Fine-Tuning 과정에서  적은 리소스를 활용하면서도 성능 하락을 방지할 수 있는 학습 방법론에

대한 필요성이 대두되었다.

 


 

PEFT는 이 문제들을 어떻게 해결했나?...

최초에 model을 학습한다는것의 의미는 전체 모델 weight를 무작위로

update하며, 목표에 적합한 형태를 띄도록 변화해나가는것을 의미한다. 

반면, Fine-Tuning을 진행할때는, 최초 학습된 어느정도 최적화된 모델의 weight를 추가적으로 update,

미세 조정의 개념으로 학습을 진행하게 되는데, 

여기서, weight를 어떤식으로 최적화나 업데이트할지, 어떤 weight를 freezing할지는 선택한 fine-tuning방식에 따라 결정된다.

 

자, 그럼 PEFT는 어떤 학습 방법으로 이 문제들을 해결하고자 했을까?

[💡 핵심 IDEA]
"Pre-Trained Model 대부분의 parameter를 freezing하고 일부 prameter만 파인튜닝한다"

 

PEFT의 아이디어를 한마디로 표현하면 위와 같은데, 

너무나 당연한 설명이라, 잘 와닿지 않는다.

물론, 위와 같이 학습하게 된다면,

➡️ 일부 파라미터만 Update하기 때문에, 계산비용과 GPU 리소스는 줄어들 수 밖에없다.

➡️ 일부 파라미터만 Update하기 때문에, 기존 학습 weight 정보를 보존할 수 있다.

➡️ 추가로, 일부 파라미터에 대한 약간의 수정내용만 발생하기 때문에, 학습파일도 Mb단위로 저장이 가능하다.

 

설명한대로라면 간단히 문제들을 해결할 수 있어 보이는데,

구체적으로 어떤 방식을 통해 효과적으로 문제를 해결했는지 자세히 알아볼 필요가 있다.

 

오늘은 PEFT으로 대표되는 학습 방법론중에서도 LoRA를 통해 그 방법을 활용해보고자 한다.


LoRA는 어떻게 동작하는가?

* LoRA
arXiv 2021 | Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen at Microsoft Corporation
17 Jun 2021
    - Paper : https://arxiv.org/abs/2106.09685
    - Github : https://github.com/microsoft/LoRA

 

먼저 LoRA는 "Low-Rank Adaptation of Large Language Models"의 약자를 의미한다.

여기서 Low-Rank Adaptation 이 LoRA의 모든것을 표현한 문장이다.

 

LoRA를 좀 더 이해하기 쉽게, 직접적으로 풀어쓴다면 이렇게 설명할 수 있을 것 같다.

"기존 weight parameter들은 freezing한 후,
별도로 구분된 저차원의 레이어(adapter, 이하 어댑터)를 학습하고, 이 어댑터를 통과한 값과,
기존 Pre-Trained model을 통과한 값을 더함(+)으로써 미세하게 조정내용을 반영하는 방법" 

LoRA Adapter 구조

 

위 그림은 LoRA가 적용, 학습되는 구조를 표현한 그림이다.

그림을 보면서 위 설명을 다시 한번 읽어보면 좀 더 이해가 잘 될것이다.

그리고, LoRA는 현재 PEFT 라이브러리에서 lora_config클래스를 통해

손쉽게 구현이 가능하지만, 내부 학습 방식을 좀 더 상세하게 이해하기 위해서는 코드레벨로 이해해보는것이 필요하다.

 

아래는 LoRA의 학습 Architecture를 code로 직접 구현한 내용이다.

class LoRA(nn.Module):
    def __init__(self, model, r=8, alpha=16, dropout_rate=0.1):
        super(LoRA, self).__init__()
        self.model = model
        self.r = r
        self.alpha = alpha
        self.scaling = alpha / r
        self.dropout_rate = dropout_rate  # 드롭아웃 비율 추가

        # 원본 모델 가중치 고정
        for param in self.model.parameters():
            param.requires_grad = False

        # LoRA 레이어를 어텐션 프로젝션에 주입
        for name, module in self.model.named_modules():
            if isinstance(module, RobertaSelfAttention):
                # 쿼리, 키, 밸류 프로젝션 수정
                self._inject_lora(module)

    def _inject_lora(self, module):
        for proj in ["query", "key", "value"]:
            weight = getattr(module, proj).weight
            bias = getattr(module, proj).bias
            out_features, in_features = weight.shape
    
            # LoRA 레이어 생성
            lora_A = nn.Linear(in_features, self.r, bias=False)
            lora_B = nn.Linear(self.r, out_features, bias=False)
            nn.init.normal_(lora_A.weight, std=0.02)
            nn.init.zeros_(lora_B.weight)
    
            # LoRAProjection에 드롭아웃 비율 전달
            lora_proj = LoRAProjection(weight, bias, lora_A, lora_B, self.scaling, dropout_rate=self.dropout_rate)
            setattr(module, proj, lora_proj)

    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)


class LoRAProjection(nn.Module):
    def __init__(self, weight, bias, lora_A, lora_B, scaling, dropout_rate=0.0):
        super(LoRAProjection, self).__init__()
        self.weight = weight
        self.bias = bias
        self.lora_A = lora_A
        self.lora_B = lora_B
        self.scaling = scaling

        # 드롭아웃 레이어 추가
        if dropout_rate > 0.0:
            self.dropout = nn.Dropout(p=dropout_rate)
        else:
            self.dropout = None

    def forward(self, x):
        # LoRA 부분 계산
        lora_out = self.lora_A(x)
        if self.dropout is not None:
            lora_out = self.dropout(lora_out)
        lora_out = self.lora_B(lora_out)

        # 원본 출력과 LoRA 보정값의 합산
        return nn.functional.linear(x, self.weight, self.bias) + self.scaling * lora_out

 

자, 이제 그림과 위 코드를 뜯어 보면서 LoRA 학습 진행과정을 차례차례 살펴보자.

(* 이제 위 그림에서 주황색 마주보고있는 사다리꼴을 어댑터로 표현하겠다.)

 

0) Parameter Settings

def __init__(self, model, r=8, alpha=16, dropout_rate=0.1):
        super(LoRA, self).__init__()
        self.model = model
        self.r = r
        self.alpha = alpha
        self.scaling = alpha / r
        self.dropout_rate = dropout_rate  # 드롭아웃 비율 추가

 

위는 LoRA 클래스 내부에서 LoRA 학습 옵션을 정의하는 부분이다.

먼저 기존 Pre-trained Model을 'model'인자로 상속받게 된다.

그리고 여기서 r(rank), alpha 인자가 LoRA의 핵심 인자값이 된다.

 

(1) R (Rank)

먼저 rank는 LoRA에서 저차원 행렬의 차원을 결정하는 주요 파라미터다.

rank는 위 그림의 어댑터에 해당하는 A, B weight matrix의 이음새 부분에 해당하는 차원을 셋팅하는 부분이다.

rank는 A, B weight matrix의 차원을 결정지으며, 

크기가 작을 수록, 학습 시간과 리소스 사용량, 학습할 수 있는 표현의 다양성(복잡성)은 감소한다.

반대로 크기가 클수록, 학습 시간, 리소스 사용량, 학습할 수 있는 표현의 다양성은 증가한다.

조금 더 자세한 내용은 아래 adapter 설명 부분에서 참고하자.

 

(2) Alpha

alpha는 학습률과 관련된 스케일링 파라미터로, 저차원 행렬로부터 얻은 변화량을 조절하는데 사용된다.

LoRA는 저차원에서 학습한 보정 값을 A x B 로 표현하는데, 이 보정 값을 그대로 사용하지 않고, 'alpha / rank' 로 스케일링해서 적용한다

 

  • 목적 : LoRA에서 학습되는 보정값을 너무 크게 적용하지 않도록, 모델이 안정적으로 학습할 수 있도록 저차원 보정값의 크기를 조절하는것을 목적으로 하며, 이를 통해 저차원의 효과로 이뤄지는 모델의 보정이 너무 크거나 작지 않도록 가중치 업데이트의 크기를 제어하는 역할을 한다.
  • 효과 : 학습의 안정성과 수렴 속도를 높이고, 저차원에서 얻은 정보를 적절하게 원래의 가중치에 반영할 수 있게 된다.

 

 

0) Origin model weight parameter Freezing

# 원본 모델 가중치 고정
for param in self.model.parameters():
    param.requires_grad = False

 

당연하게도, 새로 추가된 LoRA 파트가 아닌, 기존모델의 parameter는 update되지 않도록 freezing 처리한다.

 

1) LoRA Adapter Input

먼저, LoRA의 어댑터에 어떤값을 input vector[x]로 사용해야하는지 궁금할 것이다.

결론적으로 input vector는 "last_hidden_state"를 사용하는것이 일반적이다.

last_hidden_state을 사용하여, adapter를 학습하는 이유는,

대부분의 트랜스포머 기반 모델에서, last_hidden_state는 

input sequence가 여러 layer층을 거치면서, 가장 풍부하고 고차원적인 표현을 담고 있는

최종적인 학습 정보이므로, last_hidden_state에 미세조정을 할때 가장 효과적이기 때문이다.

 

 

2) Pre-Trained Weights & Adapter

(1) Pre-Trained Weights

먼저 input vector는 기존 weight parameter를 통과하게 되면서,

기존의 가중치 행렬을 적용받은 output A를 뱉는다.

단순히 pre-trained된 가중치 행렬인 Linear layer를 통과하는 과정을 의미한다.

 

 

(2) Adapter

먼저 weight parameter가 update되기 전 최초의 adapter

는 위 수식과 같이, 평균이 0, 분산이 σ2 인 정규분포값을 최초값(initialization)을 가진다.

최초에는 제로 행렬로 초기화 되어있지만, backpropagation을 반복하면서, 값이 update되게 된다.

 

여기서, 하단의 가중치 행렬을 A, 윗단의 가중치 행렬을 B라고할때, 

A행렬에서 고차원에서 저차원으로 투사되면서, 모델은 '중요한 특징'들만 선택적으로 학습을 진행한다.

고차원의 데이터를 간소화 하며, 정보의 중요한 부분만 추출해 학습하는 과정을 거친다.

 

이어서, B행렬을 통과하며, 저차원에서 고차원으로 다시 투사하는 과정에서 중요한 정보들이 고차원 공간에

반영되면서 최종적으로 성능이 개선되게끔 학습된다.

LoRA의 가중치 업데이트 공식은 아래와 같다.

 

 

  • W′ : 업데이트 된 가중치 행렬
  • A : rank x d 크기의 저차원 행렬을 의미한다. (저차원으로의 투사)
  • B : d x rank 크기의 저차원 행렬을 의미한다. (저차원에서 고차원으로 다시 투사)
  • Alpha : 보정값의 크기를 결정하는 파라미터로, alpha값이 클수록 보정값이 더 많이 반영된다고 이해하면 된다.

> 결국 rank설정값은 A와 B의 행렬의 크기를 결정하게 된다.

 

이 Rank개념에 대해 논문에서는 4, 8차원에서 가장 일반화된 성능을 가지고오며, 성능과 리소스, 속도 측면의 trade-off관계의

적정성을 가진다고 설명하고 있다.

 

쉽게말해,

Rank를 크게 설정하면, 좀 더 복잡한 문맥정보를 학습하겠지만,

LoRA의 장점인 빠른 학습 속도 및 낮은 리소스 사용량이라는 장점을 희생해야하고,

 

Rank를 작게 설정할 수록 장점은 극대화 되지만, 복잡한 문맥정보를 학습하는데 한계가 있다는 점이 있기 때문에,

잘 조정해서 활용해야한다.

 

여담이지만, LoRA는 이같은 특징때문에, 반복적인 fine-tuning에서 추가적인 성능향상이 어렵다는 단점을 가지고 있고, 

최근 이를 보완해서 input, output 차원이 동일한 mora라는 개념이 등장하기도 했다.

 

 

자, 그러면 LoRA Adapter가 실제로 학습 되는 부분을 살펴보도록 하자.

먼저 LoRA가 학습되는 파트에 대한 클래스('LoRAProjection' : adapter 부분에 해당)를 구현한 후,

해당 클래스를 LoRA 클래스 내부에 끼워넣게 된다.

class LoRAProjection(nn.Module):
    def __init__(self, weight, bias, lora_A, lora_B, scaling, dropout_rate=0.0):
        super(LoRAProjection, self).__init__()
        self.weight = weight
        self.bias = bias
        self.lora_A = lora_A
        self.lora_B = lora_B
        self.scaling = scaling

        # 드롭아웃 레이어 추가
        if dropout_rate > 0.0:
            self.dropout = nn.Dropout(p=dropout_rate)
        else:
            self.dropout = None

    def forward(self, x):
        # LoRA 부분 계산
        lora_out = self.lora_A(x)
        if self.dropout is not None:
            lora_out = self.dropout(lora_out)
        lora_out = self.lora_B(lora_out)

        # 원본 출력과 LoRA 보정값의 합산
        return nn.functional.linear(x, self.weight, self.bias) + self.scaling * lora_out

위는 실제 adapter의 학습 방식을 표현하는 부분이다.

nn.functional.linear(x, self.weight, self.bias) + self.scaling * lora_out

기존의 weight, bias가 적용된 linear matrix의 output을 lora output와 더해서,

lora adapter로 도출된 내용을 통해 미세조정을 적용받게 되는것을 확인할 수 있다.

 

 

(3) 학습된 weight의 병합

마지막으로, 1번과 2번의 결과로 도출된 weight를 더하게되는데, 

이것은, 만약 곱셈으로 처리될 경우, vector 공간에서 전체 weight의 방향과 위치 자체가 update되어,

기존의 학습 정보가 크게 틀어지게된다.

 

말그대로, substream task에 대한 '미세조정'의 목적을 가진 방법이기 때문에, 

덧셈을 통해, 위치 조정의 효과를 보기 위해 덧셈을 통해 처리하게 된다.

 


오늘은 최근 Fine-Tuning의 대세가 된, LoRA의 구조 및 구체적인 학습 방식에 대해서 살펴보았다.

코드레벨의 동작 방식을 살펴보면서 LoRA가 구체적으로 어떤 학습방식을 통해,

망각현상과 리소스 부하를 감소하는 효과를 내었는지 이해할 수 있었다.

 

다음에는 이러한 LoRA의 장점을 살려서, Embedding Model(RoBERTa)을 LoRA를 통해

직접 Fine-tuning했던 과정과 결과에 대해서 정리해보도록 하겠다.

 

 

 


[Reference]

<Paper>

https://arxiv.org/abs/2106.09685

 

<Github>

https://github.com/microsoft/LoRA

 

<Blog>

https://www.databricks.com/kr/blog/efficient-fine-tuning-lora-guide-llms

https://4n3mone.tistory.com/7