[yongggg's] TIES-MERGING: Resolving Interference When Merging Models Review

안녕하세요 이번 장에서는 LLM의 sub task fine-tuning model을 merge 하는 방법에 대한 내용에 대해 다뤄보겠습니다. 

Backbone LLM혹은 PLM으로 sub task를 fine-tuning했을 때, down-stream task의 성능 향상, 빠른 수렴 속도가 장점이지만, 일반적으로 단일 task만 수행할 수 있으며, 서로의 이점은 얻지 못합니다. 따라서 최근 model merge 기술은 추가 training 없이 여러 tasks의 model을 single multitask model로 결합하는 솔루션이 나왔다. 하지만, 기존 병합 방법은 서로 다른 Model의 parameter 간의 간섭을 무시하여 여러 모델을 병합할 때 큰 성능 저하를 초래한다.

이 논문에서는 이전 merge 기술이 (a) 중복 Parameter 변수 값으로 인한 간섭 (b) model 간에 주어진 매개 변수 값의 부호에 대한 불일치 라는 두가지 주요 간섭 원인으로 인해 부주의하게 귀중한 정보를 잃는 다는 것을 보여줍니다. 이를 해결하기 위해 모델을 병합할 때 세 가지 새로운 단계인 TRIM, ELECT SIGN, MERGEING(TIES-MERGING)을 도입한다. (1) fine-tuning 중에 소량만 변경된 parameter를 재설정하고, (2) 부호 충돌을 해결하며, (3) 최종 합의된 부호와 일치하는 parameter만을 병합합니다.

이 단계를 거친 Merge 방법으로 기존 여러 방법보다 merge 성능이 뛰어나며, 다양한 유형의 간섭이 모델 Parameter에 미치는 영향을 분석합니다.

이제 본격적으로 논문 Review를 시작해보겠습니다!

1. Introduction

Pre-trained models(PTMs)은 많은 real-world application에 널리 보급됐다. PTMs을 사용하는 것은 일반적으로 특정 작업에 특화되도록 fine-tuning하는 것을 포함하며, 이는 specific task에 label이 지정된 데이터를 덜 사용하여 성능을 향상시킬 수 있다. 이러한 이점으로 비전 ViT 라던가, 언어용 T5와 같은 인기 있는 PTMs에서 파생된 수천 개의 fine-tuned checkpoint가 생겨났다. 그러나 각 작업에 대해 별도의 모델을 저장하고 배포해야하며, 개별로 훈련된 모델은 관련 task의 정보를 활용하여 domain performance 또는 out-of-domain generalization을 개선할 수 없다. multi-tasking training은 이러한 우려를 해결할 수 있지만, 높은  training cost가 발생하며, 모든 작업에 대한 동시 access가 필요하다. 또한 multi-tasking 훈련이 모든 작업에 유익하도록 dataset을 잘 혼합하는 방법을 결정하는 것은 복잡하고 resource 집약적일 수 있다.

최근 model merging에 초첨을 맞춘 연구가 증가하고 있다. merging의 한 application은 추가 training을 수행하지 않고 multiple task-specific models를 단일 다중 작업 모델로 결합하는 것을 포함한다. 이전 연구에서는 매개변수 중요도를 통합하거나 순열의 불변성을 설명하는 더 정교한수단을 통해, weights 체계가 다른 개별 모델 weights을 단순 평균 합산하여 모델을 병합했다. 이러한 방식으로 fine-tuned model을 결합하는 것은 fine-tuned model의 parameter 값에서 pre-trained model의 parameter 값을 차감하여 계산된 task vector를 합산하는 것으로 볼 수 있다.

model parameter의 가중 평균이 merging 방법에 효과적인 것으로 입증되었지만, 이러한 모든 방법은 모델 간 간섭이 생겨 merged model의 성능을 손상시킬 수 있는 가능성이 발생한다. 본 논문에서는 먼저 간섭에 대한 것이 두 가지 주요 원인에서 비롯 될 수 있음을 입증한다. (아래 그림 참조) 두 가지 모두 병합된 모델에서 parameter 크기를 줄이고 값 간의 미묘한 차이를 제거할 수 있다. (1) 중복 parameter로 부터의 간섭: 모델 pruning에 대한 이전 연구에 따르면 fine-tuning 중에 많은 모델 parameter가 변경될 수 있지만, 성능에는 작은 영향만 미치는 것으로 나타났다. 그러나 한 모델에 영향력이 있지만, 다른 모델에 영향력이 없는 parameter를 병합할 때, 영향력 있는 값이 쓸모 없는 값에 의해 가려져 전체 모델 성능이 저하될 수 있다. (아래 그림의 O) (2) 반대 부호에 의한 간섭: 주어진 parameter는 일부 모델에서는 양의 값을 가질 수 있고, 다른 모델에서는 음의 값을 가질 수 있다. 결과적으로 단순 평균화를 사용하면, 두 task 모두의 성능이 저하될 수있다. (아래 그림의 육각형 모양) 이 두 가지 상황에서 모두 단순히 값을 집계하는 것은 병합된 모델에서 parameter의 값을 줄이는 간섭으로 이어진다. 영향력 있는 parameter 간의 이러한 간섭은 모델 수가 증가함에 따라 병합된 모델과 multitask trained model 모델 간의 성능 격차가 증가하는 이유를 설명할 수 있다.

평균화 또는 TIES-MERGING에 의해 생성되는 다양한 유형의 confilct 및 merged ouputs. 간섭을 일으키는 매개 변수는 점선 화살표로 표시됨.

이러한 간섭 source를 해결하기 위해 본 연구에서는 세 단계로 구성된 task vector를 결합하여 모델을 병합하는 방법인 TIES-MERGING(TRIM, ELECT SIGN, MERGE)방법을 제안한다. (아래 그림 시각화 참고) 먼저, 각 task vector의 중복 값을 0으로 설정하여 영향력 있는 매개 변수 값만 유지하도록 각 작업 벡터를 trim한다. (또는 이에 상응하여 fine-tuned parameter 값을 pre-trained model의 값으로 재설정한다.) 이 단계 이후에도 아래 보라색 그래프 그림에서 볼 수 있듯이 영향력 있는 매개 변수 값 간에 부호 충돌이 여전히 지속될 수 있다. 따라서 두 번째 단계는 서로 다른 값 간의 부호 충돌을 해결하고 마지막 단계는 부호가 모델 간의 가장 큰 총 이동 방향과 일치하는 parameter의 평균만을 한다.

TIES-MERGING과 관련된 단계를 설명한다. 모델의 각 parameter를 정사각형으로 시각화한다. 화살표는 다양한 작업(색으로 부호화)을 미세 조정하여 생성된 parameter에 대한 업데이트(작업 벡터, $\tau$)를 나타내며 방향은 부호, 길이는 크기를 나타낸다. 먼저 크기를 기준으로 작업 벡터 값을 트리밍한 다음 부호 충돌을 해결하여 각 매개변수($\gamma_{m}$, +1 또는 -1을 포함하는 녹색 벡터)의 부호를 선택한다. 마지막으로 선택한 부호와 일치하는 값만 선택하고 평균을 최종 매개변수 값으로 사용한다.

trim 이후에도 부호 충돌이 발생하며, models의 수에 따라 증가한다. trimming 후 부호 충돌이 발생하는 params의 빈도와 merge 되는 모형의 수를 plot한 그림이다.

본 연구진은 (1) 언어 및 vision benchmarks를 포함한 다양한 양식 (2) T5-base and T5-large 뿐만아니라 ViT-B/32 and BiT-L/14와 같은 서로 다른 모델 크기의 모델 (3) in-domain 및 out-domain task, (4) full fine-tuning or parameter efficient fine-tuning (5) merging hyperparameters를 설정하기 위한 유효성 검사 세트의 유무에 따라 다양한 설정에서 제안된 TIES-MERGING 방법의 효과를 입증한다. 본 연구진은 TIES-MERGING이 이러한 모든 실험 설정에서 Task Arithmetic, RegMean, Fisher Merging, and 및 weight averaging과 같은 다른 지표에서 다른 방법보다 성능이 뛰어나다는 것을 보여준다. 특히 in-domain evaluation에서 TIES-MERGING은 NLP 및 vision 설정에서 각각 평균 2.3% 및 1.7% 절대값으로 strongest baseline을 능가한다. (아래 표 1.) out of domain generalization(아래 표 2.)의 경우 TIES-MERGING은 T5-base 및 T5-large 모델에서 각각 1.0% 4.4% 절대값으로 strongest baseline을 능가한다.

2. Related Work

Loss Landscape and weight Interpolation. 신경망의 loss function은 일반적으로 non-convex이지만, 최근 연구는 다른 training runs의 parameter 값이 때때로 손실되지 않고 보간될 수 있음을 보여주었다. 예를 들어 Frankle et al. 은 최적화 궤적의 일부가 두 신경망 사이에서 공유되면, 정확도를 낮추지 않고 보간될 수 있음을 보여주었다. 반면, Neyshabur et al.은 완전히 분리된 최적화 궤적을 가진 두 신경망을 단순하게 보간하면 정확도가 크게 떨어질 수 있음을 보여준다. Ainsworth et al.는 신경망의 순열 대칭을 처리하면, 동일한 데이터 세트에서 훈련된 주어진 architecture의 모든 신경망이 선형 모드로 연결된다고 가정했다. Ainsworth et al., Singh and Jaggi, Wang et al.는 따라서 처음부터 훈련된 신경망의 손실을 증가시키지 않고 보간할 수 있도록 더 잘 align 하기 위해 순열 및 최적 전송을 찾는 데 기반한 기술을 사용했다. 

Model Merging and Different Use Cases. 동일 pre-trained model에서 초기화된 서로 다른 fine-tuned model은 최적화 궤적의 일부를 효과적으로 공유하므로 치환 대칭을 고려하지 않아도 merge되는 경우가 많다. 따라서 fine-tuning model을 병합하면 out-of-domain generalization, 다른 tasks의 multitask models 생성, federated(연합) cearning, compression, multimodal merging models, continual learning, other setings 등의 개선과 함께 단일 대상 작업의 성능을 향상시킬 수 있다. 응용 분야의 범위는 단순한 매개변수의 평균을 넘어 개선하기 위한 방법의 확산으로 이어졌다. RegMean은 모델의 각 개별 선형 계층에 대한 Local linear regression 문제를 해결하여 merged model의 Parameter에 대한 close-form solution을 제안했다. 그러나 이를 위해서는 모델과 동일한 크기의 추가 데이테 통계를 전송해야하며, 이를 계산하기 위한 추가 inference step이 필요하다. Fisher Merging은 단순히 평균을 넘어 FIsher Information Matrix를 사용하여 개별 parameter의 중요성을 식별하고, 이를 병합하는 과정에서 각 모델의 parameter에 가중치를 부여한다. 그러나 이는 여러 checkpoint를 병합할 때 거의 이점이 없으며, 메모리 비용이 높은 gradient를 계산해야 한다. Task Arithmetic에서는 task 벡터를 생성하고 덧셈과 같은 arithmetic(산술) 연산을 수행하여 multitask checkpoint를 얻는 방법을 제시했다. Ortiz-Jiménez et al.의 concurrent 연구는 pre-training 중에 발생하는 가중치 분리 속성을 기반으로 모델 병합에 대한 이론적 insights를 제공했다. 이 연구에서는 tangent space에서 fine-tuning 모델이 이러한 특성을 향상시켜 더 나은 병합 모델로 이어진다는 것을 보여주었다. 본 연구의 방법은 model merging에 대한 이러한 과거의 연구를 따르지만, 병합 중에 다른 parameter 간의 간섭을 추가로 고려한다.

3. Background and Motivation

problem setting

tasks 집합 { $t_{1}, ..., t_{n}$ }과 T5 이나 ViT같은 pre-trained model이 주어졌다고 하자, 우리는 이 전체 model을 finetune하거나 parameter efficient finetuning(PEFT)을 사용한다. 이 두 가지 경우에서 trainable parameter를 $\theta$, 이것의 initialization을 $\theta_{init}$, fine-tuned parameters를 $\theta_{ft}$로 정의한다. 본 논문에서는 multiple task에 대해 $\theta_{ft}$의 finetuned model parameters에 대한 access를 가정하고 이러한 model의 가중치를 in-domain과 out-of-domain datasets 모두에 능숙한 single multitask model로 병합하는 방법을 고안한다. 우리는 Ilharco et al.을 따르며 task vector와 병합을 수행한다. 특히 task $t$에 대한 task-vector $\tau_{t} \in R^{d}$는 $\tau_{t}=\theta_{ft}^{t} - \theta^{t}_{init}$으로 정의한다. 이 작업을 통해 각 task별 model의 fine-tuning단계에서 발생하는 변경 사항에 초점을 맞출 수 있으며, 적절한 scaling으로 모델의 가중치를 가중 평균하는 것과 동일하다.

Redundancies in Model Parameters.

먼저 본 연구진은 주어진 task-vector에서 많은 값이 중복되며(첫 번째 그림의 O으로 표시), 이를 제거하는 것은 task 성능에 영향을 미치지 않음을 보여준다. 구체적으로, 다음 아래의 그림은 상위 k% 최대 크기의 값만 유지하도록 각 task vector를 trimming하고 나머지를 초기 값으로 재설정 할 때(즉, task-vector에서 해당 값을 0으로 설정) 11개의 작업별 모델에 대한 평균 성능을 보여준다. 또 아래 그림은 다양한 K 값에 대한 평균 성능을 보여주며, 상위 20% 값만 유지해도 모든 매개변수를 유지하는 겻과 유사한 결과를 제공한다는 것을 보여준다. 

performance는 high-magnitude parameters 중 작은 부분에 의존한다. 각 task-vector에 대해 가장 큰 상위 k%의 parameter만 유지하고 11개 작업에 대한 평균 성능을 표시했다. 매개변수의 상위 20%만 유지를 해도, 성능이 저하되지 않는다.

이는 fine-tuning 중에 도입된 많은 parameter 변경 사항들이 중복됨을 보인다. 따라서 병합 중에 해당 값을 무시하면 작업의 성능을 손상시키지 않으면서 영향력 있는 parameter에 대한 간섭을 방지할 수 있다.

Disagreement between Parameter Signs.

서로 다른 fine-tuned model은 task-vector parameter에 반대되는 changes를 도입하여 충돌하는 부호로 인한 간섭을 유발할 수 있다. (그림 2의 정육각형) 아래 그림은 다양한 수의 모델을 병합할 때, 부호 충돌의 빈도에 대한 분석을 제시한다. 먼저 영향력 있는 parameter의 상위 20%만 유지하여 11개 task에 대한 task-vector를 trimming한다. 그런 다음 병합할 모델 수를 2개에서 11개로 늘릴 때 부호 충돌이 발생하는 parameter의 비율을 표시한다. 특히 서로 다른 task의 2개 모델만 병합하거나 동일한 Task의 여러 모델을 병합할 때도 부호 충돌이 발생하며, 병합되는 모델의 수에 따라 부호 충돌의 가능성이 증가한다.

Trimming 후에도 부호 충돌이 발생하며, 모델 수에 따라 증가하는 현상을 볼 수 있다. trimming 후 부호 충돌이 발생하는 parameter의 비율 vs 병합되는 모델의 수를 표시한 그림이다.

4. TIES-MERGING: TRIM, ELECT SIGN & MERGE

위의 문제를 해결하기 위해 다음의 방법을 제안한다.

4.1 Preliminaries

task vector $\tau_{t} \in R^{d}$는 task $t$에 대한 low loss region으로 이어지는 초기화에 대한 d-dimensional parameter space에서 요구되는 이동량 및 방향을 나타낸다. $tau_{t}$(특정 Parameter에 상응하는)의 각 원소들은 d 차원 공간의 한 축으로 간주될 수 있다. parameter 부호는 task $t$에서 손실을 감소시키는 이 축(positive or negative)에 따른 방향을 나타낸다. 따라서, 주어진 task-vector $tau_{t}$는 sign vector $\gamma_{t} \in R^{d}$와 크기 벡터 $\mu_{t} \in R^{d}$로 $\tau_{t}=\gamma_{t} \odot \mu_{t}$처럼 분해될 수 있다. $\odot$은 elementwise product이다. 전형적으로, $\gamma_{t} = sgn(\tau_{t})$, where $sgn(x) * |x| = x$ 이고, +1, 0, -1의 값을 반환한다. 크기 vector는 $\mu_{t} = | \tau_{t} |$로 정의하고 값 $\mu_{t}^{i}$는 initialization으로 부터의 i 번쨰 dimension에서 요구되는 movement를 알려준다.

4.2 Steps in TIES-MERGING

multiple task-specific models $\{ \theta_{t} \}^{n}_{t=1}$을 병합하기 위해서 먼저 그에 상응하는 task vectors $\{ \tau_{t} \}^{n}_{t=1}$을 생성한다. 이 task vectors가 주어졌을 떄, TIES-MERGING method는 병합을 수행하기 위해서 다음의 3개의 steps을 따른다.

1. Trim: 각 task $t$에 대해 top-k%의 값을 task vector의 크기에 따라 유지하고 중복 parameter의 하위 (100-k)%를 0으로 재설정하여 trimming함으로써 task vector $\tau_{t}$에서 중복 parameter를 trimming하여 $\hat{\tau}_{t}$를 생성한다. 이는 $\hat{\tau}_{t} = \hat{\gamma}_{t} \odot \hat{mu}_{t}$로 decompose할 수 있다.
2. Elect: 다음으로 우리는 서로 다른 모델에 걸쳐 parameter $p$에 대한 부호의 불일치를 해결하는 병합 model에 대한 aggregate elected sign vector $\gamma_{m}$을 생성한다. elected sign vector를 생성하기 위해 모든 관련 모델에서 가장 높은 total magnitude를 갖는 sign을 선택한다. 각 parameter $p \in {1,2, ..., d}$에 대해, 그 들의 sign(+1 or -1)에 근거하여 값 $\{\hat{\tau}_{t}^{p} \}^{n}_{t=1}$을 분리하고 그들의 합을 취하여 poistive와 negative 방향으로 total mass(i.e., total magnitude)를 계산한다. 그런 다음 total movement가 더 큰 부호로 $\gamma^{p}_{m}$을 할당한다. 이는 $\gamma^{p}_{m}=sgn(\sum^{n}_{t=1} \hat{\tau}^{p}_{t})$를 사용하여 효율적으로 계산할 수 있다.
3. Disjoint Merge: 그런 다음 각 parameter $p$에 대해 elected sign을 집계하고 그들의 평균을 계산한 것과 같은 signs을 갖는 models로 부터 parameter values 만을 유지함으로써 disjoint mean을 계산한다. 공식적으로 $A^{p} = \{t \in [n] | \hat{\gamma}^{p}_{t} = \gamma^{p}_{m} \}$ 이라 하면, $\tau^{p}_{m} = {1 \over | A^{p} |} \sum_{t \in A^{p}} \hat{\tau}^{p}_{t}$이다. disjoint mean은 항상 0 값을 무시한다.

최종 병합된 task-vector $\tau_{m}$이 주어지면, 이를 scaling하여 초기 parameter 값에 추가하여 병합된 모델 parameter $\theta_{m}$을 $\theta_{m}=\theta_{init}+ \lambda * \tau_{m}$으로 얻는다. 여기서 $\lambda$는 과거 works에서 사용된 바와 같이 scaling hyperparameter이다.

Experimental Setup