MaskInversion: Localized Embeddings via Optimization of Explainability Maps

1. El Problema

Los modelos fundacionales de visión-linguaje, como CLIP, han logrado avances significativos en la alineación global entre imágenes y texto. Sin embargo, presentan limitaciones críticas cuando se trata de representaciones localizadas (regiones específicas de una imagen).

• Limitación actual: Estos modelos están entrenados para alinear tokens globales ([CLS]) con descripciones de texto, ignorando detalles regionales.
• Soluciones existentes insuficientes: Métodos ingenuos como recortar la imagen (pérdida de contexto) o agregar tokens locales (que no siempre corresponden a la representación correcta) no ofrecen resultados óptimos.
• Métodos de ajuste (Fine-tuning): Enfoques como AlphaCLIP requieren reentrenar el modelo con millones de pares máscara-texto, lo que es costoso y no generaliza bien a nuevos modelos preentrenados sin adaptación.

El objetivo es generar un token de incrustación (embedding) localizado para una región específica de una imagen (definida por una máscara) sin modificar los pesos del modelo base (en configuración zero-shot o test-time).

2. Metodología: MaskInversion

La propuesta, MaskInversion, es un método de inversión que optimiza un token de incrustación en tiempo de prueba utilizando mapas de explicabilidad.

Flujo de Trabajo:

1. Inicialización: Se inicia un token de incrustación localizada ($LET_m$) copiando el token global [CLS] del modelo base (ej. CLIP).
2. Cálculo del Mapa de Explicabilidad: Se utiliza este token inicial para generar un mapa de explicabilidad (heatmap) que indica qué regiones de la imagen influyen más en la representación del token. Se emplea el método LeGrad (basado en gradientes de mapas de atención en ViT) por su eficiencia y precisión.
3. Optimización:
   • Se compara el mapa de explicabilidad generado con la máscara de consulta (el objetivo).
   • Se define una función de pérdida basada en la pérdida Dice suave para minimizar la discrepancia entre el mapa de explicabilidad y la máscara.
   • Se realiza un descenso de gradiente iterativo para actualizar únicamente el vector de incrustación ($LET_m$), manteniendo el modelo base congelado.
4. Regularización (Opcional): Se introduce una pérdida de regularización ($L_{reg}$) que alinea el token aprendido con la representación global de la imagen original. Esto permite equilibrar la información local (la región) con el contexto global de la imagen, crucial para tareas como la recuperación de expresiones referidas.

Innovación Clave: Descomposición del Gradiente

El cálculo de mapas de explicabilidad requiere derivadas de segundo orden, lo cual es computacionalmente costoso si se repite para múltiples máscaras en la misma imagen.

• Solución: Los autores proponen una estrategia de descomposición de gradientes. Dado que el token aprendido no depende de las activaciones del modelo en el paso de gradiente (es una variable externa), el gradiente se puede descomponer en un producto punto entre el token y la derivada de las activaciones.
• Beneficio: Esto permite calcular la derivada de las activaciones una sola vez y reutilizarla para todas las iteraciones de optimización y para múltiples máscaras, reduciendo drásticamente el tiempo de ejecución.

3. Contribuciones Principales

1. Aprendizaje de Incrustación Localizada en Tiempo de Prueba: Un método que genera un token único que captura las características de una región específica definida por una máscara, sin tocar los pesos del modelo preentrenado.
2. Eficiencia Computacional: La estrategia de descomposición de gradientes permite procesar múltiples máscaras en una sola imagen de manera eficiente, eliminando la necesidad de retropropagación completa en cada iteración.
3. Versatilidad y Aplicabilidad: El token aprendido actúa como un reemplazo directo ("drop-in replacement") del token [CLS] en cualquier aplicación basada en el mismo backbone (CLIP), habilitando tareas como clasificación, descripción de imágenes y generación.

4. Resultados Experimentales

El método se evaluó en múltiples tareas y conjuntos de datos (PascalVOC, MSCOCO, RefCOCO, OpenImagesV7), superando consistentemente a los métodos State-of-the-Art (SOTA).

• Recuperación de Expresiones Referidas (Referring Expression Retrieval): En datasets como RefCOCO y RefCOCO+, MaskInversion alcanzó el mejor rendimiento (SOTA), superando a métodos que modifican la entrada (como RedCircle o FGVP) y a los que requieren ajuste fino (AlphaCLIP).
• Recuperación de Clases (Class Retrieval): En tareas de clasificación zero-shot de regiones específicas, MaskInversion superó a todos los baselines en PascalVOC, PascalContext y COCO. Destaca su capacidad para manejar vocabularios grandes (OpenImagesV7 con 350 clases) sin necesidad de entrenamiento.
• Comparación con Métodos sin Entrenamiento: Superó significativamente a métodos que agregan tokens (como MaskCLIP o CLIPSurgery), especialmente al escalar a modelos de backbone más grandes (ViT-H/14), demostrando que aprovecha mejor la capacidad representativa de los modelos grandes.
• Aplicaciones Derivadas:
  • Descripción Localizada (Localized Captioning): Al reemplazar el token [CLS] en modelos como CLIPCap, logró generar descripciones precisas de regiones específicas, superando a AlphaCLIP.
  • Generación de Imágenes (Localized Diffusion): Integrado con modelos de difusión (ej. $\lambda$-ECLIPSE), permitió generar variaciones de imágenes enfocadas estrictamente en la región definida por la máscara.

5. Significado e Impacto

MaskInversion representa un avance importante en la utilización eficiente de modelos fundacionales.

• Eficiencia de Recursos: Elimina la necesidad de costosos procesos de ajuste fino (fine-tuning) o reentrenamiento para tareas de localización, permitiendo que cualquier modelo preentrenado se adapte dinámicamente a regiones específicas.
• Flexibilidad: Al ser un método de "caja negra" que solo optimiza un vector, es agnóstico al modelo base y aplicable a cualquier arquitectura ViT compatible con mapas de explicabilidad.
• Puente entre Explicabilidad y Generación: Demuestra que los mapas de explicabilidad no solo sirven para entender modelos, sino que pueden utilizarse como funciones objetivo para invertir y generar representaciones semánticas específicas de regiones, abriendo nuevas vías para la interacción humano-IA en tareas de visión por computadora.

En resumen, MaskInversion ofrece una solución elegante y computacionalmente eficiente para el problema de la alineación visión-linguaje localizada, superando las limitaciones de los enfoques actuales de recorte de imágenes o ajuste fino masivo.