[Paper review] Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time

논문 원본 : https://arxiv.org/abs/2203.05482

Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time

 

Abstract

• 모델의 정확도를 극대화하는 일반적인 방법은 다음과 같음
  • 다양한 하이퍼파라미터를 사용하여 여러 모델을 학습 시킨 후,
  • 검증 데이터셋에서 가장 좋은 성능을 보이는 단일 모델을 선택하고 나머지는 버리는 것
• 본 논문에서는 여러 모델의 가중치를 평균화하는 방법이 더 좋은 성능을 낼 수 있다고 주장
  • "Model Soups" 라고 부름
  • 추가 비용 없음 : 여러 모델을 합쳐도 속도나 메모리 문제가 생기지 않음
  • 다양한 작업에서 효과적 : 이미지와 텍스트 관련 작업에서 모두 성능이 좋아짐
  • 새로운 데이터에도 강함 : 학습에 없던 데이터나 작업에서도 잘 작동

 

1. Introduction

• 최근 연구에서는 대규모 데이터셋으로 사전 학습된 모델이 다양한 작업에 효과적으로 전이되는 representation을 학습한다는 사실을 보여줌
• 이로 인해, 머신러닝 연구자들은 downstream task를 해결하기 위해 대규모 사전 학습 모델을 fine-tuning 하는 방식이 보편화됨
  • 다양한 하이퍼파라미터 설정으로 모델을 미세 조정
  • 검증 데이터셋에서 가장 높은 정확도를 기록한 단일 모델을 선택하고 나머지는 버림
• 그러나 단일 모델을 선택하고 나머지를 버리는 방식에는 몇 가지 단점 존재
  • 여러 모델의 결과를 앙상블하면 단일 모델보다 더 나은 성능을 낼 수 있지만, 이 과정에서 inference 시 높은 계산 비용이 발생
  • downstream task에 모델을 fine tuning하면 분포 밖 데이터에서의 성능이 감소할 수도 있음
  • 검증 데이터에서 가장 좋은 모델이 실제 분포 밖 데이터에서는 최고의 모델이 아닐 수도 있음
• 본 논문에서는 독립적으로 fine-tuning된 여러 모델의 가중치를 평균화(weight averaging)하여 "Model Soup"라고 부르는 결과물을 만듦
• 이 방법은 추가적인 학습이 필요하지 않고, 추론 시 추가적인 비용도 발생하지 않음

• ImageNet과 같은 데이터셋에서 CLIP, ALIGN, ViT-G와 같은 대규모 사전 학습 모델을 사용해 Model Soup 실험
• 그 결과, Model Soup는 단일 모델보다 더 나은 성능을 자주 보여주며, ImageNet에서 90.94%의 top-1 정확도를 기록
• 이는 CoAtNet 모델이 달성한 이전 최고 기록(90.88%)을 능가했으며, 추론 시 25% 적은 FLOPs가 필요했음

• Model Soup의 다양한 구현
• 단순한 방식은 모든 모델의 가중치를 균등하게 평균화하는 것
• 그러나 본 논문은 Greedy Soup라는 방식을 사용하여 성능을 더욱 항샹시킴
  • 검증 성능이 향상되는 모델만 골라 가중치를 평균화

 

 

2. Method

 

   Model Soup의 기본 개념

• 여러 하이퍼라미터 구성 h1,h2, ..., hk로 학습된 모델 θ_i​=FineTune(θ₀​,h_i​)가 있다고 가정
• Model Soup f(x, Q_s)는 여러 모델의 θ_i의 가중치 평균을 사용

 

   Uniform Soup

• 모든 fine-tuning된 모델의 가중치를 단순 평균화하여 생성
• 즉, S = {1,...,k}로 설정하여 모든 모델을 포함시킴
• 하지만 일부 하이퍼파라미터 구성으로 생성된 모델이 낮은 정확도를 기록할 수 있음
• 이런 모델들이 포함되면 Uniform Soup의 정확도가 낮아질 수 있음

 

   Greedy Soup

• 이 문제를 해결하기 위해 Greedy Soup 사용
• Greedy Soup는 모델을 "재료"로 순차적으로 추가하며,
• 검증 데이터 성능이 향상될 경우에만 Soup에 모델을 포함
  • 모델을 검증 데이터 정확도 순으로 내림차순 정렬
  • 순서대로 모델을 하나씩 Soup에 추가
  • 추가 후 검증 데이터 성능이 향상되면 포함, 그렇지 않으면 제외
