Score Approximation for Diffusion Models on Arbitrary Low-Dimensional Structures

Planteamiento del Problema
Los fundamentos teóricos de los modelos de difusión basados en el score han dependido en gran medida de supuestos restrictivos con respecto a la distribución de datos subyacente, requiriendo específicamente densidades Lipschitz continuas o soportes de variedades suaves. Sin embargo, los datos perceptuales del mundo real (como las imágenes) a menudo exhiben singularidades, bordes afilados y grupos disjuntos donde las densidades son discontinuas o los gradientes se comportan mal. Los límites de complejidad existentes para aproximar la función de score $\nabla \log p_t(x)$ suelen sufrir una dependencia exponencial de la dimensión ambiental $n$, o bien no son aplicables a distribuciones que carecen de suavidad. Este vacío plantea una pregunta teórica fundamental: ¿Pueden los modelos de difusión tener éxito teóricamente en distribuciones de datos compactas arbitrarias sin asumir regularidad o suavidad?

Metodología
Los autores proponen un marco de aproximación de score universal que elimina todos los supuestos de continuidad y suavidad. El núcleo de su enfoque es una novedosa formulación de mezcla discreta y una estrategia intensiva en análisis centrada en los efectos de suavizado de los núcleos gaussianos.

1. Formulación de Mezcla Discreta: Los autores descomponen la distribución inicial $p_0$ en una mezcla finita de subdistribuciones $p_i$, cada una soportada en una pequeña bola $B(y_i, r)$. La verdadera función de score se aproxima reemplazando la integral sobre cada subdistribución por una evaluación de masa puntual en la media local $\tilde{y}_i$. Esto transforma el problema de la aproximación del score en la aproximación de una suma ponderada de núcleos gaussianos.
2. Análisis de Error de Divide y Vencerás: Para acotar el error de discretización, los autores particionan el espacio en grupos "cercanos" y "lejanos" relativos a un punto de consulta $x$.
   • Grupo Cercano: Para puntos cercanos al soporte, la expansión de Taylor del núcleo gaussiano alrededor de la media local muestra que el error de aproximación es pequeño.
   • Grupo Lejano: Para puntos alejados del soporte, el decaimiento del peso gaussiano es exponencial. Los autores demuestran que, incluso si el error de aproximación local es grande en estas regiones, la contribución al error total es insignificante debido al rápido decaimiento del núcleo.
3. Regularidad Universal: Un obstáculo teórico crítico es asegurar que los puntos "malos" (donde la aproximación podría fallar) tengan una medida insignificante. Los autores introducen una condición de regularidad $C_1, C_2, r$ y demuestran un teorema de "Regularidad Universal": para cualquier distribución de soporte compacto, el conjunto de puntos que no son regulares tiene una medida arbitrariamente pequeña. Esto asegura que la aproximación se cumpla casi en todas partes.
4. Implementación de Red Neuronal: Se demuestra que la fórmula de la mezcla discreta es implementable por una red ReLU. Los autores construyen una red que aproxima las operaciones necesarias (elevación al cuadrado, exponenciación, división y Softmax) con un error controlado, asegurando que la salida permanezca dentro del envolvente convexo del soporte.

Contribuciones Clave

• Teorema de Aproximación de Score Universal: El artículo establece el primer teorema que demuestra que la función de score para cualquier distribución con soporte compacto (caracterizada por una dimensión de Minkowski superior $d$) puede ser aproximada por una red ReLU. Esto elimina la necesidad de supuestos de suavidad Lipschitz, Hölder o de Sobolev.
• Límites de Complejidad: Los autores demuestran que la complejidad de la red (tamaño de parámetros) crece polinómicamente con la dimensión ambiental $n$ y exponencialmente solo con la dimensión intrínseca $d$. Específicamente, el tamaño escala aproximadamente como $O(n^3 d/4 (1/\epsilon)^{d/2})$. Esto rompe la "maldición de la dimensionalidad ambiental" que afecta a las teorías existentes.
• Condiciones de Regularidad Relajadas: Al introducir el concepto de regularidad universal, el trabajo demuestra que los conjuntos "malos" donde la aproximación falla son de medida insignificante para cualquier distribución compacta, incluyendo medidas de tipo fractal (por ejemplo, medidas de Cantor) que carecen de funciones de densidad de probabilidad.

Resultados
El resultado principal (Teorema 1) establece que para una distribución $p_0$ con soporte compacto de diámetro $R$ y dimensión de Minkowski superior $d$, la verdadera función de score $\nabla \log p_t(x)$ puede aproximarse mediante una formulación de predicción de datos $\nabla \log p_{t,appr}(x) = s_\theta(x, t) - x/t$ tal que el error cuadrático medio está acotado por $7\epsilon^2$ para todo $t \in [t_0, T]$.

• La red aproximadora $s_\theta$ tiene una profundidad de $O(\log \omega \log \log \omega)$ y un ancho de $O(C_K r^{-d} (n \log \omega + \log^2 \omega))$, donde $\omega$ depende polinómicamente de $n$ y $1/\epsilon$.
• El límite de error se mantiene globalmente, con el error de aproximación acotado por $R/t$ para puntos fuera de la distribución, lo cual es integrable sobre la pequeña región de probabilidad.

Significancia
El artículo afirma proporcionar una justificación teórica rigurosa para el éxito empírico de los modelos de difusión en estructuras de datos irregulares y no suaves. Al demostrar que es posible una estimación precisa del score para cualquier distribución compacta sin supuestos de suavidad, el trabajo cierra la brecha entre las garantías teóricas y las características de los datos del mundo real (como el recorte de píxeles y las singularidades). Los autores señalan que sus tácticas de divide y vencerás y sus resultados de regularidad pueden adaptarse para acotar derivadas de orden superior (como la Hessiana), lo que podría ayudar al trabajo teórico futuro sobre modelos de difusión que actualmente dependen de supuestos de acotación para estos términos. El trabajo se posiciona como un paso fundacional, expandiendo el alcance de la teoría de aproximación de score desde variedades suaves hacia la clase general de distribuciones de soporte compacto.