Mirai: Autoregressive Visual Generation Needs Foresight

O artigo apresenta o Mirai, um framework que melhora a geração visual autoregressiva ao injetar informações futuras durante o treinamento, acelerando significativamente a convergência e aprimorando a qualidade das imagens sem alterar a arquitetura ou adicionar sobrecarga na inferência.

O essencial

Imagine que você está tentando montar um quebra-cabeça gigante de 10.000 peças, mas com uma regra estranha: você só pode olhar para a peça que está na sua mão e para a peça imediatamente à sua frente. Você não pode olhar para a foto completa na caixa, nem para as peças que ainda estão longe.

Esse é o problema que os modelos de Geração Visual Autoregressiva (como o LlamaGen) enfrentam hoje. Eles criam imagens peça por peça, da esquerda para a direita, como se estivessem lendo um livro. O problema é que, às vezes, a "cabeça" do pássaro fica desenhada antes de saberem onde vai o "corpo", resultando em imagens onde a cabeça está flutuando ou o corpo está torto. Eles são ótimos em detalhes locais, mas péssimos em manter a coerência global.

O artigo "Mirai" (que significa "Futuro" em japonês) propõe uma solução brilhante: dar ao modelo "visão de futuro" durante o treinamento.

Aqui está a explicação simplificada, usando analogias do dia a dia:

1. O Problema: O Pintor Cego

Imagine um pintor que precisa pintar um cenário de um foguete decolando.

• O modelo antigo (LlamaGen): Ele pinta o chão, depois a base do foguete, depois a fumaça... mas como ele só olha para o que acabou de pintar, ele pode pintar a fumaça indo para a direita, quando o foguete está indo para a esquerda. Ele não tem uma "visão geral" do que está por vir.
• O resultado: Imagens que parecem boas de perto, mas que, de longe, parecem estranhas ou desconexas (como a cabeça do papagaio separada do corpo na imagem do artigo).

2. A Solução: O "Oráculo" de Futuro

Os autores do Mirai perguntaram: "E se, enquanto o pintor pinta a peça atual, alguém sussurrasse no ouvido dele como será a peça daqui a 5 passos?"

Essa é a ideia de Foresight (Previsão/Visão de Futuro). O modelo não muda a forma como ele pinta (ainda é peça por peça), mas durante o treino, ele recebe um "sussurro" do futuro para ajudar a planejar o presente.

O Mirai faz isso de duas maneiras criativas:

A. Mirai-E (O "Espelho do Futuro" - Explícito)

Imagine que o pintor tem um espelho mágico que mostra uma versão levemente adiantada do que ele vai pintar logo em seguida.

• O modelo olha para o que ele está fazendo agora e compara com uma "versão futura" que ele mesmo está criando (usando uma técnica chamada EMA, que é como uma média móvel de suas próprias previsões).
• Analogia: É como um dançarino que, enquanto faz um passo, já sente o ritmo do próximo passo no corpo, garantindo que a dança flua suavemente.
• Resultado: O modelo aprende a se alinhar com o que vem a seguir, mantendo a estrutura correta.

B. Mirai-I (O "Mestre Externo" - Implícito)

Aqui, o pintor recebe ajuda de um Mestre Externo (um modelo de IA diferente, chamado DINOv2) que já viu a foto inteira.

• Esse Mestre olha para a imagem completa e diz: "Ei, na posição onde você está pintando agora, lembre-se que o céu azul continua ali e a montanha está ali".
• Analogia: É como ter um arquiteto experiente ao lado do pintor. O pintor ainda pinta tijolo por tijolo, mas o arquiteto garante que a parede não vai cair porque ele já viu o prédio todo.
• Resultado: O modelo aprende a entender o contexto global (a "vibe" da imagem inteira) enquanto pinta os detalhes locais.

3. O Grande Truque: A "Visão em 2D"

O papel descobriu algo crucial: não basta apenas olhar para o futuro; é preciso olhar para o futuro no lugar certo.

