Resumo Técnico: STAR (SpatioTemporal Adaptive Reward Allocation)

Declaração do Problema

Os métodos existentes de pós-treinamento de Aprendizado por Reforço (RL) para geração de texto-para-imagem (T2I), como o Flow-GRPO, tipicamente convertem a recompensa final ao nível da imagem em um único escalar de vantagem. Este escalar é então aplicado com força uniforme em toda a trajetória generativa (todos os passos de denoising e todos os tokens latentes espaciais).

O artigo identifica uma desconexão de granularidade fundamental nesta abordagem:

1. Estrutura Temporal: Diferentes passos de denoising são responsáveis por distintas etapas de geração (ex: layout global vs. refinamento de detalhes).
2. Estrutura Espacial: O conteúdo que determina o alinhamento com o texto (ex: objetos específicos, atributos ou relações espaciais) frequentemente aparece apenas em regiões específicas da imagem, enquanto outras regiões (ex: fundo) contribuem menos para os requisitos semânticos do prompt.

Quando uma recompensa escalar é aplicada uniformemente, a atualização da política tem dificuldade em distinguir entre regiões que são críticas para satisfazer o prompt e aquelas que não são. Isso torna difícil para o modelo focar o aprendizado nos componentes generativos específicos que realmente impulsionam a recompensa, levando a um desempenho subótimo em generalização composicional, contagem de objetos e renderização de texto.

Metodologia: STAR

Os autores propõem a Alocação de Recompensa Adaptativa Espaço-Temporal (STAR), um framework projetado para realocar recompensas ao nível da imagem através da trajetória de denoising estruturada de modelos de difusão e fluxo. O STAR opera sem modificar a fonte de recompensa externa, aproveitando, em vez disso, a mecânica interna do modelo generativo.

1. Alocação Guiada por Atenção

O STAR constrói mapas de alocação espacial dinâmica que variam através dos passos de denoising ($t$) e rollouts ($i$). O processo envolve:

• Extração de Unidades de Texto: O prompt $c$ é decomposto em um conjunto de unidades de texto essenciais $K(c)$ (ex: sujeitos, estilos, atributos).
• Mapeamento de Atenção: O modelo utiliza camadas de atenção texto-imagem (especificamente do backbone de atenção conjunta do Stable Diffusion 3.5) para extrair sinais de roteamento. Para cada posição espacial latente $u$ e token de texto $j$, um peso de atenção $\alpha(u, j)$ é computado.
• Construção de Heatmap: Para cada unidade de texto, um mapa espacial é gerado agregando pesos de atenção sobre os tokens de texto correspondentes. Esses mapas são agregados em um heatmap específico da trajetória $H_t^i$ que destaca as regiões latentes mais ativas em relação aos componentes principais do prompt.
• Normalização: O heatmap é normalizado e escalonado para produzir um mapa de coeficientes $C_t^i(u)$. Este mapa atribui pesos mais altos às regiões relevantes ao prompt e pesos menores (mas não nulos) às regiões de fundo.

2. Objetivo de Política com Resolução Espacial

O STAR converte a vantagem escalar relativa ao grupo $A_i$ em um sinal de treinamento resolvido espacialmente $A_t^i(u) = A_i \cdot C_t^i(u)$.

• O objetivo da política é modificado para operar em contribuições de log-probabilidade espacial em vez de uma probabilidade de trajetória escalar.
• Uma razão espacial normalizada $\rho$ é computada para cada posição latente.
• A função objetivo final (baseada em clipping estilo PPO) soma as razões clipadas ponderadas pelo mapa de vantagem espacial $A_t^i(u)$ através de todos os passos de tempo e localizações espaciais. Isso garante que as atualizações de política sejam mais fortes em regiões correspondentes aos componentes principais do prompt.

Principais Contribuições

1. Framework STAR: Um método para pós-treinamento de RL de difusão/fluxo T2I que converte recompensas ao nível da imagem em sinais de treinamento distribuídos sobre a trajetória de denoising e o espaço latente, permitindo que as atualizações foquem em regiões generativas correspondentes a objetos, atributos e relações chave.
2. Interface Unificada: O uso de interações internas texto-imagem para construir mapas de alocação espacial específicos de cada passo de tempo, estendendo objetivos de RL existentes para um objetivo de política com resolução espacial sem exigir novos modelos de recompensa.
3. Validação Empírica: Demonstração de que o STAR melhora o alinhamento semântico composicional, a renderização de texto e a otimização de preferência em três tarefas distintas (GenEval, OCR, PickScore) usando o Stable Diffusion 3.5 Medium como base.

Resultados Experimentais

Os autores avaliaram o STAR no Stable Diffusion 3.5 Medium (SD3.5-M) comparando-o com o baseline Flow-GRPO e outros modelos de estado da arte.

• GenEval (Geração Composicional): O STAR alcançou uma pontuação geral de 0.9759, superando o Flow-GRPO (0.95) e o SD3.5-M base (0.63). Melhorias notáveis foram vistas em contagem (0.98), cores (0.96) e vinculação de atributos (0.94).
• Renderização de Texto OCR: O STAR alcançou uma precisão de 0.9757, superando o baseline Flow-GRPO (0.92).
• PickScore (Otimização de Preferência): O STAR alcançou uma pontuação de 23.60, comparado a 23.31 para o Flow-GRPO.
• Eficiência: O método introduz um overhead computacional mínimo, adicionando menos de 2% ao tempo de treinamento e aumentando o uso de memória de pico em apenas 1.56%, pois reutiliza os mapas de atenção da passagem forward.

Significância e Alegações

O artigo afirma que o STAR aborda o problema de "atribuição de crédito" no RL de T2I ao reconhecer a natureza espaço-temporal da geração por difusão. Ao rotear o mesmo sinal de recompensa escalar para as regiões latentes específicas que determinam o alinhamento com o texto, o STAR permite atualizações de política mais seletivas e eficazes.

Os autores enfatizam que o STAR alcança esses ganhos sem alterar a fonte de recompensa externa ou exigir novos modelos de recompensa. Ele fornece uma "interface simples, leve e eficaz" para alocar feedback ao nível da imagem para trajetórias estruturadas. Os resultados sugerem que a alocação adaptativa espaço-temporal é um fator crítico para otimizar a compreensão semântica fina, o controle de conteúdo local e o alinhamento de preferência em modelos generativos.