[논문리뷰] (24.10)LINES: POST-TRAINING LAYER SCALING PREVENTSFORGETTING AND ENHANCES MODEL MERGING

DeepLearing/NLP(FineTuning)

[논문리뷰] (24.10)LINES: POST-TRAINING LAYER SCALING PREVENTSFORGETTING AND ENHANCES MODEL MERGING

notdecidedyet 2024. 12. 6. 20:34

0 Abstract

fine-tuning 시 catastrophic forgetting 현상으로 인해 target domain 성능은 향상하나 일반화 성능은 낮아짐.
해결하는 제안책 : LiNeS (Layer-increasing Network Scaling)
- 목적: pre-trained 일반화 능력 보존 + fine-tuned task 성능 향상
- 작동 방식:
  - Layer 깊이 기반으로 parameter update를 선형적으로 scaling
  - Shallow layer: pre-trained 값에 가깝게 유지
  - Deeper layer: task-specific 표현 유지
성능 :
- Vision과 NLP benchmark에서 성능 개선
- Catastrophic forgetting 완화
- Out-of-distribution 일반화 향상

1. Introduction

Pretrain model로 현대 ML페러다임이 변화함.
- end to end 학습에서 -> finetuning으로 전환
- 소수의 foundation model에서 파생된 수천개의 fine-tuned checkpoint로 확산되었음.
모델을 finetuning시키는 것은 catastrophic forgetting과 같은 trade-off를 초래함.
- Fine-tuning 과정의 정규화 접근법 :
  - shallow layer에 더 낮은 학습률을 적용
  - 한계점: 프로세스가 복잡하고 계산 비용이 많이 듦
- Post-training 접근법
  - 방법 1: Weight space에서 직접 checkpoint 편집
  - 방법 2: Pre-trained와 fine-tuned 모델 간 weight 보간 ((1-γ)θ_pretrained + γθ_finetuned)
  - 방법 3: Multi-task 설정에서 여러 checkpoint 병합 시 task 간 간섭 감소 방법
  - 한계점: 여러 모델 병합 시 여전히 상당한 성능 저하 발생
제안 : LiNeS
- 특징 : Post-training, Residual 편집
- Layer 깊이에 따른 scaling:
  - Shallow layer (초기 층): 작은 scaling 계수 적용
    - pre-trained 모델의 일반적인 특성 보존
    - 예: 이미지의 기본적인 edge, texture 등의 특성
  - Deep layer (깊은 층): 큰 scaling 계수 적용
    - task-specific 특성 유지
    - 예: 특정 분류 작업에 필요한 고수준 특성
- 기여점 :
  - Layer별 scaling을 parameter update에 적용하여 pre-trained 모델의 zero-shot 일반화를 보존하면서 fine-tuned 지식을 유지하는 post-training editing 기법인 LiNeS를 제안
  - multi-task model merging baseline들을 크게 향상시키며, 비전과 NLP 도메인 모두에서 벤치마크와 아키텍처 전반에 걸쳐 일관되게 성능을 향상
  - out-of-distribution 일반화 향상, 다른 하이퍼파라미터 구성으로 같은 task에 fine-tuned된 여러 checkpoint 병합, 그리고 다른 reward로 aligned된 LLM policy 병합과 같은 다양한 시나리오에서 기존 weight 보간 방법을 향상

2. Related Work

Representation collapse and regularized fine-tuning

문제점 : Fine-tuning 과정에서 Representation Collapse 발생으로, 사전 학습된 feature가 왜곡되면서 zero-shot 일반화 성능이 저하됨( fine-tuning 과정에서 pre-trained feature가 왜곡)
- Catastrophic forgetting: 고양이와 개를 구분하도록 학습된 모델이 새로운 동물(예: 새)을 학습하면서 고양이와 개를 구분하는 능력이 떨어지는 현상
- Representation collapse: 다양한 객체의 일반적인 특징을 포착하던 모델의 feature extractor가 특정 객체의 특징만을 포착하도록 변화되는 현상
기존 해결 방법 :
- 정규화 도입: Fine-tuning 과정에서 사전 학습 feature 보존을 위해 다양한 정규화 방식을 사용
  - 층별 학습률 감소 (Layer-wise Learning Rate Decay)
  - Projection Invariance (사전 학습된 feature와 Fine-tuned feature 간의 표현 차이를 정규화하여 동일한 표현 공간에서 작동하도록 유도)
