Résumé Technique : x-Prediction Is All You Need

Énoncé du Problème

Les modèles de diffusion et de flow matching ont établi des performances de pointe dans les tâches génératives, mais souffrent de coûts d'inférence élevés, nécessitant typiquement des dizaines ou des centaines d'évaluations de fonctions neuronales (NFE) par échantillon. Les stratégies d'accélération existantes — telles que la rectification de trajectoire (Rectified Flow), la distillation (progressive, basée sur la cohérence, bootstrapping) et les solveurs d'ODE avancés — nécessitent généralement des modifications de la recette d'entraînement, de l'architecture du modèle ou de la trajectoire elle-même. Par conséquent, ces méthodes sont souvent incompatibles avec la vaste bibliothèque de points de contrôle pré-entraînés existants (ex: SDXL, SD3.5M, FLUX) sans engendrer de coûts de réentraînement. Le domaine manque d'une stratégie d'accélération sans entraînement qui fonctionne sur les modèles « tels qu'ils sont ».

Méthodologie : Truncated Jump Sampling (TJS)

Le papier propose le Truncated Jump Sampling (TJS), une stratégie d'inférence sans entraînement basée sur le concept de décodabilité de l'extrémité (endpoint decodability).

Insight Fondamental : Décodabilité de l'Extrémité

Les auteurs formalisent le fait que pour tout chemin de probabilité affine non dégénéré (la structure standard pour les modèles de diffusion et de flow matching définie comme $x_t = \alpha_t x_0 + \sigma_t \epsilon$), l'échantillon propre $x_0$ peut être récupéré de manière unique à partir de tout état intermédiaire $x_t$ et de sa vitesse de trajectoire $u_t$ via un décodeur en forme close.

• Récupération Algébrique : Étant donné le système linéaire couplé de position et de vitesse, $x_0$ est récupéré sous la forme $x_0 = (\sigma_t u_t - \dot{\sigma}_t x_t) / \Delta_t$, où $\Delta_t$ est le déterminant du chemin.
• Optimalité : Sous les objectifs d'entraînement $\ell_2$ standards, ce décodeur algébrique correspond à l'estimateur de l'erreur quadratique moyenne minimale (MMSE) $E[x_0 | x_t]$. Cela implique que les modèles pré-entraînés standards apprennent implicitement à prédire l'extrémité $x_0$ à chaque pas de temps, même s'ils sont entraînés pour prédire le bruit ou la vitesse.
• Indépendance de la Courbure : Crucialement, l'erreur de cette prédiction de l'extrémité dépend uniquement de la précision d'estimation du modèle et de l'incertitude irréductible de l'état $x_t$, et non de la courbure de la trajectoire ($\|\ddot{x}_t\|$). Cela remet en question l'hypothion dominante selon laquelle des trajectoires droites (comme dans Rectified Flow) sont nécessaires pour une génération précise en peu d'étapes.

L'Algorithme

Le TJS opérationnalise cet insight en arrêtant l'intégration de l'ODE plus tôt à un temps $t^*$ et en produisant immédiatement l'échantillon décodé $x_0$ plutôt que de continuer jusqu'à $t=0$.

1. Intégration Partielle : Le modèle intègre l'ODE pour une fraction $\gamma$ du nombre total d'étapes ($k^* = \lceil \gamma K \rceil$).
2. Sortie Précoce (Early Exit) : À l'étape $k^*$, l'algorithme calcule l'estimation de l'extrémité $\hat{x}_0$ en utilisant le champ de vitesse appris et le décodeur en forme close.
3. Terminaison : Le processus s'arrête, retournant $\hat{x}_0$ comme échantillon final.

Cette approche ne nécessite aucun réentraînement, aucune distillation, et aucun changement d'architecture. Elle est compatible avec n'importe quel solveur d'ODE, échelle de guidance de classe libre (CFG), ou point de contrôle existant.

Contributions Clés

