Resumo Técnico: Set Diffusion

Declaração do Problema

Modelos de difusão discretos emergiram como uma alternativa competitiva aos modelos de linguagem autorregressivos (AR), oferecendo velocidades de geração mais rápidas e melhor controlabilidade. No entanto, as abordagens de difusão existentes enfrentam limitações significativas:

1. Restrições de Comprimento Fixo e Bidirecionalidade: A maioria dos modelos de difusão discretos depende de arquiteturas que exigem contexto bidirecional completo, impedindo o uso de cache de Chave-Valor (KV) e restringindo a geração a comprimentos fixos.
2. Rigidez da Difusão em Blocos: Embora a Difusão em Blocos (BD3LM) preencha parcialmente a lacuna ao gerar blocos sequenciais de tamanho fixo da esquerda para a direita, ela sofre de inflexibilidade. Não consegue suportar inserções em posições arbitrárias (infilling) ou geração de comprimento variável dentro de um bloco, e as atualizações do cache KV são atrasadas até que um bloco inteiro seja decodificado.
3. Limitações de Trade-off: Os métodos atuais interpolam entre AR e difusão primariamente variando um tamanho de bloco fixo, o que restringe a flexibilidade de decodificação e o paralelismo.

O desafio central é desenvolver um modelo de linguagem que interpole entre autorregressão e difusão não apenas pelo tamanho do bloco, mas variando a ordem em que os conjuntos de tokens são gerados, permitindo assim geração de comprimento flexível, cache de KV e decodificação de ordem arbitrária, mantendo uma alta qualidade de amostra.

Metodologia

Os autores introduzem o Set Diffusion, um framework que fatoriza o processo de geração sobre conjuntos de tokens de posição e comprimento flexíveis, em vez de blocos fixos.

1. Modelo Probabilístico

A difusão de conjuntos define uma distribuição de probabilidade autorregressiva sobre uma sequência de conjuntos de tokens $\sigma = (\sigma_1, \dots, \sigma_N)$, onde cada $\sigma_n$ é um subconjunto de posições de tokens. A verossimilhança é fatorada como:
$$p_\theta(x) = \sum_{\sigma} \prod_{n=1}^N \pi(\sigma_n | x_{<\sigma_n}) \cdot p_\theta(x_{\sigma_n} | x_{<\sigma_n})$$
Aqui, $\pi$ é uma política de ordem que seleciona o próximo conjunto de posições, e $p_\theta$ modela a verossimilhança condicional dos tokens naquele conjunto usando difusão. O modelo otimiza um Limite Inferior de Evidência Negativa (NELBO) que soma as perdas de difusão sobre esses conjuntos.

2. Arquitetura Set-Causal e Cache de KV

Uma inovação arquitetural fundamental é o transformer set-causal. Ao contrário dos modelos de difusão bidirecionais padrão, esta arquitetura suporta cache de KV após cada etapa de inferência.

• Mecanismo: Quando um conjunto de tokens $x_{\sigma_n}$ é gerado, seus vetores de chave e valor são cacheados e reutilizados para condicionar as previsões subsequentes.
• Eficiência de Treinamento: Os autores propõem um algoritmo de treinamento vetorizado que processa contextos limpos e queries ruidosas simultaneamente em uma única passagem direta usando uma máscara de atenção especializada. Esta máscara garante que tokens ruidosos em um conjunto atentem apenas para outros tokens ruidosos no mesmo conjunto e para todos os tokens limpos em conjuntos precedentes.

3. Interpolação de Ordenações de Tokens

Para controlar o trade-off entre qualidade de amostra (estreitamento do limite) e eficiência de inferência (paralelismo), os autores introduzem esquemas de tempo de revelação dependentes da posição.

• Esquemas de Deslocamento de Posição (Position-Offset Schedules): Cada posição de token $\ell$ recebe um tempo de revelação aleatório $R_\ell$ extraído de uma distribuição definida por uma função de distribuição cumulativa (CDF) $\alpha^\ell_\tau$.
• Instanciação de Janela Deslizante (SW-SetDLM): Os autores propõem uma instanciação específica onde os tokens possuem um intervalo de geração ativo $[a_\ell, a_\ell + w]$. O tamanho da janela $w$ controla o viés "da esquerda para a direita":
  • Um $w$ pequeno (aproximando-se de $1/L$) concentra a ordem na geração da esquerda para a direita, recuperando o comportamento AR e estreitando o ELBO.
  • Um $w$ grande (aproximando-se de 1) permite a geração agnóstica à ordem, recuperando a difusão de máscara padrão.
  • Valores intermediários permitem a decodificação de janela deslizante, permitindo que o modelo preveja conjuntos de tokens de comprimento flexível em posições arbitrárias, suportando tarefas como infilling.

