LoRA: Low-Rank Adaptation of Large Language Models 논문 리뷰 (+ Adapter, Prefix Tuning)

LoRA: Low-Rank Adaptation of Large Language Models 논문 리뷰 

 

 

• 최근에 나온 MoRA를 읽어보기 전에 LoRA 논문을 올리지 않은 것 같아, 이번 기회에 정리! 
  • 사용 방법에 대한 코드도 함께 정리해볼 예정이다

 

 

들어가기 전에 : PeFT의 등장 배경   

 

 

• Fully Fine Tuning

 

• Parameter-efficient approach
  • 세타는 오리지널 파라미터보다 훨씬 적은 양의 파라미터
  • 세타_0에 아주 작은 변화량 더해준다 => 이게 LoRA의 핵심이다 

 

 

 

 

Abstract & Introduction 

 

 

• Transfer learning의 붐이 시작된 이래로 수십 개의 연구에서 parameter와 compute-efficient하게 model adaptation을 수행하는 방법을 연구
• Language Modeling의 경우 크게 두 가지 중요한 strategy가 있음
  • Add adapter layers
    • Parameter-Efficient Transfer Learning for NLP. Houlsby et al., 2019
    • Learning multiple visual domains with residual adapters. Rebuffi et al., 2017
    • AdapterFusion: Non-Destructive Task Composition for Transfer Learning. Pfeiffer et al., 2021
    • AdapterDrop: On the Efficiency of Adapters in Transformers. Ruckle et al., 2020
  • Optimizing some forms of the input layer activations (Prefix Tuning, P tuning의 경우)
    • Prefix-Tuning: Optimizing Continuous Prompts for Generation. Li & Liang, 2021
    • The Power of Scale for Parameter-Efficient Prompt Tuning. Lester et al., 2021
    • WARP: Word-level Adversarial ReProgramming. Hambardzumyan et al., 2021
    • GPT Understands, Too. Liu et al., 2021

 

 

 

• Adapter의 학습 방법 
  • 전체 Parameter는 freeze 해놓고 새롭게 추가하는 adapter layer만 새로 학습하는 것 의미 
  • 아래 그림처럼 Transformer Layer에서 FFN 통과 후 & Layer Norm. 전에 Adapter Layer를 해준다 
  • Adapter Layer 자체는 FF Down Projection과 FF Up Projection로 이뤄져 있고, 이것만 학습한다 
• Prefix Tuning 
  • In-Context Learning과 Finetunig의 문제
  • 트랜스포머 모델에서 Prefix를 Key Matrix와 Value Matirx 앞에 붙여서 학습 

가장 최초의 Adapter

 

prefix tuning

 

 

• 위 두 방법의 문제 
  • Adapter Layers introduce inference latency
    • Large-scale Neural Network는 대기 시간을 낮게 유지하기 위해 병렬 처리에 의존
    • Adapter Layer → Sequentially
      • Multi head attention의 결과값을 받아야만 Adapter에서 연산할 수 있으므로! 
  • Directly Optimizing the Prompt is Hard
    • Prefix-tuning은 최적화하기 어렵고 그 성능이 trainable parameter non-monotonically하게 변한다는 것을 관찰
      
      • non-monotonically : 우상향이나 우하향이 아닌 진동하는 경우 의미 
    • Adaptation을 위해 sequence length의 일부를 미리 떼어놔야 하기 때문에 downstream task를 처리하는데 사용할 수 있는 sequence length가 줄어듦

 

• 본 연구는 학습된 over-parameterized model이 실제로 낮은 low intrinsic dimension에 있다는 아래 두 논문에 영감을 받았음
  • Measuring the Intrinsic Dimension of Objective Landscapes Li et al., 2018a
  • Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning Aghajanyan et al., 2020
• LoRA: Low-Rank Adaptation
• LoRA를 사용하면?
  • Pre-trained weight를 고정된 상태(freeze)로 유지하면서
  • Adaptation 중 dense layer의 변화에 대한 rank decomposition matrices를 최적화
  • 이를 통해 신경망의 일부 dense layer를 간접적으로 훈련시키는 것이 가능
