PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.14942v1cs.LGcs.AI

TL;DR

Data selection seeks small training subsets that preserve full-data performance, but existing relaxations face gradient-estimation and feasible-rounding limitations, especially for multi-source problems. This paper proposes a feasible-support probabilistic relaxation with a reducible two-point estimator and ReSGD; it reports an approximation guarantee under weak submodularity and experiments validating effectiveness in data summarization and continual learning.

  • Problem

    Existing relaxations face costly or unstable gradient estimation and may yield infeasible subsets, while multi-source selection also involves heterogeneous data sources.

    Source: §1 · p. 2, §2.1 · p. 4

  • Method

    The paper parameterizes a distribution over feasible subsets with learnable categorical priors, derives a reducible two-point gradient estimator, and builds ReSGD for multi-source data selection.

    Source: §1 · p. 2, §4 · p. 7

  • Results

    ReSGD is reported effective on data summarization and continual learning, and its returned subset has an approximation guarantee under weak submodularity.

    Source: §6 · p. 9

  • Takeaways & Limitations

    The approach combines feasible-subset sampling with a reduced-retraining gradient representation for multi-source bilevel data selection.

    Source: §1 · p. 2, §6 · p. 9

    The Bernoulli-product relaxation can have an unbounded score factor near boundary probabilities and can sample subsets exceeding the K-size constraint.

    Source: §2.1 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Data selection seeks a small training subset that preserves full-dataset performance, but existing continuous relaxations face gradient, multi-source, rounding, and subset-guarantee challenges. The paper proposes a feasible-support probabilistic relaxation, a two-point estimator and ReSGD, with a weak-submodularity approximation guarantee and experiments in data summarization and continual learning.

1 Introduction

  • Data selection seeks a small training subset that achieves performance comparable to training on the full dataset, with applications including efficient training and continual learning.Continual-learning systems can retain a small rehearsal memory of previously seen data for replay on new tasks.Source: §1 · p. 1
  • 0-1 relaxation requires costly inverse-Hessian approximation, while probabilistic relaxation uses unbounded, potentially unstable, high-variance policy-gradient estimators.These existing approaches primarily address single-source selection, although heterogeneous sources can have cross-source correlations and redundant selections.Source: §1 · p. 2
  • Existing rounding is either heuristic or can return infeasible subsets, and prior convergence analyses focus on continuous stationary points rather than the quality of the selected subset.The paper frames rigorous guarantees for the returned subset as an open issue in bilevel data selection.Source: §1 · p. 2
  • The proposed probabilistic relaxation uses learnable categorical priors to place probability mass only on feasible subsets, yielding lossless rounding.This directly addresses the feasibility concern associated with sampling from prior probabilistic relaxations.Source: §1 · p. 2
  • Its two-point gradient estimator avoids inverse-Hessian computation and the unbounded score term, improving stability and reducing variance.The paper also gives a reducible representation that lowers gradient-computation retrainings from O(h + sum_i n_i) to O(h), where h is the number of sources and sum_i n_i is total sample size.Source: §1 · p. 2
  • The authors develop ReSGD from the reducible estimator and establish an approximation guarantee for its returned subset under (γ, β)-weak submodularity.The supplied passage does not state the approximation factor.Source: §1 · p. 3
  • Experiments demonstrate ReSGD's practical efficacy on data summarization and continual learning.Source: §1 · p. 3

2 Bilevel Data Selection: From Single Dataset to Multi-Source Scenarios

The section frames data selection as a bilevel combinatorial problem and reviews continuous relaxations for single-source selection before defining the multi-source setting, where joint selections must respect source-specific budgets.

2.1 Single-Source Bilevel Data Selection: Background and Relaxations

  • The bilevel selection problem is NP-hard, motivating continuous relaxations that permit approximate gradient-based optimization.The section reviews 0-1 and probabilistic relaxations as two representative approaches.Source: §2.1 · p. 3
  • 0-1 relaxation optimizes a fractional mask using implicit differentiation, but its hypergradient requires an inverse Hessian with substantial time and memory overhead.The probabilistic relaxation instead models each mask as an independent Bernoulli variable and optimizes its learnable parameters through policy gradients.Source: §2.1 · pp. 3–4
  • Bernoulli policy-gradient estimation can become unbounded near parameter boundaries, while sampling from the product distribution can return subsets exceeding the cardinality constraint.The estimator is unbounded when a selected mask value is paired with a probability approaching zero, or an unselected value with a probability approaching one.Source: §2.1 · p. 4

