ImageNet classification with deep convolutional neural networks
Alex Krizhevsky, Ilya Sutskever et al.
0
Citations
0
Influential Citations
—
Venue
2024
Year
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.
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.

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.

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.

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.

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.


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.


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.


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

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.

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.
Alex Krizhevsky, Ilya Sutskever et al.
Ashish Vaswani, Noam Shazeer et al.
Douglas M. Bates, Martin Mächler et al.
Diederik P. Kingma, Jimmy Ba