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

Zechun Liu; Changsheng Zhao; Forrest Iandola; Chen Lai; Yuandong Tian; Igor Fedorov; Yunyang Xiong; Ernie Chang; Yangyang Shi; Raghuraman Krishnamoorthi; Liangzhen Lai; Vikas Chandra

Source-linked AI summary

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

Zechun Liu, Changsheng Zhao, Forrest Iandola, Chen Lai, Yuandong Tian, Igor Fedorov, Yunyang Xiong, Ernie Chang, Yangyang Shi, Raghuraman Krishnamoorthi, Liangzhen Lai, Vikas Chandra

arXiv:2402.14905v2Feb 22, 2024cs.LGcs.AIcs.CL

TL;DR

The paper addresses the challenge of deploying capable LLMs under mobile memory, computation, and latency constraints. It develops architecturally optimized sub-billion models using deep, thin designs and weight sharing, achieving stronger benchmark performance and competitive on-device task results than prior small models.

Problem
Mobile deployment requires LLMs that fit constrained device memory and reduce the computation, energy, and latency burdens associated with cloud-scale models.
Method
The paper builds MobileLLM with deep, thin architectures, embedding sharing, grouped-query attention, and immediate block-wise weight sharing for sub-billion models.
Results
MobileLLM improves zero-shot accuracy over previous state-of-the-art 125M/350M models by 2.7%/4.3% and significantly outperforms equivalently sized models on chat and API-calling tasks.
Takeaways & Limitations
MobileLLM models substantially advance zero-shot commonsense reasoning, question answering, and reading comprehension, while fine-tuned variants support chat and API calling on-device use cases.
Takeaways & Limitations
Knowledge distillation increases training time by 2.6−3.2× and achieves comparable or inferior accuracy to label-based training.

Abstract

from arXiv · show

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.

1. Introduction

The paper motivates compact LLMs for on-device execution by highlighting cloud-scale costs, mobile memory constraints, and efficiency advantages of sub-billion models. It introduces MobileLLM and block-wise weight sharing as approaches that improve small-model performance while preserving deployment practicality.

Motivation: Cloud-scale LLM deployment could impose substantial computation, energy, and carbon costs, motivating smaller models.The paper contrasts GPT-4-scale deployment with the imperative to downsize LLMs.
Motivation: Mobile devices constrain practical LLM deployment because DRAM typically ranges from 6 to 12 GB, making 8-bit LLaMA-v2 7B prohibitively expensive.The paper identifies DRAM capacity as the relevant operational memory constraint for high-speed applications.
Motivation: A 350M 8-bit model consumes 0.035 J/token, while a 125M model reaches 50 tokens/s compared with 3∼6 tokens/s for an iPhone LLaMA 7B application.These figures illustrate the efficiency and decoding-speed advantages targeted by sub-billion models.
Contributions: Immediate block-wise weight sharing avoids weight movement by computing a shared block twice, adding only minimal latency overhead.The approach targets scenarios in which memory movement is the decoding bottleneck.
Contributions: MobileLLM outperforms previous state-of-the-art 125M/350M models by 2.7%/4.3% in a suite of zero-shot tasks.The result is reported as the model family’s headline zero-shot comparison.
Contributions: MobileLLM improves downstream chat and API-calling performance over equivalently sized models, with MobileLLM-350M comparable to LLaMA-v2 7B on API exact match.The paper also reports results for larger MobileLLM variants in the appendix.

2. Improving Sub-billion Scale LLM Design

MobileLLM designs sub-billion models by prioritizing architecture and weight utilization under tight storage constraints. Its baseline combines four architectural techniques, while layer sharing increases effective depth without increasing model size.

