Mixture of A Million Experts

https://arxiv.org/pdf/2407.04153

• 간단 요약 이 논문은 Google DeepMind의 연구로, 'PEER (Parameter Efficient Expert Retrieval)'이라는 새로운 레이어 구조를 제안합니다. PEER는 기존의 Mixture-of-Experts(MoE) 아키텍처의 확장으로, 수백만 개에 이르는 작은 전문가들로 이루어진 풀에서 선택적으로 전문가들을 활성화하여 계산 효율을 높이고자 합니다. 주요 내용은 다음과 같습니다:

  1. 문제 배경:
     기존 Transformer 모델의 피드포워드(FFW) 레이어는 모델의 크기와 함께 계산 비용이 선형적으로 증가하여 자원 소모가 큽니다. 이에 따라, MoE 아키텍처는 자원을 줄이면서도 성능을 유지하는 방법으로 제안되었습니다.

  2. PEER의 설계:

     • PEER 레이어는 입력 벡터를 '쿼리'로 변환하여 다수의 전문가 풀에서 가장 적합한 전문가들을 찾아 활성화하는 구조로 이루어져 있습니다.
     • 각 전문가(expert)는 단일 뉴런 MLP로 구성되며, 이는 효율적인 라우팅과 계산을 가능하게 합니다.
     • 제품 키(Product Key) 기술을 이용해 수백만 개의 전문가 중 효율적으로 일부를 선택하고, 활성화하는 데 필요한 계산 복잡도를 줄였습니다.

  3. 실험 결과:
     다양한 데이터셋(Curation Corpus, Lambada, Pile, Wikitext, C4)에서 PEER 레이어는 기존 밀집 FFW 레이어나 대규모 전문가 구조의 MoE와 비교하여 동일한 연산 비용 내에서 더 낮은 perplexity를 달성했습니다. 또한, 전문가 수와 활성화되는 전문가 수를 달리하며 성능에 미치는 영향을 평가한 결과, PEER 레이어의 전문가 수가 증가할수록 성능이 개선되는 것으로 나타났습니다.

  4. 결론:
     이 연구는 PEER 레이어를 통해 자원 효율성을 높이면서도 성능을 유지하거나 향상시키는 새로운 MoE 스케일링 접근 방식을 제안하며, 추후 대규모 언어 모델의 성능 최적화에 중요한 기여를 할 수 있음을 시사합니다.

     PEER는 특히, 고성능이 요구되는 언어 모델링에서 자원의 효율적 사용과 모델 확장성의 가능성을 크게 높이는 방법으로 평가됩니다.

     Mixture of Experts(MoE) 방법론은 대규모 모델에서 계산 비용을 줄이면서도 성능을 유지하기 위해 설계된 방식입니다. 핵심 아이디어는 모델의 일부 레이어(주로 피드포워드(FFW) 레이어)에 여러 개의 전문가(expert)를 배치하고, 각 입력에 대해 적합한 전문가만 선택적으로 활성화하는 방식입니다. MoE의 주요 목표는 모든 매개변수를 매번 사용하지 않고 일부 전문가만 활성화하여 모델의 효율성을 높이는 것입니다.

     MoE의 기본 개념

     MoE 아키텍처에서는 전체 모델이 여러 개의 전문가 모듈(예: FFW 레이어)로 구성됩니다. 입력 데이터가 주어지면, 이 데이터는 다음 단계에서 가장 적합한 소수의 전문가들에게만 전달됩니다. 예를 들어, 100명의 전문가가 존재하는 모델이라면, 특정 입력에 대해 상위 2-3개의 전문가만 선택해 활성화하는 방식입니다.

  5. 전문가 선택: 각 입력 데이터에 대해 적합한 전문가를 선택하는 작업을 게이팅(gating)이라고 합니다. 게이팅 함수는 입력에 대한 정보와 각 전문가의 성향을 바탕으로, 어떤 전문가가 가장 적합한지 판단하여 선택합니다.

  6. 희소 활성화: 선택된 소수의 전문가만 활성화하므로, 전체 모델을 사용하는 밀집 모델(dense model)과 달리 계산 비용을 크게 줄일 수 있습니다. 이로 인해 모델 크기가 커지더라도 모든 전문가를 활성화하지 않기 때문에 계산 자원 소모가 줄어듭니다.

     MoE의 장점

  7. 효율성: 각 입력에 대해 일부 전문가만 활성화하기 때문에 모든 매개변수를 매번 사용하는 기존 모델보다 훨씬 효율적입니다.

  8. 확장성: 전문가의 수를 늘려도 활성화되는 매개변수 수는 제한되므로, 모델을 큰 규모로 확장할 수 있습니다.

  9. 전문화: 각 전문가가 특정 입력 데이터 유형에 대해 특화된 처리를 할 수 있도록 훈련될 수 있습니다. 예를 들어, 문장 유형에 따라 서로 다른 전문가가 활성화될 수 있습니다.

     MoE의 한계와 해결방안

  • 게이팅 불균형: 일부 전문가가 자주 선택되고 다른 전문가들은 거의 사용되지 않는 불균형 현상이 발생할 수 있습니다. 이를 해결하기 위해 MoE 연구는 로드 밸런싱 기법을 사용하여 각 전문가가 골고루 선택되도록 조정합니다.

  • 훈련 및 최적화 문제: MoE 모델의 훈련과 최적화는 기존 모델보다 복잡하며, 이를 안정적으로 훈련하기 위해 다양한 최적화 기법이 필요합니다. 예를 들어, Fedus et al. (2022)에서 소개된 스위치 트랜스포머(Switch Transformer)에서는 효율적이고 안정적인 훈련을 위해 간단한 게이팅 방식을 제안했습니다.

    MoE의 최신 연구와 발전

    최근 MoE 연구에서는 토큰 선택(Token-choice)과 전문가 선택(Expert-choice) 두 가지 방법론이 개발되었습니다.

  • 토큰 선택(Token-choice): 각 토큰이 상위 전문가를 선택하도록 하는 방식으로, 각 토큰마다 별도의 전문가가 활성화될 수 있습니다.

  • 전문가 선택(Expert-choice): 전문가가 토큰을 선택하는 방식으로, 상위 토큰들을 선택하여 처리합니다. 이는 선택된 전문가가 일괄적으로 여러 토큰을 처리하므로, 계산 효율성을 높일 수 있습니다.

    또한, 일부 연구에서는 결정적 해시 테이블을 사용하여 라우팅 속도를 O(1) 수준으로 높이는 방법을 제안하기도 했습니다.

    PEER 레이어와 MoE의 차이점

    PEER(Parameter Efficient Expert Retrieval) 레이어는 MoE 방식의 확장 형태로, 제품 키 검색(Product Key Retrieval)을 통해 효율적으로 많은 전문가 중 일부를 선택합니다. 또한 각 전문가를 단일 뉴런으로 설정하여 MoE 모델에서 발생할 수 있는 계산 과부하를 크게 줄이는 방식입니다.

    PEER(Parameter Efficient Expert Retrieval) 레이어는 MoE의 개념을 확장하여 수백만 개의 작은 전문가들을 효율적으로 사용하는 방식입니다. 기본 MoE는 소수의 전문가 중에서 적합한 전문가를 선택해 활성화하지만, PEER는 훨씬 많은 전문가 풀에서 효율적으로 필요한 전문가를 검색하고 사용하는 것을 목표로 합니다.

    PEER의 작동 방식

  1. 제품 키(Product Key)를 이용한 라우팅

     PEER의 핵심은 제품 키(Product Key) 검색을 통해 입력에 가장 적합한 전문가를 효율적으로 찾는 것입니다. 제품 키는 검색 공간을 줄이기 위해 사용되는 일종의 인덱스 체계로, 쿼리 벡터와 가장 유사한 전문가를 찾아내는 데 도움을 줍니다. 각 입력에 대해 쿼리 벡터가 생성되고, 이 쿼리 벡터가 제품 키와 비교되어 내적이 높은 상위 k개의 전문가가 선택됩니다. 이를 통해 전체 전문가 수가 많아도 일부만 효율적으로 선택할 수 있습니다.

  2. 단일 뉴런 전문가

     PEER는 각 전문가가 단일 뉴런으로 구성된 매우 작은 MLP입니다. 기존 MoE는 각 전문가에 여러 개의 뉴런을 사용하는 반면, PEER는 각 전문가가 단일 뉴런만 포함하도록 설계하여 메모리 및 계산 효율을 극대화합니다.

  3. 멀티 헤드 검색(Multi-Head Retrieval)

     PEER는 멀티 헤드(Multi-Head) 방식을 채택하여 여러 개의 독립적인 쿼리 네트워크를 동시에 사용합니다. 이는 Transformer의 멀티 헤드 어텐션과 유사하게 각 헤드가 독립적으로 상위 k개의 전문가를 검색합니다. 이렇게 하면 입력 벡터를 다양한 각도에서 분석하고, 각 쿼리가 상이한 전문가를 선택할 수 있어 더 풍부한 표현력을 제공합니다.

  4. 가중 평균 결합

     선택된 전문가들의 출력은 라우터 점수(softmax 또는 sigmoid를 적용하여 계산된 점수)를 가중치로 하여 **가중 평균**으로 결합됩니다. 이를 통해 상위 전문가들의 출력을 종합해 최종 출력값을 생성하게 됩니다.

     제품 키 검색(Product Key Retrieval)의 효율성

     PEER는 일반적인 전문가 검색 방식의 계산 복잡도를 줄이기 위해 제품 키 검색 방식을 사용합니다. 검색 과정에서 쿼리 벡터를 두 개의 하위 벡터로 분할하고, 각각의 하위 벡터를 별도의 하위 키와 비교해 상위 $k$개의 전문가를 선택합니다. 이 과정은 카르테시안 곱을 통해 효율적으로 검색 범위를 좁힐 수 있으며, 전체 검색 복잡도를 O($(\sqrt{N}$$)$로 줄여 줍니다.

     요약

  • PEER 레이어는 수백만 개의 전문가 중에서 제품 키를 이용해 가장 적합한 전문가를 효율적으로 선택합니다.

  • 단일 뉴런 전문가를 사용해 메모리와 계산 자원을 절약하며, 필요 시 멀티 헤드 검색을 통해 다양한 전문가를 동시에 활용할 수 있습니다.

  • 이를 통해 모델의 확장성을 높이면서도 효율적인 검색과 계산 성능을 유지하는 것이 PEER의 주요 목표입니다.

    PEER는 이러한 방식을 통해 기존 밀집 FFW 레이어나 coarse-grained MoE보다 더 높은 성능과 효율성을 제공합니다.

• 단점(추측)? PEER 방법에도 몇 가지 단점과 제약이 존재합니다. 주요 단점은 다음과 같습니다.

  1. 고도의 하드웨어 요구사항

  • PEER는 제품 키 검색(Product Key Retrieval)과 멀티 헤드 검색(Multi-Head Retrieval)으로 수백만 개의 전문가를 효율적으로 활용하는 방식이지만, 이러한 구조는 고성능 하드웨어와 최적화된 메모리 관리가 필수적입니다. 수백만 개의 전문가를 유지하고 빠르게 검색하려면 고성능의 GPU나 특수한 가속기 등이 필요할 수 있습니다.

  • 또한, 제품 키 검색은 다수의 내적 연산이 필요하기 때문에 연산 속도를 최적화하기 위해서는 전용 하드웨어 커널이나 효율적인 메모리 접근 방식을 설계해야 합니다.

    2. 복잡한 모델 설계 및 구현

  • PEER는 다수의 전문가와 멀티 헤드 검색, 제품 키 기반 라우팅을 포함하고 있어서 모델 설계와 구현이 매우 복잡합니다. 일반 MoE와 비교해도 구조적 복잡성이 높으며, 이를 코드로 구현하고 유지하는 데 많은 노력이 필요합니다.

  • 특히, 제품 키 검색과 같은 라우팅 메커니즘은 구현 난이도가 높아 효율적인 구현을 위해서는 전문가 수준의 최적화가 필요할 수 있습니다.

    3. 게이팅 및 전문가 불균형 문제

  • PEER의 경우 수백만 개의 전문가 중 일부만 선택적으로 활성화되기 때문에 일부 전문가에 부하가 집중되는 불균형 현상이 발생할 수 있습니다. 즉, 일부 전문가들이 자주 선택되고, 다른 전문가들은 거의 선택되지 않는 문제가 생길 수 있습니다.

  • 이는 로드 밸런싱이 제대로 이루어지지 않으면 특정 전문가의 과부하나 효율적인 전문가 활용에 방해가 될 수 있습니다. 이를 해결하기 위한 추가적인 로드 밸런싱 알고리즘이나 게이팅 방식의 개선이 필요할 수 있습니다.

    4. 훈련 데이터의 제한적 활용

  • PEER는 각 입력에 대해 일부 전문가만 활성화하기 때문에, 각 전문가가 다양한 입력 데이터를 충분히 학습하기 어려울 수 있습니다. 즉, 특정 전문가가 특정 입력만을 학습하면서 다른 입력에 대해 일반화 능력이 떨어질 가능성이 있습니다.

  • 특히, 작은 데이터셋을 사용할 경우 모든 전문가가 충분히 학습되지 못해 오히려 성능이 저하될 수 있습니다. 이는 대규모 데이터셋이 필요한 대형 모델에 적합한 방식이지만, 작은 데이터셋에서는 오버헤드가 발생할 수 있습니다.

    5. 응답 지연 및 실시간 사용의 제약

  • 수많은 전문가 중 상위 전문가를 선택하고 가중 합산하는 과정에서 응답 지연(latency)이 발생할 수 있습니다. 특히 실시간 처리가 요구되는 애플리케이션에서는 복잡한 라우팅과 검색 연산이 병목이 되어 응답 속도가 떨어질 수 있습니다.

  • 실시간 응답을 필요로 하지 않는 연구나 오프라인 작업에는 문제가 없지만, 실시간 애플리케이션에 적용하기에는 개선이 필요할 수 있습니다.

    요약

    PEER는 대규모 전문가 풀을 효율적으로 활용하기 위한 강력한 방법이지만, 하드웨어 요구사항, 복잡한 모델 설계, 로드 불균형, 훈련 데이터 요구, 응답 지연 등의 단점이 있습니다. 이러한 단점을 해결하기 위해서는 효율적인 로드 밸런싱, 메모리 최적화, 고성능 하드웨어의 활용 등이 필요하며, 실시간 애플리케이션에는 추가적인 개선이 필요할 수 있습니다.

요약

표준 Transformer 아키텍처에서 피드포워드(FFW) 레이어는 숨겨진 레이어의 너비가 증가함에 따라 계산 비용과 활성화 메모리가 선형적으로 증가하는 문제를 가지고 있습니다. 이러한 문제를 해결하기 위해 희소 Mixture-of-Experts(MoE) 아키텍처가 등장했으며, 이를 통해 모델 크기를 계산 비용과 분리하여 처리할 수 있게 되었습니다. 최근에 발견된 미세-그레인 MoE 스케일링 법칙은 더 높은 세밀도(granularity)가 성능 향상에 기여한다는 것을 보여주었습니다. 그러나 기존 MoE 모델은 계산 및 최적화의 한계로 인해 소수의 전문가만 사용할 수 있었습니다. 이 논문에서는 PEER(파라미터 효율 전문가 검색)라는 새로운 레이어 디자인을 소개하며, 이 디자인은 제품 키(product key) 기술을 사용해 수백만 개 이상의 작은 전문가 풀에서 희소 검색을 수행합니다. 언어 모델링 과제에서 PEER 레이어는 밀집 FFW 및 Coarse-grained MoE 대비 성능-계산 비용 측면에서 우수한 성능을 보였습니다. 수많은 전문가를 효율적으로 활용할 수 있게 함으로써 PEER는 Transformer 모델의 추가 확장을 가능하게 하며, 동시에 계산 효율성을 유지할 수 있습니다.

서론

최근 몇 년간, 매개변수 수, 학습 데이터 양, 또는 계산 예산을 증가시키는 것이 모델 성능 향상에 중요한 방법으로 떠올랐습니다 (Kaplan et al., 2020; Hoffmann et al., 2022). 특히, Transformer에서 지식 저장을 담당하는 피드포워드(FFW) 레이어는 전체 매개변수의 3분의 2를 차지합니다 (Geva et al., 2021; Dai et al., 2022). 하지만 밀집 FFW는 매개변수 수에 따라 계산량과 메모리 소모가 선형적으로 증가하는 한계가 있습니다.

1. Mixture-of-Experts(MoE) 아키텍처

   최근 연구들(Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022; Zhou et al., 2022)은 이러한 문제를 해결하기 위해 MoE 아키텍처를 도입했습니다. MoE는 밀집 FFW 대신 희소하게 활성화된 전문가 모듈 세트를 사용하여 계산 비용과 모델 크기를 분리할 수 있습니다. 특히, Clark et al. (2022)의 연구는 MoE 모델의 확장 법칙을 규명하여, 전문가 수를 늘리는 것이 성능 향상에 효과적임을 보여주었습니다. 그러나 특정 모델 크기 이상이 되면 효율성 증가가 정체되는 현상이 나타났습니다. 최근 Krajewski et al. (2024)는 학습 토큰 수가 최적화되지 않으면 MoE 모델의 효율성이 떨어진다는 점을 발견했습니다. 그들은 '세밀도'(활성 전문가 수)를 새로운 확장 축으로 도입하며, 세밀도가 높을수록 성능이 개선됨을 실험적으로 증명했습니다.

2. 지속적 학습을 위한 MoE의 가능성

   MoE는 '지속적 학습(lifelong learning)'이라는 관점에서도 주목받고 있습니다. Chen et al. (2023)은 새로운 전문가들을 추가하고 정규화하는 방식으로 MoE 모델이 연속적인 데이터 스트림에 적응할 수 있음을 보여주었습니다. 이는 오래된 전문가를 동결하고 새로운 전문가만 업데이트하여 기억 소실을 방지하고 유연성을 유지하는 효과를 제공합니다 (Mitchell et al., 2018).

3. PEER 아키텍처 소개

   본 연구에서는 '파라미터 효율 전문가 검색(Parameter Efficient Expert Retrieval, PEER)' 아키텍처를 제안합니다. PEER는 제품 키(product key) 검색 기술(Lample et al., 2019)을 활용해 수백만 개의 전문가에 대한 효율적인 라우팅을 지원하며, 계산 비용과 매개변수 수를 분리합니다. 실험 결과, PEER는 밀집 FFW와 대규모 MoE 및 Product Key Memory(PKM) 레이어에 비해 더 나은 계산-성능 트레이드오프를 보여주며, 기초 모델 확장을 위한 경쟁력 있는 대안임을 입증했습니다.

4. 주요 기여

   • 극단적 MoE 설정 탐구: 기존 MoE 연구가 큰 전문가 소수에 집중했던 것에서 벗어나, 본 연구는 다수의 작은 전문가를 활용한 설정을 탐구합니다.
   • 라우팅을 위한 학습된 인덱스 구조: 학습된 인덱스 구조(Kraska et al., 2018)를 통해 백만 개 이상의 전문가를 효율적으로 라우팅할 수 있음을 최초로 시연합니다.
   • 새로운 레이어 디자인: 제품 키 라우팅과 단일 뉴런 전문가를 결합한 PEER 레이어는 큰 계산 비용 없이 레이어 용량을 확장합니다. 실험적으로 밀집 FFW, 대규모 MoE, PKM 레이어 대비 우수한 효율성을 입증했습니다.
   • 종합적 제거 연구: PEER의 다양한 설계 요소(예: 전문가 수, 활성 매개변수 수, 헤드 수, 쿼리 배치 정규화)가 언어 모델링 작업에 미치는 영향을 분석했습니다.

PEER 개요

PEER 레이어는 수학적으로 $f : \mathbb{R}^n \to \mathbb{R}^m$ 함수로 정의되며, 세 가지 주요 구성 요소로 이루어져 있습니다:

1. $N$ 개의 전문가 집합 $E := \{e_i\}_{i=1}^{N}$ , 여기서 각 전문가 $e_i : \mathbb{R}^n \to \mathbb{R}^m$ 는 $f$ 와 동일한 시그니처를 가집니다.
2. N 개의 제품 키 집합 $K := \{k_i\}_{i=1}^{N}$ \subset \mathbb{R}^d
3. 입력 벡터 $x \in \mathbb{R}^n$ 를 쿼리 벡터 $q(x) \in \mathbb{R}^d$ 로 매핑하는 쿼리 네트워크 $q : \mathbb{R}^n \to \mathbb{R}^d$

입력 x 가 주어지면, q(x) 와 내적이 가장 높은 상위 k 개의 제품 키에 해당하는 전문가를 검색합니다:

$I = T_k\left(\{q(x)^T k_i\}_{i=1}^{N}\right)$

여기서 $T_k$ 는 상위 $k$ 연산자를 나타내며, $k$ 개의 상위 전문가를 선택합니다.

그 후, 선택된 상위 k 개의 전문가와 쿼리 벡터 간의 내적에 비선형 활성화 함수(예: softmax 또는 sigmoid)를 적용하여 라우터 점수를 계산합니다:

$g_i(x) = s(q(x)^T k_i)$

마지막으로, 라우터 점수를 가중치로 하여 선택된 전문가들의 출력을 선형 결합하여 최종 출력을 계산합니다:

$f(x) = \sum_{i \in I} g_i(x)e_i(x)$

제품 키 검색

PEER는 수백만 개 이상의 전문가를 사용하는 것을 목표로 하므로, Eq. 1에서 단순히 $k$ 개의 상위 인덱스를 계산하는 것은 매우 비쌉니다. 이를 해결하기 위해 제품 키 검색 기술을 적용합니다. 각각의 키 $k_i$ 를 독립적인 $d$ -차원 벡터 대신 $d/2$ -차원 하위 키 집합 $C, C' \subset \mathbb{R}^{d/2}$ 에서 가져온 벡터들로 결합하여 생성합니다:

$K = \left\{\left[\begin{array}{c} c \\ c' \end{array}\right] \Bigg| c \in C, c' \in C' \right\}$

여기서 $C$ 와 $C'$ 의 크기는 $\sqrt{N}$ 이고, $c$ 와 $c'$ 의 차원은 $d/2$ 입니다. 따라서 $N$ 을 완전 제곱수로 설정하고 $d$ 를 짝수로 설정합니다.

이러한 카르테시안 곱 구조를 통해 $k$ 개의 전문가를 더 효율적으로 찾을 수 있습니다. 쿼리 벡터 $q(x)$ 를 두 개의 하위 쿼리 $q_1$ 과 $q_2$ 로 나눈 뒤, 각각 하위 쿼리와 하위 키 간의 내적을 통해 $T_k$ 연산을 적용합니다:

$I_C = T_k(q_1^T c_i), \quad I_{C'} = T_k(q_2^T c'_j)$

이 결과는 $k^2$ 개의 후보 키 집합 $K' := \left\{\left[\begin{array}{c} c_i \\ c'j \end{array}\right] \Bigg| i \in I_C, j \in I{C'}\right\}$ 를 생성하며, 이 후보 키들 중 $q(x)$ 와 가장 유사한 $k$ 개의 키가 포함됨을 수학적으로 보장합니다.

후보 키와 쿼리 벡터 간의 내적은 단순히 하위 키와 하위 쿼리 간의 내적 합으로 계산할 수 있습니다:

$q(x)^T \left[\begin{array}{c} c_i \\ c'_j \end{array}\right] = q_1^T c_i + q_2^T c'_j$

따라서 이 $k^2$ 개의 내적에 대해 다시 $T_k$ 연산자를 적용하여 최종적으로 원래 제품 키 집합 $K$ 에서 $k$ 개의 상위 매칭 키를 얻을 수 있습니다. 이를 통해 Eq. 1의 상위 $k$ 전문가 검색 복잡도를 O($( N \times d$ )에서 O(($\sqrt{N} + k^2$)$times d$)로 줄입니다.

Parameter Efficient Experts와 Multi-Head Retrieval

기존 MoE 아키텍처에서는 각 전문가의 숨겨진 레이어 크기를 다른 FFW 레이어와 동일하게 설정하는 경우가 많습니다. 그러나 PEER에서는 각 전문가 $e_i$ 가 단일 뉴런을 가진 하나의 MLP, 즉 '싱글톤 MLP'로 구성됩니다. 즉, 전문가 $e_i(x)$ 는 다음과 같이 정의됩니다:

$e_i(x) := \sigma(u_i^T x)v_i$

여기서 $v_i$ 와 $u_i$ 는 행렬이 아닌 $x$ 와 같은 차원의 벡터이며, $\sigma$ 는 ReLU 또는 GELU와 같은 비선형 활성화 함수입니다. 편의를 위해 바이어스 항은 생략했습니다.

개별 전문가의 크기를 조절하는 대신, PEER 레이어의 표현력을 Multi-Head Retrieval 방식을 통해 조정합니다. 이는 Transformer의 Multi-Head Attention 또는 PKM의 Multi-Head Memory와 유사합니다. 특히, PEER는 하나의 쿼리 네트워크 대신 h 개의 독립적인 쿼리 네트워크를 사용하여 각각 고유의 쿼리를 계산하고 별도의 전문가 집합을 검색합니다. 그러나 서로 다른 헤드는 동일한 전문가 풀과 제품 키를 공유합니다. 각 헤드의 출력은 단순히 합산됩니다:

$f(x) := \sum_{i=1}^h f_i(x) = \sum_{i=1}^h \sum_{j \in I_i} g_j(x)e_j(x)$

이 때, 각 헤드에서 하나의 전문가만 검색할 경우 (k = 1 ), h 개의 헤드를 가진 PEER 레이어는 h 개의 숨겨진 뉴런을 가진 하나의 전문가를 사용하는 것과 동일하게 작동합니다:

$f(x) = \sum_{i=1}^h e_i(x) = \sum_{i=1}^h \sigma(u_i^T x)v_i = V \sigma(W^T x)$

여기서 $W = [u_1, \cdots, u_h]$ , $V = [v_1, \cdots, v_h]$ 입니다. 즉, PEER는 공유 저장소에서 h 개의 싱글톤 MLP를 결합하여 동적으로 $h$ 개의 뉴런을 가진 MLP를 구성합니다. 기존 MoE가 여러 숨겨진 뉴런을 가진 MLP를 전문가로 사용하는 것과 달리, PEER는 전문가 간의 공유 뉴런을 가능하게 하여 지식 전이 및 파라미터 효율성을 향상시킵니다.

알고리즘 1에는 파라미터 효율적인 전문가 가중치를 임베딩 레이어에 저장하고, 이를 einsum 연산으로 결합하는 PEER의 간단한 forward 패스 구현이 나와 있습니다. 이 구현은 GLU 변형의 전문가를 추가 선형 게이트 가중치를 통해 쉽게 확장할 수 있습니다. 실제로는 효율적인 임베딩 조회 및 einsum 연산 가속화를 위해 특수 하드웨어 커널이 필요할 수 있습니다.

왜 작은 전문가가 많은가?

MoE 레이어의 경우, 총 매개변수 수 $P$ , 토큰당 활성 매개변수 수 $P_{\text{active}}$ , 전문가 크기 $P_{\text{expert}}$ 라는 세 가지 하이퍼파라미터로 특성을 정의할 수 있습니다. Krajewski et al. (2024)에 따르면 MoE 모델의 스케일링 법칙은 다음과 같습니다:

$L(P, D, G) = c + \left(\frac{g}{G^{\gamma}} + a\right)\frac{1}{P^{\alpha}} + \frac{b}{D^{\beta}}$

여기서 $L$ 은 최종 테스트 손실, $D$ 는 총 학습 토큰 수, $G$ 는 활성 전문가 수를 나타내는 세밀도입니다:

$G := \frac{P_{\text{active}}}{P_{\text{expert}}}$

모델 성능을 향상시키기 위해서는 $P, D, G$ 를 확장해야 하지만, $P_{\text{active}}$ 는 제한하는 것이 필수적입니다. 이는 훈련과 추론 중 주로 활성 매개변수로 인해 계산 및 메모리 비용이 결정되기 때문입니다.

따라서 각 전문가의 크기 $P_{\text{expert}}$ 를 줄이고 전문가 수 $N$ 을 늘려야 합니다. 일반적으로, 단일 숨겨진 레이어를 가진 전문가에 대해 $P_{\text{expert}} = (2d_{\text{model}} + 1)d_{\text{expert}}$ 이고 $P_{\text{active}} = (2d_{\text{model}} + 1)d_{\text{active}}$ 입니다. PEER의 경우 $d_{\text{expert}} = 1$ 로 설정하여 가능한 최소 전문가 크기를 사용하며, 활성화되는 뉴런 수는 검색 헤드 수와 헤드당 검색되는 전문가 수의 곱 $d_{\text{active}} = hk$ 로 설정합니다. 따라서 PEER의 세밀도는 항상 $G = \frac{P_{\text{active}}}{P_{\text{expert}}} = \frac{d_{\text{active}}}{d_{\text{expert}}} = hk$ 입니다.

알고리즘 1: PEER 레이어의 forward 패스 구현

def peer_forward(self, x):
    # 모든 전문가의 다운/업 프로젝션 가중치를 저장하는 임베딩 레이어
    self.w_down_embed = nn.Embed(num_embeddings=self.n_experts, features=self.d_model)
    self.w_up_embed = nn.Embed(num_embeddings=self.n_experts, features=self.d_model)

    # 제품 키를 사용해 상위 매칭 전문가의 가중치 검색
    # 인덱스와 점수는 'bthk' 모양을 가지며, h는 헤드 수를 나타냄
    indices, scores = self.get_indices(self.query_proj(x), self.sub_keys, top_k=self.k)
    w_down = self.w_down_embed(indices)
    w_up = self.w_up_embed(indices)

    # 전문가 출력의 가중 평균 계산
    x = jnp.einsum('btd, bthkd -> bthk', x, w_down)
    x = self.activation(x)
    x = x * nn.softmax(scores)
    x = jnp.einsum('bthk, bthkd -> btd', x, w_up)
    return x

결과

생략

관련 연구

Mixture of Experts (MoE)

Shazeer et al. (2017)은 희소하게 게이트된 MoE의 효과를 통해 GPU 클러스터에서 모델 용량을 효율적으로 증가시킬 수 있음을 보여주었으며, 이후 MoE는 대형 모델을 효율적으로 확장하는 기법으로 자리 잡았습니다. 이후 연구들(Fedus et al., 2022; Lepikhin et al., 2020; Du et al., 2022)은 MoE의 로드 밸런싱, 통신 오버헤드, 학습 불안정성 같은 문제를 해결하기 위한 다양한 변형을 제안했습니다. 이 방법들은 특정 Transformer 블록의 피드포워드(FFW) 레이어를 다수의 FFW 레이어로 구성된 희소 게이트 MoE 레이어로 대체하며, 일반적으로 각 전문가의 크기는 밀집 FFW 레이어와 동일합니다. Zhou et al. (2022)은 최근에 전문가가 상위 k개의 토큰을 선택하는 Expert Choice 라우팅 방법을 도입했지만, 토큰-선택 및 전문가-선택 방식은 N \times M 크기의 게이팅 점수 행렬에 대한 상위 k 연산자를 필요로 하므로, 계산 비용이 최소 O(N)에 이르러 전문가 수를 제한합니다.

상위 k 연산자를 사용하는 대신, Roller et al. (2021)과 dos Santos et al. (2023)와 같은 연구는 결정적 해시 테이블을 라우터로 사용하여 평균 조회 복잡도가 O(1)인 접근 방식을 제안했습니다. 그러나, 이는 고정된 라우터로 학습되지 않아 Clark et al. (2022)의 연구에 따르면 결정적 라우팅은 학습 가능한 라우터만큼 확장되지 못합니다. Dikkala et al. (2023)은 학습된 라우터가 데이터 내 잠재적 클러스터 식별과 같은 이점을 제공함을 이론적으로 증명했습니다. 본 논문에서 제안하는 PEER 레이어는 O(\sqrt{N})의 복잡도를 갖춘 학습된 라우터를 사용합니다.

파라미터 효율적인 MoE

PEER는 가벼운 전문가를 사용하므로 파라미터 효율적 MoE와 관련이 있습니다(Wang et al., 2022; Zadouri et al., 2024). 이 방법들은 전체 크기의 FFW 대신 파라미터 효율적 미세 조정(PEFT) 어댑터를 전문가로 사용하여, 큰 백본 모델을 한 번만 저장하며 미세 조정 중 업데이트할 매개변수 수를 최소화합니다. PEER에서는 MoE 레이어에서 활성 매개변수 수가 작아 학습 및 추론 중의 FLOPs와 활성화 메모리 소모에 영향을 미칩니다. PEER는 다수의 PEFT 어댑터를 검색하는 방식으로도 확장 가능할 것입니다.

검색 강화 모델 (Retrieval-Augmented Models)

PEER의 검색 메커니즘은 대형 모델에서 외부 데이터베이스에서 지식을 검색하여 성능과 효율성을 개선하는 검색 강화 모델들과도 관련이 있습니다(Khandelwal et al., 2019; Borgeaud et al., 2022a; Guu et al., 2020). 이러한 모델은 토큰, 청크, 지식 그래프 등 다양한 형태의 데이터를 검색하지만, PEER는 데이터가 아닌 학습된 함수(전문가)를 검색한다는 점에서 차별화됩니다.

효율적인 피드포워드 레이어

피드포워드 네트워크의 효율성을 높이기 위한 연구는 오래전부터 있어 왔습니다. PEER와 유사하게 대부분의 접근법은 조건부 계산(Conditional Computation)을 기반으로 하고 있습니다. 예를 들어, Davis & Arel (2013)은 선형성 이전의 활성화를 통해 활성 뉴런의 일부만 계산하는 방식으로 비용을 줄였습니다. Belcak & Wattenhofer (2023)는 균형 잡힌 이진 트리를 사용하여 계산할 뉴런 블록을 선택하는 Fast FeedForward(FFF) 레이어를 소개했습니다. 이 방식은 추론 중에 O(\log(N))의 복잡도를 가지지만 학습 중에는 모든 노드가 활성화되어 O(N)의 복잡도가 필요합니다. PEER 레이어의 경우, PKM(Product Key Memory) 레이어와 유사하게 제품 키 기반 검색을 통해 효율성을 높였으나, PKM이 메모리 벡터를 검색하는 것과 달리 PEER는 입력에 따라 달라지는 전문가 네트워크를 검색하여 더 높은 효율성을 달성합니다.

결론

이 연구는 미세-그레인 MoE 스케일링 법칙을 바탕으로 매우 넓은 밀집 피드포워드 레이어를 다수의 작은 전문가로 분해하는 세밀한 MoE 아키텍처를 제안합니다. 많은 전문가에 대한 라우팅의 계산 부담을 극복하기 위해 제품 키를 활용하여 광범위한 MLP 레이어 내에서 소수의 숨겨진 뉴런을 선택합니다. 언어 모델링 실험을 통해 동일한 계산 예산 내에서 PEER가 밀집 Transformer, Coarse-grained MoE 및 Product Key Memory 레이어보다 훨씬 우수한 성능을 발휘함을 입증했습니다.