Solving Bayesian inverse problems with diffusion priors and off-policy RL

2025-03-12

AI Summary

This work addresses challenging linear and nonlinear inverse problems in vision and scientific imaging through a novel Bayesian inference framework based on conditional diffusion posterior modeling. To overcome inherent bias in existing training-free diffusion posterior methods—particularly during latent-space inference—we introduce, for the first time, the Relative Trajectory Balancing (RTB) objective from off-policy reinforcement learning into diffusion model training. This enables efficient adaptation from a pre-trained unconditional prior to a task-driven conditional posterior without fine-tuning or additional supervision. Guided by RTB-based off-policy retrospective exploration, our method significantly improves posterior fidelity and calibration. Extensive experiments demonstrate consistent superiority over state-of-the-art training-free approaches across diverse inverse problems—including deblurring, super-resolution, and tomographic reconstruction—establishing a new paradigm for Bayesian inverse problem solving.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Relational Probabilistic Models
• Machine Learning: Calibration & Uncertainty Quantification

Application Category

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

Abstract

This paper presents a practical application of Relative Trajectory Balance (RTB), a recently introduced off-policy reinforcement learning (RL) objective that can asymptotically solve Bayesian inverse problems optimally. We extend the original work by using RTB to train conditional diffusion model posteriors from pretrained unconditional priors for challenging linear and non-linear inverse problems in vision, and science. We use the objective alongside techniques such as off-policy backtracking exploration to improve training. Importantly, our results show that existing training-free diffusion posterior methods struggle to perform effective posterior inference in latent space due to inherent biases.

Problem

• Solving Bayesian inverse problems using diffusion priors and off-policy RL.
• Training conditional diffusion models for challenging inverse problems in vision and science.
• Addressing biases in existing diffusion posterior methods for latent space inference.

Innovation

• Uses Relative Trajectory Balance for RL
• Trains conditional diffusion model posteriors
• Employs off-policy backtracking exploration