Fine-Tuning a Transformer
Transfer learning for language.
Your client wants a tool that sorts documents into categories. You've got seven hundred labelled examples. That's nowhere near enough to train a language model from scratch. But it's plenty, if you stand on the shoulders of a pretrained transformer.
- Seven hundred examples
- Not enough from scratch
- Plenty with transfer learning
A classifier from a transformer
Same recipe as with images. We take a pretrained transformer, and freeze it. We run our documents through it, keeping the first token's vector as a summary of each one. Then we train a tiny new layer, on those vectors. Run it. On held-out posts, accuracy is very strong, with only seven hundred examples, and a few seconds of training. For real projects, you'd also unfreeze the top layers. But this approach is a great, cheap start.
import torch
from torch import nn
from sklearn.datasets import fetch_20newsgroups
from transformers import AutoModel, AutoTokenizer
kinds = ["sci.space", "sci.med"]
data = fetch_20newsgroups(subset="all", categories=kinds,
remove=("headers", "footers", "quotes"))
texts = [t[:300] for t in data.data[:700]]
labels = torch.tensor(data.target[:700])
name = "distilbert-base-uncased"
tok = AutoTokenizer.from_pretrained(name)
bert = AutoModel.from_pretrained(name).eval()
def embed(batch):
enc = tok(batch, padding=True, truncation=True, max_length=64, return_tensors="pt")
with torch.no_grad():
return bert(**enc).last_hidden_state[:, 0]
X = torch.cat([embed(texts[i:i + 32]) for i in range(0, 700, 32)])
tr, te = X[:550], X[550:]
y_tr, y_te = labels[:550], labels[550:]
torch.manual_seed(0)
head = nn.Linear(768, 2)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2)
for _ in range(100):
loss = nn.functional.cross_entropy(head(tr), y_tr)
opt.zero_grad(); loss.backward(); opt.step()
print("Accuracy:", (head(te).argmax(1) == y_te).float().mean().item())Accuracy: 0.9133333563804626 [transformers] [1mDistilBertModel LOAD REPORT[0m from: distilbert-base-uncased Key | Status | | ------------------------+------------+--+- vocab_layer_norm.bias | UNEXPECTED | | vocab_transform.weight | UNEXPECTED | | vocab_transform.bias | UNEXPECTED | | vocab_projector.bias | UNEXPECTED | | vocab_layer_norm.weight | UNEXPECTED | | Notes: - UNEXPECTED: can be ignored when loading from different task/architecture; not ok if you expect identical arch.
LoRA: fine-tuning on a budget
Large models have huge layers, so fine tuning all of the weights is costly. LoRA has a clever idea. Freeze the original layer, and add a tiny pair of small matrices next to it. Only these are trained. Their product acts as a small correction. Run it. The full layer has about sixteen million weights, and LoRA trains only a few tens of thousands, a tiny fraction. That's why anyone can customise huge language models and image generators today.
import torch
from torch import nn
class LoRA(nn.Module):
def __init__(self, layer, rank=4):
super().__init__()
self.layer = layer
for p in layer.parameters():
p.requires_grad = False
self.a = nn.Parameter(torch.randn(rank, layer.in_features) * 0.01)
self.b = nn.Parameter(torch.zeros(layer.out_features, rank))
def forward(self, x):
return self.layer(x) + x @ self.a.T @ self.b.T
big = nn.Linear(4096, 4096)
lora = LoRA(big)
full = sum(p.numel() for p in big.parameters())
small = lora.a.numel() + lora.b.numel()
print(f"Full layer weights: {full:,}")
print(f"LoRA weights: {small:,}")
print(f"Share to train: {small / full:.2%}")Full layer weights: 16,781,312 LoRA weights: 32,768 Share to train: 0.20%
Summary
Freeze, replace the head, and train a little. Or use LoRA to change a huge model cheaply. Those skills cover a large share of real language projects. Now let's turn back to images, and teach Aria to do far more than simply naming them. Next, finding objects.
Key takeaways- Pretrained body, new head
- Few examples are enough
- LoRA trains tiny corrections