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

Shizhan Cai; Qixin Zhang; Siyuan Ma; Siyuan Liang; Yongcheng Jing; Dacheng Tao

cs.LGOct 10, 2026

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

Authors: Shizhan Cai, Qixin Zhang, Siyuan Ma, Siyuan Liang, Yongcheng Jing, Dacheng Tao

Organizations: College of Computing and Data Science Nanyang Technological University

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.

Figures & tables

Panel of Figure 1 — Figure 1: Performance comparison between ReSGD and other baselines on data summarization, illustrating the superior performance and stability of ReSGD under various data corruption scenarios.

Method	PermMNIST	SplitMNIST	SplitCIFAR-10	SplitCIFAR-100
Uniform	91.82 $\pm$ 0.94	78.09 $\pm$ 0.80	33.66 $\pm$ 1.11	7.06 $\pm$ 0.40
$k$ -Means Features	93.22 $\pm$ 1.47	78.29 $\pm$ 0.59	34.38 $\pm$ 1.26	8.05 $\pm$ 0.54
$k$ -Center Embedding	93.69 $\pm$ 0.59	78.70 $\pm$ 0.56	34.07 $\pm$ 2.14	8.97 $\pm$ 0.47
Hardest	88.17 $\pm$ 1.44	77.04 $\pm$ 0.84	25.35 $\pm$ 1.48	6.43 $\pm$ 0.37
iCaRL	93.68 $\pm$ 0.45	79.62 $\pm$ 0.50	33.26 $\pm$ 0.82	8.31 $\pm$ 0.22
OCS	84.86 $\pm$ 2.69	75.07 $\pm$ 0.91	34.12 $\pm$ 1.94	9.27 $\pm$ 0.45

Table 1: Performance comparison on continual learning benchmarks. We mark the Best results in bold. Our memory construction method performs on par with Greedy Coreset on the MNIST dataset and consistently outperforms other baselines on more complex datasets.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 2 — Figure 2: Performance comparison between ReSGD and baseline methods on data summarization. (a) Test accuracy across different coreset sizes for ResNet on CIFAR100. (b) and (c) evaluate the robustness of each method against varying rates of pairwise noise and symmetric noise on CIFAR100, respectively. (d) illustrates the performance under different class imbalance factors on CIFAR100. Shaded areas denote the variance across multiple runs. Overall, the results demonstrate the superior performance and stability of ReSGD under various data corruption and distribution shift scenarios.

Methods	Split CIFAR-100		Split Tiny ImageNet
	200	500	200	500
A-GEM	9.40 $\pm$ 0.05	9.42 $\pm$ 0.08	8.07 $\pm$ 0.08	8.06 $\pm$ 0.04
ER	14.18 $\pm$ 0.45	21.08 $\pm$ 0.16	8.49 $\pm$ 0.16	9.99 $\pm$ 0.29
FDR	15.32 $\pm$ 0.73	22.83 $\pm$ 0.73	8.70 $\pm$ 0.19	10.54 $\pm$ 0.21
CSReL-ER	15.35 $\pm$ 0.73	22.65 $\pm$ 0.81	8.66 $\pm$ 0.06	10.44 $\pm$ 0.17
ReSGD-ER	15.38 $\pm$ 0.25	23.21 $\pm$ 0.22	8.87 $\pm$ 0.05	10.67 $\pm$ 0.12

Appendix

Table 2: Average accuracy (%) on Split CIFAR-100 and Split Tiny-ImageNet. Results are averaged over 10 incremental tasks with memory buffers of 200 and 500 samples. We use ResNet-18 as the backbone for all experiments. Bold indicates the best performance.

Figure 3: Performance on MNIST under Asymmetric Noise. We track the evolution of the sampling probability mass on corrupted labels (red line) and the coreset purity (blue line) across 50 iterations for various coreset sizes

K\in\{50,\dots,250\}

. Despite an adversarial start where 90% of the initial mass is concentrated on noisy samples, ReSGD consistently reallocates probability to clean samples, achieving near-perfect purity across all tested coreset scales.

Figure 4: Performance on CIFAR-10 under Asymmetric Noise. Similar to the MNIST experiments, ReSGD demonstrates strong recovery capabilities on the more complex CIFAR-10 dataset. The rapid decay of the noise-associated probability mass (red) and the corresponding climb in the number of clean samples selected (blue) highlight the algorithm’s robustness to systematic label distortion across different target coreset budgets.

