Model Soups

모델 수프는 동일한 사전 학습 모델에서 서로 다른 하이퍼파라미터로 파인튜닝한 여러 모델의 가중치를 평균해 만든 단일 모델입니다. 일반적으로 하이퍼파라미터 탐색 후에는 검증 점수가 가장 높은 실행 결과만 남기고 나머지는 버립니다. 수프는 버려진 실행 결과를 재활용합니다. 각 모델이 동일한 사전 학습 모델 가중치에서 시작했으므로 파라미터가 여전히 서로 대응하며, 이들의 평균은 종종 각 모델보다 높은 점수를 내면서도 실행 비용은 단일 모델과 정확히 같습니다. 워싱턴 대학교, Google, Meta AI, 텔아비브 대학교, 컬럼비아 대학교의 Mitchell Wortsman과 동료들은 ICML 2022에서 모델 수프: 여러 파인튜닝 모델의 가중치 평균으로 추론 시간 증가 없이 정확도 향상을 발표하며 이 방법을 소개했습니다. Google의 JFT-3B 데이터셋으로 사전 학습한 ViT-G/14 모델 58개의 파인튜닝 결과를 탐욕적 수프로 결합한 결과, ImageNet에서 top-1 정확도 90.94%를 기록해 당시 최고 성능을 달성했습니다.

균일 수프와 탐욕적 수프#

파인튜닝은 대규모 사전 학습 네트워크를 가중치 공간에서 조금만 이동시키므로 서로 다른 학습률, 데이터 증강, 레이블 스무딩, 학습 기간으로 파인튜닝한 모델은 대체로 낮은 손실을 보이는 동일한 넓은 영역에 자리 잡습니다. 모델을 평균하면 그 영역의 중앙에 위치하게 되어 가장자리보다 일반화 성능이 좋아집니다. 논문은 두 가지 주요 방법을 정의합니다.

• 균일 수프: 파인튜닝한 모든 모델의 평균을 구합니다. 모든 구성 모델의 성능이 우수하면 효과적이지만, 학습률이 매우 높았던 실행 결과는 공통 영역을 벗어날 수 있고 이를 평균에 포함하면 정확도가 낮아집니다.
• 탐욕적 수프: 모델을 홀드아웃 검증 데이터의 정확도에 따라 정렬하고, 가장 우수한 모델부터 시작해 평균 수프의 검증 정확도가 떨어지지 않을 때만 다음 모델을 추가합니다. 가장 우수한 모델에서 시작하므로 탐욕적 수프의 검증 세트 성능은 해당 모델보다 나빠지지 않습니다.

CLIP ViT-B/32를 72개의 무작위 하이퍼파라미터 설정으로 ImageNet에 파인튜닝했을 때, 단일 최고 모델의 점수는 80.38%였습니다. 72개 모델 전체의 균일 수프는 79.97%를 기록했습니다. 일부 실행 결과가 공통 저손실 영역에서 벗어났기 때문입니다. 반면 탐욕적 수프는 모델 5개를 선택해 81.03%를 기록했습니다. ImageNet 스케치와 렌더링을 포함한 5개의 분포 변화 테스트 세트에서 탐욕적 수프는 평균 50.75%를 기록해 단일 최고 모델의 47.83%보다 높았습니다. 논문은 ALIGN에서도 성능 향상을 보고했으며, 텍스트 분류에 BERT와 T5를 파인튜닝했을 때는 더 작은 향상을 보고했습니다.

이 PyTorch 예제는 8x8 손글씨 숫자로 작은 분류기를 사전 학습하고, 서로 다른 학습률과 입력 노이즈 증강을 적용해 복사본 9개를 파인튜닝한 다음, 이들의 균일 수프와 9개 모델 앙상블을 비교합니다.

import copy

import torch
from sklearn.datasets import load_digits
from torch import nn

torch.manual_seed(0)
X, y = load_digits(return_X_y=True)
order = torch.randperm(len(X))
X, y = torch.tensor(X / 16, dtype=torch.float32)[order], torch.tensor(y)[order]


def train(model, first, last, lr, noise, seed):
    """Train a copy of the model on random batches from rows first to last, adding input noise."""
    model, generator = copy.deepcopy(model), torch.Generator().manual_seed(seed)
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    for _ in range(300):
        rows = torch.randint(first, last, (32,), generator=generator)
        optimizer.zero_grad()
        inputs = X[rows] + noise * torch.randn(32, 64, generator=generator)
        nn.functional.cross_entropy(model(inputs), y[rows]).backward()
        optimizer.step()
    return model


