Fine-tuning a large model for a new task can demand more memory and computation than a project can afford. Parameter-efficient fine-tuning (PEFT) adapts a model while limiting how much must change or be trained. It is particularly useful when resources are constrained or when the resulting model must run on a limited device. section 5.3.4 section 5.3.1
LoRA is one way to do this. Rather than changing the original weights directly, it trains a compact update that changes the model’s output. This makes separate adaptations for several tasks or domains more feasible: each can have its own adapter instead of folding every change into the base model. section 5.3.5
What a low-rank update saves
Imagine a weight update represented by a 100 × 100 matrix: a full update has 10,000 values. LoRA factors the update into a 100 × 2 matrix and a 2 × 100 matrix, leaving 400 values to train in this rank-two illustration. Rank is a trade-off, not a free setting: increasing it permits a more complex update but saves less memory. section 5.3.5
Figure 5.7 LoRA exemplifies the idea that you should only need to train and save the difference between where the foundation model is and where you want it to be. It does this through singular value decomposition (SVD).
The diagram shows the distinction between retaining a foundation model and training the difference needed for an adaptation. section 5.3.5
PEFT is not limited to LoRA. Prompt tuning offers another route: it focuses on how a model responds to task-specific inputs while leaving most of the foundation model unchanged. The practical choice depends on the adaptation needed and the available memory and compute; neither approach removes the need to check performance on the intended task. section 5.3.1 section 5.3.4
Listing 5.6 Example LoRA and PEFT training
import os
from datasets import load_dataset
from transformers import (
AutoModelForTokenClassification,
AutoTokenizer,
DataCollatorForTokenClassification,
TrainingArguments,
Trainer,
)
from peft import (
PeftModel,
PeftConfig,
get_peft_model,
LoraConfig,
TaskType,
)
import evaluate
import torch
import numpy as np
model_checkpoint = "meta-llama/Llama-2-7b-hf"
lr = 1e-3
batch_size = 16
num_epochs = 10
model_dir = "./models/LoRAPEFT" #1
if not os.path.exists(model_dir):
os.makedirs(model_dir)
bionlp = load_dataset("tner/bionlp2004")
seqeval = evaluate.load("seqeval")
label_list = [
"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",
]
def compute_metrics(p):
predictions, labels = p
predictions = np.argmax(predictions, axis=2)
true_predictions = [
[label_list[p] for (p, l) in zip(prediction, label) if l != -100]
for prediction, label in zip(predictions, labels)
]
true_labels = [
[label_list[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"],
}
tokenizer = AutoTokenizer.from_pretrained(
model_checkpoint, add_prefix_space=True
)
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["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 None:
label_ids.append(-100)
elif 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
tokenized_bionlp = bionlp.map(tokenize_and_align_labels, batched=True)
data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
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,
}
model = AutoModelForTokenClassification.from_pretrained(
model_checkpoint, num_labels=11, id2label=id2label, label2id=label2id
)
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)
model.print_trainable_parameters()
training_args = TrainingArguments(
output_dir=model_dir,
learning_rate=lr,
per_device_train_batch_size=batch_size,
per_device_eval_batch_size=batch_size,
num_train_epochs=num_epochs,
weight_decay=0.01,
evaluation_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_bionlp["train"],
eval_dataset=tokenized_bionlp["validation"],
tokenizer=tokenizer,
data_collator=data_collator,
compute_metrics=compute_metrics,
)
trainer.train()
peft_model_id = "stevhliu/roberta-large-lora-token-classification"
config = PeftConfig.from_pretrained(model_dir)
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)
text = (
"The activation of IL-2 gene expression and NF-kappa B through CD28 "
"requires reactive oxygen production by 5-lipoxygenase."
)
inputs = tokenizer(text, return_tensors="pt")
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]))
The Python listing provides a concrete LoRA-and-PEFT training setup using the peft and transformers libraries for biomedical token classification. section 5.3.5