• LoRA는 trainable parameter의 수가 적고 학습 처리량이 높으며 inference latency가 이전 연구 대비 적음
• 그럼에도 불구하고 RoBERTa, DeBERTa, GPT-2, GPT-3에서 fine-tuning보다 같거나 더 나은 성능을 보여줌

 

 

Methodology

 

 

 

• d_model​ : input/output dimension size
• Wq,Wk,Wv,Wo : query/key/value/output projection matrices in self-attention module
• W or W0: pre-trained weight matrix
• ΔW : accumulated gradient update during adaptation
• r : rank of a LoRA module
• Transformers 논문의 setting을 따름
  • e.g., use Adam optimizer
  • d_ffn=4×d_model​

 

 

 

 

• LoRA는 모든 dense layer에 적용 가능
• Neural Network는 Matrix multiplication을 수행하는 많은 dense layer를 포함
• 이 matrices의 weights는 일반적으로 full-rank
• Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning
  • Aghajanyan et al., 2020
  • PLM은 low intrinsic dimension을 가짐 (보통 over-parameterized model이 가지는 특징이라고 함)
  • 더 작은 subspace에 대한 random projection에도 불구하고 여전히 효율적으로 학습할 수 있음을 보임
  • 특히 RoBERTa의 경우 only 200 trainable parameters로 90%의 performance를 달성했다고 주장

 

 

 

• LoRA는 모델 뿐만 아니라, 가중치에 대한 update 또한 intrinsic rank가 낮다는 가설을 설정
  • 가중치에 대한 update도 adaptation 중 intrinsic rank가 낮다
• 변형 
  • Pre-trained weight matrix W0​
    • W0∈Rd×k에 대하여 이 행렬에 대한 update를 low-rank decomposition을 통해 아래와 같이 표현:
      • W0​+ΔW=W0​+BA
    • 여기서 B∈Rd×rB , A∈Rr×kA, r≪min⁡(d,k)
  • 학습 중 W0는 frozen:
    • gradient update를 수행하지 않음
  • A와 B
    • trainable parameters
  • Note:
    
    • W0와 ΔW=BA는 모두 동일한 입력으로 곱해짐
    • 각각의 출력 벡터는 coordinate-wise하게 summed
  • Forward pass
    • h=W0x+ΔWx=W0x+BA

 

이쯤에서 궁금한 점: Rank는 어떻게 설정? LoRA는 어디에 적용해야 하는가? 

 

• Rank는 어떻게 설정? 
  • 저자들의 실험 결과, 매애애우 작은 Rank로도 좋은 성능을 내더라 

 

 

• LoRA는 어디에 적용해야 하는가? 
  •  트랜스포머 한 블럭에는 여러 개의 가중치 매트릭스가 있는데, self-attention module에 4개의 가중치 매트릭스가 있고(𝑊𝑞,𝑊𝑘,𝑊𝑣,𝑊𝑜) 인코더, 디코더 각각에 MLP module이 있는데, 
  • 본 논문이 LoRA를 적용하는 가중치 매트릭스는 attention 가중치 매트릭스로만 제한해서 실험.
    • 업데이트하는 파라미터의 용량을 제한하기 위해 예를 들어 한 가지 종류의 가중치 매트릭스를 사용한다면 rank=8이며, 두 가지 종류의 가중치 매트릭스를 사용한다면 rank=4로 함
  • 결과 
    • 결과를 살펴보면 Wq𝑊𝑞, Wv𝑊𝑣 두 가중치 파라미터에 LoRA를 적용하는 것이 업데이트하는 가중치의 종류를 최대한 적게 가지면서 좋은 성능을 냈다. 이를 통해 알 수 있는 것은 rank를 4로 해도 충분히 ΔW의 정보를 담을 수 있다는 것이다.

 

 

 

Code  (Scratch)

 

 

• LoRA Layer Code

