No Need to Talk: Asynchronous Mixture of Language Models

TL;DR

SMALLTALK LM presents an asynchronous training approach for mixture of language models, enabling efficient specialization and routing without high communication overhead, achieving lower perplexity and strong downstream performance.

Contribution

Introduces SMALLTALK LM, a novel asynchronous training method for language model mixtures with a lightweight routing scheme that improves efficiency and performance.

Findings

Achieves lower perplexity than dense models at similar FLOPs

Uses significantly less parameters during inference

Outperforms dense baselines on 75% of downstream tasks

Abstract

We introduce SMALLTALK LM, an innovative method for training a mixture of language models in an almost asynchronous manner. Each model of the mixture specializes in distinct parts of the data distribution, without the need for high-bandwidth communication between the nodes training each model. At inference, a lightweight router directs a given sequence to a single expert, according to a short prefix. This inference scheme naturally uses a fraction of the parameters from the overall mixture model. Unlike prior works on asynchronous LLM training, our routing method does not rely on full corpus clustering or access to metadata, making it more suitable for real-world applications. Our experiments on language modeling demonstrate that SMALLTALK LM achieves significantly lower perplexity than dense model baselines for the same total training FLOPs and an almost identical inference cost. Finally, in our downstream evaluations we outperform the dense baseline on 75% of the tasks.