Love Me, Love My Label: Rethinking the Role of Labels in Prompt Retrieval for Visual In-Context Learning

Resumen Técnico: LaPR (Label-aware Prompt Retrieval)

1. Planteamiento del Problema

El Aprendizaje en Contexto Visual (VICL, por sus siglas en inglés) permite que los modelos fundacionales visuales realicen diversas tareas utilizando ejemplos demostrativos (prompts). Sin embargo, el rendimiento de VICL depende críticamente de la selección adecuada de estos prompts.

• Limitación actual: Los métodos anteriores de recuperación de prompts (como SupPR o Partial2Global) se centran casi exclusivamente en la similitud visual entre la imagen de consulta (query) y las imágenes de los prompts.
• El problema identificado: Estos enfoques a menudo recuperan prompts que son visualmente similares pero inconsistentes en sus etiquetas. Por ejemplo, para una consulta de un "gato", el sistema podría recuperar una imagen que contiene un gato pero está etiquetada como "flor". Esta inconsistencia de etiquetas degrada el rendimiento del modelo al generar predicciones erróneas.
• Observación clave: Los autores descubrieron una correlación positiva entre la consistencia de la etiqueta entre la consulta y el prompt, y el rendimiento final de VICL. Ignorar la información de la etiqueta en el proceso de recuperación es una pérdida de información crítica.

2. Metodología: El Marco LaPR

Para abordar esto, los autores proponen LaPR (Label-aware Prompt Retrieval), un nuevo paradigma que integra explícitamente la información de las etiquetas en la recuperación. La arquitectura se basa en tres pilares principales:

A. Representación Conjunta Imagen-Etiqueta

• Para los Prompts: En lugar de usar solo la imagen, LaPR fusiona las características de la imagen del prompt ($I^p$) y su etiqueta correspondiente ($L^p$) para crear una representación conjunta. Esto asegura que el espacio de embeddings capture tanto el contenido visual como el semántico (etiqueta).
• Para la Consulta: Dado que en la fase de inferencia la etiqueta de la consulta ($y_q$) es desconocida, el sistema debe inferirla implícitamente.

B. Mecanismo de Expertos Mezclados (Mixture-of-Experts - MoE)
Para manejar la incertidumbre de la etiqueta de la consulta y capturar modos específicos (ej. texturas, formas, categorías), se utiliza una arquitectura MoE en ambos lados (consulta y prompt):

• Expertos ($E_k$ y $\bar{E}_k$): Se emplean múltiples redes de expertos para capturar diferentes "modos" o características específicas dentro de las representaciones.
• Enrutador (Router): Un componente adaptativo a la consulta ($R$) analiza la imagen de entrada y genera pesos de mezcla ($\pi_q$) que determinan qué combinación de expertos es más relevante para esa consulta específica. Esto permite estimar la etiqueta implícita de la consulta y seleccionar la información del prompt más alineada con ella.

C. Optimización Alternada (Desacoplada)
Dado que los expertos y el enrutador tienen roles distintos, se propone un esquema de optimización en dos pasos por cada mini-lote para evitar conflictos de gradiente:

1. Paso de Expertos (Fijado el Router): Se optimizan los expertos utilizando una pérdida contrastiva guiada por el rendimiento ($\mathcal{L}_{PG}$). Aquí, los pares positivos/negativos se seleccionan basándose en qué prompts mejoran realmente el rendimiento de la tarea VICL (usando el modelo base MAE-VQGAN).
2. Paso del Router (Fijados los Expertos): Se optimiza el enrutador utilizando una pérdida contrastiva guiada por la etiqueta ($\mathcal{L}_{LG}$) y una pérdida de balanceo de carga ($\mathcal{L}_{LB}$). El objetivo es alinear los pesos de mezcla inferidos con las etiquetas verdaderas de los prompts y asegurar que todos los expertos sean utilizados equitativamente.

3. Contribuciones Clave

1. Primera recuperación de prompts consciente de etiquetas en VICL: LaPR es el primer marco que aborda explícitamente la inconsistencia de etiquetas, demostrando que las etiquetas son tan cruciales como las imágenes para la selección de prompts.
2. Arquitectura MoE Adaptativa: Introduce un diseño simétrico donde tanto los prompts como las consultas se codifican mediante múltiples expertos, con un enrutador que infiere dinámicamente la etiqueta implícita de la consulta para guiar la recuperación.
3. Estrategia de Optimización Desacoplada: Propone un método de entrenamiento alternado que separa el aprendizaje de representaciones específicas de modos (expertos) del aprendizaje de la asignación de esos modos (router), logrando mayor estabilidad y convergencia.
4. Generalización Robusta: El método no solo mejora el rendimiento, sino que demuestra una fuerte capacidad de transferencia entre diferentes extractores de características (CLIP vs. DINOv2) y entre diferentes particiones de datos (cross-fold).

4. Resultados Experimentales

Los autores evaluaron LaPR en tres tareas fundamentales de visión por computadora:

• Segmentación de primer plano (Dataset Pascal-5i).
• Detección de objetos individuales (Dataset Pascal VOC 2012).
• Colorización de imágenes (Dataset ImageNet).

Hallazgos principales:

• Rendimiento SOTA: LaPR superó consistentemente a los métodos más avanzados (SOTA), incluyendo variantes de reranking y voting.
  • En segmentación, superó al mejor método de recuperación (SupPR) en un 6.00% y al mejor método de reranking (RH-Partial2Global) en un 2.34% (sin votación).
  • En detección, logró mejoras significativas (ej. +3.46% sobre RH-Partial2Global).
• Transferencia Cross-Fold: En el escenario donde el modelo se entrena en un conjunto de categorías y se prueba en otras, LaPR superó a los competidores por márgenes amplios (ej. +14.8% sobre SupPR), demostrando que los expertos aprenden modos más finos que las categorías simples.
• Análisis de Ablación:
  • Eliminar la etiqueta del prompt o el enrutador causó caídas drásticas en el rendimiento.
  • La optimización alternada fue superior a la entrenamiento conjunto (joint training).
  • La pérdida guiada por el rendimiento ($\mathcal{L}_{PG}$) resultó ser el componente de supervisión más crítico.

5. Significado e Impacto

Este trabajo redefine el principio de selección de prompts en el aprendizaje en contexto visual.

• Cambio de Paradigma: Demuestra que la recuperación de prompts no debe ser solo una búsqueda de similitud visual, sino una búsqueda de coherencia semántica (etiqueta).
• Eficiencia de Datos: Al utilizar mejor la información de las etiquetas disponibles en los datos de entrenamiento, se mejora la capacidad de generalización del modelo sin necesidad de reentrenar el modelo fundacional completo.
• Futuro de la Investigación: Establece una nueva base para el desarrollo de sistemas de recuperación que integren múltiples modalidades (imagen + texto/etiqueta) y sugiere que la gestión de la incertidumbre de la etiqueta mediante mecanismos como MoE es una dirección prometedora para el VICL.

En conclusión, LaPR valida la hipótesis de que "amar la etiqueta" es esencial para el éxito del aprendizaje en contexto visual, logrando mejoras consistentes y robustas en tareas complejas de visión por computadora.