- Layer의 특성을 고려한 접근:
  - 모델의 shallow layers은 일반적인 feature를 학습하고, deep layers은 작업 특화된 표현을 학습함.
  - Layer-wise Learning Rate Decay를 적용

Weight interpolation and model merging

Weight Interpolation: 모델 파라미터(weight)를 선형/비선형 경로로 결합하여 성능을 최적화
- 각 모델은 동일한 사전 학습 모델에서 시작하여 별도로 Fine-tuning함.
- Interpolation을 통해 두 모델 간의 공통적이고 유용한 정보를 통합.
  - $θ_interpolated=α⋅θ_1+(1−α)⋅θ_2$
    - $θ_1, θ_2$ : Fine-tuned 모델
성능 :
- 단일 작업 성능 개선:단일 작업에 대해 서로 다른 하이퍼파라미터로 Fine-tuning된 여러 모델을 결합하여 성능 최적화
- 다중 작업 모델 병합: 여러 작업에 대해 학습된 모델을 통합하여 멀티태스킹 성능 제공
- OOD 일반화 : Zero-shot 학습 모델에서 다양한 데이터 분포에 대해 성능을 유지하거나 강화.
한계 :
- 단순 Interpolation은 깊은 층과 얕은 층 간의 기여도를 고려하지 않음.
- 모델 간의 간섭(Task Interference)으로 인해 성능이 감소할 가능성 존재.
- α 최적화 문제

3. POST-TRAINING LAYER-WISE SCALING MITIGATES FORGETTING

Notation :
- $θ_0 ∈ R^N$: N개의 parameter를 가진 pre-trained model
- $θ_t$: 특정 task t에 대해 fine-tuning된 weights
- $τ_t = θ_t - θ_0$: task vector 또는 task t에 대한 residual로, fine-tuning 중 이루어진 update를 나타냄
Fine-tuning leads to catastropic forgetting
- 8-task 이미지 분류 벤치마크에서 정량적으로 이 현상을 보임
- 각 task에 CLIP ViT-B/32 모델을 fine-tuning하고, target task(fine-tuned task)와 control tasks(나머지 7개 task)에서의 성능을 측정. Table 1의 결과는 fine-tuning이 target task의 정확도를 크게 향상시키지만, control tasks의 정확도를 크게 감소시켜 모델의 zero-shot 일반화 능력이 손실됨을 보임

Shallow-layer updates impact minimally on target task accuracy
- 배경 :
  - Pre-trained 모델의 모든 parameter를 업데이트할 필요가 없다
  - Task-specific feature는 deeper layer에 집중된다는 이전 연구 결과
  - 이를 통해 "shallow layer의 변화는 target task 성능에 크게 영향을 미치지 않을 것"이라는 가설을 수립
- 검증 :
  - 이를 검증하기 위해 fine-tuning 후 shallow layer의 update를 점진적으로 downscale
  - ℓ번째 layer의 update $τ^{(ℓ)}$에 대해 다음과 같은 scaling factor를 적용
    - $λ^{(ℓ)} = γ + (1 - γ)\frac{ℓ-1}{L-1}, ∀ℓ ∈ [L]$
      - 첫 번째 layer에서 γ부터 마지막 layer에서 1까지 update를 선형적으로 scaling
      - notation :
        
        ℓ: 현재 layer의 위치 (1부터 L까지)
        
        L: 전체 layer의 수
        
        γ: 최소 scaling 값 (첫 번째 layer에 적용)
        
        1: 최대 scaling 값 (마지막 layer에 적용)
    - 예)
      - 첫 번째 layer (ℓ=1): $λ^{(1)} = γ + 0 = γ$
      - 마지막 layer (ℓ=L): $λ^{(L)} = γ + (1-γ) = 1$
      - 중간 layer들: γ와 1 사이에서 선형적으로 증가
  - Shallow layer를 강하게 downscaling해도 target task 정확도가 거의 영향을 받지 않음
  - 반면, deeper layer를 downscaling하면 target task 정확도가 크게 감소

