Task Arithmetic

태스크 산술은 사전 학습된 모델에 태스크 벡터를 더하거나 빼서 모델을 편집하는 방법입니다. 태스크 벡터는 각 작업에 대한 파인튜닝으로 생성된 가중치 변화량입니다. 여러 작업의 태스크 벡터를 사전 학습된 모델 가중치에 더하면 모든 작업을 처리하는 하나의 모델을 만들 수 있습니다. 태스크 벡터를 빼면 다른 작업은 대부분 그대로 유지하면서 해당 작업의 성능을 낮출 수 있고, 유추 관계에 따라 벡터를 결합하면 레이블이 없는 환경으로 기술을 전이할 수 있습니다. 각 편집은 가중치에 대한 요소별 덧셈이므로 원래 파인튜닝 이후에는 학습 데이터나 그래디언트 단계가 필요하지 않습니다. University of Washington, Microsoft Research, Allen Institute for AI의 Gabriel Ilharco와 동료들은 ICLR 2023에서 발표한 Editing Models with Task Arithmetic 논문에서 이 접근 방식을 소개했습니다.

세 가지 연산#

세 연산 모두 θ_new = θ_pretrained + λτ_new의 새 가중치를 생성하며, 스케일 계수 λ는 홀드아웃 검증 데이터에서 선택합니다. Ilharco 등은 교통 표지판 및 위성 이미지와 같은 8개 이미지 분류 작업으로 파인튜닝한 CLIP 모델과 GPT-2 및 T5 언어 모델에서 세 연산을 테스트했습니다.

부정: τ_new = −τ를 사용해 모델이 작업을 잊게 합니다. CLIP 태스크 벡터 8개를 부정하면 ViT-L/14 모델의 해당 작업 평균 정확도가 45.8%포인트 하락하면서 ImageNet 정확도는 95% 이상 유지되었습니다. 유해 댓글로 파인튜닝한 벡터를 빼면 GPT-2 Large의 유해 생성 비율은 4.8%에서 0.8%로 줄었으며, 일반 텍스트의 퍼플렉시티는 원본 모델과 0.5포인트 이내의 차이를 유지했습니다.

덧셈: τ_new = τ_1 + τ_2 + …를 사용해 멀티태스크 모델을 만듭니다. 두 CLIP 작업의 벡터를 더하면 하나의 모델로 두 전문 모델의 평균 정확도 중 98.9%를 달성했으며, 8개 벡터를 모두 더하면 91.2%를 달성했습니다.

유추: 작업 간 관계가 “A는 B와 같고 C는 D와 같다”일 때 τ_D ≈ τ_C + (τ_B − τ_A)를 적용합니다. Ilharco 등은 Yelp 레이블 데이터 없이 Yelp 리뷰용 감정 분석 모델을 만들었습니다. 이를 위해 Amazon 리뷰 감정 벡터와 Yelp 및 Amazon 텍스트로 학습한 언어 모델링 벡터 간 차이를 더했습니다.

서로 다른 작업의 태스크 벡터는 코사인 유사도가 0에 가까운 거의 직교 방향을 가리키므로 이 연산이 가능합니다. 따라서 한 벡터를 더해도 다른 벡터가 인코딩한 내용은 거의 바뀌지 않습니다. 그래도 작업 수가 많아지면 간섭이 커집니다. 모델 병합 방법인 TIES-Merging은 벡터의 작은 변화량을 잘라내고 부호가 다른 파라미터를 조정한 다음 더하여 간섭을 줄입니다.

코드에서 기술 더하기와 빼기#

이 PyTorch 예제는 8x8 손글씨 숫자로 작은 분류기를 잠시 사전 학습한 뒤, 두 작업을 위해 첫 번째 레이어를 파인튜닝합니다. 작업 A는 숫자 0~4를, 작업 B는 숫자 5~9를 분류하며, 논문에서 각 CLIP 작업이 고유한 동결 분류 헤드를 사용하듯 각 작업도 고유한 5개 동결 출력 로짓을 사용합니다. 그런 다음 사전 학습된 가중치에 서로 다른 태스크 벡터 조합을 적용해 평가합니다.

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}")

