Diffusion Model as a Generalist Segmentation Learner

1. Declaração do Problema

Os sistemas de segmentação atuais são altamente fragmentados. Os modelos são tipicamente adaptados para tarefas específicas (semântica, instância, panóptica), vocabulários específicos (conjunto fechado vs. vocabulário aberto) ou domínios específicos (natural, médico, sensoriamento remoto). A transição entre essas configurações frequentemente exige o redesenho de arquiteturas ou o retreinamento a partir do zero.

Embora pesquisas recentes tenham tentado reaproveitar modelos de difusão texto-para-imagem pré-treinados para segmentação, extraindo mapas de atenção, essas abordagens sofrem de limitações significativas:

• Ruído e Inconsistência: Mapas de atenção brutos são frequentemente de baixa resolução, ruidosos e inconsistentes entre as camadas.
• Dependência de Pós-processamento: Frequentemente exigem pesado pós-processamento heurístico para produzir máscaras utilizáveis.
• Falta de Generalização: Métodos existentes frequentemente focam em uma única tarefa e falham em demonstrar generalização robusta através de domínios diversos e consultas de vocabulário aberto sem customização arquitetural.

A questão central abordada por este artigo é: Um único modelo unificado baseado em difusão pode servir como um aprendiz de segmentação generalista capaz de lidar com tarefas semânticas, de vocabulário aberto e de domínio cruzado sem mudanças arquitetural específicas de domínio?

2. Metodologia: Framework DiGSeg

Os autores propõem o DiGSeg, um framework que transforma um modelo de difusão pré-treinado (especificamente Stable Diffusion v2) de um prior generativo em um aprendiz de segmentação explícito por meio de um protocolo de ajuste fino unificado.

Arquitetura Central

O framework consiste em três componentes principais:

1. Caminho Latente Visual:

   • Em vez de tratar a imagem como pixels brutos, a imagem RGB de entrada ($x$) e a máscara de segmentação ground-truth ($y$) são codificadas em um espaço latente compacto usando o encoder VAE congelado do Stable Diffusion.
   • Para manter compatibilidade com o VAE de 3 canais, a máscara de canal único é replicada em três canais.
   • Isso preserva a estrutura espacial enquanto permite o aprendizado eficiente de correspondências de alta granularidade no espaço latente.

2. Condicionamento de Texto Alinhado ao CLIP:

   • Para habilitar capacidades de vocabulário aberto, um encoder de texto CLIP congelado gera embeddings ($t_{clip}$) para nomes de classes ou descrições em linguagem natural.
   • Esses embeddings são injetados no U-Net de difusão via atenção cruzada em múltiplas escalas de remoção de ruído. Isso permite que o modelo alinhe consultas textuais com representações visuais em evolução global e localmente.

3. U-Net de Remoção de Ruído Consistente com Segmentação:

   • Objetivo de Treinamento: O modelo é treinado para prever o ruído adicionado ao latente de segmentação ($z_y$), condicionado ao latente de imagem ($z_x$) e ao embedding de texto ($t_{clip}$).
   • Formulação Generativa: Diferentemente de trabalhos anteriores que usam difusão como extrator de características, o DiGSeg trata a segmentação como uma tarefa de geração condicional de remoção de ruído: $p_\theta(y|x)$. O U-Net aprende a remover ruído de ruído aleatório para uma representação latente consistente com segmentação.
   • Estratégia de Ruído: É empregado um cronograma de ruído multi-escala annealed (reduzido gradualmente). Perturbações de alta frequência dominam os passos iniciais para aprender estruturas finas, enquanto componentes de baixa frequência tornam-se influentes posteriormente para garantir coerência semântica.

Processo de Inferência

