MoEAtt: A Deep Mixture of Experts Model using Attention-based Routing Gate

하임·2024년 11월 28일

MoE

목록 보기

6/9

https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=10459810

간단 요약

논문 "MoEAtt: A Deep Mixture of Experts Model using Attention-based Routing Gate"에서는 MoE(Mixture of Experts)와 어텐션 기반의 라우팅 게이트를 결합한 새로운 아키텍처인 MoEAtt를 소개하고 있습니다. 이 아키텍처의 주요 아이디어와 연구 내용을 정리하면 다음과 같습니다:

1. MoE와 MoEAtt의 개념

MoE는 주어진 입력 공간을 여러 하위 공간으로 나누고, 각 하위 공간에 특화된 전문가(모델)를 학습하여 문제를 해결하는 방식입니다. 일반적으로 MoE에서는 라우터가 각 입력에 적합한 전문가를 선택하여, 그들의 예측을 가중 합산하여 최종 출력을 만듭니다. MoEAtt는 이러한 MoE 아키텍처에 어텐션 메커니즘을 적용하여 라우터가 입력 데이터를 기반으로 가중치를 학습하고 각 전문가의 출력을 조절하도록 설계되었습니다.

2. MoEAtt 아키텍처의 주요 구성 요소

전문가(Experts): MoEAtt의 각 전문가는 이미지 분류와 같은 특정 작업에 적합한 네트워크(예: CNN 또는 LSTM)로 구성될 수 있습니다. 각 전문가는 입력 데이터를 처리하여 예측 벡터와 임베딩 벡터를 생성합니다.
어텐션 기반 라우터(Attention-based Router): MoEAtt는 additive attention 모듈을 라우터로 사용합니다. 이 라우터는 전문가들이 생성한 임베딩 벡터를 활용하여 각 전문가의 예측 결과에 가중치를 부여하고, 이를 조합하여 최종 예측을 생성합니다. 각 전문가의 가중치는 소프트맥스 함수로 정규화되어 확률 분포로 나타나며, 해당 전문가의 중요도를 나타냅니다.

3. 손실 함수 구성

MoEAtt 모델의 손실 함수는 크게 전문가 손실과 라우터 손실로 구성됩니다.

전문가 손실: 각 전문가의 예측 오차에 대해 가중치가 적용된 크로스 엔트로피 손실을 사용하여, 각 전문가가 자신이 책임지는 샘플에 집중하도록 합니다.
라우터 손실: 분류 오류를 줄이기 위한 크로스 엔트로피 손실과 함께 KL 발산을 포함하여 특정 전문가에 편향되지 않도록 균등하게 데이터를 분배하도록 유도합니다.

4. 실험 및 성능 평가

MoEAtt는 다양한 데이터셋(MNIST, CIFAR10, CIFAR10 Rotate)에서 평가되었으며, 기존의 단일 모델 및 MoE 레이어 아키텍처와 비교해 높은 성능을 보여주었습니다.

MNIST: 단순한 네트워크를 기반으로 실험하여, 전문가의 수가 증가할수록 성능이 향상되는 것을 확인했습니다.
CIFAR10: ResNet18과 VGG16을 전문가로 사용하여 MoEAtt를 구성한 결과, 단일 모델보다 성능이 소폭 개선되었으나, 전문가 수가 일정 수준 이상 증가했을 때는 성능 향상이 제한적이었습니다.
CIFAR10 Rotate: 클러스터 구조를 가진 CIFAR10 Rotate 데이터셋에서 MoEAtt가 MoE 레이어 아키텍처보다 우수한 성능을 보여, MoEAtt의 강점을 입증했습니다.

5. 결론 및 미래 연구 방향

MoEAtt는 클러스터 구조를 가진 데이터셋에서 특히 효과적인 모델로, 설명 가능성과 확장 가능성 측면에서 높은 잠재력을 가지고 있습니다. 향후 연구로는 다양한 전문가 아키텍처를 결합하거나, 다른 어텐션 메커니즘을 시도하여 모델의 성능을 더 높이는 방안을 탐색할 수 있을 것으로 기대됩니다.

초록

