Single-Channel Tissue Segmentation via Cross-Modal Distillation from Foundation Models

Resumen Técnico: Segmentación de Tejidos de Canal Único mediante Destilación Cross-Modal de Modelos de Fundación

Planteamiento del Problema
La segmentación precisa de tejidos y células es crítica para aplicaciones biomédicas como la caracterización de enfermedades y el perfilado espacial. Si bien la microscopía de fluorescencia multiplexada (que combina los canales de núcleo DAPI y de membrana E-cadherina) proporciona un contexto espacial más rico y un rendimiento de segmentación superior en comparación con la imagen de canal único, presenta desafíos significativos de implementación. Los modelos multiplexados requieren todos los canales de entrada durante la inferencia, lo que aumenta el uso de memoria, la latencia y la complejidad. En muchos escenarios del mundo real, solo se dispone de un subconjunto de canales (por ejemplo, el canal nuclear). Los modelos de aprendizaje profundo actuales tienen dificultades para mantener una alta precisión cuando se ven restringidos a entradas de canal único, particularmente al delinear los bordes celulares en regiones densamente agrupadas.

Metodología
Los autores proponen un marco de destilación de conocimiento (KD) cross-modal diseñado para transferir información semántica desde un modelo de fundación de alta capacidad (maestro) hacia un modelo estudiante ligero que opera con entradas de canal único.

• Arquitectura:
  • Maestro (Teacher): Un modelo de fundación congelado (Segment Anything Model [SAM] ViT-H o CellSAM) procesa entradas multiplexadas (canales de núcleo + membrana) para generar objetivos de probabilidad suave por píxel.
  • Estudiante (Student): Una variante ligera de U-Net (utilizando backbones Swin-Tiny, ResNet18, EfficientNet-B0 o MobileNetV3) que se entrena exclusivamente con el canal nuclear único.
• Objetivo de Destilación:
  • Emparejamiento de Objetivos Suaves (Soft-Target Matching): En lugar de la divergencia KL estándar, el marco emplea el Error Cuadrático Medio (MSE) entre las probabilidades sigmoideas del maestro y el estudiante para alinear las predicciones por píxel.
  • Supervisión Consciente de Bordes (Boundary-Aware Supervision): Se aplica una pérdida de borde específica ($L_{bnd}$) en las regiones donde la máscara de la verdad de campo (ground truth) está morfológicamente erosionada, concentrando las señales de gradiente en los bordes celulares donde los modelos de canal único suelen fallar.
  • Ponderación por Incertidumbre: Para equilibrar la pérdida de segmentación, la pérdida de destilación y la pérdida de borde sin ajuste manual de hiperparámetros, los autores utilizan pesos de incertidumbre aprendibles ($\sigma_1, \sigma_2$). La función de pérdida total escala inversamente cada término según su incertidumbre asociada, permitiendo que el modelo priorice automáticamente los objetivos en los que tiene mayor confianza.
• Protocolo de Entrenamiento: El proceso consta de dos etapas. Primero, el maestro procesa todas las imágenes de entrenamiento para almacenar en caché los objetivos suaves, desacoplando la inferencia del maestro del entrenamiento del estudiante para reducir los costos de memoria y tiempo. Segundo, el estudiante se entrena con datos de canal único utilizando los objetivos almacenados en caché y las máscaras de verdad de campo.
• Generalización Cross-Dataset: El marco se evalúa en un entorno de cross-dataset donde los estudiantes entrenados en BBBC038 (solo canal único) son supervisados por las predicciones del maestro almacenadas en caché de TissueNet (multiplexado), sin reentrenar al maestro ni utilizar datos multiplexados del dominio objetivo.

Contribuciones Clave

1. Marco de KD Cross-Modal: Un enfoque novedoso para transferir información semántica de un modelo de fundación multiplexado congelado a un estudiante de canal único, abordando la restricción de "canal faltante" en la imagen biomédica.
2. Objetivo Ponderado por Incertidumbre: Una estrategia de optimización que integra el emparejamiento de objetivos suaves y la supervisión consciente de bordes con pesos de incertidumbre aprendibles para equilibrar señales de entrenamiento heterogéneas.
3. Evaluación Exhaustiva: Pruebas exhaustivas de dos maestros (SAM ViT-H, CellSAM) frente a cuatro arquitecturas de estudiantes (que van desde 1.5M hasta 27M de parámetros) en los conjuntos de datos TissueNet y BBBC038.
4. Protocolo Cross-Dataset: Demostración de un protocolo de destilación que aprovecha las predicciones del maestro del dominio de origen para supervisar a los estudiantes del dominio objetivo sin reentrenar al maestro.

Resultados
Los experimentos demuestran mejoras consistentes de rendimiento en todas las arquitecturas de estudiantes y conjuntos de datos:

• Ganancias de Rendimiento: En TissueNet, la destilación con SAM ViT-H mejoró las puntuaciones Dice en aproximadamente 12–13 puntos en todos los modelos de estudiante en comparación con las líneas base no destiladas. Por ejemplo, el estudiante Swin-Tiny alcanzó una puntuación Dice de 78.36 (frente a 65.31 de la línea base), recuperando el 87.9% del rendimiento de oráculo del maestro (89.12) mientras reducía los parámetros en 23×.
• Agnosticismo de Arquitectura: Las ganancias fueron consistentes independientemente de la capacidad del estudiante, siendo que MobileNetV3 (1.5M de parámetros) también experimentó mejoras significativas (~12 puntos Dice).
• Superioridad del Maestro: SAM ViT-H superó consistentemente a CellSAM como maestro, lo que sugiere que la calidad del rendimiento de oráculo del maestro es un indicador de éxito de la destilación más fuerte que la especialización de dominio por sí sola.
• Reducción de Parámetros: El marco logra hasta una reducción de parámetros de 421× en la inferencia (de 632M del maestro a 1.5M del estudiante) requiriendo solo un canal de entrada.
• Generalización: En el conjunto de datos BBBC038, los estudiantes destilados por SAM mejoraron sobre las líneas base en 7.25–8.29 puntos Dice sin reentrenar al maestro, confirmando la transferibilidad de las representaciones semánticas a través de diferentes condiciones de imagen.

Significancia y Afirmaciones
El artículo afirma que este marco ofrece una solución fundamentada para desplegar modelos de segmentación de alto rendimiento en entornos con recursos limitados o datos limitados donde solo se dispone de datos de canal único. Al aprovechar los modelos de fundación congelados, el enfoque cierra la brecha de rendimiento entre la imagen multiplexada y la de canal único sin requerir que el estudiante tenga acceso a las modalidades faltantes durante la inferencia. Los autores enfatizan que el método es particularmente efectivo para la delineación de bordes en regiones celulares congestionadas y se generaliza bien a través de los conjuntos de datos. Concluyen que la calidad del maestro es un factor crítico para una transferencia de conocimiento efectiva y sugieren trabajos futuros que extiendan esto a la segmentación de instancias/panóptica y otras modalidades como la histología H&E.