MeFEm: Medical Face Embedding model

1. Problema y Contexto

El trabajo aborda la escasez crítica de datos médicos de alta calidad y etiquetados para el entrenamiento de modelos de visión por computadora en el ámbito de la salud. Aunque los modelos fundacionales (foundation models) han tenido éxito en tareas generales, su rendimiento en tareas médicas especializadas suele ser deficiente debido a la falta de datos de entrenamiento masivos y específicos.

• Desafío específico: Las imágenes faciales contienen información fisiológica valiosa para el diagnóstico médico (biomarcadores), pero los conjuntos de datos etiquetados con condiciones médicas son pequeños, propietarios o inexistente debido a preocupaciones de privacidad y el alto costo de la anotación clínica.
• Hipótesis: Un modelo entrenado mediante aprendizaje auto-supervisado (SSL) en grandes colecciones de imágenes faciales no médicas puede aprender representaciones robustas y relevantes para la medicina, capaces de transferirse a tareas de predicción biométrica y médica con pocos datos etiquetados.

2. Metodología

El modelo propuesto, MeFEm, se basa en una arquitectura modificada de Joint Embedding Predictive Architecture (JEPA), utilizando un backbone Vision Transformer (ViT). En lugar de reconstruir píxeles (como en los autoencoders), JEPA predice representaciones latentes de regiones ocultas de la imagen basándose en las regiones visibles.

Modificaciones Clave de la Arquitectura:

1. Enmascaramiento de Tiras Axiales (Axial Stripe Masking):
   • A diferencia del enmascaramiento aleatorio por bloques estándar de JEPA (que puede dejar fuera las regiones centrales semánticamente importantes de una cara), MeFEm utiliza un enmascaramiento de tiras completas (horizontales o verticales).
   • La posición de la tira se muestrea desde una distribución gaussiana centrada en el medio de la imagen. Esto asegura que el contexto de origen (source) siempre contenga información facial relevante, evitando que el modelo aprenda características de fondo irrelevantes.
2. Asignación Probabilística del Token CLS:
   • Se reintroduce el token especial CLS (típicamente desactivado en JEPA puro) para facilitar la extracción de características globales.
   • El token CLS se asigna aleatoriamente al conjunto de origen (source) o al conjunto objetivo (target) con una probabilidad $P_{c \in S} = 0.5$. Esto permite que el token aprenda una representación global significativa sin comprometer el aprendizaje detallado de los parches locales.
3. Pesos de Pérdida Circulares (Circular Loss Weighting):
   • Se introduce un esquema de ponderación de la función de pérdida basado en la distancia al centro de la imagen.
   • Las parches cercanos al centro (donde se encuentra la cara) tienen un peso mayor, mientras que los bordes (fondo) tienen un peso menor. Esto utiliza una función sigmoide para suavizar la transición, reduciendo la influencia de regiones irrelevantes en el entrenamiento.
4. Preprocesamiento de Datos:
   • Se utiliza detección de rostros (Blazeface) seguida de un recorte cuadrado ampliado (2x el tamaño del cuadro delimitador original) para incluir características periféricas como la frente y las orejas, manteniendo una alineación espacial consistente para el ViT.

Datos de Entrenamiento:

El modelo se entrenó en un conjunto de datos consolidado de aproximadamente 6.4 millones de muestras, derivado de:

• FaceCaption-15M: Imágenes web filtradas y limpiadas.
• AVSpeech: Clips de video de personas hablando (se extrae un frame por clip).
• SFHQ: Imágenes faciales sintéticas de alta calidad.
• Nota: A pesar de ser un conjunto grande, es significativamente más pequeño que los utilizados por modelos fundacionales generales como Franca o FaRL.

3. Contribuciones Clave

• MeFEm: Un nuevo modelo fundacional específico para análisis facial médico y biométrico, disponible públicamente.
• Innovaciones en JEPA: Propuesta de estrategias de enmascaramiento y ponderación de pérdidas adaptadas a la estructura espacial de las imágenes faciales, superando las limitaciones de los métodos genéricos.
• Evaluación de BMI: Creación y evaluación en un nuevo conjunto de datos consolidado y cerrado para la estimación del Índice de Masa Corporal (BMI), abordando el sesgo de dominio presente en los datos existentes.
• Eficiencia de Datos: Demostración de que un enfoque puramente visual auto-supervisado puede superar a modelos que dependen de alineación texto-imagen (como FaRL) en tareas médicas, utilizando menos datos de entrenamiento.

4. Resultados

El modelo se evaluó en tareas de regresión y clasificación, comparándolo con baselines fuertes como FaRL (específico para rostros), Franca (modelo general de gran escala) y Buffalo (ResNet50).

• Tareas Biométricas (Edad, Género, Etnia): MeFEm superó consistentemente a todos los baselines en el conjunto de datos CelebA y FairFace.
  • En la predicción de edad (FairFace), alcanzó una precisión del 91.0% (Patches) y 90.7% (CLS), superando a FaRL (90.4%) y Franca (90.4%).
• Estimación de BMI:
  • MeFEm logró un $R^2$ de 0.506 y un Error Absoluto Medio (MAE) de 3.453, superando a FaRL ($R^2$ 0.380, MAE 3.924) y Franca ($R^2$ 0.425, MAE 3.744).
  • Esto es notable dado que MeFEm se entrenó con una fracción de los datos utilizados por sus competidores.
• Predicción de Parámetros Médicos (MCD-rPPG):
  • Se evaluó la predicción de hemoglobina, hemoglobina glicada, colesterol, presión arterial y saturación de oxígeno a partir de un solo fotograma.
  • Aunque la predicción de colesterol y presión diastólica fue pobre (como se esperaba por la dificultad del problema), el modelo mostró capacidad predictiva significativa para hemoglobina glicada y saturación de oxígeno, superando o igualando a las líneas base de fotopletismografía remota (rPPG) que utilizan secuencias de video completas.

5. Significado e Impacto

• Nuevo Estándar de Referencia: MeFEm establece una nueva línea base sólida para la investigación en análisis facial médico, demostrando que no es necesario un entrenamiento masivo con texto o datos médicos etiquetados para obtener características útiles.
• Superioridad del Enfoque Visual Puro: El trabajo sugiere que para tareas donde la alineación con el texto no aporta valor (como la detección de anomalías fisiológicas sutiles), un enfoque de SSL puramente visual especializado es más eficiente y efectivo que los modelos multimodales (CLIP-like).
• Aplicabilidad Clínica: La capacidad de estimar parámetros clínicos (como el BMI y marcadores de salud) a partir de una sola imagen facial abre puertas a aplicaciones de tamizaje médico no invasivo y de bajo costo, especialmente en entornos con recursos limitados.
• Disponibilidad: Los pesos del modelo están disponibles públicamente en Hugging Face, fomentando la reproducibilidad y el desarrollo futuro en este dominio.

En resumen, MeFEm demuestra que mediante modificaciones arquitectónicas inteligentes en el aprendizaje auto-supervisado (JEPA), es posible extraer características biomédicas robustas de imágenes faciales, superando a modelos fundacionales mucho más grandes y entrenados con más datos.