Bayesian Tensor Decomposition with Diffusion Model Prior
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Image Reconstruction | ImageNet | PSNR28.95 | 96 | |
| Image Reconstruction | FFHQ | PSNR32.13 | 28 | |
| Inpainting | FFHQ Stripe mask | PSNR27.91 | 5 | |
| Inpainting | FFHQ Uniform 0.7 mask | PSNR32.13 | 5 | |
| Inpainting | FFHQ Uniform 0.9 mask | PSNR28.28 | 5 | |
| Inpainting | FFHQ Irregular mask | PSNR30.34 | 5 | |
| High-resolution Image Recovery | Marseille Uniform 0.9 | PSNR20.15 | 3 | |
| High-resolution Image Recovery | Marseille Uniform 0.95 | PSNR18.01 | 3 | |
| High-resolution Image Recovery | Marseille Irregular | PSNR20.94 | 3 | |
| High-resolution Image Recovery | Tokyo Uniform 0.9 | PSNR20.66 | 3 |