Resumo Técnico: UNIGP – Domando o Diffusion Transformer para Geração e Percepção Unificadas com Preservação de Priors

Declaração do Problema

Avanços recentes em modelos de difusão alcançaram um sucesso notável tanto na geração de imagens controlável (ex: Texto-para-Imagem guiado por profundidade ou arestas) quanto em tarefas de predição densa (ex: estimativa de profundidade monocular e normais de superfície). No entanto, as abordagens existentes tipicamente tratam esses domínios como separados. Métodos de geração controlável (como o ControlNet) são limitados a tarefas de geração e frequentemente exigem conjuntos de dados massivos para um controle preciso, enquanto métodos de predição densa (como Marigold ou modelos baseados em regressão) são frequentemente ajustados (fine-tuned) a partir de backbones generativos, levando ao "esquecimento catastrófico" dos priors generativos versáteis do modelo.

Tentativas de unificação atuais (ex: JointNet, UniCon) frequentemente falham em explorar totalmente a relação complementar entre geração e percepção. Elas ou incorrem em um overhead computacional massivo ao duplicar backbones inteiros ou comprometem as capacidades generativas ao realizar o ajuste fino direto do backbone para percepção. O desafio central é unificar essas tarefas heterogêneas dentro de um único framework que preserve os priors generativos pré-treinados enquanto permite uma predição densa e geração controlável precisas, sem exigir arquiteturas ou perdas específicas de tarefa complexas.

Metodologia: Framework UNIGP

Os autores propõem o UNIGP, um framework construído sobre a arquitetia Multi-Modal Diffusion Transformer (MMDiT). O método consiste em duas inovações primárias: um ramo arquitetural desvinculado e uma estratégia de treinamento unificada.

1. Geração e Percepção Unificadas Desvinculadas (DUGP)

Seguindo o princípio da navalha de Occam, o UNIGP evita a duplicação de todo o backbone MMDiT. Em vez disso, introduz o ramo DUGP, que copia apenas o ramo de imagem ($\mathcal{F}_i$) do MMDiT pré-treinado.

• Arquitetura: O ramo DUGP é inicializado copiando os parâmetros do ramo de imagem. Ele é dividido em um Condition Patcher (Remendador de Condição), Stacked Control Layers (Camadas de Controle Empilhadas) e Stacked Perception Layers (Camadas de Percepção Empilhadas).
• Preservação de Prior (Isolamento de Gradiente): Uma escolha de design crítica é o "bypass" da conexão aditiva nas camadas de percepção. Nas camadas de controle, a saída do ramo DUGP é adicionada à saída do backbone via uma camada linear inicializada com zero ($\mathcal{L}_{zero}$), permitindo que a condição guie a geração. No entanto, nas camadas de percepção, este caminho aditivo é interrompido. Isso garante que os gradientes da perda de predição densa fluam exclusivamente através do ramo DUGP, impedindo que eles alterem os pesos pré-treinados do backbone. Isso efetivamente coloca em quarentena os priors generativos enquanto permite que o ramo DUGP aprenda a percepção como um "leitor" não invasivo.
• Tratamento de Múltiplas Condições: Para cenários que exigem múltiplas condições, o condition patcher do DUGP é duplicado para cada condição, com as interações sendo tratadas via concatenação de sequência para manter a eficiência de inferência.

2. Conjunto de Dados Unificado e Estratégia de Treinamento

Para unir a lacuna entre geração e percepção, o UNIGP emprega uma estratégia de treinamento unificada em um conjunto de dados heterogêneo combinando dados de geração filtrados e dados de percepção sintéticos.

• Pesagem de Perda Binária: Em vez de um balanceamento de perda complexo, o modelo utiliza um mecanismo simples de roteamento binário. Para cada amostra em um lote misto, os pesos de perda ($\lambda_g, \lambda_p$) são definidos como 1 ou 0 com base na fonte dos dados. Se a amostra for do conjunto de dados de geração, a perda de geração está ativa; se for do conjunto de dados de percepção, a perda de percepção está ativa.
• Objetivo: A rede aprende a prever o campo de velocidade para latentes de RGB, profundidade e normais simultaneamente, otimizando tanto a qualidade generativa quanto a precisão geométrica dentro de um único passe frontal (forward pass).