def accuracy(logits):
    """Return the fraction of the 997 test digits classified correctly."""
    return (logits.argmax(1) == y[800:]).float().mean().item()


base = train(nn.Sequential(nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 10)), 0, 300, 0.05, 0, 0)  # pretrained
sweep = [(lr, noise) for lr in (0.01, 0.03, 0.1) for noise in (0, 0.2, 0.4)]  # 9 hyperparameter settings
ingredients = [train(base, 300, 500, lr, noise, seed) for seed, (lr, noise) in enumerate(sweep)]

# Uniform soup: average every weight tensor of the nine fine-tuned models
soup = copy.deepcopy(base)
soup.load_state_dict({k: sum(m.state_dict()[k] for m in ingredients) / 9 for k in soup.state_dict()})

with torch.no_grad():
    singles = [accuracy(m(X[800:])) for m in ingredients]
    print(f"Fine-tuned models: {min(singles):.3f} to {max(singles):.3f}, mean {sum(singles) / 9:.3f}")
    print(f"Uniform soup: {accuracy(soup(X[800:])):.3f}")
    print(f"Ensemble of 9: {accuracy(sum(m(X[800:]).softmax(1) for m in ingredients)):.3f}")

파인튜닝된 9개 모델은 홀드아웃 숫자 997개에서 0.881부터 0.936까지의 점수를 기록하며, 평균 점수는 0.917입니다. 균일 수프의 점수는 0.934으로 단일 최고 실행 결과에 가깝고, 앙상블의 0.930보다 약간 높습니다. 앙상블은 이미지당 순전파를 9회 수행해야 하지만 수프는 1회만 수행합니다.

실제 환경에서의 수프 활용#

수프는 비전 트랜스포머(ViT) 또는 다른 파운데이션 모델을 새로운 이미지 분류 태스크에 전이 학습하는 경우처럼, 어차피 하이퍼파라미터 탐색을 수행하는 환경에서 가장 유용합니다. Ultralytics YOLO 하이퍼파라미터 튜닝도 이러한 탐색의 한 예입니다. model.tune()는 변형된 하이퍼파라미터 세트마다 모델 하나를 학습하며, 저장 공간을 절약하기 위해 현재 최고 결과가 아닌 실행 결과는 삭제하고 최고 반복의 실행 결과와 가중치만 남깁니다. 수프는 이처럼 최고 결과만 고르는 단계를 대체합니다.

구성 모델은 아키텍처와 사전 학습 시작점을 공유해야 합니다. 서로 다른 무작위 초기화에서 학습된 네트워크는 동일한 특징을 서로 다른 은닉 유닛에 저장하므로, 이들을 평균하면 정확도가 떨어집니다. 최종 가중치는 추론 시 추가 비용이 들지 않으므로 수프의 추론 지연 시간과 메모리 사용량은 단일 모델과 같습니다.

모델 수프와 관련 방법 비교#

• 모델 앙상블: 여러 모델의 예측을 평균하므로 구성 모델이 추가될 때마다 추론 비용이 증가합니다. CLIP ViT-B/32 탐색에서 로짓 앙상블은 ImageNet에서 81.19%를 기록해 탐욕적 수프의 81.03%보다 약간 높았지만, 이미지마다 모든 구성 모델을 실행해야 했습니다.
• 모델 병합: 가중치를 결합하는 방법을 포괄하는 상위 개념입니다. 수프는 모든 모델을 동일한 태스크에 파인튜닝한 경우이며, 태스크 산술은 서로 다른 태스크에 파인튜닝한 모델을 더합니다.
• 단일 실행 내 가중치 평균: 확률적 가중치 평균(SWA)과 가중치의 지수 이동 평균(EMA)은 단일 학습 궤적을 따라 체크포인트를 평균합니다. 수프는 독립적인 실행 결과를 평균합니다. 논문은 이미 EMA 또는 SWA를 사용하는 모델을 수프로 결합해도 정확도가 올라간다고 밝혔으므로 두 평균화 방식의 효과가 더해집니다.
• WiSE-FT: 제로샷 모델과 파인튜닝 버전 사이를 보간하는 두 구성 요소의 수프로, 분포 변화에 대한 사전 학습 모델의 강건성을 더 많이 유지합니다.