Model Soups

Model soup là một model duy nhất được tạo bằng cách lấy trung bình trọng số của một số model đã được fine-tune từ cùng một model tiền huấn luyện với các siêu tham số khác nhau. Thông thường, sau khi quét siêu tham số, người ta giữ lại lần chạy có điểm validation tốt nhất và loại bỏ phần còn lại. Soup tái sử dụng những lần chạy bị loại bỏ đó: vì mỗi lần đều bắt đầu từ cùng một trọng số model tiền huấn luyện, các tham số vẫn tương ứng với nhau; giá trị trung bình của chúng thường đạt điểm cao hơn bất kỳ lần chạy riêng lẻ nào, trong khi chi phí chạy đúng bằng một model. Mitchell Wortsman cùng các cộng sự tại Đại học Washington, Google, Meta AI, Đại học Tel Aviv và Đại học Columbia đã giới thiệu phương pháp này trong bài Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, được công bố tại ICML 2022. Một greedy soup gồm 58 model ViT-G/14 đã fine-tune, tiền huấn luyện trên dataset JFT-3B của Google, đạt độ chính xác top-1 90.94% trên ImageNet, một kết quả tiên tiến nhất vào thời điểm đó.

Soup đồng nhất và soup tham lam#

Fine-tuning chỉ dịch chuyển một mạng tiền huấn luyện lớn một đoạn ngắn trong không gian trọng số, vì vậy các model được fine-tune từ mạng đó với learning rate, tăng cường dữ liệu, làm mượt nhãn hoặc thời lượng huấn luyện khác nhau thường nằm trong cùng một vùng rộng có loss thấp. Lấy trung bình các model sẽ đưa kết quả về gần trung tâm vùng đó, giúp tổng quát hóa tốt hơn các điểm ở rìa. Bài báo định nghĩa hai công thức chính:

• Soup đồng nhất: Lấy trung bình mọi model đã fine-tune. Phương pháp này hiệu quả khi mọi thành phần đều tốt, nhưng một lần chạy có learning rate rất cao có thể rời khỏi vùng chung; đưa model đó vào phép tính sẽ làm giảm độ chính xác.
• Soup tham lam: Sắp xếp model theo độ chính xác trên dữ liệu validation giữ lại, bắt đầu từ model tốt nhất và chỉ thêm từng model tiếp theo nếu độ chính xác validation của soup sau khi lấy trung bình không giảm. Vì bắt đầu từ model tốt nhất, soup tham lam không bao giờ kém model đó trên tập validation.

Với CLIP ViT-B/32 được fine-tune trên ImageNet bằng 72 cấu hình siêu tham số ngẫu nhiên, model đơn lẻ tốt nhất đạt 80.38%. Soup đồng nhất gồm cả 72 model đạt 79.97% vì một số lần chạy đã rời khỏi vùng loss thấp chung, trong khi soup tham lam giữ lại năm model và đạt 81.03%. Trên năm tập kiểm thử dịch chuyển phân phối, chẳng hạn ảnh phác thảo và ảnh dựng của ImageNet, soup tham lam đạt trung bình 50.75%, so với 47.83% của model đơn lẻ tốt nhất. Bài báo cũng ghi nhận mức cải thiện với ALIGN và mức cải thiện nhỏ hơn khi fine-tune BERT và T5 cho phân loại văn bản.

Ví dụ PyTorch này tiền huấn luyện một bộ phân loại nhỏ trên chữ số viết tay 8x8, fine-tune chín bản sao với learning rate và tăng cường nhiễu đầu vào khác nhau, rồi so sánh soup đồng nhất của chúng với ensemble gồm cả chín model:

import copy

import torch
from sklearn.datasets import load_digits
from torch import nn

torch.manual_seed(0)
X, y = load_digits(return_X_y=True)
order = torch.randperm(len(X))
X, y = torch.tensor(X / 16, dtype=torch.float32)[order], torch.tensor(y)[order]


def train(model, first, last, lr, noise, seed):
    """Train a copy of the model on random batches from rows first to last, adding input noise."""
    model, generator = copy.deepcopy(model), torch.Generator().manual_seed(seed)
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    for _ in range(300):
        rows = torch.randint(first, last, (32,), generator=generator)
        optimizer.zero_grad()
        inputs = X[rows] + noise * torch.randn(32, 64, generator=generator)
        nn.functional.cross_entropy(model(inputs), y[rows]).backward()
        optimizer.step()
    return model


