Task Arithmetic

Số học tác vụ là phương pháp chỉnh sửa model tiền huấn luyện bằng cách cộng và trừ vector tác vụ, tức những thay đổi trọng số do quá trình tinh chỉnh trên từng tác vụ tạo ra. Cộng vector tác vụ của một số tác vụ vào trọng số model tiền huấn luyện sẽ tạo ra một model xử lý được tất cả các tác vụ đó; trừ một vector tác vụ khiến model kém hơn ở tác vụ tương ứng mà hầu như không ảnh hưởng các tác vụ khác; còn kết hợp vector theo phép loại suy sẽ chuyển một kỹ năng sang bối cảnh không có dữ liệu gán nhãn. Mọi chỉnh sửa đều là phép cộng từng phần tử trên trọng số, nên ngoài các lần tinh chỉnh ban đầu, phương pháp không cần dữ liệu huấn luyện hay bước gradient nào. Gabriel Ilharco cùng các cộng sự tại University of Washington, Microsoft Research và Allen Institute for AI giới thiệu phương pháp này trong Editing Models with Task Arithmetic, xuất bản tại ICLR 2023.

Ba phép toán#

Cả ba phép toán đều tạo ra trọng số mới θ_new = θ_pretrained + λτ_new, trong đó hệ số tỷ lệ λ được chọn dựa trên dữ liệu xác thực giữ lại. Ilharco và cộng sự thử nghiệm các phép toán trên model CLIP được tinh chỉnh cho tám tác vụ phân loại ảnh, chẳng hạn nhận diện biển báo giao thông và ảnh vệ tinh, cũng như trên model ngôn ngữ GPT-2 và T5.

Phủ định dùng τ_new = −τ để khiến model quên một tác vụ. Phủ định tám vector tác vụ CLIP làm giảm độ chính xác trung bình của model ViT-L/14 trên các tác vụ đó 45,8 điểm phần trăm, đồng thời giữ lại ít nhất 95% độ chính xác trên ImageNet. Trừ một vector được tinh chỉnh trên bình luận độc hại làm giảm tỷ lệ đầu ra độc hại của GPT-2 Large từ 4,8% xuống 0,8%, trong khi perplexity trên văn bản thông thường chỉ chênh lệch tối đa 0,5 điểm so với model gốc.

Cộng dùng τ_new = τ_1 + τ_2 + … để xây dựng model đa tác vụ. Cộng vector của hai tác vụ CLIP tạo ra một model có độ chính xác trung bình bằng 98,9% so với hai model chuyên biệt; cộng cả tám vector đạt 91,2%.

Phép loại suy áp dụng khi các tác vụ có quan hệ “A đối với B như C đối với D”, sử dụng τ_D ≈ τ_C + (τ_B − τ_A). Ilharco và cộng sự xây dựng model phân tích cảm xúc cho bài đánh giá Yelp mà không cần dữ liệu Yelp có nhãn, bằng cách cộng vector cảm xúc đánh giá Amazon vào hiệu giữa các vector mô hình hóa ngôn ngữ được huấn luyện trên văn bản Yelp và Amazon.

Các phép toán này hoạt động vì vector tác vụ của các tác vụ khác nhau gần như vuông góc, với độ tương đồng cosine gần bằng không, nên cộng một vector hầu như không làm thay đổi những gì vector khác mã hóa. Nhiễu vẫn tăng theo số lượng tác vụ; các phương pháp gộp model như TIES-Merging giảm nhiễu bằng cách loại bỏ các thay đổi nhỏ nhất trong từng vector và xử lý các tham số có dấu trái ngược trước khi cộng.

Cộng và phủ định kỹ năng bằng mã#

Ví dụ PyTorch này tiền huấn luyện nhanh một bộ phân loại nhỏ trên chữ số viết tay kích thước 8x8, sau đó tinh chỉnh lớp đầu tiên cho hai tác vụ: tác vụ A phân loại các chữ số 0-4, còn tác vụ B phân loại các chữ số 5-9; mỗi tác vụ đọc từ năm logit đầu ra được đóng băng riêng, giống như mỗi tác vụ CLIP dùng một head phân loại riêng được đóng băng trong bài báo. Sau đó, ví dụ đánh giá trọng số tiền huấn luyện cộng với các tổ hợp vector tác vụ khác nhau:

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)
X, y = torch.tensor(X / 16, dtype=torch.float32), torch.tensor(y)
tasks = {"A": (0, 5), "B": (5, 10)}  # task A classifies digits 0-4, task B digits 5-9, each with its own 5 logits
train, test = torch.arange(len(X)) < 1000, torch.arange(len(X)) >= 1000


def task_loss(model, task, rows):
    """Cross-entropy and accuracy of one task's 5-way classifier on the given rows."""
    lo, hi = tasks[task]
    keep = rows & (y >= lo) & (y < hi)
    logits, labels = model(X[keep])[:, lo:hi], y[keep] - lo
    return nn.functional.cross_entropy(logits, labels), (logits.argmax(1) == labels).float().mean().item()


