Geradores de imagens com IA já prontos aprendem a preencher qualquer região por uma fração do custo

Modelos autorregressivos geram imagens da mesma forma que modelos de linguagem escrevem texto: token por token. Primeiro, a imagem é dividida em códigos discretos. Depois, um transformer prevê esses códigos um a um, em geral linha por linha, da esquerda para a direita e de cima para baixo.

Essa ordem é chamada de raster scan. Pesquisadores da Stony Brook University e da ÉTS Montréal propuseram uma forma de transformar um modelo raster já treinado em um gerador de ordem arbitrária. Ele consegue preencher lacunas dentro do quadro (inpainting), completar a imagem além das bordas (outpainting) e editar regiões selecionadas.

Para isso, não é preciso retreinar do zero. Segundo estimativa dos autores, a adaptação adicional custa mais de uma ordem de grandeza menos que treinar do zero um modelo de "qualquer ordem".

Por que um modelo linha a linha não sabe consertar o meio da imagem

No treinamento raster, o modelo sempre enxerga apenas o que foi desenhado acima e à esquerda. Ele nunca aprendeu a prever um token cujos vizinhos à direita e abaixo já são conhecidos. Por isso, uma máscara no meio do quadro o deixa sem saída.

Existem duas formas de contornar a limitação, e ambas são caras:

• Treinamento do zero com uma mistura de ordens. O RandAR, por exemplo, exige 300 épocas, enquanto o LlamaGen-XL, que gera linha a linha, usa 100. O orçamento triplica.
• Difusão discreta, categoria à qual o MaskGIT pertence. O modelo restaura iterativamente os tokens cobertos pela máscara, com atenção bidirecional. A qualidade é boa, mas o KV-cache (armazenamento de dados intermediários já calculados, para não recalculá-los a cada passo), a decodificação especulativa e servidores como o vLLM não funcionam com esse esquema.

Empresas que já têm um pipeline e servidores autorregressivos acabam obrigadas a montar uma segunda stack ao lado. Os autores consideram esse gasto desnecessário.

Dois detalhes sem os quais o retreinamento falha

A ideia se apoia na matemática: pela regra da cadeia de probabilidades, a distribuição conjunta dos tokens pode ser decomposta em qualquer ordem. Logo, um modelo que aprendeu bem as imagens em uma ordem já sabe quase tudo de que precisa. Falta a ele a habilidade de usar esse conhecimento em outra ordem. Para isso, foram necessários dois componentes, e nenhum deles funciona sozinho.

Embeddings posicionais "com olhar adiante" (future-aware). Na ordem raster, o modelo sempre sabe o que vai desenhar em seguida. Na ordem aleatória, o mesmo contexto pode levar a alvos diferentes.

Um token de céu e um token de rosto são previstos de forma diferente a partir exatamente dos mesmos dados visíveis, e os gradientes passam a entrar em conflito. No modelo de linguagem XLNet, esse problema era resolvido com um segundo fluxo de atenção, que dobra o custo computacional da camada e quebra o KV-cache. Aqui, a solução é mais simples.

À entrada de cada passo, soma-se um vetor aprendido da posição que deve ser prevista em seguida. A tabela começa com zeros, e a soma é feita apenas na camada de entrada. A arquitetura continua sendo um transformer comum, de fluxo único.

A escolha da posição para a atenção é feita separadamente. No LlamaGen, que usa embeddings posicionais rotacionais (RoPE), e no RAR, que usa tabelas de posições absolutas, o que vale é o índice espacial do token, e não o número do passo. O modelo preserva tudo o que sabia sobre a vizinhança entre pixels.

Complexidade gradual das ordens (curriculum). Se permutações aleatórias forem apresentadas de uma vez, as perdas são enormes, os gradientes ficam ruidosos e o modelo sai do ponto favorável do pré-treinamento, esquecendo o que acumulou. Os autores introduzem um coeficiente que cresce linearmente: no início, quase todos os exemplos seguem a ordem raster; depois, entram cada vez mais ordens aleatórias.

Cada exemplo do batch recebe a própria permutação. Simplesmente reduzir a taxa de aprendizado não resolve: o gradiente continuaria apontando para a mesma direção desfavorável. O cronograma altera a própria composição dos dados de treinamento, e os gradientes iniciais ficam alinhados com a solução pré-treinada.

Se qualquer um dos dois componentes for removido, o treinamento diverge ou a qualidade cai de forma notável.

O que foi obtido e qual é o custo

O método foi aplicado a quatro modelos base de duas arquiteturas e dois tamanhos: LlamaGen-L, LlamaGen-XL, RAR-L e RAR-XL. Resultados:

• Em inpainting e outpainting no ImageNet, os modelos adaptados superam o MaskGIT por boa margem em FID. Essa métrica compara a distribuição das imagens geradas com a das imagens reais, e quanto menor, melhor.
• Com o MAR-L, que trabalha com tokens contínuos e uma cabeça de difusão, os resultados são comparáveis. O modelo não foi treinado especificamente para nenhuma dessas tarefas.
• Os mesmos pesos dão conta de inpainting, outpainting e edição por regiões. Muda apenas o conjunto de tokens conhecidos na entrada.
• Há um custo. No LlamaGen-XL, o FID na geração convencional condicionada por classe aumentou. Inception Score, precision e recall ficaram próximos dos valores originais. Os autores dizem com clareza: é uma troca de qualidade por flexibilidade, e não uma conversão gratuita.

A arquitetura em si continuou causal, então o KV-cache e a decodificação paralela funcionam. O modelo pode ser executado pelo vLLM com o registro de um modelo próprio. Segundo os autores, isso é inviável por princípio na difusão discreta.

O que isso significa para engenheiros

Quem já tem um modelo autorregressivo de imagens pré-treinado e a infraestrutura para ele ganha edição sem uma segunda stack e sem novo treinamento. Inpainting e outpainting passam a ser um modo do modelo já existente. As limitações também são claras. A validação se restringe ao ImageNet e a quatro modelos base, e a geração convencional perde um pouco de qualidade. Mas a direção é evidente: capacidades úteis podem ser extraídas de modelos que já foram treinados. O código está publicado no GitHub, no repositório Ordered-AR/unlockedadaptation.