Post

[Paper]LoRA

2021년에 Microsoft 연구팀에서 “LoRA: Low-Rank Adaptation of Large Language Models”라는 논문을 통해 대규모 언어 모델을 효율적으로 조정할 수 있는 기술을 제안하였습니다. 이는 사전 학습된 모델의 가중치를 고정(freeze)시키고, 학습 가능한 분해 행렬을 주입함으로써 작동합니다. GPT-3와 비교하여 학습 가능한 파라미터를 10,000배 이상 줄일 수 있었고, GPU의 메모리 사용량 또한 3배 이상 감소시킬 수 있었습니다.

Ⅰ. LoRA

1 Figure 1 : Overview of LoRA

NLP에서의 많은 응용은 하나의 거대한 모델에 의존하여 사전 학습된 언어 모델을 다양한 다운스트림으로 활용하는 것입니다. 이러한 작업은 미세 조정을 통해 조정되며, 미세 조정은 사전 학습된 모델의 모든 매개변수를 업데이트합니다. 그러나 모델의 규모가 날이 갈수록 커지며 매개변수 수의 증가는 중요한 문제로 발전하게 됩니다.

이를 극복하기 위해 매개변수의 일부를 조정하거나 새로운 작업을 위한 외부 모듈을 학습하는 등의 여러 방법들을 도입하여 추가적인 작업별 매개변수를 저장하고 로드하는 방식으로 배포 시 운영 효율을 증가시켜왔습니다. 그러나 이러한 기술들은 추론 지연 시간을 증가시켰습니다. 더 중요한 것은 이러한 방법들은 효율과 모델 성능 사이의 trade-off 관계를 가지게 되는 것입니다.

과도하게 매개변수화된 모델이 본질적인\(_{intrinsic}\) 저차원을 가지는 사실에 영감을 받아 가중치가 모델 응용 과정에서 본질적인 low-rank를 가지며 변화할 수 있다는 가설을 제시하였습니다. 이러한 LoRA는 신경망에서 여러 dense layer들을 간접적으로 행렬 분해를 통해 최적화함으로써 작동합니다.

이러한 LoRA의 주요 이점은 아래와 같습니다.

  1. 사전학습된 모델이 LoRA의 작은 모듈들로 구성
  2. 학습 효율의 증대로 하드웨어적 장벽을 낮춤.
  3. Full 미세조정 모델과 비교하여 추론 latency가 없음.
  4. 많은 기존 방법들과 직교함.

(1)을 통해 사전학습 모델을 그대로 사용하고, 작업별로 행렬 A와 B를 바꿈으로써 효율적인 전환이 가능하고, 저장 공간과 작업 전환 over-head를 상당히 줄입니다. (2)를 통해 적응형 옵티마이저를 사용할때, 상태를 유지하거나 계산하지 않아도 되고 매우 작은 저차원 행렬만 최적화 됩니다. 또한, (3)을 통해 기존의 사전학습된 모델의 가중치를 고정(freeze)시켜 배포시 병합이 가능하고 (4)를 통해 다른 작업들에 응용이 가능합니다.

Ⅱ. Problem Statement

\(P_{\Phi}(y\vert x)\)를 사전학습된 자동회귀 언어 모델로 주어졌다고 가정하면 아래와 같이 GPT 같은 포괄적인 멀티태스크 학습자가 될 수 있습니다.

\[\underset{\Theta}{max} \sum_{(x,y) \in Z}\sum_{t=1}^{\vert y \vert} log(p_{\Phi}(y_t|x,y_{\lt t}))\]

위와 같은 전체 미세 조정의 단점은 각 다운스트림 작업마다 다른 파라미터 \(\Delta\Phi\)를 학습해야 한다는 것입니다. 따라서, 만약 사전학습된 모델이 크다면 이러한 많은 독립적인 인스턴스들을 저장하고 배포하는 일은 어려울 수 있습니다.

해당 논문에서는 효율적인 매개변수 접근 방법을 채택하여 태스크 특화된 파라미터 증가 \(\Delta\Phi = \Delta\Phi(\Theta)\)가 더 작은 크기의 매개변수 세트를 통해 인코딩될 수 있음을 보여주었습니다.

\[\underset{\Theta}{max} \sum_{(x,y) \in Z}\sum_{t=1}^{\vert y \vert} log(p_{\Phi_0 + \Delta\Phi(\Theta)}(y_t|x,y_{\lt t}))\]

Ⅲ. Aren’t existing solutions good enough?