• Se você apenas olhar para a próxima peça na fila (1D), você pode estar olhando para uma peça que está longe na imagem real (como olhar para o pé de alguém enquanto pinta a cabeça).
• O Mirai organiza a "visão de futuro" em uma grade 2D (como um mapa). Ele olha para as peças vizinhas no espaço (cima, baixo, lado), não apenas na sequência de leitura.
• Analogia: É a diferença entre ler um livro linha por linha e olhar para um mapa de metrô. No mapa, você vê que a próxima estação está "ao norte", não apenas "na próxima linha". Isso mantém a geometria da imagem correta.

4. Os Resultados: Velocidade e Qualidade

O que acontece quando damos essa "visão de futuro" para o modelo?

1. Aprendizado Super Rápido: O modelo aprende 5 a 10 vezes mais rápido. Em vez de precisar de 400 "rodadas" de treino para ficar bom, ele precisa de apenas 40 ou 80. É como se um aluno de música, com um bom professor, aprendesse uma música em uma semana o que levaria um ano para aprender sozinho.
2. Imagens Perfeitas: As imagens geradas têm uma coerência global incrível. O foguete sobe reto, a fumaça segue a direção certa, e o pássaro tem a cabeça presa ao corpo.
3. Sem Custo Extra na Hora de Usar: A mágica acontece apenas no treino. Quando você vai usar o modelo para gerar uma imagem, ele é exatamente o mesmo, rápido e sem custos extras. O "Mestre" e o "Espelho" saem da sala, e o pintor trabalha sozinho, mas agora ele já aprendeu a pintar perfeitamente.

Resumo em uma frase

O Mirai ensina os modelos de IA a "pensar no futuro" enquanto aprendem a criar imagens, garantindo que cada peça do quebra-cabeça se encaixe perfeitamente no todo, resultando em imagens mais bonitas e em um tempo de treinamento muito menor.

Resumo técnico

1. O Problema

Os geradores visuais autoregressivos (AR) modernos, como o LlamaGen, modelam imagens como sequências de tokens discretos e são treinados com o objetivo de prever o próximo token (Next-Token Prediction - NTP). Embora essa abordagem tenha sido bem-sucedida em linguagem natural, ela enfrenta desafios significativos na visão computacional:

• Supervisão Causal Estrita: O treinamento baseia-se apenas no feedback do token imediato seguinte. Isso cria uma "visão de túnel", onde o modelo não tem acesso a informações globais ou de longo alcance durante o processo de geração.
• Incoerência Global: Como os tokens são gerados sequencialmente (em ordem de varredura raster), as pistas globais levam muitos passos para se propagar. Isso frequentemente resulta em imagens que são localmente consistentes, mas globalmente desalinhadas (ex: estruturas quebradas, poses impossíveis, como uma cabeça desconectada de um corpo).
• Convergência Lenta: A falta de planejamento global força o modelo a aprender a estrutura da imagem passo a passo, o que retarda significativamente a convergência do treinamento.

2. Metodologia: O Framework Mirai

Os autores propõem o Mirai (que significa "futuro" em japonês), um framework de treinamento que injeta informações de "previsão" (foresight) no processo de treinamento autoregressivo, sem alterar a arquitetura do modelo ou o custo de inferência.

A ideia central é alinhar as representações internas do modelo AR com informações derivadas de tokens futuros, permitindo que o modelo "planeje" a estrutura global enquanto mantém a decodificação causal na inferência.

Investigação Preliminar (Diagnóstico)

Antes de definir o Mirai, os autores realizaram experimentos controlados em três eixos para entender como a previsão deve ser aplicada:

1. Nível de Injeção: Injetar a previsão na saída (predição de tokens) ou nas representações internas?
   • Resultado: Injetar na saída (como em Multi-Token Prediction) causa interferência de gradientes e piora o desempenho. Injetar nas representações internas (camadas intermediárias) regulariza o estado oculto sem forçar a predição de tokens discretos, melhorando a coerência.
2. Posicionamento (Layout): A previsão deve seguir a ordem 1D (varredura raster) ou a estrutura 2D da imagem?
   • Resultado: O alinhamento em grade 2D (vizinhos espaciais) é superior ao 1D. Isso preserva a geometria visual e fornece pistas mais coerentes sobre o contexto local.
3. Fonte da Previsão: Usar um codificador unidirecional (próprio do modelo) ou bidirecional (externo)?
   • Resultado: Ambos funcionam, mas o alinhamento com um codificador bidirecional (que vê a imagem inteira) fornece um contexto global implícito poderoso.