우리는 라우팅 게이트로 주의 메커니즘을 통합한 새로운 심층 Mixture of Experts (MoE) 아키텍처인 MoEAtt를 소개합니다. 여기서 개별 전문가와 라우터는 공동으로 학습됩니다. 또한, 전문가 간의 이질성을 달성하여 입력 공간의 변별적 표현을 얻도록 설계된 새로운 학습 절차를 제안합니다. MoEAtt 아키텍처는 다양한 데이터셋에서 평가되어 여러 상황에서 그 다재다능함과 적용 가능성을 보여줍니다. 일부 데이터셋에서 최첨단 성능을 달성하며 MoEAtt의 효과를 입증합니다. 또한, MoE 아키텍처가 제공하는 추가적인 이점과 잠재적 가능성에 대해 논의합니다.

키워드 — Mixture of Experts, MoE, Attention

I. 서론

Mixture of Experts (MoE)는 기계 학습 및 통계 분야에서 오랫동안 사용된 잘 알려진 방법입니다. MoE는 입력 공간을 하위 공간으로 나누고 각 하위 공간의 문제를 해결하기 위해 특화된 전문가들을 할당하는 방식을 취합니다. 이 전문가들은 라우터(게이트 네트워크라고도 함)와 함께 학습되어 공간 분할과 모델 학습을 동시에 수행할 수 있게 합니다. 일반적으로 이러한 전문가들은 동일 유형의 서로 다른 모델들이며, 이를 결합하여 정확한 예측을 생성합니다. MoE는 앙상블 모델과 유사하지만, 모델 결합 방식에서 차이가 있습니다. 앙상블은 집합 방식을 통해 모델을 결합하는 반면, MoE는 게이트 메커니즘을 통해 모델을 결합합니다. 두 기법 모두 여러 모델의 다양성과 전문 지식을 활용하여 성능을 개선하는 것을 목표로 합니다.

MoE는 특히 클러스터 구조를 보이는 데이터셋에서 향상된 성능을 보이며 다양한 분야에서 인기를 얻고 있습니다. MoE의 사용 동기는 성능 향상 외에도 설명 가능성, 유연성, 이상 탐지에 대한 강화된 능력 등 여러 이점을 활용할 수 있다는 데 있습니다. MoE 모델은 일반적인 딥 뉴럴 네트워크보다 더 설명 가능하며, 각 전문가가 특정 데이터 하위 집합에서 학습되므로 예측을 이해하기가 더 쉽습니다. 또한 게이트 네트워크는 주어진 입력에 대해 특정 전문가가 선택된 이유를 명확하게 설명해 주어 디버깅과 해석에 유용할 수 있습니다.

Zhang 등 [2]은 심층 신경망의 비용 효율적 맞춤화를 촉진하기 위해 MoE 기법을 사용했습니다. Samsung Ltd.와 협력한 이들의 연구는 제한된 계산 자원을 가진 모바일 기기에서도 MoE가 사용자 맞춤 데이터를 크게 개선할 수 있음을 입증했습니다. Shazeer 등 [3]은 향상된 MoE 방법을 사용하여 최첨단 성능을 달성했습니다. 이 대안적 접근법은 MoE 레이어를 포함하며, 대규모 NLP 데이터셋에서 전통적인 MoE 방식을 능가하면서도 더 적은 계산 자원을 소비합니다. MoE 레이어는 모든 전문가가 전체 네트워크로 학습되는 기존의 MoE 접근과 다르며, 주요 모델 내에서 여러 MLP 전문가와 이들의 출력을 가중치로 결합하는 게이트 네트워크를 포함하여 다음 레이어에 전달되는 출력으로 만듭니다. 이 기법은 더 큰 신경망 모델 내의 작은 MoE로 볼 수 있으며, 때로는 "MLP 전문가의 혼합"이라고도 불립니다. 구현과 수식에 대한 자세한 내용은 Masoudnia 등 [4]의 연구를 참조하십시오.

Fu 등 [5]은 이미지 구조에 기반하여 다중 수준의 합성곱 특징을 결합하도록 학습하는 Mixture of Experts Scene Parsing Network (MoESPNet)을 제안하여, 일부 이미지 분할 데이터셋에서 최첨단 결과를 보고했습니다. Zhang 등 [6]은 제한된 학습 데이터에서 다양화된 ConvNet 전문가를 학습할 수 있는 미세 분류를 위한 점진적 향상 전략을 제안하여, 몇몇 도전적인 미세 분류 데이터셋에서 최첨단 성능을 달성했습니다.

