Bayesian Tensor Decomposition with Diffusion Model Prior

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.

Reconstructing a complete signal from incomplete or noisy measurements — restoring a damaged photo, completing a partial medical scan — is a central task across many fields. Classical methods impose a compact mathematical structure on the answer, which is often too restrictive; modern AI methods consult a generative model trained on millions of examples, but can hallucinate plausible-but-wrong details and require careful tuning. We propose DiffBCP, a probabilistic method that combines both kinds of knowledge. It alternates between fitting a compact structure and consulting a pretrained AI model to keep the result realistic, and automatically figures out how complex the structure needs to be. DiffBCP gives the largest advantage over either approach when the underlying signal admits such compact structure, which is typical of natural images and many scientific recordings. The gains are most pronounced in the hardest cases: extreme missingness, structured corruption, and very high-resolution images. It also reports how confident it is in each region of the result, valuable for high-stakes domains such as medical imaging.