Time-Annealed Perturbation Sampling: Diverse Generation for Diffusion Language Models

1. O Problema

Os Modelos de Linguagem de Difusão (Diffusion-LMs) emergiram como uma alternativa promissora aos modelos autoregressivos tradicionais, oferecendo dinâmicas de inferência diferentes que permitem a geração paralela e o refinamento iterativo de texto. No entanto, esses modelos enfrentam um desafio significativo: a falta de diversidade na geração.

Assim como nos modelos autoregressivos, os Diffusion-LMs tendem a produzir saídas repetitivas, conservadoras ou "seguras" quando submetidos a instruções que promovem diversidade. As técnicas existentes para controlar a diversidade (como temperature scaling, top-k, top-p ou min-p) foram desenvolvidas principalmente para modelos autoregressivos e não se transferem bem para a dinâmica de difusão. Além disso, métodos que injetam aleatoriedade apenas no nível de saída (tokens) muitas vezes comprometem a coerência, a fluência ou a aderência à instrução, criando um trade-off difícil entre diversidade e qualidade.

2. Metodologia: TAPS (Time-Annealed Perturbation Sampling)

Os autores propõem o TAPS, uma estratégia de inferência sem treinamento (training-free) que explora a estrutura temporal inerente aos modelos de difusão.

A. Inspiração Teórica: Divisão de Trabalho Temporal

O trabalho baseia-se na observação empírica de que os Diffusion-LMs exibem uma divisão de trabalho temporal:

• Etapas iniciais de denoising: Determinam a estrutura semântica global e o significado de alto nível.
• Etapas finais de denoising: Focam no refinamento léxico local, escolha de palavras e fluência.

B. O Mecanismo do TAPS

O TAPS perturba o sinal de condicionamento (o embedding do prompt) durante o processo de inferência, mas de forma controlada pelo tempo:

1. Perturbação Inicial: Nas etapas iniciais do processo de difusão (quando a estrutura semântica está sendo formada), uma perturbação estocástica significativa é injetada no embedding do contexto. Isso força o modelo a explorar diferentes trajetórias semânticas, evitando a convergência prematura para um único significado.
2. Recozimento Temporal (Annealing): A magnitude da perturbação decai monotonicamente à medida que o processo de difusão avança.
3. Refinamento Final: Nas etapas finais, a perturbação é reduzida a zero, permitindo que o modelo se alinhe novamente ao prompt original e refine os detalhes léxicos e estilísticos, garantindo fluência e aderência à instrução.

C. Estabilização e Preservação de Qualidade

Para evitar que a perturbação degrade a qualidade da geração, o TAPS emprega três mecanismos técnicos cruciais:

• Redimensionamento da Ruído (Rescaling): O embedding perturbado é reescalado para corresponder à média e ao desvio padrão dos embeddings originais, alinhando a estrutura estatística.
• Interpolação (Mixing): Uma mistura ponderada ($\psi$) é aplicada entre o embedding perturbado e o original para controlar a força efetiva da perturbação.
• Projeção Preservadora de Norma: Uma projeção é aplicada para garantir que a magnitude (norma) dos embeddings perturbados de cada token permaneça consistente com a do embedding original, preservando a geometria aprendida pelo modelo.

3. Principais Contribuições

1. Identificação Empírica: Evidenciam que os Diffusion-LMs possuem uma estrutura semântica temporal distinta, onde as etapas iniciais moldam a semântica e as finais refinam o léxico.
2. Novo Método de Inferência: Propõem o TAPS, uma estratégia simples, escalável e sem treinamento que melhora a diversidade sem sacrificar a qualidade.
3. Validação Abrangente: Demonstram a eficácia do método em diversos backbones (LLaDA e TraDo) e tarefas, incluindo geração criativa, seguimento de instruções e raciocínio matemático.

4. Resultados Experimentais

Os autores avaliaram o TAPS em benchmarks como WritingPrompts, NoveltyBench, Arena-Hard-Auto e GSM8K.

• Diversidade: O TAPS superou consistentemente todas as linhas de base (incluindo Top-k, Top-p, Min-p, EDT e Diverse Prompt) em métricas de diversidade lexical, semântica (Sent-BERT) e de embedding (EAD). A melhoria foi particularmente notável na diversidade semântica, indicando que o método explora caminhos de significado diferentes, não apenas variações superficiais de palavras.
• Qualidade: Diferente de outras abordagens que degradam a coerência ao aumentar a diversidade, o TAPS manteve ou até superou levemente a qualidade da geração (avaliada por GPT-4o) em dimensões como criatividade, coerência e relevância.
• Raciocínio (GSM8K): Em tarefas de raciocínio matemático, o TAPS mostrou-se robusto. Embora a precisão de uma única amostra (Pass@1) tenha caído ligeiramente, a precisão por votação majoritária (Majority Voting) foi superior à das linhas de base. Isso sugere que o TAPS gera um conjunto de soluções com erros menos correlacionados, permitindo que a agregação recupere a resposta correta com mais eficácia.
• Robustez: O método demonstrou ser robusto a diferentes temperaturas de decodificação e funcionou bem em modelos com arquiteturas diferentes (autoregressivo-semi vs. não-autoregressivo).

5. Significado e Impacto

Este trabalho é significativo por várias razões:

• Paradigma de Controle: Estabelece que a diversidade em modelos de difusão pode ser controlada explorando a dimensão temporal do processo de denoising, em vez de apenas manipular a distribuição de probabilidade dos tokens finais.
• Eficiência: Por ser uma estratégia de inferência sem treinamento, o TAPS pode ser aplicado imediatamente a modelos de difusão pré-treinados, sem o custo computacional de fine-tuning.
• Direção Futura: Reforça a ideia de que os Diffusion-LMs herdam características fundamentais dos modelos de difusão de imagem (separação entre formação semântica e refinamento), abrindo caminho para futuras pesquisas em geração de texto controlada e criativa baseada em difusão.

Em resumo, o TAPS oferece uma solução elegante para o problema da "homogeneidade" na geração de texto, permitindo que modelos de difusão explorem o espaço semântico de forma mais rica enquanto mantêm a fidelidade e a qualidade necessárias para aplicações práticas.