End-to-End Trained Energy-Based Model for Image Recovery

Abstract

This paper presents a novel end-to-end (E2E) empirical Bayes framework for learning the posterior distribution in linear inverse problems. The proposed framework models the log-posterior as the sum of a log-likelihood derived from the known forward model and a neural network-parameterized log-prior learned from training data. The approach enables both (a) computation of a point estimate i.e., stationary point of the learned log-posterior using a majorization minimization algorithm and (b) posterior sampling for computing the minimum mean square error (MMSE) estimate. Provided that the required surrogate function conditions are satisfied, the point estimation algorithm converges to a stationary point of the learned negative log-posterior without imposing contraction constraints. Unlike diffusion models that pre-learn the entire prior, learning the posterior directly leads to about 150× reduced training data requirements and 2× fewer parameters. Across three MRI acquisition settings, the framework improves point-estimation PSNR over PnP-ISTA by 3.07–4.37 dB while remaining competitive with existing E2E methods. Across two posterior sampling settings, it improves PSNR over DPS by 0.11–2.22 dB and over DAPS by 1.17–4.16 dB, while providing inference speedups of 2× and 5×, respectively. Furthermore, by avoiding algorithm unrolling, the proposed framework also reduces memory requirements by approximately 4× and 16× relative to E2E-MoL and E2E-EBM methods, respectively.