Depth and width: Deeper and thinner models generally outperform wider and shallower models at comparable sizes across reasoning, question answering, and reading comprehension.The study trained 19 similarly sized models spanning approximately 125M and 350M parameters.
Embedding sharing: Embedding sharing reduces parameters by 16M, approximately 11.8% of a 125M model, with a 0.2-point accuracy drop that can be recovered by adding layers.Increasing depth from 30 to 32 layers produced a 0.4-point accuracy gain while retaining 10M fewer parameters than the original 135M model.
Grouped-query attention: 16 query heads produced the best results, while reducing kv-heads from 16 to 4 preserved comparable 125M accuracy and reduced 350M accuracy by 0.2 points.The reduction also yielded almost 10% model-size reduction in the 350M setting.
MobileLLM baseline: MobileLLM combines SwiGLU FFNs, deep-and-thin architectures, embedding sharing, and grouped-query attention into a strong sub-billion baseline.These techniques target feed-forward design, depth allocation, parameter reuse, and key-value-head redundancy.
Layer sharing: Immediate block-wise layer sharing increases hidden-layer count without additional model storage and is designed for memory-constrained on-device decoding.The approach reuses adjacent-block weights and incurs only slight latency overhead.

3. Experiments

Experiments evaluate MobileLLM across reasoning, question answering, reading comprehension, chat, API calling, quantization, and mobile latency. The models consistently outperform comparable sub-billion systems, while layer sharing preserves accuracy and improves deployment efficiency.

Question answering and reading comprehension: MobileLLM-125M improves TQA accuracy by over 4.3 points, while MobileLLM-350M improves performance by approximately 10 points over other 350M models.The model family also reports significantly higher scores on reading comprehension tasks.
Chat: MobileLLM models significantly outperform previous state-of-the-art sub-billion models on AlpacaEval and MT-Bench, including models with 1 billion parameters.The evaluations use identical settings for MobileLLM and comparison checkpoints.
API calling: MobileLLM-350M achieves comparable intent and structural exact-match scores to LLaMA-v2 7B in API calling.Intent exact match measures correct API selection, while structural exact match measures correct function content.
Quantization: W8A8 post-training quantization reduces accuracy by less than 0.5 points and remains compatible with layer sharing.The quantization experiments use per-token min-max post-training quantization on 125M and 350M models.
Knowledge distillation: Knowledge distillation slows training by 2.6−3.2× and achieves comparable or inferior accuracy to label-based training.The experiments use LLaMA-v2 7B as the teacher for 125M and 350M models.
Latency: MobileLLM-LS incurs 2.2% higher loading and initialization time and 2.6% higher execution time than MobileLLM, versus 143% and 86% increases for a doubled-depth non-shared model.Measurements use FP16 models on an iPhone 13 with ExecuTorch and the Metal Performance Shaders backend.

4. Related Work

Prior work compresses models, designs compact architectures, optimizes attention, and explores weight sharing for efficient on-device inference. This paper positions its approach as complementary while emphasizing simpler intermediate-layer sharing with minimal latency increase.

Model Compression: Model compression for LLMs includes pruning, sparsity, and quantization, and the paper states its methodology is complementary to these techniques.The cited passage also indicates compatibility with quantization.
Small Model Design: Compact-model architecture research remains limited, and even the smallest TinyLLaMA exceeds 1 billion parameters.The passage frames this size as prohibitive for many on-device applications.
Neural Architecture Search: Transformer research commonly treats architecture as having limited accuracy impact when parameter counts match, unlike the paper’s focus on smaller models.The passage contrasts this consensus with the motivation for investigating architecture at sub-billion scale.
Weight Sharing: Prior weight sharing often targets input and output embeddings or uses specialized shared-layer designs.The paper instead highlights repeating transformer blocks as a simpler alternative.
Weight Sharing: Repeating transformer blocks provides improved accuracy with fixed model size and minimal latency increase.This is presented as the paper’s straightforward intermediate-layer weight-sharing contribution.
Efficient Attention and Implementation: Efficient-attention research includes low-rank and sparse attention, while other work targets hardware scheduling and weight movement.The cited passage identifies FlashAttention and FlexGen as examples of the latter direction.

5. Conclusion

The study optimizes sub-billion models for on-device applications through architecture and weight-sharing choices. MobileLLM improves several evaluation areas and is demonstrated in chat and API-calling use cases.

