[25' ICLR] Task Arithmetic in Trust Region: A Training-Free Model Merging Approach to Navigate Knowledge Conflicts

한 줄 요약

Task vector의 gradient-aligned 성분이 knowledge conflict의 주원인임을 밝히고, trust region 내의 gradient-orthogonal 방향만을 선택적으로 병합하여 training 없이도 안정적인 multi-task model merging을 달성하는 TATR 방법을 제안한다.

 

문제점

Knowledge Conflict의 발생

Task arithmetic에서 서로 다른 task의 task vector를 단순 합산하면 knowledge conflict가 발생한다.
이는 하나의 task에 대해 fine-tuning된 모델과 병합 시, 다른 task의 성능이 저하되는 현상으로 정의된다.
기존 방법들은 이 문제를 scaling coefficient α의 조정으로 해결하려 했으나, 최적의 α를 찾는 것은 비용이 높고 task 조합에 따라 달라진다.

기존 해결 방법의 한계

Training-based 방법 (예: AdaMerging)은 추가적인 학습을 요구하여 계산 비용이 크다.
Training-free 방법 (예: TIES-Merging, DARE)은 heuristic에 의존하며 이론적 근거가 부족하다.
Knowledge conflict가 parameter space의 어떤 성분에서 기인하는지에 대한 체계적 분석이 부재했다.

 

제안 방법

Knowledge Conflict의 원인 분석

저자들은 1차 Taylor 전개를 활용하여 task vector가 각 task의 loss에 미치는 영향을 분석한다.
Task vector τ를 사전학습 모델 θ₀에 더할 때, task j의 loss 변화는 근사적으로 ∇L_j(θ₀)ᵀ · τ로 표현된다.
따라서 task vector 중 gradient 방향과 정렬된 성분이 다른 task의 loss를 크게 변화시켜 conflict를 유발한다.
반면, gradient에 직교하는 방향의 성분은 loss에 미치는 영향이 미미하다.

Trust Region의 정의

Trust region은 모델 parameter space에서 task-specific loss에 작은 변화만 유발하는 차원(dimension)들의 집합이다.
구체적으로, 각 task의 loss gradient에 직교하는 방향들이 trust region을 구성한다.
이 영역 내의 task vector 성분만을 병합에 사용하면, cross-task interference를 최소화할 수 있다.

TATR 알고리즘

TATR은 두 단계로 구성된다.

Linear Layer 처리:
각 task i에 대해 removal basis B_i를 계산한다.
이는 task i의 feature projection을 최대화하면서 다른 task와의 간섭을 최소화하는 최적화 문제의 해이다.
최적화 목적 함수는 다음과 같다: max_{B_i} Σ_{j≠i} ||X_i W_j B_i B_iᵀ||²_F - λ||X_j W_j B_i B_iᵀ||²_F
이 문제는 Gram matrix의 고유벡터 분해(eigendecomposition)로 효율적으로 풀린다.
계산된 removal basis를 사용하여 conflicting 성분을 제거한다: T_i = T_i - T_i B_k B_kᵀ

Normalization/Bias Layer 처리:
Binary mask m_i를 학습하여, element-wise로 conflict를 유발하는 차원을 선택적으로 제거한다.

Plug-and-Play 호환성

TATR은 기존 task arithmetic 기반 방법에 플러그인 모듈로 적용 가능하다.
Task Arithmetic, TIES-Merging, DARE, AdaMerging 등과 결합하여 일관된 성능 향상을 제공한다.

 

학습 상세

모델 및 데이터셋

• 모델: CLIP 기반 ViT-B/32, ViT-L/14
• 데이터셋: SUN397, Cars, RESISC45, EuroSAT, SVHN, GTSRB, MNIST, DTD (8개 vision classification task)
• Exemplar 데이터: 각 task에서 소수의 unlabeled exemplar만 필요 (1~32개)

Training-Free 특성

TATR은 추가적인 gradient 기반 학습을 전혀 수행하지 않는다.
Exemplar 데이터를 forward pass하여 feature를 수집하고, 이를 기반으로 removal basis를 closed-form으로 계산한다.
따라서 학습 시간이 거의 소요되지 않으며, 기존 fine-tuned 모델만 있으면 바로 적용 가능하다.

Alpha Robustness의 이론적 정당화

Trust region 내의 성분만으로 병합하면, scaling coefficient α에 대한 민감도가 크게 감소한다.
이는 trust region 내의 task vector 성분이 다른 task의 loss에 미치는 영향이 이미 최소화되어 있기 때문이다.
따라서 α의 선택이 성능에 미치는 영향이 줄어들어, robust한 α 범위가 넓어진다.
이론적으로, trust region 내에서는 1차 근사 하에 cross-task loss 변화가 0에 수렴하므로, α에 대한 robustness가 보장된다.

 

실험 결과

ViT-B/32 결과

• Pre-trained 모델: 평균 48.0%
• Task Arithmetic: 평균 69.1%
• TATR (제안 방법): 평균 약 76.1% (Task Arithmetic 대비 +3.7%)
• TATR은 training-free 방법 중 최고 성능을 달성한다.

ViT-L/14 결과

• Task Arithmetic: 평균 84.5%
• TATR: 평균 약 86.7% (Task Arithmetic 대비 +0.8%)
• 대규모 모델에서도 일관된 향상을 보인다.

기존 방법과의 결합

• TATR + AdaMerging++: ViT-B/32에서 평균 82.5%까지 성능 향상
• Weight Averaging 대비 +10.3%, Fisher Merging 대비 +7.8% 향상
• 모든 TA-based 방법에 TATR을 결합하면 일관된 성능 개선을 확인한다.

Zero-shot TATR

• Exemplar 없이 사용 가능한 zero-shot 버전도 제공한다.
• ViT-B/32에서 Task Arithmetic 대비 +1.5%, ViT-L/14에서 +0.1% 향상을 달성한다.

Ablation Study: Exemplar 수의 영향

Exemplar 수 (per task)	평균 정확도 (ViT-B/32)
1	78.00 ± 0.2%
2	77.75 ± 0.1%
4	77.59 ± 0.1%
8	77.40 ± 0.1%
16	77.21 ± 0.1%
32	77.00 ± 0.1%

• 단 1~2개의 exemplar만으로도 충분한 성능을 달성한다.
• Exemplar 수 증가에 따른 성능 변화가 미미하여, 데이터 효율성이 높다.

Trust Region 비율 분석

• 전체 parameter의 1% 미만만 제거해도 knowledge conflict를 효과적으로 완화할 수 있다.
• 이는 conflict가 소수의 특정 parameter 차원에 집중되어 있음을 시사한다.

 

한계점 및 개인 의견

한계점

• 1차 근사의 한계: Trust region 정의가 Taylor 1차 전개에 기반하므로, 실제 loss landscape의 비선형적 특성을 완전히 포착하지 못할 수 있다.
• Training-based 방법 대비 성능 격차: AdaMerging++ 단독(test-time training 포함) 대비 여전히 성능 차이가 존재한다.
• Exemplar 데이터 필요: Training-free라고는 하지만, removal basis 계산을 위해 각 task의 unlabeled exemplar가 필요하다. 완전한 zero-shot 버전은 성능이 제한적이다.
• Layer별 처리의 독립성: 각 layer를 독립적으로 처리하여 layer 간 상호작용을 고려하지 못한다.
• Error bar 미보고: 주요 실험 결과에 대한 error bar가 보고되지 않아, 결과의 통계적 신뢰성 검증이 어렵다.

개인 의견

Knowledge conflict의 원인을 gradient alignment 관점에서 분석한 것은 매우 직관적이고 설득력 있는 접근이다.
특히 spectral decomposition을 통해 removal basis를 closed-form으로 계산하는 방식은 이론적 우아함과 실용성을 모두 갖추고 있다.
전체 parameter의 1% 미만만 제거해도 충분하다는 발견은, task vector 간 conflict의 저차원적 구조를 시사하며 흥미로운 후속 연구 방향을 제시한다.
다만, 2차 이상의 근사를 도입하거나 layer 간 상호작용을 고려하는 확장이 향후 성능 개선에 기여할 수 있을 것이다.
TATR의 plug-and-play 특성은 실용적으로 큰 장점이며, model merging 파이프라인의 표준 전처리 단계로 자리잡을 가능성이 있다.

 

논문 정보 및 리소스

• 논문 제목: Task Arithmetic in Trust Region: A Training-Free Model Merging Approach to Navigate Knowledge Conflicts
• 저자: Wenju Sun, Qingyong Li, Wen Wang, Yangliao Geng, Boyang Li
• 학회: ICLR 2025 (Conference Paper)
• 논문 링크: arXiv | OpenReview
• 페이지: 21 pages, 6 figures, 6 tables