Diffusion Model as a Generalist Segmentation Learner

A continuación se presenta un resumen técnico detallado del artículo "Diffusion Model as a Generalist Segmentation Learner" (DiGSeg).

1. Planteamiento del Problema

Los sistemas de segmentación actuales están altamente fragmentados. Los modelos suelen estar adaptados a tareas específicas (semántica, de instancias, panóptica), vocabularios específicos (conjunto cerrado frente a vocabulario abierto) o dominios específicos (natural, médico, teledetección). La transición entre estos entornos a menudo requiere rediseñar arquitecturas o volver a entrenar desde cero.

Aunque investigaciones recientes han intentado reaprovechar modelos de difusión preentrenados de texto a imagen para la segmentación extrayendo mapas de atención, estos enfoques adolecen de limitaciones significativas:

• Ruido e Inconsistencia: Los mapas de atención crudos suelen tener baja resolución, ser ruidosos e inconsistentes entre capas.
• Dependencia de Postprocesamiento: Frecuentemente requieren un pesado postprocesamiento heurístico para producir máscaras utilizables.
• Falta de Generalización: Los métodos existentes suelen centrarse en una sola tarea y fallan al demostrar una generalización robusta a través de diversos dominios y consultas de vocabulario abierto sin personalización arquitectónica.

La pregunta central abordada por este artículo es: ¿Puede un único modelo unificado basado en difusión servir como un aprendiz generalista de segmentación capaz de manejar tareas semánticas, de vocabulario abierto y de dominio cruzado sin cambios arquitectónicos específicos del dominio?

2. Metodología: Marco DiGSeg

Los autores proponen DiGSeg, un marco que transforma un modelo de difusión preentrenado (específicamente Stable Diffusion v2) de un prior generativo a un aprendiz de segmentación explícito mediante un protocolo de ajuste fino unificado.

Arquitectura Central

El marco consta de tres componentes principales:

1. Ruta Latente Visual:

   • En lugar de tratar la imagen como píxeles crudos, la imagen de entrada RGB ($x$) y la máscara de segmentación de verdad terreno ($y$) se codifican en un espacio latente compacto utilizando el codificador VAE congelado de Stable Diffusion.
   • Para mantener la compatibilidad con el VAE de 3 canales, la máscara de un solo canal se replica en tres canales.
   • Esto preserva la estructura espacial mientras permite un aprendizaje eficiente de correspondencias de alta granularidad en el espacio latente.

2. Condicionamiento de Texto Alineado con CLIP:

   • Para habilitar capacidades de vocabulario abierto, un codificador de texto CLIP congelado genera incrustaciones ($t_{clip}$) para nombres de clases o descripciones en lenguaje natural.
   • Estas incrustaciones se inyectan en el U-Net de difusión mediante atención cruzada en múltiples escalas de eliminación de ruido. Esto permite que el modelo alinee consultas textuales con representaciones visuales en evolución, tanto global como localmente.

3. U-Net de Eliminación de Ruido Consistente con Segmentación:

   • Objetivo de Entrenamiento: El modelo se entrena para predecir el ruido añadido al latente de segmentación ($z_y$), condicionado al latente de imagen ($z_x$) y a la incrustación de texto ($t_{clip}$).
   • Formulación Generativa: A diferencia de trabajos anteriores que utilizan la difusión como extractor de características, DiGSeg trata la segmentación como una tarea de generación de eliminación de ruido condicional: $p_\theta(y|x)$. El U-Net aprende a eliminar ruido desde ruido aleatorio hasta una representación latente consistente con la segmentación.
   • Estrategia de Ruido: Se emplea un novedoso programa de ruido multiescala recocido. Las perturbaciones de alta frecuencia dominan los pasos iniciales para aprender estructuras finas, mientras que los componentes de baja frecuencia se vuelven influyentes más tarde para garantizar la coherencia semántica.

Proceso de Inferencia

