Fine-tuning

Chuyển kiến thức nào sang nhiệm vụ mới?

Model pretrained gồm encoder và thường có một head phục vụ nhiệm vụ cũ. Encoder ảnh có thể học nét, cấu trúc và những tổ hợp hình dạng có ích cho nhiều nhãn khác. Nhưng head phân loại nhãn cũ không có lý do biết các nhãn mới của câu lạc bộ. Ta giữ encoder và thay head có đúng số lớp $K$ của bài mới. Cách làm này gọi là học chuyển giao (transfer learning).

Nó có ích nhất khi dữ liệu có nhãn ít mà nhiệm vụ nguồn đủ gần; ngay cả với 500 ảnh X-quang, bắt đầu từ trọng số pretrained rồi fine-tune vẫn thường tốt hơn học từ đầu.

Viết $\mathbf{z}=f_\theta(\mathbf{x})$ và logits $\boldsymbol\ell=g_\phi(\mathbf{z})$. $\theta$ là tham số encoder, $\phi$ là tham số head. Head-only training chỉ cập nhật $\phi$. Fine-tuning mở một phần hoặc toàn bộ $\theta$ cùng với $\phi$. Phần mở nhiều hơn có thể thích nghi tốt hơn nhưng cũng có thể học thuộc dữ liệu nhỏ hoặc làm mất biểu diễn có ích.

Trong mạng lớn, các block gần input thường bắt dấu hiệu chung như nét, cạnh, mảng màu; các block sau mang đặc trưng sát nhiệm vụ cũ hơn. Vì vậy thứ tự thử thực dụng là head trước, rồi mở thêm block cuối, rồi mới đến toàn bộ, đúng ba hàng của hình trên. Khi miền đích khác hẳn nguồn, như ảnh X-quang so với ảnh đời thường, đặc trưng ở các block đầu cũng có thể cần đổi; mỗi mức mở khi đó là một thí nghiệm, và baseline head-only là mốc để biết thay đổi có đáng giá.

Một mô hình nguồn tự pretrain

Để thấy mọi bước mà không cần tải trọng số, ta tự tạo một mô hình nguồn ngay trong notebook từ bộ chữ số $8\times8$ của scikit-learn. Nhiệm vụ nguồn phân loại các chữ số 0–7; nhiệm vụ đích phân biệt 8 và 9. Mọi ảnh 8 hoặc 9 đều bị loại khỏi việc học nguồn, nên encoder nguồn đóng đúng vai một checkpoint pretrained: đã học đặc trưng hữu ích trên dữ liệu khác và chưa từng thấy nhãn đích.

Với ResNet hay DistilBERT, các bước phía sau giống hệt; chỉ khác ở chỗ trọng số được nạp từ tệp checkpoint thay vì vừa học xong.

Một encoder 64 → 32 → ReLU đủ để thực nghiệm chạy trong vài giây. Pixel của bộ dữ liệu này nằm từ 0 đến 16, nên input được chia 16; mọi bước sau, kể cả nhiệm vụ đích, phải dùng đúng phép chia ấy, giống như fine-tune một checkpoint thật phải dùng đúng processor của nó (bài 4.1). Một phần ảnh nguồn được giữ làm validation để kiểm encoder thực sự đã học, thay vì gọi trọng số ngẫu nhiên là “đã học”.

import copy
import numpy as np
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.set_num_threads(1)
torch.manual_seed(3)
digits = load_digits()
source_ids = np.flatnonzero(digits.target < 8)
target_ids = np.flatnonzero(digits.target >= 8)
source_train, source_valid = train_test_split(
    source_ids, test_size=0.25, stratify=digits.target[source_ids], random_state=3
)

def inputs(ids):
    return torch.tensor(digits.data[ids] / 16.0, dtype=torch.float32)

