진화론적 최적화를 도입한 Sakana.ai의 모델 머징 방법론

배경지식

모델 머징 (model merging; 모델 합치기)

• 여러 개의 사전학습된 모델의 강점을 활용하는 방식
• 각기 다른 특정 태스크에 대해 학습한 모델들을 결합함으로써, 동시에 여러 태스크를 처리할 수 있는 통합적 모델을 만들고자 함
• 모델 머징 방법론
  • Model Soup
    • 동일한 초기 모델로부터 파인튜닝된 여러 모델의 가중치를 평균
    • 거대한 이미지 처리 혹은 이미지 분류 모델에 대해 잘 작동
  • Task Arithmetic
    • 파인튜닝된 모델의 가중치에서 사전학습 모델의 가중치를 뺌
    • 가중치를 활용하는 방식은 파라미터 사이의 간섭을 일으켜 성능 하락 야기 가능
  • TIES-Merging
    • 기존 머징 방식의 정보 손실을 줄이고자 함
    • 파라미터의 변화를 최소한으로 리셋, Sign Conflict 해결, 오직 얼라인되는 파라미터만 합침
  • DARE
    • 파인튜닝된 모델과 기존 모델 사이의 작은 차이는 0으로 만들고 큰 차이는 더욱 크게 만듦
• 모델 머징 툴킷
  • Mergekit
    • 간단한 머징 방식뿐만 아니라 위에서 언급한 방법론 또한 구현 가능
    • 가중치를 합치는 방법이 아니라 Frankenmerging (다른 레이어를 쌓는 방식) 또한 지원

진화론적 최적 신경망 구조 탐색 (Evolutionary Neural Architecture Search)

• NAS란, 최적의 신경망 구조를 찾는 연구 분야
• 진화론적 방법론이란, 자연 선택에 영감을 받은 방법으로 다양한 가능성을 탐색해 사람의 직관이 놓칠 수 있는 조합까지 확인할 수 있음

문제점

• 기존의 모델 합치기는 사람의 직관이나 도메인 지식에 의존해 최적화가 어려움

해결책: 진화적 모델 머징

파라미터 공간(parameter space)에서의 머징

• 각 레이어에서 파라미터를 혼합해 새로운 가중치를 만들어내는 방식
  • 이 논문에서 레이어란, 인풋/아웃풋 임베딩 레이어 혹은 트랜스포머 블록을 의미
• TIES-Merging과 DARE에 기반, 태스크 벡터 분석을 활용
• 선별한 태스크에 대해 CMA-ES와 같은 진화적 알고리즘을 사용해 최적화 진행

데이터 흐름 공간(data flow space)에서의 머징

• 각기 다른 레이어를 쌓아 새로운 레이어 조합을 만들어냄
  • 각 레이어 내의 가중치는 유지, 토큰이 흐르는 추론 경로를 최적화
• 탐색 공간을 줄이기 위해 모든 레이어를 순차적으로 놓고 레이어들을 포함 혹은 제외
  • 모델 i의 모든 레이어를 나열하고 그 이후에 모델 i+1의 레이어를 나열하는 식
  • 레이어를 반복적으로 위치시키거나(repetitive) 복잡하게 배열하는 것이(permuted) 오히려 성능에 부정적이라는 사전 실험을 통해 이와 같이 탐색 공간을 줄임
  • 레이어의 포함, 제외는 표시자 배열(indicator array)의 인덱스를 바탕으로 결정
    • 표시자 배열은 진화론적 탐색에 의해 최적화됨
• 행렬 W를 활용해 인풋을 적절히 스케일링하는 것이 경험적으로 비추어보았을 때 성능 하락 문제를 경감시킴
  • 행렬 W는 순방향 신경망(feed-forward)으로 표현
  • 행렬 W 또한 진화론적 탐색에 의해 최적화됨

파라미터 공간과 데이터 흐름 공간에서의 머징

• 먼저 파라미터 공간에서의 최적화를 진행
• 그 다음으로 데이터 흐름 공간에서의 최적화를 진행
• 두 방식은 orthogonal하기 때문에 이렇듯 모두 활용할 수 있음

의의

• 자동화된 모델 구성
• 다른 도메인 간 머징 가능
• 최고 성능
• 높은 효율성과 일반화 가능성
• (지역적) 문화 반영 가능

참고 링크

Sakana.ai 블로그 (https://sakana.ai/evolutionary-model-merge/)
Evolutionary Optimization of Model Merging Recipes 논문 (https://arxiv.org/pdf/2403.13187.pdf)