Sintesi Tecnica: x-Prediction Is All You Need

Problematica

I modelli di diffusione e flow matching hanno stabilito prestazioni allo stato dell'arte nei compiti generativi, ma soffrono di elevati costi di inferenza, richiedendo tipicamente da decine a centinaia di valutazioni della funzione neurale (NFE) per campione. Le attuali strategie di accelerazione — come la rettifica della traiettoria (Rectified Flow), la distillazione (progressiva, basata sulla coerenza, bootstrapping) e gli avanzati solver ODE — richiedono generalmente modifiche alla ricetta di addestramento, all'architettura del modello o alla traiettoria stessa. Di conseguenza, tali metodi sono spesso incompatibili con la vasta libreria di checkpoint pre-addestrati esistenti (es. SDXL, SD3.5M, FLUX) senza incorrere in costi di ri-addestramento. Il campo manca di una strategia di accelerazione "training-free" che funzioni sui modelli "così come sono".

Metodologia: Truncated Jump Sampling (TJS)

Il paper propone il Truncated Jump Sampling (TJS), una strategia di inferenza training-free basata sul concetto di endpoint decodability (decodificabilità del punto finale).

Intuizione Centrale: Endpoint Decodability

Gli autori formalizzano che per qualsiasi percorso di probabilità affine non degenere (la struttura standard per i modelli di diffusione e flow matching definita come $x_t = \alpha_t x_0 + \sigma_t \epsilon$), il campione pulito $x_0$ può essere recuperato univocamente da qualsiasi stato intermedio $x_t$ e dalla sua velocità di percorso $u_t$ tramite un decoder in forma chiusa.

• Recupero Algebrico: Dato il sistema accoppiato lineare di posizione e velocità, $x_0$ viene recuperato come $x_0 = (\sigma_t u_t - \dot{\sigma}_t x_t) / \Delta_t$, dove $\Delta_t$ è il determinante del percorso.
• Ottimalità: Sotto gli standard degli obiettivi di addestramento $\ell_2$, questo decoder algebrico corrisponde allo stimatore MMSE (Minimum Mean Square Error) $E[x_0 | x_t]$. Ciò implica che i modelli pre-addestrati standard apprendono implicitamente a predire l'endpoint $x_0$ ad ogni timestep, anche se sono addestrati per predire rumore o velocità.
• Indipendenza dalla Curvatura: Fondamentalmente, l'errore di questa predizione dell'endpoint dipende solo dall'accuratezza di stima del modello e dall'incertezza irriducibile dello stato $x_t$, non dalla curvatura della traiettoria ($\|\ddot{x}_t\|$). Questo sfida l'assunto prevalente che traiettorie dritte (come in Rectified Flow) siano necessarie per una generazione accurata in pochi step.

L'Algoritmo

TJS operazionalizza questa intuizione interrompendo l'integrazione ODE precocemente a un tempo $t^*$ e restituendo immediatamente il $x_0$ decodificato invece di continuare fino a $t=0$.

1. Integrazione Parziale: Il modello integra l'ODE per una frazione $\gamma$ dei passi totali ($k^* = \lceil \gamma K \rceil$).
2. Early Exit (Uscita Anticipata): Al passo $k^*$, l'algoritmo calcola la stima dell'endpoint $\hat{x}_0$ utilizzando il campo di velocità appreso e il decoder in forma chiusa.
3. Terminazione: Il processo si interrompe, restituendo $\hat{x}_0$ come campione finale.

Questo approccio richiede zero ri-addestramento, zero distillazione e zero modifiche all'architettura. È compatibile con qualsiasi solver ODE, schema di rumore (noise schedule), scala di classifier-free guidance (CFG) o checkpoint esistente.

Contributi Chiave

