Active Diffusion Subsampling

TL;DR

Active Diffusion Subsampling (ADS) proposes a white-box, diffusion-model–based approach to adaptive subsampling that jointly designs measurement masks and reconstructs signals by guiding the reverse diffusion with maximum-entropy information gains. By maintaining a belief over the true state $\mathbf{x}$ via a particle ensemble and using Tweedie-based denoising to estimate $\hat{\mathbf{x}}_0$, ADS computes a tractable entropy of the predicted measurements $p(\mathbf{y}_t|A_t,\mathbf{y}_{t-1})$ and selects the next sampling region that maximizes information about $\mathbf{x}$. The method requires no task-specific retraining and yields interpretable, information-theoretic acquisition decisions; it achieves substantial improvements over fixed masks on MNIST, CelebA, and ultrasound, and competitive performance with supervised MRI methods. Limitations include inference-time constraints and reduced gains at very low sampling rates, motivating future work on accelerated posterior sampling and batch-design extensions across more domains.

Abstract

Subsampling is commonly used to mitigate costs associated with data acquisition, such as time or energy requirements, motivating the development of algorithms for estimating the fully-sampled signal of interest $x$ from partially observed measurements $y$. In maximum entropy sampling, one selects measurement locations that are expected to have the highest entropy, so as to minimize uncertainty about $x$. This approach relies on an accurate model of the posterior distribution over future measurements, given the measurements observed so far. Recently, diffusion models have been shown to produce high-quality posterior samples of high-dimensional signals using guided diffusion. In this work, we propose Active Diffusion Subsampling (ADS), a method for designing intelligent subsampling masks using guided diffusion in which the model tracks a distribution of beliefs over the true state of $x$ throughout the reverse diffusion process, progressively decreasing its uncertainty by actively choosing to acquire measurements with maximum expected entropy, ultimately producing the posterior distribution $p(x \mid y)$. ADS can be applied using pre-trained diffusion models for any subsampling rate, and does not require task-specific retraining - just the specification of a measurement model. Furthermore, the maximum entropy sampling policy employed by ADS is interpretable, enhancing transparency relative to existing methods using black-box policies. Code is available at https://active-diffusion-subsampling.github.io/.

Figures (16)

• Figure 1: Active Diffusion Subsampling jointly designs a subsampling mask and reconstructs the target signal in a single reverse diffusion process.
• Figure 2: Schematic overview of the proposed Active Diffusion Sampling (ADS) method.
• Figure 3: Illustration of a single action selection using ADS. ① shows the current batch of partially-denoised images $\{{\mathbf x}_{\tau'}^{(i)}\}$ at diffusion step $\tau'$. In ②, this batch of particles is mapped using Tweedie's formula to a batch of fully-denoised particles $\{\hat{{\mathbf x}}_{0}^{(i)}\}$, constituting the belief distribution at time $\tau'$. The forward model $f$ is then applied in ③ to simulate the set of measurements that result from the belief distribution, used to approximate the measurement posterior as a GMM. Given this GMM, the measurement entropy is computed in ④ using Equation \ref{['eq:final_policy']}. Finally, in ⑤, the maximum entropy line is selected as the next measurement location.
• Figure 4: Comparison of ADS (ours) with two non-adaptive baselines. Evaluated based on reconstruction Mean Absolute Error (MAE) on $N=500$ unseen samples from the MNIST test set. Note that MAE is plotted on a log scale.
• Figure 5: In (a), PSNR ($\uparrow$) scores for ADS vs DPS with random and data variance sampling on N=100 unseen samples from the CelebA dataset are plotted, for increasing numbers of measurements. A measurement here is a $4\times4$ box of pixels. (b) shows some examples of ADS inference versus DPS with random measurements on the CelebA dataset from the evaluation with 200 boxes sampled.