Resumo Técnico: BlockGen: Modelagem de Sequência por Blocos Flexível com Amostradores Híbridos

Declaração do Problema

Transformers autorregressivos (AR) dominam a modelagem de sequências, mas sofrem com limites de geração sequencial (throughput/latência) e restrições de atenção causal que dificultam tarefas de raciocínio que exigem contexto bidirecional. Modelos de difusão discreta oferecem uma alternativa ao refinar iterativamente sequências ruidosas, permitindo atualizações paralelas. No entanto, comparações recentes entre Modelos de Difusão de Máscara (MDMs) e Modelos de Difusão de Estado Uniforme (USDMs) produziram insights conflitantes dependendo do regime de amostragem:

1. Amostragem Ancestral: USDMs frequentemente superam MDMs em tarefas de downstream apesar de uma perplexidade pior, atribuído à capacidade dos USDMs de autocorrigir erros (tokens podem transitar entre quaisquer estados), enquanto MDMs fixam tokens uma vez desmascarados.
2. Geração por Blocos: Trabalhos anteriores comparam majoritariamente a difusão de sequência completa. Permanece incerto se a vantagem do USDM se mantém quando os tokens são gerados bloco a bloco (semi-autorregressivo), um cenário que permite o uso de KV caching e inferência mais rápida.
3. Corretores Informados: Comparações existentes frequentemente utilizam amostradores preditor-corretor (PC) "não informados" que reinjetam ruído em posições aleatórias. Não se sabe como corretores "informados" (visando erros prováveis) afetam o ranking MDM vs. USDM, particularmente sem modelos de verificador auxiliares.

Metodologia: BlockGen

Os autores introduzem o BlockGen, um modelo de sequência por blocos treinado sobre uma mistura de tamanhos de bloco, instanciado com difusão mascarada ou uniforme dentro dos blocos.

1. Formulação de Mistura e Treinamento

O BlockGen define a densidade da sequência como uma mistura sobre $M$ tamanhos de bloco $S = \{s_1, \dots, s_M\}$:
$$p_{\theta}^{\text{BlockGen}}(x) = \sum_{i=1}^M \gamma_i p_{\theta}^{(s_i)}(x)$$
onde $p_{\theta}^{(s_i)}$ é uma densidade válida fatorizada em blocos de tamanho $s_i$.

• Denoiser Compartilhado: Em vez de treinar modelos separados, um único denoiser $\hat{x}_\theta$ recebe o tamanho do bloco $L'$ como entrada, ajustando os padrões de atenção de acordo.
• Limites de Verossimilhança: O artigo deriva dois limites para a densidade de mistura. Um limite log-sum-exp mais justo é usado para avaliação, enquanto um limite de verossimilhança de mistura mais frouxo (uma expectativa sobre tamanhos de bloco) é usado para o treinamento para manter a eficiência computacional.
• Amostragem Estratificada: Para reduzir a variância durante o treinamento multi-GPU, os tamanhos de bloco são selecionados via amostragem estratificada em vez de amostragem i.i.d.

2. Preditor-Corretor Informado por AR (ARPC)

Uma contribuição fundamental é o amostrador ARPC, que aproveita a mistura do BlockGen para realizar correções informadas sem um verificador auxiliar.

• Mecanismo: Durante a geração, após os passos de difusão ancestral padrão, o amostrador executa um passo de "proposta limpa". Ele então realiza uma segunda passagem forward em modo AR ($L'=1$) usando o mesmo denoiser para pontuar os tokens propostos.
• Pontuação: Tokens com baixa log-verossimilhança AR são identificados como "improváveis" e são re-ruídos (corrigidos) para um nível de ruído anterior.
• Vantagem: Isso permite que o modelo use suas próprias previsões AR para guiar as correções de difusão, intercalando efetivamente os paradigmas AR e de difusão dentro de um único modelo.

3. Atenção ao Nível de Bloco

O BlockGen utiliza atenção causal de bloco: tokens dentro de um bloco atendem bidirecionalmente, enquanto o bloco atende causalmente a todos os blocos precedentes. Isso permite o KV caching entre blocos decodificados, preservando os benefícios de eficiência da geração semi-autorregressiva.

Resultados Principais

1. Raciocínio Matemático (GSM8K)

Avaliado no GSM8K usando modelos treinados em TinyGSM:

• Amostragem Ancestral: No cenário bloco a bloco, USDMs superam MDMs, especialmente no regime de poucos passos (baixo NFE). Isso confirma que a vantagem do USDM de sequência completa se transfere para a geração por blocos.
• Amostragem ARPC: A lacuna de desempenho entre USDMs e MDMs diminui e se inverte em alto NFE. Sob ARPC com blocos de 16 tokens, MDMs superam levemente os USDMs na acurácia do GSM8K.
• ARPC vs. Baselines: O ARPC supera consistentemente a amostragem ancestral e o EIPC (Entropy-Informed PC) na maioria dos orçamentos de NFE. O ARPC com MDMs aproxima-se do desempenho da amostragem AR com decodificação greedy.

2. Modelagem de Linguagem (OpenWebText & LM1B)

• Verossimilhança: O treinamento sobre uma mistura de tamanhos de bloco melhora significativamente a perplexidade em comparação com modelos de bloco fixo. O BlockGen com difusão mascarada alcança 17.5 Val. PPL no OpenWebText (vs. 16.7 para AR e 21.6 para o melhor BDM de bloco fixo), estreitando a lacuna entre modelos de difusão e autorregressivos.
• Qualidade de Amostragem:
  • Sob amostragem ancestral, USDMs alcançam menor Perplexidade Generativa (Gen. PPL) que MDMs.
  • Sob ARPC, a tendência se inverte em alto NFE (ex: 64 NFE/bloco), onde o masked-ARPC alcança uma fronteira de Gen. PPL melhor que o uniform-ARPC.
  • O ARPC reduz a lacuna para AR, mas não ultrapassa o AR em Gen. PPL dentro do regime prático de baixa temperatura.

Significância e Alegações

O artigo afirma resolver duas questões abertas sobre difusão discreta:

1. Validade por Blocos: A vantagem da difusão de estado uniforme sobre a difusão de máscara observada na geração de sequência completa se transfere para a geração bloco a bloco sob amostragem ancestral.
2. Correção Informada: O ranking entre MDMs e USDMs é dependente de regime. Enquanto USDMs se destacam com amostradores ancestrais simples ou não informados, corretores informados (especificamente ARPC) podem fechar essa lacuna e até inverter o ranking, favorecendo a difusão de máscara em orçamentos de computação mais altos (NFE).

Os autores enfatizam que o BlockGen fornece um framework flexível que interpola entre AR e difusão pura, permitindo que um único modelo atue tanto como um amostrador por blocos quanto como um verificador AR. Eles observam que, embora suas alegações empíricas sejam robustas dentro dos orçamentos e amostradores testados, elas não estabelecem se esses trade-offs persistem na escala de grandes modelos de linguagem (LLMs), já que seus experimentos são limitados a backbones de 170 milhões de parâmetros. O trabalho sugere que a escolha do paradigma de difusão (máscara vs. uniforme) não é universalmente superior, mas depende fortemente da estratégia de amostragem e do orçamento computacional.