Active Data Acquisition with Side Information via Discrete Diffusion Priors

An Vuong; Thinh Nguyen

eess.IVSep 26, 2026

Active Data Acquisition with Side Information via Discrete Diffusion Priors

Authors: An Vuong, Thinh Nguyen

Organizations: Department of Electrical Engineering and Computer Science Oregon State University Corvallis, OR 97331, USA.

Abstract

Acquiring data is costly: higher measurement fidelity costs power and storage and risks collecting irrelevant content, while aggressive cost reduction can discard information that later analysis needs. We address this trade-off with an information-theoretic framework that acquires data relevant to a broad set of tasks rather than to one model. A mask policy, conditioned on side information, chooses which pixels to measure so as to maximize the mutual information between a discrete image and its partial observation under a budget; since the image entropy does not depend on the mask, this is equivalent to minimizing the conditional entropy. A frozen discrete denoising diffusion model (D3PM) supplies the posterior, and we use it in two ways: as an entropy surrogate for training a one-shot mask generator, and as the criterion for sequential greedy acquisition. The one-shot generator outperforms random masks only with care, including an unbiased gradient estimator for binary masks. With sequential acquisition, on MNIST the prior makes $8\times$ fewer errors than random at a $10\%$ budget, and on CIFAR-10 it gains $0.9$ -- $3.4$ ~dB. On fastMRI, our proposed technique using a static mask outperforms the well-known methods such as variable density and LOUPE.

Figures & tables

Figure 1: The two uses of the prior, on one MNIST test digit at $s=0.1$ . (a) Training the one-shot mask generator. From the side information $Z$ , the generator gives each pixel a keep probability $p_{i}$ ; a sampled mask $X$ yields the observation $Y$ , in which unobserved pixels take the mask state (grey). The frozen D3PM, queried at the timestep $t(s)$ whose survival matches the budget, gives per-pixel entropies $H_{i}$ , and the loss is their mean plus the budget penalty. The generator is trained with DisARM from the loss of antithetic mask pairs; no gradient passes through the prior. (b) Sequential acquisition with the same prior: starting from no observations, with only the label, each of 16 steps measures the unobserved pixels of highest entropy and updates the posterior. Insets are outputs of the trained generator and prior.

	MNIST (errors/img)		CIFAR-10 (PSNR, dB)
Summary	$s=0.1$	$0.3$	$s=0.1$	$0.3$	$0.5$
target’s own	23.0	4.74	20.45	24.80	28.01
another image’s	26.8	5.22	20.12	24.25	27.40
none	26.4	5.31	19.89	23.78	26.42
random (no generator)	46.7	24.0	20.14	24.33	27.47

Table 1: What the one-shot generator gets from its side information: the target’s own

8\times 8

summary, another test image’s summary, or none (an all-zero summary). Budget-exact sampled masks, hold-out test split, 5 seeds; seed spread at most

0.10

errors per image and

0.02

dB. On MNIST the generator outperforms random even without the true summary; on CIFAR-10 its gain needs it.

Figure 2: MNIST test images at $s=0.1$ and $0.3$ , one row per method; masks are tinted by method, and each reconstruction keeps the observed pixels and fills the rest from the prior. The one-shot generator, trained with DisARM and sampled on budget, covers the region where the digit lies, densely by $s=0.3$ ; sequential acquisition (label-greedy, 16 steps; Section 6 ) places its pixels on and around the stroke, because the prior’s entropy concentrates there, as Figure 1 (b) shows step by step. On all $2560$ test images the one-shot generator makes $22.9$ errors per image at $s=0.1$ against $46.6$ for random ( $4.7$ against $24.0$ at $s=0.3$ ), and sequential acquisition makes $5.6$ and $0.01$ .

Figure 3: Sequential acquisition against the baselines on the test split; every mask is reconstructed by the same prior except “LOUPE, own U-Net”. MNIST: label-greedy with 16 steps, errors per image (log scale). CIFAR-10: variance-greedy with 16 steps, alone and after a jittered-grid probe, PSNR. The grid probe is used on CIFAR-10 only: on MNIST, starting from the label is better than starting from a probe (Table 2 ). LOUPE is shown at the better of its top-

k

and budget-exact sampled masks (Appendix C ). The one-shot mask sees the target’s own

8\times 8

summary.

