[논문리뷰] Model Stock: All We Need Is Just a Few Fine-Tuned Models

논문 정보

• Title: Model Stock: All We Need Is Just a Few Fine-Tuned Models
• Authors: Dong-Hwan Jang, Sangdoo Yun, Dongyoon Han (NAVER AI Lab)
• Published: 2024.09.30
• Pages: pp 207–223
• Conference: ECCV 2024
• Link: https://doi.org/10.1007/978-3-031-72784-9_12

 

Background

Merge란?

모델 N개의 가중치를 결합하여 하나의 모델로 만드는 방법론입니다.

class SimpleModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super(SimpleModel, self).__init__()
        # 첫 번째 레이어: Linear 레이어
        self.fc1 = nn.Linear(784, 64)
        # 두 번째 레이어: Linear 레이어
        self.fc2 = nn.Linear(64, 2)

    def forward(self, x):
        # 첫 번째 레이어에 ReLU 활성화 함수 적용
        x = F.relu(self.fc1(x))
        # 두 번째 레이어에 소프트맥스 활성화 함수 적용
        x = F.log_softmax(self.fc2(x), dim=1)
        return x

1개의 모델은 N개의 레이어로 이루어져 있습니다.

위의 예시를 기반으로 첫번째 레이어를 출력해보면 다음과 같이 이루어져 있습니다.

print("Weights:", linear_layer.weight)
print("Bias:", linear_layer.bias)

Weights: tensor([[ 0.0123, -0.0345, ...,  0.0567],...,[0.0789, -0.0123, ...,  0.0456]], requires_grad=True)
Bias: tensor([ 0.0123, -0.0345, ...,  0.0567], requires_grad=True)

가중치는 (64,784)의 크기를 가지는 2차원 tensor로 이루어져있고,
Bias는 (64)의 크기를 가지는 1차원 tensor로 이루어져 있습니다.

Model Merge는 N개의 모델에서 같은 레이어를 기준으로 가중치(Weights)를 결합하는 방법이고,

단순히 가중치의 평균으로 만들지, 수식을 적용할지 등등이 각각의 merge 방법론들 입니다.

 

Model Soup이란?

Model Stock을 알기 전, Soup 방법론을 알아야합니다.

Model Merge 방법론 이전에는 Model Ensemble 방법론이 존재했습니다.

단순히 모델들을 ensemble 하는 방안은 모델이 증가할 수록 시간복잡도가 증가되는 O(n) 형식이였습니다.

이를 해결하기 위해 나온 방법론이 Model Soup이며, 다수의 파인튜닝 모델들의 가중치를 단순히 평균화 하여 하나의 모델로 만드는 방법입니다.

 

k = merge할 모델 개수

• Basic Soup
  • weight_1/k +weight2/k + ... +weight_k/k
• Greedy Soup
  • 모델들을 val acc 순 내림차순 진행
  • sota 모델부터 하나씩 weight_1/k +weight2/k 진행
  • 성능이 낮아지기 전까지 병합
• Learned Soup(Advanced)
  • 모든 모델을 메모리에 올려서 조합을 한번에 시도(요구 자원 높음)

 

Soup 방법론은 기존 Ensemble의 단점을 해결하였으나 다음과 같은 한계가 있습니다.

• 많은 모델 필요 (논문에서는 48개의 모델 사용)
• 튀는 가중치 있다면 성능이 하락함 (같은 loss basin에 있어야함)

 

Model STOCK이 생겨난 과정

Soup의 방법론은 기존 Ensemble의 단점을 해결하고 성능 향상을 이루었습니다.

다만 앞서 말한 한계점들이 존재하였고, 이는 곧 실제 현실세계에서 사용하기 비효율적인 방법론입니다.

따라서, 여러개의 모델을 평균화해서 중심을 찾는 것 말고, 최소한의 모델(2개)를 가지고 수학적인 수식을 이용하여 최대한 가중치 중심에 가깝게 만들어버리는 것을 목표로 하는게 Stock 방법론입니다.

 

관측 1

Fig.2는 랜덤시드로 학습된 파인튜닝된 모델들의 레이어를 시각화한 사진입니다.
각 레이어에서 50개의 별개의 웨이트를 뽑아서 각도랑 Norm을 비교해보니 , 각 레이어별로 가중치 벡터의 각도와 norm이 거의 일정하게 유지되는 것을 확인할 수 있습니다.

즉, 특정 task를 해결하기위해 똑같은 데이터로 학습을 진행한다면, 시드에 관계없이 가중치의 크기나 각도가 특정 방향으로 일관되게 학습됨을 의미합니다.

 

관측 2

랜덤시드로 같은 데이터를 학습한 50개의 모델을 만들고, 50개의 가중치 평균을 중심(u)이라고 했을 때, 미세조정된 모델들은 가중치 중심으로부터 일정한 거리를 유지하고 표준 편차가 매우 적다는 것을 볼 수 있습니다.

논문의 저자들은 u를 중심으로 일정한 거리 내에 가중치들이 분포되어 있는 형태를 얇은 띠 형태라고 표현하였습니다.

즉, 비슷한 데이터로 학습한 모델들의 가중치는 들쭉날쭉이 아닌 크게 벗어나지 않는 범위 내에 존재하며,

이는 얇은 껍질 형태로 보이고, 모델이 많아질 수록 중심에 가까워진다를 의미합니다.

 

관측 3

fig3은 test 데이터를 기반으로 학습된 모델이 얼마나 틀렸는지의 대한 landscape입니다.

w_0 → 학습이 안된 pre-trained model

w_1→ 랜덤으로 선택된 fine-tuned model

w_50→ 가중치 중심

