[Paper Review] MobiLlama: Towards Accurate and Lightweight Fully Transparent GPT

Omkar Thawakar, Ashmal Vayani+7 more|arXiv (Cornell University)|Feb 26, 2024

Advanced Data Storage Technologies4 citations

TL;DR

MobiLlama introduces a fully transparent, 0.5B-parameter Small Language Model (SLM) that achieves state-of-the-art performance among sub-1B models by employing a shared feed-forward network (FFN) across all transformer blocks, reducing redundancy and improving efficiency. It outperforms existing SLMs by 2.4% on average across nine benchmarks while enabling on-device inference with low memory, compute, and energy consumption.

ABSTRACT

"Bigger the better" has been the predominant trend in recent Large Language Models (LLMs) development. However, LLMs do not suit well for scenarios that require on-device processing, energy efficiency, low memory footprint, and response efficiency. These requisites are crucial for privacy, security, and sustainable deployment. This paper explores the "less is more" paradigm by addressing the challenge of designing accurate yet efficient Small Language Models (SLMs) for resource constrained devices. Our primary contribution is the introduction of an accurate and fully transparent open-source 0.5 billion (0.5B) parameter SLM, named MobiLlama, catering to the specific needs of resource-constrained computing with an emphasis on enhanced performance with reduced resource demands. MobiLlama is a SLM design that initiates from a larger model and applies a careful parameter sharing scheme to reduce both the pre-training and the deployment cost. Our work strives to not only bridge the gap in open-source SLMs but also ensures full transparency, where complete training data pipeline, training code, model weights, and over 300 checkpoints along with evaluation codes is available at : https://github.com/mbzuai-oryx/MobiLlama.

Motivation & Objective

• To address the lack of fully transparent, open-source SLMs that are both accurate and efficient for resource-constrained environments.
• To overcome the performance limitations of conventional SLMs that scale down large models by reducing hidden dimensions or number of layers.
• To design a lightweight yet powerful SLM framework that maintains high performance through architectural innovation, not just parameter count.
• To democratize access to SLMs by releasing the complete training pipeline, code, model weights, and over 300 checkpoints.
• To enable on-device processing for improved privacy, security, and energy efficiency in real-world applications.

Proposed method

• Propose a shared feed-forward network (FFN) architecture where a single FFN is shared across all transformer blocks, reducing parameter redundancy.
• Train the SLM from scratch using a 1.2T-token dataset, with careful initialization and optimization to maintain performance at low parameter count.
• Leverage a parameter sharing scheme that reduces both pre-training and deployment costs while preserving model capacity.
• Fine-tune the language model end-to-end with a vision encoder (CLIP) to create a multimodal version, MobiLlama-V, for visual reasoning.
• Release the full training pipeline, including data, code, model weights, and 300+ intermediate checkpoints, to ensure full transparency and reproducibility.
• Evaluate efficiency across diverse hardware (GPU, CPU, mobile) using metrics like tokens/sec, memory usage, and battery consumption per 1000 tokens.

Experimental results

Research questions

• RQ1 Can a shared FFN design across all transformer blocks improve the performance of small language models without increasing parameter count or computational cost?
• RQ2 How does a fully transparent SLM with complete training and inference code release impact reproducibility and community adoption in the SLM space?
• RQ3 To what extent can a 0.5B-parameter SLM match or exceed the performance of larger SLMs and even some 7B LLMs in downstream benchmarks?
• RQ4 How efficient is the MobiLlama model in real-world low-resource environments, including mobile and edge devices?
• RQ5 Can a shared FFN-based SLM be effectively extended to multimodal tasks with strong visual reasoning capabilities?

Key findings

• MobiLlama 0.5B outperforms existing 0.5B SLMs by 2.4% on average across nine downstream benchmarks, demonstrating superior accuracy.
• The model achieves an average inference speed of 125 tokens/sec on an RTX-2080Ti GPU, with only 1.2 GB memory usage and 12.5 mAh/1000 tokens battery consumption on a smartphone.
• On a laptop with an i7 CPU, MobiLlama 0.5B achieves 45 tokens/sec with 1.1 GB memory and 15.3 mAh/1000 tokens battery usage.
• The multimodal MobiLlama-V 0.8B model achieves strong performance on vision-language benchmarks, including 62.1% on MME and 58.3% on GQA, showing strong visual reasoning capabilities.
• The large-base MobiLlama 1.2B model achieves an average score of 49.06 across nine benchmarks, outperforming other fully transparent SLMs pre-trained on larger datasets.
• The model’s efficiency and transparency enable on-device deployment with minimal energy use, supporting privacy-preserving, low-latency applications.