1. A imagem de entrada é codificada em uma representação latente.
2. Um latente de segmentação é inicializado como ruído gaussiano.
3. Sob condicionamento de texto, o U-Net remove ruído iterativamente do latente (usando um amostrador DDIM com uma estratégia de passo de tempo "trailing" para eficiência).
4. O latente limpo final é decodificado de volta ao espaço de pixels usando o decoder VAE congelado para produzir a máscara de segmentação.
5. Ensemble em Tempo de Teste: Múltiplas passagens de inferência com diferentes sementes de ruído são médias no espaço latente para melhorar a consistência espacial e reduzir artefatos.

3. Contribuições Principais

• Framework Unificado: O DiGSeg é o primeiro framework a reaproveitar com sucesso um backbone de difusão pré-treinado em um aprendiz de segmentação generalista que lida com tarefas de vocabulário fechado, vocabulário aberto e de domínio cruzado sem mudanças arquitetural específicas de tarefa.
• Aprendizado Latente Direto: Em vez de depender de mapas de atenção ruidosos, o método ajusta finamente o U-Net de remoção de ruído para gerar diretamente latentes consistentes com segmentação, produzindo máscaras de alta qualidade sem pesado pós-processamento.
• Condicionamento de Texto Multi-escala: Introduz um mecanismo para injetar características de linguagem através de múltiplas escalas de remoção de ruído, permitindo alinhamento robusto entre prompts de texto arbitrários e características visuais.
• Estratégia de Ruído Novel: Propõe um cronograma de ruído multi-escala annealed que equilibra o aprendizado de limites de alta granularidade com a coerência semântica global.

4. Resultados Experimentais

Os autores avaliaram o DiGSeg em benchmarks padrão e domínios especializados, demonstrando desempenho State-of-the-Art (SOTA).

• Segmentação de Vocabulário Aberto:
  • Superou 13 baselines (incluindo ODISE, OVSeg, CAT-Seg) em A-847, PC-459, A-150, PC-59 e Cityscapes.
  • Alcançou ganhos significativos (por exemplo, +2,8 mIoU em PC-59, +2,2 em PC-459) sobre os métodos de segundo melhor.
• Segmentação Semântica Padrão:
  • Alcançou SOTA em COCO (50,8 mIoU) e ADE20K (58,6 mIoU), superando modelos fortes baseados em transformadores como Mask2Former e OneFormer.
• Transferência de Domínio Cruzado:
  • Demonstrou forte transferência zero-shot para conjuntos de dados médicos (REFUGE-2), agrícolas (PhenoBench) e de sensoriamento remoto (DeepGlobe) sem modificação arquitetural.
  • Nota: O desempenho em domínios médicos foi menor do que em cenas naturais, atribuído à lacuna de domínio nos dados de pré-treinamento do CLIP, mas ainda competitivo.
• Eficiência e Uso de Dados:
  • O modelo é altamente eficiente em dados, alcançando desempenho quase idêntico com apenas 50% dos dados de treinamento.
  • A velocidade de inferência é melhorada via DDIM-trailing (atingindo saturação em 1 passo), embora permaneça mais lenta do que modelos feed-forward devido à natureza da difusão.

5. Significado

Este trabalho fundamentalmente une a lacuna entre geração visual e compreensão visual.

• Mudança de Paradigma: Desafia a noção de que modelos de difusão são exclusivamente para síntese, provando que podem servir como aprendizes discriminativos generalistas poderosos quando devidamente condicionados.
• Generalização: Demonstra que os priores espaciais e semânticos ricos codificados em modelos generativos de grande escala são transferíveis para tarefas de predição densa através de domínios diversos (natural, médico, agrícola).
• Simplicidade: Ao evitar pipelines multi-estágio complexos ou cabeças especializadas, o DiGSeg oferece uma interface unificada e simplificada para segmentação, sugerindo um futuro onde um único backbone generativo pode lidar com uma ampla gama de tarefas de visão computacional.

Em conclusão, o DiGSeg estabelece que backbones de difusão modernos, quando ajustados finamente com um protocolo condicionado a latentes e alinhamento de texto, podem superar modelos de segmentação especializados enquanto oferecem flexibilidade superior em cenários de vocabulário aberto e de domínio cruzado.