Scope: MobileLLM focuses on optimizing sub-billion language models for on-device applications.The conclusion frames this as the study’s central scope.
Design: Deeper architectures and embedding, grouped-query, and block-wise weight sharing improve weight utilization under storage constraints.These design choices are presented together as the study’s approach to constrained deployment.
Evaluation: MobileLLM advances zero-shot commonsense reasoning, question answering, and reading comprehension compared with previous state-of-the-art methods.The conclusion reports improvements across these task categories.
Applications: Fine-tuned MobileLLM models are evaluated on chat and API calling as prevalent on-device use cases.The conclusion highlights the models’ performance on these applications.

Impact Statement

The paper advocates sub-billion LLMs for on-device applications to address inference energy consumption. It presents the approach as promising for reducing deployment-related computational costs.

Impact: The paper advocates sub-billion large language models for on-device applications.This is the impact statement’s stated direction.
Impact: The stated aim is to mitigate energy consumption during LLM inference.The passage links the proposed deployment setting with this objective.
Impact: The proposed approach is described as promising for alleviating computational costs associated with LLM deployment.The statement does not quantify the reduction.

A. Scaling Up to Larger Model Architectures

The paper extends MobileLLM’s sub-billion design principles to larger model sizes and reports stronger zero-shot performance than similarly scaled prior models. Its experiments also support deeper architectures as beneficial across evaluated tasks.

A. Scaling Up to Larger Model Architectures: MobileLLM extends its SwiGLU, deeper-architecture, grouped-query-attention, and embedding-sharing principles to 600M, 1B, and 1.5B variants.The expansion targets applications spanning different memory constraints.
A. Scaling Up to Larger Model Architectures: 59.4 average accuracy points on zero-shot commonsense reasoning tasks is achieved by MobileLLM-1.5B.This result is reported in comparison with general-purpose pretrained models.
A. Scaling Up to Larger Model Architectures: 2.9 points higher than Qwen1.5-1.8B, MobileLLM-1.5B outperforms a model with more parameters.The passage identifies Qwen1.5-1.8B as the previous state-of-the-art comparison.
Evaluation Setup: Table 8 compares MobileLLM-600M, 1B, and 1.5B on zero-shot common-sense reasoning tasks, highlighting the highest and second-highest average scores by size category.The caption specifies the table’s comparison and highlighting scheme.
Architecture Specifications: Table 9 lists MobileLLM architecture specifications, including embedding and feed-forward hidden dimensions.The caption defines “Emb Dim” and “Hidden Dim.”
Depth Versus Width: Models with fewer than 10 layers perform poorly, while increasing depth beyond 10–20 layers continues to benefit reasoning, question answering, and reading comprehension.The passage reports consistent gains from deeper architectures across the three task groups.

E. Layer-Sharing Number Ablation

The layer-sharing ablation finds that doubling depth by sharing weights across every two transformer blocks improves accuracy, while further repetition yields diminishing returns. The experiments also examine quantization compatibility for MobileLLM and MobileLLM-LS.

Experimental setup: The ablation compares an 8-layer 125M model with embedding dimension 896 and a 15-layer 350M model with embedding dimension 1280.These models were used to investigate the optimal number of layer repetitions.
Layer repetition: 0.4–0.6% accuracy improvement follows doubling the layer number with every two transformer blocks sharing weights.Further tripling or quadrupling layer repetitions reduces the incremental benefit.
Layer repetition: The adopted configuration shares weights across every two blocks, effectively doubling the total number of layers.This choice follows the observed accuracy gains and diminishing returns from additional repetition.
Quantization compatibility: W8A8 post-training quantization causes an accuracy drop of within 0.5% and remains compatible with MobileLLM and MobileLLM-LS.The experiments quantize both weights and activations to 8-bit using per-token min-max quantization.

G. Knowledge Distillation

Knowledge distillation uses LLaMA-v2 7B as a teacher for 125M or 350M students, but its accuracy is comparable to or worse than hard-label training while substantially increasing training time.

