Unrolled denoising networks provably learn optimal Bayesian inference
This addresses a foundational issue in Bayesian inference for inverse problems, providing theoretical guarantees for a widely used deep learning method, though it is incremental as it builds on existing unrolling and AMP frameworks.
The paper tackles the problem of whether algorithm unrolling can provably achieve optimal Bayesian inference when the prior is unknown, and proves that unrolled networks converge to optimal denoisers for compressed sensing with product priors, with numerical experiments showing advantages in low dimensions and non-Gaussian settings.
Much of Bayesian inference centers around the design of estimators for inverse problems which are optimal assuming the data comes from a known prior. But what do these optimality guarantees mean if the prior is unknown? In recent years, algorithm unrolling has emerged as deep learning's answer to this age-old question: design a neural network whose layers can in principle simulate iterations of inference algorithms and train on data generated by the unknown prior. Despite its empirical success, however, it has remained unclear whether this method can provably recover the performance of its optimal, prior-aware counterparts. In this work, we prove the first rigorous learning guarantees for neural networks based on unrolling approximate message passing (AMP). For compressed sensing, we prove that when trained on data drawn from a product prior, the layers of the network approximately converge to the same denoisers used in Bayes AMP. We also provide extensive numerical experiments for compressed sensing and rank-one matrix estimation demonstrating the advantages of our unrolled architecture - in addition to being able to obliviously adapt to general priors, it exhibits improvements over Bayes AMP in more general settings of low dimensions, non-Gaussian designs, and non-product priors.