4. Inferência

Durante a inferência, o modelo amostra posições candidatas com base na ordenação atual e no esquema de deslocamento de posição. Em seguida, ele remove o ruído dessas posições em paralelo usando o cache de KV. Os tokens aceitos são confirmados e o cache é atualizado imediatamente, permitindo a geração de comprimento variável e decodificação eficiente.

Principais Contribuições

1. Framework Set Diffusion: Uma nova classe de modelos de linguagem que fatoriza a geração sobre conjuntos de tokens de posição e comprimento flexíveis, generalizando autorregressão, difusão agnóstica à ordem e difusão em blocos.
2. Arquitetura Set-Causal: Uma arquitetura transformer que suporta atualizações de cache de KV após cada etapa de inferência, permitindo a geração eficiente de comprimento variável.
3. Interpolação de Ordem: Um método principiado para interpolar entre AR e difusão ajustando esquemas de ordenação dependentes da posição (especificamente esquemas de janela deslizante) em vez de tamanhos de bloco fixos.
4. Desempenho de Estado da Arte: Demonstração de que o Set Diffusion alcança fronteiras de trade-off velocidade-qualidade superiores em raciocínio matemático, sumarização e geração incondicional comparado a modelos de difusão anteriores, e supera significativamente a difusão em blocos em tarefas de infilling.

Resultados Experimentais

Os autores avaliaram o Set Diffusion (especificamente o SW-SetDLM) em vários benchmarks:

• Raciocínio Matemático (GSM8K): O Set Diffusion alcança uma fronteira Pareto de velocidade-precisão estritamente melhor que a Difusão em Blocos. Por exemplo, com um tamanho de janela de decodificação $S \le 8$, ele atinge 66,41% de precisão pass@1 com um throughput de $\ge 60,42$ tokens/seg, superando a Difusão em Blocos ($S=8$) que atinge 63,53% de precisão a $\ge 55,39$ tokens/seg.
• Infilling (ROCStories): O Set Diffusion supera substancialmente a Difusão em Blocos em tarefas de infilling (preenchimento de 1 ou 3 sentenças intermediárias), alcançando pontuações ROUGE mais altas com até ~25% de decodificação mais rápida. A difusão em blocos tem dificuldades aqui devido à sua estrutura de bloco rígida da esquerda para a direita.
• Sumarização (CNN/DailyMail): O Set Diffusion atinge pontuações ROUGE competitivas em comparação com baselines de difusão e até 10% de decodificação mais rápida que a Difusão em Blocos. Também supera um baseline AR, que os autores observam que tende a sofrer overfitting neste dataset.
• Geração Incondicional e Verossimilhança: No OpenWebText e LM1B, o Set Diffusion melhora as pontuações MAUVE sobre a Difusão em Blocos enquanto decodifica mais rápido. No LM1B, atinge a melhor perplexidade entre os modelos de difusão (27.83 PPL para $S \le 8$) mantendo um throughput maior que modelos de difusão anteriores.

Significância e Alegações

O artigo afirma que o Set Diffusion representa um avanço significativo ao desacoplar a interpolação entre autorregressão e difusão de estruturas de blocos rígidas. Ao interpolar via ordenação de tokens, o framework oferece:

• Maior Flexibilidade: Suporte para inserções em posições arbitrárias e geração de comprimento variável, que são críticos para tarefas como infilling.
• Eficiência: A capacidade de atualizar caches de KV após cada etapa, um recurso anteriormente indisponível em modelos de difusão, levando a uma inferência mais rápida.
• Desempenho: Um trade-off superior entre qualidade de amostra e velocidade de inferência comparado aos modelos de linguagem de difusão existentes.

Os autores posicionam o Set Diffusion como um framework unificador que recupera AR, difusão agnóstica à ordem e difusão em blocos como casos especiais, fornecendo um mecanismo ajustável para equilibrar paralelismo e qualidade de amostra com base nos requisitos específicos da tarefa de inferência.