🤖 machine learning

Sticky Routing: Training MoE Models for Memory-Efficient Inference

이 논문은 StickyMoE를 제안하는데, 이는 미세한 퍼플렉시티 저하를 최소화하면서 엣지 디바이스에서의 메모리 집약적인 가중치 스와핑을 크게 줄이기 위해, 의미론적으로 일관된 토큰 구간 전반에 걸쳐 전문가 할당을 유지하도록 Mixture-of-Experts 모델을 학습시키는 미분 가능한 라우팅 일관성 손실 함수이다.

원저자: Ali Kayyam

게시일 2026-07-13

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Ali Kayyam

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신에게 64명의 서로 다른 전문가 사서들이 있는 거대한 도서관이 있다고 상상해 보세요. 하지만 당신의 아주 작은 엣지 디바이스(예: 스마트폰이나 스마트 스피커)에는 책상 위에 한 번에 2명의 사서만 올려둘 수 있는 정도의 공간밖에 없습니다. 나머지 사서들은 느리고 먼지 쌓인 지하실에 보관되어 있습니다.

표준적인 "전문가 혼합(Mixture-of-Experts, MoE)" 모델에서는, AI가 새로운 단어를 읽을 때마다 어떤 사서에게 물어볼지 결정하기 위해 동전을 던집니다. 문제는 이 동전을 매 단어마다 던진다는 것입니다. 그래서 AI는 "사과(apple)"를 위해 사서 A에게 물어본 다음, "파이(pie)"를 위해 사서 B에게 물어보고, "크러스트(crust)"를 위해 사서 C에게 물어봅니다. 매번 바뀔 때마다 AI는 지하실로 달려가서 사서 B를 쫓아내고 사서 C를 끌고 올라와야 합니다. 이렇게 "왔다 갔다 하는" 과정이 너무 느려서, AI는 이야기를 읽는 시간보다 사서를 데려오는 데 더 많은 시간을 쓰게 됩니다.

StickyMoE는 모델이 이렇게 동전을 너무 무작정 던지지 않도록 학습시키는 새로운 방법입니다. 대신 모델은 "끈기 있게(sticky)" 행동하는 법을 배웁니다. 만약 모델이 "사과"를 위해 사서 A에게 물어봤다면, 주제가 논리적으로 이어지는 한 "파이"와 "크러스트"에 대해서도 사서 A에게 계속 물어보도록 권장되는 방식입니다.

핵심 발견: 제대로 훈련시켜라, 나중에 고치려 하지 말고

저자들은 이 "왔다 갔다 하는" 문제를 해결하기 위해 몇 가지 방법을 시도했으며, 다음과 같은 결과를 얻었습니다.

"사후(Post-Hoc)" 수정의 실패: 그들은 이미 훈련된 모델을 가져와서 동전 던지기 메커니즘(라우터)을 덜 무작위적이도록 "미세 조정(fine-tuning)"하는 시도를 했습니다. 이를 ReMoE라고 불렀습니다. 결과는? 실패했습니다. 전환율은 거의 변하지 않았습니다(절대 변화량 0.5% 미만). 저자들은 이미 사서들(전문가들)이 특정하고 고립된 단어들을 위한 전문가로 훈련되었기 때문이라고 주장합니다. 만약 사서들 자체가 긴 연속적인 주제를 다룰 준비가 되어 있지 않다면, 동전 던지기 기계에게 "더 끈기 있게" 행동하라고 말하는 것은 소용이 없습니다. 이는 마치 단거리 선수의 출발 자세만 바꾼다고 해서 그를 마라톤 선수로 만들 수 없다고 말하는 것과 같습니다. 근육(전문가 가중치) 자체가 짧은 폭발력을 위해 만들어졌기 때문입니다.
"강제적(Hard)" 수정의 경직성: 그들은 또한 엄격한 규칙("Hard-Window")을 사용하여 모델이 정해진 단어 수 동안 한 명의 사서에게 붙어 있도록 강제하는 방법도 시도했습니다. 이 방법은 전환을 줄였지만, AI의 답변 품질을 크게 떨어뜨렸습니다(중형 모델에서 퍼플렉시티(perplexity)가 최대 6.8% 증가). 이는 마치 한 사서에게 장르가 완전히 바뀌더라도 소설 한 권을 통째로 읽으라고 강요하는 것과 같았으며, 그 결과 이야기의 질이 저하되었습니다.
"Sticky" 솔루션의 성공: 승리한 방법인 StickyMoE는 훈련의 매우 첫 단계에 부드러운 "일관성 페널티(consistency penalty)"를 추가합니다. 이는 모델에게 이렇게 말하는 것과 같습니다: "이봐, 만약 바로 옆에 있는 두 단어 사이에서 전문가를 바꾼다면, 그것은 약간의 페널티 대상이야." 이를 통해 모델은 언어를 배우는 동시에 '끈기'를 함께 학습하게 됩니다.

