[논문 리뷰] Model Fusion via Optimal Transport

Sidak Pal Singh, Martin Jaggi|arXiv (Cornell University)|2019. 10. 12.

Advanced Neural Network Applications참고 문헌 32인용 수 7

한 줄 요약

이 논문은 최적 운반(OT)을 사용하여 서로 다른 모델 간의 뉴런을 정렬한 후 가중치를 평균화함으로써 일회성 모델 융합 방법을 제안한다. 계층별 가중치 분포의 워셔스타인 바리센터를 계산함으로써, 기존의 단순 평균화나 디스티illation보다 뛰어난 성능을 달성하며, 비동일한 데이터, 이질적인 데이터에 대해서도 효과적인 지식 전이가 가능하다. 특히 모델 압축 및 앙상블의 대체 수 Mittel로 강력한 성능을 보인다.

ABSTRACT

Combining different models is a widely used paradigm in machine learning applications. While the most common approach is to form an ensemble of models and average their individual predictions, this approach is often rendered infeasible by given resource constraints in terms of memory and computation, which grow linearly with the number of models. We present a layer-wise model fusion algorithm for neural networks that utilizes optimal transport to (soft-) align neurons across the models before averaging their associated parameters. We show that this can successfully yield "one-shot" knowledge transfer (i.e, without requiring any retraining) between neural networks trained on heterogeneous non-i.i.d. data. In both i.i.d. and non-i.i.d. settings , we illustrate that our approach significantly outperforms vanilla averaging, as well as how it can serve as an efficient replacement for the ensemble with moderate fine-tuning, for standard convolutional networks (like VGG11), residual networks (like ResNet18), and multi-layer perceptrons on CIFAR10, CIFAR100, and MNIST. Finally, our approach also provides a principled way to combine the parameters of neural networks with different widths, and we explore its application for model compression. The code is available at the following link, https://github.com/sidak/otfusion.

연구 동기 및 목표

• 자원 제약 환경에서 앙상블 방법의 높은 메모리 및 계산 비용으로 인한 비효율성을 해결하기 위해.
• 서로 다른 가중치 파arameterization을 가진 신경망을 융합할 때 단순 가중치 평균화의 열악한 성능을 극복하기 위해.
• 재학습이나 데이터 공유 없이 이질적이고 비i.i.d. 데이터로 학습된 모델을 일회성으로 융합할 수 있도록 하기 위해.
• 특히 모델 압축을 위한 목적에 맞게 다른 너비를 가진 모델을 원칙적으로 융합할 수 있는 방법을 제공하기 위해.
• 연합 및 분산 학습 환경에서 디스티illation 및 앙상블 방법의 효율적인, 자료 없는 대안을 제공하기 위해.

제안 방법

• 이 방법은 활성화 유사도 또는 입력 가중치 분포를 기반으로 하여 두 개 이상의 부모 모델의 대응 계층 내 뉴런 간에 부드럽고 구조적인 정렬을 최적 운반(OT)을 통해 계산한다.
• 각 계층의 융합을 부모 모델의 가중치 벡터에 의해 정의된 확률 측도의 워셔스타인 바리센터를 계산하는 것으로 공식화한다.
• 융합 과정은 계층별로 적용되며, 재학습 없이 일회성 융합이 가능하다.
• OT를 통한 뉴런 정렬을 통해 너비가 다른 모델을 융합할 수 있도록 지원한다.
• 구조적 프루닝의 후처리 도구로 사용되며, 디스티illation 또는 파인튜닝의 초기화 전략으로도 활용된다.
• VGG11, ResNet18, MLP를 사용하여 CIFAR-10, CIFAR-100, MNIST에서 평가되었으며, 하이퍼파ram터 및 초기화 전략에 대한 분석도 수행되었다.

실험 결과

연구 질문

• RQ1 비i.i.d. 데이터로 학습된 신경망을 융합할 때, 최적 운반 기반의 계층별 융합이 단순 평균화보다 우수한가?
• RQ2 다른 작업이나 레이블 세트로 학습된 모델 간에 지식 전이가 성공적으로 이루어지는가?
• RQ3 재학습이나 데이터 접근 없이 OT 융합이 앙상블 평균화의 효과적인 대체 수 Mittel이 될 수 있는가?
• RQ4 특히 디스티illation이 불가능한 상황에서 OT 융합은 모델 압축에 얼마나 효과적인가?
• RQ5 최적 운반 기반 초기화는 무작위 또는 모델 기반 초기화에 비해 디스티illation 성능을 향상시키는가?

주요 결과

• VGG11을 사용한 CIFAR-10에서 OT 융합은 98.30%의 테스트 정확도를 기록하여, 단순 평균화(98.11%)와 디스티illation(평균 98.26%)를 뛰어넘었다.
• CIFAR-100에서 ResNet18을 사용한 경우, OT 융합은 94.45%의 정확도를 달성하여 단순 평균화(93.89%)를 초월했으며, 디스티illation 결과와 유사하거나 이를 뛰어넘었다.
• 더 큰 교사 모델을 반으로 압축하는 경우(ρ=2), OT 융합은 98.26%의 정확도를 기록하여, 무작위 초기화(98.13%)와 모델 $M_B$ 초기화(98.17%)보다 뛰어났다.
• 극한의 압축 설정(ρ=10)에서, OT 융합은 97.59%의 정확도를 기록하여 무작위 초기화(97.21%)와 $M_B$ 초기화(97.55%)보다 높은 성능을 보였으며, 일관된 성능 향상을 보였다.
• 파인튜닝을 통한 OT 융합은 모든 설정에서 디스티illation 성능을 따라하거나 뛰어넘었으며, 경량 대안으로서의 실현 가능성을 입증했다.
• 다른 레이블 세트로 학습된 모델 간에 지식 전이가 성공적으로 이루어졌으며, 각 부모 모델의 고유한 기능을 습득하면서도 개별적으로 학습된 경우보다 뛰어난 성능을 보였다.