2.2 Multi-Source Bilevel Data Selection

  • Multi-source selection allocates at most K_i samples from each heterogeneous source to form a jointly evaluated subset.A source may represent a task, domain, or data provider; selecting samples independently per source can ignore cross-source correlations and create redundant selections.Source: §2.2 · p. 4

3 A New Probabilistic Relaxation via Learnable Categorical Priors

The paper constructs a multi-source relaxation whose sampling distribution is supported only on feasible subsets. It then derives a two-point estimator and a reducible representation to make gradient computation more stable and less costly.

3 A New Probabilistic Relaxation via Learnable Categorical Priors

  • The categorical-prior construction restricts probability mass to feasible subsets, so sampling satisfies each source's budget and preserves the relaxed objective in expectation.For a sampled subset S, the paper states E_S[L(θ∗(S))] = Φ(π).Source: §3 · p. 5

3.1 Variance-Reduced Two-point Gradient Estimator and Reducible SGD

  • The two-point estimator avoids inverse-Hessian computation and unbounded score factors, while paired subsets differing by one data point reduce gradient fluctuation.The paper attributes the variance-reduction effect to the strong coupling between the two subset objective values.Source: §3.1 · pp. 5–6
  • Directly evaluating every point's marginal contribution would require O(h + Σ_i n_i) lower-level retrainings, making the estimator impractical at large data scale.Here h is the number of sources and Σ_i n_i is the total number of data points.Source: §3.1 · p. 6
  • The Bayesian reducible representation cuts the required lower-level trainings from O(h + Σ_i n_i) to O(h).The representation computes the two-point differences using posterior predictive quantities and models trained on source-excluded subsets and the full dataset.Source: §3.1 · pp. 6–7

4 Stochastic Gradient Descent via Reducible Two-point Representation

ReSGD optimizes categorical priors by projected stochastic gradient descent using the reducible two-point representation. Under weak submodularity, the paper establishes an approximation guarantee for its returned subset.

4 Stochastic Gradient Descent via Reducible Two-point Representation

  • ReSGD initializes one categorical prior per source, computes gradients from the reducible two-point representation, and updates the priors by projected gradient descent before sampling the returned subset.The algorithm also trains a holdout model on the full dataset and source-excluded models for its gradient computation.Source: §4 · p. 7

4.1 Approximation Guarantee under Weak Submodularity

  • For Bayesian V-optimal design, the induced utility is monotone and (1/c, c)-weakly submodular, with c determined by the data and model parameters.The stated c is 1 + s^2/σ^2 · λ_max(Λ), where s = max_i ||x_i||_2 and λ_max(Λ) is the largest eigenvalue of Λ.Source: §4.1 · p. 8
  • For a (γ, β)-weakly monotone submodular utility, ReSGD returns a subset with expected loss bounded by ρL(θ∗(S∗)) + (1 − ρ)L(θ_init) + ϵ.With T = O(nK_max^2/ϵ^2), the utility guarantee is ρU(θ∗(S∗)) − ϵ, where ρ = γ^2/[β + β(1 − γ) + γ^2].Source: §4.1 · p. 8

5 Experiments

Across data summarization benchmarks, ReSGD outperforms the reported baselines in test accuracy and shows robustness to corruption. In continual learning, it matches Greedy Coreset on MNIST and outperforms other baselines on the more complex datasets.

5 Experiments

  • ReSGD consistently outperforms all baselines in test accuracy across coreset sizes and shows stronger robustness to label noise and class imbalance.The paper reports similar trends on CIFAR-100; Figure 1 presents the data-summarization comparisons under corruption.Source: Figure 1 · p. 8, §5 · p. 9
  • On continual-learning benchmarks, ReSGD performs on par with Greedy Coreset on the MNIST datasets and outperforms all other baselines on Split CIFAR-10 and Split CIFAR-100.The table caption identifies Table 1 as the continual-learning benchmark comparison.Source: §5 · p. 9, Table 1 · p. 9
  • The continual-learning experiments use memory sizes of 100 for Split MNIST and Perm MNIST, and 200 for Split CIFAR-10 and Split CIFAR-100.For Split CIFAR-100, the 100 classes are divided into 10 disjoint tasks and ResNet-18 is used as the backbone.Source: §5 · p. 9
  • The paper also reports effectiveness and sensitivity analyses in appendices.The excerpt does not provide their findings.Source: §5 · p. 9

6 Conclusion

