Resumo Técnico: Modelagem Generativa de Ponto Fixo com Máscara (CoFRe)

Declaração do Problema

Modelos Generativos com Máscara (MGMs), como o MDLM para texto e o MaskGIT para imagens, permitem a decodificação paralela e uma forte qualidade de geração em diversas modalidades. No entanto, eles enfrentam gargalos de eficiência significativos. Os MGMs padrão exigem uma passagem bidirecional completa da sequência em cada etapa de refinamento, levando a altos custos de treinamento (VRAM e tempo) e à degradação da qualidade de geração sob orçamentos de amostragem baixos (número limitado de passagens de ida totais).

As melhorias de eficiência existentes dividem-se em três categorias, nenhuma das quais resolve totalmente o problema:

1. Melhorias no amostrador (Sampler): Alteram as estratégias de seleção de tokens, mas deixam o custo do denoiser por etapa fixo.
2. Arquiteturas de profundidade fixa eficientes: Reduzem o desperdício arquitetural (ex: pulando tokens mascarados), mas ainda dependem de denoisers de capacidade fixa.
3. Modelos Adaptativos/Em Loop: Fornecem profundidade controlável, mas frequentemente falham em transferir efetivamente para a geração mascarada discreta devido a mudanças abruptas de estado por token (tokens revelados vs. mascarados) que interrompem as suposições de difusão contínua.

O desafio central é reduzir o custo computacional do denoiser e permitir profundidade adaptativa sem sacrificar a qualidade de geração, particularmente em regimes de baixo orçamento onde os MGMs padrão têm dificuldade.

Metodologia: Estrutura CoFRe

Os autores introduzem os Modelos Generativos de Ponto Fixo com Máscara (FP-MGMs) e uma estrutura completa de treinamento-inferência chamada CoFRe (Consistência, Ponto Fixo e Reuso). A abordagem substitui a pilha central de camadas de transformer distintas em um MGM padrão por um bloco de ponto fixo de compartilhamento de pesos.

1. Arquitetura de Denoising de Ponto Fixo

Em vez de uma pilha profunda de camadas únicas, o denoiser é decomposto em:

• Pilha de Pré-processamento ($P_{\theta_P}$): Camadas explícitas para a geração inicial do estado oculto.
• Bloco de Ponto Fixo ($F_{\theta_F}$): Um bloco compartilhado aplicado iterativamente para encontrar um estado de equilíbrio $h^*_t = F_{\theta_F}(h^*_t; \tilde{h}_t, t)$. A profundidade efetiva é controlada pelo número de iterações do solver ($N$), não pelo número de parâmetros.
• Pilha de Pós-processamento ($H_{\theta_H}$): Camadas explícitas que mapeiam o equilíbrio para logits.

Este design reduz a contagem de parâmetros (através do compartilhamento de pesos) e permite que o modelo adapte sua profundidade efetiva dinamicamente durante a inferência. O treinamento utiliza Retropropagação Estocástica Livre de Jacobiano (SJFB) para evitar o armazenamento de ativações para toda a trajetória do solver, aproximando os gradientes ao desenrolar um número aleatório de iterações com e sem rastreamento de gradiente.

2. Reuso de Três Estados (3SR)

O aquecimento padrão (usar a solução do passo anterior como o estado inicial para o próximo) é ineficiente em MGMs porque os estados dos tokens mudam de forma não uniforme: alguns tokens permanecem visíveis, outros permanecem mascarados e outros são recém-revelados.

• Mecanismo: A inicialização do solver $h^0_t$ é uma combinação ponderada da solução de ponto fixo anterior ($h^*_{t+1}$) e a saída da camada prévia atual ($h_{pre,t}$).
• Modulação por Token:
  • Tokens visíveis inalterados: Reutilizam totalmente a solução anterior ($\gamma = 1$).
  • Tokens ainda mascarados: Reutilizam parcialmente a solução, com o coeficiente de reuso aumentando conforme o denoising progride.
  • Tokens recém-revelados: Dependem minimamente da solução anterior ($\gamma = 0.2$), dependendo mais da injeção de entrada atual.
    Esta estratégia reconhece que os tokens recém-revelados passam pelas maiores mudanças de representação, enquanto os tokens estáveis requerem recomputação mínima.

3. Regularização de Consistência entre Passos ($L_{CONS}$)

Arquiteturas de ponto fixo sozinhas podem sofrer de deriva de representação entre os passos, especialmente sob orçamentos baixos.

