🤖 machine learning

MoE-nD: Per-Layer Mixture-of-Experts Routing for Multi-Axis KV Cache Compression

이 논문은 KV 캐시 압축을 위해 각 레이어가 서로 다른 토큰 제거 비율과 정밀도를 가질 수 있도록 하는 MoE-nD 프레임워크를 제안하여, 긴 컨텍스트 LLM 추론 시 메모리 사용량을 획기적으로 줄이면서도 정확도를 유지하거나 향상시킨다는 점을 보여줍니다.

원저자: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

게시일 2026-04-21

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

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

🧠 배경: AI 의 '작업대' 문제

인공지능이 긴 글을 읽거나 수학 문제를 풀 때, 중요한 정보들을 작업대 (KV Cache) 위에 펼쳐놓고 생각합니다. 하지만 이 작업대가 너무 커지면 컴퓨터 메모리가 부족해져서 AI 가 멈추거나 (OOM), 매우 느려집니다.

기존 방법들은 이 작업대를 정리할 때 모든 층 (Layer) 에 똑같은 규칙을 적용했습니다.

"모든 층에서 50% 는 버려라." (일괄 삭제)
"모든 층의 숫자 크기를 반으로 줄여라." (일괄 압축)

하지만 이 방법은 **"모든 층이 똑같다"**는 잘못된 가정에서 시작합니다.

🔍 발견: 층마다 성격이 다릅니다!

저자들은 AI 의 각 층을 다양한 직업을 가진 직원들로 비유할 수 있다고 말합니다.

초기 층 (신입): 중요한 핵심 정보를 많이 담고 있어서, 정보를 지우면 (삭제) 큰 실수가 나옵니다. 하지만 숫자 크기를 줄이는 것 (압축) 은 잘 견딥니다.
후기 층 (베테랑): 세부적인 계산에 집중하므로, 정보를 조금 지우는 것은 괜찮지만 숫자 크기를 줄이면 큰 오류가 생깁니다.

즉, 어떤 층은 '삭제'가 나쁘고, 어떤 층은 '압축'이 나쁜 것입니다. 그런데 기존 방법은 모든 직원에게 똑같은 규칙을 강요했기 때문에, AI 의 성능이 크게 떨어졌습니다.

💡 해결책: MoE-nD (맞춤형 관리자)

이 논문이 제안한 MoE-nD는 바로 이 문제를 해결하는 **'스마트한 관리자'**입니다.

전문가 팀 (Mixture-of-Experts):
- '삭제 전문가'와 '압축 전문가'가 있습니다.
- 관리자는 각 층 (직원) 의 성격을 미리 파악해 둡니다.
맞춤형 지시 (Routing):
- A 층: "너는 기억력이 중요하니까 정보를 지우지 말고, 대신 숫자 크기를 줄여라."
- B 층: "너는 세부 사항이 중요하니까 숫자 크기는 그대로 두고, 불필요한 정보는 50% 정도 지워라."
- 이렇게 층마다 가장 적합한 조합을 찾아서 적용합니다.
전체 예산 관리:
- 전체 메모리 예산 (작업대 크기) 을 정해두면, 관리자는 그 안에서 각 층에 가장 효율적으로 자원을 배분합니다.

📊 성과: 놀라운 결과

실험 결과, 이 방법을 쓰면 다음과 같은 기적이 일어납니다.

14 배 압축: 메모리 사용량을 1.9GB 에서 **136MB(약 14 배 줄임)**로 줄였는데도, 성능은 줄지 않았습니다.
기존 방법과의 차이: 같은 메모리 크기를 쓴 다른 방법들은 성능이 80% 이상 떨어졌지만, 이 방법은 원래 성능을 그대로 유지했습니다.
복잡한 문제 해결: 수학 문제 (AIME) 같은 어려운 추론 작업에서도 기존 방법보다 훨씬 높은 점수를 받았습니다.

🚫 언제 효과가 없을까요? (한계)

이 방법은 긴 글이나 긴 대화가 필요할 때 가장 빛을 발합니다.