1. Formalizzazione Teorica: Il paper stabilisce l'Endpoint Decodability (Teorema 5) come una proprietà universale dei percorsi affini. Dimostra che il decoder indotto è lo stimatore MMSE-ottimale (Teorema 6) e che la decomposizione dell'errore è indipendente dalla curvatura della traiettoria (Teorema 9).
2. Sfida alla Rettifica della Traiettoria: Gli autori dimostrano che le traiettorie dritte sono sufficienti ma non necessarie per una precisa predizione dell'endpoint (Proposizione 12). Ciò decoppia l'accelerazione dell'inferenza dalla necessità di rettifica della traiettoria, offrendo una base teorica al motivo per cui TJS funziona su percorsi curvi dove l'integrazione di Eulero fallirebbe.
3. Algoritmo Training-Free (TJS): Un campionatore pratico che riduce le NFE del 20-70% attraverso diverse famiglie di modelli senza modificare il modello o i dati di addestramento.
4. Validazione Empirica Completa: Esperimenti estesi su sei famiglie di modelli (SDXL, SD3.5M, Z-Image-Turbo, ImageNet-256, CIFAR-10, MNIST) dimostrano che TJS raggiunge una qualità quasi equivalente o superiore rispetto alle traiettorie ODE complete.

Risultati Sperimentali

Gli autori hanno valutato TJS attraverso benchmark class-conditional e text-to-image:

• Riduzione NFE: TJS riduce le NFE del 20–70% mantenendo una qualità quasi identica.
  • SDXL/SD3.5M: Raggiunge il 95% della qualità dell'ODE completa (misurata tramite ImageReward) con circa il 37–40% di risparmio NFE (fermandosi a $k^* \approx 18$ di 30 step).
  • Z-Image-Turbo: Su un modello distillato a 10 step, TJS ottiene una ritenzione della qualità del 95% con un ulteriore risparmio del 70% (fermandosi a $k^*=2$, ovvero 3 NFE totali), dimostrando che la decodificabilità dell'endpoint persiste anche nelle traiettorie distillate.
  • Class-Conditional: Su MNIST e CIFAR-10, TJS talvolta supera l'ODE completa (il fenomeno dell' "overshoot") perché evita gli errori di discretizzazione accumulati negli ultimi step dell'integrazione di Eulero.
• Gerarchia delle Metriche: Gli esperimenti rivelano una gerarchia di convergenza in cui il contenuto semantico (CLIP, PickScore) si risolve precocemente (spesso entro $k^* \approx 12$), mentre i dettagli estetici fini (ImageReward, AES) richiedono un'integrazione più profonda.
• Robustezza: Le prestazioni di TJS sono invarianti rispetto alla scelta del solver ODE (DDIM, DPM++, UniPC, ecc.), dello schema di rumore (Beta, Karras, ecc.) e della scala CFG, confermando la sua ortogonalità ai componenti di inferenza esistenti.

Significato e Rivendicazioni

Il paper sostiene che un'accelerazione moderata di pochi step non richiede la rettifica della traiettoria, la distillazione o il ri-progettazione dell'addestramento. Il significato risiede nel riconoscere l'endpoint decodability come una proprietà strutturale inerente ai modelli generativi su percorsi affini.

• Impatto Pratico: TJS fornisce un percorso di accelerazione immediato e a costo zero per la libreria di checkpoint fine-tuned esistente della comunità, che spesso non è percorribile tramite ri-addestramento o distillazione.
• Svolta Teorica: Dimostrando che le traiettorie dritte non sono necessarie per la predizione dell'endpoint, il lavoro sfida la motivazione fondamentale di Rectified Flow e dei Modelli di Coerenza riguardo all'accelerazione dell'inferenza.
• Natura Complementare: Gli autori posizionano TJS come ortogonale alla distillazione. Mentre la distillazione comprime la lunghezza totale della traiettoria, TJS elimina la coda ridondante di qualunque traiettoria rimanga, permettendo un'accelerazione additiva (ad esempio, applicando TJS a un modello distillato come Z-Image-Turbo).

Il paper conclude che riconoscere questa proprietà "non costa nulla" e che ignorarla significa lasciare prestazioni inutilizzate, offrendo un modo rigoroso per scambiare la profondità di integrazione con la velocità di inferenza senza sacrificare la qualità del campione generato.