Resumen Técnico: Modelado Generativo mediante Deriva (Drifting)

Planteamiento del Problema

El modelado generativo es fundamentalmente la tarea de aprender un mapeo $f$ tal que el pushforward de una distribución previa $p_{prior}$ (por ejemplo, ruido gaussiano) coincida con la distribución de los datos $p_{data}$, denotado como $f\#p_{prior} \approx p_{data}$. Mientras que los paradigmas predominantes como los Modelos de Difusión y el Flow Matching logran esto refinando iterativamente las muestras durante el tiempo de inferencia (descomponiendo el complejo mapeo de pushforward en una cadena de transformaciones factibles), este enfoque requiere múltiples evaluaciones de la red neuronal (NFE) por muestra, lo que conlleva altos costos computacionales.

Los intentos existentes para crear generadores de un solo paso suelen basarse en destilar modelos pre-entrenados de múltiples pasos o aproximar las trayectorias de ecuaciones diferenciales estocásticas (SDE) o de ecuaciones diferenciales ordinarias (ODE) durante el entrenamiento. Los autores proponen un nuevo paradigma que desplaza fundamentalmente el proceso iterativo del tiempo de inferencia al tiempo de entrenamiento, permitiendo una generación de alta calidad con una sola evaluación de la función neuronal (1-NFE) sin depender de formulaciones SDE/ODE.

Metodología: Modelos de Deriva (Drifting Models)

La contribución central es el Modelo de Deriva (Drifting Model), que formula el modelado generativo como la evolución de la distribución de pushforward durante el propio proceso de entrenamiento.

1. Evolución en el Tiempo de Entrenamiento

En lugar de ver el entrenamiento simplemente como la actualización de pesos para minimizar una pérdida estática, los autores ven la secuencia de modelos $\{f_i\}$ producidos por la optimización iterativa (por ejemplo, SGD) como una secuencia de distribuciones de pushforward $\{q_i\}$. En cada iteración, una muestra $x_i = f_i(\epsilon)$ es implícitamente "derivada" hacia $x_{i+1} = f_{i+1}(\epsilon)$ debido a las actualizaciones de los parámetros. El objetivo es definir un mecanismo donde esta deriva guíe naturalmente la distribución $q$ hacia $p_{data}$.

2. El Campo de Deriva (Drifting Field)

Para gobernar esta evolución, los autores introducen un campo de deriva $V_{p,q}(x)$, una función que computa un vector de desplazamiento $\Delta x$ para una muestra $x$ basado en la distribución generada actual $q$ y la distribución de datos objetivo $p$.

• Condición de Equilibrio: El campo está diseñado de tal manera que $V_{p,q}(x) = 0$ si y solo si $p = q$. Cuando las distribuciones coinciden, las muestras dejan de derivar, alcanzando un equilibrio.
• Anti-simetría: El campo satisface $V_{p,q}(x) = -V_{q,p}(x)$. Esta propiedad asegura que si las distribuciones se intercambian, la dirección de la deriva se invierte, garantizando que $V=0$ cuando $p=q$.

3. Instanciación mediante Atracción y Repulsión

El campo de deriva específico utilizado en el artículo está inspirado en el algoritmo de mean-shift y se construye utilizando muestras positivas de $p$ y muestras negativas de $q$:
$$V_{p,q}(x) = V^+_p(x) - V^-_q(x)$$
Donde:

• $V^+_p(x)$ es una fuerza atractiva que atrae a $x$ hacia la distribución de datos $p$.
• $V^-_q(x)$ es una fuerza repulsiva que aleja a $x$ de la distribución generada $q$.
• Estas fuerzas se computan utilizando una función de kernel $k(x, y)$ (específicamente un kernel exponencial basado en la distancia $\ell_2$) y se normalizan mediante operaciones softmax para asegurar la estabilidad.

4. Objetivo de Entrenamiento

El objetivo de entrenamiento se deriva de la relación de punto fijo en el equilibrio: $f(\epsilon) = f(\epsilon) + V_{p,q}(f(\epsilon))$. La función de pérdida minimiza la distancia entre la predicción actual y un "objetivo congelado" (la predicción más la deriva):
$$L = \mathbb{E}_{\epsilon} \left[ \| f_\theta(\epsilon) - \text{stopgrad}(f_\theta(\epsilon) + V_{p,q}(f_\theta(\epsilon))) \|^2 \right]$$
Esta formulación permite al optimizador evolucionar la distribución moviendo las muestras hacia sus versiones derivadas. Crucialmente, el gradiente no se propaga hacia atrás a través de la distribución $q$ en sí, evitando las complejidades de diferenciar a través de distribuciones.

