Resumo Técnico: Edit-R2 – Aprendizado por Reforço Sensível ao Contexto para Edição de Imagens Multi-turno

Definição do Problema

Os sistemas atuais de edição de imagens guiados por texto, impulsionados por modelos de difusão e modelos fundacionais multimodais unificados, operam predominantemente em cenários de turno único. Nesses cenários, um modelo recebe uma imagem de entrada e uma única instrução para produzir uma saída editada. Essa formulação falha em capturar fluxos de trabalho realistas de edição in-context multi-turno, onde os usuários refinam iterativamente uma imagem através de uma sequência de instruções.

Em tais sessões multi-turno, um modelo enfrenta dois modos de falha acoplados:

1. Diluição de Contexto Longo: À medida que a sessão progride, o histórico intercalado de imagens e textos cresce. Restrições textuais esparsas de turnos iniciais tornam-se difíceis de recuperar devido à decadência da atenção e ao pesado orçamento de tokens visuais necessário para codificar as imagens geradas.
2. Contaminação de Estado: Erros cometidos em turnos iniciais (por exemplo, posicionamento incorreto de objetos ou modificação de atributos) são carregados adiante, degradando as gerações subsequentes e tornando o rollout de treinamento não confiável.

Benchmarks e paradigmas de treinamento existentes são amplamente agnósticos ao histórico, tratando cada turno como uma tarefa isolada dependente apenas da imagem imediatamente precedente e da instrução atual, ignorando restrições acumuladas em nível de sessão.

Metodologia: Edit-R2

Os autores propõem o Edit-R2, um novo framework de pós-treinamento de aprendizado por reforço (RL) projetado para modelos multimodais unificados (especificamente construído sobre o BAGEL). O Edit-R2 formula a edição in-context multi-turno como um Processo de Decisão de Markov (MDP) em nível de sessão e introduz três mecanismos principais para lidar com a perda de contexto e a contaminação de estado:

1. Reconstrução do Estado da Sessão via Cadeia de Pensamento In-Context (IC-CoT)

Para mitigar a diluição de contexto longo, o Edit-R2 não passa o histórico intercalado bruto diretamente para o gerador. Em vez disso, ele emprega um módulo IC-CoT (um "especialista em compreensão") que realiza o raciocínio sobre todo o histórico da sessão antes de cada etapa de geração.

• Função: O módulo reconstrói a intenção operacional da sessão, consolidando restrições históricas dispersas (ex: regras globais como "todos os objetos adicionados devem ser roxos") e resolvendo pronomes (ex: "mude a cor dele") em uma instrução explícita e aprimorada ($\tilde{T}_t$).
• Benefício: Isso cria um rastro de raciocínio compacto que encurta a distância entre o módulo de geração e a informação crítica de que ele necessita, fornecendo uma representação de estado estável.

2. Objetivo de RL Multi-Turno Unificado

O Edit-R2 otimiza conjuntamente o rastro de raciocínio IC-CoT e o processo de geração visual.

• Espaços Acoplados: O framework acopla o espaço textual discreto da reconstrução de intenção com o espaço latente contínuo do gerador de imagem de flow-matching.
• Estrutura de Recompensa: Uma função de recompensa composta ($r_t$) é definida em cada turno, combinando:
  • Seguimento de Instrução (IF - Instruction Following): O modelo executou a instrução atual?
  • Consistência de Conteúdo (CC - Content Consistency): O conteúdo não-alvo foi preservado?
  • Consciência Global (GA - Global Awareness): O output adere às restrições de nível de sessão estabelecidas em turnos anteriores?
• Otimização: A política é atualizada usando Unified-GRPO, que trata a sequência de tokens de CoT e a trajetória de flow como uma decisão conjunta, garantindo que o módulo de raciocínio seja otimizado especificamente para as capacidades de geração visual subsequentes.

3. Filtragem de Trajetória (Refinamento de Vantagem Válida de Prefixo)

Para abordar a contaminação de estado, o Edit-R2 implementa um mecanismo de filtragem de trajetória.

• Mecanismo: Se a recompensa para um turno específico cair abaixo de um limiar ($\delta$), o rollout é terminado imediatamente. Os turnos subsequentes naquela trajetória são descartados.
• Refinamento de Vantagem: As estimativas de vantagem para atualizações de política são normalizadas exclusivamente sobre rollouts ativos (aqueles que não foram terminados). Isso evita que trajetórias corrompidas e de baixa qualidade distorçam a linha de base de normalização e contaminem o sinal de treinamento.

Benchmark: MICE-Bench

Para permitir uma avaliação sistemática, os autores introduzem o MICE-Bench (Multi-turn In-Context Editing Benchmark), um dataset de larga escala composto por 720 instâncias multi-turno (3 turnos cada).

• Tarefas:
  • Memória de Conteúdo: Testa a adesão a restrições de nível de sessão (ex: "todos os objetos adicionados devem ser vermelhos").
  • Compreensão de Conteúdo: Testa a resolução de correferências pronominais entre turnos (ex: "remova-o").
• Métricas: Introduz a Consciência Global (GA) como uma nova métrica para quantificar a conformidade com as restrições acumuladas da sessão, juntamente com o Seguimento de Instrução (IF) e a Consistência de Conteúdo (CC) padrão.

Principais Resultados

Experimentos no MICE-Bench demonstram que o Edit-R2 supera substancialmente os baselines de código aberto e proprietários:

• Ganhos de Desempenho: Comparado ao modelo base BAGEL, o Edit-R2 alcança melhorias significativas, particularmente nos turnos tardios, onde a perda de contexto é mais severa:
  • Turno 2: +18% IF, +18% GA.
  • Turno 3: +8% IF, +15% GA.
• Estudos de Ablação:
  • Remover o IC-CoT leva à degradação do desempenho no Turno 2 e 3, confirmando seu papel na mitigação da diluição de contexto.
  • Remover o Text-GRPO (otimização conjunta) reduz os ganhos, destacando a necessidade de otimizar os módulos de raciocínio e geração juntos.
  • Remover o Refinamento de Vantagem Válida de Prefixo resulta em um desempenho geral inferior, validando a importância de filtrar rollouts corrompidos.
• Análise Qualitativa: Visualizações de atenção mostram que o Edit-R2 mantém uma atenção fortemente focada nos tokens de objetos chave de turnos anteriores, enquanto modelos baselines (como o BAGEL) exibem atenção difusa e mal direcionada devido à decadência de distância do RoPE.

Significância e Alegações

O artigo afirma que o Edit-R2 representa um passo significativo para tornar os sistemas de edição de imagem capazes de consistência em nível de sessão.

• Mudança de Paradigma: Ele vai além de edições locais isoladas para um problema de RL em nível de sessão, reconhecendo que o valor de uma edição depende de sua capacidade de preservar o contexto para turnos futuros.
• Primeiro de seu Tipo: O Edit-R2 é apresentado como um dos primeiros frameworks de pós-treinamento de RL projetados especificamente para edição in-context, abordando os desafios únicos de dependências de longo horizonte e contaminação de estado.
• Padrão de Avaliação: Ao introduzir o MICE-Bench e a métrica GA, o trabalho fornece o primeiro protocolo automatizado e sistemático para avaliar a edição in-context multi-turno, uma capacidade anteriormente pouco explorada tanto em modelos quanto em datasets.

Os autores concluem que, embora o MICE-Bench cubra dois cenários representativos, o framework lança as bases para a exploração futura de intenções de usuários mais amplas em edições interativas do mundo real. Eles reconhecem os riscos da geração de imagens controlável (ex: desinformação) e se comprometem a lançar diretrizes de uso junto com os pesos do modelo.