Unit 09.02: Using pretrained models for feature extraction
The cheapest useful thing you can do with a pretrained model: run it once, keep the vectors, and fit something small on top.
Embed once, reuse many times
The encoder is frozen and switched to eval mode. Every document is passed through once, pooled to a single vector, and those vectors are cached.
From there a logistic regression — or any classical model — trains in milliseconds:
import torch
from torch import nn
from sklearn.linear_model import LogisticRegression
torch.manual_seed(0)
# Stand-in for a pretrained encoder: frozen, used only to produce features.
encoder = nn.TransformerEncoderLayer(d_model=16, nhead=4, dim_feedforward=32,
batch_first=True)
for p in encoder.parameters():
p.requires_grad = False
encoder.eval()
tokens = torch.randn(120, 6, 16)
labels = (tokens[:, 0, 0] > 0).long()
with torch.no_grad(): # no gradients needed at all
features = encoder(tokens).mean(dim=1) # mean-pool over tokens
print("features:", tuple(features.shape), "one vector per document")
clf = LogisticRegression(max_iter=1000).fit(features[:90].numpy(), labels[:90].numpy())
acc = clf.score(features[90:].numpy(), labels[90:].numpy())
print("logistic head on frozen features:", round(float(acc), 3))
print("trainable transformer parameters :",
sum(p.numel() for p in encoder.parameters() if p.requires_grad))
# This is the cheapest useful thing you can do with a pretrained model: embed
# once, cache the vectors, and fit a small classical model on top.
Zero trainable transformer parameters, and a working classifier. The entire cost was one forward pass per document, done inside torch.no_grad() so no graph was built.
This approach scales well operationally: the expensive step happens once, and downstream experiments are cheap enough to iterate on freely.
The mistake this prevents
Re-encoding the same documents for every experiment. The embedding pass dominates the runtime, and caching it turns a slow loop into a fast one.
Takeaway
Freeze, embed once, cache the vectors, and iterate on a cheap head.
