ShuffleFlow: Scalable Posterior Inference for Bayesian Inverse Imaging

IEEE International Conference on Computational Photography (ICCP), 2026
Best Poster Honorable Mention, Midwest Machine Learning Symposium (MMLS), 2025

Tianao Li1,3, Tjitske Starkenburg1,3, Yu Sun2, Emma Alexander1,3
1Northwestern University, 2John Hopkins University, 3NSF-Simons AI Institute for the Sky (SkAI)
Paper arXiv Code
ShuffleFlow pipeline overview
Overview of ShuffleFlow. We propose a new VI framework for inverse problems that decomposes the posterior modeling of a large image into three parts: coordinate sampling, partial image representation, and light-weight uncertainty quantification. Specifically, a pixel-unshuffling-based coordinate sampler partitions an image into a stack of sub-images, a neural field encodes spatial features for each sub-image, and a conditional normalizing flow estimates the joint distribution of the sub-image stack. Our framework's modular structure is agnostic to specific network architectures and allows easy integration of advances in neural fields, normalizing flow networks, and generative priors.

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.

ShuffleFlow scalability
Scalability comparison. (a) VI methods scale dramatically better than diffusion samplers for high-resolution posterior estimation. (b,c,d) Flow network size, optimization time, and VRAM vs. image dimension across VI methods. When the computation is dominated by the flow network (TV prior, dashed lines), our method (purple) shows the theoretically predicted growth trend compared to DPI (black). Even when an expensive score prior dominates (solid lines), our method still shows a substantial efficiency advantage.
Posterior quality comparison
Posterior sample efficiency
Results on nonlinear Fourier phase retrieval. (a) We draw 128 samples from each method and filter the samples into different modes if a bimodal posterior exists. We show the mean image from each mode, and PSNR, SSIM, and LPIPS are calculated with respect to the ground truth of each mode. Our method correctly captures the bimodal posterior among VI methods with the best image quality. DAPS stands out with the best image quality among all. (b) We compare our method with DAPS and DPS in a large sample count (10,000, left) and a small inference time (16 min, right) scenario, showing the t-SNE visualization for the posterior samples (top) and pixel histograms (bottom). Our method is able to generate 10,000 posterior samples in 16 minutes (including both training and sampling time), resulting in a high-resolution bimodal posterior density in the histogram. In contrast, DAPS and DPS generate spiky and noisy posterior densities in a similar time and need 20× more time to generate a high-resolution posterior.

Citation

@inproceedings{li2026shuffleflow,
  title={ShuffleFlow: Scalable Posterior Inference for Bayesian Inverse Imaging},
  author={Li, Tianao and Starkenburg, Tjitske and Sun, Yu and Alexander, Emma},
  booktitle={IEEE International Conference on Computational Photography (ICCP)},
  year={2026},
  organization={IEEE}
}

Acknowledgements

We gratefully acknowledge the support of the NSF-Simons AI Institute for the Sky (SkAI) via grants NSF AST-2421845 and Simons Foundation MPS-AI-00010513. This material is based upon work supported by the U.S. National Science Foundation under Award No. 2542022. The authors would like to thank Bryan Pardo, He Sun, and Yi-Chun Hung for helpful discussions.