Resumo Técnico: Escalonamento Recursivo em Modelos de Difusão Mascarada

Declaração do Problema

Os Modelos de Difusão Mascarada (MDMs) emergiram como um paradigma promissor para a geração de sequências discretas, oferecendo capacidades de geração totalmente paralelas que contrastam com a natureza sequencial dos modelos autorregressivos (AR). No entanto, o escalonamento de MDMs para alcançar uma geração de alta qualidade tem dependido tradicionalmente de dois eixos: o aumento da contagem de parâmetros (largura/profundidade) ou o aumento do número de passos de denoising ($T$). Esta abordagem acarreta custos computacionais significativos, tanto no treinamento (modelos maiores) quanto na inferência (mais passagens forward). O desafio central abordado é se a profundidade recursiva pode servir como um terceiro eixo de escalonamento, ortogonal, que melhore a qualidade da geração sem aumentar a contagem de parâmetros, e se pode substituir passos de denoising adicionais para reduzir o computação em tempo de inferência.

Metodologia

Os autores introduzem os Modelos de Difusão Mascarada Recursivos (R-MDMs), que incorporam um mecanismo de looping recursivo ao paradigma padrão de treinamento de MDM.

Arquitetura

Em vez de aplicar uma camada transformer $KL$ apenas uma vez, o R-MDM utiliza um bloco transformer de $K$ camadas compartilhado, aplicado $L$ vezes sequencialmente dentro de um único passo de denoising.

• Eficiência de Parâmetros: O modelo possui a mesma contagem de parâmetros que um modelo de passagem única de $K$ camadas, apesar de possuir uma profundidade efetiva de $K \times L$.
• Passagem Forward: Dada uma entrada ruidosa $x_t$, o estado oculto é atualizado iterativamente: $h^{(\ell)} = f_\theta(h^{(\ell-1)}, \ell, L)$. A saída do loop $\ell$ serve como entrada para o loop $\ell+1$.
• Embedding de Passo: Para permitir que o modelo distinga entre diferentes estágios de recursão, um embedding de passo $v^{(\ell)}$ é adicionado ao estado oculto. Este embedding é derivado do progresso de loop normalizado $s_\ell = (\ell-1)/(L-1)$ via um pequeno MLP, permitindo que o modelo especialize sua computação com base na profundidade da recursão.

Objetivos de Treinamento

O artigo avalia quatro estratégias de supervisão para os loops recursivos:

1. Todos os passos (ALL): Médias a perda de entropia cruzada através de todos os $L$ loops, fornecendo sinais de gradiente densos (inspirado em aprendizado supervisionado profundo).
2. Passo final (FINAL): Supervisiona apenas os logits do último loop, estendendo os objetivos padrão de MDM.
3. Ponderado (Weighted): Atribui pesos não uniformes às perdas em diferentes loops.
4. Truncado (Truncated): Supervisiona apenas os últimos $k$ loops.

Os autores descobrem que a perda de Todos os passos (All-steps loss) produz o desempenho mais robusto, particularmente em regimes de baixo computo.

Inferência e Agendamento

No tempo de inferência, o modelo pode ser avaliado com uma profundidade de recursão de amostragem $L_s$ que difere da profundidade de treinamento $L_t$. Os autores exploram vários esquemas de treinamento para $L_t$, incluindo profundidade fixa, amostragem estocástica e agendamento baseado em currículo. Um currículo reverso (começando com alta profundidade de recursão e diminuindo para 1) é considerado superior, forçando o modelo a aprender estratégias de refinamento iterativo de alta capacidade precocemente no treinamento.

Principais Contribuições

