# Flow-Based Generative Modeling for Optimizing Sampling Policies in Compressed Sensing Applications

arXiv: https://arxiv.org/abs/2606.00078 · HTML: https://mlsanity.com/paper/2606.00078/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 39 references that are in the map.
- First posted on arXiv 22 May 2026.

## Builds on (most cited first, 21 of 21)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [Exploiting Diffusion Prior for Real-World Image Super-Resolution](https://mlsanity.com/paper/2305.07015/) (2023)
- [fastMRI: An Open Dataset and Benchmarks for Accelerated MRI](https://mlsanity.com/paper/1811.08839/) (2018)
- [MoDL: Model Based Deep Learning Architecture for Inverse Problems](https://mlsanity.com/paper/1712.02862/) (2017)
- [Inpaint Anything: Segment Anything Meets Image Inpainting](https://mlsanity.com/paper/2304.06790/) (2023)
- [Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets](https://mlsanity.com/paper/1903.05662/) (2019)
- [Plug-and-Play Methods Provably Converge with Properly Trained Denoisers](https://mlsanity.com/paper/1905.05406/) (2019)
- [A Survey on Diffusion Models for Inverse Problems](https://mlsanity.com/paper/2410.00083/) (2024)
- [Hierarchical Integration Diffusion Model for Realistic Image Deblurring](https://mlsanity.com/paper/2305.12966/) (2023)
- [A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning](https://mlsanity.com/paper/2110.01515/) (2021)
- [The Principles of Diffusion Models](https://mlsanity.com/paper/2510.21890/) (2025)
- [Data-Driven Sparse Sensor Placement for Reconstruction](https://mlsanity.com/paper/1701.07569/) (2017)
- [Deep Compressed Sensing](https://mlsanity.com/paper/1905.06723/) (2019)
- [A Theoretical Justification for Image Inpainting using Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2302.01217/) (2023)
- [Learning-based Optimization of the Under-sampling Pattern in MRI](https://mlsanity.com/paper/1901.01960/) (2019)
- [Active Diffusion Subsampling](https://mlsanity.com/paper/2406.14388/) (2024)
- [Differentiable Fast Top-K Selection for Large-Scale Recommendation](https://mlsanity.com/paper/2510.11472/) (2025)

## Abstract (arXiv)

Numerous modern applications in signal processing and medical imaging necessitate acquiring high-dimensional signals under tight resource constraints. Traditional sampling theory suggests that accurate signal reconstruction requires a number of measurements proportional to the signal's ambient dimension, a requirement often too expensive or impractical. Compressed sensing challenges this notion by demonstrating that sparse signals can be recovered with fewer measurements, provided the measurement operator meets certain conditions. This proof-of-concept study presents a task-aware flow-based generative framework -- a reformulation of the conventional Flow Matching training paradigm with a flow model trained to optimize subsampling in compressed sensing applications. We establish the fundamental feasibility of the proposed framework of learning subsampling masks that substantially enhance the performance of compressed sensing for image classification, image reconstruction, and MRI acceleration. For the image reconstruction task, our method demonstrated state-of-the-art performance, achieving Peak Signal-to-Noise Ratio of 25.17 dB at the subsampling rate of 5% on the CelebA dataset and 29.24 dB when reconstructing 8× accelerated MRI measurements (fastMRI dataset) with the minimal computational overhead. These results highlight the effectiveness of task-conditioning within generative flow models and reveal a promising direction for representation learning strategies. Overall, the proposed framework offers a unified, flexible approach to designing data- and task-driven sensing schemes that can be potentially adapted to a broad range of inverse problems.