Resumen Técnico: x-Prediction Is All You Need

Planteamiento del Problema

Los modelos de difusión y de ajuste de flujo (flow matching) han establecido un rendimiento de vanguardia en tareas generativas, pero sufren de altos costes de inferencia, requiriendo típicamente decenas o cientos de evaluaciones de la función neuronal (NFEs) por muestra. Las estrategias de aceleración existentes —como la rectificación de trayectoria (Rectified Flow), la destilación (progresiva, basada en consistencia, bootstrapping) y los solvers de ODE avanzados— generalmente requieren modificaciones en la receta de entrenamiento, la arquitectura del modelo o la propia trayectoria. En consecuencia, estos métodos suelen ser incompatibles con la vasta biblioteca de checkpoints preentrenados existentes (p. ej., SDXL, SD3.5M, FLUX) sin incurrir en costes de reentrenamiento. El campo carece de una estrategia de aceleración libre de entrenamiento que funcione con los modelos "tal como son".

Metodología: Truncated Jump Sampling (TJS)

El artículo propone el Truncated Jump Sampling (TJS), una estrategia de inferencia libre de entrenamiento basada en el concepto de decodificabilidad del punto final (endpoint decodability).

Idea Central: Decodificabilidad del Punto Final

Los autores formalizan que para cualquier trayectoria de probabilidad afín no degenerada (la estructura estándar para los modelos de difusión y ajuste de flujo definida como $x_t = \alpha_t x_0 + \sigma_t \epsilon$), la muestra limpia $x_0$ puede recuperarse de forma única de cualquier estado intermedio $x_t$ y su velocidad de trayectoria $u_t$ mediante un decodificador de forma cerrada.

• Recuperación Algebraica: Dado el sistema lineal acoplado de posición y velocidad, $x_0$ se recupera como $x_0 = (\sigma_t u_t - \dot{\sigma}_t x_t) / \Delta_t$, donde $\Delta_t$ es el determinante de la trayectoria.
• Optimalidad: Bajo los objetivos de entrenamiento $\ell_2$ estándar, este decodificador algebraico corresponde al estimador de Error Cuadrático Medio Mínimo (MMSE) $E[x_0 | x_t]$. Esto implica que los modelos preentrenados estándar aprenden implícitamente a predecir el punto final $x_0$ en cada paso de tiempo, incluso si están entrenados para predecir el ruido o la velocidad.
• Independencia de la Curvatura: Crucialmente, el error de esta predicción del punto final depende únicamente de la precisión de estimación del modelo y de la incertidumbre irreducible del estado $x_t$, no de la curvatura de la trayectoria ($\|\ddot{x}_t\|$). Esto desafía la suposición prevalente de que las trayectorias rectas (como en Rectified Flow) son necesarias para una generación precisa de pocos pasos.

El Algoritmo

TJS operacionaliza esta idea deteniendo la integración de la ODE prematuramente en un tiempo $t^*$ y emitiendo inmediatamente el $x_0$ decodificado en lugar de continuar hasta $t=0$.

1. Integración Parcial: El modelo integra la ODE para una fracción $\gamma$ de los pasos totales ($k^* = \lceil \gamma K \rceil$).
2. Salida Temprana (Early Exit): En el paso $k^*$, el algoritmo calcula la estimación del punto final $\hat{x}_0$ utilizando el campo de velocidad aprendido y el decodificador de forma cerrada.
3. Terminación: El proceso se detiene, devolviendo $\hat{x}_0$ como la muestra final.

Este enfoque requiere cero reentrenamiento, cero destilación y cero cambios en la arquitectura. Es compatible con cualquier solver de ODE, escala de guía de clasificación libre de clase (CFG), esquema de ruido o checkpoint existente.

Contribuciones Clave