source_encoder = nn.Sequential(nn.Linear(64, 32), nn.ReLU())
source_head = nn.Linear(32, 8)
source_model = nn.Sequential(source_encoder, source_head)
source_optimizer = torch.optim.Adam(source_model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
source_x = inputs(source_train)
source_y = torch.tensor(digits.target[source_train], dtype=torch.long)
for _ in range(80):
    source_model.train()
    source_optimizer.zero_grad(set_to_none=True)
    loss = criterion(source_model(source_x), source_y)
    loss.backward()
    source_optimizer.step()

source_model.eval()
with torch.no_grad():
    source_prediction = source_model(inputs(source_valid)).argmax(dim=1).numpy()
print("source train/valid:", len(source_train), len(source_valid))
print("source valid accuracy:", round(float((source_prediction == digits.target[source_valid]).mean()), 4))
print("source gap target:", len(set(source_ids) & set(target_ids)) == 0)

Kết quảsource train/valid: 1082 361
source valid accuracy: 0.9806
source gap target: True

Accuracy 0.9806 trên 361 ảnh nguồn giữ lại cho thấy encoder đã học được đặc trưng chữ số có ích. Dòng thứ ba xác nhận không ảnh 8 hay 9 nào lọt vào việc học nguồn; nếu có, điểm của nhiệm vụ đích sau này sẽ đo cả phần encoder đã thuộc lòng, không chỉ phần kiến thức được chuyển.

Split ở đây chia theo từng ảnh; với dữ liệu có nhiều ảnh của cùng một người hay cùng một buổi, split phải chia theo nhóm ấy (bài 1.3).

Bắt đầu bằng một head mới

Ta dành 48 ảnh đích làm train có nhãn; các ảnh đích còn lại làm validation. Nhãn 8 đổi thành 0, nhãn 9 đổi thành 1 để khớp head hai lớp. Việc đổi nhãn phải lưu lại để trả output đúng tên lớp sau khi dự đoán.

target_train, target_valid = train_test_split(
    target_ids, train_size=48, random_state=5, stratify=digits.target[target_ids]
)
target_x = inputs(target_train)
target_y = torch.tensor(digits.target[target_train] - 8, dtype=torch.long)
valid_x = inputs(target_valid)
valid_y = torch.tensor(digits.target[target_valid] - 8, dtype=torch.long)

torch.manual_seed(10)
initial_head = nn.Linear(32, 2)
head_only = nn.Sequential(copy.deepcopy(source_encoder), copy.deepcopy(initial_head))
for parameter in head_only[0].parameters():
    parameter.requires_grad_(False)
optimizer = torch.optim.SGD(head_only[1].parameters(), lr=0.05)
before = {name: value.detach().clone() for name, value in head_only[0].named_parameters()}

head_only.train()
head_only[0].eval()
optimizer.zero_grad(set_to_none=True)
loss = criterion(head_only(target_x), target_y)
loss.backward()
print("encoder grad None:", all(p.grad is None for p in head_only[0].parameters()))
print("head co gradient:", head_only[1].weight.grad.norm().item() > 0)
optimizer.step()
print("encoder khong doi:", all(torch.equal(before[n], p) for n, p in head_only[0].named_parameters()))
print("target train/valid:", len(target_train), len(target_valid))

Kết quảencoder grad None: True
head co gradient: True
encoder khong doi: True
target train/valid: 48 306

Ở mạng này, encoder chỉ có Linear và ReLU; eval() chưa đổi phép tính của nó. Ta vẫn gọi để tạo thói quen đúng khi sau này dùng encoder có Dropout hoặc BatchNorm.

Dòng encoder grad None kiểm việc không tạo gradient cho tham số đã đóng băng. Dòng thứ ba so từng giá trị tensor trước và sau một bước, kiểm việc optimizer không sửa encoder. Nếu chỉ nhìn loss giảm, ta không biết phần nào đã được cập nhật.

Đóng băng nghĩa là đặt requires_grad_(False) cho tham số encoder: chúng không nhận gradient và optimizer không sửa. Số tham số được học khi đó chỉ là của head, đếm bằng sum(p.numel() for p in model.parameters() if p.requires_grad): ResNet-18 đóng băng hết rồi thay fc bằng Linear(512, 10) còn $512 \cdot 10 + 10 = 5130$; ResNet-50 với head $2048 \to K$ là $2048K + K$.

Vì sao encoder dùng bước nhỏ hơn

Xét mô hình một chiều $\hat y=a(wx)$, encoder có tham số $w$, head có tham số $a$, loss $L=(\hat y-y)^2$. Với $x=1$, $w=2$, $a=1$, $y=3$, ta có $\hat y=2$, phần dư $-1$, nên

Nếu learning rate head là 0.1, encoder là 0.01, cập nhật cho $a=1.4$ và $w=2.02$. Nếu encoder đóng băng thì $w$ vẫn bằng 2.

x, w, a, y = 1.0, 2.0, 1.0, 3.0
residual = a * w * x - y
grad_a = 2 * residual * w * x
grad_w = 2 * residual * a * x
print("gradient:", grad_a, grad_w)
print("sau buoc:", a - 0.1 * grad_a, w - 0.01 * grad_w)

Kết quảgradient: -4.0 -2.0
sau buoc: 1.4 2.02

Để ý $\partial L/\partial w$ chứa $a$: gradient về encoder phải đi qua head. Head mới khởi tạo ngẫu nhiên, nên trong những bước đầu, hướng gradient mà encoder nhận được cũng mang theo sự ngẫu nhiên đó. Nếu encoder đi bước lớn như head, nó có thể bị kéo xa chỉ để bù cho một head còn sai, làm hỏng những đặc trưng tốt đã học. Việc mô hình mất dần kiến thức nguồn khi được huấn luyện tiếp cho nhiệm vụ mới gọi là quên thảm khốc (catastrophic forgetting); bước lớn ở phần pretrained làm hiện tượng này đến nhanh hơn.

Hai cách hạn chế: learning rate nhỏ cho encoder, hoặc học head-only vài epoch rồi mới mở encoder. Tỉ lệ 10 lần ở ví dụ là một lựa chọn, không phải hằng số. Với BERT, bài báo gốc fine-tune toàn bộ mô hình với learning rate từ $2\cdot10^{-5}$ đến $5\cdot10^{-5}$ trong 2–4 epoch, nhỏ hơn mức $10^{-4}$ dùng khi pretrain chính BERT.

So sánh head-only và mở encoder

Hai thực nghiệm bắt đầu từ cùng encoder nguồn và cùng head mới, nên khác biệt chỉ nằm ở việc mở encoder. Ta chạy 60 epoch full-batch, chấm validation sau mỗi epoch và giữ checkpoint có validation loss thấp nhất. Head dùng learning rate 0.05; khi mở, encoder có nhóm riêng với 0.003, nhỏ hơn khoảng 17 lần. Vì validation vừa dùng để chọn epoch vừa để so hai mode, con số báo cáo cuối cùng nên lấy từ một test giữ kín.

def run_transfer(open_encoder):
    transferred = nn.Sequential(copy.deepcopy(source_encoder), copy.deepcopy(initial_head))
    for parameter in transferred[0].parameters():
        parameter.requires_grad_(open_encoder)
    groups = [{"params": transferred[1].parameters(), "lr": 0.05}]
    if open_encoder:
        groups.append({"params": transferred[0].parameters(), "lr": 0.003})
    optimizer = torch.optim.Adam(groups)
    best_loss = float("inf")
    best_state, best_epoch = None, None
    history = []
    for epoch in range(1, 61):
        transferred.train()
        if not open_encoder:
            transferred[0].eval()
        optimizer.zero_grad(set_to_none=True)
        train_loss = criterion(transferred(target_x), target_y)
        train_loss.backward()
        optimizer.step()
        transferred.eval()
        with torch.no_grad():
            val_loss = criterion(transferred(valid_x), valid_y).item()
        history.append((train_loss.item(), val_loss))
        if val_loss < best_loss:
            best_loss = val_loss
            best_epoch = epoch
            best_state = copy.deepcopy(transferred.state_dict())
    transferred.load_state_dict(best_state)
    with torch.no_grad():
        accuracy = (transferred(valid_x).argmax(dim=1) == valid_y).float().mean().item()
    trainable = sum(p.numel() for p in transferred.parameters() if p.requires_grad)
    return transferred, history, best_epoch, best_loss, accuracy, trainable

for opened in [False, True]:
    net, history, epoch, val_loss, accuracy, count = run_transfer(opened)
    print("open encoder =", opened, "| trainable =", count, "| best epoch =", epoch)
    print("best val loss =", round(val_loss, 4),
          "| epoch 60 val loss =", round(history[-1][1], 4),
          "| val accuracy =", round(accuracy, 4))

Kết quảopen encoder = False | trainable = 66 | best epoch = 36
best val loss = 0.2247 | epoch 60 val loss = 0.2366 | val accuracy = 0.902
open encoder = True | trainable = 2146 | best epoch = 60
best val loss = 0.157 | epoch 60 val loss = 0.157 | val accuracy = 0.9379

Đọc loss và accuracy ở checkpoint được chọn, không chỉ ở epoch cuối:

• Head-only học 66 tham số ($32\cdot2+2$); validation loss thấp nhất 0.2247 ở epoch 36, đến epoch 60 là 0.2366 trong khi train loss vẫn giảm. Giữ checkpoint epoch 36 nghĩa là không phải trả giá cho 24 epoch học thêm vô ích.
• Mở encoder học 2146 tham số (thêm $64\cdot32+32=2080$ của encoder), loss 0.1570 và accuracy 0.9379, so với 0.9020 của head-only: đúng thêm 11 trong 306 ảnh validation. Checkpoint tốt nhất của nó rơi đúng epoch cuối, tức validation còn đang giảm; ngân sách 60 epoch có thể chưa đủ, và đó là tín hiệu để chạy dài hơn, không phải bằng chứng đã hội tụ.

Trên một split nhỏ, chênh lệch 11 ảnh là bằng chứng có ích, chưa phải định luật “mở encoder luôn thắng”. Loss và accuracy cũng có thể chọn hai epoch khác nhau: cross-entropy còn phạt sự tự tin sai, còn accuracy chỉ nhìn argmax.

Ở đây encoder chỉ có một tầng Linear nên mở nó là mở toàn bộ; với ResNet, “mở block cuối” là một tập tham số riêng, có số lượng và bộ nhớ riêng phải đo lại. Mỗi tham số được mở kéo theo một gradient và, với Adam, hai moment $m$, $v$: thêm ba số float32 cho mỗi tham số, chưa kể activation phải giữ để backward qua phần mở (bài 3.7). Với 66 tham số, phần này không đáng kể; với hàng chục triệu tham số, nó lên tới hàng trăm MB.

Script của chương đưa đúng thí nghiệm này lên DistilBERT với model và dữ liệu lưu sẵn trên máy. Mode head đóng băng backbone và giữ nó ở eval; mode full học backbone với learning rate $10^{-5}$ và head với $10^{-3}$. Cả hai chọn checkpoint theo macro-F1 validation và ghi số tham số học, lịch sử validation, thời gian vào RESULT.json; so hai tệp ấy giống như so hai cặp dòng output trên.

Freeze, eval và BatchNorm là ba chuyện

BatchNorm (bài 3.6) có tham số affine như weight và bias, đồng thời có buffer running_mean, running_var. Trong train mode, batch mới có thể cập nhật buffer dù không tính gradient; chuyển sang eval thì dùng buffer đã tích lũy. Với encoder cần giữ cố định, ta giữ cả buffer bằng eval.

Khi batch fine-tune nhỏ, như vài ảnh độ phân giải cao mỗi batch, thống kê batch rất nhiễu, nên nhiều hệ thống phát hiện vật thể giữ BatchNorm của backbone cố định ngay cả khi các trọng số khác được học.

Thí nghiệm sau đặt mọi tham số BatchNorm thành không học, nhưng vẫn cho một batch qua train mode. Trung bình batch là $(3,6)$, momentum=0.5, trung bình cũ là $(0,0)$, nên trung bình mới là $(1.5,3)$. Trong BatchNorm, momentum này là hệ số cập nhật thống kê, khác nghĩa với momentum của optimizer.

bn = nn.BatchNorm1d(2, momentum=0.5)
for parameter in bn.parameters():
    parameter.requires_grad_(False)
batch = torch.tensor([[2., 4.], [4., 8.]])
bn.train()
with torch.no_grad():
    bn(batch)
after_train = bn.running_mean.clone()
bn.eval()
with torch.no_grad():
    bn(batch + 100)
print("running mean sau train:", after_train.tolist())
print("eval giu buffer:", torch.equal(after_train, bn.running_mean))
print("khop tinh tay:", torch.allclose(after_train, batch.mean(dim=0) * 0.5))

Kết quảrunning mean sau train: [1.5, 3.0]
eval giu buffer: True
khop tinh tay: True

Dòng đầu trực tiếp phản bác hiểu lầm “không có gradient thì không có gì thay đổi”. Dropout lại không có running mean; eval làm nó bỏ thao tác che ngẫu nhiên. Khi gọi model.train() ở đầu epoch, các module con cũng chuyển train; nếu encoder phải đứng yên, gọi encoder.eval() sau đó. Còn no_grad() không được bọc phần đang học (head, adapter): autograd phải ghi lại các phép tính ấy thì backward mới tới được tham số của chúng.

encoder.eval()            # gọi một lần trước vòng lặp
for epoch in range(10):
    model.train()
    for x, y in loader:
        ...               # forward, backward, step cho head