Bayesian Tensor Decomposition with Diffusion Model Prior

2026-06-02

AI Summary

This work addresses the significant performance degradation of low-rank tensor decomposition under high missing rates or strong noise, which stems from the inability of conventional priors to capture the complex statistical structures of real-world data. To overcome this limitation, we propose DiffBCP, a novel framework that, for the first time, incorporates a pre-trained diffusion model as an implicit data prior within Bayesian CP decomposition. DiffBCP automatically infers the tensor rank through an accumulated shrinkage process and employs a split Gibbs sampler to jointly optimize the likelihood, low-rank constraints, and diffusion prior. Furthermore, a noise-adaptive coupling scheduling mechanism is introduced to enable low-rank-guided denoising. Experiments demonstrate that DiffBCP consistently outperforms existing Bayesian, nonlinear, and plug-and-play tensor completion methods on image inpainting and denoising tasks, including high-resolution out-of-distribution data.

Abstract

Low-rank tensor decomposition (TD) is usually effective on clean, fully observed data, but it often degrades under severe missingness or noise. Low-rankness is itself a useful but limited structural prior, and additional handcrafted priors (e.g., sparsity or smoothness) still fall short of capturing the rich statistics of real-world data. To compensate for this weak inductive bias under heavy corruption, one would like to inject a learned, data-driven prior; however, the state-of-the-art diffusion models are not readily compatible with current TD and tractable posterior inference. To address these challenges, we introduce DiffBCP, a hybrid-prior Bayesian CP decomposition framework that couples a cumulative shrinkage process prior over the CP factors for automatic rank selection with an off-the-shelf pre-trained diffusion model as an implicit data prior on the reconstructed tensor. To make posterior inference tractable despite the coupling among the likelihood, low-rank constraint, and diffusion prior, we develop a split Gibbs sampler: CP factors admit conjugate updates, while the diffusion block is sampled via low-rank-guided denoising. A noise-adaptive coupling schedule further reduces sensitivity to hand-tuned annealing. Experiments on image inpainting and denoising, including high-resolution out-of-distribution images, show consistent gains over Bayesian, nonlinear, and plug-and-play TD baselines.

Problem

• tensor decomposition
• missing data
• noise robustness
• Bayesian inference
• data-driven prior

Innovation

• Bayesian tensor decomposition
• diffusion model prior
• CP decomposition
• split Gibbs sampler
• low-rank-guided denoising