짧은 질문 (예: "1+1 은?"): 정보가 너무 적어서 처음부터 작업대가 작습니다. 이 경우 모든 층에 '지우기'를 적용할 필요가 없으므로, 관리자가 "아, 다 지울 필요 없네?"라고 판단하고 아무것도 안 합니다. 이때는 기존 방법과 차이가 없습니다.
핵심: 정보가 길고 복잡할수록, 각 층에 맞는 '맞춤형 정리'가 필요하다는 뜻입니다.

🎯 한 줄 요약

"AI 의 기억 공간 (메모리) 을 정리할 때, 모든 층에 똑같은 규칙을 적용하지 말고, 각 층의 성격에 맞춰 '무엇을 지울지'와 '얼마나 압축할지'를 똑똑하게 배분하면, 메모리는 14 배 줄이면서도 성능은 그대로 유지할 수 있다!"

이 기술은 앞으로 AI 가 더 긴 책을 읽거나 더 복잡한 문제를 풀 수 있게 해주는 핵심 열쇠가 될 것입니다.

1. 문제 정의 (Problem)

KV 캐시 병목: 현대 LLM 은 긴 문맥을 처리할 때 시퀀스당 수십 GB 의 KV 캐시 메모리를 소모하며, 이는 추론 지연을 유발하고 상용 하드웨어에서 메모리 부족 (OOM) 을 초래합니다.
기존 방법의 한계: 기존 압축 기법들은 주로 4 차원 KV 텐서의 단일 축 (Sequence, Precision, Head dimension, Cross-layer) 중 하나에만 작용합니다.
- 동질성 문제 (Homogeneity): 모든 트랜스포머 레이어에 동일한 압축 비율 (Eviction ratio) 과 정밀도 (Quantization bits) 를 적용합니다.
- 비효율성: 저자들에 따르면, 서로 다른 레이어는 압축 작업 (토큰 제거 또는 양자화) 에 대해 매우 다르게 반응합니다. 즉, 모든 레이어에 동일한 '레시피'를 적용하는 것은 정확도 손실을 초래하거나 잠재적인 압축 기회를 놓치게 합니다.

2. 방법론 (Methodology: MoE-nD)

MoE-nD 는 전문가 혼합 (Mixture-of-Experts, MoE) 프레임워크를 KV 캐시 압축에 적용하여, 각 레이어가 글로벌 메모리 예산 내에서 최적의 압축 조합을 선택하도록 합니다.

핵심 아이디어: 압축 축 (Eviction 비율, K-정밀도, V-정밀도) 을 '전문가 (Expert)'로 간주하고, 각 레이어에 가장 적합한 조합을 라우팅합니다.
3 가지 조정 가능한 파라미터 (Knobs) per Layer:
1. Keep Ratio (Eviction): 캐시 토큰 중 유지할 비율 (예: 0.1 ~ 1.0).
2. K-bit Width: Key 텐서의 양자화 비트 수 (16, 8, 4 비트).
3. V-bit Width: Value 텐서의 양자화 비트 수 (16, 8, 4 비트).
오프라인 보정 (Offline Calibration):
- 모델의 각 레이어에 대해 다양한 압축 설정을 적용했을 때의 **정밀도 손실 (L2 오차)**을 측정하여 '민감도 테이블 (Sensitivity Table)'을 생성합니다.
- 이 데이터는 추론 시 라우팅 결정을 위한 기준이 됩니다.
탐욕적 솔버 (Greedy Solver):
- 주어진 글로벌 메모리 예산 내에서 전체 모델의 예측된 품질 손실을 최소화하는 레이어별 (keep_ratio, k_bits, v_bits) 튜플을 선택합니다.
- 단위 메모리 증가당 민감도 감소량 ( $\Delta S / \Delta m$ ) 이 가장 큰 레이어를 순차적으로 업그레이드하는 방식으로 동작합니다.
이질적 어텐션 패치 (Heterogeneous Attention Patch):
- 추론 시 각 레이어마다 다른 토큰 길이와 비트 폭을 처리할 수 있도록 어텐션 커널을 수정했습니다.
- 레이어별 토큰 위치를 추적하고 RoPE (Rotary Positional Embedding) 역변환을 수행하여, 레이어 간 토큰 위치 불일치를 해결합니다.

