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

Technical Summary: Flow-Based Generative Modeling for Optimizing Sampling Policies in Compressed Sensing

Problem Statement
Modern applications in signal processing and medical imaging often require acquiring high-dimensional signals under strict resource constraints. Classical sampling theory dictates that accurate reconstruction requires measurements proportional to the signal's ambient dimension, a requirement that is frequently impractical. Compressed Sensing (CS) addresses this by enabling the recovery of sparse signals from fewer measurements, provided the measurement operator satisfies specific conditions (e.g., incoherence). However, the quality of reconstruction is highly dependent on the choice of the sampling scheme. Finding an optimal subsampling pattern is an NP-hard combinatorial problem because the sampling mask is discrete and non-differentiable, preventing naive end-to-end optimization. Furthermore, the search space grows exponentially with the number of sensor locations, making exhaustive exploration infeasible. Existing generative approaches, such as Deep Probabilistic Subsampling (DPS) or Active Diffusion Subsampling (ADS), often produce fixed masks or suffer from high computational latency, limiting their adaptability and real-time clinical utility.

Methodology
The authors propose a task-aware, flow-based generative framework that reformulates the conventional Flow Matching (FM) training paradigm. Instead of modeling the data manifold directly, the model operates on the manifold of the sampling space to learn optimal subsampling masks.

• Flow Matching on the Sampling Manifold: The framework utilizes a neural Ordinary Differential Equation (ODE) to transport points from a standard Gaussian noise distribution ($p_0$) to a target distribution of sampling operators ($p_1$). The model learns a time-dependent vector field $v_\phi(x, t)$ that defines a generative trajectory.
• Task-Conditioned Optimization: Unlike standard FM which relies on linear coupling between noise and data, this framework couples the generative flow with a downstream CS task solver ($f_\theta$). The training objective minimizes a composite loss function combining the standard FM regression loss ($L_{FM}$) and a task-specific loss ($L_{task}$), weighted by learnable uncertainty parameters ($\sigma_1, \sigma_2$). This allows the model to jointly optimize the sampling policy and the task performance.
• Differentiable Sigmoid Gating: To address the non-differentiability of discrete binary masks, the continuous output of the flow model ($x_t$) is passed through a custom differentiable sigmoid gate. This gate employs a time-dependent steepness function $s(t)$ and a percentile threshold $\tau$ to produce soft masks ($g(x_t) \in (0, 1)$). As the time step $t$ increases, the transition becomes sharper, effectively regularizing the entropy of the mask generation process.
• Inference Modes: The framework supports two operational modes:
  1. Random-to-Optimized: Initial random measurements are taken, features are extracted, and the flow model generates an optimized mask for subsequent measurements. This is suitable for scenarios where measurements are inexpensive or signals are static (e.g., MNIST, CelebA).
  2. Adaptive Sensing: Low-frequency (LF) k-space data is acquired first. A feature extractor conditions the flow model to determine the optimal sampling pattern for the remaining high-frequency bands. This mode is designed for costly measurements, specifically MRI acceleration.

Key Contributions

• Novel Paradigm: The paper introduces a proof-of-concept where Flow Matching is trained within the manifold of sampling operators rather than the signal space, tailored specifically for downstream CS tasks.
• Continuous and Adaptive Policy: The framework provides a continuous, time-continuous parametrization of the sampling scheme, offering a smooth evolution from coarse structures to task-specific refinements, unlike static mask generation methods.
• Joint Optimization: The method enables end-to-end joint optimization of the sampling strategy and the reconstruction/classification task, leveraging the task prior to guide the generative process.
• Unified Framework: It offers a flexible approach applicable to various inverse problems, including image classification, image reconstruction, and MRI acceleration, without requiring retraining for different measurement operators.

Experimental Results
The framework was evaluated on MNIST, CelebA, and the fastMRI knee dataset:

• MNIST Classification: The flow-based approach improved classification accuracy by over 20% compared to classifiers trained on random masks. While it trailed the A-DPS model at specific fixed sampling ratios, it demonstrated greater versatility across a range of ratios (1% to 8%) without needing ratio-specific training.
• Image Reconstruction (CelebA & MNIST): The method achieved state-of-the-art performance in image reconstruction. On the CelebA dataset, it reached a Peak Signal-to-Noise Ratio (PSNR) of 25.17 dB at a 5% subsampling rate. It outperformed the state-of-the-art ADS method, particularly at aggressive (low) sampling rates.
• MRI Acceleration: In the fastMRI experiments, the framework achieved a Structural Similarity Index (SSIM) of 0.7163 at 8× acceleration, outperforming all baselines (LOUPE, DPS, ADS, A-DPS) at this high acceleration factor. At 4× acceleration, it performed comparably to other methods.
• Computational Efficiency: The inference time for generating a sampling mask was approximately 127–130 ms on an NVIDIA GeForce RTX 4070. This is significantly faster than the acquisition time per k-space line in MRI (500–2000 ms) and drastically faster than the ADS method's reported acquisition time per line (3040 ms), suggesting feasibility for real-time clinical deployment.

Significance and Claims
The authors claim that this work establishes the fundamental feasibility of using task-conditioned flow-based generative models to optimize sampling policies for compressed sensing. The study highlights that:

1. Task-Conditioning is Effective: Integrating a downstream task prior directly into the flow matching training process significantly enhances performance in inverse problems.
2. Sampling Manifold Constraints: The sampling space possesses intrinsic structural constraints (binarity, sparsity) that can be leveraged to improve training efficiency compared to modeling the natural image manifold.
3. Clinical Viability: The low computational overhead and adaptive nature of the framework make it a promising candidate for real-world medical imaging applications, particularly in reducing MRI scan times and improving diagnostic data quality.
4. Generalizability: The approach offers a unified, flexible strategy for designing data- and task-driven sensing schemes applicable to a broad range of inverse problems beyond the specific tasks tested.

The paper positions itself as a proof-of-concept study, acknowledging that while it achieves state-of-the-art results in specific reconstruction tasks, the primary contribution is the demonstration of a new, flexible framework for learning sampling policies that can be adapted to various inverse problems.