JEPAMatch: Geometric Representation Shaping for Semi-Supervised Learning

1. Planteamiento del Problema

El aprendizaje semi-supervisado (SSL) busca aprovechar grandes cantidades de datos no etiquetados para mejorar el rendimiento de los modelos cuando los datos etiquetados son escasos. Aunque métodos basados en FixMatch (que combinan aumentación de datos y auto-entrenamiento con pseudo-etiquetas) han logrado resultados de vanguardia, sufren de dos cuellos de botella críticos:

1. Sesgo de clases mayoritarias: La selección de pseudo-etiquetas basada en umbrales de confianza fijos o adaptativos tiende a favorecer a las clases mayoritarias. Esto genera un desequilibrio donde las clases dominantes refuerzan su presencia en el espacio de representación, mientras que las minoritarias reciben una supervisión escasa, lo que lleva a modelos sesgados.
2. Convergencia lenta y ruido temprano: Las pseudo-etiquetas generadas en las etapas iniciales del entrenamiento suelen ser ruidosas e incorrectas. Esto impide que el modelo forme fronteras de decisión claras, obligando a realizar un gran número de iteraciones (a menudo más de 220k) para converger, lo que incrementa drásticamente el costo computacional.

2. Metodología Propuesta: JEPAMatch

El artículo propone JEPAMatch, un nuevo paradigma que se aleja de la optimización puramente lógica basada en umbrales de salida para enfocarse en la configuración explícita de la geometría de las representaciones latentes. La metodología se inspira en las arquitecturas LeJEPA (Latent-Euclidean Joint-Embedding Predictive Architectures), que postulan que las representaciones de alta calidad deben seguir una estructura de Gaussiana isotrópica en el espacio latente.

El enfoque se divide en dos niveles complementarios:

A. Nivel de Currículo (Curriculum Level)

Este nivel se encarga de la selección de pseudo-etiquetas, manteniendo la lógica de FlexMatch (una extensión adaptativa de FixMatch):

• Utiliza vistas débiles y fuertes de las imágenes no etiquetadas.
• Emplea umbrales de confianza dinámicos por clase ($\tau_c$) para mitigar el dominio de las clases mayoritarias.
• Genera máscaras de confianza para decidir qué pseudo-etiquetas son lo suficientemente fiables para supervisar el aprendizaje.

B. Nivel de Representación (Representation Level)

Este es el núcleo de la innovación, diseñado para estructurar el espacio de características y evitar el colapso dimensional. Utiliza una estrategia de aumentación multi-vista (vistas globales débiles/fortes y vistas locales) y se basa en tres componentes de pérdida:

1. Pérdida de Predicción Auto-supervisada (JEPA Loss):

   • El modelo predice las representaciones de las vistas fuertes y locales basándose en la vista global débil (que actúa como objetivo).
   • Esto fuerza al modelo a aprender características geométricamente significativas e invariantes sin reconstruir la entrada cruda.

2. SIGReg Adaptativo por Clase (Adaptive Class-wise SIGReg):

   • Inspirado en LeJEPA, busca que las representaciones sigan una distribución Gaussiana isotrópica.
   • Fase de Calentamiento (Warmup): Se aplica una regularización global (SIGReg) sobre todas las vistas locales para asegurar un espacio de características no colapsado, ignorando temporalmente las asignaciones de clase para evitar sesgos de confirmación.
   • Fase Principal: Se transiciona a una regularización por clase. Se fuerza a que cada clase forme su propia "isla" gaussiana aislada con una media específica ($\mu_c$) y una varianza controlada.
   • Annealing de Varianza: La varianza objetivo ($\sigma$) se reduce suavemente de 1.0 a 0.1 a medida que avanza el entrenamiento, compactando los clusters de clase.

3. Pérdida de Repulsión Activa (Active Repulsion Loss):

   • Para evitar que el modelo minimice la pérdida de regularización colapsando todas las medias de clase en un solo punto, se añade una pérdida que penaliza la similitud coseno entre las medias de clases distintas, manteniendo los clusters separados.

Pérdida Total: La función de objetivo combina la pérdida supervisada, la pérdida no supervisada (consistencia) y la pérdida de representación ($L_{rep}$), equilibradas por hiperparámetros.

3. Contribuciones Clave

• Cambio de Paradigma: Propone un enfoque que desacopla la tarea de clasificación discreta de la organización geométrica del espacio latente, integrando la teoría de JEPAs en el SSL.
• Mitigación del Desequilibrio: Al estructurar el espacio latente en clusters gaussianos isotrópicos por clase, se reduce el sesgo hacia las clases mayoritarias sin depender exclusivamente de umbrales de confianza.
• Aceleración de la Convergencia: La estructura geométrica impuesta permite que el modelo aprenda representaciones informativas mucho más rápido, reduciendo la necesidad de iteraciones excesivas.
• Versatilidad: El método es compatible con variantes de FixMatch como FlexMatch, FreeMatch y SoftMatch.

4. Resultados Experimentales

Los autores evaluaron JEPAMatch en los benchmarks estándar CIFAR-100, STL-10 y Tiny-ImageNet.

• Rendimiento Superior: JEPAMatch superó consistentemente a los métodos de vanguardia (FixMatch, FlexMatch, FreeMatch, SoftMatch, etc.).
  • En CIFAR-100 con 400 etiquetas, logró un error del 34.25% (vs. 37.98% de FreeMatch).
  • En STL-10 con 1000 etiquetas, alcanzó un error del 4.28%, superando a todos los baselines.
• Eficiencia Computacional:
  • Logró converger en 217 iteraciones, frente a las 220+ iteraciones típicas de otros métodos, reduciendo el costo computacional.
  • En términos de velocidad de convergencia, JEPAMatch alcanzó una precisión del 46.4% aproximadamente 50,000 iteraciones antes que FlexMatch (en escenarios de pocas etiquetas), convergiendo 3 veces más rápido.
• Calidad de Pseudo-etiquetas: La precisión de las pseudo-etiquetas generadas por JEPAMatch fue significativamente mayor (alcanzando 50% en CIFAR-100) en comparación con FlexMatch (35%), lo que indica una mejor filtración de ruido desde las etapas tempranas.
• Robustez ante Desequilibrio: En experimentos con ratios de desequilibrio alto ($\gamma=100$), JEPAMatch mantuvo un rendimiento superior, reduciendo el error en un 1-2% frente a los métodos base.

5. Significado e Impacto

JEPAMatch representa un avance significativo en el aprendizaje semi-supervisado al demostrar que la estructura geométrica del espacio latente es tan crucial como la precisión de las pseudo-etiquetas. Al integrar principios teóricos de las arquitecturas predictivas de incrustación conjunta (JEPAs) con estrategias de auto-entrenamiento, el método resuelve problemas fundamentales de sesgo de clase y lentitud de convergencia.

Esto sugiere que el futuro del SSL no reside solo en refinar la selección de etiquetas, sino en dar forma activamente a la geometría de las representaciones, permitiendo modelos más robustos, precisos y eficientes en términos de recursos computacionales, lo cual es vital para aplicaciones en entornos con datos limitados y costosos.