Technical Summary: Bayesian Tensor Decomposition with Diffusion Model Prior

Problem Statement

Low-rank tensor decomposition (TD) is a standard tool for multi-dimensional data analysis, effective when data is clean and fully observed. However, performance degrades significantly under severe missingness or noise. While the low-rank assumption serves as a structural prior, it is often insufficient for real-world data recovery. Previous attempts to augment TD with handcrafted priors (e.g., sparsity, smoothness) fail to capture the rich statistics of complex data distributions. Conversely, while state-of-the-art diffusion models offer powerful, learned data priors, they are not readily compatible with current TD frameworks or tractable posterior inference, particularly when coupled with low-rank constraints and automatic rank selection.

Methodology: DiffBCP

The authors propose DiffBCP, a hybrid-prior Bayesian CP decomposition framework designed to integrate a learned diffusion prior with structural low-rank constraints.

1. Probabilistic Formulation

The model defines a joint probability distribution over observed data $Y$, the reconstructed tensor $X$, and latent CP factors ($\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)$$

Key components include:

• Likelihood: Assumes observed entries follow a Normal distribution with precision $\tau$.
• Structural Constraint: $p(X | A^{(1:N)}, \lambda)$ enforces the CP decomposition structure ($X = \text{CP}(\lambda, A^{(1:N)})$) via a Dirac delta function.
• Latent Factor Priors:
  • CUSP Prior: The CP weights $\lambda$ follow a Cumulative Shrinkage Process (CUSP) prior. This hierarchical prior enables automatic rank selection by shrinking redundant components toward a small constant $\theta_\infty$ as the rank increases.
  • Gaussian Priors: Factor matrices $A^{(n)}$ follow standard Gaussian priors to regularize norms.
• Diffusion Prior: The reconstructed tensor $X$ is assigned a powerful, pre-trained diffusion model prior $p(X)$. This acts as an implicit, data-driven regularizer capable of capturing complex textures and structures that low-rank assumptions alone cannot.

2. Inference: Split Gibbs Sampler (SGS)

Direct posterior inference is intractable due to the coupling between the likelihood, the low-rank constraint, and the implicit diffusion prior. The authors employ a Split Gibbs Sampler with an auxiliary variable $Z$ to decouple these terms:

• Augmentation: An auxiliary variable $Z$ is introduced with a coupling term $\phi(Z, X; \rho) = \frac{1}{2\rho^2}\|Z - X\|_F^2$.
• Latent Variable Updates: Given $Z$, the CP factors ($\lambda, A^{(n)}$) and hyperparameters ($\tau, \theta, \nu$) admit closed-form conjugate updates (Gaussian, Gamma, Beta, etc.).
• Diffusion Block Update: Given $X$, the auxiliary variable $Z$ is sampled from a distribution defined by the diffusion prior and the coupling term. This is solved as a denoising problem using the pre-trained score network $s_\psi$ via the Stochastic Differential Equation (SDE) reverse process, guided by the low-rank estimate $X$.
• Noise-Adaptive Coupling: A schedule is introduced where the coupling parameter $\rho$ is dynamically adjusted based on the inferred noise precision $\tau$ (specifically $\tau\rho^2 = c$). This reduces sensitivity to manual annealing schedules.

3. Rank Adaptation

The framework includes a bidirectional rank adaptation scheme. The CUSP prior shrinks inactive components, while a separate mechanism adds new components if the current rank is insufficient. This ensures robustness against initial rank misspecification.

Key Contributions

1. Hybrid-Prior Framework: The first Bayesian tensor decomposition model to integrate an off-the-shelf pre-trained diffusion model as a data prior on the reconstructed tensor, combined with a structural low-rank CP prior.
2. Tractable Inference: Development of a Split Gibbs Sampler that enables posterior inference despite the complex coupling between likelihood, low-rank constraints, and diffusion priors. This includes closed-form updates for latent factors and low-rank-guided denoising for the diffusion block.
3. Automatic Rank and Noise Inference: The use of the CUSP prior and fully Bayesian noise estimation eliminates the need for manual tuning of rank and noise levels.
4. Robustness: The method demonstrates robustness to initial rank misspecification and varying noise levels through its adaptive coupling schedule.

Experimental Results

The method was evaluated on image inpainting and denoising tasks using FFHQ and ImageNet datasets, as well as high-resolution out-of-distribution (OOD) images.

• Performance: DiffBCP consistently outperformed baselines, including traditional Bayesian TD (BCP, BTR), deep learning-based TD (HLRTF, DeepTensor), and plug-and-play methods (GLON, PnP-DM).
  • On FFHQ with 90% missing data (Uniform mask), DiffBCP achieved 28.28 dB PSNR, improving over the best baseline (DeepTensor) by 2.17 dB.
  • On irregular masks, gains were even more pronounced (+2.33 dB over DeepTensor).
• OOD and High-Resolution: The method successfully reconstructed 2048×2048 images (e.g., Marseille, Tokyo) using a diffusion model trained on 256×256 images, maintaining coherence and detail where other methods (like PuTT) introduced artifacts or failed.
• Ablation Studies:
  • Removing the diffusion prior significantly degraded performance on structured masks (Stripe, Irregular), confirming the diffusion prior's role in recovering spatial continuity.
  • The CUSP prior proved critical for recovery when the initial rank was severely under-specified (e.g., starting at rank 10 vs. 200).
• Comparison with Diffusion Samplers: Compared to PnP-DM, DiffBCP showed superior stability regarding noise level scheduling and better performance on high-resolution images, though PnP-DM could achieve competitive results on uniform masks if its noise parameters were manually tuned to optimal values.

Significance and Claims

The paper claims that DiffBCP successfully bridges the gap between the structural inductive bias of low-rank tensor decomposition and the expressive power of modern generative diffusion models. By formulating this integration within a fully probabilistic framework, the authors achieve:

• State-of-the-art recovery for noisy and incomplete tensor data.
• Automatic adaptation to unknown ranks and noise levels, removing the need for extensive hyperparameter tuning.
• Robustness to distribution shifts and high-resolution inputs, leveraging the low-rank structure as a stabilizing inductive bias when the diffusion prior faces OOD data.

The authors note that the performance relies on the underlying signal admitting a meaningful low-rank structure; if this assumption is violated, the structural contribution diminishes, though the diffusion prior remains active. Future work is suggested for extending the framework to other inverse problems (e.g., compressed sensing) and other tensor decompositions (Tensor Train, Tensor Ring).