결과: 끈기 있는 것이 더 낫다

그들이 WikiText-2 데이터셋(위키피디아 기사 모음)을 사용하여 소형 및 중형 모델에 테스트했을 때, 결과는 명확했습니다.

적은 전환: 단 하나의 노브(λ라고 불리는 숫자)를 조절함으로써, 전문가를 교체하는 횟수를 최대 **59%**까지 줄일 수 있었습니다. 예를 들어, 중형 모델의 경우 전환율이 0.71(거의 매번 전환함)에서 0.29로 떨어졌습니다.
더 나은 속도 (시뮬레이션): 모델이 덜 전환했기 때문에, 2명의 전문가를 담을 수 있는 캐시(책상 선반)를 시뮬레이션했습니다. "캐시 적중률(Cache Hit Rate, 맞는 사서가 이미 책상 위에 있는 비율)"이 0.54에서 0.88로 급증했습니다. 이는 모델이 이전보다 지하실로 달려가야 하는 횟수가 3.92배 적게 된다는 것을 의미합니다.
더 나은 품질: 놀랍게도, 중형 모델에서 적당한 양의 "끈기"(λ = 0.05)를 사용했을 때 AI는 오히려 더 똑똑해졌으며, 퍼플렉시티(혼란도를 나타내는 척도)를 4.1% 낮추었습니다. 모델은 주제에 집중하는 것이 문맥을 이해하는 데 도움이 된다는 것을 배웠습니다.
붕괴 없음: 모델이 게을러져서 하나의 사서만 모든 것에 사용하는 것이 아니냐는 큰 우려가 있었으나, "이용 엔트로피(Utilisation Entropy)"는 높은 수준(최대 2.0 비트 중 1.92 비트 이상)을 유지했습니다. 즉, 4명의 전문가가 여전히 공평하게 사용되고 있었습니다. 모델은 게을러진 것이 아니라, 조직적으로 변한 것입니다.

무엇이 효과가 없었나 (그리고 왜인가)

논문은 모델이 구축된 후에 이 문제를 해결할 수 있다는 생각을 명시적으로 부정합니다. ReMoE 방식(기존 모델을 미세 조정하는 것)은 효과가 없는 것으로 나타났는데, 이는 "전문가 표현(expert representations, 사서들의 뇌)"이 이미 주제를 지속하는 데 신경 쓰지 않는 훈련 과정에 의해 형성되었기 때문입니다. 국소성(locality)은 사후에 덧붙일 수 있는 것이 아니라, 처음부터 심어 넣어야 하는 것입니다.

또한, 모델의 가장 첫 번째 레이어(Layer 0)가 가장 "끈기 있게" 만들기 어렵다는 점을 언급했습니다. 최상의 설정을 사용하더라도 이 레이어의 전환율은 0.49 이상으로 유지되었습니다. 이는 첫 번째 레이어가 문맥과 섞이기 전의 가공되지 않은 단어들을 보기 때문에, "사과"와 "파이"가 이 단계에서는 실제로 매우 다르게 보일 수 있기 때문입니다.

결론

이 논문은 만약 당신이 메모리가 제한된 기기에서 빠르게 작동하는 AI를 원한다면, 단순히 똑똑한 캐싱 규칙을 적용하거나 나중에 모델을 고치려고 애쓰지 말아야 한다고 제안합니다. 대신, 훈련 첫날부터 모델이 "끈기 있게" 행동하도록 훈련시켜야 합니다. 인접한 단어 사이에서 전문가를 너무 자주 바꾸는 것에 대해 간단한 페널티를 부여함으로써, 모델은 자연스럽게 주제에 집중하고, 느린 지하실로 가는 횟수를 줄이며, 실제로 이야기를 더 잘 이해하게 됩니다.

저자들은 이 연구가 시뮬레이션에 기반하며, 소형 및 중형 모델(최대 22M 파라미터)의 특정 데이터셋(WikiText-2)에서 수행되었다는 점을 인정합니다. 그들은 모델이 커질수록 이 "끈기"가 더 잘 작동할 수 있다고 제곡하지만, 아직 수십억 개의 파라미터를 가진 거대 모델에 대해서는 테스트하지 못했습니다. 또한, 현재 방식은 문장이 끝나는 시점을 알지 못하므로, 문단 구분 사이에서도 "너무 끈끈하게" 행동할 수 있다는 점을 언급하며, 이는 향후 해결해야 할 과제로 남겨두었습니다.

기술 요약: 메모리 효율적인 MoE 추론을 위한 스티키 라우팅(Sticky Routing)