언어 모델을 예시로 들어 두 가지 주요 전략으로 어댑터 레이어를 추가하는 방법입력 레이어 활성화의 일부 형태(프롬프트)를 최적화하는 방법입니다. 어댑터 레이어를 포함시키는 방식의 디자인은 추가 계산을 우회하는 직접적인 방법이 없어 추론 시간이 증가합니다. 또한, 프롬프트를 직접 최적화하는 것은 어려운 작업입니다. 즉, 이 두 전략에는 특히 대규모 모델이나 지연시간에 민감한 시나리오에서는 한계가 존재합니다.

Ⅳ. Method

ⅰ. Low-Rank Parameterized update matrices

사전학습된 언어 모델은 “본질적인 저차원”이며 효율적으로 학습 될 수 있는 저차원의 부분 공간으로도 학습 될 수 있습니다. 이에 영감을 받아 가중치를 또한 “본질적인 low-rank”를 가지도록 업데이트합니다. 학습 동안 \(W_0\)는 고정\(_{freeze}\)되어 있고 A와 B를 학습 가능한 매개변수로 포함됩니다. \(W_0\)와 \(\Delta W = BA\) 는 같은 입력을 받아서 곱해지며 그 결과 출력 벡터의 표현은 coordinate-wise로 합산됩니다.

\[h = W_0x + \Delta Wx = W_0x + BAx\]

위의 적응형 매개변수 행렬 A와 B를 랜덤 가우시안 초기화를 적용하고, 학습 초기에는 \(\Delta W\) 또한 0으로 지정합니다. 그러면 \(\Delta Wx\) 가 \(\frac{\alpha}{\gamma}\) 를 통해 스케일링 되고, \(\alpha\)는 \(\gamma\)의 상수입니다.

  • 전체 미세조정의 일반화

    LoRA는 그래디언트 업데이트가 전체 차원에서 일어날 필요가 없어 더 작은 차원의 LoRA Rank r을 설정하여 학습된 가중치 행렬의 랭크와 일치하도록 근사해 더 적은 수의 매개변수로도 모델의 표현력을 유지하면서 학습이 가능합니다.

  • 추가 추론 latency 없음

    LoRA는 추가적인 inference latency 없이, 가중치 행렬을 명시적으로 계산하고 저장하여 모델을 작업별로 쉽게 전환하고 재사용이 가능하게 되어 매우 적은 메모리 오버헤드와 빠른 수행이 가능합니다.

ⅱ. Applying LoRA To Transformer

트랜스포머 아키텍처에서는 셀프 어텐션 모듈에 네 개의 가중치 행렬(Wq, Wk, Wv, Wo)이 존재하는데 이 중 Wq는 단일 행렬로 다루어지며 차원은 d_model × d_model입니다. 다운스트림 작업에서는 어텐션 가중치만을 적응시켜 상당한 메모리와 저장 공간을 절약할 수 있습니다. 예를 들어, GPT-3의 경우 VRAM 소비량이 1.2TB에서 350GB로 줄어들었고, 학습 중에는 25%의 속도 향상이 관찰되었습니다.

Ⅴ. Code

즉, LoRA는 가중치 행렬을 작은 차원의 A와 B 행렬을 통해 기존의 가중치 행렬로 근사화하고 업데이트하는 방법이라 할 수 있습니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
############################
###  Applying to Linear  ###
############################
class Linear(nn.Linear, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(...):
        ...

        # 학습 가능한 적응형 파라미터 선언
        if r > 0:
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
            self.scaling = self.lora_alpha / self.r

    def train(self, mode: bool = True):
        ...
        
        # 기존 가중치는 freeze 시켜놓고, 
        # LoRA의 BA 곱셈행렬을 통해 생성된 행렬을 더하여 업데이트 진행
        if self.r > 0:
            self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
        
    
    def forward(self, x: torch.Tensor):
        ...
        
        # 기존 선형 함수
        result = F.linear(x, T(self.weight), bias=self.bias)    
        # LoRA를 적용 시킨 함수
        result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) \
                  * self.scaling
        return result


#################################
###  Applying to Transformer  ###
#################################

class MergedLinear(nn.Linear, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(...):
        ...
        # 학습 가능한 매개변수 선언
        if r > 0 and any(enable_lora):
            self.lora_A = nn.Parameter(
                self.weight.new_zeros((r * sum(enable_lora), in_features)))
            self.lora_B = nn.Parameter(
                self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r))
            )
            self.scaling = self.lora_alpha / self.r
            
class Attention(nn.Module):
    def __init__(self, ...):
        ...
        self.c_attn = lora.MergedLinear(...)

    def forward(self, ...):
        hidden_states = x

        # 각 Q, K, V 행렬을 LoRA를 통해 분해
        x = self.c_attn(x)
        query, key, value = x.split(self.split_size, dim=2)

Ⅵ. REFERENCES

  1. LoRA: Low-Rank Adaptation of Large Language Models
  2. microsfot LoRA github



This post is licensed under CC BY 4.0 by the author.