Deep End-to-End Posterior ENergy (DEEPEN) for image recovery

2025-03-21

AI Summary

Existing image reconstruction methods face a fundamental trade-off: end-to-end and plug-and-play approaches only approximate maximum a posteriori (MAP) estimation and cannot sample from the full posterior distribution; while diffusion models enable posterior sampling, they are incompatible with efficient end-to-end training. Method: We propose the first end-to-end trainable energy-based model (EBM) that directly learns the posterior energy function—without algorithmic unrolling or restrictive regularization constraints. Our approach jointly models data consistency and the negative log-prior, optimizing the energy function via maximum likelihood. Contribution/Results: The framework unifies exact MAP inference and efficient posterior sampling. It achieves state-of-the-art MAP performance, surpassing leading end-to-end and plug-and-play methods; enables orders-of-magnitude faster sampling than diffusion models; and exhibits superior robustness to variations in imaging parameters. To our knowledge, this is the first method to simultaneously deliver precise Bayesian inference and high-fidelity generative capability within a single, end-to-end trainable architecture.

Technology Category

• Machine Learning: Bayesian Learning
• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Current end-to-end (E2E) and plug-and-play (PnP) image reconstruction algorithms approximate the maximum a posteriori (MAP) estimate but cannot offer sampling from the posterior distribution, like diffusion models. By contrast, it is challenging for diffusion models to be trained in an E2E fashion. This paper introduces a Deep End-to-End Posterior ENergy (DEEPEN) framework, which enables MAP estimation as well as sampling. We learn the parameters of the posterior, which is the sum of the data consistency error and the negative log-prior distribution, using maximum likelihood optimization in an E2E fashion. The proposed approach does not require algorithm unrolling, and hence has a smaller computational and memory footprint than current E2E methods, while it does not require contraction constraints typically needed by current PnP methods. Our results demonstrate that DEEPEN offers improved performance than current E2E and PnP models in the MAP setting, while it also offers faster sampling compared to diffusion models. In addition, the learned energy-based model is observed to be more robust to changes in image acquisition settings.

Problem

• Enables MAP estimation and posterior sampling for image recovery
• Reduces computational and memory costs compared to E2E methods
• Improves robustness to changes in image acquisition settings

Innovation

• E2E posterior energy learning for image recovery
• Combines MAP estimation with posterior sampling
• Reduced computational footprint without unrolling