1. 문제 정의: 전문가 교체(Expert-Swapping)의 병목 현상

혼합 전문가(Mixture-of-Experts, MoE) 모델은 토큰당 활성화되는 전문가의 하위 집합만을 사용하여 대규모 언어 모델(LLM)을 확장하도록 설계되었습니다. 이론적으로 이는 방대한 파라미터 수를 보유하면서도 토큰당 계산량은 낮게 유지할 수 있게 해주며, 이는 특정 시점에 소수의 전문가만이 빠른 메모리(VRAM/SRAM)에 상주하면 되는 엣지 배포 환경에서 매우 매력적입니다.

그러나 본 논문은 실제 적용 시 발생하는 결정적인 실패 모드를 식별했습니다: 바로 **시간적 라우팅 불일치(temporal routing inconsistency)**입니다. 표준 MoE 라우터는 오직 해당 토큰의 은닉 표현(hidden representation)에만 기반하여 각 토크에 대해 독립적인 결정을 내립니다. 결과적으로, 연속된 토큰들이 완전히 다른 전문가들을 빈번하게 활성화하게 됩니다. 메모리가 제한된 하드웨어에서, 이러한 전환(switch)이 발생할 때마다 캐싱된 전문가 가중치를 비우고 새로운 가중치를 느린 저장 장치(CPU RAM/Flash)로부터 로드해야 합니다. PCIe/NVMe 대역폭은 GPU 메모리 대역폭보다 몇 자릿수나 낮기 때문에, 이러한 빈번한 "전문가 교체"는 추론 지연 시간을 지배하며, 종종 MoE 모델을 유사한 크기의 밀집(dense) 모델보다 더 느리게 만듭니다.

기존 솔루션들은 불충분한 것으로 간-주됩니다:

시스템 수준 접근 방식 (예: 캐싱 휴리스틱, 프리페처): 이들은 존재하는 국부성(locality)을 활용하는 반응적(reactive) 방식일 뿐, 본질적으로 불일치한 라우터를 수정할 수는 없습니다.
사후 처리(Post-hoc) 접근 방식 (예: 라우터만 미세 조정): 라우터가 국부성을 유도하도록 설계되지 않은 표현 공간 내에서 결정을 재정렬하기 어렵기 때문에, 전문가 표현이 국부성을 고려하지 않은 학습 체계에 의해 형성되었다면 실패하게 됩니다.
아키텍처 재설계 (예: Oracle-MoE): 모델이나 학습 파이프라인에 상당한 구조적 변경을 요구합니다.

2. 방법론: StickyMoE

저자들은 StickyMoE를 제안하며, 이는 라우팅의 시간적 국부성을 명시적인 목적 함수로 취급하는 학습 시점의 접근 방식입니다. 핵심 혁신은 인접한 토큰 간의 전문가 할당 변화를 징벌하는 **미분 가능한 라우팅 일관성 손실(differentiable routing consistency loss, $L_{cons}$ )**입니다.

라우팅 일관성 손실

손실 함수는 연속된 토큰의 소프트맥스 게이트 확률 벡터( $g_t$ ) 사이의 $\ell_2$ 거리를 최소화합니다:
$L_{cons} = \frac{1}{T-1} \sum_{t=2}^{T} \| g_t - g_{t-1} \|_2^2$
여기서 $g_t \in \mathbb{R}^N$ 은 토큰 $t$ 에 대한 라우터의 출력입니다.

이 손실은 표준 교차 엔트로피( $L_{CE}$ ) 및 부하 분산( $L_{bal}$ ) 손실과 함께 전체 학습 목적 함수에 통합됩니다:
$L = L_{CE} + \lambda L_{cons} + \mu L_{bal}$

주요 메커니즘

공동 적응(Co-adaptation): 사후 처리 방식과 달리, StickyMoE는 라우터와 전문가 표현을 동시에 학습시킵니다. $L_{cons}$ 를 최소화함으로써 라우터 가중치( $W_r$ )를 업데이트하고 은닉 상태( $h_t$ )를 통해 역전파하여, 모델이 의미론적으로 일관된 구간 동안 자연스럽게 일관성을 유지하는 표현을 학습하도록 유도합니다.
Soft-Hard 변형: "장거리 표류(long-range drift)"(작은 단계별 변화가 누적되어 먼 거리의 전환으로 이어지는 현상)를 해결하기 위해, 저자들은 앵커 제약 조건을 도입했습니다. 시퀀스를 윈도우 단위로 나누고, 윈도우의 첫 번째 토큰이 앵커를 설정하며, 이후의 토큰들은 이 앵커로부터 벗어나는 것에 대해 윈도우 시작점으로부터의 거리에 따라 선형적으로 증가하는 페널티를 받습니다.
호환성: 이 방법은 아키텍처 변경을 요구하지 않으며 계산 오버헤드가 미미합니다. 또한 모든 표준 top- $k$ MoE 및 기존 부하 분산 기술과 호환됩니다.