• 이 방식으로 만들어진 Greedy Soup는 검증 데이터에서 최고의 단일 모델 성능보다 나쁘지 않음

 

   Learned Soup

• 모델의 가중치 평균화 비율을 최적화하는 방법
• 이를 위해 gradient-based minibatch optimization을 사용
• 이 방법은 더 정교하지만, 모든 모델을 메모리에 동시에 로드해야 하기 때문에 대규모 네트워크에서는 사용이 어려움

 

3. Experiments

 

   3.1 Experimental setup

• CLIP,  ALIGN, BASIC
  • 이들은 이미지-텍스트 쌍에서 contrastive supervision으로 사전 학습된 모델
• ViT-G/14 
  • JFT-3B에서 사전 학습된 모델
• Transformer 
  • 텍스트 분류를 위한 사전 학습 모델
• 실험에서는 대부분 CLIP ViT-B/32 모델 사용
• fine-tuning은 end-to-end 방식으로 수행되며, 이는 마지막 선형 계층만 학습하는 방식보다 더 높은 정확도 제공

 

• 마지막 선형 계층의 초기화 방법
  • Linear Probe (LP) 초기화
    • 사전 학습된 모델의 기존 계층은 frozen하고, 마지막 선형 계층만 학습하는 초기화 방법
  • Zero-shot 초기화
    • CLIP이나 ALIGN 모델의 Text Tower에서 생성된 분류기 사용
  • 두 초기화 방법 모두 비슷한 결과를 보였으므로, 특별한 경우가 아니면 LP 초기화 사용

 

• 앙상블 기준에서는 logit (정규화되지 않은 출력)을 평균화
• fine-tuning은 cross-entropy loss를 사용하며, 기본적으로 ImageNet에서 수행
• ImageNet fine-tuning 시, 다음 다섯 가지 distribution shift에서 평가를 진행
  • ImageNetV2 
  • ImageNet-R 
  • ImageNet-Sketch 
  • ObjectNet
  • ImageNet-A 

• 결과는 이 다섯 가지 distribution shift 평균으로 보고하는 경우가 많음

 

• Greedy Soup 구성
• ImageNet의 공식 검증 데이터셋은 테스트 셋으로 사용되기 때문에, ImageNet 학습 데이터셋의 약 2%를 따로 떼어 검증 데이터셋으로 사용해 Greedy Soup를 구성

 

   3.2 Intuition and motivation

 

   Error landscape visualizations

• CLIP 모델을 ImageNet에서 fine-tuning하는 과정에서, 학습 손실과 테스트 오류의 error landscape를 시각화
• 사전 학습된 초기 가중치 θ0를 사용해 두 번 독립적으로 fine tuning하여 θ1과 θ2를 얻었음
• 이 세 점(θ0, θ1, θ2)을 연결하는 평면 상에서 학습 손실과 테스트 오류 계산

• 오류/손실 그래프가 "basin-shaped" 형태를 보임
  • 이는 두 모델의 가중치를 선형 보간(linear interpolation)하면 개별 모델보다 성능이 향상될 가능성을 시사
• 가중치 보간의 정확도는 두 모델 간의 각도에 따라 영향을 받음
  • 두 모델이 서로 독립적 (90도에 가까운 각도)일수록 정확도가 더 높아질 가능성이 큼

 

• 학습률, 랜덤 시드, 데이터 증강을 다르게 설정하여 여러 모델을 학습
• 각 모델 쌍 (θ1, θ2)에 대해, 가중치 평균화 후의 정확도와 개별 모델의 평균 정확도를 비교
• 두 모델 간의 각도(φ)와 정확도 개선 간의 상관관계 분석
• 가중치 평균화 후 정확도의 개선 정도는 두 모델 간 각도와 강하게 상관됨

 

 

   Ensemble comparison

• 모델 쌍 (θ1, θ2)을 무작위로 선택하고,
• 학습률에 따른 앙상블(로짓 평균화)과 Soup(가중치 평균화) 성능을 비교

