Resumo Técnico: Aproximação de Rollout Paralelo para Geração de Imagens Autorregressiva em Espaço de Pixels

1. Definição do Problema

A geração de imagens autorregressiva (AR) em espaço de pixels visa modelar imagens como sequências de patches de pixels brutos sem depender de tokenização discreta ou tokenizadores pré-treinados separadamente. Embora conceitualmente direta, essa abordagem enfrenta dois desafios acoplados que historicamente limitaram seu desempenho em comparação com modelos AR em espaço latente ou modelos de difusão:

1. Desafio do Lado da Saída (Predição de Alta Dimensionalidade): Predizer patches de pixels brutos de alta dimensão (ex: vetores de 768 dimensões para patches de 16×16) em um único passo AR é difícil, levando a grandes erros de geração de passo único.
2. Desafio do Lado da Entrada (Descompasso entre Treino e Inferência): O treinamento padrão via teacher forcing condiciona o modelo em prefixos de verdade fundamental (ground-truth) limpos. No entanto, durante a inferência, o modelo deve condicionar-se em seus próprios patches de pixels previamente gerados (e potencialmente errôneos). Esse descompasso faz com que os erros se acumulem e se propaguem através dos passos AR, um fenômeno conhecido como viés de exposição (exposure bias).

Estratégias de mitigação existentes, como a predição de $x$ (eficaz em modelos de difusão) ou a injeção de ruído na entrada, abordam apenas parcialmente esses problemas. O treinamento de rollout exato, que exporia o modelo aos seus próprios prefixos gerados durante o treino, seria computacionalmente proibitivo para AR de tokens contínuos, pois exige amostragem sequencial para cada token.

2. Metodologia: Aproximação de Rollout Paralelo (PRA)

Os autores propõem a Aproximação de Rollout Paralelo (PRA), um framework escalável que aborda conjuntamente os desafios de saída e entrada, mantendo uma interface "pixel na entrada, pixel na saída".

2.1 Lado da Saída: Alvos Intermediários de Ponta a Ponta

Para reduzir a dificuldade da geração de passo único, o PRA não prediz diretamente patches de pixels de alta dimensão. Em vez disso, ele gera estados intermediários de baixa dimensão ($z_i \in \mathbb{R}^{d_z}$, onde $d_z \ll d$).

• Construção do Alvo: Durante o treinamento, o token de pixel de verdade fundamental $x_i$ e a representação do prefixo causal $h_{i-1}$ são mapeados para um estado intermediário alvo $z_i$ por meio de uma rede de projeção leve $g_\psi$. Este alvo é sensível ao contexto, incorporando informações do prefixo.
• Decodificação: Um decodificador de pixels causal $D_\phi$ mapeia o estado intermediário gerado $\hat{z}_i$ (e estados anteriores $\hat{z}_{<i}$) de volta para o espaço de pixels, produzindo o token de saída $\hat{x}_i$.
• Treinamento: O modelo AR é treinado para predizer $z_i$ usando uma cabeça de difusão de nível de token (objetivo de fluxo retificado/ rectified-flow), enquanto o decodificador é treinado para reconstruir $x_i$ a partir de $z_i$. Isso permite que o backbone AR opere em um espaço de menor dimensão enquanto a interface externa permanece no espaço de pixels.

2.2 Lado da Entrada: Entradas Semelhantes à Inferência em Paralelo

Para reduzir o descompasso entre treino e inferência sem o rollout sequencial, o PRA constrói entradas de pixels semelhantes à inferência em paralelo.

• Mecanismo: Em vez de treinar em pixels de verdade fundamental limpos, o PRA perturba os estados intermediários alvo $z_i$ (usando interpolação de fluxo retificado) para criar estados ruidosos $\tilde{z}_t$. Estes são passados através do mesmo decodificador de pixels usado na inferência para gerar tokens de pixels reconstruídos $\bar{x}_i$.
• Paralelismo: Como cada $\bar{x}_i$ é construído independentemente de seu respectivo estado alvo (em vez de uma trajetória autorregressiva sequencial), este processo permanece totalmente paralelizável.
• Loop de Treinamento: O modelo AR é então treinado para predizer o alvo $z_i$ condicionado ao prefixo desses pixels reconstruídos e semelhantes à inferência $\bar{x}_{<i}$. Isso aproxima a estrutura de erro do rollout em tempo de inferência, mantendo a eficiência do treinamento via teacher forcing.

