RBF-Solver: A Multistep Sampler for Diffusion Probabilistic Models via Radial Basis Functions

Resumen Técnico: RBF-Solver

1. El Problema

Los Modelos Probabilísticos de Difusión (DPMs) han demostrado un rendimiento excepcional en tareas generativas (imagen, video, audio), pero sufren de un alto costo computacional debido a la necesidad de realizar muchos pasos de denoising iterativos para generar muestras de alta calidad.

• Limitaciones de los métodos actuales: Los muestreadores multietapa (multistep) basados en polinomios (como DPM-Solver++ o UniPC) aceleran la inferencia utilizando aproximaciones de Taylor o interpolación de Lagrange. Sin embargo, estos métodos siguen un esquema predefinido y rígido.
• Inestabilidad de alto orden: A medida que se aumenta el orden de precisión (orden 4 o superior) para reducir aún más los pasos de evaluación de la función (NFE), los métodos basados en polinomios tienden a volverse numéricamente inestables, sufriendo del fenómeno de Runge y degradando la calidad de la imagen.
• Falta de flexibilidad: Las aproximaciones polinómicas no tienen parámetros aprendibles que permitan adaptar la trayectoria de muestreo a las características específicas del modelo o del dato.

2. Metodología: RBF-Solver

Los autores proponen RBF-Solver, un muestreador multietapa que utiliza interpolación con Funciones de Base Radial (RBF) Gaussianas en lugar de polinomios algebraicos.

• Interpolación con RBFs: En lugar de aproximar la función de predicción del modelo $\hat{x}_\theta$ con un polinomio, RBF-Solver la interpola utilizando una combinación de funciones Gaussianas:
  $$R(\lambda) = \sum w_{j} \phi_{j}(\lambda) + w_{const}$$
  Donde $\phi_j$ son funciones Gaussianas con un parámetro de forma ($\gamma$) aprendible. Este parámetro controla la "anchura" de la base gaussiana.
• Adaptabilidad: La clave del método es que los parámetros de forma $\gamma$ se optimizan para seguir trayectorias de muestreo óptimas. Esto permite que el interpolante se adapte a las características locales de los datos, evitando la oscilación típica de los polinomios de alto orden.
• Convergencia Teórica:
  • Orden 1: Se reduce al método de Euler (equivalente a DDIM).
  • Orden $\ge$ 2: A medida que el parámetro de forma $\gamma \to \infty$, las RBFs se vuelven planas y la interpolación converge a la interpolación de Lagrange, haciendo que RBF-Solver sea equivalente al Método de Adams. Esto garantiza compatibilidad con los muestreadores existentes.
  • Orden $\to$ 0: Cuando $\gamma \to 0$, converge a un muestreador con coeficientes iguales.
• Optimización de Parámetros: Se introduce un algoritmo (Algoritmo 2) para optimizar los parámetros de forma $\gamma$ en tiempo de ejecución (o pre-calculado) minimizando el error cuadrático medio entre la trayectoria predicha y una trayectoria objetivo de alta calidad (generada con un muestreador de referencia con muchos NFE).
• Estabilidad: Gracias a la propiedad de localización de las RBFs Gaussianas, el método mantiene alta fidelidad incluso en órdenes superiores (4, 5, 6), donde los métodos polinómicos fallan.

3. Contribuciones Clave

1. Nueva Formulación de Muestreo: Introducción de la interpolación RBF Gaussiana como mecanismo central para resolver las EDOs de difusión, superando las limitaciones de rigidez de los métodos polinómicos.
2. Garantías Teóricas: Demostración de que el método satisface la condición de suma de coeficientes (consistencia), converge al método de Adams cuando $\gamma \to \infty$, y ofrece garantías de precisión de orden $p$.
3. Optimización de Trayectoria: Desarrollo de un esquema para aprender los parámetros de forma que permiten al muestreador adaptarse dinámicamente a la complejidad de la trayectoria de difusión, mejorando la precisión en pasos grandes.
4. Estabilidad en Alto Orden: Resolución del problema de inestabilidad numérica en órdenes superiores, permitiendo muestreadores de orden 4+ que mantienen la calidad de la imagen.

4. Resultados Experimentales

El método se evaluó en tareas de generación incondicional y condicional frente a muestreadores de última generación (DPM-Solver++, UniPC, DPM-Solver-v3, DC-Solver).

• Generación Incondicional (CIFAR-10, ImageNet 64x64):
  • En el régimen de NFE alto ($\ge$ 15), RBF-Solver supera consistentemente a los métodos basados en polinomios.
  • En CIFAR-10 (Score-SDE), logra un FID de 2.87 con 15 NFE y 2.48 con 40 NFE, superando a UniPC y DPM-Solver++.
  • Muestra una mejora notable en la estabilidad: mientras UniPC degrada su rendimiento drásticamente en órdenes $\ge$ 5, RBF-Solver mantiene o mejora su FID en órdenes 4, 5 y 6.
• Generación Condicional (ImageNet 128x128, 256x256, Stable Diffusion):
  • En el régimen de NFE bajo (5-10), especialmente con escalas de guía altas (6.0 y 8.0), RBF-Solver logra reducciones significativas en el FID (hasta un 33.73% de mejora relativa) comparado con los baselines.
  • En Stable Diffusion v1.4, supera a los competidores en similitud coseno (CLIP) y error cuadrático medio (RMSE) en escalas de guía altas.
• Eficiencia: El proceso de optimización de parámetros es ligero, requiriendo solo 128 pares objetivo y completándose en tiempos que van de 1 minuto a 1 hora en una GPU, sin necesidad de reentrenar la red neuronal.

5. Significado e Impacto

• Superación de Barreras de Precisión: RBF-Solver demuestra que es posible utilizar muestreadores de muy alto orden sin sacrificar la estabilidad numérica, algo que los métodos polinómicos no logran.
• Flexibilidad vs. Rigidez: Introduce un paradigma donde los muestreadores no son fijos, sino que pueden "aprender" parámetros de interpolación para adaptarse a la dinámica específica del modelo de difusión, acercándose a una solución óptima de la EDO.
• Aceleración Práctica: Permite generar imágenes de alta calidad con menos pasos de inferencia (NFE), reduciendo significativamente el costo computacional y el tiempo de generación, lo cual es crucial para la viabilidad comercial de los modelos de difusión en aplicaciones en tiempo real.
• Generalización: Al ser un método "sin entrenamiento" (training-free) una vez optimizados los parámetros, puede aplicarse a cualquier modelo de difusión pre-entrenado sin modificar sus pesos.

En conclusión, RBF-Solver representa un avance significativo en la teoría de muestreo de difusión, reemplazando la aproximación polinómica tradicional con una basada en RBFs que ofrece mayor flexibilidad, estabilidad en alto orden y rendimiento superior en escenarios de inferencia rápida.