D avid helps Goliath: Inference-Time Collaboration Between Small Specialized and Large General Diffusion LM s (2024)

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.

## introduction following the footsteps of diffusion-based generative models for continuously valued data such as images, audio, and video @xcite @xcite , recent works have attempted to replicate these successes on discrete text data @xcite @xcite @xcite . several studies have shown that diffusion-based language models (lms) perform competitively to their autoregressive counterparts, and even surpass 1 https://github.com/xhan77/ssd-2 . figure 1 : inference-time collaboration between a large general model and a small user model that incorporates user-specified knowledge. the collaboration between autoregressive models performs decoding tokenby-token, while the collaboration between diffusion models refines a block of generated tokens iteratively with bi-directional contexts ( §3). meanwhile, autoregressive language models @xcite have emerged as general-purpose solutions capable of holding conversations with humans and solving tasks by following instructions @xcite @xcite @xcite . their abilities are primarily due to two factors: scaling the model parameters as well as pretraining datasets, and instruction finetuning with carefully curated datasets @xcite . however, as the models become increasingly immense and proprietary, it is difficult for individual users to customize the system with their own data (e.g., specialized knowledge) due to cost or privacy reasons ( §3). a primary contribution of this work is to illustrate a novel setup of inference-time collabo-ration between lms and show a unique advantage of diffusion lms compared to autoregressive models in this scenario. with diffusion models' iterative generation design over a span of bi-directional contexts, multiple diffusion lms with different capabilities can be easily ensembled at the sequence level at test time, leveraging advantages of each lm in the ensemble. as a preliminary to our experiments, we first present an exploratory study to scale and incorporate instruction-following and conversational capabilities in diffusion-based lms. we introduce ssd-2, an improved version of recently introduced simplex-based diffusion lm ssd-lm @xcite proposing several modifications to its training and inference procedures. we incorporate these improvements in scaling ssd-2 to 13b parameters, up from 0.4b in ssd-lm. we show that similarly to autoregressive lms, by finetuning with curated instruction datasets, ssd-2 is well-suited to follow chat-style instructions. we then present our main case study highlighting the setup of inference-time collaboration: we augment a general-purpose large ssd-2 model with 13b parameters with a 100x smaller, useraccessible model. this setup allows incorporating user-provided knowledge into the generation process without directly inputting it into the large model (which can be undesirable due to cost or privacy reasons, more details in §3). we show that ssd-2's instruction finetuned model is substantially more effective at this collaboration than the autoregressive baselines, leveraging bi-directional contexts in the ensemble. ## background semi-autoregressive simplex-based diffusion lm (ssd-lm) is trained to generate text in blocks of tokens by performing diffusion in the simplex space of the model vocabulary @xcite . for text continuation tasks, it has shown competitive performance against autoregressive models (e.g., gpt-2; @xcite when trained with a similar number of model parameters and pretraining data. furthermore, it naturally enables post-hoc control in generated text using off-the-shelf classifiers, outperforming prior approaches to controlling autoregressive models. below we briefly overview the training and decoding algorithm of ssd-lm. ## ssd-2 in the age of llms, individual users of nlp models may often face a dilemma when they wish to customize the system with their own data. on one hand, it is difficult for user-owned devices to fit very large models and smaller models are usually not powerful enough. on the other hand, uploading their data to a commercial host of large models for finetuning or long in-context learning is expensive and also may not be desirable due to privacy risks. we aim to address this dilemma in this work by proposing a collaborative inference-time algorithm between two diffusion models: a large general-purpose model (such as ones only accessible through an api) and a small model which a user can customize ( §3.2). we first present ssd-2 building on top of ssd-lm with several modifications to improve its training and decoding efficiency ( §3.1). we train ssd-2 with a larger pretraining corpus and more parameters (ranging from 0.1b to 13b) than ssd-lm and fine-tune it to follow instructions ( §4). next, we present how different versions of ssd-2 (general-purpose large models and user-enhanced small models) can be effectively interpolated at inference time, outperforming their autoregressive counterparts ( §5). ## experimental setup pretraining existing work on diffusion lms is limited to modest model sizes below the order of 1b parameters @xcite @xcite . for example, ssd-lm has the same size as roberta-large @xcite with 0.4b parameters. it is unclear whether diffusions lms have the ability to scale like autoregressive lms. 8 to answer this question, we pretrain three versions of ssd-2 with 0.1b, 2.7b, and 13b parameters, on a subset of a large corpus c4 @xcite . instead of pretraining from scratch, we initialize these models using publicly available opt models @xcite . 9 we consider a maximum sequence length of 500 (up from 200 in ssd-lm), with a diffusion block size @xmath0. on the 13b ssd-2 model for our main evaluation, we first do 50k warmup steps without self-conditioning and then start a 100k-step pretraining with the full algorithm. it uses approximately 38b tokens from the c4 data in total. other pretraining hyperparameters can be found in §b. we show the pretraining losses of ssd-2 over time in §c. based on the trend of pretraining losses and the scale of our pretraining data compared to recent work, 10 we conjecture that our ssd-2 models are still considerably undertrained. due to our computing budget, we leave to future work a potential continued pretraining over current ssd-2 models on larger and better curated data. instruction finetuning while @xcite show the effectiveness of pretrained ssd-lm in general-domain text continuation, in this work, we primarily investigate the use of ssd-2 in downstream fine-tuning tasks, particularly on chat-style instruction following. 11 we finetune the models with the dolly dataset 12 containing 15k humancollected instructions and responses @xcite . dolly covers categories like open/closed-qa, brainstorming, and creative writing, though it may still be less powerful than the distillationbased data in terms of size and quality. 13 we finetune on 95% of the dolly data and use the rest for held-out evaluation. we finetune with a batch size of 384 and for 500 or 1000 steps for the 0.1b/2.7b/13b models. as a baseline, we finetune the autoregressive model opt (0.1b/2.7b/13b) on dolly with the same setup. ## related work diffusion-based language models have been receiving increasing attention as a potential alternative to autoregressive language models. we identify three win rate ppl (gpt-neo/gpt2) bertscore (precision/f1) ssd-2 13b 52.3% 7.58 / 9.62 85.9 / 85.2 opt 13b 47.7% 8.44 / 10.08 85.3 / 84.9 with autoregressive language models, various efforts have been made towards building chat-style instruction following models based on open source language models @xcite to replicate strong production-level closed source counterparts @xcite . many of such work are concurrent to ours and collect high-quality finetuning datasets by distilling prompts and responses from openai models @xcite @xcite . in this work, we deliberately seek fully open source data not depending on openai models and adopt the dolly data @xcite . we expect our models can be further improved with future releases of more curated chat-style instruction tuning datasets @xcite . one novel setup we explored in this work is the inference-time collaboration between a large, general-purpose diffusion chat model and small, user-specific models. inference-time collaboration has been generally explored in autoregressive models via ensembles of logits, either in an interpolation or contrastive manner @xcite @xcite @xcite . in diffusion models, classifierfree guidance in image generation @xcite contrastively reconstruct representations with and without a controlling attribute using a single model, whereas our work collaboratively decode with models with different sizes and inputs. we show an unique advantage of simplex-based diffusion language models in such inference-time collaboration compared to autoregressive language models. ## conclusion we present an exploratory step towards pretraining a large simplex-based diffusion language model ssd-2 and finetuning it with an open-source chatstyle instruction dataset. in a motivated setup where large general models and small user models are to collaborate with each other at inference time, we find ssd-2 substantially more collaborative than its autoregressive counterparts. these findings show the promise of diffusion language models as an instruction-following chat model and a worthy alternative to autoregressive language models.