Shallow-layer updates undermine zero-shot generalization

Shallow-layer의 update은 target task 정확도에 미미한 영향을 미치는 반면, pre-training 중에 학습된 일반적인 feature를 많이 변화시킴.(일반화 성능 낮아짐)
-> 근간을 흔드는 feature는 주로 shallow layer에 있다
특정 태스크에 파인튜닝한 모델의 일반화 성능 저하가 대부분 shallow layer에서 발생한 변화 때문으로 생각한다.
control task (다른 7개의 unseen task)에서의 zero-shot 성능을 평가.
- γ = 1 : 모든 layer가 원래 Update 유지
- Figure 1 (왼쪽)
  - γ = 0일 때 : shallow layer부터 deep layer까지 점진적으로 가중치를 주며 변화시킨 모델에서는 Fine-tuned모델과 비교했을 때 그렇게 많은 하락을 보이지는 않음.
- Figure 1 (오른쪽)
  - γ = 0일 때 : shallow layer부터 deep layer까지 점진적으로 가중치를 주며 변화시킨 모델에서는, 일반화 성능이 그렇게 떨어지지 않음(Pre-trained와 비슷함)

Improved trade-off between target and control performance

최적화하는 scaling 계수 γ를 탐색 (자세한 내용은 Appendix C.1 참조).
최적의 scaling 계수를 선택한 후의 테스트 결과는 Table 1의 마지막 행에 표시, post-training 방법은 fine-tuned 모델에 비해 target task 정확도를 최소 0.2% 차이로 유지하면서 control task 성능을 10% 향상시킴

Wang이 제안한 20개의 컴퓨터 비전 task 벤치마크 사용
각 task를 평가할 때는 하나의 task를 target task로, 나머지 19개를 control task로 설정
성능 측정을 위해 normalized accuracy 사용
Target Task Normalized Accuracy :
기준점: Fine-tuned 모델의 성능(100%로 설정)
결론 : LiNeS 방법의 성능
- Target Task: 99.8% (fine-tuned 모델 성능의 99.8% 유지)
- Control Tasks: 97.9% (pre-trained 모델 zero-shot 성능의 97.9% 유지)

4. METHOD

LiNeS 제안

전체적인 프로세스

Pre-trained 모델 (7B) → Fine-tuning → Fine-tuned 모델 (7B)
(Fine-tuned 모델 - Pre-trained 모델) = Task vector 계산
Task vector에 layer별 scaling 적용
Pre-trained 모델 + Scaled task vector = 최종 모델 (7B)

네트워크의 layer 깊이에 기반하여 서로 다른 layer를 선형적으로 rescale하는 간단한 post-training 기법
shallow layer에서 일반적인 feature를 유지하면서 deeper layer에서 task-specific 적응을 보존하도록 설계
L개의 layer block을 가진 task vector τ가 주어졌을 때, 다음 공식을 사용하여 shallow와 deep layer의 기여도를 조정하기 위한 layer별 선형 scaling을 적용
- 더보기
  전체적인 프로세스
  1. Task Vector계산
  τ = $θ_t$ - $θ_0$ (fine-tuned weights - pretrained weights)
  2. Layer별 Scaling 적용:
  $τ_{LiNeS} = concat(λ(1)τ(1), ..., λ(L)τ(L)) where λ(ℓ) = α + β(ℓ-1)/(L-1)$
  3. 최종 모델 weights 계산:
  $θ_{final} = θ_0 + τ_{LiNeS}$
