Resumen Técnico: Descomposición de Tensores Bayesiana con Prior de Modelo de Difusión

Declaración del Problema

La descomposición de tensores (TD) de bajo rango es una herramienta estándar para el análisis de datos multidimensionales, efectiva cuando los datos están limpios y completamente observados. Sin embargo, su rendimiento se degrada significativamente bajo una ausencia de datos severa o ruido. Aunque la suposición de bajo rango sirve como un prior estructural, a menudo es insuficiente para la recuperación de datos del mundo real. Los intentos previos de aumentar la TD con priors diseñados manualmente (por ejemplo, dispersión o suavidad) no logran capturar las ricas estadísticas de las distribuciones de datos complejos. Por el contrario, aunque los modelos de difusión de vanguardia ofrecen priors de datos aprendidos muy potentes, no son fácilmente compatibles con los marcos actuales de TD ni tratables para la inferencia posterior, particularmente cuando se acoplan con restricciones de bajo rango y selección automática de rango.

Metodología: DiffBCP

Los autores proponen DiffBCP, un marco de descomposición CP bayesiana con prior híbrido diseñado para integrar un prior de difusión aprendido con restricciones estructurales de bajo rango.

1. Formulación Probabilística

El modelo define una distribución de probabilidad conjunta sobre los datos observados $Y$, el tensor reconstruido $X$, y los factores CP latentes ($\lambda, A^{(1:N)}$):
$$p(Y, X, \lambda, A^{(1:N)}, \tau) \propto p(Y | X, \tau) \cdot p(X | A^{(1:N)}, \lambda) \cdot p(A^{(1:N)}) \cdot p(\lambda) \cdot p(X) \cdot p(\tau)$$

Los componentes clave incluyen:

• Verosimilitud (Likelihood): Asume que las entradas observadas siguen una distribución Normal con precisión $\tau$.
• Restricción Estructural: $p(X | A^{(1:N)}, \lambda)$ impone la estructura de la descomposición CP ($X = \text{CP}(\lambda, A^{(1:N)})$) mediante una función delta de Dirac.
• Priors de Factores Latentes:
  • Prior CUSP: Los pesos de CP $\lambda$ siguen un prior de Proceso de Contracción Acumulativa (CUSP). Este prior jerárquico permite la selección automática de rango al contraer los componentes redundantes hacia una pequeña constante $\theta_\infty$ a medida que el rango aumenta.
  • Priors Gaussianos: Las matrices de factores $A^{(n)}$ siguen priors gaussianos estándar para regularizar las normas.
• Prior de Difusión: Al tensor reconstruido $X$ se le asigna un potente prior de modelo de difusión pre-entrenado $p(X)$. Esto actúa como un regularizador implícito y basado en datos capaz de capturar texturas y estructuras complejas que las suposiciones de bajo rango por sí solas no pueden capturar.

2. Inferencia: Muestreador Gibbs Dividido (SGS)

La inferencia posterior directa es intratable debido al acoplamiento entre la verosimilitud, la restricción de bajo rango y el prior de difusión implícito. Los autores emplean un Muestreador Gibbs Dividido con una variable auxiliar $Z$ para desacoplar estos términos:

• Aumentación: Se introduce una variable auxiliar $Z$ con un término de acoplamiento $\phi(Z, X; \rho) = \frac{1}{2\rho^2}\|Z - X\|_F^2$.
• Actualizaciones de Variables Latentes: Dado $Z$, los factores CP ($\lambda, A^{(n)}$) y los hiperparámetros ($\tau, \theta, \nu$) admiten actualizaciones conjugadas en forma cerrada (Gaussiana, Gamma, Beta, etc.).
• Actualización del Bloque de Difusión: Dado $X$, la variable auxiliar $Z$ se muestrea de una distribución definida por el prior de difusión y el término de acoplamiento. Esto se resuelve como un problema de eliminación de ruido utilizando la red de puntuación (score network) pre-entrenada $s_\psi$ a través del proceso de SDE inversa, guiado por la estimación de bajo rango $X$.
• Acoplamiento Adaptativo al Ruido: Se introduce un esquema donde el parámetro de acoplamiento $\rho$ se ajusta dinámicamente basado en la precisión de ruido inferida $\tau$ (específicamente $\tau\rho^2 = c$). Esto reduce la sensibilidad a los esquemas de templado manual.

