FineRMoE: 세밀한 전문가 설계를 위한 차원 확장 및 업사이클링 접근법

기존 MoE 모델은 전문가의 크기를 줄여 효율성을 높였으나 특정 지점에서 성능 향상이 멈추는 한계가 있었다. 이 논문은 전문가의 출력 차원까지 세분화하는 새로운 구조를 통해 파라미터 효율성을 6배 높이고 추론 지연 시간을 획기적으로 단축했다.

Mixture-of-Experts(MoE)는 전체 파라미터 중 일부만 활성화하여 연산량을 줄이는 기법이다. 기존의 'Fine-grained MoE'는 전문가 내부의 중간 차원 크기를 줄여 전문가 수를 늘리는 방식에 집중했다. 하지만 연구 결과, 중간 차원만 줄이는 방식은 특정 임계값을 넘으면 오히려 모델 성능이 저하되는 한계에 부딪혔다. FineRMoE는 이 문제를 해결하기 위해 전문가의 '출력 차원'까지 세분화하는 아이디어를 도입했다. 이는 Multi-Head Attention에서 각 헤드가 서로 다른 특징을 추출하도록 출력 차원을 나누는 것과 유사한 원리다. 전문가의 출력 차원을 줄임으로써 각 전문가가 더 독립적이고 특화된 표현을 학습하도록 유도하여 중복성을 억제하고 전문성을 높였다. 이렇게 세분화된 전문가들을 효율적으로 관리하기 위해 '차원 축소 후 복원'하는 이단계 구조를 사용한다. 결과적으로 기존 방식보다 훨씬 적은 파라미터 활성만으로도 동등하거나 더 높은 성능을 내며, 특히 추론 시 데이터 처리량을 100배 이상 개선했다.

FineRMoE 아키텍처는 공유 전문가와 희소 세분화 전문가로 구성된다. 전문가의 구조는 Up-projection, Gate, Down-projection 가중치 행렬로 이루어지며, 입력 x에 대해 행렬 곱과 활성화 함수를 거쳐 출력을 산출한다. [입력 벡터 x → 가중치 행렬과 곱셈 및 활성화 함수 적용 → 중간 결과 생성 → 출력 가중치 행렬과 곱셈 → 최종 전문가 출력 산출] 핵심은 Sparse Concatenation Layer와 Sparse Sum Layer의 이단계 구조다. Sparse Sum Layer에서는 각 그룹 내의 세분화된 전문가들의 출력을 가중 합산하여 차원이 축소된 후보 벡터들을 생성한다. Sparse Concatenation Layer는 이 후보 벡터들 중 최적의 것을 선택하여 연결함으로써 원래의 은닉 차원 크기를 복원한다. [세분화된 전문가 출력들 → 가중치 기반 합산 → 차원 축소된 벡터 생성 → 선택된 벡터들 연결 → 원래 차원의 출력 완성] 효율적인 제어를 위해 단일 라우터 메커니즘을 사용한다. 하나의 라우터가 계산한 점수를 바탕으로 두 레이어의 활성화를 동시에 결정하여 활성화 충돌을 방지하고 파라미터 오버헤드를 최소화한다. [입력 토큰 → 단일 라우터 네트워크 통과 → 전문가 및 벡터 선택 마스크 생성 → 두 레이어 동시 제어] 사전 학습된 모델을 활용하는 업사이클링 기법을 통해 기존 FFN의 가중치를 중간 차원과 출력 차원 방향으로 유연하게 분할하여 전문가를 초기화한다. 이를 통해 무작위 초기화 상태에서 학습을 시작하는 막대한 비용을 피하고 기존 모델의 지식을 보존했다.

Qwen 2.5 모델을 기반으로 50B 토큰의 추가 학습을 진행한 결과, 10개의 표준 벤치마크에서 기존 MoE 기법들을 압도하는 성능을 보였다. 특히 7B 모델 기준 FineRMoE는 평균 점수 70.04를 기록하여 강력한 베이스라인인 NVShard(65.18)를 크게 앞질렀다. 추론 효율성 측면에서 비약적인 발전을 이루었다. 가장 강력한 비교 대상인 C32A2 모델과 비교했을 때, 첫 토큰 생성 시간(TTFT)은 281배 빨라졌으며(178.3ms vs 50245.9ms), 초당 토큰 생성량은 136배 향상되었다(27.3 tokens/s vs 0.2 tokens/s). Ablation Study를 통해 중간 차원과 출력 차원 모두를 세분화했을 때 가장 낮은 전문가 간 유사도를 보임을 확인했다. 이는 각 전문가가 서로 중복되지 않는 고유한 정보를 학습하고 있음을 증명하며, 단일 라우터 설계가 분리된 라우터 설계보다 성능이 우수함도 입증되었다.

FineRMoE는 전문가의 세분화 단위를 정의하기 위해 4개의 하이퍼파라미터(중간 차원 세분도, 중간 차원 확장률, 출력 차원 세분도, 출력 차원 확장률)를 도입한다. 이를 통해 전문가의 수와 활성화 규모를 정밀하게 조절한다. Sparse Concatenation Layer의 출력은 출력 차원 세분도만큼의 구성 요소를 연결하여 형성된다. 각 구성 요소는 후보 벡터 중 라우터 점수가 가장 높은 하나를 선택하는 Top-1 Select 방식을 취한다. Sparse Sum Layer에서는 전문가들이 여러 그룹으로 나뉘며, 각 그룹은 세분화된 전문가들로 구성된다. 후보 벡터는 해당 그룹 내 전문가 출력들의 가중 합산으로 계산되며, 이때 가중치는 라우터가 계산한 소프트맥스 확률값을 사용한다. 단일 라우터 메커니즘은 점수 행렬을 생성한 후 이를 각 레이어에 필요한 마스크로 재구성한다. 두 마스크의 요소별 논리곱(AND) 연산을 통해 최종적으로 활성화될 전문가를 결정함으로써 연산 일관성을 유지한다. 업사이클링 과정에서 기존 출력 가중치 행렬은 중간 차원과 출력 차원 양방향으로 분할되어 세분화된 전문가의 가중치로 할당된다. 이는 기존 모델의 가중치 분포를 MoE 구조에 맞게 재배치하는 핵심 알고리즘이다.