Unrolled denoising networks provably learn optimal Bayesian inference

📄 arXiv: 2409.12947v1 📥 PDF

作者: Aayush Karan, Kulin Shah, Sitan Chen, Yonina C. Eldar

分类: cs.LG, cs.DS, stat.ML

发布日期: 2024-09-19

备注: 32 pages

💡 一句话要点

提出基于展开降噪网络的贝叶斯推断学习框架，可证明地学习最优贝叶斯推断。

🎯 匹配领域: 支柱八：物理动画 (Physics-based Animation)

关键词: 算法展开 贝叶斯推断 近似消息传递 深度学习 压缩感知 逆问题 神经网络

📋 核心要点

  1. 现有贝叶斯推断方法依赖于已知先验，但在先验未知的情况下，其最优性保证失效，面临挑战。
  2. 论文提出基于展开近似消息传递（AMP）的神经网络，通过数据驱动的方式学习最优贝叶斯推断。
  3. 实验结果表明，该方法在压缩感知和秩一矩阵估计等任务上，优于传统的贝叶斯AMP方法。

📝 摘要（中文）

贝叶斯推断的核心在于为逆问题设计估计器，这些估计器在数据来自已知先验的假设下是最优的。但如果先验未知，这些最优性保证意味着什么？近年来，算法展开已成为深度学习解决这个古老问题的方法：设计一个神经网络，其层原则上可以模拟推理算法的迭代，并在由未知先验生成的数据上进行训练。尽管它在经验上取得了成功，但这种方法是否能可靠地恢复其最优的、先验已知的对应物的性能，仍然不清楚。在这项工作中，我们证明了基于展开近似消息传递（AMP）的神经网络的第一个严格的学习保证。对于压缩感知，我们证明了当在从乘积先验中抽取的数据上训练时，网络的层近似收敛到贝叶斯AMP中使用的相同降噪器。我们还为压缩感知和秩一矩阵估计提供了广泛的数值实验，证明了我们的展开架构的优势——除了能够盲目地适应一般先验之外，它还在低维度、非高斯设计和非乘积先验的更一般设置中表现出优于贝叶斯AMP的改进。

🔬 方法详解

问题定义：论文旨在解决在先验信息未知的情况下，如何进行最优贝叶斯推断的问题。传统贝叶斯方法依赖于精确的先验模型，当先验未知或模型不准确时，性能会显著下降。现有算法展开方法虽然在实践中表现良好，但缺乏理论保证，无法证明其能够恢复最优性能。

核心思路：论文的核心思路是利用算法展开技术，将近似消息传递（AMP）算法展开为一个深度神经网络。通过在数据上训练该网络，使其能够学习到适应未知先验的最优降噪器。这种方法结合了深度学习的自适应性和贝叶斯推断的理论基础。

技术框架：整体框架包括以下几个主要步骤：1) 选择AMP算法作为基础。2) 将AMP算法的迭代过程展开为神经网络的层。3) 使用数据驱动的方式训练网络，学习每一层的降噪器参数。4) 在测试阶段，使用训练好的网络进行贝叶斯推断。

关键创新：论文最重要的创新在于提供了基于展开AMP网络的学习保证。证明了在特定条件下（如乘积先验），网络的层可以收敛到最优贝叶斯AMP中使用的降噪器。这是第一个为展开网络提供严格学习保证的工作。

关键设计：关键设计包括：1) 选择合适的AMP算法变体。2) 设计合适的网络结构，使其能够模拟AMP算法的迭代过程。3) 选择合适的损失函数，例如均方误差，以驱动网络学习最优降噪器。4) 针对不同的应用场景，调整网络参数和训练策略。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文通过数值实验验证了所提出方法的有效性。在压缩感知和秩一矩阵估计任务上，与传统的贝叶斯AMP方法相比，该方法在低维度、非高斯设计和非乘积先验等更一般的情况下表现出显著的性能提升。实验结果表明，该方法能够有效地学习未知先验，并实现接近最优的贝叶斯推断。

🎯 应用场景

该研究成果可应用于各种逆问题，例如图像重建、信号恢复、压缩感知等。在医学成像、遥感、无线通信等领域具有广泛的应用前景。该方法能够自适应地学习未知先验，提高逆问题求解的鲁棒性和准确性，具有重要的实际价值和潜在的未来影响。

📄 摘要（原文）

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.