LoRA

LoRA (Low Rank Adaption for Large Language Model) 는 아래 논문에서 처음 발표되었습니다.

https://arxiv.org/abs/2106.09685

LoRA: Low-Rank Adaptation of Large Language Models

 

LoRA는 LLM fine tuning 시, 전체 파라미터를 조정하는 것이 아닌, 일부 파라미터 조정 후 기존 LLM 모델과 합치는 방식을 제안합니다.

 

논문 그림에도 나와 있듯, "We only train A and B" 가 해당 논문의 가장 핵심 문구입니다. 기존 LLM 모델 (파란 부분) 은 그대로 두고, fine-tuning을 위한 새로운 부분 A,B만 따로 학습 후 기존 LLM과 합치겠다. 라는 의미입니다.

(구체적으로 왜 되는 지는 잘 모릅니다...)

 

huggingface의 LoRA 튜토리얼을 따라 bionlp2004 라는 데이터셋을 이용해 text 내에서 단백질, dna 등을 의미하는 것을 분류하도록 해보겠습니다.

https://huggingface.co/docs/peft/main/en/task_guides/token-classification-lora

LoRA for token classification

 

bionlp 2004 데이터셋 라벨은 다음과 같습니다.

"O", "B-DNA", "I-DNA", "B-protein", "I-protein", "B-cell_type", "I-cell_type", "B-cell_line", "I-cell_line", "B-RNA", "I-RNA"

 

생물학 전공이 아니다 보니 무엇인지는 잘 모르겠지만, gpt에 따르면 다음과 같습니다.

(즉 "O"는 분류 대상이 아닌 것을 의미하고, 나머지 라벨은 특별한 의미를 갖는 것 같습니다.)

 

Env : Google Colab

 

1. 필요 패키지 설치

!pip install -q peft transformers datasets evaluate seqeval

 

2. 필요 패키지 import

from datasets import load_dataset
from transformers import AutoModelForTokenClassification, AutoTokenizer, DataCollatorForTokenClassification, TrainingArguments, Trainer
from peft import get_peft_config, PeftModel, PeftConfig, get_peft_model, LoraConfig, TaskType
import evaluate
import torch
import numpy as np

 

- datasets : bionlp2004 데이터셋을 불러오기 용

 

 - transformers : 대부분 LLM에서 따르는 트랜스포머 모델

--- AutoModelForTokenClassification : classification (분류) 모델 

--- AutoTokenizer : 데이터셋을 토큰화하기 위한 용도

--- DataCollatorForTokenClassification : tokenizer에 대해 padding, batch 작업 

--- TrainingArguments : training 을 위한 파라미터 설정

--- Trainer : training 도구

 

- peft : PEFT 작업 위한 패키지

--- get_peft_config : peft 작업 위한 설정값 세팅 (PeftConfig)

--- PeftModel : peft 모델 생성용

--- PeftConfig : peft 모델에 필요한 설정값 

--- get_peft_model : 학습 완료된 모델 불러오기

--- LoraConfig : LoRA 모델에 대한 설정값 세팅

--- TaskType : 어떤 작업할 것인지 (LoraConfig)

 

- evaluate : precision, recall, f1 score, accuracy 평가용

 

- torch : 학습용 딥러닝 프레임워크 (huggingface에서 만든 transformers, peft 등 라이브러리가 pytorch 기반으로 동작)

 

3. 모델 및 파라미터 세팅

model = "roberta-base"
lr = 1e-3
batchSize = 16
epoches = 3

 

huggingface 튜토리얼에서는 roberta-large 모델을 사용하였고, epoch을 10으로 잡았지만, 학습 속도를 비교적 빠르게 테스트하기 위해 base 모델을 사용하였고, epoch을 3으로 잡았습니다.

더 좋은 성능을 사용하고 싶으면 large 등 다른 모델을 사용하거나 epoch을 키우면 됩니다.

learning rate (lr) 과 batch size는 튜토리얼과 동일한 값으로 사용하였습니다.

 

 4. bionlp 2004 데이터셋 불러오기