Strategy	$s=0.05$	$s=0.1$	$s=0.2$	$s=0.3$
random	60.5 (0%)	46.7 (0%)	32.1 (0%)	24.0 (0%)
probe + greedy, 1 step	40.0 (0%)	26.1 (1%)	9.22 (14%)	1.97 (50%)
probe + greedy, 16 steps	26.9 (4%)	8.45 (17%)	0.42 (83%)	0.02 (99%)
label-greedy, 4 steps	26.8 (4%)	10.0 (14%)	0.64 (81%)	0.03 (99%)
label-greedy, 16 steps	22.9 (6%)	5.63 (26%)	0.19 (92%)	0.01 (100%)

Table 2: MNIST errors per test image (lower is better; fraction of images reconstructed exactly in parentheses).

Mask	Zero-filled	U-Net B	Gain
ACS + equispaced	0.03820	0.02623	31.3%
ACS + random	0.03819	0.02696	29.4%
static search (ours)	0.03283	0.02389	27.2%
ACS + variable density	0.03382	0.02474	26.8%
learned policy, static (ours)	0.03431	0.02887	15.9%
learned policy, adaptive (ours)	0.03337	0.02840	14.9%

Table 3: Val NMSE at

s=0.25

. Gain is the relative NMSE reduction from zero-filled to U-Net B reconstruction,

1-\mathrm{NMSE}_{\mathrm{B}}/\mathrm{NMSE}_{\mathrm{zf}}

. Center-concentrated masks gain least from a reconstructor that carries a prior.

$s$	Static search (ours)	ACS + VD	LOUPE	Sequential
0.25	0.02389	0.02471 $\pm$ 0.00006	0.02759	0.02589
0.40	0.01860	0.01889 $\pm$ 0.00005	0.02213	0.01992
0.50	0.01523	0.01546 $\pm$ 0.00007	0.01995	0.01559
0.75	0.00751	0.00764 $\pm$ 0.00001	0.00966	0.00807

Table 4: Hold-out val NMSE under the judge U-Net B. ACS + VD: mean

\pm

sd over 5 seeds. Sequential: rows ranked by the coarse row-profile prior’s entropy, 16 steps.

Figure 4: Phase-encode rows sampled by the static search mask, LOUPE and ACS + variable density (

f=0

is DC); Figure 15 in Appendix D shows every static mask.

Figure 5: A validation slice of median difficulty for ACS + variable density, at $s=0.25$ , reconstructed by the hold-out judge. Top: reconstructions. Bottom: absolute error $|$ reconstruction $-$ fully sampled $|$ , from 0 (white) to a quarter of the slice’s 99.5th-percentile intensity (dark blue). PSNR and NMSE are per slice.

Figure 6: Sequential acquisition on the slice of Figure 5 at $s=0.25$ : from the ACS block, each of 16 steps adds the unmeasured k-space rows of highest prior entropy. Top: per-row entropy before the step (DC in the middle; blue: measured rows). Bottom: U-Net B reconstruction after the step, with its NMSE.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	MNIST	CIFAR-10
Batch size	512	128
Learning rate	$3\times 10^{-4}$	$2\times 10^{-5}$
Sparsity weight $\lambda$	10	10
Budget range	$[0.05,0.95]$	$[0.1,0.9]$
Gradient estimator	DisARM	DisARM
Epochs	200	200

Appendix

Table 5: Mask-generator training; DisARM uses one antithetic pair per image.

$s$	$\mathcal{H}_{\mathrm{all}}$	Full acc.	Fg recovery	Informative	Hidden-fg acc.	Errors/img
0.1	0.0564 / 0.0685	0.978 / 0.954	0.903 / 0.777	0.41 / 0.09	0.829 / 0.753	23.0 / 46.7
0.3	0.0495 / 0.0527	0.995 / 0.977	0.978 / 0.894	0.28 / 0.09	0.813 / 0.849	4.74 / 24.0
0.5	0.0476 / 0.0488	0.999 / 0.986	0.996 / 0.939	0.19 / 0.09	0.734 / 0.879	0.68 / 14.4
0.7	0.0467 / 0.0471	1.000 / 0.992	1.000 / 0.968	0.13 / 0.09	0.578 / 0.893	0.06 / 7.88

Appendix

Table 6: MNIST mask generator trained with

\mathcal{H}_{\mathrm{all}}

