Scaling Inference-Efficient Language Models

TL;DR

This paper introduces inference-aware scaling laws for large language models, optimizing model architecture and training to improve inference efficiency without sacrificing accuracy, demonstrated by the Morph-1B model.

Contribution

It extends existing scaling laws to include inference costs, proposing a method to train inference-efficient models and releasing a new model with improved latency and maintained accuracy.

Findings

Model architecture significantly impacts inference latency.

Wider and shallower models can be more efficient while preserving accuracy.

The Morph-1B model achieves 1.8x faster inference with comparable performance.

Abstract

Scaling laws are powerful tools to predict the performance of large language models. However, current scaling laws fall short of accounting for inference costs. In this work, we first show that model architecture affects inference latency, where models of the same size can have up to 3.5x difference in latency. To tackle this challenge, we modify the Chinchilla scaling laws to co-optimize the model parameter count, the number of training tokens, and the model architecture. Due to the reason that models of similar training loss exhibit gaps in downstream evaluation, we also propose a novel method to train inference-efficient models based on the revised scaling laws. We perform extensive empirical studies to fit and evaluate our inference-aware scaling laws. We vary model parameters from 80M to 1B, training tokens from 1.6B to 30B, and model shapes, training 63 models. Guided by our inference-efficient scaling law and model selection method, we release the Morph-1B model, which improves inference latency by 1.8x while maintaining accuracy on downstream tasks compared to open-source models, pushing the Pareto frontier of accuracy-latency tradeoff. Notably, our experiments reveal that wider and shallower models can yield efficiency gains while preserving accuracy.