Dense2MoE: Pushing the Pareto Frontier of On-Device LLMs via Unified Pruning and Upcycling

기술 요약: Dense2MoE

문제 제기
자율주행, 엣지 AI 등 리소스가 제한된 엣지 장치에 대규모 언어 모델 (LLM) 을 배포하는 것은 '추론 메모리 벽'으로 인해 심각하게 저해됩니다. 하드웨어 루프라인 (Roofline) 이론에 따르면, 밀집형 LLM 은 자기회귀적 디코딩 중 메모리 병목 현상에 직면하며, 이는 단순히 이론적 FLOPs 을 줄이는 것만으로는 지연 시간을 개선하기에 불충분함을 의미합니다. 기존 솔루션들은 근본적인 트레이드오프에 직면해 있습니다:

1. 구조적 가지치기 (Structured Pruning): UIDL 과 같은 방법은 중복된 계층을 제거하여 지연 시간을 단축하지만, LLM 의 주요 지식 저장소 역할을 하는 멀티 레이어 퍼셉트론 (MLP) 을 제거함으로써 필연적으로 모델 정확도를 희생합니다.
2. MoE 업사이클링 (MoE Upcycling): 밀집형 모델을 전문가 혼합 (Mixture-of-Experts, MoE) 아키텍처로 전환하면 선형적으로 증가하지 않는 동적 연산으로 용량을 늘릴 수 있습니다. 그러나 표준 업사이클링은 동질적인 전문가 초기화 (원래 MLP 의 복제) 로 인해 '표현 붕괴 (representation collapse)'를 겪는 경우가 많으며, 전문화를 유도하기 위해 대규모 지속적 사전 학습이 필요합니다. 또한 최근의 희소 재사용 방법들은 주로 계층 내 중복성에 초점을 맞추어, 중복된 어텐션 및 KV 캐시 로딩 오버헤드를 유발하는 거시적 계층 간 구조적 중복을 해결하지 못합니다.

방법론: Dense2MoE
저자들은 Layer-Fusion Upcycling (LF-UC) 이라는 새로운 기법을 통해 구조적 계층 가지치기와 MoE 업사이클링을 통합한 통합 프레임워크인 Dense2MoE를 제안합니다. 이 방법은 세 단계로 운영됩니다:

1. 계층 간 유사성 분석: 이 프레임워크는 디코더 출력과 MLP 입력 모두에서 인접 계층 간의 높은 코사인 유사성을 분석하여 구조적 중복성을 식별합니다. 특히, 중복된 MLP 를 재사용할 때 활성화 이상이나 스케일 변화를 유발하지 않도록 엄격한 특징 노름 (feature norm) 매칭 제약을 부과합니다.
2. 레이어 퓨전 업사이클링 (LF-UC): Dense2MoE 는 중복 계층을 완전히 폐기하는 대신, 메모리 대역폭을 많이 소모하는 연산을 제거하기 위해 이러한 계층의 어텐션 모듈 (MHA) 만 물리적으로 가지치기합니다. 가지치기된 계층에서 사전 학습된 MLP 들은 보존되어 유지된 베이스 계층 내에서 이종 전문가 (heterogeneous experts) 로 재사용됩니다.
   • 이는 '베이스 전문가'가 복제된 원래 MLP 이고 '보조 전문가'가 가지치기된 계층에서 보존된 MLP 인 전문가 풀을 생성합니다.
   • 동적 토큰 수준 라우터는 이러한 전문가들에게 토큰을 선택적으로 배분하거나 잔여 경로를 통해 우회시킵니다.
   • 이 설계는 동질적 업사이클링에서 흔히 발생하는 '승자 독식 (Winner-Takes-All)' 라우팅 붕괴를 완화하기 위해 계층 간 구조적 이질성을 활용합니다.