테스트 숫자 797개에서 사전 학습 모델은 작업 A에서 0.802, 작업 B에서 0.897의 점수를 얻습니다. 작업 A의 벡터만 더하면 작업 A의 점수는 0.925로 올라가지만 작업 B의 점수는 0.850으로 낮아집니다. 이는 일반적인 파인튜닝의 파국적 망각이며, 작업 B의 벡터만 더하면 그 반대가 일어납니다. 두 벡터를 모두 더하면 하나의 모델로 0.922와 0.940의 점수를 얻어 각 작업에서 전문 모델에 가까운 성능을 냅니다. 작업 A의 벡터를 빼면 작업 A 점수는 0.616로 낮아지는 반면, 작업 B 점수는 0.897에서 0.887로만 바뀝니다.

요건 및 한계#

태스크 산술을 사용하려면 각 파인튜닝을 동일한 아키텍처의 동일한 사전 학습 가중치에서 시작해야 합니다. 새 데이터셋으로 yolo26n.pt에서 학습하는 Ultralytics YOLO26은 형태가 맞는 사전 학습 텐서를 모두 복사하고, 클래스 이름을 기준으로 분류 헤드 행을 매칭한 다음 “Transferred N/M items from pretrained weights”라는 수를 로그에 기록합니다. 클래스 수가 다른 데이터셋에서는 최종 분류 컨볼루션의 형태가 달라지므로, 해당 실행에서 얻은 태스크 벡터는 공통으로 사용하는 레이어에서만 더할 수 있습니다.

벡터가 서로 충돌하거나 지나치게 큰 계수로 스케일링되면 태스크 산술은 제대로 작동하지 않습니다. 큰 λ 값은 파인튜닝된 모델이 잘 작동하는 영역 밖으로 가중치를 밀어낼 수 있습니다. 부정은 파인튜닝 데이터가 다룬 모든 항목의 성능을 낮추므로 경량 머신 언러닝의 한 형태입니다. Guillermo Ortiz-Jimenez, Alessandro Favero, Pascal Frossard의 접공간에서의 태스크 산술은 성공적인 산술의 원인을 가중치 분리에서 찾았습니다. 가중치 공간의 개별 방향이 입력 공간의 서로 겹치지 않는 영역에 대한 모델 출력에 각각 영향을 준다는 설명이며, 모델의 선형화된 버전을 파인튜닝하면 이러한 분리가 강화된다는 점을 보였습니다.

태스크 산술과 관련 개념 비교#

• 태스크 벡터: 태스크 산술에서 사용하는 가중치 차이입니다. 태스크 산술은 덧셈, 부정, 유추로 이루어진 연산 집합입니다.
• 모델 수프: 하나의 태스크에서 여러 파인튜닝 모델의 평균을 구해 해당 태스크의 성능을 개선합니다. 태스크 산술은 서로 다른 태스크의 파인튜닝 결과를 더하며, 두 태스크 벡터를 λ = 0.5로 더하면 파인튜닝된 두 모델의 평균을 구한 것과 같습니다.
• 모델 편집: 특정 사실 하나 또는 제한된 동작을 대상으로, 일반적으로 폐쇄형 가중치 업데이트를 적용해 변경합니다. 태스크 산술은 전체 태스크를 한 번에 편집합니다.
• 멀티태스크 학습: 모든 태스크의 데이터를 함께 사용해 하나의 모델을 학습하므로 간섭을 방지하지만 모든 데이터셋을 한꺼번에 준비해야 합니다. 태스크 산술은 데이터에 더 이상 접근할 수 없을 수도 있는 모델을 결합하며, 태스크를 하나씩 차례로 파인튜닝할 때 발생하는 망각을 방지합니다.