Preprint
Large Language Models

MobileLLM

February 1, 2024

0

Citations

0

Influential Citations

Venue

2024

Year

Abstract

Leverages various architectures and attention mechanisms to achieve a strong baseline network, which is then improved upon by introducing an immediate block-wise weight-sharing approach, resulting in a further accuracy boost.

Analysis

Why This Paper Matters

As AI moves onto smartphones, wearables, and edge devices, the demand for small yet capable language models has surged. MobileLLM tackles the core challenge: how to design sub-billion parameter transformers that deliver strong performance without the computational luxury of cloud-scale models. The paper systematically overturns the assumption that architecture matters little for small models, showing that careful design choices can yield significant gains. This is crucial for practitioners building on-device assistants, real-time translation, or privacy-preserving AI.

How It Works

Figure 1

MobileLLM starts by establishing a strong baseline through three key architectural decisions:

Deep and Thin Architecture: Contrary to the common belief that parameter count alone determines performance, the authors find that increasing depth (more layers) improves accuracy more than increasing width (larger hidden dimensions) for small models. This is illustrated in Figure 3.

Figure 3

Embedding Sharing: In small models, the embedding layer can consume over 20% of parameters (e.g., 512-dim embedding with 32k vocab in a 125M model). By reusing input embeddings as the output projection weights, the model saves parameters that are then reinvested into additional layers. Figure 4 shows this design.

Figure 4

Grouped-Query Attention (GQA): Experiments show that using 16 query heads with only 4 key-value heads yields comparable accuracy while reducing model size by 10%. This is especially effective for small models, as shown in Figure 5.

Figure 5

The baseline also adopts SwiGLU activation in the feed-forward network, which proves beneficial even at small scales.

Immediate Block-Wise Weight-Sharing: To further boost performance without increasing model size, the authors propose sharing weights across consecutive transformer blocks. Among three strategies (immediate block-wise, repeat-all-over, reverse), immediate block-wise sharing performs best. The shared weights stay in cache, enabling fast computation. This yields the MobileLLM-LS variant, which adds 0.7-0.8% accuracy over the baseline.

(a) Baseline model without layer sharing; (b) Immediate block-wise sharing; (c) Repeat-all-over sharing; (d) Reverse sharing. A transformer block contains the multihead self-attention (MHSA) and the feed-forward network (FFN).

Figure 7

Results

MobileLLM-125M and 350M outperform prior SOTA models (OPT, GPT-Neo, Pythia) by 2.7% and 4.3% on zero-shot commonsense reasoning. The LS variants add further gains. On the TQA benchmark, MobileLLM-125M improves by 4.3 points over previous SOTA. On RACE reading comprehension, the entire MobileLLM family scores significantly higher.

Zero-shot performance on Common Sense Reasoning tasks.

Performance on Trivia QA and RACE datasets for question answering and reading comprehension tasks.

In chat evaluation (AlpacaEval), MobileLLM-LS-350M achieves a 48.2% win rate against GPT-3 (text-davinci-001), nearly matching GPT-3's self-win rate of 50%. For API calling, MobileLLM-350M matches LLaMAv2 7B in intent and structure exact match, despite being 20x smaller.

Benchmark results on AlpacaEval (Evaluator: GPT4; Reference model: text-davinci-001) and MT-Bench.

API calling evaluation score.

The models also tolerate 8-bit weight and activation quantization with less than 0.5 point accuracy drop, making them deployment-friendly.

Comparison between BFloat16 model and 8-bit weight 8-bit activation post-training quantized model.

Scaling up to 600M, 1B, and 1.5B, MobileLLM continues to outperform comparable models. MobileLLM-1.5B achieves 59.4% average accuracy on commonsense reasoning, beating Qwen1.5-1.8B by 2.9 points.

Zero-shot performance on Common Sense Reasoning tasks for MobileLLM-600M, 1B and 1.5B.

Significance

MobileLLM provides a practical recipe for building high-performance small language models. Its design principles—deep and thin, embedding sharing, GQA, and block-wise weight-sharing—are immediately actionable for practitioners targeting on-device deployment. The paper also demonstrates that small models can match or exceed larger ones on specific tasks like API calling, opening doors for specialized, efficient AI agents. The code is open-source, enabling rapid adoption and further innovation in the edge AI space.