OmniAlpha: Aligning Transparency-Aware Generation via Multi-Task Unified Reinforcement Learning

1. Declaração do Problema

Os modelos atuais de IA generativa, particularmente os sistemas baseados em difusão, excelam na criação de imagens RGB de alta fidelidade. No entanto, muitos fluxos de trabalho visuais do mundo real (por exemplo, efeitos visuais, design gráfico, edição de imagens) exigem geração consciente de transparência (RGBA), que envolve modelar não apenas a cor, mas também a opacidade por pixel (canais alfa) e a composição entre camadas.

As soluções existentes sofrem com duas limitações principais:

• Fragmentação: Os métodos atuais são silos específicos de tarefas. Existem pipelines separados para recorte de imagem (matting), remoção de objetos, decomposição de camadas e geração de texto para imagem. Não há um modelo unificado capaz de lidar com fluxos de trabalho RGBA diversos.
• Limitações do Ajuste Fino Supervisionado (SFT): Embora escalar o SFT entre tarefas seja um passo lógico para a unificação, ele depende de objetivos de regressão localizados (por exemplo, perda L1/L2 pixel a pixel). Esses objetivos falham em otimizar diretamente propriedades críticas da geração consciente de transparência, como:
  • Consistência e física entre camadas.
  • Precisão nas fronteiras do alfa (por exemplo, cabelo, pelos).
  • Fidedignidade composicional quando as camadas são mescladas.
  • Consequentemente, modelos baseados apenas em SFT frequentemente estagnam, lutando para capturar detalhes estruturais finos necessários para uma transparência perfeita.

2. Metodologia: OMNIALPHA

Os autores propõem o OMNIALPHA, um framework unificado multi-tarefa que migra do aprendizado puramente supervisionado para o alinhamento via Aprendizado por Reforço (RL). A arquitetura consiste em três componentes principais:

A. Arquitetura Unificada

• VAE Consciente de Alfa: Um Variational Autoencoder (VAE) de ponta a ponta inicializado a partir de um VAE RGB pré-treinado usando uma estratégia de "inicialização opaca". O codificador zera os pesos do alfa inicialmente para focar no RGB, enquanto o decodificador é inicializado para prever um canal alfa totalmente opaco ($\alpha=1.0$). Isso permite uma adaptação eficiente ao espaço RGBA de 4 canais.
• Transformador de Difusão Sequência-a-Sequência (DiT): A espinha dorsal de remoção de ruído processa múltiplas entradas e saídas simultaneamente.
  • Eixo de Camada Bi-direcional: Para lidar com múltiplas imagens (entradas e alvos) em uma única passagem direta, os autores estendem os Embeddings de Posição Rotacional 2D padrão (RoPE) com um eixo z bi-direcional.
  • Os latentes de entrada recebem índices não negativos ($0$a$n-1$), os latentes de alvo recebem índices negativos ($-1$a$-m$) e os embeddings de texto recebem índices positivos distintos. Isso permite que o modelo diferencie as camadas espacialmente dentro da sequência.

B. Paradigma de Treinamento

O treinamento segue um processo de duas etapas:

1. Início Frio (SFT Multi-Tarefa): O modelo é primeiro treinado em um conjunto de dados unificado cobrindo cinco categorias de tarefas usando objetivos de Flow Matching. Isso estabelece uma capacidade básica em diversas tarefas.
2. Pós-Treinamento (Otimização de Política Relativa por Grupo - GRPO): Para superar as limitações do SFT, os autores formulam o processo de geração como um Processo de Decisão de Markov (MDP).
   • Estado: Condições multimodais, passo de tempo e sequência latente intermediária.
   • Ação: O passo de remoção de ruído guiado pelo campo de velocidade previsto.
   • Recompensa: Uma Recompensa Consciente de Camada é computada apenas no final da trajetória (baseada em resultado).

C. Design de Recompensa Consciente de Camada

