Resumen Técnico: ExPLoRe (Expert Patch-Level Loss Routing for Multi-Objective Masked Image Modeling)

Declaración del Problema

El Modelado de Imágenes Enmascaradas (MIM) multiobjetivo busca aprender representaciones visuales robustas combinando señales de aprendizaje complementarias, tales como la destilación a nivel de token, el alineamiento global de CLS y la reconstrucción a nivel de píxel. Aunque los métodos existentes integran con éxito estos objetivos, típicamente los equilibran utilizando pesos escalares globales aplicados uniformemente en todas las ubicaciones espaciales.

Este enfoque ignora la heterogeneidad espacial: diferentes regiones de la imagen se benefician de diferentes señales de aprendizaje. Los parches de primer plano (foreground) con riqueza semántica son mejor servidos por la destilación del profesor, mientras que las regiones de fondo (background) con mucha textura se benefician más de la reconstrucción de píxeles. Los métodos de equilibrio global (por ejemplo, GradNorm, ponderación por incertidumbre) no logran capturar esta variación e incluso pueden provocar una degradación catastrófica del rendimiento cuando los objetivos presentan una importancia espacialmente heterogénea. Hasta la fecha, ningún trabajo previo ha abordado la ponderación de pérdida adaptativa por parche para el modelado de imágenes enmascaradas.

Metodología: ExPLoRe

Los autores proponen ExPLoRe (Expert Patch-Level Loss Routing), un marco que reconvierte los pesos de despacho (dispatch weights) de Soft Mixture of Experts (Soft-MoE) para que sirvan como coeficientes de pérdida aprendidos por cada parche.

Mecanismo Central: Acoplamiento de Pérdida (Loss-Coupling)

A diferencia de las aplicaciones tradicionales de MoE que utilizan el enrutamiento para el escalado de capacidad o la asignación de expertos por tarea, ExPLoRe introduce el acoplamiento de pérdida:

1. Pesos de Despacho como Coeficientes de Pérdida: El marco utiliza los pesos de despacho ($D$) generados por el enrutador de Soft-MoE. Estos pesos, que determinan cuánto contribuye cada parche a cada experto, se utilizan para escalar la pérdida de ese parche específico.
2. Flujo de Gradiente: Crucialmente, los gradientes de los objetivos de pérdida fluyen de regreso a través de estos pesos de despacho hacia los parámetros del enrutador. Esto permite que el enrutador aprenda una especialización dependiente del contenido, ajustando dinámicamente la importancia de diferentes objetivos (por ejemplo, destilación vs. reconstrucción) para cada parche basándose en su contenido.
3. Estrategia de Normalización: Los autores utilizan pesos de despacho (normalizados sobre parches por experto) en lugar de pesos de combinación (normalizados sobre expertos por parche). Esta elección evita un problema de degeneración donde el enrutador podría minimizar la pérdida eliminando los pesos para parches difíciles. La normalización por imagen asegura que el enrutador deba distribuir la atención entre los parches en lugar de colapsar la pérdida.

Arquitectura y Entrenamiento

• Marco: El codificador estudiante (ViT-Base) procesa parches visibles (codificación dispersa) mientras que un profesor CLIP congelado proporciona objetivos semánticos.
• Objetivos: El sistema combina Destilación a Nivel de Token (pérdida Huber), Alineamiento Global de CLS (similitud de coseno) y Reconstrucción de Píxeles (MSE).
• Regularización de Entropía: Para prevenir el colapso de los expertos (donde un experto maneja todos los parches), se aplica un término de regularización de entropía a los pesos de despacho. Se encontró que el peso óptimo para esta regularización depende del número de expertos (por ejemplo, $\lambda=0.5$ para 2 expertos, $\lambda=5.0$ para 64 expertos).
• Ubicación: Las capas Soft-MoE reemplazan a las MLPs en bloques alternos (1, 3, 5, 7, 9, 11) del ViT-Base, aplicando el acoplamiento de pérdida en el bloque final de MoE (Bloque 11).

Recetas de Transferencia Downstream

Para abordar el desafío de transferir modelos MoE a tareas downstream sin sobrescribir los patrones de enrutamiento aprendidos, los autores proponen tres estrategias de adaptación complementarias:

1. Congelar Enrutamiento (Freeze Routing - FR): Congela los parámetros del enrutador ($\Phi, s$) para preservar las asignaciones de parche-experto preentrenadas.
2. Congelar Atención (Freeze Attention - FA): Congela los pesos de atención para preservar estructuras visuales transferibles.
3. Dropout de Expertos (Expert Dropout - ExD): Aplica dropout a las salidas de los expertos para evitar la dependencia excesiva de los expertos dominantes.
   La combinación de FR + FA + ExD es fundamental para cerrar la brecha de rendimiento entre los modelos MoE y los modelos no-MoE en tareas de predicción densa.

Contribuciones Clave

1. Ponderación de Pérdida por Peso de Despacho: La introducción de un mecanismo que reconvierte los pesos de despacho de Soft-MoE como coeficientes de pérdida por parche, permitiendo una especialización dependiente del contenido a través de múltiples objetivos.
2. Validación del Acoplamiento de Pérdida: Demostración empírica de que el flujo de gradientes a través de los pesos de despacho es el mecanismo central que impulsa la especialización. Una ablación de "desconexión" (bloqueo de gradientes) degrada el rendimiento en un 1.6%, y las variantes de MoE sin ponderación rinden por debajo de las bases no-MoE.
3. Análisis de la Dinámica de Enrutamiento: Caracterización del comportamiento de enrutamiento, revelando una transición brusca entre el enrutamiento estable y el colapso catastrófico gobernado por la regularización de entropía. El estudio identifica que la regularización óptima depende del número de expertos.
4. Recetas de Adaptación Downstream: Desarrollo de una estrategia compuesta (Congelar Enrutamiento, Congelar Atención, Dropout de Expertos) que permite a los modelos MoE igualar o superar a las bases no-MoE tanto en clasificación como en segmentación semántica.

Resultados Experimentales

Evaluado en ImageNet-1K con ViT-Base y profesores CLIP congelados:

• Rendimiento de Clasificación:
  • Token+CLS: ExPLoRe alcanza un 80.6% de precisión de linear probe (64 expertos) y un 85.3% de precisión de finetuning, siendo competitivo con métodos de vanguardia como CAE v2 y BEiT v2.
  • Token+Pixel: Alcanza un 73.7% de precisión de k-NN con 32 expertos, una mejora de +2.2% sobre la base no-MoE.
  • Aislamiento del Mecanismo: El mecanismo de ponderación de pérdida por sí mismo contribuye con un +1.3% (2 expertos) y un +0.9% (64 expertos) sobre las variantes de MoE sin ponderación, confirmando el valor del mecanismo de enrutamiento más allá del simple escalado de parámetros.
• Predicción Densa (ADE20K):
  • Sin las recetas de adaptación, los modelos MoE rinden por debajo de las bases no-MoE por 2.5–2.9 mIoU.
  • Con la receta completa FR+FA+ExD, los modelos MoE cierran esta brecha, alcanzando 52.8 mIoU (Token+Pixel) y 51.1 mIoU (Token+CLS), igualando o superando ligeramente las bases no-MoE.
• Eficiencia: A pesar de tener hasta 1.86B de parámetros (64 expertos), Soft-MoE desacopla el conteo de parámetros del costo de inferencia. Las configuraciones de ExPLoRe requieren menos GFLOPs (11.93–13.86) que los comparadores estándar de ViT-B/16 (17.45), ofreciendo una relación costo-calidad favorable.

Significado y Reivindicaciones

El artículo afirma que ExPLoRe aborda una brecha fundamental en el aprendizaje multitarea para la visión: la incapacidad del equilibrio de pérdida global para manejar la heterogeneidad espacial. Al demostrar que el acoplamiento de pérdida es el motor esencial de la especialización, el trabajo establece que los mecanismos de enrutamiento pueden reconvertirse no solo para la capacidad, sino para la modulación dinámica de la pérdida.

Los autores enfatizan que, si bien las ganancias absolutas de k-NN son modestas, las mejoras significativas en la precisión de linear probe sugieren que el enrutamiento acoplado a la pérdida moldea el espacio de características de formas que no son capturadas totalmente por la evaluación de vecinos más cercanos. Además, el trabajo destaca que los modelos MoE no son inherentemente superiores para la transferencia downstream; requieren recetas de adaptación específicas (Freeze Routing, etc.) para alcanzar su potencial, particularmente en tareas de predicción densa donde el finetuning estándar suele fallar.

El estudio conclula que ExPLoRe proporciona un énfasis de pérdida por parche dependiente del contenido que los métodos previos carecen, logiendo resultados competitivos en clasificación y predicción densa sin sacrificar la calidad de la representación, al tiempo que ofrece un perfil de inferencia más eficiente en comparación con el escalado de parámetros denso.