Sintesi Tecnica: Distillazione di Transformer Drifting con Representation Autoencoders

Definizione del Problema

Sebbene i Representation Autoencoders (RAE) abbiano migliorato significativamente le prestazioni dei modelli generativi basati su flow, utilizzando spazi latenti semanticamente ricchi e clusterizzati per etichetta derivati da feature DINO pre-addestrate, la distillazione di questi modelli in generatori efficienti a uno o pochi step rimane una sfida.

Gli autori identificano un disallineamento geometrico fondamentale: gli spazi latenti RAE esibiscono una severa anisotropia e alte curvature nelle loro traiettorie di Ordinaria Equazione Differenziale (ODE) sottostanti rispetto agli spazi dei tradizionali Variational Autoencoder (VAE). I metodi di distillazione basati sulle traiettorie esistenti si affidano all'assunzione di traiettorie del docente lisce e quasi rettilinee. Le traiettorie altamente curve negli spazi RAE causano instabilità o inefficienza in tali metodi. Inoltre, il recentemente proposto Drifting Model offre un approccio a livello di distribuzione che evita il matching di traiettorie curve, ma la sua formulazione originale richiede un Masked Autoencoder (MAE) ausiliario, addestrato separatamente, per fornire feature concentrate, introducendo un significativo overhead computazionale.

Metodologia

Il paper propone Drift-RAE, un framework di distillazione che adatta il paradigma del Drifting Model specificamente agli spazi latenti RAE, eliminando la necessità di estrattori di feature ausiliari.

1. Analisi Teorica di Geometria e Dinamica

Gli autori analizzano quantitativamente gli spazi latenti, dimostrando che le feature RAE sono sostanzialmente più anisotrope (minore Participation Ratio e Spectral Entropy) rispetto alle feature SD-VAE. Dimostrano teoricamente (Teorema 1) che in spazi ad alta dimensione, se i campioni positivi sono eccessivamente dispersi, la componente attrattiva del campo di drift tende a svanire, portando a un campo degenerato.

• Insight Chiave: Mentre l'originale Drifting Model richiede un MAE per concentrare le feature disperse, gli spazi latenti RAE sono naturalmente concentrati dal punto di vista semantico. Ciò consente al Drifting Model di operare direttamente sui latenti RAE senza il modulo MAE ausiliario.

2. Il Framework Drift-RAE

Il metodo distilla un modello di flow pre-addestrato $v_\theta$ in un generatore a un solo step $G_\theta(z) = z - v_\theta(z, 1)$. Invece di eseguire il matching delle traiettorie ODE, il generatore è addestrato per minimizzare un obiettivo di drifting che spinge i campioni generati verso le distribuzioni dei dati reali e lontano dagli altri campioni generati.

• Obiettivo: La perdita è calcolata token-wise sulle rappresentazioni latenti RAE, utilizzando un campo di drifting normalizzato $\tilde{V}$ derivato da campioni positivi (reali) e negativi (generati).

3. Modifiche Pratiche per la Stabilità

Per allineare il Drifting Model al framework Diffusion-GAN e migliorare la stabilità dell'addestramento, gli autori introducono tre modifiche specifiche:

• Dimensione Softmax: Gli autori rimuovono il trucco del softmax bi-direzionale utilizzato nel Drifting Model originale. Mantengono un singolo softmax sugli indici dei campioni, il che si allinea meglio con la direzione del gradiente del potenziale log-sum-exp sottostante.
• Perturbazione dell'Input: Per prevenire la scomparsa del gradiente (vanishing gradient) causata da manifold non intersecanti tra dati reali e generati, i campioni generati vengono perturbati con rumore campionato da una distribuzione di Laplace prima di calcolare il campo di drifting.
• Distacco Parziale dei Campioni Negativi: Per aumentare il numero di campioni negativi utilizzati per l'approssimazione della distribuzione senza superare i budget di memoria GPU, gli autori propongono l'uso di un pool più ampio di campioni negativi ($N_{total\_neg}$) dove solo un sottoinsieme ($N_{neg}$) trasporta i gradienti, mentre il resto è staccato (detached). Ciò migliora l'approssimazione empirica della distribuzione generata.

Contributi Chiave

1. Analisi Geometrica: Una dimostrazione quantitativa che gli spazi latenti RAE possiedono un'alta curvatura e anisotropia, rendendo instabile la distillazione basata sulle traiettorie standard.
2. Giustificazione Teorica: Una prova che mostra come i Drifting Models falliscano su spazi eccessivamente dispersi (necessitando un MAE nella formulazione originale) ma prosperino negli spazi semantici naturalmente concentrati degli RAE.
3. Drift-RAE: Un nuovo metodo di distillazione che raggiunge lo stato dell'arte nella generazione a un solo step negli spazi RAE senza richiedere un estrattore di feature MAE ausiliario.
4. Miglioramenti della Stabilità: Insight teorici ed empirici nell'aggiornamento dell'obiettivo di Drifting (dimensione softmax, perturbazione dell'input e distacco dei campioni) per allinearsi ai principi di Diffusion-GAN e migliorare la convergenza.

Risultati Sperimentali

Il metodo è stato valutato sul dataset ImageNet 256 $\times$ 256 con backbone DiTDH-XL e DiTDH-L.

• Performance: Drift-RAE raggiunge un FID di 1.77 (con DiTDH-XL) e 2.12 (con DiTDH-L) utilizzando un singolo step di generazione.
• Efficienza: Il modello raggiunge questi risultati in soli 10.000 step di distillazione (circa 16 epoche).
• Confronto:
  • Supera i precedenti metodi di distillazione negli spazi RAE (es. MF-RAE).
  • Ottiene un FID comparabile e un punteggio FDDINOv2 migliorato rispetto all'originale Drifting Model, nonostante non utilizzi l'ausiliario e costoso MAE.
  • Rimane competitivo con i generatori a uno o pochi step addestrati in altri spazi latenti (es. Spazio dei Pixel, SD-VAE).

Significato e Rivendicazioni

Il paper sostiene che il Drifting fornisca un framework efficace e promettente per distillare modelli generativi negli spazi di rappresentazione. Il significato primario risiede nel risolvere l'instabilità della distillazione dei modelli basati su RAE sfruttando la loro inerente concentrazione semantica, eliminando così la dipendenza da estrattori di feature ausiliari.

Gli autori notano con modestia i limiti, riconoscendo che la loro analisi teorica si basa su modelli semplificati ad alta dimensione e che l'addestramento dei Drifting Models da zero (senza distillazione) negli spazi RAE rimane un problema aperto. Evidenziano inoltre come il requisito di molti campioni positivi della stessa classe possa porre sfide di scalabilità per la generazione text-to-image, suggerendo questo come una direzione per la ricerca futura.