class Linear(nn.Linear, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(
        self, 
        in_features: int, 
        out_features: int, 
        r: int = 0, 
        lora_alpha: int = 1, 
        lora_dropout: float = 0.,
        fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
        merge_weights: bool = True,
        **kwargs
    ):
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
                           merge_weights=merge_weights)

        self.fan_in_fan_out = fan_in_fan_out
        # Actual trainable parameters
        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
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
        self.reset_parameters()
        if fan_in_fan_out:
            self.weight.data = self.weight.data.transpose(0, 1)

    def reset_parameters(self):
        nn.Linear.reset_parameters(self)
        if hasattr(self, 'lora_A'):
            # initialize B the same way as the default for nn.Linear and A to zero
            # this is different than what is described in the paper but should not affect performance
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

    def train(self, mode: bool = True):
        def T(w):
            return w.transpose(0, 1) if self.fan_in_fan_out else w
        nn.Linear.train(self, mode)
        if mode:
            if self.merge_weights and self.merged:
                # Make sure that the weights are not merged
                if self.r > 0:
                    self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = False
        else:
            if self.merge_weights and not self.merged:
                # Merge the weights and mark it
                if self.r > 0:
                    self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = True       

    def forward(self, x: torch.Tensor):
        def T(w):
            return w.transpose(0, 1) if self.fan_in_fan_out else w
        if self.r > 0 and not self.merged:
            result = F.linear(x, T(self.weight), bias=self.bias)            
            result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling
            return result
        else:
            return F.linear(x, T(self.weight), bias=self.bias)

 

 

•  중요한 부분은 아래다 
  • A matrix는 kaiming uniform으로 구현됨
  • B matrix는 0 행렬로 초기화

        # Actual trainable parameters
        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
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
        self.reset_parameters()
        if fan_in_fan_out:
            self.weight.data = self.weight.data.transpose(0, 1)

 

 

 

• Inference시에는 학습한 A Matrix와 B Matrix 불러와서 쓰기만 하면 됨!
  • 근데 이제 adapter와는 다르게 adapter는 또 시퀀셜하게 연산해줘야 되는데,
  • 얘는 그냥 원래 모델 파라미터에 BA를 합쳐주면 된다 ㅇㅇ 

    def train(self, mode: bool = True):
        def T(w):
            return w.transpose(0, 1) if self.fan_in_fan_out else w
        nn.Linear.train(self, mode)
        if mode: # train
            if self.merge_weights and self.merged:
                # Make sure that the weights are not merged
                if self.r > 0:
                    self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = False
        else: # eval 
            if self.merge_weights and not self.merged:
                # Merge the weights and mark it
                if self.r > 0:
                    self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = True

 

 

 

 

huggingface peft 구현체 

 

 

 

from transformers import AutoModelForSeq2SeqLM
from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, TaskType
import torch
from datasets import load_dataset
import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"
from transformers import AutoTokenizer
from torch.utils.data import DataLoader
from transformers import default_data_collator, get_linear_schedule_with_warmup
from tqdm import tqdm
from datasets import load_dataset

device = "cuda"
model_name_or_path = "bigscience/mt0-large"
tokenizer_name_or_path = "bigscience/mt0-large"

checkpoint_name = "financial_sentiment_analysis_lora_v1.pt"
text_column = "sentence"
label_column = "text_label"
max_length = 128
lr = 1e-3
num_epochs = 3
batch_size = 8

# creating model
peft_config = LoraConfig(task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)

model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()
model

# loading dataset
dataset = load_dataset("financial_phrasebank", "sentences_allagree")
dataset = dataset["train"].train_test_split(test_size=0.1)
dataset["validation"] = dataset["test"]
del dataset["test"]

classes = dataset["train"].features["label"].names
dataset = dataset.map(
    lambda x: {"text_label": [classes[label] for label in x["label"]]},
    batched=True,
    num_proc=1,
)

dataset["train"][0]

 

 

깃헙과 잘 정리된 블로그 글을 참고하고 정리한다 

 