bioNlp = load_dataset("tner/bionlp2004")

 

bioNlp["train"][0] 로 학습용 데이터 중 첫 번째 값 확인 시 다음과 같습니다.

HUVECs 는 label 7, VCAM-1 은 label 3, 나머지는 모두 label 0 입니다. 

구성된 token으로 보아, 데이터셋은 text의 띄어쓰기, '.', ',' 를 기준으로 토큰화가 이루어진 것을 알 수 있습니다.

 

5. sequence 라벨링 평가에 사용될 모듈 생성 ( precision, recall, f1 score, accuracy 평가)

seqEval = evaluate.load("seqeval")

 

6. 데이터셋의 라벨값 분류

labelList = [
    "O",
    "B-DNA",
    "I-DNA",
    "B-protein",
    "I-protein",
    "B-cell_type",
    "I-cell_type",
    "B-cell_line",
    "I-cell_line",
    "B-RNA",
    "I-RNA",
]

 

7. precision, recall, f1 score, accuracy 측정을 위한 metrics

def compute_metrics(p) :
  predictions, labels = p
  predictions = np.argmax(predictions, axis=2)

  true_predictions = [
      [ labelList[p] for (p, l) in zip(prediction, label) if l != -100 ]
      for prediction, label in zip(predictions, labels)
  ]

  true_labels = [
      [ labelList[l] for (p, l) in zip(prediction, label) if l != -100 ]
      for prediction, label in zip(predictions, labels)
  ]

  results = seqEval.compute(predictions = true_predictions, references = true_labels)

  return {
      "precision" : results["overall_precision"],
      "recall" : results["overall_recall"],
      "f1" : results["overall_f1"],
      "accuracy" : results["overall_accuracy"]
  }

 

8. tokenizer 생성

tokenizer = AutoTokenizer.from_pretrained(model, add_prefix_space=True)

 

9. 데이터셋을 tokenizer에 토큰화로 적용하기 위한 함수

def tokenize_and_align_labels(examples) :
  tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
  labels = []
  for i, label in enumerate(examples[f"tags"]) :
    word_ids = tokenized_inputs.word_ids(batch_index = i)
    previous_word_idx = None
    label_ids = []

    for word_idx in word_ids :
      if word_idx is not None and word_idx != previous_word_idx :
        label_ids.append(label[word_idx])
      else :
        label_ids.append(-100)
      previous_word_idx = word_idx
    labels.append(label_ids)

  tokenized_inputs["labels"] = labels

  return tokenized_inputs

 

10. bionlp 2004 데이터셋 tokenization

tokenizedBioNlp = bioNlp.map(tokenize_and_align_labels, batched=True)

 

11. tokenization 된 데이터셋을 학습에 맞게 padding 등 진행

data_collator = DataCollatorForTokenClassification(tokenizer = tokenizer)

 

12. 기존 데이터셋 라벨링에 index 부여 

id2label = {
    0: "O",
    1: "B-DNA",
    2: "I-DNA",
    3: "B-protein",
    4: "I-protein",
    5: "B-cell_type",
    6: "I-cell_type",
    7: "B-cell_line",
    8: "I-cell_line",
    9: "B-RNA",
    10: "I-RNA",
}
label2id = {
    "O": 0,
    "B-DNA": 1,
    "I-DNA": 2,
    "B-protein": 3,
    "I-protein": 4,
    "B-cell_type": 5,
    "I-cell_type": 6,
    "B-cell_line": 7,
    "I-cell_line": 8,
    "B-RNA": 9,
    "I-RNA": 10,
}

 

13. classification 모델 생성

model = AutoModelForTokenClassification.from_pretrained(model, num_labels = 11, id2label = id2label, label2id = label2id)

 

14. peft 모델 설정

peft_config = LoraConfig(task_type = TaskType.TOKEN_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias="all")
model = get_peft_model(model, peft_config)

 

 

단순 classification 이므로 답변 시 필요한 inference 는 false

r, alpha 는 LoRA 모델에서 사용할 파라미터

 