- recab : task vector - pretrain과 새로운 weight의 차이인 Residual이다.
- $τ_{LiNeS} = concat(λ^{(1)}τ^{(1)}, ..., λ^{(L)}τ^{(L)})$ where $λ^{(ℓ)} = α + β\frac{ℓ-1}{L-1}, ∀ℓ ∈ [L]$
  - Task Vector의 구조
    - $τ^{(ℓ)}$: ℓ번째 layer의 weight 변화량
    - L: 전체 layer 수
  - Scaling Factor의 구조($λ^{(ℓ)}$)
    - α: 첫 번째 layer (ℓ=1)의 최소 scaling 값
    - β: 마지막 layer와 첫 layer의 scaling 차이
    - $\frac{ℓ-1}{L-1}$: 0부터 1까지 선형적으로 증가하는 값
    - 더보기
      예: L=4인 경우
      
      ℓ=1: $λ^{(1)} = α + β(0/3) = α$
      
      ℓ=2: $λ^{(2)} = α + β(1/3)$
      
      ℓ=3: $λ^{(3)} = α + β(2/3)$
      
      ℓ=4: $λ^{(4)} = α + β(3/3) = α + β$
    - α = β = 0: pre-trained 모델과 동일 : 모든 layer의 update가 0이 되어 원래 pre-trained 상태 유지
    - α = 1, β = 0: fine-tuned 모델과 동일 : 모든 layer가 동일하게 scaling되어 기존 fine-tuning 유지
    - 파라미터별 성능은 Appendix에서 찾아볼 수 있음.
  - τ의 layer들은 첫 번째 layer의 α부터 마지막 layer의 α+β까지 선형적으로 증가하는 계수로 점진적으로 scaling

5. MODEL MERGING EXPERIMENTS

5.1 IMPROVING ROBUST FINE-TUNING FOR OOD GENERALIZATION

LiNeS의 효과를 실증적으로 검증하기 위해, 비슷한 방법론인 WiSE-FT와 비교함.

WiSE-FT
- 더보기
  Pre-trained 모델과 fine-tuned 모델의 중간 지점을 찾도록 하는 방법
  $(1-γ)θ_0 + γθ = θ_0 + γτ$
  - γ=0이면 완전히 pre-trained 모델
  - γ=1이면 완전히 fine-tuned 모델
  - θ_0: pre-trained 모델의 가중치
  - θ: fine-tuned 모델의 가중치
  - τ: fine-tuning으로 인한 변화량 (θ - θ_0)
  - γ: 0과 1 사이의 mixing 비율
- OOD에서도 잘 작동하는 모델임.
- 해당 연구에서 제공하는 70개의 체크포인트를 사용
실험 모델 : ImageNet에 fine tuning된 CLIP모델( WiSE-FT 70개 체크포인트에서 진행)
5개의 OOD데이터에 테스트 : ImageNetSketch / ImageNet-A / ImageNet-R / ObjectNet / ImageNet-V2

70개의 체크포인트에 (α = β = 0.5) LiNeS 적용
결과 (Figure3)
- LiNeS만 적용했을 때도 In-distribution(ID, 원래 학습된 데이터 분포) 성능 향상, Out-of-distribution(OOD, 새로운 데이터 분포) 성능 향상

5.2 IMPROVING MULTI-TASK MODEL MERGING

LiNeS를 통해 multi-task merging 알고리즘을 향상시키는 방법을 설명
목표는 서로 다른 task에 독립적으로 fine-tuning된 여러 모델을 단일 모델로 결합하는 것
과거 연구
- Task arithmetic(Ilharco et at., 2023)은 pre-trained 모델과 개별 task vector의 기여도를 분리하는 것을 제안
  - 더보기
    1. Multi-task Vector 생성 (τ_{MTL} = g(τ_1, ..., τ_T))
    - $τ_i$는 각 task i에 대해 fine-tuning된 모델과 pre-trained 모델의 차이를 나타냄 : task vector
    - Merging 함수 g는 이러한 개별 task vector들을 하나의 통합된 vector로 결합
    - 예) Merging함수는 단순평균, 가중평균 등이 될 수 있음.
    
    2. Multi-task 모델 생성 (θ = θ_0 + λ·τ_{MTL})
    - θ_0: 원래의 pre-trained 모델 weights
    - λ: Scaling factor로 merged vector의 영향력을 조절
    - 여기서 λ는 그리드 서치를 통해서 최적화됨
  - 한계점 :
    - 그러나 merged multi-task 모델과 원래 fine-tuned checkpoint 사이에 상당한 성능 차이가 발생.
    - 이는 부분적으로 task vector 간의 interference에서 비롯되며, 한 task의 기여가 다른 task의 성능에 부정적인 영향을 미치기 때문이다.
    - Task interference는 catastrophic forgetting과 연관되어 있다 - 개별 task vector들이 다른 task에 대한 일반화 능력을 상당히 잃어버리기 때문에 이들을 병합하면 서로 간섭이 발생
  - 극복 방법 :
    - 각각의 task vector를 개별적으로 LiNeS를 적용한 후 병합하는 방법(task 에 대한 일반화 작업 능력을 복원)
    - 여러 task vector를 먼저 병합하여 하나의 multi-task vector를 만든 후, 이 병합된 vector에 한 번만 LiNeS를 적용하는 방법
