MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases

This paper addresses the growing need for efficient large language models (LLMs) on mobile devices, driven by increasing cloud costs and latency concerns. We focus on designing top-quality LLMs with fewer than a billion parameters, a practical choice for mobile deployment. Contrary to prevailing belief emphasizing the pivotal role of data and parameter quantity in determining model quality, our investigation underscores the significance of model architecture for sub-billion scale LLMs. Leveraging deep and thin architectures, coupled with embedding sharing and grouped-query attention mechanisms, we establish a strong baseline network denoted as MobileLLM, which attains a remarkable 2.7%/4.3% accuracy boost over preceding 125M/350M state-of-the-art models. Additionally, we propose an immediate block-wise weight-sharing approach with no increase in model size and only marginal latency overhead. The resultant models, denoted as MobileLLM-LS, demonstrate a further accuracy enhancement of 0.7%/0.8% than MobileLLM 125M/350M. Moreover, MobileLLM model family shows significant improvements compared to previous sub-billion models on chat benchmarks, and demonstrates close correctness to LLaMA-v2 7B in API calling tasks, highlighting the capability of small models for common on-device use cases.

한 줄 요약

MobileLLM은 10억 미만 파라미터 모델에서 모델 아키텍처 최적화를 통해 모바일 기기에서의 성능을 획기적으로 향상시킨다.

핵심 기여도

125M/350M 모델 대비 2.7%/4.3% 정확도 향상.
즉각적 블록 간 가중치 공유를 제안, 모델 크기 증가 없이 0.7%/0.8% 추가 성능 향상.
30~42층의 깊은 구조가 12층 모델보다 125M 규모에서 훨씬 우수한 성능을 보임.
API 호출 작업에서 LLaMA-v2 7B 모델과 유사한 정확도 달성.

핵심 아이디어

기존 연구는 파라미터 수와 데이터 양이 모델 성능을 결정한다고 가정하지만, 본 연구는 10억 미만 파라미터 모델에서 **모델 아키텍처**가 훨씬 더 중요한 역할을 한다는 점을 밝혔다. 특히, **깊고 얇은**(deep and thin) 구조가 추상적 개념을 더 잘 포착하여 성능 향상에 기여한다는 것을 실험적으로 입증했다. 이는 기존 125M 모델이 12층에 머물렀던 것에 비해, MobileLLM은 30~42층의 깊은 구조를 채택함으로써 ARC-easy, PIQA, HellaSwag 등 다양한 zero-shot 추론 작업에서 우수한 성능을 보였다. 또한, **임베딩 공유**(embedding sharing)와 **그룹 쿼리 어텐션**(grouped-query attention)을 통해 가중치 활용률을 극대화했으며, **즉각적 블록 간 가중치 공유**(immediate block-wise weight sharing)를 통해 메모리 이동 비용을 줄이며 성능을 추가적으로 향상시켰다.

기술적 접근법

**모델 아키텍처**: 125M/350M 규모의 모델에서 30~42층의 깊은 구조를 채택.
**임베딩 공유**: 입력과 출력 임베딩을 공유하여 파라미터 수를 줄임.
**그룹 쿼리 어텐션**: 어텐션 헤드를 그룹화하여 계산 효율성 향상.
**즉각적 블록 간 가중치 공유**: 인접 블록 간 가중치를 공유하여 메모리 이동 비용 최소화.
**하드웨어 제약 고려**: 모바일 기기의 DRAM(6~12GB)과 배터리 소모(0.035 J/token)를 고려한 모델 크기 및 에너지 효율 최적화.

주요 결과

**zero-shot 추론 작업**: 125M/350M 모델 대비 2.7%/4.3% 정확도 향상.
**Chat 및 API 호출 작업**: 기존 125M/350M 모델 대비 0.7%/0.8% 추가 향상.
**LLaMA-v2 7B 대비 성능**: API 호출 작업에서 유사한 정확도 달성.
**속도**: 125M 모델은 50 tokens/s 처리, LLaMA 7B 기반 모델 대비 8~16배 빠름.

의의 및 한계

MobileLLM은 모바일 기기에서의 LLM 실행을 가능하게 하며, 에너지 소모와 배터리 소모를 대폭 줄이는 데 기여한다. 특히, 350M 모델은 iPhone 한 번 충전으로 하루 동안 대화형 작업을 지속할 수 있다. 그러나 모델 크기가 10억 미만이므로, 복잡한 추론 작업이나 대규모 데이터셋에서의 학습에는 한계가 있을 수 있다. 또한, 모델의 깊이가 증가함에 따라 학습 시간과 메모리 요구량이 증가할 수 있으며, 이는 모바일 환경에서의 추가 최적화 필요성을 시사한다.

실용적 활용

MobileLLM은 스마트폰, 태블릿, IoT 기기 등에서 실시간 대화, 추천 시스템, API 호출 등 다양한 on-device 작업에 적용 가능하다. 특히, 배터리 효율과 처리 속도가 중요한 모바일 앱 개발 및 엣지 컴퓨팅 분야에서 유용하게 활용될 수 있다.