Using Transformers for Classification
Fine-tuning in practice.
Earlier, we built a strong classifier with TF-IDF. Let's now see what a transformer can do. Remember our rule. Never assume that the complex model wins. Measure it.
- TF-IDF was strong
- Now try a transformer
- Measure, don't assume
BERT features and a linear head
The quickest approach. Use BERT as a frozen reader, and take the first token's vector as the summary of the text. Then train one linear layer on top. We use only six hundred posts. Run it. A strong score with very little data, and training took seconds, since BERT's weights stayed frozen. For even better results, you can unfreeze the model, and fine tune it, using a library like transformers' Trainer, on a GPU.
import torch
from torch import nn
from sklearn.datasets import fetch_20newsgroups
from transformers import AutoModel, AutoTokenizer
cats = ["rec.autos", "sci.med", "sci.space", "comp.graphics"]
data = fetch_20newsgroups(subset="all", categories=cats,
remove=("headers", "footers", "quotes"))
texts, y = [t[:300] for t in data.data[:800]], torch.tensor(data.target[:800])
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, 800, 32)])
head = nn.Linear(768, 4)
opt = torch.optim.AdamW(head.parameters(), lr=0.01)
for _ in range(150):
loss = nn.functional.cross_entropy(head(X[:600]), y[:600])
opt.zero_grad(); loss.backward(); opt.step()
print("Trained on 600 posts")
print("Accuracy:", (head(X[600:]).argmax(1) == y[600:]).float().mean().item())Trained on 600 posts Accuracy: 0.8149999976158142 [transformers] [1mDistilBertModel LOAD REPORT[0m from: distilbert-base-uncased Key | Status | | ------------------------+------------+--+- vocab_layer_norm.bias | UNEXPECTED | | vocab_projector.bias | UNEXPECTED | | vocab_layer_norm.weight | UNEXPECTED | | vocab_transform.weight | UNEXPECTED | | vocab_transform.bias | UNEXPECTED | | Notes: - UNEXPECTED: can be ignored when loading from different task/architecture; not ok if you expect identical arch.
Which tool for classification?
Here's how to choose. With a few examples and a small budget, use sentence embeddings with logistic regression. With thousands of labelled examples, fine tune a transformer. If you need the cheapest and fastest, use TF-IDF. And with almost no labels at all, use a large language model, with a good prompt.
| Situation | Use |
|---|---|
| Few labels | Embeddings + logistic |
| Thousands of labels | Fine-tune a transformer |
| Very high volume, cheap | TF-IDF + logistic |
| No labels, quick start | LLM with a prompt |
Summary
Frozen BERT gives you a fast, strong classifier. And you know how to pick the right tool. Next, let's give Scribe new powers, using ready trained models for extraction, questions, summaries, and translation.
Key takeaways- BERT as a frozen reader
- Train only the head
- Pick by data and budget