최근 Chen 등 [1]은 클러스터 구조를 가진 데이터셋에서 MoE 모델의 우수성을 수학적으로 입증한 연구를 수행하여, 이러한 특정 유형의 데이터셋에서 MoE 모델의 효과에 대한 이론적 증거를 제시했습니다.

이상 탐지는 데이터에서 표준 또는 예상되는 행동과 크게 벗어나는 패턴이나 인스턴스를 식별하는 것을 의미합니다. Zhao 등 [7]은 현실 세계에서 흔히 발생하는 노이즈 레이블 상황에서 MoE 기반 접근법이 이상 탐지에 대해 갖는 잠재력을 강조했습니다.

본 논문에서는 주의 메커니즘을 라우터로 통합한 새로운 MoE 아키텍처인 MoEAtt를 소개합니다. 또한, 전문가 간의 이질성을 달성하고 전체 모델을 효율적으로 학습할 수 있도록 설계된 새로운 학습 과정을 제안합니다.

II. 배경

Mixture of Experts (MoE)는 입력 특징과 출력 사이의 복잡한 관계를 여러 전문가 모델을 결합하여 모델링하는 기계 학습 기법입니다. 이 기법은 다양한 분야에서 적용되었습니다. MoE는 '분할 정복' 원리에 기반하며, 목적은 입력 공간을 하위 공간으로 나누어 각 전문가가 특정 하위 공간에 전문화되도록 하는 것입니다 [8], [9].

MoE 개념에는 두 가지 핵심 요소가 있습니다: 전문가와 라우터(게이트 네트워크라고도 함)입니다. 전통적인 MoE 프레임워크에서는 입력이 동시에 전문가들과 라우터에 전달됩니다. 전문가들은 다양한 유형의 분류기, 신경망, SVM, 의사 결정 트리 등일 수 있습니다. 라우터의 역할은 주어진 입력에 가장 관련이 깊은 전문가를 식별하고, 해당 입력을 해당 전문가에게 가중치와 함께 전달하는 것입니다. 이러한 가중치들은 전문가의 결정에 가중 투표를 수행하는 데 사용됩니다. 라우터는 동적으로 가장 적합한 전문가를 선택하여 게이트나 스위치 역할을 합니다. 주어진 입력에 대해 어떤 전문가를 사용할지 결정하기 위해 라우터로 신경망을 사용하는 것이 일반적이며, 입력을 받아 각 전문가에 해당하는 가중치를 반환합니다.

일반적인 MoE 목적 함수는 다음과 같습니다:

$L(\theta) = \frac{1}{n} \sum_{i=1}^{n} CE(y_i, g(x_i, \theta))$

여기서, $CE$ 는 교차 엔트로피 손실이며, $\theta$ 는 전문가의 가중치와 파라미터를 포함하는 MoE 모델의 파라미터를 나타냅니다. n은 훈련 데이터셋의 샘플 수이며, $y_i$ 는 i번째 샘플의 실제 클래스 레이블의 원-핫 인코딩 표현입니다. $g(x_i, \theta)$ 는 입력 $x_i$ 와 파라미터 $\theta$ 로 주어진 i번째 샘플의 예측 출력입니다.

$g(x_i, \theta)$ 는 MoE 모델의 예측 함수로, 개별 전문가의 예측을 라우터가 지정한 가중치 $w_j(x_i, \theta)$ 로 결합하여 입력 공간을 부드럽게 (soft) 분할합니다. 이는 앙상블 방식과 다르며, MoE 모델에서는 각 샘플 $x_i$ 에 대해 특정 전문가가 더 큰 가중치로 참여합니다.

우리의 제안 방법인 MoEAtt에서는 주의 메커니즘을 라우터 구성 요소로 사용했습니다. 주의 메커니즘 [10]은 학습된 가중치 행렬을 사용하여 입력 특징들의 가중 합을 계산해 시퀀스에서 각 요소의 관련성을 평가합니다. 다만, 주의는 단일 입력만을 처리할 수 없고, 벡터로 표현된 여러 자원이 필요합니다. 이를 해결하기 위해 4절에서 해결책을 제시할 것입니다.

