FA-Seg: A Fast and Accurate Diffusion-Based Method for Open-Vocabulary Segmentation

1. Declaração do Problema

Segmentação Semântica de Vocabulário Aberto (OVSS) visa segmentar objetos em imagens com base em categorias de texto arbitrárias, sem exigir anotações densas em nível de pixel para cada classe possível.

• Limitações dos Métodos Existentes:
  • Modelos Visão-Linguagem (e.g., CLIP): Embora eficazes para reconhecimento zero-shot, eles dependem de embeddings globais imagem-texto, resultando frequentemente em baixa precisão espacial e limites difusos em nível de pixel.
  • Modelos Baseados em Difusão: Embora codifiquem naturalmente características espaciais de alta granularidade por meio de mecanismos de atenção, os métodos existentes frequentemente sofrem com altos custos computacionais. Muitos exigem múltiplos passos de inferência por classe, otimização extensiva ou geração de dados sintéticos, criando um compromisso entre qualidade de segmentação e eficiência de inferência.
• A Lacuna: Há uma necessidade crítica de um método OVSS livre de treinamento que alcance alta precisão de segmentação (precisão espacial de alta granularidade) enquanto mantém eficiência computacional adequada para implantação no mundo real.

2. Metodologia: FA-Seg

O FA-Seg é um framework livre de treinamento que aproveita modelos de difusão imagem-para-texto pré-treinados (especificamente Stable Diffusion) para realizar segmentação em uma única execução de inferência. O pipeline consiste em três etapas principais:

A. Construção de Duplo Prompt

Para permitir tanto a reconstrução fiel da imagem quanto a extração de atenção específica de classe, o FA-Seg constrói dois prompts complementares:

1. Prompt de Legenda ($\mathcal{P}$): Gerado por um modelo de geração de legendas de imagem (e.g., BLIP) para guiar a inversão DDIM na reconstrução da imagem de entrada.
2. Prompt de Classe ($\mathcal{P}'$): Uma lista de categorias candidatas de objetos (gerada via TagCLIP + BLIP) anexada à legenda.

• Mecanismo: O modelo de difusão utiliza $\mathcal{P}$ para reconstrução, mas utiliza $\mathcal{P}'$ para gerar mapas de atenção cruzada especificamente para as classes candidatas. Isso permite que o modelo foque em conceitos semânticos específicos sem re-inferência iterativa.

B. Inversão Eficiente e Extração de Atenção

• Inversão Rápida: Em vez da inversão DDIM padrão de múltiplos passos, o FA-Seg emprega 2-Rectified Flow (2-RF). Este modelo destilado permite reconstrução precisa de imagem em um processo de (1+1)-passo (um passo para frente, um passo para trás), reduzindo drasticamente o tempo de inferência.
• Extração de Atenção: Durante o processo de remoção de ruído, o método extrai:
  • Mapas de Atenção Cruzada: Ligando tokens de texto (do prompt de classe) a regiões da imagem.
  • Mapas de Atenção Self: Capturando a estrutura espacial interna e similaridades em pares da imagem.
  • Estes são extraídos em múltiplas resoluções ($r \in \{8, 16, 32, 64\}$).

C. Método de Refinamento Hierárquico de Atenção (HARD)

Para produzir máscaras de alta qualidade, o FA-Seg funde mapas de atenção de diferentes resoluções:

1. Fusão de Atenção Cruzada: Mapas de atenção cruzada de diferentes resoluções são aumentados e agregados usando uma soma ponderada. Resoluções mais baixas (e.g., 16) recebem pesos maiores para contexto semântico, enquanto resoluções mais altas são ponderadas para reduzir ruído.
2. Refinamento de Atenção Self: O mapa de atenção cruzada fundido é refinado adicionalmente usando matrizes de afinidade de atenção self. Mapas de atenção self fornecem detalhes estruturais (limites e formas) que a atenção cruzada carece. O método transforma tensores de atenção self para corresponder à resolução do mapa de atenção cruzada e aplica-os como pesos multiplicativos para afiar os limites dos objetos.
3. Viragem no Tempo de Teste (TTF): Para melhorar a consistência espacial, o processo é repetido em uma versão horizontalmente flipada da imagem de entrada. Os mapas de atenção resultantes são realinhados e médios com os mapas originais para reduzir ruído e melhorar a robustez.

3. Principais Contribuições

1. Inferência Unificada em Execução Única: Ao contrário de métodos anteriores que exigem execuções de inferência separadas para cada classe, o FA-Seg gera máscaras de segmentação para todas as classes candidatas simultaneamente em uma única inferência de (1+1)-passo.
2. Mecanismo de Duplo Prompt: Uma estratégia inovadora que separa o prompt de reconstrução do prompt consciente de classe, permitindo que o modelo de difusão extraia mapas de atenção discriminativos para categorias específicas sem perder fidelidade da imagem.
3. Refinamento Hierárquico de Atenção (HARD): Uma estratégia de fusão multi-escala que combina atenção cruzada (para localização semântica) e atenção self (para refinamento estrutural) através de resoluções para alcançar limites precisos de objetos.
4. Eficiência Livre de Treinamento: O método não requer ajuste fino, não requer geração de dados sintéticos e utiliza recursos computacionais mínimos, alcançando desempenho state-of-the-art com menos de 1 segundo de tempo de inferência.

4. Resultados Experimentais

Os autores avaliaram o FA-Seg em três benchmarks padrão: PASCAL VOC, PASCAL Context e COCO Object.

• Desempenho (mIoU):
  • O FA-Seg alcançou um mIoU médio de 43,8% em todos os três conjuntos de dados.
  • Estabeleceu um novo state-of-the-art (SOTA) para métodos baseados em difusão livres de treinamento, superando concorrentes como InvSeg (42,4%), DiffSegmentor (41,8%) e várias abordagens baseadas em CLIP.
  • Pontuações específicas: 64,0% (PASCAL VOC), 31,3% (PASCAL Context), 36,2% (COCO).
• Eficiência:
  • Tempo de Inferência: 0,36 segundos por imagem em uma GPU RTX 4090. Isso é significativamente mais rápido que o InvSeg (7,9s em H100) e o DiffSegmentor (tempo escala com o número de classes).
  • Memória: Consome apenas 13,4 GB de VRAM, comparado a 32,4 GB para o InvSeg.
• Estudos de Ablação:
  • A remoção do refinamento de atenção self causou a maior queda de desempenho (-15,1%), provando seu papel crítico na precisão espacial.
  • A fusão ponderada de mapas multi-resolução superou a fusão por média simples.
  • A combinação TagCLIP+BLIP para geração de classes candidatas provou ser a mais robusta.

5. Significado e Impacto

• Ponteando a Lacuna: O FA-Seg conecta com sucesso a lacuna entre qualidade de segmentação e eficiência de inferência, demonstrando que segmentação de vocabulário aberto de alta precisão não requer treinamento pesado ou otimização multi-passo.
• Implantação Prática: Ao reduzir o tempo de inferência para menos de um segundo e o uso de memória para ~13GB, o FA-Seg torna a segmentação de vocabulário aberto em tempo real viável para aplicações do mundo real (e.g., robótica, direção autônoma, ferramentas interativas).
• Fundação para Trabalhos Futuros: O artigo estabelece uma linha de base forte para segmentação por difusão livre de treinamento, sugerindo que melhorias futuras poderiam focar em geração adaptativa de candidatos e segmentação em nível de instância.

Em conclusão, o FA-Seg representa um avanço significativo no campo da segmentação de vocabulário aberto, aproveitando as capacidades espaciais inerentes dos modelos de difusão através de um framework altamente otimizado, livre de treinamento e computacionalmente eficiente.