Multi-task vector τ_{MTL}에 LiNeS를 적용
- 두 파라미터 중 β만 튜닝하고, α는 heuristic으로 설정
- α 설정 방법
  - Task Arithmetic의 경우:
    - $τ_{sum}$ = Στ_i (단순 합산) where i from 1 to $N_{models}$
    - α = 1/Nmodels
  - 다른 Merging 전략:
    - $τ_{MTL}$의 norm이 merging 방법에 따라 다름
    - 정규화를 위해 $α = (1/N_{models})(∥τ_{sum}∥/∥τ_{MTL}∥)$ 사용

5.2.1 Computer Vision

LiNeS를 적용했을 때, 모든 모델, 모든 테스크에서 성능 향상이 있었음.

5.2.2 Natural Language Processing

Appendix C.5의 결과로, T5의 Decoder부분만 사용함.

LiNeS는 baseline merging 방법과 벤치마크 전반에 걸쳐 multi-task 성능을 일관되게 향상시킴

5.3 IMPROVING MODEL SOUPS FOR MERGING SINGLE-TASK MODELS

동일한 pre-trained 모델에서 파생된 같은 task에 fine-tuned된 여러 모델을 weight space에서 평균화 및 병합하면 성능이 향상됨을 보였다. 이 섹션에서는 LiNeS가 single-task 모델을 병합할 때 test 성능을 향상시킬 수 있는지 연구한다.

ImageNet에 fine-tuning된 70개의 CLIP ViT-B/32 checkpoint와 pre-trained checkpoint를 병합한다. 이때 Wortsman 등이 소개한 두 가지 변형을 모두 고려한다.
- Uniform soup
- Greedy soup
- 자세한 방법은 Appendix B.3
두가지 경우 모두 weight-averaging 다음과 같이 분해될 수 있다.
- $θ_{soup} = θ_0 + τ_{soup}$, where $τ_{soup} = \frac{1}{N_{models}} \sum_{i=1}^{N_{models}} (θ_i - θ_0)$
- α=1로 고정하고 β는 탐색함.

Table4에서 성능이 미비하게 향상됨.

5.4 IMPROVING REWARDED SOUPS (생략)

6. Discussion

LiNeS를 layer별 scaling coefficient를 backpropagation을 통해 최적화하는 이전 연구들과 비교

Layer별 가중치 얼마나 줄 것인지 과거 연구들과 비교

과거 연구
- Ada-merging : Test set의 entropy를 최소화
  - Unlabeled test data에서 예측의 entropy를 최소화하는 방향으로 layer별 scaling coefficient를 학습
  - Test데이터 분포에 맞춰 자동으로 적응 하도록함.
  - Task별 Layer별 서로 다른 coefficient 학습
- aTLAS
  - Validation 데이터에서 cross-entropy loss를 최소화하도록 layer별 coefficient를 최적화

훈련 없이도 LiNeS는 neural network의 inductive bias를 활용하여 Ada-merging이나 aTLAS와 매우 유사한 scaling을 달성한다.

허나, 계산 비용이 훨씬 작다. aTLAS와 Ada-merging은 과도한 메모리 오버헤드 외에도 여러 번의 훈련 epoch가 필요하여 큰 모델에 대해 확장하기 어려운 단점이 있으나, LiNeS는 없다.