우리의 방법에서는 다음 세 가지 단계로 예측을 수행합니다. 먼저, 모든 전문가 $m$ 을 통해 입력을 전달하여 숨겨진 표현 $z_1, ..., z_m$ 을 추출합니다. 그런 다음 이 숨겨진 표현을 더하기 주의 모듈인 라우터에 전달하여 크기 m의 가중치 벡터를 출력합니다. 마지막으로, 모든 전문가의 예측 벡터를 라우터가 생성한 가중치 벡터와 곱하여 최종 출력을 계산합니다. 해당 네트워크 구조는 그림 2에 나와 있습니다.

주의 메커니즘에는 여러 종류가 있으며, 예를 들어 스케일드 닷 프로덕트 주의나 다중 헤드 주의 [11]가 있습니다. 하지만 우리의 아키텍처에서 라우터로는 더하기 주의를 선택했으며, 이는 유사한 작업에서 우수한 성능을 보여 주었고 여러 벡터를 하나의 출력으로 융합하는 데 적합합니다.

III. 관련 연구

최근 몇 년간 MoE에 대한 연구는 점점 더 중요해지고 있으며, 다양한 분야에서 널리 활용되어 많은 작업을 효과적으로 처리하고 있습니다. 컴퓨터 비전에서 Enzweiler와 Gavrila [12]는 보행자 분류를 위해 다중 MoE 프레임워크를 제안하여 전문가 분류기를 사용해 여러 모달리티를 결합했습니다. Fu 등 [5]은 MoE-SPNet 기법을 사용하여 이미지 분할 작업을 처리했으며, MoE 접근법이 기존 방법을 능가하는 성능을 보임을 입증했습니다.

Zhang 등 [6]은 제한된 훈련 데이터에서 다양한 ConvNet 전문가를 학습하는 전략을 제안했습니다. 이 전략은 이미지에서 특정 관심 영역에 집중하도록 하는 주의 메커니즘과 다양한 관심 영역을 결합하는 전문가의 혼합을 포함합니다. Liu 등 [13]은 형광 조직 이미지의 핵과 세포질 분할을 위해 Attention-based Mixture of Experts (Att-MoE) 아키텍처를 제안했으며, 이 아키텍처는 여러 전문가 네트워크를 통합하여 적응적으로 다수의 네트워크를 사용하는 방법을 제시했습니다.

Chen 등 [1]은 클러스터 구조를 가진 데이터셋에서 MoE 모델의 효용성을 수학적으로 증명한 연구를 수행하여 MoE 모델이 이와 같은 구조의 데이터셋에 적합하다는 이론적 근거를 제시했습니다.

우리 연구는 이들의 연구와 비교하여 MoE의 라우터 역할에 주의를 사용함으로써 더욱 효과적인 결과를 달성하는지 평가합니다.

IV. MoEAtt 모델

MoEAtt는 입력 공간 $x_1, ..., x_n \in X$ 을 하위 공간으로 효과적으로 분할하는 새로운 아키텍처로, 각 전문가가 특정 하위 공간 학습에 집중할 수 있게 합니다. 각 전문가는 자신에게 할당된 하위 공간에서 뛰어난 분류기를 수행하지만, 다른 하위 공간의 샘플에서는 성능이 저조할 수 있습니다 (그림 1 참조). MoEAtt는 새로운 MoE 아키텍처, 참신한 주의 기반 라우터, 그리고 설계 목표에 맞춘 맞춤형 손실 함수를 결합합니다. 이 아키텍처 개발에서 세 가지 주요 과제를 해결했습니다:

단일 인스턴스를 처리하는 주의 메커니즘 통합
전문가들 간의 이질성 보장
단일 전문가가 다른 전문가를 지배하지 않도록 함

A. 전문가들

MoEAtt에서 각 전문가는 기본적으로 분류기로, 이미지 분류 작업에는 CNN 기반 네트워크, 시계열 작업에는 LSTM 기반 네트워크를 사용할 수 있습니다. 각 전문가는 입력 $x_i$ 를 받아 표현 벡터 $z_{ij}$ 와 분류 벡터 $\hat{y}{ij}$ 를 생성합니다. 전문가 손실은 가중 교차 엔트로피를 사용하여 계산되며, 라우터가 할당한 가중치 $w{ij}$ 를 곱하여 샘플마다 페널티를 다르게 부여합니다.

