Transferência de Aprendizado para Modelos de Difusão de Texto

Transfer Learning for Text Diffusion Models

January 30, 2024

Autores: Kehang Han, Kathleen Kenealy, Aditya Barua, Noah Fiedel, Noah Constant

cs.AI

Resumo

Neste relatório, exploramos o potencial da difusão de texto para substituir a decodificação autoregressiva (AR) no treinamento e implantação de grandes modelos de linguagem (LLMs). Estamos particularmente interessados em verificar se modelos AR pré-treinados podem ser transformados em modelos de difusão de texto por meio de um procedimento de adaptação leve que chamamos de "AR2Diff". Começamos estabelecendo uma configuração de linha de base robusta para o treinamento de modelos de difusão de texto. Comparando várias arquiteturas e objetivos de pré-treinamento, descobrimos que treinar um modelo apenas de decodificador com um objetivo de LM prefixado é o melhor ou quase o melhor em várias tarefas. Com base nessa descoberta, testamos várias configurações de transferência de aprendizagem para modelos de difusão de texto. Na tradução automática, observamos que a difusão de texto tem desempenho inferior à abordagem AR padrão. No entanto, na síntese de código e em QA extrativo, descobrimos que modelos de difusão treinados do zero superam os modelos AR em muitos casos. Também observamos ganhos de qualidade com o AR2Diff — adaptando modelos AR para usar decodificação por difusão. Esses resultados são promissores, considerando que a difusão de texto é relativamente pouco explorada e pode ser significativamente mais rápida que a decodificação AR para geração de textos longos.

English

In this report, we explore the potential for text diffusion to replace autoregressive (AR) decoding for the training and deployment of large language models (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.