가중치 중심 근처에서 테스트 오류가 낮으며, 중심에서 멀어질수록 성능이 점점 낮아집니다.

즉, 가중치 중심에 가까울 수록 성능이 향상된다는 것을 알 수 있습니다.

 

관측 4

fine-tuned models의 가중치를 평균내었을 때, 많은 모델을 묶을 수록 중심에 가까워지고, 중심에 가까워질 수록 성능이 향상됩니다.

랜덤 가중치를 만들었을 때도 중심 가중치에 가깝게 생성된 가중치들은 높은 성능을 내며, 멀어질수록 성능이 하락합니다.

여기서 랜덤 가중치는 가우시안 분포에서 무작위로 샘플링된 가중치라고 합니다.

즉, 레이어의 가중치를 랜덤으로 만들어도 가중치의 중심에 근접하면 성능이 높다는 것을 의미합니다.

 

위의 관측들을 기반으로

저자들은 가중치 중심이 가장 높은 성능을 내는 가중치로 정의하고,

최소한의 모델을 사용하여 효율적으로 가중치 중심에 근접하도록 Merge하는 방안인 Model Stock을 고안하게 되었다고합니다.

 

Model STOCK Main Method

 

저자들이 제안하는 핵심 수식은 위와 같습니다.

가중치의 중심인 u가 어디있을지는 모르겠으나,

u를 찾는 것이 아닌, u와 가장 가까울거라 예상되는 가중치인 W_H를 찾는 방법임을 명심해주세요.

 

작동방식

1. 랜덤시드로 학습한 두 개의 fine-tuned models를 가져온다.
2. 두 개의 fine-tuned models의 평균 근처에 가중치 중심이 있을 것이라 가정한다.
   (N개의 튜닝된 모델들을 평균내었을 때가 가중치의 중심이라 하였으니..!)
3. 두 개의 fine-tuned models의 가중치가 있을 테니 두 가중치(벡터)의 각도를 알 수 있습니다.
   코싸인 유사도인 cos A = (w1 * w2) / (|w1| * |w2|)를 이용하여 계산을 합니다.
4. 만약 $cosA$=1 일경우 두 각도가 0도로 작다는 거고, 이는 특정 가중치에 동일하게 편향되어 있으니 가중치의 중심에 가장 가까운 지점은 w_12에 더 가깝게 존재한다는 것입니다.
   두 가중치의 각도가 큰 반대의 경우는 학습이 안된 사전학습된 모델에 더 가까이 존재한다는 것을 의미합니다.

 

코싸인 유사도 수식 및 의미

 

이렇게 가중치의 중심과 가장 가까울거라 예상되는 가중치 W_H로 설정하여 이 가중치를 가진 새로운 모델로 만듭니다.

 

가설 검증 및 결과

사용 모델: CLIP ViT-B/32, CLIP ViT-B/16, CLIP ViT-L/14

데이터셋: ImageNet

위쪽 → 10에포크 약한 데이터 증강 (soup zero-shot 환경과 동일하게)

아래쪽 → 16에포크 강한 데이터 증강 (soup lp 환경과 동일하게)

• 여기서 LP란 Label Propagation를 의미하며, 레이블이 있는 데이터의 정보를 사용하여 레이블이 없는 데이터에도 라벨을 추정하고 이를 모델 학습에 활용하는 방식으로 더 학습을 잘하게 만드는 방법입니다.
• ‘Avg. shifts’ 는 5개의 타 데이터셋의 점수 평균을 의미합니다.
  즉, 왼쪽 점수는 In-distribution(ID), 오른쪽 점수는 Out-of-distribution(OOD)를 의미합니다.

Table2를 보면 Cost 대비 Model Stock이 우수한 성능을 보임을 알 수 있습니다.

Table3에서는 거의 Stock 방법론이 최고 성능을 달성한 것을 볼 수 있습니다.

Table5는 fine tuned models를 결합하면 결합할 수록 In-Distribution score가 향상되며, Out-of-Distribution score도 향상되고, 가중치 중심에 점점 가까워짐을 보입니다.

Table6은 병합 주기에 대해서 말하는데, 1epoch는 약 2500 iters 이라고 합니다.

즉, 병합을 언제하든 Model Stock 방법론은 성능에 큰 변화없이 안정적인 결과를 보인다를 의미합니다.

 

Table7은 왼쪽에 단순히 평균내는 것 대비해서 Stock을 적용한 것이 성능이 더 좋았다를 의미합니다.

위와 마찬가지로, 두 가중치를 단순 평균내는건 성능 향상없이 거의 동일하거나 떨어지는데, stock 방법론은 성능이 향상됨을 검증하였습니다.

결론 및 생각

• 특정 Task를 해결하기 위해 유사한 데이터로 학습한 모델은 가중치 또한 해당 task를 잘 풀기 위한 가중치로 편향되어 있을 것입니다.
• 저자들은 이런 파인튜닝된 모델의 가중치는 랜덤시드에 영향을 거의 안받고 가우시안 분포 특성을 보이며, 가중치 중심에 가까울수록 성능이 개선됨을 검증하였습니다.
• 따라서 파인튜닝한 모든 모델을의 가중치 중심을 찾으면 성능이 향상될텐데, 모델 몇 십개를 만들어 merge한다는게 쉽지 않으니, 사전학습된 모델을 기준 앵커 포인트로 잡고 최소한의 파인튜닝된 모델로 가중치의 중심을 예측해서 해당 가중치를 가진 모델을 만들어 성능 향상을 이루는 Stock 방법론을 제시했습니다.
• 본 방법론을 실제로 적용해보니 분명히 효과가 있었고, 이 방법론에서 아이디어를 얻어 머지 방법론을 만들어 내보는 것도 좋은 연구가 될 것 같다고 생각했습니다.