David helps Goliath: Inference-Time Collaboration Between Small Specialized and Large General Diffusion LMs

TL;DR

This paper scales diffusion-based language models from 0.4B to 13B parameters, introduces SSD-2 for inference-time ensemble collaboration with smaller models, and demonstrates improved response quality over autoregressive models.

Contribution

It presents a scalable training approach for diffusion LMs, introduces SSD-2 for inference-time ensemble, and shows enhanced performance through collaborative diffusion models.

Findings

Scaling diffusion LMs improves their capabilities.

Ensembling small models with a large diffusion LM enhances response quality.

Diffusion LMs outperform autoregressive models in collaborative settings.

Abstract

Diffusion-based language models are emerging as a promising alternative to autoregressive LMs: they approach the competence of autoregressive LMs while offering nuanced controllability at inference time. While autoregressive LMs have benefited immensely from scaling and instruction-based learning, existing studies of diffusion LMs have been conducted on a smaller scale. Starting with a recently proposed diffusion model SSD-LM, in this work we first explore methods to scale it from 0.4B to 13B parameters, proposing techniques to improve its training and inference efficiency, and to finetune the model to follow instructions. Armed with a more powerful, general purpose diffusion LM, we introduce the primary contribution of this work -- SSD-2 -- an approach to easily ensemble at inference time a large general-purpose diffusion LM with smaller, but specialized and contextualized diffusion LMs. We show that SSD-2 facilitates novel ensembles with 100x smaller models that can be customized and deployed by individual users. We find that compared to autoregressive models, the collaboration between diffusion LMs is more effective, leading to higher-quality model responses due to their ability to dynamically incorporate bi-directional contexts.