LoRAConfig (https://github.com/huggingface/peft/blob/main/src/peft/tuners/lora/config.py)

config = LoraConfig(
    r=16, #attention heads
    lora_alpha=32, #alpha scaling
    # target_modules=["q_proj", "v_proj"], #if you know the 
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM" # set this for CLM or Seq2Seq
)

 

 

 

• r:
  • 타입: int
  • 설명: LoRA의 주의(attention) 차원(랭크)을 설정합니다. 기본값은 8입니다.
• target_modules:
  • 타입: Optional[Union[List[str], str]]
  • 설명: LoRA 어댑터를 적용할 모듈의 이름 목록입니다. 문자열을 전달하면 정규 표현식이 적용되고, 문자열 목록을 전달하면 정확히 일치하는 모듈 이름에 적용됩니다. 기본값은 None입니다.
• lora_alpha:
  • 타입: int
  • 설명: LoRA의 스케일링 파라미터를 설정합니다. 기본값은 8입니다.
• lora_dropout:
  • 타입: float
  • 설명: LoRA 레이어의 드롭아웃 확률을 설정합니다. 기본값은 0.0입니다.
• fan_in_fan_out:
  • 타입: bool
  • 설명: 레이어가 (fan_in, fan_out) 형태로 가중치를 저장하는 경우, 이를 True로 설정합니다. 예를 들어, GPT-2는 Conv1D를 사용하므로 이 값을 True로 설정해야 합니다

 

 

 

+) 어떤 글에 따르자면 경험적으로 r과  target_modules가 adaptation 품질에 큰 영향을 미친다고 함 (https://www.databricks.com/kr/blog/efficient-fine-tuning-lora-guide-llms)

 

#If only targeting attention blocks of the model
target_modules = ["q_proj", "v_proj"]

#If targeting all linear layers
target_modules = ['q_proj','k_proj','v_proj','o_proj','gate_proj','down_proj','up_proj','lm_head']

 

 

 

*) 지원하는 Task Type은 아래와 같다 

https://github.com/huggingface/peft/blob/v0.8.2/src/peft/utils/peft_types.py#L68-L73

class TaskType(str, enum.Enum):
    """
    Enum class for the different types of tasks supported by PEFT.

    Overview of the supported task types:
    - SEQ_CLS: Text classification.
    - SEQ_2_SEQ_LM: Sequence-to-sequence language modeling.
    - Causal LM: Causal language modeling.
    - TOKEN_CLS: Token classification.
    - QUESTION_ANS: Question answering.
    - FEATURE_EXTRACTION: Feature extraction. Provides the hidden states which can be used as embeddings or features
      for downstream tasks.
    """

    SEQ_CLS = "SEQ_CLS"
    SEQ_2_SEQ_LM = "SEQ_2_SEQ_LM"
    CAUSAL_LM = "CAUSAL_LM"
    TOKEN_CLS = "TOKEN_CLS"
    QUESTION_ANS = "QUESTION_ANS"
    FEATURE_EXTRACTION = "FEATURE_EXTRACTION"

 

 

 

 

 

 

get_peft_model

model = get_peft_model(model, peft_config)

• AutoModelForCasualLM을 통해 불러오는 Model을 get_peft_model 함수에 인자로 넣어준다. 이를 통해 peft model을 준비함을 알 수 있다. 

 

 

def get_peft_model(
    model: PreTrainedModel, peft_config: PeftConfig, adapter_name: str = "default", mixed: bool = False
) -> PeftModel | PeftMixedModel:
    """
    Returns a Peft model object from a model and a config.

    Args:
        model ([`transformers.PreTrainedModel`]):
            Model to be wrapped.
        peft_config ([`PeftConfig`]):
            Configuration object containing the parameters of the Peft model.
        adapter_name (`str`, `optional`, defaults to `"default"`):
            The name of the adapter to be injected, if not provided, the default adapter name is used ("default").
        mixed (`bool`, `optional`, defaults to `False`):
            Whether to allow mixing different (compatible) adapter types.
    """
    model_config = getattr(model, "config", {"model_type": "custom"})
    if hasattr(model_config, "to_dict"):
        model_config = model_config.to_dict()

    peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)

    if mixed:
    	#return값이 PeftMixedModel 또는 PeftModel이다. get_peft_model의 인자가 거의 그대로 들어간다.
        return PeftMixedModel(model, peft_config, adapter_name=adapter_name)

    if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not peft_config.is_prompt_learning:
        return PeftModel(model, peft_config, adapter_name=adapter_name)

 

 