Distillation setup: Knowledge distillation computes cross-entropy between LLaMA-v2 7B teacher logits and the 125M or 350M student network.The teacher provides soft labels for small-model pre-training.
Notation: In the distillation notation, T and S denote the teacher and student networks, while c is the vocabulary size.The loss is defined over samples in the batch and vocabulary classes.
Training cost: 2.6–3.2× slower training results from adding knowledge-distillation loss compared with training from scratch using labels.All models were trained on 32 A100 80G GPUs with batch size 32 for 120k iterations.
Results: Knowledge distillation is comparable to or lower in performance than using next-token labels alone.The experiments therefore use labels rather than KD in the reported models.

H. Datasets and Benchmarks

The paper evaluates MobileLLM across reasoning, question answering, reading comprehension, chat, and API-calling settings, while also testing 8-bit quantization and knowledge distillation.

Evaluation suite: MobileLLM is evaluated on seven zero-shot common-sense tasks, TriviaQA, RACE, MT-Bench, AlpacaEval, and an API-calling dataset.The evaluation spans reasoning, question answering, reading comprehension, chat, and API-calling use cases.
Quantization: 8-bit weight-and-activation post-training quantization achieves an accuracy gap within 0.5% of the full-precision BF16 counterpart.This result is reported for zero-shot common-sense reasoning tasks.
Knowledge distillation: Knowledge distillation with LLaMA-v2 7B soft labels slightly degrades performance relative to using hard labels alone.The comparison is presented as an ablation of the training approach.

H.1. Zero-shot Common Sense Reasoning tasks

The zero-shot evaluation covers diverse commonsense, question-answering, reading-comprehension, and chat benchmarks, alongside generated API-calling examples. The benchmarks target physical and social reasoning, factual answering, multi-turn dialogue, and tool-oriented interactions.

Commonsense reasoning: BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, ARC, and OBQA cover yes/no comprehension, physical and social commonsense, inference, pronoun resolution, science, and open-book answering.These are among the zero-shot common-sense reasoning benchmarks used in the evaluation.
Question answering: TriviaQA is a closed-book question-answering benchmark with over 650,000 question-answer evidence triples and an average of six evidence documents per question.The triples derive from 95,000 question-answer pairs authored by Trivia enthusiasts.
Reading comprehension: RACE contains more than 28,000 passages and nearly 100,000 questions from English examinations for middle- and high-school students in China.It is used to assess reading comprehension.
Chat evaluation: MT-Bench uses 160 multi-turn questions across eight domains, with GPT-4 scoring two responses per question from 1 to 10.AlpacaEval instead uses 805 single-turn questions and GPT-4-rated pairwise win rate against text-davinci-001.
API calling: The API-calling dataset is generated from simulated human-agent conversations in which a language model explicitly identifies tasks for designated APIs.Examples include news, sports, messaging, finance, and local-information requests.

J. Chat Generations from MobileLLM

The section presents MobileLLM outputs for gift wrapping, beginner exercises, and mindful walking, alongside LLaMA-v2 7B responses for comparison. The examples show both coherent practical guidance and factual or procedural weaknesses in some small-model generations.

Mindful walking: MobileLLM-125M describes mindful walking as offering posture, flexibility, sleep, stress, and mental-health benefits.
Mindful walking: MobileLLM-350M connects mindful walking with reduced muscle and joint tension, cardiovascular fitness, core strength, flexibility, sleep, stress reduction, and well-being.
Gift wrapping: MobileLLM-350M provides multi-step gift-wrapping guidance, including choosing decorations, securing them, and adding personal touches.
Comparison examples: The LLaMA-v2 7B examples provide longer, more elaborated responses for gift wrapping, beginner exercises, and mindful walking.
Beginner exercises: MobileLLM-125M lists beginner exercises but includes a factual error by describing dumbbell rows as building leg and core strength.
Beginner exercises: MobileLLM-350M gives structured beginner-exercise advice and recommends starting with exercises requiring limited strength or endurance.