ShuffleFlow: Scalable Posterior Inference for Bayesian Inverse Imaging

2026-06-19

AI Summary

This work addresses the limited scalability of existing flow-based variational inference methods in large-scale image inversion. To overcome this challenge, the authors propose ShuffleFlow, a novel framework that partitions an image into a stack of sub-images via pixel unshuffling. By integrating neural fields to encode spatial structure and employing shared conditional normalizing flows to model channel-wise dependencies, ShuffleFlow enables efficient Bayesian posterior inference. The method achieves significantly improved computational efficiency while preserving high sample quality, outperforming diffusion-based samplers on both linear and nonlinear imaging inverse problems. Consequently, ShuffleFlow offers a scalable and practical solution for large-scale image reconstruction tasks.

Abstract

Variational inference (VI) is a powerful method for principled posterior inference for scientific inverse imaging. VI learns the posterior distribution, often with a flow-based network, which can cheaply generate posterior samples upon optimization, and can flexibly incorporate score-based or classic priors. However, its application to large-scale image reconstruction is severely hindered by the poor scalability of the flow-based networks. In this work, we introduce ShuffleFlow, a scalable VI framework to address this challenge. Our method breaks down the problem into three parts: a pixel-unshuffling-based image coordinate sampler, a neural field as feature encoder, and a conditional normalizing flow (CNF) as posterior estimator. Specifically, our framework partitions an image into a stack of sub-images with pixel-unshuffling and uses a shared CNF to model the joint distribution of the sub-image stack. We condition the CNF on the output of a neural field, which embeds feature vectors corresponding to pixel-unshuffling sample locations to capture spatial structures, and share the flow's latent variable across the channels to model their correlations. We demonstrate our method's effectiveness and efficiency on both linear and nonlinear imaging inverse problems, and show its ability to more rapidly generate a high-sample-count posterior than diffusion samplers.

Problem

• variational inference
• flow-based networks
• scalability
• Bayesian inverse imaging
• posterior inference

Innovation

• ShuffleFlow
• scalable variational inference
• pixel-unshuffling
• conditional normalizing flow
• neural field