• 학습률이 낮은 경우, 앙상블과 Soup 성능이 비슷하지만 둘 다 최적의 성능은 아님
• 학습률이 적당히 높을 경우, 앙상블과 Soup 모두 높은 성능을 보임
• 학습률이 지나치게 높으면 , 앙상블 성능이 Soup보다 우수함
• ImageNet에서는 앙상블이 더 높은 성능을 보이지만, 분포 이동 데이터셋에서는 Soup가 더 나은 성능을 보임

 

   One dimensional hyperparameter grids

• 특정 하이퍼파라미터(ex.학습률) 범위 ha,....,hb에서,
  • 범위의 양 끝 값(ha와 hb)으로 학습된 모델 평균과,
  • 범위 안에서 가장 좋은 단일 모델의 성능을 비교
• 대부분의 경우, 범위 양 끝 점에서 학습한 모델의 평균이 단일 모델보다 성능이 우수

 

   3.3 Model soups

• 다양한 하이퍼파라미터를 사용해 미세 조정된 모델을 평균화하면, 검증 데이터에서 가장 좋은 단일 모델을 선택하는 기존 방법보다 우수한 성능을 보임

 

• CLIP, ALIGN 모델: ImageNet에서 미세 조정.
• ViT-G 모델: JFT-3B로 사전 학습, ImageNet에서 미세 조정.
• Transformer 모델: 텍스트 분류 작업에서 미세 조정

 

   Fine-Tuning CLIP And ALIGN

• 사전 학습된 CLIP ViT-B/32와 ALIGN EfficientNet-L2 모델을 ImageNet에서 미세 조정
• 하이퍼파라미터 탐색을 통해 여러 개의 미세 조정된 모델을 생성하고, 이를 조합한 Greedy Soup가 개별 모델보다 더 나은 성능을 낼 수 있는지 확인

 

 

• CLIP 모델
  • 하이퍼파라미터 설정
    • 학습률, weight decay, 학습 epoch 수, label smoothing, 데이터 증강.
  • random search을 통해 72개의 fine tuning 모델 생성
• ALIGN 모델
  • 하이퍼파라미터 설정
    • 학습률, 데이터 증강, mixup.
  • grid search을 통해 12개의 fine tuning 모델 생성
• Greedy Soup 구성
  • 모든 모델을 검증 데이터셋 기준으로 정확도 내림차순으로 정렬
  • Recipe 1(Greedy Soup 방법)을 적용하여 정확도가 향상되는 모델만 선택
  • 결과적으로, CLIP와 ALIGN 모두 5개의 모델을 선택

 

 

• CLIP와 ALIGN의 성능 비교
  • Greedy Soup가 하이퍼파라미터 탐색에서 가장 좋은 개별 모델보다 성능이 개선됨
    • CLIP: +0.7%
    • ALIGN: +0.5%
• 다양한 모델 수에서 성능 비교
  • CLIP 모델을 랜덤 순서로 추가하며 Uniform Soup와 Greedy Soup 성능 비교
  • 결과
    • Greedy Soup는 ImageNet에서 Uniform Soup보다 더 나은 성능
    • 분포 이동 데이터셋에서는 둘의 성능이 비슷
• Logit 앙상블과의 비교
  • ImageNet에서는 Logit 앙상블이 Greedy Soup보다 더 나은 성능
  • 하지만, 분포 이동 데이터셋에서는 Greedy Soup가 우세

 

 

• Greedy Soup의 효과
  • 추가 학습 없이, 기존 개별 모델보다 성능이 향상됨
• Uniform Soup 한계
  • 개별 모델 성능이 높을 때만 유의미한 성능 향상
  • 학습률이 높은 모델들은 성능이 낮아 Greedy Soup에서 제외됨

 

 

 

 

   

   Fine-Tuning A ViT-G Model Pre-Trained On JFT-3B

 

 

• JFT-3B에서 사전 학습된 ViT-G/14 모델을 사용해 ImageNet에서 fine tuning
• Greedy Soup가 최첨단 기법과 결합했을 때 성능 향상이 가능한지 확인

 

 

• 하이퍼파라미터 탐색
  • 학습률, decay 스케줄, 손실 함수, 데이터 증강(minimum crop size, RandAugment, mixup, CutMix).
  • 58개의 모델 학습
  • 일부 모델에는 Sharpness-Aware Minimization(SAM) 적용

