Regional Mass Audits for Correcting Neural Samplers

Neural samplers can generate samples from every mode of a target distribution while producing them in the wrong proportions. We propose regional mass auditing to guide correction using target probabilities or reference data. The audit selects an under-represented region and estimates its probability deficit on data held out from selection. This deficit guides where to obtain additional samples and which examples to use for fine-tuning. Valid probability bounds determine how many samples can be added without over-filling the region at a specified confidence level. When the additional samples follow the target distribution conditioned on the region, the same bounds give a computable lower bound on the reduction in Kullback-Leibler divergence. We compare audit-guided and untargeted correction at matched sampling or training budgets on synthetic distributions, particle systems, and proteins. Audit-selected simulation starts improve particle-distance and protein free-energy estimates. Fine-tuning neural samplers with selected data improves regional probabilities in newly generated samples. For proteins, fine-tuning coarse-grained potentials with selected reference structures also improves several structural and free-energy estimates.