The conclusion presents a feasible-subset probabilistic relaxation, a two-point estimator, and ReSGD, with a theoretical guarantee under weak submodularity and experiments on data summarization and continual learning. The related work situates this approach among bilevel optimization, continual-learning memory selection, and weakly submodular maximization.

6 Conclusion

  • The relaxation assigns probability only to feasible subsets, while its two-point estimator avoids inverse-Hessian computation and the unbounded score of prior probabilistic relaxation.A reducible representation reduces gradient-computation cost from O(h + sum_i n_i) to O(h) retrainings.Source: §6 · p. 9
  • ReSGD's returned subset has a theoretical approximation guarantee under (γ, β)-weak submodularity, and experiments validate the method on data summarization and continual learning.The conclusion passage does not provide the approximation value.Source: §6 · p. 9

A Related Work

  • Related work frames bilevel continuous optimization through hypergradient and penalty-based methods, and reviews replay-based continual-learning coresets and weakly submodular maximization.The related-work discussion also notes that distorted local search can incur substantial overhead from repeated swaps.Source: A Related Work · p. 18

B.1 Proof of Theorem 2

These proof passages establish properties of Bayesian V-optimal utility and outline gradient-estimator derivations for the probabilistic relaxation. They also connect that criterion to total prediction error and derive a reducible representation using predictive probabilities.

B.1 Proof of Theorem 2

  • The induced Bayesian V-optimal utility is monotone and (1/c, c)-weakly submodular, with c depending on the maximum feature norm, noise variance, and prior covariance.The passage identifies s as the maximum squared feature norm and λmax(Λ) as the largest eigenvalue of Λ.Source: §B.1 · p. 21

B.2 Equivalence of the V-Optimal Criterion

  • When n is sufficiently large, the total prediction error and Bayesian V-optimal criterion differ only by the additive constant σ^2/2.The passage says the criterion's established weak submodularity therefore implies that of total prediction error to some extent.Source: §B.2 · p. 21

C.1 Policy Gradient Estimator

  • The policy-gradient estimator for the probabilistic relaxation still involves an unbounded score factor.The section contrasts policy-gradient and two-point estimators, but the supplied passage does not give the estimator formulas.Source: §C.1 · p. 22, §B.2 · p. 22

C.2 Two-point Gradient Estimator

  • The two-point estimator derivation uses independent draws and evaluates a subset with one draw excluded, as indicated by the proof passages.The supplied excerpts do not state a complete estimator formula or its full properties.Source: §C.2 · p. 23

C.3 Reducible Two-point Representation

  • The reducible representation uses predictive-distribution log probabilities, expressing their change between the full data and selected subset.The loss is specified as the negative log-likelihood of the predictive distribution.Source: §C.3 · p. 24

D Proof of Theorem 3

The proof develops a guarantee for ReSGD under weak submodularity, describes simplex projection, and establishes a reservoir-rounding scheme that returns exact per-source budgets while preserving expected utility under monotonicity.

D Proof of Theorem 3

  • The proof of Theorem 3 establishes a guarantee for ReSGD under (γ, β)-weak submodularity, using a first-order inequality for feasible subsets.The supplied passages state the condition and inequality but do not expose the full approximation expression.Source: D Proof of Theorem 3 · pp. 24, 25, 28

E Projection Calculation

  • The simplex projection either keeps the nonnegative part when its ℓ1 norm is at most one or shifts coordinates by a threshold λ* until they sum to one.The threshold equation is solved by bisection when the nonnegative part exceeds the simplex constraint.Source: E Projection Calculation · p. 29

F Effective Reservoir Rounding

  • For monotone utility, Reservoir Rounding returns exactly K_i distinct candidates from each source and has expected utility at least Ψ(π).It samples up to the prior's support size without replacement, then fills any remaining budget uniformly from unselected points.Source: F Effective Reservoir Rounding · p. 30

G Data Summarization on Cifar-100 Dataset

Figure 2 compares ReSGD with baselines on CIFAR100 data summarization across coreset sizes and distribution shifts. Its caption reports superior performance and stability across corruption and imbalance scenarios.

G Data Summarization on Cifar-100 Dataset

  • Figure 2 reports superior performance and stability for ReSGD across data corruption and distribution-shift scenarios.Shaded regions denote variance across multiple runs.Source: Figure 2 · p. 32
  • Panel (a) compares ResNet test accuracy on CIFAR100 across different coreset sizes.The caption does not provide individual accuracy values.Source: Figure 2 · p. 32
  • Panels (b)–(d) assess robustness to pairwise noise, symmetric noise, and class-imbalance factors on CIFAR100.The caption identifies these conditions but gives no panel-specific numerical outcomes.Source: Figure 2 · p. 32

