Sintesi Tecnica: WeightCLIP

Problematica

L'apprendimento nello spazio dei pesi (Weight Space Learning - WSL) tratta i parametri delle reti neurali (NN) come una modalità di dati per l'apprendimento di rappresentazioni, abilitando compiti quali la predizione delle proprietà dei modelli partendo dai pesi o la sintesi di nuovi modelli. Sebbene gli approcci esistenti abbiano mostrato potenziale, essi soffrono di un limite critico: le rappresentazioni apprese nello spazio dei pesi mancano di un esplicito riferimento semantico. Di conseguenza, gli spazi latenti addestrati puramente per la ricostruzione o la predizione di proprietà sono difficili da navigare; la prossimità nello spazio latente non corrisponde necessariamente alla variazione a livello di dataset, e la "navigazione" non si mappa intuitivamente su specifiche distribuzioni di dati. Ciò limita la capacità di eseguire il trasferimento zero-shot, la generazione controllata di pesi o il raffinamento efficace senza un costoso campionamento nell'alto-dimensionale spazio dei pesi.

Gli autori postulano che i modelli addestrati portino una "didascalia" implicita: il dataset su cui sono stati addestrati. Il problema centrale affrontato è come sfruttare questa informazione sul dataset per rimodellare la rappresentazione dello spazio dei pesi, creando uno spazio latente condiviso dove i dataset fungono da segnali di supervisione naturali, analogamente a come il testo guida i modelli visione-linguaggio.

Metodologia: WeightCLIP

WeightCLIP propone un metodo per allineare le rappresentazioni dei dataset con le rappresentazioni dei pesi del modello utilizzando un obiettivo di apprendimento contrastivo. Il framework è composto da tre componenti primarie:

1. Autoencoder dello Spazio dei Pesi: Seguendo il lavoro precedente (Schürholt et al., 2024), le reti neurali sono rappresentate come sequenze di "token di peso" (sottoinsiemi appiattiti di parametri di strato). Un encoder $g_\phi$ comprime questi token in rappresentazioni latenti, e un decoder $h_\psi$ ricostruisce i pesi.
2. Encoder del Dataset: Un encoder del dataset $d_\theta$ (basato sull'architettura DeepSets) mappa un piccolo sottoinsieme di immagini da un dataset in un embedding del dataset $e_D$.
3. Allineamento Contrastivo: L'innovazione principale è una perdita contrastiva a livello di token che allinea le rappresentazioni latenti dei pesi del modello con gli embedding dei dataset. A differenza dei metodi che apprendono uno spazio di proiezione separato, WeightCLIP rimodella direttamente lo spazio latente dell'autoencoder. La funzione di perdita include un termine di ricostruzione ($L_{recon}$) e un termine di allineamento contrastivo bidirezionale ($L_{align}$) che massimizza la similarità tra le rappresentazioni dei token del modello e l'embedding del suo dataset di addestramento, minimizzando al contempo la similarità con altri dataset.

Componenti a valle (Downstream)

Una volta allineato lo spazio latente, WeightCLIP introduce due meccanismi per i compiti a valle:

• Mappatura Dataset-Modello: Per generare modelli partendo da un nuovo prompt di dataset, un mapper leggero traduce l'embedding del dataset $e_D$ in una sequenza di rappresentazioni dei token del modello $Z$. Gli autori propongono due mapper:
  • Linear Mapper: Una semplice trasformazione affine che regredisce la sequenza di token dal dataset embedding.
  • Memory-Bank Mapper: Un metodo che polarizza le predizioni verso il manifold dei modelli reali attraverso la sintesi di token tramite una combinazione convessa di rappresentazioni dei token di addestramento memorizzate, pesate da una distribuzione di vicinato "soft".
• Raffinamento Latente: Per affrontare i disallineamenti (ad esempio, differenze nel classifier head) dopo la generazione, gli autori propongono di raffinare il modello direttamente nello spazio latente piuttosto che effettuare il fine-tuning dei pesi. Questo processo prevede:
  1. Prior Globale: Proiettare i token latenti iniziali su un guscio ipersferico (stimato dai dati di addestramento) per garantire che giacciano sul manifold valido del modello.
  2. Ottimizzazione Locale: Ottimizzare i token latenti $Z$ per minimizzare la perdita del task su un piccolo dataset target, con vincoli che mantengono gli aggiornamenti vicini al manifold appreso. Questo è strutturalmente distinto dal classico fine-tuning dei pesi nello spazio dei pesi.

Contributi Chiave

Il documento presenta tre contributi primari:

1. Spazio Latente Allineato ai Dataset: Gli autori dimostrano che l'uso dei dataset come punti di riferimento tramite allineamento contrastivo organizza lo spazio latente dei pesi. Ciò consente una navigazione intuitiva in cui i modelli addestrati su dataset simili si raggruppano insieme.
2. Generazione di Pesi Guidata dal Dataset: Introducono metodi per mappare un embedding di dati (prompt di dati) direttamente in una sequenza di embedding di modelli, che possono essere decodificati in pesi di reti neurali. Ciò consente la generazione di modelli specificamente adattati a un target dataset, inclusi dataset out-of-distribution (OOD).
3. Raffinamento del Manifold Latente: Propongono una procedura di raffinamento che ottimizza la perdita del task attraverso il decoder, vincolando al contempo i latenti a un guscio ipersferico e a un vicinato locale. Ciò consente un adattamento al tempo di test (test-time adaptation) efficace che spesso supera il fine-tuning standard dei pesi generati.

Risultati Sperimentali

Gli autori valutano WeightCLIP su due zoo di modelli: una vasta collezione di 10.000 checkpoint CNN e 1.000 checkpoint ResNet-18, addestrati su diversi dataset di classificazione di immagini.

• Organizzazione dello Spazio Latente: Le visualizzazioni t-SNE e la classificazione k-NN mostrano che l'allineamento raggruppa significativamente le rappresentazioni dei modelli per il loro dataset di addestramento (migliorando l'accuratezza della classificazione dal 77% al 98% rispetto alle baseline non allineate). L'interpolazione tra i prompt dei dataset produce trade-off fluidi nelle prestazioni del modello, confermando la navigabilità.
• Retrieval: WeightCLIP migliora significativamente il recupero dataset-modello. In contesti in-distribution, raggiunge il 95% di Recall@1, superando sia l'addestramento non allineato che la baseline TANS. In contesti OOD, i modelli recuperati ottengono un'accuratezza zero-shot e di fine-tuning superiore rispetto alle baseline.
• Generazione: Quando mappano i prompt dei dataset per generare nuovi modelli, WeightCLIP supera l'addestramento da zero, TANS e le baseline di hypernetwork (Text2Model). Anche all'inizializzazione (0 epoche di fine-tuning), i modelli generati mostrano un forte allineamento con il target task. Il Memory-Bank Mapper generalmente fornisce le migliori prestazioni iniziali, mentre il Linear Mapper rimane competitivo.
• Raffinamento: Il raffinamento latente raggiunge prestazioni comparabili o superiori al fine-tuning standard nello spazio dei pesi utilizzando lo stesso budget computazionale. Gli studi di ablazione confermano che sia il vincolo ipersferico che l'allineamento iniziale sono necessari per questi guadagni.

Significato e Rivendicazioni

Il paper sostiene che WeightCLIP affronti un limite chiave dei metodi WSL esistenti incorporando esplicitamente le informazioni sul dataset per rimodellare la rappresentazione dello spazio dei pesi. Trattando i dataset come un riferimento semantico, il metodo abilita:

• Navigazione Strutturata: Lo spazio latente diventa interrogabile e navigabile in base alla semantica del dataset.
• Sintesi Efficiente di Modelli: Permette la generazione di forti inizializzazioni di modelli per nuovi dataset senza campionamento costoso o addestramento da zero.
• Adattamento Migliorato: Il raffinamento latente proposto offre un'alternativa strutturalmente diversa e spesso più efficace al fine-tuning standard.

Gli autori posizionano questo lavoro come un passo verso il rendere l'addestramento delle reti neurali più efficiente in termini di tempo e computazione per architetture note su nuovi dataset. Riconoscono i limiti, notando che l'attuale ambito è limitato ai dataset di computer vision e ad architetture di modelli omogenee, e che il metodo si affida ad aumentazioni di addestramento per gestire le simmetrie dello spazio dei pesi invece di imporle esplicitamente.