5. Espacio de Características y Guía

• Espacio de Características: Para manejar datos de alta dimensión, la pérdida de deriva se computa en un espacio de características utilizando un codificador auto-supervisado pre-entrenado (por ejemplo, un codificador tipo MAE de ResNet). Esto asegura que las muestras semánticamente similares permanezcan cerca, proporcionando señales de gradiente más ricas.
• Guía Libre de Clasificador (CFG): El método soporta naturalmente CFG mezclando muestras de datos no condicionadas en la distribución de muestras negativas durante el entrenamiento, lo que permite una generación controlable en la inferencia sin necesidad de reentrenamiento.

Contribuciones Clave

1. Nuevo Paradigma: Introduce los "Modelos de Deriva", un marco donde la distribución de pushforward evoluciona durante el tiempo de entrenamiento, eliminando la necesidad de inferencia iterativa.
2. Campo de Deriva: Propone un mecanismo de pérdida novedoso basado en un campo de deriva que logra el equilibrio cuando las distribuciones generada y de datos coinciden, gobernado por la atracción hacia los datos y la repulsión de la generación actual.
3. Generación de Un Solo Paso de Estado del Arte: Demuestra que una red de un solo paso puede lograr un rendimiento competitivo con, o superior a, los modelos de difusión y de flujo de múltiples pasos.
4. Robustez al Colapso de Modo: Experimentos con modelos de juguete sugieren que el método es robusto al colapso de modo; si la distribución generada colapsa en un modo, la atracción de otros modos en la distribución de datos continúa impulsando la deriva, empujando a la distribución a evolucionar.

Resultados Experimentales

Los autores evalúan los Modelos de Deriva en ImageNet 256×256 tanto en espacios latentes como de píxeles.

• Espacio Latente: Utilizando un tokenizador SD-VAE estándar, el modelo logra un FID de 1-NFE de 1.54 (con un modelo L/2 más grande). Esto supera a todos los métodos previos de un solo paso (por ejemplo, iMeanFlow, AdvFlow) y es competitivo con los modelos de difusión de múltiples pasos (por ejemplo, SiT-XL/2 con 250 pasos logra un FID de 2.06).
• Espacio de Píxeles: Sin un tokenizador latente, el modelo logra un FID de 1-NFE de 1.61, superando sustancialmente a los métodos previos de un solo paso en el espacio de píxeles (por ejemplo, StyleGAN-XL con 2.30 FID) y compitiendo con los modelos de múltiples pasos en el espacio de píxeles.
• Control Robótico: En el entorno de Política de Difusión (Diffusion Policy), el Modelo de Deriva de un solo paso iguala o supera el rendimiento de las políticas de difusión de múltiples pasos (100 NFE) en diversas tareas de manipulación robótica.

Significancia y Reivindicaciones

El artículo afirma abrir nuevas oportunidades para el modelado generativo de alta calidad y eficiente. Al reformular el entrenamiento iterativo de redes neuronales como un mecanismo de evolución de la distribución, el trabajo desafía la necesidad de formulaciones SDE/ODE para tareas generativas.

Los autores enfatizan que su método:

• Admite naturalmente la inferencia de un solo paso (1-NFE), reduciendo significativamente el costo computacional en el despliegue.
• No depende de la destilación de modelos de múltiples pasos pre-entrenados, entrenando desde cero.
• Ofrece una alternativa conceptualmente distinta a la difusión y al flujo de emparejamiento (flow matching), centrándose en la regla de actualización ($x_{i+1} = x_i + \Delta x_i$) durante el entrenamiento en lugar de la inferencia.

Los autores mantienen la modestia respecto a las garantías teóricas, señalando que, si bien $p=q \implies V=0$ está probado, el caso inverso ($V \to 0 \implies p \to q$) se cumple empíricamente y bajo supuestos de no degeneración específicos, pero no está garantizado para campos vectoriales arbitrarios. Consideran su trabajo como un punto de partida para explorar otras realizaciones de mecanismos de evolución de distribuciones.