$L_{e_j} = \frac{1}{n} \sum_{i=1}^{n} w_{ij} \cdot CE(y_i, \hat{y}_{ij})$

모든 전문가의 총 손실은 개별 전문가 손실의 합으로 표현됩니다.

$L_e = \sum_{j=1}^{m} L_{e_j}$

B. 주의 기반 라우터

MoEAtt 모델에서는 더하기 주의 모듈을 라우터로 사용합니다. 이 아키텍처에서 입력은 모든 전문가를 통과하고, 각 전문가는 샘플 xi에 대해 표현 벡터 z{ij}와 예측 벡터 $\hat{y}{ij}$ 를 생성합니다. 라우터는 각 전문가의 예측 $\hat{y}{i1}, ..., \hat{y}{im}$ 의 중요도를 결정하며, 이를 기반으로 가중치 벡터 $w{i1}, ..., w_{im}$ 을 통해 최종 예측을 만듭니다.

최종 MoEAtt 모델의 출력은 각 전문가의 출력을 가중 합하여 계산됩니다:

$\hat{y}i = \text{argmax}\left(\sum{j=1}^{m} w_{ij} \cdot \hat{y}_{ij}\right)$

라우터 손실 함수는 분류 오류에 대한 손실과 Kullback-Leibler(KL) 발산으로 구성됩니다:

$L_r = \frac{1}{n} \sum_{i=1}^{n} CE(y_i, \hat{y}i) + D{KL}(P || Q)$

이 KL 발산은 가중치 벡터와 균등 분포 간의 발산을 측정하여, 입력 데이터가 전문가들 사이에 고르게 분포되도록 합니다.

V. 실험

MoEAtt 방법을 평가하기 위해 일련의 실험을 수행했습니다. 초기 실험에서는 MNIST 데이터셋을 사용하여 MoEAtt 아키텍처가 간단한 모델을 확장하는 효과를 보여주었습니다. 이후 CIFAR10 및 CIFAR10 Rotate 데이터셋에서 MoEAtt 성능을 평가했습니다. 이 실험에서 VGG16 및 ResNet18과 같은 백본 전문가를 사용했습니다. 모든 실험에서는 SGD 옵티마이저와 CosineAnnealingLR 스케줄러를 사용했습니다.

A. 전문가 수 증가의 영향

MNIST 데이터셋에서 ‘초보’ 완전 연결 네트워크(784, 392, 2, 10)를 학습하여 전문가 수가 증가함에 따라 성능이 개선됨을 확인했습니다. 특정 전문가가 레이블에 따라 데이터를 분할하는 라우터의 능력 덕분에 MoE 모델의 성능이 향상되었습니다.

B. CIFAR10 이미지 분류 성능 비교

이 실험에서는 ResNet18 및 VGG16을 MoEAtt 아키텍처의 전문가 네트워크로 사용하여 성능을 평가했습니다. 결과는 단일 전문가 모델에 비해 성능이 약간 향상됨을 보여줍니다.

C. CIFAR10 Rotate에서의 MoEAtt 및 MoE Layer 아키텍처 비교

이 실험에서는 MoE 설정에 적합한 CIFAR10 Rotate 데이터셋을 사용하여 MoEAtt와 MoE Layer 아키텍처를 비교했습니다. 실험 결과, MoEAtt 아키텍처가 MoE Layer 아키텍처를 능가하는 성능을 보여주었습니다.

VI. 요약

본 연구에서는 주의 기반 라우팅 게이트와 맞춤형 손실 함수를 결합한 새로운 MoE 아키텍처인 MoEAtt를 소개했습니다. MoEAtt 아키텍처는 CIFAR10 Rotate와 같은 클러스터 구조를 가진 데이터셋에서 단일 모델과 MoE Layer 아키텍처보다 우수한 성능을 보였습니다. MoEAtt는 더 많은 파라미터를 포함하지만, 성능 향상과 설명 가능성 덕분에 경쟁력 있는 선택이 될 수 있습니다. 추가 연구 방향으로는 동일한 MoEAtt 아키텍처 내에서 다양한 전문가 아키텍처를 사용하는 것, 다른 주의 메커니즘의 사용, 그리고 도메인 적응 작업에서의 MoEAtt 적용을 탐색하는 것이 포함됩니다.