3. 주요 기여 (Key Contributions)

다중 축 (Multi-Axis) 레이어별 라우팅: 기존 연구는 단일 축 (Eviction 또는 Quantization) 만 레이어별로 조정하거나, 모든 축을 균일하게 적용했습니다. MoE-nD 는 Eviction 비율과 Quantization 비트를 동시에 레이어별로 최적화하는 최초의 방법입니다.
레이어별 이질성 (Heterogeneity) 발견: 실험을 통해 심층 레이어와 얕은 레이어가 압축에 반응하는 방식이 극명하게 다름을 입증했습니다. 특히, 공격적인 토큰 제거 (Eviction) 의 경우 레이어 간 민감도 차이가 최대 689 배까지 발생하여 레이어별 라우팅이 필수적임을 보였습니다.
효율적인 라우팅 전략: 실험 결과, 레이어별 Eviction 라우팅이 성능 향상의 주된 동력임을 확인했습니다. 반면, 레이어별 Quantization 라우팅만으로는 평균적인 성능 향상이 거의 없었습니다.

4. 실험 결과 (Results)

DeepSeek-R1-Distill-Qwen-7B 모델을 사용하여 LongBench-v1 및 AIME 벤치마크에서 평가했습니다.

LongBench-v1 (긴 문맥):
- **압축률 14 배 (136 MB)**에서 압축 없는 풀 캐시 (1.9 GB)baseline 과 동등한 정확도를 달성했습니다.
- 기존 단일 축 (1d) 또는 균일 2 축 (2d_uniform) 방법들은 동일한 메모리에서 8/100 미만의 점수를 기록하며 정확도가 급격히 하락했습니다.
- MoE-nD 의 이질적 (Hetero) 변형은 균일 2 축 방법보다 1.6~2 배 높은 F1 점수를 기록했습니다.
AIME (수학 추론):
- AIME-24 및 AIME-25 벤치마크에서 모든 예산 조건에서 기존 최강의 베이스라인보다 +6 ~ +27 포인트의 성능 향상을 보였습니다.
- 특히 예산이 매우 제한적일 때 (예: 64 토큰), 다른 방법들은 0 점에 수렴했으나 MoE-nD 는 유의미한 성능을 유지했습니다.
한계 및 범위 (Null Results):
- MATH-500 및 LongBench 의 TREC 과 같은 짧은 문맥 데이터셋에서는 성능 향상이 없었습니다. 이는 솔버가 대부분의 레이어에서 keep=1.0 (제거 없음) 을 선택하기 때문이며, 이는 방법론의 한계가 아니라 긴 문맥/긴 생성 작업에 특화된 방법임을 명확히 합니다.

5. 의의 및 결론 (Significance)

새로운 압축 패러다임: KV 캐시 압축이 단순히 "모든 레이어를 동일하게 줄이는 것"이 아니라, **"각 레이어의 특성에 맞춰 다른 축을 다르게 압축하는 것"**이어야 함을 입증했습니다.
실용성: 14 배의 압축률을 달성하면서도 정확도 손실 없이 긴 문맥 추론을 가능하게 하여, 상용 하드웨어에서 고해상도 LLM 추론의 실용성을 크게 높였습니다.
확장성: 이 프레임워크는 다양한 압축 축을 결합하고 단일 보정된 라우터로 관리하는 새로운 압축 방법론의 템플릿을 제공합니다.

요약하자면, MoE-nD 는 LLM 의 긴 문맥 처리 능력을 유지하면서 메모리 효율성을 극대화하기 위해, 레이어별 민감도를 기반으로 Eviction 과 Quantization 을 동적으로 조율하는 지능형 압축 시스템입니다.