, learned / random, hold-out test split, 5 seeds, budget-exact sampled masks (bold = better). Hidden-fg acc. scores only the digit pixels left hidden; it is biased against good masks, since a mask that observes most of the digit leaves only the pixels the prior finds hardest.

$s$	$\mathcal{H}_{\mathrm{all}}$	Full acc.	Non-black rec.	PSNR (dB)
0.1	0.790 / 0.786	0.675 / 0.670	0.676 / 0.670	20.45 / 20.14
0.3	0.420 / 0.435	0.824 / 0.816	0.824 / 0.815	24.80 / 24.33
0.5	0.236 / 0.250	0.901 / 0.894	0.901 / 0.893	28.01 / 27.47
0.7	0.116 / 0.126	0.951 / 0.946	0.951 / 0.946	31.54 / 30.87

Appendix

Table 7: CIFAR-10 mask generator trained with

\mathcal{H}_{\mathrm{all}}

, learned / random, hold-out test split, 5 seeds, budget-exact sampled masks (bold = better). Non-black rec. is the fraction of non-black pixels reconstructed correctly.

Figure 7: Top: the value of the whole-image surrogate objective the mask is trained on, with DisARM. Bottom: a criterion the mask is not optimized for, on the same checkpoints, with budget-exact sampled masks. On MNIST at

s=0.7

both improve together, and digit recovery stays above random throughout (

0.992

at epoch 10 and

1.000

at 200, against

0.968

). On CIFAR-10 at

s=0.1

