논문 리뷰) LM-Cocktail: Resilient Tuning of Language Models via Model Merging

[총평]
- 간결하고, 심플한 구현 방법인 merge로 새로운 도메인 학습 추가가 가능함을 실험적으로 보임.
- FlagEmbedding에서 라이브러리로 구현이 되어 있음. 심지어 잘 먹힘.
- 임베딩과 LLM으로 모두 실험하였는데, LLM에서 조금 더 기존 도메인 성능 유지가 잘된 것 같음. 

---

[Abstract]
사전 학습된 언어 모델은 다운스트림 애플리케이션을 더 잘 지원하도록 지속적으로 미세 조정(fine-tuning)됩니다. 그러나 이 과정은 특정 도메인을 목표로 할 경우, 일반적인 작업에서 성능 저하를 초래할 수 있습니다. 이러한 문제를 해결하기 위해 LM-Cocktail을 제안하며, 이를 통해 미세 조정된 모델이 일반적인 작업에서도 강건함을 유지할 수 있도록 합니다.

해당 방법은 모델 병합(model merging) 방식으로 이루어지며, 미세 조정된 언어 모델을 사전 학습된 기본 모델 또는 다른 도메인의 모델과 가중 평균(weighted average) 방식으로 병합합니다. 이러한 방식은 단순하지만 놀랍도록 효과적입니다. 병합된 모델은 특정 도메인에서의 우수한 성능을 유지하면서도, 전체적인 일반 작업에서도 강력한 성능을 발휘할 수 있습니다.

 LLama 및 BGE 모델을 대상으로 FLAN, MMLU, MTEB 등의 대표적인 벤치마크에서 광범위한 실험을 수행하였으며, 그 결과는 해당 방법론이 효과적임을 입증합니다.

 

 

[Introduction]
본 연구에서는 파국적 망각(catastrophic forgetting) 문제를 해결할 수 있는 효과적인 프레임워크를 설계하는 것을 목표로 합니다. 이를 통해 미세 조정된 언어 모델이 일반적인 작업에서도 강건성을 유지할 수 있도록 합니다. 또한, 새로운 프레임워크가 더 실용적이어야 한다는 점도 고려하여, 단순하면서도 일반적인 모델 학습 워크플로우와 완벽하게 호환될 수 있도록 설계하고자 합니다.

이러한 고려를 바탕으로, LM-Cocktail이라는 새로운 접근 방식을 제안합니다. 이 방법은 미세 조정된 언어 모델을 모델 병합(model merging) 방식으로 지속적으로 적응시켜 강건성을 확보하는 전략입니다. LM-Cocktail은 **일반적인 패러다임(general paradigm)**으로, 여러 가지 조건에서 적용될 수 있습니다.

가장 간단한 형태에서는 미세 조정된 모델을 사전 학습된 기본 모델과 직접 병합하여, 미세 조정된 모델의 일반적 성능을 향상시킵니다.
더 나아가, **다른 일반 도메인에서 미세 조정된 모델들(peer models)**을 추가로 포함시켜, **few-shot 검증 샘플(few-shot validation examples)**을 활용하여 병합 가중치를 추정함으로써 더욱 강력한 성능을 발휘할 수 있습니다.
심지어 미세 조정 데이터가 없는 경우에도, 사전 학습된 기본 모델과 다른 일반 도메인에서 미세 조정된 모델들을 병합하는 방식으로 경쟁력 있는 강건성을 유지할 수 있습니다.

 

 

2. LM-Cocktail 방법론 소개
[일반적인 형태의 모델 병합 (General Paradigm)]

미세 조정된 모델은 목표 도메인 t 이외의 일반적인 도메인에서 성능 저하(파국적 망각, catastrophic forgetting) 현상을 겪을 가능성이 높습니다. LM-Cocktail의 목표는 특정 작업에 대한 미세 조정을 수행하면서도 모델의 일반적 성능을 유지하는 것입니다.

LM-Cocktail의 핵심은 동일한 아키텍처를 가지지만 서로 다른 가중치를 지닌 여러 모델을 병합하여 단일 모델로 통합하는 것입니다. 이러한 방식을 통해, 병합된 모델은 개별 모델들의 장점을 효과적으로 결합하여 강건한(resilient) 성능을 유지할 수 있습니다.

LM-Cocktail에서 효과적인 모델 병합을 위해 해결해야 할 두 가지 핵심 문제가 있습니다.

1. 어떤 모델 그룹을 병합할 것인가?
2. 병합 가중치는 어떻게 결정할 것인가?

LM-Cocktail의 병합 함수는 다음과 같이 정의됩니다.

여기서 Mr은 **강건하게 조정된 모델(resilient-tuned model)**입니다.
α는 하이퍼파라미터로, 기본값은 0.5입니다.
𝑤𝑖는 **병합 가중치(merging weight)**이며,  𝑤𝑖의 총합은 1로 정규화됩니다.

이때, 각 𝑤𝑖는 Target Domain의 평가 데이터 점수를 기반으로 가중치를 결정합니다.
가중치는 아래와 같이 계산됩니다.
목표 도메인에서 손실이 클수록(성능이 낮을수록), 해당 모델에 할당되는 가중치가 작아지는 방식입니다.

위 식에서 𝐸𝑡는 Few-shot 샘플로 목표 도메인에서 따로 보관한 작은 데이터셋입니다. 

이는 병합 가중치를 조정하는 데 사용됩니다. 

실험적으로 5-shot 샘플만으로도 다양한 설정에서 충분히 강력한 성능을 발휘합니다.

 

 

[변형된 형태의 모델 병합 (Variation)]

 LM-Cocktail에는 앞서 설명한 형태 외에 추가적인 2가지의 Variation이 있을 수 있습니다.

 

