Meta's MobileLLM trades width for depth and gains 2.7 and 4.3 points
Meta Reality Labs researchers published MobileLLM on February 22, 2024, a family of language models built to run on a phone rather than in a data centre. The paper argues that below a billion parameters the shape of the network matters more than the parameter count or the amount of training data. Its 125M and 350M models beat the previous best models at those sizes by 2.7 and 4.3 points on eight zero-shot common sense reasoning tasks.
The central finding is that depth beats width. The authors trained 19 models of roughly equal size but different proportions, 9 around 125M parameters and 10 around 350M, and report that 30-layer and 42-layer designs scored higher than the 12-layer shape used by OPT, GPT-Neo and Pythia at the same size. Sharing one weight matrix between the input embedding and the output layer removes 16M parameters, about 11.8% of a 125M model, for a 0.2 point accuracy drop the authors then recover by spending the freed budget on two more layers. Grouped query attention, where several query heads read one shared set of key and value heads, cuts 16 key-value heads to 4 at close to the same accuracy.
The paper sets out what a phone can actually spare. The authors put DRAM on current handsets at 6 to 12 GB, shared with the operating system and every other app, and argue that an app should not claim more than 10% of it. For energy they use a rule of thumb of 0.1 J per token per billion parameters, which puts a 7B model at 0.7 J/token, drains 0.2% of an iPhone battery every 64 tokens and allows under 2 hours of conversation at 10 tok/s. A 350M model with 8-bit weights costs 0.035 J/token by the same arithmetic, which the authors say covers a full day of use.
Meta also profiled the models on hardware rather than estimating. Running the 125M model in FP16 through ExecuTorch on an iPhone 13 with iOS 17.2.1 and the Metal Performance Shaders backend, the authors measured 39.2 ms to load, 1361.7 ms to initialise and 15.6 ms to execute, the last averaged over 50 runs. MobileLLM-LS, the layer-sharing variant, computes each transformer block twice in a row so that 30 blocks of weights behave like 60 layers at the same file size, a pattern the team picked over the alternatives because the repeated weights stay in the roughly 20 MB of on-chip cache instead of being fetched from DRAM again. That costs 2.2% more loading and initialisation time and 2.6% more execution time, against 143% and 86% for a model that really has 60 layers, and adds 0.7 and 0.8 points of accuracy at 125M and 350M.
After fine-tuning, MobileLLM-350M scores 3.28 on the MT-Bench chat benchmark against 1.37 for OPT-350M and 2.24 for the 1.3B OPT, and the layer-shared 350M wins 48.2% of AlpacaEval comparisons against text-davinci-001, a model that wins 50% against itself. For API calling the team built a synthetic set of 5,000 training and 2,500 test conversations that turn requests such as setting an alarm into a function call, and MobileLLM-350M scored 65.3 on intent exact match against 62.8 for Llama 2 7B, while trailing the larger model on the structure of the call and on the wording of the reply. Quantising weights and activations to 8 bits costs under 0.5 points. The work was published at ICML 2024 and the code is in Meta’s repository.