1. Introdução da Profundidade Recursiva como um Eixo de Escalonamento: O artigo propõe o looping recursivo dentro dos passos de denoising como um mecanismo para aumentar a profundidade efetiva do modelo sem aumentar a contagem de parâmetros.
2. Demonstração de Eficiência de Parâmetros: Demonstra-se que os R-MDMs igualam ou superam baselines não recursivos significativamente maiores (até $5\times$ mais parâmetros) em tarefas de raciocínio estruturado.
3. Eficiência de Tempo de Inferência: O estudo demonstra que o refinamento recursivo pode substituir parcialmente os passos de denoising, permitindo que os modelos alcancem geração de alta qualidade com menos passagens forward totais ($T \times L$).
4. Análise Dependente de Tarefa: O trabalho avalia sistematicamente os trade-offs entre recursão, contagem de parâmetros e passos de denoising através de tarefas estruturadas (Sudoku, Countdown) e não estruturadas (Text8).

Resultados Experimentais

Experimentos foram conduzidos em três conjuntos de dados: Sudoku (9x9 e 25x25), Countdown (planejamento aritmético com 3, 4 e 5 operandos) e Text8 (modelagem de linguagem ao nível de caractere).

Tarefas de Raciocínio Estruturado (Sudoku & Countdown)

• Substituição de Parâmetros (RQ1): Um modelo recursivo $(K \otimes L)$ consistentemente iguala ou supera um baseline de passagem única iso-FLOP $(KL \otimes 1)$.
  • Exemplo: No Sudoku 9x9, um modelo $(6 \otimes 5)$ (10.6M parâmetros) alcança 93.6% de Taxa de Puzzle Válido (VPR) em $T=5$, igualando o desempenho de um modelo de passagem única de 30 camadas (53.1M parâmetros) que alcança apenas 77.8% no mesmo número de passos.
  • Exemplo: No Countdown-4, um modelo $(3 \otimes 5)$ (5.5M parâmetros) atinge 70.2% de Taxa de Alcance do Alvo (RTR) em $T=20$, superando um baseline de 15 camadas (26.7M parâmetros) que atinge 69.4%.
• Redução de Passos (RQ2): Modelos recursivos requerem menos passos de denoising para atingir um limiar de qualidade alvo.
  • Um modelo em loop ($L_s=3, T=5$) alcança uma taxa de puzzle válido de 95% em 15 passagens forward totais, enquanto um modelo de passagem única comparável requer 40 passos (uma redução de 2.7x no computo).
  • No Sudoku 25x25 com 80% de mascaramento, um modelo recursivo ($L=5$) alcança 91.8% VPR em $T=1$, enquanto o baseline falha quase inteiramente neste número de passos.
• Flexibilidade: Modelos treinados com um currículo reverso ($L_t$ decrescente) podem ser efetivamente avaliados com diferentes profundidades de amostragem ($L_s$), permitindo uma alocação flexível de computo durante a inferência.

Modelagem de Linguagem Não Estruturada (Text8)

• No conjunto de dados Text8, modelos recursivos desempenham-se abaixo dos baselines de passagem única em termos de métricas baseadas em verossimilhança (NLL, Perplexidade).
• No entanto, análises qualitativas sugerem que modelos recursivos produzem texto mais coerente, apesar de pontuações de verossimilhança piores, indicando que as métricas padrão podem ser enganosas neste regime. Os autores observam que os benefícios da recursão são mais pronunciados em tarefas com estrutura iterativa explícita.

Significância e Alegações

O artigo afirma que a profundidade recursiva é um mecanismo de escalonamento praticamente útil para MDMs, particularmente para tarefas de geração estruturada. A significância reside em duas áreas principais:

1. Eficiência de Parâmetros: R-MDMs podem alcançar o desempenho de modelos muito maiores com uma fração dos parâmetros, tornando a geração de alta qualidade viável em cenários de recursos limitados.
2. Eficiência de Computo: Ao substituir o refinamento recursivo por passos de denoising adicionais, os R-MDMs podem reduzir o número total de passagens forward necessárias na inferência, melhorando o trade-off entre qualidade de geração e latência.

Os autores concluem que, embora a recursão atue como um poderoso mecanismo para propagação de restrições iterativas em domínios estruturados, sua utilidade na modelagem de linguagem não estruturada permanece uma questão aberta que requer investigação adicional. O trabalho enriquece a compreensão do escalonamento de modelos generativos além da visão tradicional centrada na contagem de parâmetros.