이것은 get_peft_model 함수 정의입니다. 네 개의 인자를 받습니다:

1. model: transformers 라이브러리의 PreTrainedModel 인스턴스.
2. peft_config: PEFT 모델의 매개변수를 포함하는 PeftConfig 인스턴스.
3. adapter_name: 기본값이 "default"인 선택적 문자열 인자. 사용될 어댑터의 이름을 지정합니다.
4. mixed: 기본값이 False인 선택적 부울 인자. 서로 다른 어댑터 유형의 혼합을 허용할지 여부를 나타냅니다.

이 함수는 PeftModel 또는 PeftMixedModel을 반환합니다.

 

 

 

PeftModel

class PeftModel(PushToHubMixin, torch.nn.Module):
    """
    Base model encompassing various Peft methods.

    Args:
        model ([`~transformers.PreTrainedModel`]): The base transformer model used for Peft.
        peft_config ([`PeftConfig`]): The configuration of the Peft model.
        adapter_name (`str`,  *optional*): The name of the adapter, defaults to `"default"`.
    """

    def __init__(self, model: PreTrainedModel, peft_config: PeftConfig, adapter_name: str = "default") -> None:
        super().__init__()
        self.modules_to_save = None
        self.active_adapter = adapter_name
        #위 get_peft_model에서 넣어줬다. 안넣어줬으면 'default'로 들어왔다. 
        self.peft_type = peft_config.peft_type
        self._is_prompt_learning = peft_config.is_prompt_learning
        if self._is_prompt_learning:
            self._peft_config = {adapter_name: peft_config}
            self.base_model = model
            self.add_adapter(adapter_name, peft_config)
        else:
            self._peft_config = None
            cls = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]
            # base model 선언.
            self.base_model = cls(model, {adapter_name: peft_config}, adapter_name)  
            self.set_additional_trainable_modules(peft_config, adapter_name)

        if getattr(model, "is_gradient_checkpointing", True):
            model = self._prepare_model_for_gradient_checkpointing(model)

        # the `pretraining_tp` is set for some models to simulate Tensor Parallelism during inference to avoid
        # numerical differences, https://github.com/pytorch/pytorch/issues/76232 - to avoid any unexpected
        # behavior we disable that in this line.
        if hasattr(self.base_model, "config") and hasattr(self.base_model.config, "pretraining_tp"):
            self.base_model.config.pretraining_tp = 1

 

 

 

 

 

Ref. 

 

https://arxiv.org/abs/2106.09685

LoRA: Low-Rank Adaptation of Large Language Models

 

https://www.youtube.com/watch?v=BJqwmDpa0wM

 

https://pytorch.org/torchtune/stable/tutorials/lora_finetune.html

Finetuning Llama2 with LoRA — TorchTune documentation

https://arxiv.org/abs/2110.04366

Towards a Unified View of Parameter-Efficient Transfer Learning

https://underline.io/lecture/26070-prefix-tuning-optimizing-continous-prompts-for-generation

Prefix-Tuning: Optimizing Continuous Prompts for Generation

https://github.com/microsoft/LoRA/blob/main/loralib/layers.py

LoRA/loralib/layers.py at main · microsoft/LoRA

https://hi-lu.tistory.com/entry/%EA%B5%AC%ED%98%84%EC%B2%B4%EB%A5%BC-%ED%86%B5%ED%95%B4-PEFT-lora-%EB%A5%BC-%EC%95%8C%EC%95%84%EB%B3%B4%EC%9E%90

구현체를 통해 PEFT lora 를 알아보자

https://medium.com/data-science-in-your-pocket/lora-for-fine-tuning-llms-explained-with-codes-and-example-62a7ac5a3578

LoRA for Fine-Tuning LLMs explained with codes and example

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

LoRA를 통한 효율적인 미세 조정: 대규모 언어 모델을 위한 최적의 파라미터 선택 가이드