Em vez de recompensas específicas de tarefas, o OMNIALPHA utiliza quatro objetivos universais definidos na sequência RGBA decodificada para guiar a política de RL:

1. Fidelidade da Camada ($r_{layer}$): Mede a similaridade RGB (SSIM/PSNR) em todas as camadas válidas.
2. Fidelidade da Composição ($r_{comp}$): Avalia a imagem composta globalmente contra a verdade fundamental (ground truth).
3. Estrutura de Fundo ($r_{bg}$): Penaliza desvios estruturais nas camadas de fundo usando LPIPS.
4. Fronteira de Primeiro Plano ($r_{fg}$): Impõe precisão e recall nas fronteiras do alfa (crucial para cabelo/pelos) usando máscaras de fronteira suave.

Essas recompensas são agregadas dinamicamente com base na tarefa específica (por exemplo, texto-para-imagem ignora recompensas alinhadas a pixels, pois a verdade fundamental não está disponível, enquanto tarefas de recorte ponderam fortemente as recompensas de fronteira).

3. Contribuições Principais

1. Primeiro Framework Unificado de RL para RGBA: O OMNIALPHA é o primeiro framework a aplicar aprendizado por reforço multi-tarefa para unificar a geração e manipulação conscientes de transparência, indo além de pipelines fragmentados e específicos de tarefas.
2. Paradigma de Recompensa Estrutural: Os autores introduzem recompensas fundamentadas estruturalmente que otimizam explicitamente a coerência entre camadas e a precisão da fronteira alfa, transcendendo com sucesso o teto de desempenho do SFT padrão.
3. Codificação de Camada Bi-direcional: Uma extensão novel de embeddings de posição que permite que um único DiT modele conjuntamente múltiplas camadas RGBA de entrada e saída em uma única passagem direta.

4. Resultados Experimentais

O modelo foi avaliado em cinco categorias de tarefas: Texto-para-Imagem, Remoção de Objetos, Recorte Automático, Recorte Referenciado e Decomposição de Camadas.

• Desempenho Geral: O OMNIALPHA supera consistentemente sua própria linha de base SFT e modelos especialistas especializados em diversos benchmarks.
• Métricas Chave:
  • Decomposição de Camadas: Alcançou uma redução relativa de 9,07% no L1 RGB em comparação com a linha de base SFT.
  • Recorte Automático: Superou ferramentas convencionais com uma melhoria de 74% no SAD (Soma das Diferenças Absolutas) e 68% de melhoria no Grad (Diferença de Gradiente).
  • Recorte Referenciado: Alcançou resultados state-of-the-art no benchmark RefMatte-RW100 (SAD: 14,768), superando significativamente modelos especializados como o CLIPMat.
  • Remoção de Objetos: Demonstrou forte generalização em conjuntos de dados fora do domínio (OBER-Wild), alcançando pontuações ReMOVE comparáveis ao modelo altamente especializado ObjectClear.
• Estudos de Ablação: A remoção de qualquer componente único de recompensa ($r_{layer}, r_{comp}, r_{bg}, r_{fg}$) resultou em degradação de desempenho, confirmando que o conjunto completo de recompensas é necessário para um alinhamento robusto multi-tarefa.

5. Significado

O OMNIALPHA representa uma mudança de paradigma significativa na IA generativa para edição visual. Ao demonstrar que o Aprendizado por Reforço pode alinhar efetivamente modelos de difusão para tarefas complexas de RGBA multi-camada, o artigo preenche a lacuna entre a geração RGB opaca e os requisitos nuances dos fluxos de trabalho visuais profissionais.

Isso sugere que os futuros modelos fundacionais para edição de imagens devem afastar-se do ajuste fino isolado e específico de tarefas em direção a políticas unificadas otimizadas via recompensas baseadas em resultado. Essa abordagem não apenas melhora a fidelidade dos canais alfa e da composição de camadas, mas também permite que um único modelo lide com um amplo espectro de tarefas criativas, desde a geração de objetos transparentes até a decomposição de cenas complexas.