1. Formalisation Théorique : Le papier établit la décodabilité de l'extrémité (Théorème 5) comme une propriété universelle des chemins affines. Il prouve que le décodeur induit est l'estimateur optimal MMSE (Théorème 6) et que la décomposition de l'erreur est indépendante de la courbure de la trajectoire (Théorème 9).
2. Contestation de la Rectification de Trajectoire : Les auteurs prouvent que les trajectoires droites sont suffisantes mais non nécessaires pour une prédiction d'extrémité précise (Proposition 12). Cela découple l'accélération de l'inférence du besoin de rectification de trajectoire, offrant une base théorique expliquant pourquoi le TJS fonctionne sur des chemins courbes là où l'intégration d'Euler échouerait.
3. Algorithme Sans Entraînement (TJS) : Un échantillonneur pratique qui réduit les NFE de 20 à 70 % sur diverses familles de modèles sans modifier le modèle ou les données d'entraînement.
4. Validation Empirique Complète : Des expériences étendues sur six familles de modèles (SDXL, SD3.5M, Z-Image-Turbo, ImageNet-256, CIFAR-10, MNIST) démontrent que le TJS atteint une qualité proche ou supérieure à celle des trajectoires ODE complètes.

Résultats Expérimentaux

Les auteurs ont évalué le TJS sur des benchmarks de conditionnement de classe et de texte-vers-image :

• Réduction de NFE : Le TJS réduit les NFE de 20 à 70 % tout en maintenant une qualité quasi identique.
  • SDXL/SD3.5M : Atteint 95 % de la qualité de l'ODE complète (mesurée par ImageReward) avec environ 37-40 % d'économies de NFE (en s'arrêtant à $k^* \approx 18$ sur 30 étapes).
  • Z-Image-Turbo : Sur un modèle distillé de 10 étapes, le TJS atteint 95 % de rétention de qualité avec 70 % d'économies supplémentaires (en s'arrêtant à $k^*=2$, soit 3 NFE au total), démontrant que la décodabilité de l'extrémité persiste même dans les trajectoires distillées.
  • Conditionnement de Classe : Sur MNIST et CIFAR-10, le TJS surpasse parfois l'ODE complète (le phénomène de "dépassement" ou overshoot) car il évite les erreurs de discrétisation accumulées lors des dernières étapes de l'intégration d'Euler.
• Hiérarchie des Métriques : Les expériences révèlent une hiérarchie de convergence où le contenu sémantique (CLIP, PickScore) se résout tôt (souvent dès $k^* \approx 12$), tandis que les détails esthétiques fins (ImageReward, AES) nécessitent une intégration plus profonde.
• Robustesse : La performance du TJS est invariante au choix du solveur d'ODE (DDIM, DPM++, UniPC, etc.), du planning de bruit (Beta, Karras, etc.), et de l'échelle de CFG, confirmant son orthogonalité aux composants d'inférence existants.

Signification et Revendications

Le papier affirme que l'accélération modérée en peu d'étapes ne nécessite ni rectification de trajectoire, ni distillation, ni redéfinition de l'entraînement. La signification réside dans la reconnaissance de la décodabilité de l'extrémité comme une propriété structurelle inhérente aux modèles génératifs sur des chemins affines.

• Impact Pratique : Le TJS offre une voie d'accélération immédiate et sans coût pour la bibliothèque existante de points de contrôle affinés de la communauté, qui sont souvent impossibles à réentraîner ou à distiller.
• Changement Théorique : En prouvant que les trajectoires droites ne sont pas nécessaires pour la prédiction de l'extrémité, ce travail remet en question la motivation fondamentale de Rectified Flow et des Consistency Models concernant l'accélération de l'inférence.
• Nature Complémentaire : Les auteurs positionnent le TJS comme orthogonal à la distillation. Alors que la distillation compresse la longueur totale de la trajectoire, le TJS élimine la queue redondante de n'importe quelle trajectoire restante, permettant une accélération additive (ex: appliquer le TJS à un modèle distillé comme Z-Image-Turbo).

Le papier conclut que reconnaître cette propriété « ne coûte rien » et que l'ignorer laisse des performances inexploitées, offrant une manière fondée sur des principes de troquer la profondeur d'intégration contre la vitesse d'inférence sans sacrifier la qualité de l'échantillon généré.