Enhancing Spatial Understanding in Image Generation via Reward Modeling

O artigo apresenta um novo método que aprimora a compreensão espacial em modelos de geração de imagens, criando o conjunto de dados SpatialReward-Dataset e o modelo de recompensa SpatialScore para avaliar e otimizar, via aprendizado por reforço online, a precisão das relações espaciais complexas descritas em prompts de texto.

O essencial

Imagine que você pediu a um pintor de IA para desenhar uma cena muito específica: "Uma mesa de escritório moderna. No centro, um laptop. À direita do laptop, uma xícara de café. À esquerda, um caderno. E atrás de tudo, uma estante cheia de livros."

O problema é que, até agora, esses pintores de IA eram ótimos em desenhar objetos bonitos, mas péssimos em entender onde colocar as coisas. Eles podiam desenhar um laptop e uma xícara lindos, mas muitas vezes colocavam a xícara em cima do laptop, ou do lado errado, ou esqueciam a estante. Para conseguir o desenho perfeito, você teria que pedir 10 ou 20 vezes, sorteando qual ia ficar certo.

Os autores deste artigo decidiram resolver esse problema criando um "Professor de Geografia Visual" para ensinar a IA a entender espaço.

Aqui está a história de como eles fizeram isso, explicada de forma simples:

1. O Problema: O Pintor Confuso

As IAs atuais são como alunos que memorizaram a lista de ingredientes de uma receita, mas não sabem como montar o prato. Elas sabem o que é um "cachorro" e o que é um "passeio", mas se você pedir "um cachorro correndo atrás de uma bicicleta", elas podem colocar o cachorro em cima da bicicleta ou na frente.

Os modelos de avaliação que existiam antes eram como juízes que só olhavam se a foto era bonita (cores, iluminação), mas não se a lógica espacial fazia sentido. Eles davam nota alta para fotos erradas se a foto fosse "bonita".

2. A Solução: Criando o "Espaço-Reward-Dataset" (O Banco de Questões)

Para ensinar a IA, os pesquisadores precisaram de um material de estudo. Eles criaram um banco de dados gigante com 80.000 pares de imagens.

• Como funcionava: Eles pegavam um pedido complexo (ex: "A cadeira está à esquerda da mesa") e geravam duas imagens:
  1. A Perfeita: Onde a cadeira estava realmente à esquerda.
  2. A Perturbada (Errada): Onde eles trocavam a cadeira para a direita, mas mantinham o resto igual.
• O Truque: Humanos especialistas revisaram tudo para garantir que a imagem "perfeita" estivesse realmente certa e a "errada" estivesse realmente errada. Isso criou um "livro de exercícios" de alta qualidade.

3. O Professor: O Modelo "SpatialScore"

Com esse banco de dados, eles treinaram um novo modelo chamado SpatialScore. Pense nele como um professor rigoroso de geometria.

• Enquanto outros modelos diziam: "Nossa, que cores vibrantes! Nota 10!", o SpatialScore olhava e dizia: "Espera aí! O prompt disse que o gato estava embaixo da mesa, mas na imagem ele está em cima. Nota 2!"
• O incrível é que esse professor de 7 bilhões de parâmetros (um modelo de tamanho médio) ficou mais inteligente em entender espaço do que os "gigantes" pagos e proprietários (como o GPT-5 ou Gemini), que muitas vezes alucinavam e erravam a posição dos objetos.

4. O Treinamento: A "Aula de Reforço" (Reinforcement Learning)

Agora que eles tinham o professor, precisavam ensinar o pintor (o modelo de geração de imagens). Eles usaram uma técnica chamada Aprendizado por Reforço Online.

• O Processo: A IA tenta desenhar a cena. O professor (SpatialScore) olha, dá uma nota baseada na precisão espacial e diz: "Isso está bom, mantenha!" ou "Isso está errado, tente de novo!".
• O Filtro Top-K (A Peneira Inteligente): Aqui entra uma ideia genial. Às vezes, o pedido é tão fácil que a IA acerta quase tudo. Se o professor der nota alta para todos, a IA não aprende nada novo.
  • Os pesquisadores criaram uma estratégia de filtragem: eles pegavam apenas os melhores exemplos (os que a IA acertou) e os piores (os que a IA errou feio) para treinar. Eles ignoravam os "mais ou menos".
  • Analogia: É como um treinador de futebol que não perde tempo com jogadores que já sabem jogar ou que não têm potencial. Ele foca nos que estão no limite do sucesso e nos que estão no limite do fracasso para extrair o máximo de aprendizado.

