JEPAMatch: Geometric Representation Shaping for Semi-Supervised Learning

1. Il Problema

L'apprendimento semi-supervisionato (SSL) mira a sfruttare grandi quantità di dati non etichettati per migliorare le prestazioni dei modelli quando i dati etichettati sono scarsi. Sebbene metodi basati su FixMatch (e le sue varianti come FlexMatch) abbiano raggiunto lo stato dell'arte, soffrono di due colli di bottiglia critici:

• Dominanza delle classi maggioritarie: La selezione dei pseudo-label basata su una soglia di confidenza fissa o adattiva tende a favorire le classi più frequenti. Gli pseudo-label errati delle classi maggioritarie rinforzano il loro dominio nello spazio delle rappresentazioni, portando a modelli distorti e a un collasso dimensionale (dimensional collapse).
• Convergenza lenta: I metodi basati sulla consistenza richiedono un numero elevato di iterazioni (spesso >220k) per convergere. All'inizio dell'addestramento, la confidenza del classificatore è bassa, quindi la maggior parte dei dati non etichettati viene scartata, rendendo inefficiente l'uso delle risorse computazionali.

2. Metodologia: JEPAMatch

Il paper propone JEPAMatch, un nuovo paradigma che sposta il focus dalla logica di soglia delle uscite all'modellazione esplicita della geometria delle rappresentazioni. Il metodo integra i principi delle Joint-Embedding Predictive Architectures (JEPA), in particolare il framework teorico LeJEPA, che afferma che le rappresentazioni di alta qualità dovrebbero seguire una struttura di Gaussiana isotropa nello spazio latente.

L'approccio scompone l'apprendimento in due livelli complementari:

A. Livello Curriculum (Selezione dei Pseudo-Label)

• Utilizza il framework FlexMatch per la selezione dei pseudo-label.
• Impiega soglie di confidenza dinamiche per classe ($\tau_c$) per mitigare il problema dello squilibrio delle classi.
• Utilizza viste globali (deboli e forti) per generare pseudo-label affidabili e applicare la loss di cross-entropy non supervisionata.

B. Livello Rappresentazione (Geometria dello Spazio Latente)

Questo è il cuore dell'innovazione, che introduce una regolarizzazione geometrica basata su JEPA:

1. Multi-view Augmentation: Oltre alle viste globali (deboli/forti), vengono estratte viste locali ($K$ crop locali) dall'immagine.
2. Loss di Predizione Self-Supervised ($L_{pred}$): Basata su JEPA, forza le rappresentazioni delle viste forti e locali a predire la rappresentazione della vista debole globale. Questo insegna invarianza strutturale e semantica.
3. Adaptive Class-wise SIGReg (Sketched Isotropic Gaussian Regularization):
   • Fase Warmup: Si applica una regolarizzazione SIGReg globale (distribuzione $N(0, I)$) per prevenire il collasso dimensionale e garantire che tutte le dimensioni siano utilizzate, ignorando temporaneamente le etichette di classe.
   • Fase Principale: La regolarizzazione diventa adattiva per classe. Si forza ogni classe a formare un proprio cluster gaussiano isotropo isolato (con media specifica $\mu_c$ e varianza controllata $\sigma$).
   • Active Repulsion Loss: Una loss aggiuntiva che penalizza la similarità tra le medie delle diverse classi, impedendo che tutti i cluster collassino in un unico punto.

Il training avviene in due fasi: un warmup per stabilizzare la geometria globale, seguito dalla fase principale dove si strutturano i cluster di classe.

3. Contributi Chiave

• Svolta Paradigmatica: Passaggio dalla logica di soglia delle probabilità di uscita alla modellazione geometrica esplicita dello spazio latente.
• Integrazione JEPA-SSL: Prima applicazione dei principi LeJEPA (Gaussiana isotropa) all'interno di un framework SSL basato su pseudo-labeling (FlexMatch).
• Decoupling: Separazione del compito di classificazione discreta (Curriculum Level) dall'organizzazione geometrica dello spazio delle caratteristiche (Representation Level).
• Efficienza Computazionale: Il metodo permette di utilizzare tutti i dati non etichettati fin dalle prime iterazioni, accelerando drasticamente la convergenza.
• Robustezza allo Squilibrio: La regolarizzazione adattiva per classe riduce il bias verso le classi maggioritarie, migliorando le prestazioni su dataset sbilanciati.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su CIFAR-100, STL-10 e Tiny-ImageNet.

• Prestazioni: JEPAMatch supera costantemente gli stati dell'arte (inclusi FixMatch, FlexMatch, FreeMatch, SoftMatch).
  • Su CIFAR-100 con 400 etichette, ottiene un tasso di errore del 34.25% (vs 37.98% di FreeMatch).
  • Su STL-10 con 1000 etichette, ottiene il 4.28% di errore.
• Efficienza e Velocità di Convergenza:
  • JEPAMatch converge in 217 iterazioni, contro le 220+ richieste dai metodi basici (risparmio computazionale significativo).
  • La figura di convergenza mostra che JEPAMatch raggiunge prestazioni superiori di FlexMatch circa 50.000 iterazioni prima.
• Qualità dei Pseudo-Label: JEPAMatch mantiene un'accuratezza dei pseudo-label molto più alta durante tutto il training (es. ~50% su CIFAR-100 vs ~35% di FlexMatch), indicando una migliore filtrazione del rumore.
• Robustezza allo Squilibrio: Su scenari con forte squilibrio di classe ($\gamma=100$), JEPAMatch riduce l'errore di circa 1-2% rispetto alle migliori baseline.

5. Significato e Impatto

JEPAMatch dimostra che la struttura geometrica dello spazio latente è un fattore cruciale spesso trascurato nell'SSL. Integrando la teoria delle rappresentazioni isotrope (LeJEPA) con le strategie di pseudo-labeling adattivo, il metodo risolve il compromesso tra la necessità di selezionare solo dati "sicuri" (che rallenta l'apprendimento) e la necessità di strutturare lo spazio delle feature.

Il lavoro apre la strada a futuri approcci SSL che non si basano solo sulla consistenza delle previsioni, ma sulla modellazione attiva della geometria delle rappresentazioni, offrendo soluzioni più robuste, veloci e efficienti dal punto di vista computazionale, specialmente in scenari con dati etichettati molto scarsi o dataset sbilanciati.