H Coreset Selection for Existing Continual Learning Methods

The continual-learning evaluation compares ReSGD and CS-Rel as replacements for reservoir sampling in ER. Table 2 reports average accuracy across two benchmarks, two memory budgets, and 10 incremental tasks.

H Coreset Selection for Existing Continual Learning Methods

  • ReSGD and CS-Rel replace reservoir sampling in ER and are compared with ER, A-GEM, and FDR.Both benchmarks use 10 class-incremental tasks; all samples from the current task are used for training.Source: H Coreset Selection for Existing… · p. 32
  • Table 2 reports average accuracy on Split CIFAR-100 and Split Tiny-ImageNet with memory buffers of 200 and 500 samples.All experiments use ResNet-18, and results are averaged over 10 incremental tasks.Source: Table 2 · p. 33

I Experiment with Adversarial Initialization

This experiment tests ReSGD’s ability to recover a high-quality coreset from an adversarial probability initialization on MNIST with asymmetric label noise. Despite starting with 90% of probability mass on noisy samples, ReSGD reallocates mass toward clean samples and reaches near-perfect purity across tested coreset sizes.

I Experiment with Adversarial Initialization

  • The MNIST test uses 10,000 pristine holdout samples and a separate 10,000-sample training pool, with 1,000 training labels deterministically corrupted.The pristine holdout trains a frozen reference model used to compute reference loss during ReSGD.Source: I Experiment with Adversarial… · p. 33
  • Evaluation tracks corrupted-sample probability mass, final Top-K purity, and test accuracy after training from scratch on the selected coreset.Purity is the fraction of clean samples among the final Top-K; downstream accuracy is compared with uniform random sampling.Source: I Experiment with Adversarial… · pp. 33–34
  • ReSGD reaches near-perfect coreset purity across K values from 50 to 250 despite an initialization placing 90% of probability mass on noisy samples.Figure 3 tracks corrupted-label probability mass and coreset purity over 50 iterations.Source: I Experiment with Adversarial… · p. 33, Figure 3 · p. 34

J Data Corruption

The data-corruption experiments test summarization under label noise and class imbalance across MNIST, CIFAR-10, and CIFAR-100. ReSGD selects fewer noisy samples across noise levels, while the reported evaluations retain each baseline’s original protocol for comparison.

J Data Corruption

  • The corruption tests cover symmetric and pairwise label noise with ratios from 0.2 to 0.8, and class imbalance factors from 50 to 200.The coreset size is 200 for MNIST and CIFAR-10; Figure 4 also reports recovery behavior on CIFAR-10 under asymmetric noise.Source: J Data Corruption · p. 34, Figure 4 · p. 35
  • ReSGD selects the fewest noisy samples across noise levels compared with CSRel [104], Greedy Coreset [8], and PBCS [116] on MNIST, CIFAR-10, and CIFAR-100.Figure 5 compares selected noisy-sample counts under pairwise noise.Source: Figure 5 · p. 35

K Experimental Details

  • Results follow each baseline’s original evaluation protocol rather than imposing one unified protocol.The authors note that training epochs, data repetition, and evaluation schedules differ across methods.Source: K Experimental Details · pp. 35–36

K.1 Data Summarization Experiment Details

The appendices specify the data-summarization and continual-learning experiment settings, including a Permuted MNIST loss-factor comparison. The paper also identifies random historical-iterate selection as a practical limitation of its current analysis.

K.1 Data Summarization Experiment Details

  • Data-summarization experiments use MNIST, CIFAR-10, and CIFAR-100 with CNN or ResNet-based setups and baseline-specific training procedures.The appendix lists implementation settings for Greedy Coreset, Uniform Sampling, PBCS, CSRel, and ReSGD.Source: §K.1 · p. 36

K.2 Continual Learning Experiment Details

  • On Permuted MNIST, α = 0.15 gives the optimal balance of average accuracy and forgetting, while α = 0.2 also performs robustly.Continual-learning evaluations cover four benchmarks, with dataset-specific architectures, training settings, and loss factors.Source: §K.2 · pp. 37–38

L Limitations and Broader Impact

  • The current analysis requires sampling a historical prior uniformly over training iterations, which may require storing or recovering an iterate.The authors leave a last-iterate variant that returns a subset from the final prior as future work.Source: L Limitations and Broader Impact · p. 39

Related papers

Loading 2610.14942v1…