def accuracy(logits):
    """Return the fraction of the 997 test digits classified correctly."""
    return (logits.argmax(1) == y[800:]).float().mean().item()


base = train(nn.Sequential(nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 10)), 0, 300, 0.05, 0, 0)  # pretrained
sweep = [(lr, noise) for lr in (0.01, 0.03, 0.1) for noise in (0, 0.2, 0.4)]  # 9 hyperparameter settings
ingredients = [train(base, 300, 500, lr, noise, seed) for seed, (lr, noise) in enumerate(sweep)]

# Uniform soup: average every weight tensor of the nine fine-tuned models
soup = copy.deepcopy(base)
soup.load_state_dict({k: sum(m.state_dict()[k] for m in ingredients) / 9 for k in soup.state_dict()})

with torch.no_grad():
    singles = [accuracy(m(X[800:])) for m in ingredients]
    print(f"Fine-tuned models: {min(singles):.3f} to {max(singles):.3f}, mean {sum(singles) / 9:.3f}")
    print(f"Uniform soup: {accuracy(soup(X[800:])):.3f}")
    print(f"Ensemble of 9: {accuracy(sum(m(X[800:]).softmax(1) for m in ingredients)):.3f}")

Chín model được fine-tune đạt điểm từ 0.881 đến 0.936 trên 997 chữ số giữ lại, trung bình là 0.917. Soup đồng nhất của chúng đạt 0.934, gần bằng lần chạy đơn lẻ tốt nhất và nhỉnh hơn mức 0.930 của ensemble; ensemble cần chín lượt forward cho mỗi ảnh, còn soup chỉ cần một lượt.

Sử dụng soup trong thực tế#

Soup hữu ích nhất khi vốn đã cần chạy một lượt quét, chẳng hạn khi transfer learning một vision transformer (ViT) hoặc foundation model khác cho tác vụ phân loại ảnh mới. Tinh chỉnh siêu tham số Ultralytics YOLO cũng là một lượt quét như vậy: model.tune() huấn luyện một model cho mỗi bộ siêu tham số đã biến đổi và xóa các lần chạy không còn tốt nhất để tiết kiệm dung lượng, chỉ giữ lại lần chạy cùng trọng số của vòng lặp tốt nhất — bước chọn model tốt nhất mà soup thay thế.

Các thành phần phải dùng chung kiến trúc và điểm khởi đầu tiền huấn luyện, vì các mạng được huấn luyện từ những lần khởi tạo ngẫu nhiên khác nhau lưu trữ cùng một đặc trưng ở các đơn vị ẩn khác nhau, khiến việc lấy trung bình làm mất độ chính xác. Trọng số cuối không tốn thêm chi phí khi suy luận, vì vậy soup giữ nguyên độ trễ suy luận và bộ nhớ của một model.

Model soup so với các phương pháp liên quan#

• Ensemble model: Lấy trung bình dự đoán của một số model, khiến chi phí suy luận tăng theo số thành viên. Trong lượt quét CLIP ViT-B/32, ensemble logit đạt 81.19% trên ImageNet, nhỉnh hơn một chút so với 81.03% của soup tham lam, nhưng phải chạy mọi model thành viên cho mỗi ảnh.
• Hợp nhất model: Nhóm phương pháp rộng hơn để kết hợp trọng số. Soup là trường hợp mọi model được fine-tune trên cùng một tác vụ, còn phép toán tác vụ cộng các model được fine-tune trên những tác vụ khác nhau.
• Lấy trung bình trọng số trong một lần chạy: Trung bình trọng số ngẫu nhiên (SWA) và trung bình trượt hàm mũ (EMA) của trọng số lấy trung bình các checkpoint dọc theo một quỹ đạo huấn luyện duy nhất. Soup lấy trung bình các lần chạy độc lập; bài báo phát hiện rằng lấy soup từ các model vốn đã dùng EMA hoặc SWA vẫn cải thiện độ chính xác, nghĩa là hai kiểu lấy trung bình cộng hưởng với nhau.
• WiSE-FT: Nội suy giữa model zero-shot và phiên bản đã fine-tune, tạo thành soup gồm hai thành phần giúp giữ lại nhiều hơn khả năng chống chịu dịch chuyển phân phối của model tiền huấn luyện.