Principais Contribuições

1. Framework UNIGP: Um modelo de Diffusion Transformer unificado que integra geração controlável e predição densa. Utiliza um ramo DUGP que copia apenas o ramo de imagem do MMDiT, aderindo à navalha de Occam para minimizar o overhead computacional enquanto preserva os priors versáteis do backbone.
2. Estratégia de Treinamento Unificada: Um novo conjunto de dados e abordagem de treinamento que aprende conjuntamente tarefas heterogêneas usando um esquema simples de pesagem de perda binária, eliminando a necessidade de designs arquiteturais específicos de tarefa ou balanceamento de perda complexo.
3. Sinergia Demonstrada: O artigo fornece evidências de que o treinamento conjunto gera benefícios complementares: os priors generativos enriquecem os detalhes perceptuais (capturando texturas finas), enquanto o aprendizado perceptual impõe um alinhamento espacial estrito na geração.

Resultados Experimentais

Experimentos extensivos demonstram que o UNIGP supera abordagens unificadas anteriores e performa em nível de métodos "especialistas".

• Geração Controlável: O UNIGP supera modelos de controle único (ControlNet) e multi-controle (UniControlNet) em FID e pontuações CLIP em várias condições (profundidade, normal, esboço, pose, Canny). Notavelmente, atinge o menor RMSE em tarefas de profundidade-para-imagem e normal-para-imagem, indicando um alinhamento estrutural superior.
• Predição Densa: Em estimativa de profundidade e normais de superfície zero-shot, o UNIGP atinge um desempenho comparável a modelos de regressão de estado da arte (ex: DepthAnything V2, DSINE), apesar de ser treinado em significativamente menos imagens (59K vs. 62.6M para DepthAnything). Ele supera significativamente outros métodos de percepção generativa (Marigold, GeoWizard) e modelos unificados (JointNet, UniCon).
• Geração Conjunta: O UNIGP suporta a geração simultânea de mapas de RGB, profundidade e normais de superfície a partir de prompts de texto, superando baselines como LDM3D e JointNet tanto em qualidade de imagem (FID) quanto em consistência geométrica (métricas GD/GN).
• Estudos de Ablação:
  • Sinergia: Remover o treinamento de percepção degrada o alinhamento estrutural das imagens geradas (maior RMSE), enquanto remover o treinamento de geração resulta em mapas de profundidade excessivamente suavizados e sem detalhes finos.
  • Escolhas de Design: Substituir o DUGP por uma duplicação total do backbone (estilo JointNet) ou ajuste fino total (estilo Marigold) resulta em desempenho significativamente inferior, confirmando que o fluxo de gradiente desvinculado é a chave para o sucesso.
  • Equilíbrio de Camadas: Uma alocação igual de camadas de controle e percepção empilhadas (12:12) produz o equilíbrio ideal entre controle generativo e precisão de percepção.

Significância e Alegações

O artigo afirma que o UNIGP consegue unir os campos separados de síntese baseada em difusão e predição de geometria densa. Sua principal significância reside em demonstrar que os priors generativos e as restrições perceptuais não são mutuamente exclusivos, mas sim mutuamente reforçadores.

• Preservação de Prior: Ao isolar o fluxo de gradiente, o UNIGP retém as capacidades generativas de alta qualidade do backbone MMDiT pré-treinado, que frequentemente são perdidas no ajuste fino específico para percepção.
• Eficiência: O método atinge resultados de estado da arte com um aumento relativamente pequeno de parâmetros (de 2B para 3B) em comparação com a duplicação de todo o backbone, tornando-o uma solução unificada computacionalmente eficiente.
• Versatilidade: O framework suporta uma ampla gama de tarefas dentro de um único modelo, incluindo texto-para-imagem, estimativa de profundidade/normal e geração de qualquer condição, ofereção uma alternativa mais versátil à manutenção de múltiplos modelos especializados.

Os autores reconhecem que, embora o UNIGP introduza um overhead computacional adicional em comparação com o modelo base (aumentando o tempo de inferência de ~5s para ~9.5s), esse compromisso é justificado pela capacidade de substituir múltiplas arquiteturas especializadas por um único modelo unificado altamente versátil.