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
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 ± 0.94
78.09 ± 0.80
33.66 ± 1.11
7.06 ± 0.40
k -Means Features
93.22 ± 1.47
78.29 ± 0.59
34.38 ± 1.26
8.05 ± 0.54
k -Center Embedding
93.69 ± 0.59
78.70 ± 0.56
34.07 ± 2.14
8.97 ± 0.47
Hardest
88.17 ± 1.44
77.04 ± 0.84
25.35 ± 1.48
6.43 ± 0.37
iCaRL
93.68 ± 0.45
79.62 ± 0.50
33.26 ± 0.82
8.31 ± 0.22
OCS
84.86 ± 2.69
75.07 ± 0.91
34.12 ± 1.94
9.27 ± 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
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 ± 0.05
9.42 ± 0.08
8.07 ± 0.08
8.06 ± 0.04
ER
14.18 ± 0.45
21.08 ± 0.16
8.49 ± 0.16
9.99 ± 0.29
FDR
15.32 ± 0.73
22.83 ± 0.73
8.70 ± 0.19
10.54 ± 0.21
CSReL-ER
15.35 ± 0.73
22.65 ± 0.81
8.66 ± 0.06
10.44 ± 0.17
ReSGD-ER
15.38 ± 0.25
23.21 ± 0.22
8.87 ± 0.05
10.67 ± 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∈{50,…,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 α on Average Accuracy and Forgetting. Comparison of performance metrics under different loss factors α on the Permuted MNIST dataset. Results show that α∈[0.15,0.20] provides an optimal trade-off for mitigating forgetting while maintaining overall accuracy.
Data selection chooses informative, representative, or reliable examples to improve model training and evaluation while reducing data, compute, or labeling costs.
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.
Saeed Masiha, Zebang Shen, Negar Kiyavash +1
EPFL School of Management of Technology, Station 5, 1015 Lausanne, Switzerland · ETH Department of Computer Science, Universitätstrasse 6, 8092 Zürich, Switzerland
In recent years, federated stochastic bilevel optimization has attracted increasing attention due to its wide range of applications in machine learning. To reduce the computational overhead associated with second-order Hessian and Jacobian matrices, several first-order methods have been proposed. However, existing methods typically impose restrictive assumptions on the lower-level function, suffer from a strong dependence on the condition number in their convergence rates, and require different learning-rate scales for variables across the upper- and lower-level problems, limiting their practical applicability and complicating hyperparameter tuning. To address these challenges, we propose a stochastic doubly smoothed gradient method for nonconvex federated stochastic bilevel optimization problems, which decouples the learning rates of upper- and lower-level variables and does not require a strongly-convex lower-level loss function. We establish rigorous theoretical guarantees for the proposed algorithm, demonstrating an improved convergence rate of O(κ15/2/ε5) and a communication complexity of O(κ4/ε3), where κ denotes the condition number and ε represents the solution accuracy. Notably, these bounds exhibit significantly better dependence on the condition number κ than those of existing methods. Extensive experiments validate the effectiveness of our algorithm.
Xinwen Zhang, Peiran Yu, Zhaosong Lu +1
Temple University · University of Texas at Arlington · University of Minnesota
Federated stochastic bilevel optimization has been actively studied in recent years due to its widespread applications in machine learning. However, most existing federated stochastic bilevel optimization algorithms require the computation of second-order Hessian and Jacobian matrices, which leads to longer running times in practice. To address these challenges, we propose a novel federated stochastic variance-reduced bilevel gradient descent algorithm that relies solely on first-order oracles. Specifically, our approach does not require the computation of second-order Hessian and Jacobian matrices, significantly reducing running time. Furthermore, we introduce a novel learning rate mechanism, i.e., a constant single-timescale learning rate, to coordinate the update of different variables. We also present a new strategy to establish the convergence rate of our algorithm. Finally, the extensive experimental results confirm the efficacy of our proposed algorithm.