3. 지속적 사전 학습: 수정된 아키텍처는 라우팅 메커니즘을 최적화하고 퓨전된 전문가들을 정렬하기 위해 경량화된 지속적 사전 학습 단계 (225B 토큰) 를 거칩니다. 이 예산은 처음부터 MoE 모델을 학습하는 데 필요한 연산량의 1% 미만을 소모합니다.

최적의 가지치기 깊이를 결정하기 위해 저자들은 MnasNet 에서 영감을 받은 하드웨어 인식 보상 지표를 사용하여 가지치기 제약을 특정 대상 하드웨어 (예: NVIDIA Jetson Thor-U) 에 대한 지연 시간 대 정확도의 파레토 프론티어에 매핑합니다.

주요 기여

• 통합 프레임워크: Dense2MoE 는 구조적 계층 가지치기와 MoE 업사이클링을 원활하게 통합하여 추론 효율성과 모델 정확도 간의 트레이드오프를 해결합니다.
• 새로운 아키텍처 (LF-UC): 어텐션 병목 현상을 가지치기하면서 MLP 를 이종 전문가로 재활용함으로써, 표현 능력을 파괴하지 않고 메모리 대역폭 소비를 줄입니다.
• 최첨단 파레토 최적성: 이 프레임워크는 온디바이스 추론 지연 시간 대 모델 정확도에 대한 파레토 프론티어를 진전시켜, 밀집형 베이스라인, 최첨단 압축 방법 및 표준 업사이클링 접근법보다 우수한 성능을 보입니다.
• 강건한 확장성: 이 방법은 다양한 모델 규모 (0.5B~7B) 와 서로 다른 아키텍처 계열 (Qwen2.5 및 LLaMA2) 에서 검증되었습니다.

실험 결과
Qwen2.5 및 LLaMA2 모델을 사용하여 MMLU, ARC, GSM8K, HumanEval 을 포함한 9 가지 다양한 벤치마크에서 평가한 결과:

• 성능 대 지연 시간: 최적의 Dense2MoE 구성 (24 계층 Qwen2.5-0.5B 시드에서 19 계층 유지) 은 Thor-U SoC 에서 271.0ms의 물리적 추론 지연 시간을 달성했습니다 (시드 모델의 307.5ms 에서 감소). 동시에 평균 벤치마크 점수 47.31을 기록하여 밀집형 베이스라인 (37.53) 을 크게 능가했습니다.
• 베이스라인과의 비교: 추론 및 코딩 작업에서 심각한 용량 붕괴를 겪었던 전통적인 가지치기 방법 (UIDL, LLM-Streamline 등) 과 희소 재사용 접근법 (ToMoE, LLaMA-MoE 등) 과는 달리, Dense2MoE 는 견고한 표현 유지 능력을 유지했습니다. 평균 점수에서 밀집형 베이스라인보다 9.78 퍼센트 포인트 이상 높은 성능을 보이면서 0.42B 파라미터만 활성화했습니다.
• 확장성: 더 큰 모델 (Qwen2.5-1.5B 및 LLaMA2-7B) 에 적용했을 때, Dense2MoE 는 동일한 225B 토큰 데이터셋으로 파인튜닝 후 각각 3.04 및 3.54 퍼센트 포인트씩 밀집형 대응 모델보다 일관되게 우수한 성능을 보였습니다.

의의
이 논문은 Dense2MoE 가 오픈소스 밀집형 LLM 을 고효율 온디바이스 아키텍처로 재사용하기 위한 실용적이고 산업 표준 수준의 레시피를 확립한다고 주장합니다. 계층 간 구조적 중복성을 체계적으로 활용함으로써 추론 지연 시간과 모델 정확도의 공동 최적화를 가능하게 합니다. 이 연구는 MoE 모델을 처음부터 학습하는 prohibitive 한 비용이나 공격적인 가지치기로 인한 정확도 저하 없이도 추론 메모리 벽을 깨고 복잡한 추론 능력을 보존할 수 있음을 보여줍니다. 이는 지연 시간에 민감한 엣지 배포에 최적화된 유연한 모델 연속체를 제공합니다.