they part ways: the objective rises above random’s after epoch 40 while PSNR keeps improving and stays above random (

20.31

dB at epoch 40,

20.45

at 200, against

20.15

Figure 8: The MNIST advantage in digit recovery (left) holds at every budget and narrows as the budget grows (

0.903

vs.

0.777

s=0.1

1.000

vs.

0.968

s=0.7

); accuracy on the digit pixels the mask leaves hidden (centre) is worse than random for every

s\geq 0.3

, a measure biased against a mask that observes more of the digit. Mask trained with DisARM, final checkpoint, budget-exact sampled masks. On CIFAR-10 (right) the mask is

0.3

0.7

dB above random at every budget.

Figure 9: CIFAR-10 test images at $s=0.1$ and $0.3$ , one row per method; masks are tinted by method, and each reconstruction keeps the observed pixels and fills the rest with the prior’s posterior mean. The one-shot generator, trained with DisARM and sampled on budget, scatters its pixels over the whole image much as a random mask does; sequential acquisition (variance-greedy, 16 steps; Section 6 ) spreads them more evenly and more densely where the image has structure.

Figure 10: Sequential acquisition on MNIST (exact values in Table 2 ). Left: errors per image, log scale. Right: fraction of test images reconstructed exactly.

Figure 11: MNIST test images at $s=0.1$ and $0.2$ : each method’s mask, tinted in its colour, and the reconstruction under the prior (LOUPE’s own U-Net in its own row). LOUPE’s masks are sampled from its learned map.

Figure 12: CIFAR-10 test images at $s=0.1$ and $0.3$ , as in Figure 11 ; the sequential-acquisition row is variance-greedy with 16 steps. The one-shot mask, trained with DisARM and sampled on budget, scatters its pixels over the image much as a random mask does.

Mask	$s=0.25$	$s=0.40$	$s=0.50$	$s=0.75$
ACS + variable density	0.02471	0.01889	0.01546	0.00764
static search (ours)	0.02389	0.01860	0.01523	0.00751
sequential acquisition, entropy, 1 step	0.03031	0.02571	0.02061	0.00751
sequential acquisition, entropy, 4 steps	0.02670	0.02254	0.01916	0.00941
sequential acquisition, entropy, 16 steps	0.02589	0.01992	0.01559	0.00807
control, expected energy, 1 step	0.03359	0.02928	0.02381	0.01145

Appendix

Table 8: Hold-out val NMSE under the judge U-Net B for sequential row acquisition with the coarse row-profile prior: starting from the ACS block, rows are added in 1, 4 or 16 steps, ranked by the prior’s predictive entropy or, as a control, by its expected energy. ACS + variable density is the mean over 5 seeds. Bold: the lowest NMSE per budget, and any entry within two standard deviations of it, measured by the seed spread of ACS + variable density.

Figure 13: The same masks under zero-filled reconstruction, with no network. The sequential-acquisition line has lower NMSE than variable density here at

s\leq 0.50

Figure 14: Row masks on the 199 hold-out validation slices, reconstructed by the hold-out judge U-Net B. Top: NMSE (log scale), PSNR and SSIM against the budget. Bottom: gain over ACS + variable density. The static search mask is best on NMSE and PSNR at every budget and LOUPE is last; the center-concentrated relaxed row profile and energy oracle, with sequential acquisition, lead on SSIM at

s=0.40

–

0.75

. The sequential-acquisition line (16 steps; Table 8 ) trails variable density on NMSE at every budget except

s=0.50

, where it is level.

Figure 15: Phase-encode rows sampled by every static mask, including the relaxed row profile and the energy oracle (

f=0

is DC).

Figure 16: A validation slice of 75th-percentile difficulty for ACS + variable density, at $s=0.25$ , reconstructed by the hold-out judge. Top: reconstructions. Bottom: absolute error $|$ reconstruction $-$ fully sampled $|$ , from 0 (white) to a quarter of the slice’s 99.5th-percentile intensity (dark blue). PSNR and NMSE are per slice.

$s$	Without	With mask state	Fewer errors	Seeds
0.10	$28.72\pm 0.36$	$\mathbf{25.73}\pm 0.44$	$10.4\pm 0.9\%$	5/5
0.20	$10.08\pm 0.22$	$\mathbf{8.83}\pm 0.33$	$12.4\pm 2.8\%$	5/5
0.25	$4.98\pm 0.14$	$\mathbf{4.28}\pm 0.25$	$14.2\pm 4.5\%$	5/5
0.30	$2.07\pm 0.08$	$\mathbf{1.76}\pm 0.15$	$15.0\pm 7.4\%$	5/5
0.40	$0.22\pm 0.01$	$\mathbf{0.19}\pm 0.02$	$14.4\pm 10.8\%$	4/5

Appendix

Table 9: MNIST errors per test image with and without the separate mask state, probe + one greedy step, mean

\pm

standard deviation over

5

fine-tuning seeds. The last column counts the seeds in which the mask state has fewer errors.

Active Feature Acquisition

Active Feature Acquisition studies how models can selectively obtain costly or missing inputs to improve predictions within limited budgets.

See all 19 papers

Improving Sampling for Masked Diffusion Models via Information Gain
Feb 20, 2026Kaisen Yang, Jayden Teoh, Kaicheng Yang +2Diffusion Sampling Diffusion Language Model Decoding
Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University +2
Policy-Mask: Budget-Aware Adaptation of Diffusion Features to Discriminative Tasks
Oct 20, 2026Ivan Listopadov, Maksim Nakhodnov, Aibek AlanovDiffusion Model Adaptation Feature Selection
HSE University · MSU · Constructor University +1
Noise Scheduling as Information-Guided Allocation in Diffusion Training
Feb 20, 2026Gabriel Raya, Bac Nguyen, Georgios Batzolis +6Diffusion Timestep Scheduling
Tilburg University & JADS · Sony AI · University of Cambridge +2

Feb 20, 2026J/K move · Enter open · S save

Improving Sampling for Masked Diffusion Models via Information Gain

Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang+1

Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University+2

Diffusion Sampling Diffusion Language Model Decoding Masked Diffusion Models

Open paper AskPDF arXiv

Masked Diffusion Models (MDMs) enable flexible decoding orders, yet existing samplers remain largely greedy, selecting locally certain tokens without accounting for their downstream effects. We show that this myopia can increase cumulative uncertainty and lead to suboptimal generation. To address this, we propose the Info-Gain Sampler, a training-free decoding method that uses the bidirectional structure of MDMs to balance immediate uncertainty with the information gained over remaining masked positions. Across reasoning, coding, creative writing, and image generation tasks, Info-Gain Sampler consistently outperforms existing MDM samplers, improving average reasoning accuracy by 2.9--11.6 percentage points and achieving a 62.8% average win rate in creative writing. The code is available at https://github.com/yks23/Information-Gain-Sampler.

Abstract

Active Feature Acquisition

Explore similar work

Improving Sampling for Masked Diffusion Models via Information Gain

Policy-Mask: Budget-Aware Adaptation of Diffusion Features to Discriminative Tasks

Noise Scheduling as Information-Guided Allocation in Diffusion Training