• EMA(Exponential Moving Average)
  • 모델 학습 중 EMA로 가중치를 저장
    • 낮은 EMA(0.999)와 높은 EMA(0.9999999) 비교
  • 높은 EMA는 단일 모델 성능이 더 좋지만, Greedy Soup는 낮은 EMA로 학습된 모델에서 더 나은 성능
• 평가 방식
  • ImageNet 검증 데이터와 다섯 가지 분포 이동 데이터셋 사용
  • 추가로 ReaL 및 multilabel ImageNet 데이터셋에서도 평가

 

• Greedy Soup의 성능
  • 58개 모델 중 14개의 모델을 선택해 구성한 Greedy Soup는, 개별 모델보다 거의 모든 데이터셋에서 더 나은 성능을 보임
  • 단, ObjectNet 데이터셋에서는 개별 모델이 더 나은 성능
• Oracle(성능 기준 비교)
  • 개별 모델을 각 테스트 셋 성능 기준으로 선택해도, Greedy Soup가 대부분의 데이터셋에서 우수한 성능.
• Logit 앙상블과의 비교
  • ImageNet에서 Greedy Soup와 Logit 앙상블은 비슷한 성능
  • 그러나 분포 이동 데이터셋에서는 Greedy Soup가 더 나은 성능

 

 

 

   

   Fine-Tuning On Text Classification Tasks

• 이미지 분류 외에 자연어 처리(NLP) 에서도 Model Soup가 효과적인지 확인

 

• 사용한 모델
  • BERT와 T5 모델
• 데이터셋
  • GLUE 벤치마크의 4개 텍스트 분류 작업
    • MRPC, RTE, CoLA, SST-2.
• 하이퍼파라미터 탐색
  • 학습률, 배치 크기, 학습 epoch, 랜덤 시드
  • 각 데이터셋마다 32개 모델 학습
• 텍스트 분류 작업에서도 Greedy Soup가 개별 모델보다 성능이 향상됨
• 성능 향상 폭은 이미지 분류보다 작지만, 유의미한 결과를 보임

 

 

 

 

4. Analytically comparing soups to ensembles

• Soup
• 모델의 가중치를 선형 보간

• Ensemble
• 모델의 출력(로짓)을 선형 결합

 

• Soup error
• θα​로 생성된 수프의 분류 오류율

• Ensemble error
• 문헌에 따르면, 앙상블 오류는 보통 min{err0,err1}보다 작음
• 따라서, Soup error가 ensemble error와 비슷하면, Soup가 개별 모델보다 성능이 좋을 가능성이 높음

 

• 여러 학습률, 데이터 증강 전략, 랜덤 시드를 사용해 모델을 학습
• α 값을 조정하며 Soup와 Ensemble의 성능 비교

• 학습률이 낮거나 보통일 때, Soup와 Ensemble 성능이 유사
• 학습률이 높으면, Soup가 덜 효과적

 

 

 

5. Scope and limitations

 

   Applicability

 

• Model Soup는 대규모, 다양한 데이터셋으로 사전 학습된 모델에서 주로 효과적
• ImageNet-22k 사전 학습 모델에서는 여전히 성능 향상이 있었지만, CLIP 및 ALIGN처럼 큰 향상을 보이지는 않음

 

   Calibration

• Ensemble은 모델의 calibration을 개선(예측 확률의 신뢰도 향상)하는 반면, Model Soup는 동일한 효과를 가지지 않음
• 동일한 하이퍼파라미터를 사용하고 랜덤 시드만 다른 20개의 모델로 Soup와 Ensemble을 비교
• Soup는 Cailbration에서 앙상블보다 성능이 낮음

 

7. Conclusion

• 본 논문의 결과는 fine tuning 시 검증 세트에서 가장 성능이 좋은 모델 하나를 선택하는 기존의 접근 방식이 항상 최선은 아닐 수 있다는 점을 지적
• 대신, 여러 번 fine tuning한 모델의 가중치를 평균해서 새로운 모델을 만들면, 별도의 계산 비용을 추가하지 않으면서도 더 나은 성능을 보이는 경우가 있는 것을 확인