3. Principais Contribuições

1. Framework Unificado: O PRA é o primeiro framework escalável que aborda conjuntamente os erros de saída de alta dimensão e o descompasso treino-inferência na geração de imagens em espaço de pixels sem exigir um tokenizador pré-treinado separadamente.
2. Aprendizado de Estado Intermediário: O artigo introduz um espaço de alvo intermediário aprendido de ponta a ponta que é otimizado especificamente para o modelo AR causal, em vez de depender de espaços latentes congelados de VAEs.
3. Aproximação de Rollout Paralelo: O método aproxima com sucesso os benefícios do treinamento de rollout (expondo o modelo a contextos gerados) através de uma construção paralela e não sequencial de entradas de pixels decodificadas.
4. Interface Pixel-In, Pixel-Out: Diferente dos modelos AR em espaço latente, o PRA mantém uma interface direta com patches de pixels brutos, facilitando potenciais aplicações em compreensão e geração de imagens unificadas.

4. Resultados Experimentais

Os autores avaliaram o PRA na geração de ImageNet-1K condicional de classe na resolução de 256×256.

• Desempenho de Estado da Arte:
  • PRA-S (135M parâmetros): Alcançou um FID de 2.58, superando modelos anteriores de AR em espaço de pixels de escala bilionária (ex: FARMER-1.9B com FID 3.60).
  • PRA-L (511M parâmetros): Alcançou um FID de 1.94, estabelecendo um novo estado da arte para modelos AR em espaço de pixels.
• Estudos de Ablação:
  • Alvos Intermediários: Alvos sensíveis ao contexto (incorporando a representação do prefixo) superaram alvos puramente locais ou codificadores LDM congelados. Uma dimensão intermediária de $d_z=16$ foi considerada ideal.
  • Construção de Entrada: O uso de entradas de pixels decodificadas (PRA) superou significativamente o treinamento em pixels de verdade fundamental limpos, estados intermediários de verdade fundamental ou simples injeção de ruído, confirmando a necessidade de combinar o loop de feedback de pixel da inferência.
• Compreensão de Imagem:
  • Quando aplicado ao linear probing de ImageNet, o PRA-L alcançou 68.80% de acurácia top-1, superando tanto os baselines de AR em espaço latente (SphereAR-L: 52.19%) quanto os baselines de difusão em espaço de pixels (JiT-L: 42.76%). Isso sugere que o modelo aprende representações visuais transferíveis.

5. Significância e Alegações

O artigo afirma que o PRA resolve os gargalos fundamentais que impedem o AR em espaço de pixels de competir com modelos de difusão e AR em espaço latente. Ao reduzir a dificuldade da geração de passo único através de estados intermediários de baixa dimensão e mitigar o acúmulo de erro através da aproximação de rollout paralelo, o PRA permite a geração de imagens de alta qualidade diretamente a partir de pixels brutos.

Os autores destacam que a natureza "pixel-in, pixel-out" do PRA, combinada com seu forte desempenho em linear probing, sugere um potencial para modelos de espaço de pixels unificados capazes tanto de geração de imagens de alta fidelidade quanto de compreensão visual robusta, eliminando a necessidade de tokenizadores ou autoencoders separados.

Limitações: Os autores observam que o PRA introduz componentes adicionais (codificador de estado intermediário, decodificador de pixels e um passo adicional de forward AR paralelo durante o treinamento). Embora estes sejam treinados de ponta a ponta, trabalhos futuros são necessários para validar o framework em domínios de dados e tarefas de geração mais amplos.