3. Adaptación de Rango

El marco incluye un esquema de adaptación de rango bidireccional. El prior CUSP contrae los componentes inactivos, mientras que un mecanismo separado añade nuevos componentes si el rango actual es insuficiente. Esto asegura la robustez contra la errónea especificación inicial del rango.

Contribuciones Clave

1. Marco de Prior Híbrido: El primer modelo de descomposición de tensores bayesiano que integra un modelo de difusión pre-entrenado listo para usar como un prior de datos sobre el tensor reconstruido, combinado con un prior de CP de bajo rango estructural.
2. Inferencia Tratable: Desarrollo de un Muestreador Gibbs Dividido que permite la inferencia posterior a pesar del complejo acoplamiento entre la verosimilitud, las restricciones de bajo rango y los priors de difusión. Esto incluye actualizaciones en forma cerrada para los factores latentes y eliminación de ruido guiada por bajo rango para el bloque de difusión.
3. Inferencia Automática de Rango y Ruido: El uso del prior CUSP y la estimación bayesiana completa del ruido elimina la necesidad de ajustar manualmente el rango y los niveles de ruido.
4. Robustez: El método demuestra robustez contra la errónea especificación inicial del rango y niveles variables de ruido a través de su esquema de acoplamiento adaptativo.

Resultados Experimentales

El método fue evaluado en tareas de inpainting de imágenes y eliminación de ruido utilizando los conjuntos de datos FFHQ e ImageNet, así como imágenes de alta resolución fuera de la distribución (OOD).

• Rendimiento: DiffBCP superó consistentemente a los baselines, incluyendo la TD bayesiana tradicional (BCP, BTR), la TD basada en aprendizaje profundo (HLRTF, DeepTensor) y los métodos plug-and-play (GLON, PnP-DM).
  • En FFHQ con un 90% de datos faltantes (máscara uniforme), DiffBCP alcanzó 28.28 dB PSNR, mejorando sobre el mejor baseline (DeepTensor) en 2.17 dB.
  • En máscaras irregulares, las ganancias fueron aún más pronunciadas (+2.33 dB sobre DeepTensor).
• OOD y Alta Resolución: El método reconstruyó con éxito imágenes de 2048×2048 (por ejemplo, Marseille, Tokyo) utilizando un modelo de difusión entrenado en imágenes de 256×256, manteniendo la coherencia y el detalle donde otros métodos (como PuTT) introdujeron artefactos o fallaron.
• Estudios de Ablación:
  • Eliminar el prior de difusión degradó significativamente el rendimiento en máscaras estructuradas (Stripe, Irregular), confirmando el papel del prior de difusión en la recuperación de la continuidad espacial.
  • El prior CUSP resultó crítico para la recuperación cuando el rango inicial estaba severamente subespecificado (por ejemplo, comenzando en rango 10 vs. 200).
• Comparación con Muestreadores de Difusión: Comparado con PnP-DM, DiffBCP mostró una mayor estabilidad respecto al esquema de nivel de ruido y un mejor rendimiento en imágenes de alta resolución, aunque PnP-DM pudo lograr resultados competitivos en máscaras uniformes si sus parámetros de ruido se ajustaban manualmente a valores óptimos.

Significado y Reivindicaciones

El artículo afirma que DiffBCP logra cerrar la brecha entre el sesgo inductivo estructural de la descomposición de tensores de bajo rango y el poder expresivo de los modelos generativos de difusión modernos. Al formular esta integración dentro de un marco totalmente probabilístico, los autores logran:

• Recuperación de vanguardia (state-of-the-art) para datos tensoriales ruidosos e incompletos.
• Adaptación automática a rangos y niveles de ruido desconocidos, eliminando la necesidad de un ajuste extensivo de hiperparámetros.
• Robustez ante cambios de distribución y entradas de alta resolución, aprovechando la estructura de bajo rango como un sesgo inductivo estabilizador cuando el prior de difusión enfrenta datos OOD.

Los autores señalan que el rendimiento depende de que la señal subyacente admita una estructura de bajo rango significativa; si esta suposición se viola, la contribución estructural disminuye, aunque el prior de difusión permanece activo. Se sugiere como trabajo futuro la extensión de este marco a otros problemas inversos (por ejemplo, compressed sensing) y otras descomposiciones de tensores (Tensor Train, Tensor Ring).