3. 주요 기여

라우팅 일관성 손실: 단일 하이퍼파라미터( $\lambda$ )만을 필요로 하며, 인접 토큰 간의 전문가 전환을 직접적으로 징벌하는 단순하고 미분 가능한 학습 목적 함수를 제공합니다.
부하 분산과의 상보성: 실험적 증거는 일관성 손실과 부하 분산 손실이 상보적임을 보여줍니다. 일관성 손실은 "전문가 붕괴(expert collapse)"를 일으키지 않으면서(이용 엔트로피가 이론적 최대치 근처로 높게 유지됨) 언제 전문가가 사용될지를 형성합니다.
Soft-Hard 변형: 단계별 소프트 페널티와 세그먼트 수준의 앵커 제약을 결려하여, 추가적인 품질 비용 없이 강력한 국부성을 보장합니다.
포괄적 평가: 소규모 및 중규모 모델에 대해 Vanilla MoE, 사후 라우터 미세 조정(ReMoE), 아키텍-구조적 베이스라인(Oracle-MoE)과 비교하는 통제된 실험을 수행했습니다.

4. 실험 결과

WikiText-2 데이터셋을 사용하여 소규모(~~8.8M 파라미터) 및 중규모(~~22M 파라미터) MoE 모델에 대한 실험을 진행했습니다.

전문가 전환율 (Expert Switch Rate, SR): StickyMoE는 전문가 전환율을 최대 **59%**까지 감소시켰습니다 (예: 중규모 모델에서 0.71에서 0.29로 감소).
캐시 적중률 (Cache Hit Rate, CHR): 전환율의 감소는 중규모 모델에서 2-전문가 슬롯 캐시를 사용할 때 캐시 미스를 3.92배 감소시켰습니다 (CHR이 약 0.54에서 0.88로 증가).
퍼플렉시티 (품질):
- 적절한 $\lambda$ 값(예: 0.05)에서, StickyMoE는 베이스라인 대비 퍼플렉시티를 개선했습니다 (예: 중규모 모델에서 -4.1%). 이는 일관성 손실이 유익한 정규화 역할을 함을 시사합니다.
- 공격적인 $\lambda$ 값은 트레이드오프를 유발하지만, StickyMoE의 파레토 최적선(Pareto frontier)은 다른 방법들을 엄격히 압도합니다.
사후 미세 조정의 실패: ReMoE 베이스라인(사전 학습된 모델의 라우터만 미세 조정)은 전환율을 유의미하게 줄이는 데 실패했습니다 (<0.5% 변화), 이는 표현이 고정된 후에는 국부성을 사후에 주입할 수 없음을 확인해 줍니다.
계층별 분석: 일관성 손실은 깊은 층(L1–L3)에서 가장 효과적이었습니다. 입력 층(L0)은 토큰 임베딩에 컨텍스트 혼합이 부족하기 때문에 국부성 개선에 저항성을 보였습니다.

5. 의의 및 주장

본 논문은 라우팅의 시간적 국부성이 시스템 수준의 사후 고려 사항이 아니라, 사전 학습 중에 반드시 심어져야 하는 정당한 학습 목적 함수라고 주장합니다.

패러다임 전환: 양자화 인식 학습(Quantization-aware training)이 표준이 된 것처럼, 저자들은 메모리 제한적 배포를 목적으로 하는 MoE 모델의 기본값으로 "국부성 인식 학습(locality-aware training)"이 되어야 한다고 단언합니다.
근본 원인 해결: StickyMoE는 증상(캐시 미스)이 아닌 근본 원인(불일치하는 라우팅 결정)을 해결합니다. 표현과 라우팅 결정을 공동 학습함으로써, 사후 처리 방식이 겪는 표현 불일치 문제를 피합니다.
실용성: 이 방법은 아키텍처에 무관하며 추론 파이프라인의 변경을 요구하지 않습니다. 학습된 라우터는 자연스럽게 일관된 게이트 분포를 생성하여, 추가적인 엔지니어링 없이도 표준 캐싱 전략(LRU, LFU)이 훨씬 더 효율적으로 작동할 수 있게 합니다.

저자들은 시스템 수준의 최적화와 사후 조정이 유용한 보완책이 될 수는 있지만, 애초에 캐시 친화적으로 학습되지 않은 모델을 대체할 수는 없다고 결론짓습니다. StickyMoE는 엣지 디바이스에서 MoE 아키텍처의 이론적 메모리 효율성을 실현하기 위한 경로를 제시합니다.