Editing Models with Task Arithmetic

#multi-task-learning #task-vectors #model-editing #task-arithmetic

Link

Abstract

Changing how pre-trained models behave -- e.g., improving their performance on a downstream task or mitigating biases learned during pre-training -- is a common practice when developing machine learning systems. In this work, we propose a new paradigm for steering the behavior of neural networks, centered around \textit{task vectors}. A task vector specifies a direction in the weight space of a pre-trained model, such that movement in that direction improves performance on the task. We build task vectors by subtracting the weights of a pre-trained model from the weights of the same model after fine-tuning on a task. We show that these task vectors can be modified and combined together through arithmetic operations such as negation and addition, and the behavior of the resulting model is steered accordingly. Negating a task vector decreases performance on the target task, with little change in model behavior on control tasks. Moreover, adding task vectors together can improve performance on multiple tasks at once. Finally, when tasks are linked by an analogy relationship of the form ``A is to B as C is to D", combining task vectors from three of the tasks can improve performance on the fourth, even when no data from the fourth task is used for training. Overall, our experiments with several models, modalities and tasks show that task arithmetic is a simple, efficient and effective way of editing models.

Synth

Problem:: 하나의 모델에 대해 세상에 사용 할 수 있는 Weight가 매우 많음

Solution:: 하나의 모델의 Weight들을 단순 연산을 통해 향상시키자

Novelty:: 세상에 많은 Weight들을 단순한 연산을 통해 효과적으로 이용하는 방식

Note:: 유사한 Task간의 Weight 유사도가 높다는 것을 보임 → 서로 다른 Task는 거의 Orthogonal → 직교하므로 더하더라도 벡터 내적 결과에 큰 영향이 없음 → 동일 Feature에 대해 서로 다른 Branch의 Feature가 유사하면 서로 상충된다? /Figure 7은 좀 의심스러움 → 아주 적은 Step으로 학습된 Task Vector를 더한 것으로 높은 성능이 달성된다? (아주 적은 학습으로 애초에 해당 Task의 성능이 높았다면 말이됨. 단, 이 경우 최종 성능을 달성하는것까지의 시간이 짧아진다가 주장이 될텐데 글 작성이 모호함)

Summary

Motivation

Pre-trained 모델 행동 수정(downstream task 성능 향상, bias 완화 등)은 ML 시스템 개발에서 일반적 작업
인기 있는 모델은 사용 할 수 있는 Pre-trained Weight가 아주 많음 → 이거랑 Weight Interpolation을 잘 엮어보자

Method

Task Vector $τ_{t}$ 는 Pre-trained 모델 가중치( $θ_{p r e}$ )와 같은 모델을 task t에 Fine-tuning한 가중치( $θ_{f t}^{t}$ ) 차이로 정의
- $τ_{t} = θ_{f t}^{t} - θ_{p r e}$ 로 계산
Task Vector는 같은 아키텍처의 모델 파라미터 $θ$ 에 요소별 덧셈과 옵션 스케일링 계수 $λ$ 를 통해 적용
- $θ_{n e w} = θ + λ τ$
- $λ = 1$ 일 경우 Fine-tuned 모델과 동일

Task Arithmetic 연산

Negation(부정): $τ_{n e w} = - τ$
- 원치 않는 행동 제거 또는 특정 Ttask 잊기에 활용
Addition(덧셈): $τ_{n e w} = \sum_{i} τ_{i}$
- 다중 task에 대한 모델 생성
- 경우에 따라 개별 Fine-tuned 모델보다 우수한 성능 달성 가능
Analogy(유추): $τ_{D} = τ_{C} + (τ_{B} - τ_{A})$
- "A는 B에 대해 C는 D와 같다" 형태의 관계 활용
- D task에 대한 데이터가 적거나 없어도 성능 향상 가능

Method 검증

Forgetting via Negation

이미지 분류: CLIP 모델에 8개 Task에 대한 실험 수행
- Task Vector 부정 시 ViT-L/14에서 Target Task 정확도 45.8%p 감소, Control Task(ImageNet) 정확도는 유지 → 특정 Task 선택적 망각 가능성 입증
- Random Vector는 Target 정확도에 영향 적음, Gradient Ascent는 Control Task 성능도 심각하게 저하
텍스트 생성: GPT-2 모델에 Toxic Content 생성 감소 실험
- Civil Comments 데이터 중 Toxicity Score 0.8 이상인 데이터로 Fine-tuning 후 Vector 부정
- 독성 생성 비율 4.8%에서 0.8%로 감소(6배 감소), WikiText-103 Perplexity는 0.5 이내 변화 → Control Task 능력 유지하며 특정 행동만 제거 가능

Learning via Addition

이미지 분류: 8개 Task Vector 조합 실험
- 두 개 Task Vector 더했을 때 평균 98.9% Normalized Accuracy 달성 → 개별 모델 사용과 유사한 성능을 단일 모델로 구현
- Task Vector 수 증가에 따라 평균 성능도 향상 → 8개 모두 사용 시 91.2% 달성
자연어 처리: T5 모델에서 GLUE Benchmark Task Vector 결합
- Hugging Face Hub의 427개 호환 Checkpoint에서 최적 조합 탐색
- MRPC, RTE, CoLA, SST-2 등 모든 Task에서 Fine-tuning 대비 성능 향상(평균 +0.5%p)

Task Analogies

도메인 일반화: 새 도메인에 대한 라벨 데이터 없이 성능 향상
- Yelp 감성분석 Task Vector: ${\hat{τ}}_{y e l p; s e n t} = τ_{a m a z o n; s e n t} + (τ_{y e l p; l m} - τ_{a m a z o n; l m})$ 수식으로 구성
- 다양한 T5 모델(Small, Base, Large)에서 Auxiliary Task만으로 Fine-tuning한 것보다 높은 성능 달성 → Target Task 라벨 없이도 도메인 일반화 가능함 입증
데이터 부족 하위집단: 데이터 희소성 문제 해결
- 125개 클래스를 가진 ImageNet과 스케치 데이터셋에서 유추 관계 활용
- 예: ${\hat{τ}}_{s k e t c h, l i o n} = τ_{s k e t c h, d o g} + (τ_{r e a l, l i o n} - τ_{r e a l, d o g})$ → 스케치 사자 (데이터 없음) = 스케치 개 (Weight 1) + 진짜 사자 (Weight 2) - 진짜 개 (Weight3) → 스케치 사자의 성능이 좋음
- Pre-trained 모델 대비 평균 3.4%p 향상 → 약 100개 훈련 샘플 직접 수집 효과와 동등

Ablation

Task Vector 간 유사도 분석 결과 대부분 직교에 가까움 → 덧셈 연산 시 최소한의 간섭으로 효과적 작동
Learning Rate 실험에서 낮은 Learning Rate가 Task Vector에 더 유리함을 확인
Intermediate Task Vector 실험 결과 Fine-tuning 초기 단계(200-500 스텝)에서 추출한 Vector도 완전 훈련된 Vector와 유사한 성능 제공 → 계산 자원 절약 가능 → 해당 실험 방식이 너무 모호하게 작성됨

Limitation

동일 아키텍처 모델에만 적용 가능
같은 Pre-trained 초기화에서 Fine-tuning된 모델 간에만 작동
가중치 공간의 선형 연산만 지원

Summary

Motivation

Method

Task Arithmetic 연산

Method 검증

Forgetting via Negation

Learning via Addition

Task Analogies

Ablation

Limitation

Ablation