model.print_trainable_parameters() 로 기존 roberta-base 모델에 비해 fine-tuning 해야 할 파리미터 비교 시 다음과 같습니다.

기존 roberta-base 모델의 경우, 124,661,782 개의 파라미터를 갖지만, LoRA를 이용한 fine-tuning에 사용될 parameter는 700, 427 개로, 전체 파라미터의 0.5 % 만 fine-tuning 시킬 수 있습니다.

 

15. train config 설정

training_args = TrainingArguments(
    output_dir = "roberta-base-lora-token-classification",
    learning_rate = lr,
    per_device_train_batch_size = batchSize,
    per_device_eval_batch_size = batchSize,
    num_train_epochs = epoches,
    weight_decay = 0.01,
    evaluation_strategy = "epoch",
    save_strategy = "epoch",
    load_best_model_at_end = True,
)

 

16. train 시작

trainer = Trainer(
    model = model,
    args = training_args,
    train_dataset = tokenizedBioNlp["train"],
    eval_dataset = tokenizedBioNlp["validation"],
    tokenizer = tokenizer,
    data_collator = data_collator,
    compute_metrics = compute_metrics,
)

trainer.train()

 

학습이 완료된 후, colab 왼쪽 파일 아이콘 클릭 시, training_args 에서 output_dir 에 설정한 디렉토리 아래에 학습된 모델이 기록되어 있습니다.

해당 checkpoint 중 가장 마지막 checkpoint 를 기준으로 이후 학습된 모델을 불러와 테스트를 진행합니다.

 

17. 학습된 모델 불러오기

peft_model_id = "./roberta-base-lora-token-classification/checkpoint-3117"
config = PeftConfig.from_pretrained(peft_model_id)
inference_model = AutoModelForTokenClassification.from_pretrained(
    config.base_model_name_or_path, num_labels=11, id2label=id2label, label2id=label2id
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
model = PeftModel.from_pretrained(inference_model, peft_model_id)

 

여기서 학습 진행 시점에 따른 오픈 소스 모델 업데이트, 파라미터 수에 따라 checkpoint 값이 다르므로 확인 후 해당 값을 입력하면 됩니다. 

 

18. 테스트용 text

test_text = "The activation of IL-2 gene expression and NF-kappa B through CD28 requires reactive oxygen production by 5-lipoxygenase."
inputs = tokenizer(test_text, return_tensors="pt")

 

19. 테스트

with torch.no_grad():
    logits = model(**inputs).logits

tokens = inputs.tokens()
predictions = torch.argmax(logits, dim=2)

for token, prediction in zip(tokens, predictions[0].numpy()):
    print((token, model.config.id2label[prediction]))

 

20. 테스트 결과

위 그림은 제가 테스트한 결과, 아래는 huggingface 튜토리얼 페이지 결과입니다.

 

<s>, </s> 는 tokenization을 진행하면서 생긴 구분자이므로 크게 신경을 안 써도 됩니다. 

위 test_text 를 대상으로 띄어쓰기 부분 대신 G' 과 같은 문자가 추가되었습니다. (인코딩 문제인 듯)

해당 값을 무시하고 판단하면, 분류 시 "O" 로 분류된 값 외에는 huggingface 튜토리얼과 거의 유사한 값을 가지게 되었습니다. 

 

21. 결론

저렇게 분류한 것이 옳게 분류가 된 것인지는 생물학 전공이 아니다 보니 잘 모르겠지만, huggingface 튜토리얼과 거의 흡사한 값을 갖게 되었으므로 roberta-large 대신 roberta-base 모델, 에폭 수 작게 등 빠르게 학습시켜도 어느 정도 효과를 갖는 것으로 판단됩니다.

 

따라서 이후 본인이 원하는 데이터셋으로 fine-tuning을 할 때,

해당 데이터셋을 peft에 적용 가능 하도록 전처리만 진행한다면 (이것도 매우 힘든 작업일 듯...) 

peft에서 제공하는 라이브러리를 바탕으로 LoRA 등 여러 fine-tuning 기법을 몇 가지 config 설정만으로 진행할 수 있습니다.