Transfer Learning for Text Diffusion Models (2401.17181v1)

Abstract: In this report, we explore the potential for text diffusion to replace autoregressive (AR) decoding for the training and deployment of LLMs. We are particularly interested to see whether pretrained AR models can be transformed into text diffusion models through a lightweight adaptation procedure we call ``AR2Diff''. We begin by establishing a strong baseline setup for training text diffusion models. Comparing across multiple architectures and pretraining objectives, we find that training a decoder-only model with a prefix LM objective is best or near-best across several tasks. Building on this finding, we test various transfer learning setups for text diffusion models. On machine translation, we find that text diffusion underperforms the standard AR approach. However, on code synthesis and extractive QA, we find diffusion models trained from scratch outperform AR models in many cases. We also observe quality gains from AR2Diff -- adapting AR models to use diffusion decoding. These results are promising given that text diffusion is relatively underexplored and can be significantly faster than AR decoding for long text generation.

• The paper introduces AR2Diff, a novel method that adapts autoregressive models into non-autoregressive text diffusion models for efficient text generation.
• It demonstrates improved performance on tasks like code synthesis and extractive question-answering through iterative adaptation.
• The study highlights substantial efficiency gains, with speedups of 10×–30× for long texts, despite slower individual diffusion steps.

Introduction

In the landscape of LLMs, autoregressive (AR) decoding has long been the established method for generating text sequences. However, this traditional left-to-right approach is inherently sequential and can be inefficient, particularly when generating longer spans of text. This limitation has given rise to exploration into non-autoregressive (non-AR) decoding methods that offer the potential for more flexible and efficient text generation by allowing for parallel token generation or iterative text refinement.

Non-Autoregressive Decoding and AR2Diff

One emerging alternative is text diffusion models, which show promise in accelerating the decoding process and perhaps even enhance quality on certain tasks. This paper introduces a novel training paradigm termed "AR2Diff," which seeks to transition pre-existing AR models into text diffusion models capable of utilizing non-AR decoding. The core hypothesis is that leveraging the iterative refinement akin to human writing strategies could lead to improvements over the conventional AR approach.

Initial findings confirm that decoder-only models with a prefix LLMing objective stand out across several tasks, including machine translation, extractive question-answering, and code synthesis. Surprisingly, despite utilizing fewer parameters, these models outperform their encoder-decoder counterparts, highlighting the efficacy of the simplified model structure.

Comparative Performance Evaluation

In a series of experiments comparing AR and text diffusion models across various architectures, tasks, and learning setups, it was observed that while text diffusion models lagged behind AR models in machine translation tasks, they showed promising results in both code synthesis and extractive question-answering. Notably, AR2Diff's method of adapting AR models to use diffusion decoding brought forth quality improvements, particularly at larger model sizes and with increased adaptation steps prior to fine-tuning. The performance gains suggest that there is indeed merit in the iterative adaptation from AR models to text diffusion models.

Efficiency and Future Directions

Critical to the practical application of these models is their efficiency. The paper performs an analysis of inference speeds, revealing a significant advantage of diffusion models in generating longer texts, with speedups ranging from 10× to 30× depending on the text length. Yet, each individual diffusion step remains slower compared to an AR step, underscoring the need for further optimization in diffusive models.

Conclusion

The comprehensive analysis of non-AR decoding paradigms, such as text diffusion models, suggests that we are on the cusp of a paradigm shift in LLMs. The research presented in this paper could potentially lower the cost and barriers of entry for developing more efficient and competitive text generation models at scale. Notably, the AR2Diff method stands as a promising avenue to refine and re-purpose existing autoregressive LMs towards this new generation paradigm, signaling a potential shift in preference for future LLM deployments.