• Mecanismo: Uma perda de consistência alinha as representações ocultas de um estado estudante "mais ruidoso" ($z_{t_s}$) com um estado professor "mais limpo" ($z_{t_c}$) derivado da mesma sequência subjacente (usando máscaras aninhadas).
• Objetivo: $L = L_{MGM} + \lambda \|h_s - \text{sg}(h_c)\|^2_2$, onde $h_s$ e $h_c$ são os estados ocultos finais antes dos logits.
• Efeito: Isso atua como uma autodestilação cross-time, afunilando as previsões e estabilizando o solver de ponto fixo, o que é crítico para a geração de baixo orçamento.

4. Conversão de Modelo Pré-treinado

Os autores demonstram que checkpoints de MDLM pré-treinados podem ser convertidos em FP-MDLMs sem retreinamento total. Ao mapear camadas específicas para os blocos de pré/pós-processamento e o bloco de ponto fixo (guiados pela análise CKA que mostra redundância em camadas mais profundas), e aplicar um estágio curto de adaptação professor-aluno com perda de divergência KL, o modelo alcança um desempenho forte com treinamento adicional mínimo.

Principais Contribuições

1. FP-MGMs: Uma nova classe de arquitetura que substitui camadas de transformer distintas por um bloco de ponto fixo compartilhado, reduzindo parâmetros e permitindo profundidade adaptativa.
2. Estrutura CoFRe: Uma receita completa que combina denoising de ponto fixo com Reuso de Três Estados (3SR) e Consistência de Passo Cruzado ($L_{CONS}$) para abordar os desafios específicos da geração mascarada discreta.
3. Adaptação Eficiente: Um método para converter MGMs pré-treinados em FP-MGMs eficientes via curta destilação, evitando o custo de treinar do zero.
4. Validação Empírica: Demonstrar que esses componentes coletivamente melhoram a relação qualidade-custo em modalidades de texto e imagem.

Resultados Experimentais

A estrutura foi avaliada no OpenWebText (OWT) para linguagem e ImageNette para imagens.

Linguagem (OpenWebText)

• Eficiência: Comparado ao baseline MDLM, o CoFRe reduz os parâmetros em 38,8% (170M $\to$ 104M), o tempo de treinamento em 11,5% e o uso de VRAM em 16,9%.
• Qualidade: Com um baixo orçamento de amostragem de 96 passagens de ida de blocos de transformer, o CoFRe melhora a perplexidade generativa (Gen. PPL) de 830,8 (MDLM) para 101,8. Mesmo comparado ao MDLM com Self-Distillation Through Time (SDTT), o CoFRe alcança uma perplexidade significativamente menor (101,8 vs 193,1) no mesmo orçamento.
• Ablações: $L_{CONS}$ é identificado como o principal motor dos ganhos de qualidade em baixo orçamento, enquanto o 3SR fornece melhorias adicionais ao otimizar a inicialização do solver.

Imagem (ImageNette)

• Eficiência: O CoFRe reduz o tempo de treinamento em 48,6% e a VRAM em 50,7% em comparação com o MaskGIT-Large.
• Qualidade: O CoFRe melhora o Fréchet Inception Distance (FID) em todos os orçamentos testados. Por exemplo, no orçamento 48, o FID melhora de 174,1 (MaskGIT-Large) para 96,7 (CoFRe).

Significância e Alegações

O artigo afirma que o CoFRe oferece uma estrutura prática para alcançar treinamento mais barato e geração mais forte sob orçamentos de amostragem limitados.

• Relação Qualidade-Custo: A principal contribuição é deslocar a fronteira de Pareto para modelos generativos mascarados, permitindo geração de alta qualidade com significativamente menos parâmetros e recursos computacionais.
• Agnóstico a Modalidade: A abordagem é validada tanto em tokens de texto discretos quanto em latentes de imagem quantizados, sugerindo generalização.
• Praticidade: A capacidade de converter modelos pré-treinados via curta adaptação torna a abordagem acessível sem exigir recursos massivos de pré-treinamento.
• Limitações: Os autores notam modestamente que os resultados atuais são limitados à escala OWT e ImageNette, e que o solver de ponto fixo adiciona um overhead de controle de fluxo que pode nem sempre se traduzir diretamente em acelerações de tempo real (wall-clock speedups) dependendo da implementação do hardware. Eles também observam que o método é projetado para decodificação mascarada monotônica e pode exigir adaptação para samplers que re-mascaram tokens.

Em resumo, o artigo argumenta que, ao tratar o denoiser como um problema de ponto fixo e gerenciar cuidadosamente o reuso de estados e a consistência entre os passos, os modelos generativos mascarados podem se tornar significativamente mais eficientes sem sacrificar, e muitas vezes melhorando, a qualidade da geração.