1. Formalización Teórica: El artículo establece la Decodificabilidad del Punto Final (Teorema 5) como una propiedad universal de las trayectorias afines. Demuestra que el decodificador inducido es el estimador óptimo de MMSE (Teorema 6) y que la descomposición del error es independiente de la curvatura de la trayectoria (Teorema 9).
2. Desafío a la Rectificación de Trayectoria: Los autores demuestran que las trayectorias rectas son suficientes pero no necesarias para la predicción precisa del punto final (Proposición 12). Esto desacopla la aceleración de la inferencia de la necesidad de rectificación de trayectoria, ofreciendo una base teórica de por qué TJS funciona en trayectorias curvas donde la integración de Euler fallaría.
3. Algoritmo Libre de Entrenamiento (TJS): Un sampler práctico que reduce las NFEs entre un 20 y un 70% en diversas familias de modelos sin modificar el modelo ni los datos de entrenamiento.
4. Validación Empírica Exhaustiva: Experimentos extensos en seis familias de modelos (SDXL, SD3.5M, Z-Image-Turbo, ImageNet-256, CIFAR-10, MNIST) demuestran que TJS logra una calidad casi idéntica o superior a las trayectorias completas de la ODE.

Resultados Experimentales

Los autores evaluaron TJS en benchmarks de condicionalidad de clase y texto-a-imagen:

• Reducción de NFE: TJS reduce las NFEs entre un 20 y un 70% manteniendo una calidad casi equivalente.
  • SDXL/SD3.5M: Logra el 95% de la calidad de la ODE completa (medido por ImageReward) con aproximadamente un 37–40% de ahorro de NFE (deteniéndose en $k^* \approx 18$ de 30 pasos).
  • Z-Image-Turbo: En un modelo destilado de 10 pasos, TJS logra una retención de calidad del 95% con un ahorro adicional del 70% (deteniéndose en $k^*=2$, es decir, 3 NFEs en total), demostrando que la decodificabilidad del punto final persiste incluso en trayectorias destiladas.
  • Condicionalidad de Clase: En MNIST y CIFAR-10, TJS a veces supera a la ODE completa (el fenómeno de "sobreimpulso" o overshoot) porque evita los errores de discretización acumulados en los pasos finales de la integración de Euler.
• Jerarquía de Métricas: Los experimentos revelan una jerarquía de convergencia donde el contenido semántico (CLIP, PickScore) se resuelve temprano (a menudo para $k^* \approx 12$), mientras que los detalles estéticos finos (ImageReward, AES) requieren una integración más profunda.
• Robustez: El rendimiento de TJS es invariante a la elección del solver de ODE (DDIM, DPM++, UniPC, etc.), el esquema de ruido (Beta, Kasarras, etc.) y la escala de CFG, confirmando su ortogonalidad a los componentes de inferencia existentes.

Significado y Reivindicaciones

El artículo sostiene que la aceleración moderada de pocos pasos no requiere rectificación de trayectoria, destilación o rediseño de entrenamiento. La importancia radica en reconocer la decodificabilidad del punto final como una propiedad estructural inherente de los modelos generativos en trayectorias afines.

• Impacto Práctico: TJS proporciona una vía de aceleración inmediata y sin coste para la biblioteca de checkpoints ajustados (fine-tuned) existentes de la comunidad, que a menudo es inviable reentrenar o destilar.
• Cambio Teórico: Al demostrar que las trayectorias rectas no son necesarias para la predicción del punto final, este trabajo desafía la motivación fundamental de Rectified Flow y los Modelos de Consistencia respecto a la aceleración de la inferencia.
• Naturaleza Complementaria: Los autores posicionan a TJS como ortogonal a la destilación. Mientras que la destilación comprime la longitud total de la trayectoria, TJS elimina la cola redundante de cualquier trayectoria restante, permitiendo una aceleración aditiva (por ejemplo, aplicando TJS a un modelo destilado como Z-Image-Turbo).

El artículo concluye que reconocer esta propiedad "no cuesta nada" e ignorarla deja rendimiento sobre la mesa, ofreciendo una forma fundamentada de intercambiar profundidad de integración por velocidad de inferencia sin sacrificar la calidad de la muestra generada.