동적 전문가 공유: 병렬 디퓨전 LLM에서 메모리와 병렬성 분리

초록

본 논문은 병렬 디코딩을 사용하는 확산 대형 언어 모델(dLLM)에서 Mixture‑of‑Experts(MoE) 구조가 초래하는 “전문가 폭발” 문제를 해결한다. 저자는 토큰‑단위가 아닌 시퀀스‑레벨에서 전문가 집합을 선택하는 Dynamic Expert Sharing(DES) 기법을 제안하고, 두 가지 구현인 DES‑Seq와 Saliency‑Aware Voting(DES‑Vote)을 통해 고유 활성 전문가 수를 55 % 이상 감소시키고 지연 시간을 최대 38 % 줄이면서도 정확도는 99 % 수준으로 유지한다.

상세 분석

이 논문은 병렬 디퓨전 LLM(dLLM)과 MoE 결합 시 발생하는 메모리 병목 현상을 정량적으로 분석한다. 기존 MoE는 토큰당 K개의 전문가를 독립적으로 선택하므로, 병렬 블록 크기 N이 증가할수록 고유 활성 전문가 집합 |∪ₙSₙ|이 거의 선형적으로 늘어나 HBM→SRAM 전송 비용(b·|∪ₙSₙ|)이 지배적인 메모리‑바운드 상황을 만든다. 저자는 이 비용을 토큰‑레벨에서 시퀀스‑레벨로 전환함으로써 해결하고자 한다. 핵심 아이디어는 “코어셋(coreset)”이라 불리는 작은 전문가 서브셋 C⊂{E₁…E_M}을 런타임에 동적으로 선정하고, 모든 토큰이 Top‑K 라우팅을 이 코어셋 안에서만 수행하도록 강제하는 것이다. 이렇게 하면 메모리 전송 비용은 b·|C| 로 축소되고, 연산 비용 a·(N·K)는 기존과 동일하게 유지된다.

두 가지 선택 전략이 제안된다. 첫 번째인 DES‑Seq는 각 토큰이 독립적으로 선택한 전문가들의 합집합을 그대로 코어셋으로 사용한다. 이는 구현이 간단하고, 토큰 간 중복을 자연스럽게 활용하지만, 불필요한 전문가가 포함될 가능성이 있다. 두 번째인 DES‑Vote는 각 토큰의 라우터 로짓을 집계해 가중 평균을 구하고, 이 집계값을 기반으로 가장 높은 “saliency”를 보이는 전문가들을 투표 방식으로 선정한다. 이 과정은 토큰 전체의 의미적 복잡성을 반영해 실제로 필요한 전문가만을 남기므로 코어셋 크기를 더 작게 만들 수 있다.

수학적으로는 코어셋 선택 함수를 Φ(I)라 두고, 다음 최적화 문제를 정의한다.
 Φ* = arg min_{Φ} |Φ(I)|  subject to A(Φ(I)) ≥ A_base − ε
여기서 A(·)는 모델 정확도, ε는 허용 오차이다. 논문은 이 문제를 근사적으로 해결하기 위해 greedy‑based top‑k 선택과 라우터 가중치 정규화를 결합한 알고리즘을 제시한다.

실험에서는 LLaD‑A‑MoE와 LLaD‑A2.0‑mini 두 모델에 DES를 적용했으며, HumanEval, MBPP, MATH500, GSM8K 등 네 가지 벤치마크에서 평균 고유 전문가 수를 55 % 이상 감소시켰다. 지연 시간은 블록 크기 16~64 토큰에서 최대 38 % 감소했으며, 정확도는 99 % 수준(≤1 % 손실)으로 유지되었다. 특히 DES‑Vote는 활성 전문가 수와 정확도 사이의 파레토 프론티어를 크게 확장해, 동일한 정확도를 유지하면서도 메모리 트래픽을 크게 절감한다는 점이 강조된다.

이 접근법은 기존의 토큰‑중심 전문가 스키핑(예: expert skipping, OEA)과는 근본적으로 다르다. 토큰별 독립 라우팅을 포기하고 시퀀스 전체의 컨텍스트를 활용함으로써, 메모리‑바운드 구간에서의 병목을 직접 해소한다. 또한 코어셋 선택은 오프라인 프루닝이나 양자화와도 호환 가능하므로, 향후 하드웨어‑특화 최적화와 결합할 여지가 크다.