1. La imagen de entrada se codifica en una representación latente.
2. Un latente de segmentación se inicializa como ruido gaussiano.
3. Bajo condicionamiento de texto, el U-Net elimina ruido iterativamente del latente (utilizando un muestreador DDIM con una estrategia de paso de tiempo "trailing" para la eficiencia).
4. El latente limpio final se decodifica de nuevo al espacio de píxeles utilizando el decodificador VAE congelado para producir la máscara de segmentación.
5. Ensamblaje en Tiempo de Prueba: Se promedian múltiples pasadas de inferencia con diferentes semillas de ruido en el espacio latente para mejorar la consistencia espacial y reducir artefactos.

3. Contribuciones Clave

• Marco Unificado: DiGSeg es el primer marco que reaprovecha con éxito una arquitectura base de difusión preentrenada en un aprendiz generalista de segmentación que maneja tareas de vocabulario cerrado, vocabulario abierto y de dominio cruzado sin cambios arquitectónicos específicos de la tarea.
• Aprendizaje Latente Directo: En lugar de depender de mapas de atención ruidosos, el método ajusta fino el U-Net de eliminación de ruido para generar directamente latentes consistentes con la segmentación, produciendo máscaras de alta calidad sin un pesado postprocesamiento.
• Condicionamiento de Texto Multiescala: Introduce un mecanismo para inyectar características lingüísticas a través de múltiples escalas de eliminación de ruido, permitiendo una alineación robusta entre indicaciones de texto arbitrarias y características visuales.
• Estrategia de Ruido Novel: Propone un programa de ruido multiescala recocido que equilibra el aprendizaje de límites de alta granularidad con la coherencia semántica global.

4. Resultados Experimentales

Los autores evaluaron DiGSeg en benchmarks estándar y dominios especializados, demostrando un rendimiento de vanguardia (SOTA).

• Segmentación de Vocabulario Abierto:
  • Superó a 13 líneas base (incluyendo ODISE, OVSeg, CAT-Seg) en A-847, PC-459, A-150, PC-59 y Cityscapes.
  • Logró ganancias significativas (por ejemplo, +2.8 mIoU en PC-59, +2.2 en PC-459) sobre los métodos de segundo mejor rendimiento.
• Segmentación Semántica Estándar:
  • Logró SOTA en COCO (50.8 mIoU) y ADE20K (58.6 mIoU), superando a modelos potentes basados en transformadores como Mask2Former y OneFormer.
• Transferencia de Dominio Cruzado:
  • Demostró una fuerte transferencia zero-shot a conjuntos de datos médicos (REFUGE-2), agrícolas (PhenoBench) y de teledetección (DeepGlobe) sin modificación arquitectónica.
  • Nota: El rendimiento en dominios médicos fue inferior al de las escenas naturales, atribuido a la brecha de dominio en los datos de preentrenamiento de CLIP, pero sigue siendo competitivo.
• Eficiencia y Uso de Datos:
  • El modelo es altamente eficiente en datos, logrando un rendimiento casi idéntico con solo el 50% de los datos de entrenamiento.
  • La velocidad de inferencia se mejora mediante DDIM-trailing (alcanzando la saturación en 1 paso), aunque sigue siendo más lenta que los modelos de alimentación directa debido a la naturaleza de la difusión.

5. Significado

Este trabajo cierra fundamentalmente la brecha entre la generación visual y la comprensión visual.

• Cambio de Paradigma: Cuestiona la noción de que los modelos de difusión son exclusivamente para síntesis, demostrando que pueden servir como poderosos aprendices discriminativos generalistas cuando se condicionan adecuadamente.
• Generalización: Demuestra que los ricos priores espaciales y semánticos codificados en modelos generativos a gran escala son transferibles a tareas de predicción densa a través de diversos dominios (natural, médico, agrícola).
• Simplicidad: Al evitar pipelines complejos de múltiples etapas o cabezales especializados, DiGSeg ofrece una interfaz unificada y simplificada para la segmentación, sugiriendo un futuro donde una única arquitectura base generativa puede manejar una amplia gama de tareas de visión por computadora.

En conclusión, DiGSeg establece que las arquitecturas base de difusión modernas, cuando se ajustan fino con un protocolo condicionado latente y alineación de texto, pueden superar a modelos de segmentación especializados mientras ofrecen una flexibilidad superior en escenarios de vocabulario abierto y de dominio cruzado.