Unrolled denoising networks provably learn optimal Bayesian inference

TL;DR

This paper proves that unrolled denoising neural networks trained on data can asymptotically learn optimal Bayesian denoisers, demonstrating theoretical guarantees and empirical advantages over traditional methods in inverse problems.

Contribution

It provides the first rigorous theoretical guarantees for unrolled AMP-based neural networks and shows they can approximate optimal Bayesian inference under unknown priors.

Findings

Networks converge to Bayes AMP denoisers when trained on product priors.

Unrolled networks outperform Bayes AMP in low-dimensional, non-Gaussian, and non-product prior settings.

Empirical results confirm the ability of unrolled networks to adapt to various priors and improve inference quality.

Abstract

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.