Probabilistic Multi-Source Data Selection via Learnable Categorical Priors: Variance Reduction and Guaranteed Derandomization

Abstract

Bilevel data selection is a fundamental task with numerous applications in machine learning and statistics. A common strategy for tackling this problem is to relax the original bilevel combinatorial optimization into a continuous one. However, existing methods, including 0-1 relaxation and probabilistic relaxation, still suffer from several limitations, such as unstable gradient estimation and infeasible rounding. To address these issues, this paper proposes a new probabilistic relaxation for the multi-source bilevel data selection problem. Our key idea is to introduce multiple learnable categorical priors to parameterize a probability distribution supported only on feasible data subsets, thus naturally yielding a lossless rounding scheme. We further show that our proposed relaxation admits an effective two-point gradient estimator, which can simultaneously avoid the inverse-Hessian computation required by existing 0-1 relaxation and the unbounded score term appearing in prior probabilistic relaxation. Furthermore, we develop a reducible representation of this estimator, which substantially reduces the number of lower-level retrainings required for gradient computation. Built upon this reducible two-point estimator, we then design a novel stochastic gradient algorithm, termed ReSGD, for multi-source bilevel data selection problem, and establish a rigorous approximation guarantee for the subset returned by ReSGD under weak submodularity. Finally, extensive experiments are conducted to validate the effectiveness of our proposed ReSGD.