O BERT é apenas uma etapa única de difusão de texto

• Modelos de linguagem mascarada como BERT e RoBERTa também podem ser interpretados como o conceito de modelo de difusão de texto
• Diferente de auto-regressivos (por exemplo, GPT), foi demonstrado que é possível aplicar geração por blocos e o modo de restauração progressiva
• Confirmamos experimentalmente que, ao treinar ajustando a taxa de mascaramento etapa por etapa, o RoBERTa também consegue gerar texto natural
• Mesmo sem mudar a estrutura existente, é possível obter capacidade de geração ao ajustar apenas o objetivo de treinamento
• Comparado ao GPT-2, um modelo de difusão baseado em RoBERTa também produz resultados de geração de texto com consistência em um certo nível

---

Visão geral

O Gemini Diffusion divulgado pelo Google DeepMind é um modelo de linguagem baseado em difusão que, diferente da família GPT, gera texto em blocos de uma vez. Essa abordagem é uma forma de gerar texto que refina ruído aleatório gradualmente. Ao analisar o paper relacionado (Large Language Diffusion Models), constatou-se que a difusão de linguagem discreta é uma generalização da modelagem de linguagem mascarada (MLM). Ou seja, foi investigado experimentalmente se essa ideia permitiria geração de texto também em modelos da família BERT.

Nota: posteriormente, o paper do DiffusionBERT valida essa ideia de forma mais rigorosa.

História resumida dos Transformers

O Transformer, proposto inicialmente em 2017, tinha estrutura encoder-decoder. Em 2018, surgiram famílias especializadas com o separo entre encoder (BERT: bidirecional, com foco em reconstrução mascarada) e decoder (GPT: autoregressivo, com foco em predição sequencial).

• Apenas encoder (família BERT)
  • Recebe o contexto completo da sentença e mascara certas partes com <MASK>, restaurando-as com o restante como contexto
  • Ponto forte em representação de frase, classificação etc.
• Apenas decoder (família GPT)
  • Prediz o próximo token a partir da sequência dada
  • Se destacou em geração, resumo, tradução etc.

O BERT foi usado diretamente para tarefas de classificação, mas depois a família GPT ganhou mais espaço com a melhora em capacidades de geração.

Modelo de difusão de linguagem discreta

Modelos de difusão popularizaram-se originalmente em geração de imagens.
No caso de imagens:

• Processo direto: adiciona-se ruído gaussiano progressivamente a uma imagem limpa até chegar a um estado de ruído puro
• Processo reverso: repetição de desruido progressivo com um modelo de deep learning para restaurar o dado original

Ao aplicar para texto, a forma mais simples é um processo de ruído baseado em mascaramento.

• Direto (mascaramento)
  • Em t=0 há texto original; à medida que a etapa aumenta, parte dos tokens é substituída aleatoriamente por <MASK>
  • Na etapa final, tudo é preenchido por <MASK>
• Reverso (desruído)
  • O transformer encoder aprende a restaurar os tokens originais a partir de uma configuração de mascaramento dada
  • Com baixa taxa de mascaramento, a restauração é mais fácil; com taxa alta, fica mais difícil
  • Repetindo do percentual de mascaramento mais alto para o mais baixo, é possível gerar uma sequência inteira

Nesse framework de difusão, o modelo é treinado somando a perda de desruído para diferentes níveis de mascaramento em cada etapa. O objetivo de restauração mascarada do BERT é, na prática, parte de uma difusão de texto. Combinando agenda de taxa de mascaramento e desruído iterativo, é possível expandir o objetivo do BERT para um procedimento gerativo em linguagem natural.

Experimento RoBERTa Diffusion

O RoBERTa, apresentado em 2019, se caracteriza por uma expansão de hiperparâmetros e dados em relação ao BERT, além de um objetivo de treino mais simples (MLM Only).
No experimento, foram usados os pesos-base do RoBERTa, tokenizer e Trainer das bibliotecas transformers e datasets da HuggingFace.
Com base no conjunto de dados WikiText, o ajuste fino foi feito com o seguinte processo:

• Em cada batch, amostra-se uma de 10 etapas da agenda de difusão (mask_probs: 1.0~0.1) para mascarar
• Usa-se um diffusion_collator customizado que, após escolher a probabilidade de mascaramento, aplica <MASK> estocasticamente em cada token
• Para manter o contexto do prompt, os primeiros 16 tokens são sempre preservados

Máscara de dados (custom collator):

• Para cada lote de tokens do exemplo, faz-se padding e escolhe-se aleatoriamente uma probabilidade de mascaramento
• Aplica-se <MASK> de forma probabilística em todos os tokens, exceto os 16 iniciais
• Retorna o conjunto de dados mascarado e os rótulos-alvo

Geração (inference):

• Entrada com sequência de 256 tokens, nos quais os 16 primeiros são prompt e o restante está como <MASK>
• Em cada etapa, o modelo amostra os tokens previstos e preenche, depois remasca-se novamente uma fração dos tokens
• Repete-se reduzindo progressivamente a taxa de mascaramento; ao final, tudo é restaurado

Exemplo de saída gerada:

• Com prompt e texto subsequente natural, é possível gerar texto com boa consistência
• Algumas particularidades vêm do formato de pré-processamento do WikiText (por exemplo, hífen @-@)

Comparação com GPT-2

• Embora GPT-2 seja um pouco mais rápido e apresente consistência maior, o RoBERTa Diffusion também mostra desempenho acima do esperado (há potencial para melhoria incremental)
• Qualidade e velocidade podem melhorar com novas estratégias como AR-Diffusion, Skip-Step Diffusion e otimizações

Conclusão

• Modelos de linguagem mascarada como o RoBERTa, ao treinar com taxa de mascaramento ajustada progressivamente, também podem se tornar motores de geração de linguagem natural
• Só com o treinamento que vai degradando e restaurando texto por tokens <MASK> já é possível confirmar a viabilidade de convertê-los em modelos totalmente generativos
• É possível obter capacidade de geração com a alteração apenas do objetivo de treino, sem mudar a própria estrutura
• No fim, os modelos da família BERT também são, essencialmente, equivalentes a um único modelo de difusão de texto