5. O Resultado: O Mestre da Composição

Depois desse treinamento, a IA mudou drasticamente:

• Ela começou a entender prompts longos e complexos.
• Se você pedisse "três maçãs em uma mesa, duas vermelhas à esquerda e uma verde à direita", ela desenhava exatamente isso, sem precisar de 20 tentativas.
• Ela superou modelos que usavam regras rígidas (que quebravam quando havia objetos escondidos ou oclusos) e modelos que só olhavam para a beleza da imagem.

Resumo em uma frase

Os pesquisadores criaram um professor especialista em espaço (SpatialScore) e um banco de provas rigorosas (Dataset) para treinar IAs de geração de imagens, fazendo com que elas deixem de ser apenas "pintoras bonitas" e se tornem "arquitetas precisas" que entendem exatamente onde cada objeto deve ficar no mundo.

Isso abre portas para criar imagens muito mais complexas, úteis para design, arquitetura e até simulações para robôs, onde a posição exata das coisas é crucial.

Resumo técnico

1. O Problema

O avanço recente em modelos de geração de imagem a partir de texto (Text-to-Image ou T2I) melhorou significativamente a fidelidade visual e a criatividade. No entanto, esses modelos ainda enfrentam dificuldades críticas ao interpretar relações espaciais complexas entre múltiplos objetos em prompts detalhados.

• Falha na Composição: Modelos atuais frequentemente falham em posicionar objetos corretamente (ex: "à esquerda de", "atrás de", "alinhado com") quando o prompt contém múltiplas restrições espaciais.
• Limitações das Abordagens Atuais:
  • Modelos de Recompensa Existentes: Modelos baseados em preferência humana (como HPS, ImageReward) ou alinhamento texto-imagem (como VQAScore) focam na estética e semântica geral, falhando em avaliar a precisão espacial fina. Eles muitas vezes atribuem pontuações altas a imagens espacialmente incorretas.
  • Avaliação Baseada em Regras (GenEval): Métodos que usam detectores de objetos e classificadores de cores falham em prompts longos e complexos, sofrendo com oclusão e falta de generalização.
  • APIs Proprietárias e VLMs Abertos: Embora modelos de Linguagem Visual (VLMs) proprietários (como GPT-5) tenham bom raciocínio espacial, são caros demais para o treinamento online por reforço (RL). VLMs abertos (como Qwen2.5-VL) sofrem de alucinações e não são otimizados especificamente para raciocínio espacial complexo.

2. Metodologia

Os autores propõem uma abordagem de Aprendizado por Reforço Online (Online RL) guiada por um modelo de recompensa especializado, composto por três pilares principais:

A. SPATIALREWARD-DATASET (Construção do Dataset)

Para treinar um modelo de recompensa confiável, foi criado um dataset de 80.000 pares de preferência adversariais:

• Geração Adversarial: Utilizando GPT-5, foram criados prompts com relações espaciais complexas. A partir desses prompts "perfeitos", geraram-se prompts "perturbados" alterando intencionalmente uma ou mais relações espaciais (ex: trocar "esquerda" por "direita").
• Geração de Imagens: Imagens foram geradas a partir dos prompts perfeitos e perturbados usando modelos de ponta (Qwen-Image, HunyuanImage, Seedream).
• Validação Humana: Pares foram rigorosamente revisados por especialistas para garantir que a imagem "perfeita" seguisse o prompt e a imagem "perturbada" violasse a relação espacial, filtrando casos onde o modelo de geração falhou em ambos.

B. SPATIALSCORE (Modelo de Recompensa)

Um modelo de recompensa especializado foi desenvolvido para avaliar a precisão das relações espaciais:

• Arquitetura: Baseado no Qwen2.5-VL-7B (um modelo de linguagem visual), com a cabeça de linguagem substituída por uma cabeça linear de recompensa.
• Treinamento: Utiliza uma distribuição Gaussiana para modelar a pontuação de recompensa (em vez de um valor determinístico) e é treinado com a função de perda de Bradley-Terry para maximizar a probabilidade de escolher a imagem "vencedora" (perfeita) sobre a "perdedora" (perturbada).
• Desempenho: O SpatialScore supera modelos proprietários e VLMs abertos de grande porte na tarefa de distinguir imagens com relações espaciais corretas das incorretas.

C. Treinamento Online RL com GRPO e Filtragem Top-k

O SpatialScore foi utilizado como sinal de recompensa para afinar o modelo base Flux.1-dev usando o algoritmo GRPO (Group Relative Policy Optimization):

• Desafio do Viés de Vantagem: Em prompts fáceis, a maioria das amostras gera recompensas altas, o que pode resultar em vantagens negativas para amostras de alta qualidade devido à normalização do grupo.
• Estratégia Top-k Filtering: Para mitigar esse viés e melhorar a eficiência, os autores propõem selecionar apenas os top-k (melhores) e bottom-k (piores) exemplos de um grupo de amostras para calcular a média e o desvio padrão da vantagem. Isso equilibra a distribuição de recompensas e reduz o número de avaliações de função (NFE) necessárias, acelerando o treinamento.

3. Contribuições Principais

1. SPATIALREWARD-DATASET: A introdução de um dataset de 80k pares de preferência adversariais, cuidadosamente curado por humanos, focado especificamente em relações espaciais complexas.
2. SPATIALSCORE: O desenvolvimento de um modelo de recompensa especializado que supera modelos proprietários de ponta (como GPT-5 e Gemini-2.5 Pro) na avaliação de precisão espacial, oferecendo um sinal de recompensa confiável e de baixo custo para RL.
3. Estratégia de Filtragem Top-k: Uma técnica inovadora para treinamento de RL em geração de imagens que mitiga o viés de vantagem em prompts de dificuldade variável e reduz custos computacionais.
4. Validação Empírica: Demonstração de que o uso desse modelo de recompensa leva a melhorias significativas e consistentes na compreensão espacial, superando abordagens baseadas em regras (GenEval) e modelos base não treinados.

4. Resultados

Os experimentos foram conduzidos em múltiplos benchmarks (DPG-Bench, TIIF-Bench, UniGenBench++):

• Avaliação do Modelo de Recompensa: O SpatialScore (7B parâmetros) alcançou 95,8% de precisão na previsão de preferência, superando o GPT-5 (93,3%) e o Gemini-2.5 Pro (95,1%) no benchmark de avaliação espacial.
• Geração de Imagens:
  • No benchmark in-domain (SpatialScore), a pontuação do modelo Flux.1-dev aumentou de 2,18 para 7,81 após o treinamento com RL.
  • Em benchmarks de fora do domínio (DPG-Bench), o modelo treinado superou o Flux.1-dev original e o modelo treinado com GenEval, alcançando desempenho comparável ao modelo proprietário GPT-Image-1.
  • Qualidade Visual: Comparativos qualitativos mostram que o modelo treinado com SpatialScore consegue seguir prompts longos e complexos com múltiplos objetos, enquanto o modelo treinado com GenEval falha em manter a consistência espacial e introduz artefatos visuais (ex: objetos flutuantes).

5. Significado e Impacto

Este trabalho aborda uma lacuna fundamental na geração de imagem: a capacidade de raciocínio espacial composicional.

• Superação das Limitações Atuais: Demonstra que a dependência de modelos de recompensa genéricos ou regras baseadas em detectores é insuficiente para prompts complexos.
• Viabilidade do RL Online: Prova que é possível realizar RL online eficiente e de baixo custo para tarefas de geração de imagem complexas, desde que se tenha um modelo de recompensa especializado e uma estratégia de filtragem adequada.
• Futuro: Abre caminho para aplicações que exigem precisão espacial rigorosa, como design de interiores, arquitetura, e simulações para robótica (sim-to-real), onde a posição relativa dos objetos é crítica.

Em resumo, o artigo estabelece um novo paradigma onde a compreensão espacial é tratada como uma tarefa distinta e otimizável através de modelos de recompensa especializados, superando as limitações de modelos de geração de imagem atuais e de métodos de avaliação tradicionais.