dMoE: dLLMs with Learnable Block Experts

당신이 엄청나게 똑똑하고 거대한 도서관(즉, Diffusion Large Language Model, 또는 dLLM)을 가지고 있다고 상상해 보세요. 이 도서관은 이야기를 쓰고, 수학 문제를 풀고, 질문에 답할 수 있습니다. 이 도서관을 너무 무겁게 만들지 않으면서도 더 똑똑하게 만들기 위해, 설계자들은 **Mixture-of-Experts (MoE)**라는 특별한 기능을 추가했습니다.

여기서 "전문가(Experts)"는 100명의 전문 사서 팀이라고 생각하면 됩니다. 질문을 던질 때, 도서관은 100명의 사서 모두에게 질문을 던지지 않습니다. 대신 "라우터(Router)"(사서 매니저)가 당신의 특정 질문에 가장 적합한 상위 8명의 사서를 선발합니다. 이 방식 덕분에 작업이 빠르고 효율적으로 진행됩니다.

문제점: "솔로 하이킹" vs. "그룹 하이킹"

여기서 문제가 발생했습니다.

• 기존 방식 (Autoregressive Models): 마치 한 명의 하이커가 산을 오를 때 한 걸음씩 이동하는 것과 같습니다. 매 걸음마다 하이커는 매니저에게 묻습니다. "지금 당장 누구의 도움을 받아야 하나요?" 그러면 매니저는 8명의 전문가를 뽑고, 그들이 도와주면, 하이커는 다음 걸음을 내딛습니다.
• 새로운 방식 (dLLMs): 새로운 dLLM은 밀집된 그룹으로 그룹 하이킹을 하는 것과 같습니다. 이들은 한 명씩 움직이지 않습니다. 한 번에 32명의 인원이 하나의 큰 블록(block)으로 움직입니다. 이들은 앞길 전체를 내다보고 여러 단계를 동시에 수정할 수 있습니다. 훨씬 빠르죠!

불일치(Mismatch):
문제는 "매니저(라우터)"가 여전히 솔로 하이커를 상대하는 것처럼 행동하고 있었다는 점입니다. 그룹이 32명 단위의 블록으로 움직임에도 불구하고, 매니저는 이 블록 안의 모든 사람을 각각 별개의 솔로 하이커로 취급했습니다.

• 블록의 1번 사람이 도움을 요청합니다: 매니저가 8명의 전문가를 뽑습니다.
• 2번 사람이 도움을 요청합니다: 매니저가 다른 8명의 전문가를 뽑습니다.
• ...
• 32번 사람이 도움을 요청합니다: 매니저가 또 다른 8명의 전문가를 뽑습니다.

결과적으로 그룹이 함께 움직이고 있음에도 불구하고, 매니저는 단 하나의 움직임을 처리하기 위해 수십 명의 서로 다른 전문가를 호출하게 됩니다. 이것은 마치 버스 운전사가 32대의 서로 다른 차에 기름을 채우기 위해 32개의 서로 다른 주유소에 들르는 것과 같습니다. 시스템은 한 번에 너무 많은 전문가를 불러오려고 시도하기 때문에, 메모리(연료)가 바닥나고 전체 프로세스가 느려지며 버스가 멈춰 서게 됩니다.

해결책: dMoE (The "Block Manager")

저자들은 dMoE라고 불리는 새로운 시스템을 제안합니다. 그들은 그룹이 함께 움직인다면, 그들을 선택할 때 하나의 단일 단위로 취급해야 한다는 것을 깨달았습니다.

dMoE가 어떻게 작동하는지 쉬운 비유를 통해 설명하겠습니다.

1. 그룹 투표: 32명의 사람에게 각자 누구의 도움이 필요한지 개별적으로 묻는 대신, dMoE는 그룹 전체에게 물어봅니다. "이 블록 전체가 집단적으로 어떤 도움이 필요한가요?"
2. 통합 리스트: 매니저는 모든 개인의 투표를 모아 하나의 "블록 점수(Block Score)"로 결합합니다. 만약 블록 내의 20명이 "수학 전문가 E1"을 필요로 하고, 10명이 "수학 전문가 E2"를 필요로 한다면, 매니저는 이 블록 전체가 정말로 E1과 E2를 필요로 한다는 것을 파악합니다.
3. 스마트한 요약 목록: 매니저는 이 결합된 리스트를 보고 다음과 같이 결정합니다. "좋아, 이 블록 전체를 위해서, 우리는 그룹의 니즈를 90% 충족할 수 있는 상위 전문가들만 있으면 돼." 그리고 이 블록을 위한 작은 고정 리스트(코어셋, coreset)를 만듭니다.
4. 결과: 이제 시스템은 한 블록을 처리하기 위해 60개 이상의 서로 다른 전문가를 불러오는 대신, 약 14명의 전문가만을 불러옵니다. 이는 버스 운전사가 "아, 이 버스에 탄 사람들은 모두 같은 3곳의 주유소에서 기름을 채워야 하는구나"라고 깨닫는 것과 같습니다.

이것이 왜 중요한가 (결과)

저자들은 이 새로운 "블록 매니저"를 최첨단 모델인 LLaDA2.0-mini에 적용하여 까다로운 수학 및 논리 테스트(MATH500, GSM8K 등)를 테스트했습니다.

• 적은 메모리 사용량: 방대한 양의 혼란스러운 전문가 리스트를 불러오는 대신, 컴퓨터의 메모리 사용량이 약 77%에서 80% 감소했습니다. 이는 100가지의 다양한 도구를 실은 트럭에서 필수적인 14가지만 담은 배낭으로 교체한 것과 같습니다.
• 빠른 속도: 모델이 1.14배에서 1.66배 더 빨라졌습니다. 버스가 주유소에 덜 들르게 된 것입니다.
• 지능의 손실 없음: 가장 놀라운 점은 모델이 멍청해지지 않았다는 것입니다. 모델은 원래의 지능을 99.11% 유지했습니다. 동일한 수의 수학 문제를 정확하게 풀어냈으며, 단지 훨씬 더 효율적으로 수행했을 뿐입니다.

핵심 요약

이 논문은 다음과 같이 말합니다: "우리는 이 새로운 빠른 AI 모델들이 그룹으로 작동할 때, 이들을 개별 단위로 취급하면 교통 체증이 발생한다는 것을 발견했습니다. 그룹이 필요한 전문가에 대해 함께 투표하도록 함으로써, 우리는 호출하는 작업자의 수를 거의 5배 가까이 줄이고, 메모리를 대폭 절약하며, 지능의 손실 없이 AI를 더 빠르게 실행할 수 있습니다."

이것은 "학습 가능한(learnable)" 전략입니다. 즉, AI는 훈련 과정에서 이러한 투표를 그룹화하는 방법을 학습하며, 이를 통해 스마트하고 자동화된 해결책을 만들어냅니다.