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

Shizhan Cai, Qixin Zhang, Siyuan Ma, Siyuan Liang, Yongcheng Jing +1

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…