As Duas Variantes do Mirai

O framework Mirai adiciona uma função de perda de alinhamento ($L_{Foresight}$) à perda de previsão de próximo token ($L_{NTP}$):

• Mirai-E (Explicit Foresight):
  • Utiliza uma versão com Média Móvel Exponencial (EMA) do próprio modelo AR como codificador de previsão.
  • A EMA é unidirecional, fornecendo uma previsão explícita e indexada por posição de um pequeno conjunto de locais futuros na grade 2D.
  • O modelo AR alinha suas representações internas com essas previsões futuras específicas.
• Mirai-I (Implicit Foresight):
  • Utiliza um codificador bidirecional pré-treinado (como o DINOv2) que processa a imagem completa.
  • Como o codificador é bidirecional, cada token de saída contém informações implícitas sobre todo o contexto da imagem (incluindo o futuro).
  • O modelo AR alinha suas representações internas com os recursos do codificador bidirecional na mesma posição espacial.

Inferência: Durante a inferência, os cabeçotes de projeção e os codificadores de previsão são descartados. O modelo gera tokens um por um, de forma estritamente causal, com o mesmo custo computacional do modelo base.

3. Contribuições Principais

1. Investigação Sistemática da Previsão: O trabalho demonstra que a previsão de tokens futuros, quando aplicada corretamente (nível interno, layout 2D), não viola a causalidade, mas sim a fortalece, permitindo um planejamento global.
2. Proposta do Mirai: Um framework simples e eficaz que melhora modelos AR visuais sem modificar a arquitetura de inferência.
3. Aceleração e Qualidade: Demonstração de que o Mirai acelera drasticamente a convergência e melhora a qualidade de geração em benchmarks padrão.

4. Resultados Experimentais

Os experimentos foram realizados no conjunto de dados ImageNet (256x256) com modelos da família LlamaGen (B, L, XL).

• Aceleração de Treinamento:
  • O Mirai-I acelerou a convergência do LlamaGen-B em até 10x. Um modelo treinado com Mirai-I por 40 épocas atingiu uma qualidade (FID) comparável ao modelo base treinado por 400 épocas.
  • O Mirai-E mostrou uma aceleração de 5x.
• Qualidade de Geração (FID-50K):
  • No LlamaGen-B (300 épocas), o FID caiu de 5.34 (base) para 4.34 (Mirai-I) e 4.49 (Mirai-E).
  • No LlamaGen-XL, o Mirai-I alcançou um FID de 2.59, superando todos os métodos baseados em AR e competindo com modelos de Difusão e GANs.
• Coerência Espacial: Visualizações de t-SNE das representações internas mostram que o Mirai produz campos de cor mais suaves e espacialmente coerentes na grade 2D, indicando que o modelo aprendeu melhor a estrutura espacial da imagem.
• Eficiência Computacional: Embora o treinamento tenha um custo ligeiramente maior devido à perda de alinhamento (6.6% a mais para Mirai-I e 38.2% para Mirai-E), a redução no número de épocas necessárias resulta em uma redução total de FLOPs de 9.4x (Mirai-I) e 3.6x (Mirai-E) para atingir o mesmo nível de qualidade.

5. Significado e Conclusão

O artigo "Mirai" desafia a noção de que modelos autoregressivos visuais devem ser estritamente limitados ao feedback imediato do próximo token. Ele estabelece que:

• A Geração Visual AR Precisa de Previsão: A ausência de sinais de planejamento global é uma limitação fundamental dos modelos AR puros.
• Causalidade e Previsão são Compatíveis: É possível usar informações futuras durante o treinamento para guiar a formação de representações internas, sem comprometer a capacidade de geração causal na inferência.
• Impacto Prático: O método oferece uma rota eficiente para melhorar modelos de geração de imagem baseados em AR, tornando-os competitivos com modelos de difusão em termos de qualidade e velocidade de treinamento, sem aumentar o custo de inferência.

Em resumo, o Mirai introduz um mecanismo de "olhar para o futuro" durante o treinamento, permitindo que os modelos autoregressivos construam uma compreensão global da imagem desde os estágios iniciais, resultando em gerações mais estruturadas e coerentes.