Figure 5: Comparison of the number of selected noisy samples across different noise ratios. We evaluate our method against baseline approaches (CSRel [ 104 ] , Greedy Coreset [ 8 ] , and PBCS [ 116 ] ) on three datasets: (a) MNIST using CNN, (b) CIFAR-10 using ResNet, and (c) CIFAR-100 using ResNet. ReSGD consistently selects the fewest noisy samples across all noise levels, demonstrating its superior robustness and ability to distinguish clean data from pairwise noisy samples.

Figure 6: Impact of

\alpha

on Average Accuracy and Forgetting. Comparison of performance metrics under different loss factors

\alpha

on the Permuted MNIST dataset. Results show that

\alpha\in[0.15,0.20]

provides an optimal trade-off for mitigating forgetting while maintaining overall accuracy.

Data Selection

Data selection chooses informative, representative, or reliable examples to improve model training and evaluation while reducing data, compute, or labeling costs.

See all 84 papers

Select-then-differentiate: Solving Bilevel Optimization with Manifold Lower-level Solution Sets
May 9, 2026Saeed Masiha, Zebang Shen, Negar Kiyavash +1Nonconvex Optimization Bilevel Optimization
EPFL School of Management of Technology, Station 5, 1015 Lausanne, Switzerland · ETH Department of Computer Science, Universitätstrasse 6, 8092 Zürich, Switzerland
Smoothed Gradient Method for Nonconvex Federated Stochastic Bilevel Optimization
Oct 4, 2026Xinwen Zhang, Peiran Yu, Zhaosong Lu +1Gradient Descent Nonconvex Stochastic Optimization
Temple University · University of Texas at Arlington · University of Minnesota
Federated stochastic bilevel optimization with fully first-order gradients
Sep 14, 2026Yihan Zhang, Rohit Dhaipule, Chiu C Tan +2Stochastic Optimization Bilevel Optimization
1Temple University · 2Stony Brook University

May 9, 2026J/K move · Enter open · S save

Select-then-differentiate: Solving Bilevel Optimization with Manifold Lower-level Solution Sets

Saeed Masiha, Zebang Shen, Negar Kiyavash, Niao He

EPFL School of Management of Technology, Station 5, 1015 Lausanne, Switzerland · ETH Department of Computer Science, Universitätstrasse 6, 8092 Zürich, Switzerland

Nonconvex Optimization Bilevel Optimization Differentiable Optimization

Open paper AskPDF arXiv

We study optimistic bilevel optimization when the lower-level problem has a non-isolated manifold of minimizers. In this setting, the hyper-objective may be non-differentiable because the upper-level criterion must choose among multiple lower-level solutions. Under a local Polyak--Łojasiewicz (PŁ) condition, we show that differentiability does not require the lower-level solution set to be a singleton: uniqueness of the optimistic selection is sufficient. This yields an explicit pseudoinverse-based hyper-gradient formula extending the classical singleton-minimizer result. We further characterize the regularity of the hyper-objective: non-degeneracy of the selected minimizer along the solution manifold yields local smoothness, while failure of uniqueness can create many non-differentiable points and failure of non-degeneracy can destroy all positive Hölder regularity of the hyper-gradient. Motivated by this theory, we propose HG-MS, a select-then-differentiate method combining explicit optimistic selection with efficient pseudoinverse-based hyper-gradient computation. Despite the nonconvex nature of optimistic selection over the lower-level solution manifold, we show that HG-MS converges to a stationary point of the optimistic objective with complexity governed by the intrinsic dimension of the solution manifold rather than its ambient dimension. Empirically, we test a practical variant of HG-MS for matched-budget LLM source reweighting. This variant preserves the select-then-differentiate principle and obtains the best GSM8K/MATH scores across the tested backbones, along with competitive or best MT-Bench instruction-following results.

Abstract

Data Selection

Explore similar work

Select-then-differentiate: Solving Bilevel Optimization with Manifold Lower-level Solution Sets

Smoothed Gradient Method for Nonconvex Federated Stochastic Bilevel Optimization

Federated stochastic bilevel optimization with fully first-order gradients