def fit(model, params, steps, loss_fn):
    """Run full-batch SGD on the given parameters."""
    optimizer = torch.optim.SGD(params, lr=0.05, momentum=0.9)
    for _ in range(steps):
        optimizer.zero_grad()
        loss_fn().backward()
        optimizer.step()
    return model


# Pretrained model: a few steps on 200 images of all ten digits
base = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 10))
fit(base, base.parameters(), 50, lambda: nn.functional.cross_entropy(base(X[:200]), y[:200]))

# Fine-tune the first layer for each task (heads stay frozen) and keep the task vectors
vectors = {}
for task in tasks:
    tuned = copy.deepcopy(base)
    fit(tuned, tuned[0].parameters(), 100, lambda: task_loss(tuned, task, train)[0])
    vectors[task] = {k: v - base.state_dict()[k] for k, v in tuned.state_dict().items()}

edits = {"Pretrained": {}, "+A": {"A": 1}, "+B": {"B": 1}, "+A +B": {"A": 1, "B": 1}, "-A": {"A": -1}}
for name, coefs in edits.items():
    # Pretrained weights plus a weighted sum of task vectors
    model = copy.deepcopy(base)
    weights = {k: v + sum(c * vectors[t][k] for t, c in coefs.items()) for k, v in base.state_dict().items()}
    model.load_state_dict(weights)
    with torch.no_grad():
        print(f"{name}: task A {task_loss(model, 'A', test)[1]:.3f}, task B {task_loss(model, 'B', test)[1]:.3f}")

Trên 797 chữ số kiểm thử, model tiền huấn luyện đạt 0.802 ở tác vụ A và 0.897 ở tác vụ B. Chỉ cộng vector của tác vụ A làm điểm tác vụ A tăng lên 0.925 nhưng khiến điểm tác vụ B giảm xuống 0.850; đây là hiện tượng quên thảm họa trong tinh chỉnh thông thường. Chỉ cộng vector của tác vụ B tạo ra tác động ngược lại. Cộng cả hai tạo ra một model đạt 0.922 và 0.940, gần với kết quả của từng model chuyên biệt trên tác vụ tương ứng. Trừ vector của tác vụ A làm điểm tác vụ A giảm xuống 0.616, trong khi điểm tác vụ B chỉ thay đổi từ 0.897 thành 0.887.

Yêu cầu và giới hạn#

Số học tác vụ yêu cầu mọi lần tinh chỉnh đều bắt đầu từ cùng trọng số tiền huấn luyện và dùng cùng kiến trúc. Khi Ultralytics YOLO26 huấn luyện từ yolo26n.pt trên tập dữ liệu mới, model sao chép mọi tensor tiền huấn luyện có hình dạng phù hợp, khớp các hàng của head phân loại theo tên lớp và ghi số lượng theo dạng “Đã chuyển N/M mục từ trọng số tiền huấn luyện”. Tập dữ liệu có số lượng lớp khác nhau khiến các lớp tích chập phân loại cuối cùng có hình dạng khác nhau, vì vậy chỉ có thể cộng vector tác vụ từ các lần chạy đó ở những lớp chung.

Phép tính cũng không còn hiệu quả khi các vector xung đột hoặc bị nhân tỷ lệ quá lớn, vì giá trị λ lớn đẩy trọng số ra khỏi vùng mà model tinh chỉnh hoạt động tốt. Phủ định làm giảm hiệu năng trên những gì dữ liệu tinh chỉnh bao phủ, khiến đây trở thành một dạng quên máy học nhẹ. Bài báo Task Arithmetic in the Tangent Space của Guillermo Ortiz-Jimenez, Alessandro Favero và Pascal Frossard truy nguyên thành công của phép tính này đến sự tách biệt trọng số; trong đó các hướng riêng biệt trong không gian trọng số làm thay đổi đầu ra model trên những vùng không chồng lấn riêng biệt của không gian đầu vào. Bài báo cũng cho thấy tinh chỉnh một phiên bản tuyến tính hóa của model giúp tăng cường sự tách biệt này.

Số học tác vụ và các khái niệm liên quan#

• Vector tác vụ: Các độ chênh lệch trọng số mà số học tác vụ thao tác. Số học tác vụ là tập hợp các phép toán: cộng, phủ định và phép loại suy.
• Model soup: Lấy trung bình một số model được fine-tune cho cùng một tác vụ để cải thiện tác vụ đó. Phép toán tác vụ cộng các bản fine-tune của những tác vụ khác nhau; cộng hai vector tác vụ với λ = 0.5 tương đương với lấy trung bình hai model đã fine-tune.
• Chỉnh sửa model: Thay đổi một sự thật hoặc hành vi giới hạn bằng một lần cập nhật trọng số có mục tiêu, thường có dạng đóng. Phép toán tác vụ chỉnh sửa toàn bộ tác vụ cùng lúc.
• Huấn luyện đa tác vụ: Huấn luyện một model trên dữ liệu của tất cả tác vụ cùng lúc, giúp tránh nhiễu lẫn nhau nhưng đòi hỏi phải có mọi dataset cùng lúc. Phép toán tác vụ kết hợp các model mà dữ liệu có thể không còn sẵn có, đồng thời tránh hiện tượng quên do fine-tune các tác vụ lần lượt.