Why Prototypes Collapse: Diagnosing and Preventing Partial Collapse in Prototypical Self-Supervised Learning

1. Il Problema: Il Collasso Parziale dei Prototipi

Il paper affronta un problema fondamentale nell'apprendimento auto-supervisionato (SSL) basato su prototipi, utilizzato in framework come DINO, DINOv2 e CARP.

• Definizione: Il "collasso parziale dei prototipi" si verifica quando più prototipi (vettori appresi che agiscono come ancoraggi di cluster) convergono verso rappresentazioni quasi identiche durante l'addestramento.
• Conseguenze: Questo fenomeno mina lo scopo centrale dell'SSL basato su prototipi, che è fornire target diversificati e informativi per guidare l'encoder verso rappresentazioni ricche. Di conseguenza, la diversità delle rappresentazioni apprese si riduce drasticamente.
• Soluzioni attuali insufficienti: La comunità ha risposto a questo problema sovraparametrizzando i set di prototipi (usandone molti di più del necessario) o aggiungendo regolarizzatori ad-hoc (come la regolarizzazione KoLeo). Tuttavia, queste soluzioni mitigano solo i sintomi senza affrontare la causa radice, aumentando i costi computazionali e non garantendo robustezza.

2. Diagnosi della Causa Radice

Gli autori identificano la causa principale del collasso nella ottimizzazione congiunta (joint optimization) dell'encoder e dei prototipi sotto un'unica funzione di perdita condivisa.

• Meccanismo di "Shortcut Learning": Inizialmente, l'ottimizzazione congiunta spinge i prototipi verso rappresentazioni ridondanti che minimizzano rapidamente la perdita (loss) senza necessariamente migliorare la diversità o la qualità semantica delle rappresentazioni dell'encoder. I prototipi "scorrono" verso configurazioni facili da ottimizzare ma povere di informazione.
• Evidenza Empirica: L'analisi mostra che questo collasso non è limitato alla famiglia DINO, ma è un fenomeno diffuso in molti framework SSL basati su prototipi, specialmente nelle formulazioni a livello di istanza.

3. Metodologia: Apprendimento Completamente Disaccoppiato

Per rompere il ciclo di ottimizzazione congiunta, gli autori propongono una strategia di addestramento completamente disaccoppiato (fully decoupled training).

• Separazione degli Obiettivi: Invece di ottimizzare encoder ($f_\theta$) e prototipi ($C$) insieme, il metodo separa i due processi:

  1. Stima dei Prototipi: I prototipi vengono aggiornati indipendentemente dalla perdita dell'encoder, risolvendo un problema di stima non supervisionato sui feature latenti.
  2. Aggiornamento dell'Encoder: L'encoder viene aggiornato con i prototipi fissi per minimizzare la perdita di consistenza tra vista studente e insegnante.

• Implementazione Tecnica (GMM Online):

  • I prototipi sono modellati come le medie di un Modello a Mixture Gaussiana (GMM).
  • L'aggiornamento avviene tramite una procedura di tipo Expectation-Maximization (EM) online.
  • Vantaggi del GMM Online:
    • Aggiorna i componenti della mixture in modo incrementale man mano che arrivano nuovi dati, evitando di basarsi solo su mini-batch isolati (che causerebbero rumore).
    • È computazionalmente efficiente rispetto al ri-clustering completo (es. K-Means) ad ogni iterazione.
    • Include meccanismi di "forgetting" pesati sulla responsabilità e "annealing" per gestire spazi ad alta dimensionalità e grandi numeri di prototipi, prevenendo lo spostamento dei componenti.

4. Risultati Chiave

Gli esperimenti sono stati condotti su dataset come ImageNet-1k e iNaturalist-2018, confrontando il metodo proposto (es. su CARP + Decoupling) con baseline come DINO, DINOv2, CARP e varianti con regolarizzazione KoLeo.

• Eliminazione del Collasso: Il metodo disaccoppiato elimina completamente il collasso parziale dei prototipi. Mentre i metodi tradizionali (come DINO) mantengono solo l'1,5% - 40% dei prototipi unici, la versione disaccoppiata mantiene il 100% dei prototipi unici, indipendentemente dalla soglia di distanza ($\epsilon$) utilizzata per definire l'unicità.
• Diversità delle Rappresentazioni: La distribuzione dei prototipi rimane uniforme e diversificata per tutta la durata dell'addestramento, come mostrato dalle visualizzazioni KDE e dalle distribuzioni angolari.
• Prestazioni a valle (Downstream):
  • Classificazione Lineare e k-NN: Si osservano miglioramenti consistenti nella qualità delle rappresentazioni (es. +1.8% in k-NN su ImageNet per CARP).
  • Robustezza su Dati Squilibrati: Su iNaturalist-2018 (dataset a coda lunga), il metodo disaccoppiato migliora le prestazioni su tutte le categorie (Head, Medium, Tail), dimostrando che una maggiore diversità dei prototipi aiuta a modellare meglio le classi rare (tail), a differenza di alcune regolarizzazioni che possono destabilizzare l'addestramento su architetture sensibili come DINO.
• Efficienza: Il metodo non introduce un sovraccarico computazionale significativo (aumento di soli 0.2 ore di tempo di addestramento) e, sorprendentemente, riduce l'uso di memoria GPU (fino a 1.5GB in meno per batch size di 2048) perché non deve conservare il grafo di calcolo per l'ottimizzazione congiunta dei prototipi.

5. Significato e Contributi

Questo lavoro offre un contributo fondamentale alla comprensione teorica e pratica dell'SSL:

1. Diagnosi Sistematica: È il primo studio sistematico che identifica l'ottimizzazione congiunta come la causa principale del collasso dei prototipi, estendendo il problema ben oltre la famiglia DINO.
2. Soluzione Principale: Introduce un framework di addestramento disaccoppiato che risolve il problema alla radice senza bisogno di regolarizzatori ad-hoc o iperparametri delicati per bilanciare la diversità.
3. Impatto Pratico: Dimostra che la diversità dei prototipi è un indicatore chiave per la robustezza e la qualità delle rappresentazioni, specialmente in scenari di dati reali e sbilanciati.
4. Efficienza: Offre un approccio che migliora le prestazioni riducendo al contempo il consumo di memoria, rendendo l'addestramento di modelli SSL più scalabile.

In sintesi, il paper dimostra che separare l'apprendimento dei prototipi da quello dell'encoder non solo previene il collasso, ma genera rappresentazioni più ricche, robuste e trasferibili, offrendo una nuova direzione per la progettazione di algoritmi di apprendimento auto-supervisionato.