LM-Cocktail에서 제안하는 Merge 방식 3가지.
1) General 모델 병합 (앞에서 설명한 방식)
2) 단일 전문가(Mono-Specialist) 모델 병합
3) 미세 조정 없이 모델만 병합 (Without Fine-tuning)

 

 1. 단일 전문가(Mono-Specialist) 모델 병합
일반 도메인에서 다양한 미세 조정 모델이 없는 경우, merge 함수는 기본 모델 
𝑀𝑏과 목표 도메인의 미세 조정 모델 𝑀𝑡만을 결합하는 형태로 단순화됩니다.

 

 

2. 미세 조정 없이 모델 병합 (Without Fine-tuning)

목표 도메인에서 미세 조정을 수행할 데이터나 연산 자원이 부족한 경우, 
merge 함수는 기본 모델과 일반 도메인에서 미세 조정된 모델 그룹을 그대로 활용하는 형태가 됩니다.

이 경우 목표 도메인에서 학습된 모델 없이도, 
기존 모델들을 효과적으로 결합하여 다운스트림 작업에 적합한 모델을 생성할 수 있습니다.
여기서도 General Merge 포맷과 동일하게, Few-shot 샘플로 Merge 가중치(𝑤𝑖)를 결정합니다.

 

 

 

[3. 실험 설정 (Experimental Setup)]
디코더 기반 LM(Decoder-based LM)과 인코더 기반 LM(Encoder-based LM) 두 가지 유형의 모델을 사용하여 실험을 진행했습니다.
각 모델 유형별로 n개의 미세 조정된 모델로 학습한 후,
기존 미세 조정 모델의 성능과 LM-Cocktail을 적용한 모델(Resilient-tuned model)의 성능을 비교 평가했습니다.


3.1 디코더 기반 LM (Decoder-based LM)의 실험 설정
• 기본 모델 (Base Model)
  - Llama-2-chat-7b (Touvron et al., 2023) 모델을 사용
  - 다양한 작업에서 뛰어난 zero-shot 성능을 보이는 모델

• 미세 조정 (Fine-tuning)
  - (Cheng et al., 2023; Wang et al., 2023)에서 수집한 FLAN (Wei et al., 2022) 데이터셋 활용
  - FLAN 데이터셋 내 30개 작업 중 9개 작업을 선정하여 미세 조정
    - NQ (Natural Questions)
    - SQuAD (Stanford Question Answering Dataset)
    - Hellaswag
    - SST2 (Stanford Sentiment Treebank)
    - Winogrande
    - CommonGen
    - MRPC (Microsoft Research Paraphrase Corpus)
    - AG News
    - MNLI (Multi-Genre Natural Language Inference)
  - 학습 코드: FastChat 패키지 기반
  - 학습 하이퍼파라미터 
     - 학습률: 2e-5
     - 배치 크기 128
     - 최대 에포크 수: 3

• 평가 (Evaluation)
 - (Cheng et al., 2023; Wang et al., 2023)에서 수집한 30개 테스트 데이터셋으로 평가
 - 미세 조정된 작업(위 9개)의 테스트 데이터도 포함
 - 개별 작업의 평가 지표는 (Wang et al., 2023)을 참고
 - 추가적으로 MMLU 데이터셋(대형 언어 모델 성능 평가의 표준 벤치마크)에서 추가 실험 수행



3.2 인코더 기반 LM (Encoder-based LM)
• 기본 모델 (Base Model)
  - bge-base-v1.5 (Xiao et al., 2023) 모델을 사용
  - 텍스트를 임베딩(embedding) 표현으로 변환하는 임베딩 모델

• 미세 조정 (Fine-tuning)
  - Sentence Transformers 저장소에서 9개 데이터셋을 선정하여 미세 조정
    - GooAQ
    - Yahoo Answers
    - MSMarco
    - Stack Exchange
    - ELI5
    - SQuAD
    - AmazonQA
    - Quora
    - HotpotQA
  - FlagEmbedding 도구를 사용하여 BGE 모델 미세 조정
  - 학습 하이퍼파라미터
    - 최적화 알고리즘: AdamW
    - 학습률: 2e-5
    - 배치 크기: 256
 - 대조 학습(Contrastive Learning) 온도 값: 0.02

• 평가 (Evaluation)
 - MTEB 벤치마크 (Muennighoff et al., 2022)의 15개 검색(retrieval) 작업에서 모델 평가
 - 평가 지표: NDCG@10
 - 미세 조정에 사용된 3개 작업(MSMarco, HotpotQA, Quora)의 테스트 데이터도 포함
 - 다양한 작업 간의 학습 및 평가를 용이하게 하기 위해 (Xiao et al., 2023)의 기본 쿼리 지시문(query instruction)은 사용하지 않음

 

 

[4. 평가 결과(Experimental Results)]

디코더 모델에서 효과가 좋았음.
방법 1과 방법 2 모두 망각이 없거나 적은 상태로, 새로운 도메인의 성능이 올라감.

 

LLM 만큼은 아니지만, 효과가 있었음. 
개인적으로 스코어링으로 확인해 볼 도메인이 적은게 아쉬움.

 

각 도메인 별로 특화된 Merging을 구현하지 않고, 그냥 Unified시킨 모델도 성능이 좋음.
(Besides, we aggregate all examples from each task to compute merging weights, and produce a unified model named LM-Cocktail(u) for all tasks, rather than generate a separate model for each task)

 

 

[라이브러리 구경가기..]
FlagEmbedding에서 아주 심플하게, 잘 구현되어 있다.
merge 가중치(w) 별로 모델 parameter를 진짜 그냥 합친다...