Conditional Injective Flows for Bayesian Imaging

TL;DR

This paper introduces C-Trumpets, a novel class of conditional injective flows that efficiently model the posterior distribution in Bayesian imaging, enabling fast, accurate image reconstruction and uncertainty quantification with reduced computational resources.

Contribution

The paper presents C-Trumpets, a new architecture of conditional injective flows that outperform existing models in imaging tasks by reducing memory, computation, and handling nonlinear problems effectively.

Findings

C-Trumpets outperform regular conditional flow models on imaging tasks.

They enable fast approximation of point estimates and uncertainty quantification.

They require less memory and computation than previous variational inference methods.

Abstract

Most deep learning models for computational imaging regress a single reconstructed image. In practice, however, ill-posedness, nonlinearity, model mismatch, and noise often conspire to make such point estimates misleading or insufficient. The Bayesian approach models images and (noisy) measurements as jointly distributed random vectors and aims to approximate the posterior distribution of unknowns. Recent variational inference methods based on conditional normalizing flows are a promising alternative to traditional MCMC methods, but they come with drawbacks: excessive memory and compute demands for moderate to high resolution images and underwhelming performance on hard nonlinear problems. In this work, we propose C-Trumpets -- conditional injective flows specifically designed for imaging problems, which greatly diminish these challenges. Injectivity reduces memory footprint and training time while low-dimensional latent space together with architectural innovations like fixed-volume-change layers and skip-connection revnet layers, C-Trumpets outperform regular conditional flow models on a variety of imaging and image restoration tasks, including limited-view CT and nonlinear inverse scattering, with a lower compute and memory budget. C-Trumpets enable fast approximation of point estimates like MMSE or MAP as well as physically-meaningful uncertainty quantification.