A Stochastic Approach to the Subset Selection Problem via Mirror Descent
Authors: Dan Greenstein
Elazar Gershuni
Ilan Ben-Bassat
Yaroslav Fyodorov
Ran Moshe
Fiana Raiber
Alex Shtoff
Oren Somekh
Nadav Hallak
Venue: ICLR 2025 Poster (2025)
Keywords: Nonconvex Optimization
Subset Selection
Stochastic
Mirror Descent
Stochastic Mirror Descent
Primary area: optimization
TL;DR: This paper proposes two stochastic formulations for the subset selection problem based on parametric distributions, derives unbiased stochastic gradient formulas, and establishes a rate result for the stochastic mirror descent in our setting.

Abstract: The subset selection problem is fundamental in machine learning and other fields of computer science.
We introduce a stochastic formulation for the minimum cost subset selection problem in a black box setting, in which only the subset metric value is available.
Subsequently, we can handle two-stage schemes, with an outer subset-selection component and an inner subset cost evaluation component. We propose formulating the subset selection problem in a stochastic manner by choosing subsets at random from a distribution whose parameters are learned. Two stochastic formulations are proposed.
The first explicitly restricts the subset's cardinality, and the second yields the desired cardinality in expectation.
The distribution is parameterized by a decision variable, which we optimize using Stochastic Mirror Descent.
Our choice of distributions yields constructive closed-form unbiased stochastic gradient formulas and convergence guarantees, including a rate with favorable dependency on the problem parameters.
Empirical evaluation of selecting a subset of layers in transfer learning complements our theoretical findings and demonstrates the potential benefits of our approach.

Published as a conference paper at ICLR 2025
A STOCHASTIC APPROACH TO THE SUBSET SELECTION PROBLEM VIA MIRROR DESCENT
Dan Greenstein
Elazar Gershuni lan Ben-Bassat Yaroslav Fyodorov
Ran Moshe
Fiana Raiber! Alex Shtoff Oren Somekh* Nadav Hallak!
I Technion Israel Institute of Technology, Faculty of Data and Decision Sciences
Technion Israel Institute of Technology, Faculty Computer Science
Yahoo Research Technology Innovation Institute Private Consultant sdngzcene www.tetw.......c........ clazargggmail. com, ilan. byw..baaaaaaaaaaaooo..ccccc com, yaroslav. fyodorov@tii ae, rani rani.moshe@gmm meaiiee@mmail.. com, flanaftechnion.ac..1 alexander whwaanderrssoffffiii.... oren www.ssssee.eii...mm ndvhllketechnion. ac.
ABSTRACT
The subset selection problem fundamental in machine learning and other fields of computer science. We introduce stochastic formulation for the minimum cost subset selection problem in a black box setting, in which only the subset metric value is available. Subsequently, we can handle two-stage schemes, with an outer subset-selection component and an inner subset cost evaluation component. We propose formulating the subset selection problem in stochastic manner by choosing subsets at random from a distribution whose parameters are learned. Two stochastic formulations are proposed. The first explicitly restricts the subset's 's cardinality, and the second yields the desired cardinality in expectation. The distribution is parameterized by decision variable, which we optimize using Stochastic Mirror Descent. Our choice of distributions yields constructive closed- form unbiased stochastic gradient formulas and convergence guarantees, including a rate with favorable dependency on the problem parameters. Empirical evaluation of selecting a subset of layers in transfer learning complements our theoretical findings and demonstrates the potential benefits of our approach.
INTRODUCTION
This paper proposes stochastic optimization approach to the Subset Selection Problem in which the goal is to choose subset of size k that attains the minimal loss defined by
min ((C), CEC
(P)
where C* = (L,,,,),, - C³K and C
set loss function. We make no assumptions regarding the loss function other than having known lower and upper bounds. In particular, we neither assume it is defined outside of C nor require it to be differentiable, unlike previous related research Ahmed et al. (2022); Pervez et al. (2022); Sander et al. (2023); Xie and Ermon (2019). As result, our formulation can accommodate loss functions that others cannot, such those arising in feature selection However, it is important to note that while our convergence results hold for arbitrary f, the properties of l significantly influence both the convergence rate and the algorithm's 's practical performance.
transfer learning via fine-tuning subset of neural layers.
Central in many fields of science, such as machine learning and theoretical computer science, the subset selection problem is NP-hard (see Bar- Yehuda and Even (1981)). In this paper, propose
Research partly conducted during Dan and Elazar's 's internships at Yahoo Research. Research conducted while Yaroslay, Ran, Fiana, Alex and Oren worked at Yahoo Research.

---

Published as a conference paper at ICLR 2025
two distributional--asse variants for (P) both optimize distribution parameters to minimize the expected subset loss. The size of the subsets in the first variant (surely). while in the second, is k in expectation. We prove that our algorithm converges to a stationary point for both formulations. To the best of our knowledge, this is the first proof that does not require the differentiability of the loss function &.
Both formulations are of the general form
inf WER'
{s0 (wr) == RC-D.. [((C))]
(1.1)
where Dw is a distribution on subsets parameterized by w, l loss function associated with the subsets, fD is the expected loss, and X' C R=.
The first formulation we consider limits the size of the subset to be exactly via the distribution Dw Vw. and is intended for cases where k = O(1). It is formulated by
inf ( (w) == Ecrrr [((C)]] WEAR
(P&)
where A := (w E An w; 2, w0 > Vi for 0 S c < 1, An is the n-dimensional unit simplex. and t'w is unordered choice without replacement according to weights Note that this a distribution over subsets of size k.
Remark 1.1 (on the closeness of the feasible set of (Px)). The feasible set in (Pk) is not closed this hindrance is addressed in the theoretical analysis of the formulation in Section 5.
The second formulation we consider allows different sizes of the selected subsets via random Bernoulli-bar mechanism where the expected subset size k. It defined by
min wEAD
{s (w) = Re-o [(()))}
(PB)
where مr := { E R" w E [c, 1c,1- E)
3 wy k) for O S c (n- ()-1, In this formulation, @w is the result of heterogeneous Bernoulli trials, where the ith element included in the subset if the jth trial is successful, and the individual element inclusion probabilities sum up to k.
The choice between fixed (Px) and expected (PB) cardinality depends on whether cardinality is a strict requirement (e.g., Subset Sum) or guideline (c.g., feature selection). While expected cardinality can approximate fixed cardinality by penalizing subsets of undesired sizes, this may hinder convergence due to frequent low-quality samples.
The detailed distributions are given in Section 4, and the purpose of the approximation parameter will become evident in Section 5. In both formulations, (Px) and (Pn), selecting subset boils down to determining a distribution parameterized by the decision variable w. To determine w, we utilize a nonconvex Stochastic Mirror-Descent (SMD) Zhang and He (2018) based method detailed in Algorithm 1. Accordingly, we measure the first-order optimality of solution w using the standard Bregman stationarity measure Zhang and He (2018); we elaborate such preliminaries in Section 2.
Contribution. We briefly summarize the main contributions this work:
We propose two continuous stochastic formulations for the Subset Selection Problem with possibly discrete loss, and prove that they can encode the optimal selection. Moreover, by choosing appropriate distributions for these formulations, we derive constructive gradient estimators, and the relatively weak convexity of the expected loss functions.
We prove convergence rate via an SMD-based method, without any assumption the underlying loss. We are not aware of such proof even for differentiable loss functions.
We derive concentration bound for the loss of the subsets sampled using our method.
We demonstrate the efficacy the proposed algorithm in a deep transfer learning setting.
Related Work. Recent research stochastic subset selection emphasizes its integration into computation graph (CG) frameworks, widely used in modern machine learning applications Sander al. (2023); Xie and Ermon (2019); Jang al. (2016); Maddison et al. (2016); Ahmed al.
2

---

Published as a conference paper at ICLR 2025
(2022); Pervez et al. (2022). These methods aim to compute gradient estimators for subset selection parameters by encoding subsets as k-hot vectors vectors-binary binary vectors with & non-zero entries representing the chosen subsets. While CG frameworks assume the loss is differentiable almost everywhere, address discrete loss functions where this assumption fails, as with metrics like classification accuracy Spearman correlation.
Sander et al. (2023) provide theoretical guarantees for their gradient estimator, introducing a general- ized top-k operator with a subgradient and relaxed version with an exact gradient. Their approach, however, does not apply to non-differentiable losses.
The Softma trick, relaxation of top-k sampling Xie and Ermon (2019), is suited for differentiable losses. In our setting, it has two limitations: (i) defining loss for continuous relaxations of k-hot vectors is challenging for tasks like feature selection, and (ii) the relaxation depends on temperature parameter t. While the relaxed distribution converges to weighted size-k subset sampling as + 0, the distributions differ for any fixed >
Other notable approaches for stochastic subset sampling in CG settings include those by Ahmed et al. (2022) and Pervez al. (2022). Ahmed et al. (2022) investigate a Poisson-Binon -inomial--ike distribution conditioned on selecting exactly & elements. They propose dynamic programming algorithms to compute marginal selection probabilities and enable sampling, along with a theoretically intuitive gradient estimator based on these probabilities, though lacks rigorous proof.
Pervez et al. (2022) propose a sampling mechanism that reduces subset size variance in the Poisson- Binomial distribution. Using Bernoulli random variables, they generate subsets near size k through sequential trials and apply the Throgghh gradient estimator Bengio et al. (2013), though without guarantees on bias or variance. While their method is primarily relevant to work from practical perspective, combining it with ours could yield synergistic and innovative approaches, as discussed in Appendix F.
Outside the CG framework, the REINFORCE gradient estimator Williams (1992) does not require differentiable losses. While provably unbiased, it has high variance, and we are unaware of theo- retical bounds on its variance or moments, which are crucial for first-order stochastic optimization convergence.
Our work intersects with stochastic smoothing techniques, smoothing for set-valued functions, and variational optimization. For foundational works, see Lovász (1983) for deterministic smoothing of set-valued functions, Nesterov (2005) for smoothing in nonsmooth continuous functions, Duchi et al. (2012) for stochastic smoothing, and Staines and Barber (2012) for variational optimization.
Finally, Combinatorial Multi-Armed Bandits (CMAB), in the stochastic Agarwal and Aggarwal (2018); Rejwan and Mansour (2020) and adversarial Han et al.; Audibert al. (2014) settings, are somewhat related to our research. Similar to CMAB, we select a subset of size k and receive subsequent reward. However, in CMAB, it is assumed that each element has a non-constant intrinsic score, and the total reward is a function of the selected elements, which can be linear Audibert et al. (2014); Rejwan and Mansour (2020), monotonically increasing in the clement scores Agarwal and Aggarwal (2018), elements with better scores is preferable. In our setting, there are no intrinsic element scores, two similar subsets have vastly different rewards, and the reward for given subset is constant.
known in advance Han et al.. The underlying assumption is that selecting
2 PRELIMINARIES
Relative weak convexity Zhang and He (2018); Davis and Drusvyatskiy (2019) will be used to obtain the Mirror Descent method's 's desired guarantees.
Definition 2.1 (Relative Weak Convexity). Let f X' - R. We say that ff is (p, p) relatively weakly convex (RWC) if: (i) μ differentiable and 1-strongly convex over X;; (ii) The function of + pp is convex over X.
The following definitions are taken from (Zhang and He, 2018, Section 2.3).
Definition 2.2 (Bregman Divergence). Let X + R be differentiable and convex. Then for x, y E X. the Bregman Divergence is defined
B₁ (x., y) 1a (a:) (y) (Vp (v). -y)
3

---

Published as a conference paper at ICLR 2025
Definition 2.3 (Bregman Proximal Operator). Let X - R differentiable, and X R. differentiable and strongly convex, such that f (p, P)) RWC. Let < p Then the Bregman proximal operator is defined as
xrox (z) == argmin f(w) + - B,,,,,)
wEX
The Bregman Proximal Operator is well-defined and unique for every E X.
Lemma 2.1 ((Zhang and He, 2018, Lemma 2.2)). Suppose a function f is (p. R on X and 0 À < p=1. Then for any input X, the function + Ba(;, z) is (X-1 p)-strongly convex. Moreover, the Bregman proximal operator prox (z)) is unique.
The Bregman proximal operator is necessary optimality condition, and natural when discussing Mirror Descent based algorithms. Since the Bregman proximal operator is unique, the following notion of Bregman Stationarity Measure is well defined.
Definition 2.4 (Bregman Stationarity Measure). Let [, defined as in Definition 2.3. Then,
Ax(w) x
x (B, prox,, + ( w))
When z Lipschitz continuous, measure and the more traditional measure of distance between the subdifferential set and zero (see (Zhang and He, 2018, Equation 2.12)); We provide detailed proof for completeness.
derive connection between the Bregman stationarity
Lemma 2.2 (optimality measure). Let ú prox (un) and assume that p has M Lipschitz gradient. Then, (dist (0, a (ʃ + 6x) (wb)))² < MAx(w), where dist (w,S)) inf || yl| yES
ALGORITHM
To obtain the desired distribution in two stochastic formulations (Px) and (Px)* we propose nonconvex Stochastic Mirror Descent (SMD) -based method detailed Algorithm 1.
For both (Px) and (Px) the algorithm provide: (1) Maintains a parametric distribution of subsets Dw (Dw = the or Dw = Φw ); (2) Samples a subset C' ~ Dw: (3) Computes stochastic gradient estimator (C') given the subset C; and (4) Performs a SMD update using the gradient estimator.
The main challenges in implementing distribution. In particular, the distribution must induce an unbiased stochastic gradient estimator with explicit formula that be bounded. Moreover, the expected loss incurred by the distribution must be Relatively Weakly Convex (cf. Section 2). As we shall demonstrate in the sequel, the distributions propose in Section overcome these challenges.
SMD approach strongly linked to the choice of the
Algorithm 1: Stochastic Subset Learner
Input: c >>>>>>>0000 0, X C R", wo E X' p: + R, (ar)Tm
for ( (= 1........ 1..., do
sample C Dun,
evaluate VD (C) // Evaluate the gradient estimator
Wi++1 argmin V., (C), w) + B,(w, (w) // Mirror Descent step WEX
a
end
sample R (1, (1,..., T), such that TP (R at/ Σ a
Return wR
"The properties of the formulations and their approximations discussed Section and Section 5.
4

---

Published as a conference paper at ICLR 2025
Remark 3.1 (inputs of Algorithm 1). Applied setting and based the requirements of the SMD, we assume that the inputs Algorithm obey the following: X is either A, or A and is differentiable and 1-strongly convex. The step size a is discussed in Theorem 7.1 and Theorem 7.2.
In Theorem 3.1, we state the convergence guarantees for Algorithm informally and concisely. The full version of this thcorem is established Section 7.
Theorem 3.1 (Algorithm guarantees (informal)). xn the outpu of Algorithm for one of the supported formulations, with an appropriately chosen step size. Then there exists > O such that
122.5
E [1/222 (wn) SO
VT
where the expectation taken the random choices of the algorithm.
We are not aware of any other convergence result in our setting. Even in CG setting, where the loss function is assumed to be differentiable rather than discrete, we are not aware of any convergence results to stationary points. We note that many applications might find the best sampled subset to be of interest, and indeed, our experiments Section 8 use this measure.
There are several heuristics which improve the performance of Algorithm 1. We refer the interested reader to Appendix F.
4 PROBABILITY DISTRIBUTIONS OVER SUBSETS
In both formulations, (P:) and (Pn), selecting the best subset involves determining distribution using the decision variable w, which we optimize with Algorithm 1. This section presents the proposed probability distributions for these formulations.
For (Px), we define Vhe an unordered weighted-choice without replacement distribution, ensuring subsets of size k. To support this, we introduce some notation.
For subset C' E C* let कं :.... "K. denote all permutations of C. We simplify to m1..... Fh when C' is clear and omit indices when arbitrary. The 7th element of a permutation æ is æ [T]. or simply [1] when the permutation is clear.
CI,..., CN (N = (+)) ) be some enumeration subsets in C*. The distribution the defined
by the probability vector nd for a weights vector w e A constructed via the probability of choosing the permutation xCr given by
p. = PT C. 2
W
(4.1)
w
The probability of choosing the subset C₁ the sum of the probabilities of choosing its ordered tuples
k!
k!
p == P(C; w) = Σ P(
(w) Σ P
(1
For every choice of i, kc n. ग and w, computing P. is possible in O (k) operations. Given that there are k! possible permutations, computing p takes O (k k!) = O ((k + 1)!) operations. Therefore, computing this probability is feasible only for relatively small values of On the other hand, note that the probability computations are completely independent of n.
We now move define the distribution &w for the problem (Pg). Let Ci...., Cn be some enumer- ation of the subsets in C. Given weights vector w E A n.k. k the distribution &n is defined by the probability function
PCC, w)
II w II (1-w;). JEC
(6ww
jæC
For every n, k. and weights vector w, calculating o takes 0 (n) operations.
5

---

Published as a conference paper at ICLR 2025
5 RELATIONS BETWEEN THE STOCHASTIC FORMULATIONS AND THEIR APPROXIMATIONS
This section justifies stochastic formulations by connecting the solution sets of (P), (Px), (P and their approximations. These approximations provide bounds on the gradient estimator's 's moments and the Hessian cigenvalues, essential for the convergence result in Section 7, with proofs in Appendix B.
While it is tempting to use (Px) and (P directly, the resulting gradients and gradient estimators can have arbitrarily large norm. Therefore, tackle formulations with relaxed upper/lower -bounds on the decision variables instead, that is, we minimize over the sets A and A for e > 0. The c-approximated problems are:
VK min BC-i' [e(C)], WEAr
(PP)
VB ;= min Bena [e(C'].
(PP)
wEA.
Theorem 5.1 provides a bound on the effect of the constraint relaxation in (PE) on the optimal value.
Theorem 5.1 (approximation gap in k-cardinality subsets). Denote I.mar = max ((C), L
CECk
min (CC), and let O Limax L be some suboptimalliyy gap. CEe
The minimum element weight is defined via an auxiliary variable ē c(n k) and à
II (1 1) (k
-))) Define 7t as
( kd (1 (1(1-0))) (Imas 12) -
C = argmin 05
-1'-
Then for every constraint relaxation of 0 < c c² (n-k)" k) it holds that Vg S La +T. Note that the upper bound of e depends
Theorem 5.2 (approximation gap in k-cardinality expected value subsets). Let Limaz inf EEsple(Css)) and let O 5 Lmax L' be some PEAN
max ((C), L min {(C)
CC(1.2.....11
CEC
suboptimality gap. Let and 1 be the lower and upper bound of the element weights, respectively. We define e² as
c² (n k)- (1
1,²)-1)-
Then for any constraints relaxation 0 < < c* it holds that Vg < L +T.
Remark 5.1. Note that in both Theorem 5.1 and Theorem 5.2, the upper bound c* depends on Moreover, e* is strictly monotonically increasing for (0, Liaas - L*).
The elements weight bound enables to fine-tune the tradcoff between the suboptimality gap of the approximate problems, and the gradient estimator norm bounds.
THE UNBIASED STOCHASTIC GRADIENT
now turn to obtain our gradient estimator. is derived by analytically calculating the derivative of our proposed distributions. The constructive derivation of the gradient estimator allows us to bound the estimator's size and variance.
We note that the CG framework literature offers several gradient estimators, all of which assume that the loss function is differentiable. To name a few: Sander et al. (2023) introduce an unbiased gradient estimator, Xie and Ermon (2019) present Gumbel Softmxx gradient estimator, and examples of heuristic gradients that are not related to the Gumbel Softmax trick can be found in Ahmed et al. (2022); Pervez et al. (2022). None of the aforementioned gradient estimators provides any theoretical bounds for the size or variance of the gradient estimator.
The stochastic gradient formulas for (Px) and (P defined next, and then proved to be unbiased.

---

Published as a conference paper at ICLR 2025
Definition 6.1 (stochastic gradients).
Define the stochastic gradient in the following manner:
Dw = duu Then for C ~ the set
1 Opt pi dr
V
(C:)
((C.),
Vg E f1,2....,,A
(6.1)
The partial derivative is given by
Op
k! Opp
(6.2)
ow
Ow
where TI,
Thk are all possible permutations C; and
j-1
Op..
(m)))
+
Σ
Σ
lil
P
q C [m]]]] =: [m]
(6.3)
I jmm++
0,
otherwise.
Dw = ow. Then for Ow set
w (0())(()) [0) -1 (1-ww, E(C).
jEC otherwise.
(6.4)
The fact that both of the gradient estimators are well-defined and unbiased is established next; it is proved separately for each estimator in Appendix C.
Lemma 6.1 (gradient estimators properties). Suppose that Du- she or Dw Qw- Then the gradient estimators (6.1) defined in Definition 6.1 and (6.4) defined in Definition 6. respectively are well-defined and unbiased, that is, BCnDw
4 (C) ViC-D= [((C)), where D (wws Dw).
We conclude this section with an informal remark regarding the way in which the properties of affect the convergence rate of Algorithm empirically.
Remark 6.1. Definition 6.1 suggests that the subset loss E(C) for any sampled subset C E C influences the sampling probabilities of all subsets that share elements with C' in subsequent rounds, implicitly assuming that ((C) provides insight into the losses of subsets overlapping with C. Experi- mental results in Appendix G indeed indicate that when this assumption holds more strongly, the algorithm performs better.
7 CONVERGENCE RESULTS
The SMD convergence analysis using the framework proposed in Zhang and He (2018) suggests that the main challenges in deriving the convergence results lie establishing the RWC of the objective functions and in obtaining (stochastic) bound expression for the gradient estimators in Section 6.
The following lemma provides bound for the gradient and the gradient estimator.
Lemma 7.1 (gradient estimator bounds). Let 4. Then for every C C* it holds that:
1. th ( < k.s. -I((C)) "ijn wi
-2 (min w,
2. 0 maz (V [e(C)) n (k2 + k +1+ max max| [((.))
We now establish bounds utilized when proving the RWC in (PE).
Lemma 7.2, Let up E A Then,
| <
1 Wi
1. Ee~or
5
max ((C)2
w.
CEC
7

---

Published as a conference paper at ICLR 2025
1
2. Let pB(w) (n
max max my
max
[((C))
Wm w (1 Wm) w,) Wm
w+)
CEC
Then Omar (CR (e(C]) (p()))
The RWC of the objectives of the approximations is stated next; it is a corollary of the bounds proved separately for each model in Appendix D.
Corollary 7.1 (RWC). For W E (47,kr Ar) define p W - IR to be p(m)) 1
/0ol -u(w)
Σ ln;w;), or p(w) = Σ In(w;), Then EC-D [(C)) is (PB: H)-RWCC where
1. D = Own W = A k.k. c > defined in Theorem 5.2, PB = (n 1 1) ② Oel
max (e(C));
2. D Vw. W = 2 0 defined in Theorem 5.1, P = 2 (A²2 + k 1) max; I((C.))
Finally, by utilizing the the gradient estimator's bounds derived above and the RWC property in the problems, can conclude with the rate guarantees for (Pş) and (P%). While the rate depends on both n and k, the dependency is disjoint, opposed the (?)) In particular, for small values of k, the rate is of order O (n2. /VTT as opposed to O (n²))
Theorem 7.1 ((P)) rate result). Let wr be the output of Algorithm I for (PE) with at = (n VT) -1 and let p n- (k² + k- - 11) max |e(C)), where c is defined as in (PE). Let G == max ((C) CEe
CEEK
min ((C) and M == max 1((C)). Then. CEC
CEC
A A1/2) S (12-5 1. (122 (k2+ -12 2n.nnnnn .111555)) /VT//. E
In particular, for ē= c² and ct in Theorem 5.1. It holds that
n n
( s ((G + (M²) n2 /25//T).
where è does not depend on n.
The rate result for (P%) independent k, is is therefore suitable for large values of
Theorem 7.2 ((P)) rate result). Let w be the output of Algorithm for (Pb) with ar = (n2.5/T) -1 Suppose that p
max [Ie(C'] C C} where is defined in Theorem 5.2. Let G max ((C) min {CC) and M == max |((C)), Then, Cec
CEC
CEC
2.5
2M
E [w)))[WRR] <
+
VT VnvT-C.
In particular, setting e² where c* is defined in Theorem 5.2. It holds that
E (1)) s 2 2G + M38--3 nnn))n222)) 2.5
where è does not depend on n or
We conclude this section with high probability bound for the value of the sampled subset. Theorem 7.3. Let n be the distribution parameter returned from Algorithm applied either (PP) or (PE), and let Dwn Ving or Dwg Øung R accordingly. Let Ci...., Cm ~ Dwug be independent samples. Denote max ((C) and min E(C). Then for every 8 0.
CEC
CEC
m8²
P min E(C,) - EEwDwo [E(C)] 28 S exp iE)
2 (u - 2)2
8

---

Published as a conference paper at ICLR 2025
Layer Suboet Choloe
Cridmn Model
Layarr to LOne
Aoconding layer waights
fully connected layer
modo
porformance helld out d60952t
proportional 20 parformano
Uporate weigghts
host 127m
Figure 1: Layer Subset Selection for Transfer Learning Diagram.
EXPERIMENTS
We conducted experiments in two setups: (i) Transfer Learning (TL) to showcase practical applica- tion, and (ii) Synthetic experiments for direct evaluation subset selection tasks. TL. experiments are discussed below, with synthetic experiment details in Appendix G.
Experimental Setup. A natural application of our model Transfer Learning (TL), where knowledge from one task improves performance on a related task Tan et al. (2018); Zhuang et al. (2020). In our setup, the subset selection algorithm selects layers from frozen pre-trained neural network. At each iteration, a subset of l layers is sampled, connected to trainable fully connected layer, and the resulting model trained and evaluated on a held-out dataset. The algorithm computes loss based on this performance to update selection probabilities. After a preset number of iterations, the process outputs the best subset observed, as shown in Figure 1.
Concretely, we experiment on TL in vision classification tasks based on two models, Vision former (ViT) Dosovitskiy et al. (2020) and ResNet18 He et al. (2016), and two datasets, CIFARIO IO Krizhevsky (2009) and SVHN Netzer et al. (2011). The model weights are pre-trained on Ima- geNetl Deng et al. (2009). The experiments are restricted to a setting in which the target dataset is small so using frozen model is preferable to partial or full fine-tunning (see Plested and Gedeon (2022)). Three layer-choosing methods act as benchmarks, the first two we use hyperparameter tnning (HPO) the validation set tune the learning rate and batch size, and in the third HPO from the second method. The methods are:
Lastl Using the last layer with learning rate [0.001, 0.1] and batch size in (32, 64, 128).
Last4 Using the last four layers with the same learning rate and batch size ranges.
URand4 Sampling four layers uniformly at random for 50 iterations, training the resulting model on the training set, evaluating on the calibration set, and then selecting the one with the highest accuracy on the calibration set.
We propose the following layer selection methods based on Algorithm 1: For any set of Algorithm 1 hyperparameters (7 and a, where a for any t), we sample layers according to Dw- For each layer-sample, we construct a model, and train it on the training set with the learning rate and batch size selected by HPO for Last4. The model then evaluated on the calibration set, and we set the loss in Algorithm 1 to calibration accuracy. We update the layer weights and repeat for T iterations according to Algorithm (I' 50). Upon termination of Algorithm 1, the best layers-subset is evaluated on the validation set to update the HPO process for Algorithm 1. Once HPO is finished, we repeat Algorithm for the chosen values r, a, and evaluate the best in terms of calibration accuracy on the test set. We set ((w)= = | and use the following to sample layers at each iteration:
CWR4 Sampling four layers with weighted Choice Without Replacement (#w)).
HIB4 Sampling layers with Heterogeneous Independent Bernoulli experiments (dw). Due to memory considerations, samples with more than five layers assigned accuracy zero.
The performance of each experiment is evaluated for the best set of hyperparameters. Additionally, the three sampling experiments are evaluated on calibration accuracy. For the best set of hyperparameters according to the evaluation accuracy, show the best observed calibration accuracy per iteration.
The experiments are carried out on AWS Sagemaker, with the instance types "ml.g4dn. "ml.g44dn.16xlarge" with 64 vCPU, Nvidia t4 tensor core GPU, and Intel Xeon Family physical processor.

---

Published as a conference paper at ICLR 2025
ViT-CIFARIO
ViT-SVHN
ResNet18-
ResNet
Mean Ace %
Acc Std
Mean Acc %
Acc Std
Mean Acc %
Acc Std
Mean Acc
Acc Std
Lastl
91.654
0.365
50.0369
0.971
66.698
1.134
34.700
1.345
Last4
86.414
0.551
72.523
2.316
70.794
0.836
49.975
0.689
URand4
89.362
1.085
75.053
2.553
74.052
1.797
53.317
12.577
CWR4
90.692
0.559
75.827
3.611
71.554
2.296
60.867
4.962
HIB4
90.724
1.045
77.781
2.036
71.336
4.102
60.478
60.478
Table Accuracy (Acc) averaged over the subsamples on the original test set. Bold indicates the best.
SVHN With ResNet18
Srst Olservee Accune))
Bust Olsurve AcOuRiC)
0.70
0.85
0.3
0.55
0.50
URand1
URandi
CWN4
CNNA
HIB4
Figure 2: Calibration accuracy per iteration for representative subsamples.
Experimental Results. The results in terms of test accuracy are displayed in Table 1. In all cases, selecting the last four layers deterministically (Last4) is inferior to the layer-selection algorithms URand4, CWR4, and HIB4. Among these algorithms, both methods outperform the random greedy method URand4 in three out of four cases.
The comparison to Lastl is of interest only in the context of TL, this is the simplest way to use a frozen pre-trained model. In this context, Lastl is the superior method in single case (VIT with CIFARIO) 10) and inferior to all other methods in all other cases. In the general context of subset selection, the Lastl experiment holds little relevance our experiments aim to compare methods of selecting subsets of size 4, either deterministically or in expectation.
In broader sense, our results support the superiority of layer selection algorithms over a deterministic selection in TL on frozen models. This conclusion is compounded by the use of learning rate and batch size selected for Last4, instead of optimizing them directly. We speculate that these results can be extended to more general settings, where the choice of layers can take different meaning. One example is using a decaying learning rate between layers, as suggested in Plested and Gedeon (2022).
We found that in the calibration accuracy comparisons, one of our methods outperforms the random greedy approach in 14 out of 20 experiments (2 models x 2 datasets x 5 subsamples). Actual numbers are omitted as they provide no further insights. The effect of layer selection algorithms is much more pronounced in ResNet18 8 than in ViT representative example for each model is provided in Figure 2. Our experiments also show that CWR4 tends to reach its plateau performance more quickly. We believe this is because its search space is smaller than that of HIB4, and unlike URand4, the initial iterations of CWR4 actively encourage exploration.
CONCLUSIONS
We introduced a stochastic approach to non-continuous subset selection formulation facilitating novel theoretical guarantees. Our results motivate both future theoretical and experimental investi- gations. In particular, further theoretical study of the stochastic model approach, formulation, and possible distributions, well other optimization frameworks. Numerical experimentation TL and other subset selection applications such as Vertex Cover or Independent Set, especially those with challenging non-differentiable metrics, also provide intriguing research prospects.
10

---

Published as a conference paper at ICLR 2025
10 ACKNOWLEDGEMENTS
Dan Greenstein's 's work was partially supported by the ISRAEL SCIENCE FOUNDATION (grant No. 637/21). Nadav Hallak's work was supported by the ISRAEL SCIENCE FOUNDATION (grant No. 637/21).
REFERENCES
Mridul Agarwal and Vaneet Aggarwal. Regret bounds for stochastic combinatorial multi-armed bandits with linear space complexity. arXiv preprint arXiv:1811.11991 2018.
Kareem Ahmed, Zhe Zeng. Mathias Niepert, and Guy Van den Broeck. Simple: A gradient estimator for k-subset sampling. arXiv preprint arXiv:2210.01941, 2022.
Jean- Yves Audibert, Sébastien Bubeck, and Gábor Lugosi. Regret in online combinatorial optimiza- tion. Mathematics of Operations Research, 2014.
R Bar- Yar-dehudaaaa and S Even. A linear-time approximation algorithm for the weighted vertex cover problem. Journal of Algorithms, 1981.
Amir Beck. First-order methods in optimization. SIAM, 2017.
Amir Beck and Yonina C Eldar. Sparsity constrained nonlinear optimization: Optimality conditions and algorithms. SIAM Journal on Optimization, 2013.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 1eXiv:1308.00 www.I11388333333 34322, 2013.
Damek Davis and Dmitriy Drusvyatskiy. Stochastic model-based minimization of weakly convex functions. SIAM Journal on Optimization, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Conference on Computer Vision and Pattern Recognition, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 1929, 2020.
John C Duchi, Peter L Bartlett, and Martin J Wainwright. Randomized smoothing for stochastic optimization. SIAM Journal on Optimization, 2012.
Yanjun Han, Yining Wang, and Xi Chen. Adversarial combinatorial bandits with general non-linear reward functions. International Conference on Machine Learning.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition, 2016.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax, arXiv preprint urXiv11611.01144 2016.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
L. Lovász. Submodular functions and convexity. In Achim Bachem, Bernhard Korte, and Martin Grötschel, editors. Mathematical Programming The State of the Art: Bonn 1982, pages 235-257. Springer Berlin Heidelberg, 1983.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.000 1007121 2016.
Yu Nesterov. Smooth minimization of non-smooth functions. Mathematical programming, 2005.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng, al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature learning, 2011.
11

---

Published as a conference paper at ICLR 2025
Adeel Pervez, Phillip Lippe. and Efstratios Gavves. Scalable subset sampling with neural conditional poisson networks. In International Conference on Learning Representations, 2022.
Jo Plested and Tom Gedeon. Deep transfer learning for image classification: a survey, ar preprint DxXiv:225500900 2022.
Idan Rejwan and Yishay Mansour. Top-k combinatorial bandits with full-bandit feedback. In International Conference on Algorithmic Learning Theory. 2020.
Michael Eli Sander, Joan Puigcerver, Josip Djolonga, Gabriel Peyré, and Mathieu Blondel. Fast, differentiable and sparse top-k: convex analysis perspective. In International Conference Machine Learning, 2023.
Joe Staines and David Barber. Variational optimization. arXiv preprint 121 2012.
Stefan Szeider. The parameterized complexity of k-tlip local search for sat and max sat. Discrete Optimization, February 2011.
Chuanqi Tan, Fuchun Sun, Tao Kong, Wenchang Zhang, Chao Yang, and Chunfang Liu. A survey on deep transfer learning. In International Conference on Artificial Neural Networks, 2018.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 1992.
Sang Michael Xie and Stefano Ermon. Reparameterizable subset sampling via continuous relaxations. arXiv preprint arXiv: 1901 1et...9911...., 10517, 2019.
Siqi Zhang and Niao He. On the convergence rate of stochastic mirror descent for nonsmooth nonconvex optimization, arXiv preprint arXiv:1806.04788 2018.
Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey on transfer learning. Institute of Electrical and Electronics Engineers, 2020.
A PROOFS OF SECTION 2
Proof. By the optimality conditions of
1 (w) argmin f(z) B,, w) argmin f(z) 1 (m(z) m(w) - (Vu(z), -w)).
1
X
2eX
zex
we have
1
0 (EV()) + (1)) + (Vp(z)) vu(w)) x (Vu(w) - Vu(()) a 6x) (10).
Hence,
1
(dist (0, 8 (ʃ + 6x) (1)))²² < x |Vb(t) - vp(w)]l²
By (Beck, 2017, Theorem 5.8, part 4), since μ is assumed to have M Lipschitz gradient over X,,
M
Vu(10)
Vp(ww) (1)2))1 <
(Vp(w) - Vu(z)), w w)
xZ
x
Since
A Pz (V(ww) vu(w)), 10)
this concludes the proof.
12

---

Published as a conference paper at ICLR 2025
B PROOFS OF SECTION 5
The following Lemma formally establishes that the minimum of (P) is equal to the infimum of (Px),, and is lower bounded by the minimum of (P
N
Lemma B.1 (basic relations between formulations). Denote L := inf [Smp]e(C.)]
Σ Ps
pEAN
i=1 R(C.), and let L'g be the optimal value of (P') and L be the optimal value of (Px)) Suppose that C1....: On is some enameration C Then:
1. L min ((C+)) S l i=1....N
2. There exists convergent sequence {wili=) of feasible solutions to (Pa) such that Ec [e(C)] 1-0,, L
3. LB S Li
Proof.
1. Consider the optimization problem
inf E [((;)) PEAN
Σ P ((C.).
Set
1, if
(.
otherwise
where
E argmin E(C.).
We note that
N
N
Σ P
((C)) 2
Σ
P
min EI.....N
((C)) min ((C.),
i=l....N
1=1
1=1
where the second equality follows from E AN. Additionally, for p",
Σ pt ((C,)) = Ds. ((C+)) min ((C)). i=1....N
The relation L < LŹ follows immediately from the fact that for every w E A, 0 R4++ D/ E AN, and that the target functions are identical for fixed distributions.
2. Let i* E argmin (CC.). In part of Lemma B.1 was shown that
p" E argmin PEAN
where p" is given by
if if
p
0.
otherwise.
We construct convergent sequence { ) such that
T 1 p²
Since Exp ((C:)] is continuous in p and argmin Expp [(():)] such construction is sufficient for the proof.
PEAN
13

---

Published as a conference paper at ICLR 2025
Let
j E C.-
k
2k.1'
(wi))
1
otherwise
2 2 k)-
Clearly, for every I,
(wi) Σ
(wi);
Σ (wi))
jEC,
jdC,
1
1 2(n () k:) k))
k:
+ (n - k)
1,
K
2k
and for every and] j
(wi) > 0.
Therefore, wi E An RR++. Furthermore, wi 1.00 w", where w" is given by
if E C-- otherwise
w]
k''
0,
Next, we calculate p.a
1/k 1/2kl
m P--
1)
-12kk/
2kl - (j--jj--y 1)
j=1
21 -1
ΣI
2 -
Clearly,
21 1
10000
2(k - + 1)1 +
k: k-j+
Therefore it follows that
II j=1
21
1-000
k
2
2
1|
k - j +
k!
j=1
and
A k
EI
21
1+00
2(k j 1)1+ÿ-1
1000 and n A n. it follows that for every # 1*
Since p.
P mi
1+000
Hence,
10000 p". n
Since Esppr [e(C.)) = Ec [((C)] is continuous in nd, immediately follows that
Eawirr [((C))
1+000
Einp [((C,)) L..
14

---

Published as a conference paper at ICLR 2025
3. Let Ci...., CN be some enumeration of C*, and let i" argmin {(C)).
Define vp as
; = 11
JE Ci-
0. otherwise
Then w E {w RT Σ w, k} and
1 P-- = 1.
Therefore, is feasible solution to (PB) whose value is L.g Hence,
LB S Li
Proof of Theorem 5.1. Let Ci...., Cn be some enumeration of C and let E argmin ((C.).
Define
on k)), jEC C.C.
C,
otherwise
and note that An Nz Z
for S c S (n k)
The probability of choosing the optimal subset Ct- according to the distribution Dw is
k! (1 (x (1 c(n k:))
P1
(j 1) (
{ (1 c k))
k
Denoting é (n k).
k!! (-(1 (1 2) )
k
1 1- 1- 1)
(+ 1 (1- 2)
II
Note that this is strictly monotonic decreasing function 2 0 S ē 1, whose value is for ē = and 0 for ē = 1. Therefore, 1 D. is strictly monotonic increasing in č. Since 0 < 7 < Lmax L*, follows that T_max 1. E (0, 1). and therefore there exists single solution č* [0, to
r
k
k! (1 -1-
k
(Lmax L²) = 0.
k
(j - 1) (#) --)))
1=1
Hence, the same ō minimizes
k! 1 (1 (1-2))
(Lman --)- -
1) ( (1 2)
We conclude that for every << n1
15

---

Published as a conference paper at ICLR 2025
min wEA,N
Bc [((C)] - p.r + (1 Lmar S L' + T.
Proof of Theorem 5.2. Note that c² is the only root in [0, (n k)- ] of
(1- (1)))((n-k)) k)))) (Lmax L") 0.
Let CI...., ON be some enumeration of subsets of C. and let i* E argmin (C,).
Define w as
c(n k)
j E Cj.
h
C,
otherwise
and e. (n k).
We can see that
1
c(n pan-hi k)
nnk (1 -1- -(1-c)=-
-1-m
(1
n-
k
k
C.
m
The sequence Um
1
is increasing, and therefore for every k. n such that <
m
n-k
2 2--
n k
and
(1- k
2 - 2,
and at least one of the inequalities is strict.
Therefore,
D (1 21= = (1 c(n -k))
Note that A n follows that,
EcNdr (1- (1-1++11111111111
(1 (1-))) Lmaz.
min MEAC
For < e
min RC-ae [((C)] < L weAs,
C PROOFS OF SECTION
The properties of the gradient estimators for our two formulations summarized in Lemma 6.1 are proved separately in the following two lemmas.
Lemma C.1 establishes that the stochastic gradient (6.1) well-defined and unbiased.
Lemma C.1 (gradient estimator for (Px) Let Cı,... CN be some enumeration of C* and suppose
that Dan = v/a- Then the gradient estimator (6.1) defined in Definition 6.1 is well-defined and unbiased, that is,
( (C) = VEC [e(C)).
16

---

Published as a conference paper at ICLR 2025
Proof. Let C C Consider element index q C and denote by some permutation of the elements of Ct. Since q E C2 there exists an index m, such that π [m = [m]
For j m,
a p10mm
[
0.
1
For j m,
i
W m
m-1
Σ
1
Σ
1=1
For j > m,
a
w (
3-1
j-1 Σ lil
lill
Recall the derivative rule for product (assuming fi(w). ..... fa( + 0)
[ [[)
k
af((w)
k
k (f((w)
II
fy(
II (w)
(ml
mm
f((w)
j=1
Combining all the above,
Σ
Öp:. @10mm
1=1
w [m]
Σ
1=1
Σ
w
w
(1
10
jmm+t
1
Σ Wi
(iil
Equivalently,
"p
Σ
w J
Ow
[mm
1-1
j-1
j=m+1
1
Σ
li
Σ
[
1=1
Since
w
80
P
j=1
wi
(6.3) for C, follows.
On the other hand, if q & C₁, then p. is not a function of wg, and therefore (6.3) for & C, holds.
3p. w om,
The result that the partial derivative
given by
Op w
k!
Fw.
Ow
r=1
17

---

Published as a conference paper at ICLR 2025
then follows immediately from the fact that p. and that (6.1) well-defined.
Σ
P which verifies the correctness of (6.2)
To prove that (6.1) is unbiased, note that the gradient at index q is given by
Σ Eppe(ct))
0yp l(C)). Ero
V (Eorpw [((C)),
dw
On the other hand,
Ec
V [C)]
Σ i=1
Pi
V (Ci)
Σ
Pi
1 0p
((C;)
Σ
N Op: ((C.).
p Or
Ow,
1=1
Hence,
BC
A. (C)
V (Bowew [((C))).
The following lemma establishes that the stochastic estimator of the gradient is unbiased for (PB).
Lemma C.2 (gradient estimator for (PB)). Suppose that D = tre- Then the gradient estimator (6.4) defined in Definition 6.1 is unbiased, that is,
Ecnbo V (C( VEC [((C)].
Proof of Lemma C.2. Let Ci...., O2n E C be an enumeration of all possible subsets of {1...., n2..
Let i be the index of the subset C. We reformulate Vá (C) as
1
((C)
P
JEC
Ve
(c)]
D
w
1
E(C)
otherwise.
D
1 - wi
Since
II wi jEC,
II (1 wk) kgC,
it follows that
j jEC,
Op.
P otherwise.
wi
Hence,
Op Iw;
2 Σ
2" Σ
VEC-be [(CT]ll = Σ
(Ct)
p."h ((C.) +
D ((C.).
w
w
1=1
i=1, j4C;
On the other hand,
2 Σ
2"
2
Eowdrr ( w
20
(C)) = Σ P.,
Σ
1
p, ((C,)
wy
1 w,
1=1
JEC,
f=1, J&C;
Consequently,
( [C) VEO [((C)].
EOndw
18

---

Published as a conference paper at ICLR 2025
D PROOFS OF SECTION 7
k
1
We start with a technical lemma to bound the component
+
Σ
[[m]
j=m+1
Σ
wiy
Lemma D.I. Let w E An O + C E C and let x. be some permutation on the indices of the elements of C. Then,
k min (1.....)
W [m]
min w.
Wg
1mm++1
Σ [
DEC
qe
1=1
Proof. The fact that
[mm]
min W
min
Wy
JEC
oE
n
follows trivially from the definition of the minimum.
k Define wcC)) = Σ
[mm] Note that w(C) S
since w E An- For every permutation æ of the
m=
indices of the elements of C and every S k.
j-1
1-1
Σ
fn 2 (CC
Σ
wly 2 min wy
2 min
Wg.
9EC gE(l......)
E
Since un E An O 1 we have
1-1 1
min Wg
min Wg (E(I.....)
Σ
Wi
gEC
[m1
The sum
has most elements. Therefore,
j=m+1
k
Σ
W [m)
min wi 9EC
min W. ge (e(....,,n)
j 1m
-
Σ [
Lemma D.2. Let {w E R
< wj
< Vj. Σ w k}. Then,
max
[((C)
max max
GEC
jE (1......))
W w
19

---

Published as a conference paper at ICLR 2025
Proof. By the triangle inequality,
2
2
[VC- [((C)11]=
=
Σ
Pi ((C:) +
Σ
Pi e(Ci)
W
1 wi
1= jEC,
jæC;
pi [((()))))
+
Σ
p |E(C,))
Wi
Wj
i= JEC;
1= JEC.
1
max CEC
[(C)
p
1
Wj
w
JEC,
i=1, jdC
1
Using the fact that 0 < wy <
have that
and
Since p. is
wy
w,
W
w
probability, |,:00 p. Therefore,
2
VEC- [e(c))l,
max [((C)]
Pr
Pi
CEC
w
1 W
JEC,
(=1, J&C,
2
2
max Cec
[(C))
max
P
Σ
P
W; 1 W;
F= JEC,
14C
2
Therefore, since
Σ JEC,
1 +
Σ
D. = 1. we conclude that
J&C,
1 1
IVEErro [((C)))
S max [((C))
max
Cec
Wj" 1 w
Finally, we conclude that
VEC~ow
[(
< max |((C)
max (e.......))
max
CEC
W
w
Next, provide an entrywise bound for the gradient estimator.
Lemma D.3. Let w E {w E R
0 w
<
Vj
Σ
wi = k}. Then, for every C' E C,
[0)
S
max
((C)
max jE(L....,.
max
Wi - w }.
1 1
CEe
Proof. By the triangle inequality and the fact that < wj < 1, for every j E [I,...,
[[,1 = Ljec wj
1 E(C) < (((C)) max
1
((C) - liec 1- wi
w 1 w }.
Hence,
A (c)
< max
[((C)
max (e
max
Cec
w
w
An immediate corollary follows from the relation
||al|2 S / |alloc for every E R
20

---

Published as a conference paper at ICLR 2025
Corollary D.I. Let u fw Rd 00 wy <1 Vj. Σ w k). Then, for every C C,
| S vn max
1
[((C)
max jE(I....m) 10 1 20
max
Cee
Proof of Lemma 7.1. Let C' E C We now prove the two bounds.
1. In order to establish the first bound, we first show that
k
|E(C)).
min w gEC
Let Ci...., CN be some enumeration of C Denote the index of C' i. Let E nI. By the triangle inequality.
],
1 p
k!
Sp Sw,
k! Σ
ap...
V の w
E(C)
<
[(())
P
Oee
Since p!" is the probability of a subset that was sampled, p." 0, and therefore
> 0.
p
By the first part of Definition 6.1, which was proven Lemma C.I, and by w E A. 0 R
it follows that dp..
2 0. Hence,
rwg
Op Iw
3p..
and
p
Furthermore, if q E Cj then by Lemma D.I,
pptr. Ow
min w P.
If q & C₁ then it trivially holds that
Op..
k
min un Ph... )
Ow,
Therefore,
Op.
k
<
Pi
Dw,
min we gEC
and consequently,
A (C)
1
P |((C))
p
min W gEC
Since Σ ph
p.
( (c)
((()),
min wy DEC
and we can conclude that
Su (C)
k
[((C))
(D.I)
min W. QEC
21

---

Published as a conference paper at ICLR 2025
OpT, for every q & C, it follows that Vy has at most & nonzero
By the fact that ww
k
elements. By (D.1), it follows that each entry if bound by
[((C)).
Therefore,
min W
DEC
V% (
V
(0)
[
k-3
|((0)12
k1.5
|((C))
min w gEC
min wy DEC
2. If C2 or $
O.
iw, ow
Otherwise, there exist indices m, such that g π [m]. q T[o].
Denote d' (w) =
Σ
Note that
f=m+1
1
W
(din (w))
Σ
(D.2)
@w
j
j max [m,o)+1
If m = o, using the derivative rule for multiplication, it follows that
Opp @wimm
1
dai (w)
(D.3)
[w]j)
Owiml
W m)
1
Tw m)
+
d (w)
Pi.
Op.
die
(w)
[mm]
By (6.3),
Op
+
Lhat (w)
+ dir (ww)
pi..
W[m]
mm
W/mi
Plugging into (D.3),
0
0p
(D.4)
@wim]
1
Odi. (w)
1
2
d.. (wz) + (dhhn (mo))²
+
Pi.
w
@rv(m)
? [m]
[m]
(D.5)
- (adio (ux)
- @m
di. (w) + (dit (w)))2
p...
W[m]
j-1
Using the definition of d= (w), the fact that 1 - min wi and (D.2), follows that
Σ
wlg
2 min wy > and wim) 2
(a) ddi. (w) @w
< (k: 1) min my
22

---

Published as a conference paper at ICLR 2025
(b) 2 d (w) S 2k min w
-2
4imm]
(c) (d (w)) 2 1 (k22 2k + 1)
min w0
Combining with (D.4), we get
Op [m]m)
(mij ung
-2
$ (A² + 1:)
Pi.s
(D.6)
aw [m]
If o o m, the same arguments lead to Op..
+ Feat (no)
(D.7)
Ow @wm(m)
[m(m
Odht (w)
1
(
1
Pi,a +
+ dia (wz)
+ d= (w) pi.
[mmm
[1))
j-1
By the definition of d' (w), (D.2),
Σ
2 min wy > and w [m) W(o 2
8
min w > O.
(a) Odh (w)
S (ks 1]
min w
(b)
min w
[mm] [[]
-2
(c)
die (w) + Wiml
d/ (w) S 2k
min w 3
-2
(d) d (w) d= (w) I (122 2k + 1) min wj
Plugging into (D.7), it follows that
0p..
S (1.2 + k + (min un
-2
(D.8)
Tw
@wmm)
ppp θqoq
By (D.4) and (D.7), it also follows that for every á E Ci.
0. Since if either
pppt 0qoq
q $ C. or q & C it holds that
= 0, we conclude
20 0 Vg. E (1....,n))
(D.9)
Ōqoğ
Using the triangle inequality, it immediately follows that
a²
N Al
(ECrvr [{(C]) =
(D l(Ct))
@w,ow,
Swoōw
A1
(pro.))
[((C))
Twoory
82
(p
max (E(C,))
awyewy
< (k22 + k + 1) (min ung
max (((C.)
23

---

Published as a conference paper at ICLR 2025
Let pa(())) (h22 + k 1)
(min mii
max; |((C:)]. Since
-2 (Re- (((C))) 2 (k<())- -(h² + k 1). (min (in
pk(w)) + ow,
max (((C,)],
and
8
pk(w) Ow,
(Bonvw [((C]]) 2 (2( - (1.2 +k + min wz m0;
max (((C.))
and in addition
(min -2
max
Σ
P [((C')] S (n- (k2 (62+))++++00 +k
-max ((()))
it follows that for all i
(w(w)) + Bcww [((C)] 2 Σ
V BCnw [((C)]]
j=1, jzi
and
(=(w)) - 1 [((C))]
1, 2
Σ
BCrpp [((C)] i
Therefore, both (k(w) + VLEcmw [((C)] and ph (w)/ - BC [((C'] are diagonally
dominant. Since diagonally dominant matrices are positive definite, follows that
P + V2EC [((C]] > 0 = V'Ecatn [((C)] > (w(w)I
p: (= - v B [((C)) > 0 = pk(w)] > V-Ecs (((C))
In conclusion,
Pa(w) > V [((C)] > -px(w)l.
Proof of Lemma 7.2. We prove the two parts.
1. For every j E [l,...,,
n2) by the law of total expectation and the definition of V (C),
Ec~ov
= P (j is chosen)
EEnno 1j chosen [
w
+ (1 (j chosen))
EErow is not chusen [((C)²].
(1 )-
Since
IP (j is chosen) W
and
max (EC chosen
[E(C)?] E.orr i chosen [((C)
max ((C)2
CEE
can derive that
1
Ecag
5
+
max I(C)
1
w.
CEC
Therefore,
Eondw
|
5 Σ
1
max ((C)
wi
w
CEC
24

---

Published as a conference paper at ICLR 2025
2. Recall that
2"
Ecnew
[[[C]] = Σ
II i=1 jEC;
w II
(1 - w,)
jEC,
and that
P4
II w
(1 - m0
JEC.
JØC,
For every and every j.
8 D.
= 0,
Om
and hence for every j [I....,m.]
V2 BEne [e(C)]l = 0.
For i E Co.
1
P,
If C,14 Ci.
0ppr
Pi
Own
wy (1 wi)
If & C, C..
pr
p.
ow Ow:
(1 - wi wi
i & C..
82p
D.:0
iw,
(1 - (1-w;) (1 w01
Therefore,
0²
Benera [R(C)]
Ovy, Brn
Saec
Lec,ee wy (1 (l:w;))
1e0,1e,, (1 - w.) ii
1,060 (1 - (1-w;) (1 wi)
p." ((C))
wyw
Using the facts that
1
max m. Wm ww,)
<0
Wj (1 wi)
and
1
< max wywi" (1 w,) wi)
max max
m,
Wm w. (1 wm)) (1 w.
well as the fact that for every
((CA) max ((())
CEC
and that
2
Pr
1,
1=1
can deduce that
82
max "2,""
max
[(C)) S ow
Ecroin [((C']
Wm (1 w.)
CEC
max
max
[(C)) CEC
m,r"
Wt, Wy (1 Wm)) (1 w,
25

---

Published as a conference paper at ICLR 2025
Consequently,
0² own
Bcnow [((C)]
(D.10)
1
1
< max m
max
|((C))
mar
1m
(1 Wmm) (1 wr) Wm (1
w.))
OEC
By (D.10), for every j.
PB(W)
ow, Ea [[C)]]
1=1,#
and hence Vrwar [((C)] + pB(w)L is diagonally dominant. Since every diagonally dominant matrix is positive semi-definite, it follows that
V2Ea- [((C)] + pn(w) 0
and therefore
Amain ( [(())) -pB(w)
For similar reasons, (B(w)
BCEndu [((C)] is also diagonally dominant, and hence
PB(W)) Ba [((C) ł 0,
which leads to
Amar EV [((C)]) & PB (w))..
Proof of Theorem 7.1. The result follows immediately from (Zhang and He, 2018, Corollary n (k² + - max I((C) (as proven in Corollary 7.1), L = CEC*
3.1), with the parameters
c²
k1.5
max I((C') (proven in Lemma 7.1), (wo) < max ((C), Tmin 2 min ((C).
Ceck
Cec
Ceck
ē n k.'
Using the fact that
where é is independent of n, yields that particular result.
Proof Theorem 7.2. The result follows immediately from (Zhang and He, 2018, Corollary 3.1),
2n
with the parameters p = n max (((C') (as proven in Corollary 7.1), L2
max maxx(())² (proven CEC
c2 Cec
c
in Lemma 7.2), (1/2p(wo) S max E(C),, Tmin 2 min E(C).
CeC
CEC*
Proof of Theorem 7.3. Note that min ((C.) BENDwN [e(C)) 2 8 + ((C.) - EC [e(C)) 2 8 Vi. iE(1.....m
By Hoeffding's inequality, for á,
P (((C.))- - EonDuN [e(C)] s Eexp {
82
)
2 - -tu-
Since Or...., Om are i.i.d.
P min ((C.) BONDing [[(C)]2 =P (((CC))) - -ECDDw [E(C)]
m²
S exp {
mo²
2(u - ()
26

---

Published as a conference paper at ICLR 2025
E MIRROR DESCENT STEP CALCULATION ALGORITHMS
Let <<< -. We consider the mirror descent step for all the combinations of
1. X1 : A E cs 2j Vj E [I....,n]] ,.]
2. 82 = A = for ER c x3 S 1 vj (1, (n......
Σ Σx zj = k}
and :s: + for E (1,2,3} and j E 23. where
1
1. m1(2) |||||
2. p2(2) = Σ In;))).
1=1
3. ua(x)= = Σ
I; In(x;).
1=1
Lemma E.1. Let A + R be defined as m1(x)= = ||=11² Then,
1
argmin =ED;
2) + Bp, (I) = {=²},
where 2" is given by
1² < I2 -
Ti a +
9b
otherwise
and μ²' is the unique solution of
Σ
c+
Σ
(x-- a(u 9i)) 1.
(=>(B_-///-9 i: (((ln-c)/a-g
Furthermore, μº can be found using bisection.
Proof. We start by proving that μ² exists, is unique, and can be found using bisection. The function continuous, strietly decreasing when there exists least one index such that μ S (1 c)/a 9s. and
Σ c+
Σ
(x+ a(k + gi))
00,
is =
Σ c+ Σ
(m--a(f - + 9i)) m
Therefore, 1² exists, is unique, and can be found using bisection.
Bug (x,, z) continuous and A, is compact, minimizer exists. Since (g x) +Bpp (x,=) strongly convex, the minimizer is unique. Furthermore, due to the convexity, the KKT conditions are sufficient.
Since 2),x)+B,aa
Define the Lagrangian L Rd x RT x I,
||2
xll +
Σ
2j
+
2a
j=1
The KKT conditions are
1
1. 93
+
2,
+ 4 Ag = for all f1.....n).
a
27

---

Published as a conference paper at ICLR 2025
2. Ay(c- 2,) = 0 for all j ml).
3. Σ 2j =1
/=1
4. 2 2 e for all n..
5. 020.
We can see that x², , X², where X is given by
Ij
0
14
5
9j
1
2
gj otherwise
is a valid solution to the KKT system. Hence, =* is the optimal solution.
"
Lemma E.2. Let u2 A + R be defined as p2(x) = - Σ (x(x)). Assume that a E A. Then,
i=1
argmin (g, 2) + a Bpg = {="}.
=EA:
where 2" is given by
2 = max
1/20 a (gi wry.
and μ² is the unique solution of
max
c) =1
1/2, + a (g: + u)
on (max-
azj 9j. 00). Furthermore, μº' can be found using bisection.
Proof. We start by proving that μz exists, is unique, and can be found using bisection. The function continuous. Additionally,
max
1
1-000
n n.cc 1
1/x;
(gi 1c)
and
(
-,,**
Σ m
1 1/x4 + a (g. + 1*)
axt )= 00.
Therefore, solution μ² exists.
For μ max- 1 Ij. 1/2; + a (g. + u*) is strictly decreasing.
) aT;
Therefore, the solution μ² is unique on (max-
i az, Ij. 00), and can be found using bisection.
Define the Lagrangian L Rd x RT x IR,
1 + + 21 2; a In(zs) p
B
L(z, 11
0i az;
X
The problem is continuous and the feasible set is compact, therefore optimal solution exists. Since the problem is also convex, the KKT conditions are sufficient for optimality. The KKT conditions are given by
28

---

Published as a conference paper at ICLR 2025
1
1 11 -0 for for e n... 0
1. 9j +
+ X,-
az)
2. Ay(e- - 2,) = 0 for all j nL.
3. Σ =1
/=1
4. 2j 2 e for all j E nn.
5. 0.
solution to the KKT conditions is given by (z", (=,,,,,," ,²",,, where =" and μ=" were defined above and
0, 9j + 11
X;
otherwise
ac
To complete the proof, we will prove that X 2
0. If 23 > c. then X7 2 trivially. If 23 = then
1
S C.
(E.1)
1/2, (gj 1x*)
Since both sides of the inequality are non-negative, it follows that
1 ac
S 1 + +r,
a.jjj
and therefore
1
= 9j
1rt +
2 0.
a3,
ac
11
Lemma E.3. Let μ3 A, + R be defined as p2(x)) = Σ
23 In(x;). Assume that x E A Then,
i=1
argmin(g, =EA:
2) + a (p = {2"}.
where 2n is given by and p² is the unique solution of
z max ((artt))))))
max - a(g+ta) .c = 1.
Furthermore, p² can be found using bisection.
Proof. Note that as long as there exists an index i such that
a(g:ta") > c,
I&e
the function
max (ax (bie a(sit()) .c)
strictly decreasing in μ. Furthermore,
max a(ortt)) -afxrtt) p--00 0o.
29

---

Published as a conference paper at ICLR 2025
and
max (exx(1e))) [re (aae-alattt .) - += n <1
Therefore, the value u² can be efficiently calculated using bisection.
The Lagrangian L 12"h x 124 x R is given by
M) (a.2)
Σ #s In(a,) Σ z, In(2:) Σ ann + 1) (2()-)))) -.)) +u
E 24 -
+ Σ a (a 2.)
Rearranging,
12,,,11) =
In
Zz
9i + 1
X1
2i 1
Ae + const.
The problem is continuous over a compact set, and therefore, by Weierstrass' theorem, a minimum exists. Since the target function is convex, the KKT conditions sufficient. The KKT conditions are given by
1.
+ Di + A; for all n..
Fi
2. Ay(c- - 2 = for for all j nn.
3. Σ 2j 1.
,=1
4. 2/ 2 for all n).
5. A2 0.
The KKT conditions are fulfilled by (=", X², ,²), where 2" and n² are as defined above, and X is given by
0,
otherwise
If =] > then AA 2 0 trivially. Otherwise, by our choice of aJ. it follows that
3,6 a(gttt) <
and hence that X3 2 0.
1 Lemma E.4. Let #1 Ai, + be defined as pı(x) ||0 and min(0.5, 1/n). Then,
argmin 2) + a Bp, Bp, (x, = {2*}.
=EAr
where is given by and n² is the unique solution of
= min =in
max a + - c).
cJ k.
Furthermore, n² can be found using bisection.
30

---

Published as a conference paper at ICLR 2025
Proof. Note that as long as there exists an index j such that
c < xz a (gg x²) --c,
the function
min m a (g) + ,,), - c}
is strictly decreasing in Furthermore, the function is continuous, and
min (maxf2, a (0) ,,) 1 cf - n.c n.cklsh
and
Σ
min a (g) (0),, 10 n (1 - n 12k 12h. j=1
Therefore, a solution n² exists and can be found using bisection, proving the latter part of the lemma. To see that the proposed 2" indeed the optimal solution, note that
1 || x=|2
B₁g (z, x)
Therefore,
(g. 2) + Bu (I, 2) = (g, 2) + 20x
|| x//2
a
The Lagrangian L 1²"
R4 x Rº IR given by
1
=(z,A,n,,)=
+
|| - x1| + Σ
(c 2,) +
7p (2 - 1 1c +
Σ - k
2a
1=1
1=1
Rearranging,
L(z,A,7,P)= Σ
1 +
2
g1
Ij
+
10
2i +
Σ
(c-1))
20 |
a
The function is continuous and the feasible set compact, and hence, by Weierstrass' theorem, minimizer exists. The feasible set and target function convex, and therefore the KKT conditions are sufficient for optimality.
The KKT conditions given by
1. - 2j + 9,
\ + 7 + p = for all j E ni.
a
2. Ay(c- 2,) 0 for all j mn..
3. (22 - 1 + c) 0 for all j E n).
4. Σ 2/ = k.
j=l
5. 2, 2 for all n1.
6. 2 < 1 - for all j E nI....,n..
7. 722 0.
8. 1) 2 0.
31

---

Published as a conference paper at ICLR 2025
The vectors (2*, A", ,n,, p²) are solution to the KKT system, where 2" and f were defined above, X is given by
>
A3 =
c Øj
+ 9j + ,
otherwise 1
and r7" is given by
0
< <
71
Ij 1 +
9) u
otherwise
a
1 2 0 trivially when 23 When 2] it implies
Tj a (gj 1² c,
and hence
リj + p 0. 2j Ijj
Likewise, vi 2 0 trivially when zj < When 2j = 1 c, it implies
2j (و) + (²) 2 -
and hence
nj 3 1
9j 11 O.
Lemma E.5. Let P2
An - R be defined as p2(x)
Σ In(2;). Assume that
min 1/n]. Then,
argmin
2) + a (2, {
=EAg
where: z is given by
1
mixx
1/2; a (gi + pr) 00,1 cf
and n² is the unique solution of
1 1/x; (gi + ct,
Σ min
on max 1
0., 9j. 00
Farthermore, p² can be found using bisection.
Proof. Note that as long as there exists an index j such that
1
<1 <1-c. Cc
1/2, + (g. + p*)
the function
n
min (max(x; - a agg + n²)) el,
max 1
is strictly decreasing in μ
az, -j.. -9j.00 Furthermore, the function is continuous in
max 1 91 0 and
ax;
p max ax, -9;
min min(max{2 (9) + +,)),,,111
(1 c) 12
32

---

Published as a conference paper at ICLR 2025
and Σ mix - (0z + (,,,),
cf m n < S k.
Therefore, solution μ exists and can be found using bisection, proving the latter part of the lemma.
Note that
Byz (2c, 2) In (zz)
In(ars)
n1
f=1
Therefore, the Lagrangian L: RR
R
R is given by
(a. 2)
hn(2s) +
Σ Acle 2) Σ 1+0
Σ 151
26 2x-1)
const.
1=1
ze
Rearranging,
M(2 = Σ (
-
1 az;
2 - In(zi) + const
The feasible set compact and the target function is continuous, and therefore, by Weierstrass' theorem, a minimizer exists. Since the problem is convex, the KKT conditions are suffcient for optimality. The KKT conditions are given by
d + + 1 0 for all j E nl. O.2;
1. 9:
03;
2. 2,(c- - 2,) = 0 for all j E mI.....m..
3. ns(2) 1+ 0 for all j E nI.
4. Σ 2j k. f=1
5. 23 2 c for all n..
6. 2 S - c all j E (I,...,n).
7. 0. 0.
8. 1)20 0. 0.
The vectors (2",^", X given by
>3, nn" 1²)) are solution to the KKT system, where 2" and μ were defined above,
1
9j +
+ +
otherwise
ax,
ac
and 17" is given by
0,
<1-c
n
1
9j ax,
1² otherwise
a (1 c)
Note that since < 0.5, 27 c and 27 are mutually exclusive, and therefore least of Xj. n3 is zero.
If 2 c. X3 2 trivially. Otherwise, by the choice of 27.
1 1/2, + a (g. + )()
S.c.
and therefore
27 gj
+ 11
2 0.
OT
ac
33

---

Published as a conference paper at ICLR 2025
If 2] < then nj 2 0 triviallyy. Otherwise, by the choice of ej.
1 C, 1/2; + a (gi + 11))
and therefore
73
91 azj
1¹ 2 0.
(1 c)
Lemma E.6. Let M3 A R be defined as pz(x)= Σ 2; In(x;)) Assume that æ E D. and
< <minn0.5, 1/n). Then,
argmin 2) + B Bus (x, = {2"},
where z" is given by
27 - mix -aa a(att()) ,0,, 1 cc
and n² is the unique solution of
mma alx cc,, - c) = k.
Furthermore, μ* can be found using bisection.
Proof. Note that as long as there exists an index i such that the function
TGe a(gta") 1 C,
min (max (a ,ct, cc
is strictly decreasing in j.. Furthermore,
max alscta), -alette-alattt 1ー n . (1 c) > n 1,
and
max max(x,e a(la c) n=0 n
Therefore, the value n² can efficiently calculated using bisection.
Note that
Bug (z) 2)
24 In (2)
2i In(xx.) Σ (In(x:) + (2+ xz)
Zi (ln(())) -
(In(x;) + 1) z; +
12
i=1
i=1
1=1
1
2
24
2
1=1
Therefore, the Lagrangian L : 1.
x
R
X PT
x
is given by
Σ
24 In
Σ
24 +
Σ Iz
i=1
i=1
+ μ
+
X (c - 21) + Σ
1h (z+ - 1+ c).
i=1
i=1
1=1
34

---

Published as a conference paper at ICLR 2025
Rearranging,
(00 - -n (z,) - (2, = 1=1
= +
1 - In (zi) + const.
The feasible set is compact and the target function is continuous, and therefore, by Weierstrass' theorem, a minimizer exists. Since the problem is convex, the KKT conditions are suffcient for optimality. The KKT conditions are given by
1. 9: + + 7i
In (x:) + (2)) for j nI.
a
2. Ay(e- 2)) = 0 for all j E n..
3. ny(2 - 1+ c) 0 for all j E {I..... nn.
4. Σ 2j = k. /=1
5. 2 e for all j E n.).
6. 2 S 11 1 e-efor for all j E (1......} n).
7. 0.
8. 1) 1) 0.
The vectors (2", x3, ,7" 1²)) are solution to the KKT system, where 2" and μr" were defined above, X is given by
0,
>
In
24e
otherwise
and ry* is given by
0,
< 1
ZiC -alsitu")
In
c
otherwise
Note that since < 0.5, 23 = and 2j = 1 - are mutually exclusive, and therefore at least one of ^ラ. 7 is zero. If 23 c. then X3 0 trivially. Otherwise, by the choice of 2j.
T&e olgctta) S c.
Hence,
In X5
TGe a( a(g
0.
If 2 < 1 then n 2 0 trivially. Otherwise, by the choice of 27.
24e aagitta)
2 c
and it follows that
In WGE alsitta"
C
2 0.
Remark E.1. If 0.5, the feasible set Ai is always empty. If 0.5, feasible solution exists if and only if n = 2, in which case the only feasible solution is 0.5).
35

---

Published as a conference paper at ICLR 2025
F SCALABILITY AND PERFORMANCE CONSIDERATIONS
Regarding the practicality and relevance to large-scale applications, below we provide adaptations which may improve the performance of Algorithm 1.
F.1 THEORY-PRESERTE ADAPTATIONS
1. Permutation sampling:
Sampling a permutation instead of a subset and adapting the gradient estimator accordingly reduce the computational cost of gradient evaluation. While we have not yet proven this rigorously, we believe the gradient estimator bounds would still hold. This approach may, however, increase the gradient's variance.
2. Caching subset evaluations: Storing values of previously evaluated subsets in a hash table can reduce redundant calculations, especially when the algorithm close to convergence and repeatedly samples the same subsets.
F.2 HBURISTIC ADAPTATIONS
Combining sampling techniques: For large k; we can integrate our approach for sampling subsets with expected cardinality k with the method in Pervez al. (2022). This would yield subsets of exact size k2 albeit at the cost of slightly biased gradient estimators. Given the design of the sampling process in Pervez et al. (2022), expect the bias of the gradient estimator to remain low.
2. Objective evaluation for hyperparameter tuning: Direct evaluation of the objective function is computationally expensive for large-scale problems. Instead, sampling can be used to estimate the objective efficiently. With upper and lower bounds available, the sampling average converges at sub-Gaussian rate.
3. Early stopping criterion: If a sufficiently good subset is encountered during sampling, we can terminate the process early. This approach avoids unnecessary computation and provides practical stopping condition for the algorithm.
G SYNTHETIC EXPERIMENTS
This section presents synthetic experiments for three Subset Selection tasks: Subset Sum, Un- structured Subset Selection, and Sparse Least Squares. The Subset Sum experiment evaluates the performance our algorithm on a combinatorial problem without an evident continuous analog. The Unstructured Subset Selection experiments assess method's effectiveness in setting devoid of underlying structure. Lastly, the Sparse Least Squares experiment benchmarks our approach against continuous optimization algorithm suited explicitly for that problem.
G.1 SUBSET SUM
G.1.1 SETTING
The first set of synthetic experiments focuses on the Subset Sum problem. In this problem, given n numbers Y "y....... Un E R), a subset size k, and a target value E R. The goal is determine whether there is subset of elements & from Y such that the sum of the elements equals The Subset Sum problem is well-known to be NP-Complete. In our synthetic setup, we set and y/ Uniform(0 [0, 10] for all e {l...., n). The target chosen the sum of a randomly selected subset of size k from Y, ensuring that a solution always exists.
10,
We evaluate the Subset Sum problem under the settings proposed in this paper: selecting subsets of size k = 4 exactly (formulated by (Pt)) and selecting subsets of size k = 4 in expectation (formulated (P))). In both cases, for any given subset C S Y. the loss function defined
((C) =
Σ DEC
Clearly, C' solves the Subset Sum problem if and only if (CC) 0.
36

---

Published as a conference paper at ICLR 2025
In the (Px)) setting, Algorithm with three Bregman divergence variants outlined in Corol- lary 7.1. For comparison, we implement two baseline stochastic algorithms:
Uniform Sampling: each iteration, subset chosen uniformly at random, and the best-observed subset is returned.
2. Stochastic 1-Flip: It is based on the 1-Flip algorithm, which is discussed, for example, in Szeider (2011), At each step, this method randomly selects two candidate elements, one not in the subset to replace the other in the subset.
For the (Pz) model, only compare with the Uniform Sampling procedure due to the lack of straightforward analog for the Stochastic I-Flip algorithm in this setting.
G.1.2 PARAMETER TUNING
In both settings, experiments have shown that the theoretical values of a derived Theorem 7.1 and Theorem 7.2 yield a very slow convergence rate when the Bregman divergence base function is 42 Σi= Iú (mx,)) or pa(x) = = -i In(x.). Consequently, we adopt larger values of a:
n 2.55. 0.5 for (2(2x), and a n..570..5 for (z)()))
where e and T are defined as in Theorem 7.1 and Theorem 7.2. Although these step sizes are not directly supported by our theoretical analysis, they are strongly inspired by it. The Hessians of (2())) and (z((2) are:
1
V u2(x1) diag
V (((()) = diag
21
x
The bound on p derived in Lemma 7.1 and Lemma 7.2 is O (m. xz)) -2) condition:
relying on the
(
-2)
(=)D(x++0)
I O. min 34 (=l,..,,
we assume the stronger result:
v (² diag
1 0,
then the values of p improve to O (2) for pa(x)) and 0(1) for (z(x)) Based the proofs of Lemma 7.1 and Lemma 7.2, we believe that this heuristic-type assumption is reasonable. Indeed, under this assumption, the chosen values align with the convergence guarantees in (Zhang and He, 2018, Corollary 3.1).
G.1.3 RESULTS
The results of the Subset Sum experiments under the (P:) setting are presented in Figure 3 and Table 2. Figure 3 plots the function values and best-observed subsets for two representative instances, while Table 2 summarizes the outcomes of 1000 experiments. Across all configurations, the methods implemented using Algorithm significantly outperform the Uniform Sampling baseline.
Among the three Bregman divergence-based approaches, those utilizing /1 (x) = ||~ (Euclid- Squared) and (22(z) = Σ I₁ In x; (Negative Entropy) as the base functions outperform the Stochastic 1-Flip algorithm. However, the approach based on pa(x)) Σn= In(xx) (Minus Ln) performs slightly worse than Stochastic 1-Flip.
Interestingly, the Negative Entropy method outperforms EuclidSquared in minimizing the function value despite achieving poorer results in identifying the best observed subset. We hypothesize that this diserepancy arises from the higher variability in function values exhibited by the Negative Entropy approach during optimization. This increased fluctuation may lead to sampling subsets from less favorable distributions at specific iterations compared to the EuclidSquared method under equivalent conditions.
37

---

Published as a conference paper at ICLR 2025
Subset Sum
Exactly 4
Function Valun Eurlidiquared
00
Best Subset Value EuchiSqquares
14
Negative Entropy
Best Subset Value Negative Entropy
12
Function Value Minus Ln
Subset
Functoon/ Best Obsurves
Subset Value Uniform Sampling
Best Subset Value Stothastic 1-
Optimal
20 Feration
40
Subset Sum Exactly 4
Function Value EuclidSqurred
Best Subset Value EuclidSquaree
Function Value Nogative
Function BBst Obsrvvdd Subset
15.0
Negative Entropy
Function Value Minus
12.5
Blest Subset Value Minus
Sampling
Best Subset Value Stachastic Flip
Optimal Value
7.5
5.0
0.0
10
20
30
40
Iteratio
Figure 3: Function values and best observed subset values per iteration representative samples for (Px) setting.
Best Observed Subset
Last Iterate Function Value
Mean
Std
Mcan
CWR4 EuclidSquared
0.145
0.270
3.607
1.161
CWR4 Negative Entropy
0.160
0.294
2.888
1.358
CWR Minus Ln
0.218
0.467
4.156
2.022
Uniform Sampling
0.244
0.474
NA
NA
Stochastic 1-Flip
0.208
0.313
NA
NA
Table 2: Average value of best observed subset and function value for 1000 subset sum instances (P&) setting, Bold indicates the best in "Mean" column.
The results for the (Pn) setting are presented Figure and Table Figure displays two representative problem instances, while Table provides summary of the results of 1000 experiments. The three variants of Algorithm outperform the stochastic sampling baseline, with the Minus Ln approach achieving the most significant improvement.
In terms of function value, the Minus Ln approach achieves substantially lower values compared to the Negative Entropy method, with both approaches significantly outperforming the EuclidSquared method. We hypothesize that these differences are primarily attributable to variations in the step size parameter, which may impact the optimization dynamics for each divergence function.
38

---

Published as a conference paper at ICLR 2025
Subset Sum In Expectatian
Function Valun Furlidiquared
Best Subset Value EucliSqquares
Negative Entropy
30
Best Subset Value Negative Entropy
Function Value Minus Ln
25
Subset
Functonn Bast Observes
Subset Value Uniform Sampling
20
Optimal Value
15
10
20 heration
40
Subset
Expectation
Function Value EuclidSqurred
Functioo Becst Observed Subset Value
30
Best Subset Value EuclidSquaree
Function Value Nogative
25
Negative Entropy
Function Value Minus
Blest Subset Value Minus
20
Sampling
Optimal Value
15
10
20
30
40
Iteration
Figure 4: Function Values and Best Observed Subset Values Per Iteration Representative Samples For (P&) Setting.
Best Observed Subset
Last Iterate Function Value
Mean
Std
Mean
Std
HIB4 EuclidSquared
0.269
0.313
7.815
1.363
HIB4 Negative Entropy
0.267
0.307
6.432
1.287
HIB4 Minus Ln
0.239
0.265
5.760
1.895
Uniform Sampling
0.278
0.332
NA
NA
Table 3: Average value of best observed subset and function value for 1000 subset sum instances (Pg) setting. Bold indicates the best in "Mean" column.
G.2 UNSTRUCTURED SUBSET SELECTION
G.2.1 SETTING
In the unstructured experiment, the value of each subset is assigned independently, with values sampled uniformly random from the range [0, 10). To ensure that the optimal subset value is O, subtract the value of the optimal subset from all subset values. The same algorithms, iteration counts, and step sizes described in Appendix G.1 are used in this setting.
The unstructured nature this experiment poses a significant challenge for algorithms that rely on structural patterns, little structure is present.
39

---

Published as a conference paper at ICLR 2025
G.2.2 RESULTS
The results for the (P:.) setting are summarized in Table 4, and two representative instances are visualized in Figure 5. The results show that the Stochastic 1-Flip algorithm, which assumes a strong underlying structure, struggles significantly in this setting. While the variants of Algorithm I demonstrate better performance, only the EuclidSquared version marginally outperforms the Uniform Sampling baseline, and even this improvement is negligible.
These findings highlight the inherent difficulty the unstructured problem, where information about subset offers predictive insight into other subsets and very little insight regarding the value of individual elements. This lack of interdependence severely limits the ability of algorithms to exploit any structural advantages, resulting in diminished overall performance.
Unstructured Exactly
Function Value EuclidSquared
Functton/ Best Obsurved Suboit Value
Best Subset Value Euclieeuurres
Negative Entropy
Best Subset Value Negative Entropy
Function Value Minus
Subset
Subset Value Uniform Sampling
Best Subset Value Stathastic 1-
Optimal
20 Beration
40
Unstructured Exactly
10
Function Value EuclidSquared
Function Besr Otterves Subsot value
Best Subset Value EuclidSquaree
Function Value Negative Entropy
Nogative Entropy
Function Value Minus Ln
Best Subset Value Minus
Sampling
Best Subset Value Stochastic Fip
Optimal Value
10
20
30
40
Figure 5: Function values and best observed subset values per iteration representative samples for (Pş) setting.
Best Observed Subset
Last Iterate Function Value
Mean
Std
Mean
Std
CWR4 FuelidSquared
0.1
0.235
4.788
0.330
CWR4 Negative Entropy
0.228
0.256
4.426
0.840
CWR Minus Ln
0.247
0.278
4.781
0.524
Uniform Sampling
0.203
0.238
NA A
NA
Stochastic 1-Flip
0.380
0.405
NA
.AA
Table 4: Average value best observed subset and function value for 1000 subset sum instances (P&) setting. Bold indicates the best in "Mean" column.
40

---

Published as a conference paper at ICLR 2025
In the (PD) setting, Figure and Table 5 illustrate that all variants of Algorithm 1 struggle to significantly reduce the function value within the allotted number of iterations. The function value remains nearly identical to that of the uniform distribution, and the average best observed subset is similarly close to the one produced by the Uniform Sampling method.
Unstructured Expectation
Tunction Value EuclidSquered
Function Bostt Otsered Suboot Value
Subset Value EucliSSquares
Nogative Entropy
Best Subset Value Negalive Entropy
Function Value Minus
Best Subset Value Unifom Sampling
20 Iteration
40
Unstructured Expectation
Function Value EuclidSqured
Function Best Otserved Suboet Valae
Best Subset Value EuclidSquared
Function
Subsel Negalive Entropy
Function Value Minus
Subset Value
Sampling
10
20
30
40
Iteraiio
Figure 6: Function values and best observed subset values per iteration representative samples for (Pa) setting.
Best Observed Subset
Last Iterate Function Value
Mean
Std
Mcan
Std
H1B4 EuclidSquared
0.234
0.235
4.980
0.111
III44 Negative Entropy
0.240
0.244
4.952
0.155
HIB
Minus Ln
0.239
0.249
4.901
0.336
Uniform Sampling
0.238
0.246
NA
NA
Table 5: Average value of best observed subset and function value for 1000 subset sum instances (PB) setting. Bold indicates the best in "Mean" column.
G.3 SPARSE LEAST SQUARES
G.3.1
SETTING
The final set of experiments is conducted in Spansity Least Squares setting, defined as
min {ILAc - b1]2 subject to ||=||o s k) 2ER
41

---

Published as a conference paper at ICLR 2025
where the I-llo "norm", which counts the number of nonzero entries, is given by
|
Σ
21
In each problem instance, we sample the matrix A E R$x10 where each entry is independently sampled from standard normal distribution: A.g (0, 1). The ground truth vector is generated by first sampling entries æ; ~ (0, 1) independently, and then setting all but randomly selected entries of è zero. Finally, we set Aø + e, where the noise vector has entries sampled independently from e, ~ (0,,0.1).
We apply all the algorithms and settings discussed in Appendix G.I. Additionally, in the (Px) setting, we also test a projected gradient descent algorithm, with step size of 1, where the Lipschitz constant of (xx) || bll². The convergence properties of projected gradient descent in this setting are discussed in Beck and Fldar (2013).
G.3.2 RESULTS
The results for the (Px) setting are provided in Figure and Table 6. Two of the three variants of Algorithm outperform other methods, despite the relatively structured nature of the problem. which we believe to be advantageous for an algorithm such as Stochastic I-Flip.
Least Squares Exactly
Luclidiquared
Fundtion Bust Observed Subset Value
Best Subset Value EuclieSquared
Function Value Negative Entropy
Subset
40
Function
Best Subset Value Minus in
Best Subset Value Uniform Sampling
30
1-Flip
Optimal Value
20
Projected Gradient Descent
10
Sparse Least Squares Exactly
175
Function Value EuclidSquered
Function Bust Observed Subsut Value
EaclidSquared
Function Value. Nogative Entropy
250
Best Subset Value Negalive Entropy
Minus
325
Best Subset Value Minus
Best Subset Value Uniform Sampling
Subset Stachastic
Optimal
75
Projected Gradient Descent
50
25
Figure 7: Function values and best observed subset values per iteration representative samples for (P&) setting.
42

---

Published as a conference paper at ICLR 2025
Observed Subset
Last Iterate Function Value
Mean
Std
Mean
Std
CWR4 EuclidSquared
1.056
1.380
8.208
8.986
CWR4 Negative Entropy
1.026
1.505
7.493
11.311
CWR Minus Ln
1.222
1.601
10.772
10.312
Uniform Sampling
1.331
1.658
NA
NA
Stochastic I-Flip
1.067
1.532
NA
NA
Projected Gradient Descent
3.690
4.438
NA
NA
Table Average value best observed subset and function value for 1000 subset sum instances (P&) setting. Bold indicates the best in "Mean" column.
In the (PB) setting, the variants of Algorithm 1 outperform the Uniform Sampling approach. Similarly to the results for the (Pæ) in Appendix G.I, we can see that the best average function value result does not necessarily translate to the best observed subset result. As before, speculate that this is due to fluctuations in the function value in the Minus Ln approach which are very apparent in the sample instances illustrated in Figure 8.
Expectation
175
Luclidsquared
Fundion Bust Obser Subset Value
Best Subset Value EuclieSquaree
Function Value Negative Entropy
Subset
Best Subnet Value Minus in
100
Best Subset Value Uniform Sampling
75
50
Sparse Least Squares Expectation
10
Function Value EuclidSquared
Bundtoon/ B8st Observed Subzet Value
EaclidSquared
70
Function Value: Nogative Entropy
Best Subset Value Negalive Entropy
60
Minus
Best Subset Value Minus
50
Best Subset Value Uniform Sampling
40
30
20
10
heration
Figure 8: Function values and best observed subset values per iteration representative samples for (Pz) setting.
43

---

Published as a conference paper at ICLR 2025
Best Observed Subset
Last Iterate Function Value
Mean
Std
Mean
Std
HIB4 EuelidSquared
8,982
9.140
34.605
27.865
HIB4 Negative Entropy
8.688
8.372
32.084
32.375
HIB4 Minus Ln
10. 10.483
10.748
28.663
30.810
Uniform Sampling
10.863
9.798
NA
NA
Table 7: Average value of best observed subset and function value for 1000 subset sum instances (Pn) setting. Bold indicates the best in "Mean" column.
44


Reviews:
[{"review_id": "NRa7wpSWRs", "reviewer": "Reviewer_cTYy", "summary": "This paper studies the subset selection problem. A two stochastic formulations of the \nproblem are proposed, as well as a computable approximations, and solved via mirror descent, and \na $1/\\sqrt{T}$ rate of convergence to a stationary point is proven. The methods are tested \nempirically in a transfer learning task.", "questions": "- As mentioned above, the approach seems to work even when the loss is defined in such a way that \nevery subset maps to a distinct and unrelated loss. This suggests that the problem setting should be significantly harder \nthan e.g. the combinatorial bandits setting, where we just need to find all the \"best\" indices. In fact, this difficulty seems to be reflected directly in the upperbound: the numerator is a factor of $n^2$ larger than what is attainable in the combinatorial bandits setting ($\\sqrt{n}$). Do you expect that this is the optimal dependence?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It is rather surprising to me that the approach is seemingly able to \nget around making assumptions on the underlying structure of the losses that \nthat one would see in the combinatorial bandits literature --- ie, that \nthe losses take the form $\\langle\\ell_t, w\\rangle= \\sum_i \\ell_{ti}\\mathbb{I}(w=1)$, whereas in this \nwork it is possible to have losses defined in such a way that every possible subset maps to independent values. \n\n- I really like the idea of second approach, in which the condition that $\\\\|w\\\\|_1\\le k$ is relaxed \nto $\\\\|w\\\\|_1\\le k$ *on average*, in exchange for a more readily computable update\n\n- The paper is generally well-written and well explained", "weaknesses": "- Lines 33/34: we make *no assumptions* regarding the loss function other than having lower and upper bounds\n    - This is not true; the results rely on relative weak convexity of $f$ wrt $\\mu$, and moreover any assumptions made about $\\mu$ \n    will imply certain restrictions on the curvature of $f$ as well. For instance, the stationarity measure is related to \n    the more traditional one for $\\mu$ with Lipschitz gradients, and hence the results are not necessarily \n    meaningful unless $f$ is RWC wrt some smooth $\\mu$, which implicitly limits what $f$ can be\n    - The approach critically relies on the approximations to the stochastic problems, and choosing a valid approximation constant \n    seems to require prior knowledge of $\\max_C \\ell(C)$ and $\\min_C \\ell(C)$. So we must not only assume that $\\ell$ is bounded, but \n    that the bounds are *known* to the user. This is arguably a very strong assumption to make.\n\n- The experiments feature an interesting *application* of the subset selection problem, but given that this paper is about a new approach to a rather fundamental problem, it would have been more illuminating to see results in more standard subset selection problems such as Vertex Cover and comparisons against existing methods. As it stands, the experiments currently in the paper do not shed any light on how well the new approach works for subset selection, because it's not clear to me how important subset selection really is for transfer learning.\n\n- Theorem 7.3 seems unnecessary --- it appears to just be a straight-forward application of Azuma-Hoeffding. \n\n- There is a long history of using a smoothing of the loss function to relax convexity assumptions; perhaps there should be more discussion on this front in the related work\n\n### Minor Details\n- Line 199: This is a very atypical way to write the mirror descent update; is there\n  a reason the update is not written in terms of the Bregman divergence?\n\n- It looks like you might be using \\Cref in some of your appendix names, you might want to consider using something more based on the standard \\ref call, like \"Proofs of Section~ \\ ref{sec:gradientestimators}\", in these places. This will still put an in-text reference with hyperlink, but now the index in the pdf table of contents will correctly say \"Proofs of Section 7\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XLfABu8FoL", "reviewer": "Reviewer_eDJC", "summary": "The authors address the subset selection problem, a fundamental task in machine learning and computer science, by introducing a stochastic formulation for selecting minimum-cost subsets in a black-box setting. Here, only the metric value for a subset is accessible, limiting the information available for the selection process. The authors propose a two-stage framework where subset selection is decoupled into an outer selection component and an inner cost evaluation component. The paper parameterizes the subset distribution using a decision variable, optimized via Stochastic Mirror Descent (SMD). The proposed distributions enable constructive, closed-form, unbiased stochastic gradient formulas with favorable convergence guarantees. The approach is empirically validated on a subset selection task involving layer selection in transfer learning, showcasing both theoretical and practical benefits.", "questions": "1. What are the specific properties of the chosen distributions that make them suitable for subset selection? Are these distributions known to have properties that benefit the subset selection problem? \n\n2. How does the approach handle non-smooth or noisy subset metrics? Would there be significant degradation in performance, or could the method be adapted to handle irregular cost functions?\n\n3. Can you discuss the relative benefits of each stochastic formulation (fixed vs. expected cardinality)? Could the authors discuss specific scenarios where one formulation might be preferable over the other?\n\n4. What potential approximations could make this approach more feasible for large-scale problems? Have the authors considered any approximation techniques, such as subsampling or early stopping, that could reduce the computational cost? Discussing potential modifications or approximations would provide insights into how the method could be adapted for resource-constrained environments.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper presents an approach to the subset selection problem by proposing a stochastic formulation optimized using Stochastic Mirror Descent (SMD). Traditional subset selection techniques often involve deterministic methods, such as greedy algorithms or combinatorial optimization, which lack the flexibility and exploration capabilities needed for high-dimensional or complex subset spaces.\n\n2. They derive closed-form, unbiased stochastic gradients for the subset distribution parameters, which makes the method practical and computationally efficient, especially compared to approaches requiring approximate gradients.\n\n3. Overall, the paper is well-structured and clear, with a logical progression from the problem formulation to the proposed approach, theoretical analysis, and empirical evaluation. \n\n4. The significance of this work is substantial, both theoretically and practically. The subset selection problem is fundamental in machine learning, impacting areas like feature selection, model compression, neural architecture search, and sensor placement. By proposing a stochastic approach, this paper introduces a new direction for subset selection that emphasizes flexibility, exploration, and computational efficiency, which is especially relevant in high-dimensional or complex problem spaces.", "weaknesses": "1. While the layer selection experiment in transfer learning provides an interesting and relevant use case, the empirical evaluation is limited to a single task. Subset selection has a wide range of applications across machine learning, including feature selection, model pruning, sensor placement, and combinatorial optimization. Evaluating the proposed method on a broader range of tasks (as stated above) would strengthen the paper and better demonstrate the method’s generalizability.\n\n2. There are various existing subset selection methods, both deterministic (e.g., greedy algorithms, combinatorial optimization) and stochastic (e.g., Monte Carlo methods, stochastic optimization techniques). The paper lacks a comparison with these methods, which would provide readers with more context on the advantages and disadvantages of the proposed approach and help highlight its unique contributions.\n\n3. The paper does not clearly address how the assumptions of smoothness and favorable structure in the subset metric or cost function impact the theoretical results’ robustness in settings with irregular cost functions. Extending the analysis to cover more realistic conditions or discussing the limitations of these assumptions would provide a clearer understanding of the method’s applicability in diverse practical scenarios.\n\n4. The choice of specific distributions for sampling subsets is not well-justified in the paper. While the authors mention that the distribution parameters are optimized, they do not explain why these particular distributions were chosen or how they might perform compared to other candidate distributions. This choice is critical, as different distributions could have substantial effects on the subset sampling’s efficiency and coverage. A clearer explanation of the rationale behind the selected distributions or a comparison with alternative options would strengthen the argument and clarify the trade-offs involved in the design.\n\n5. Despite the use of stochastic methods, the proposed approach may still involve significant computational costs, particularly when repeatedly sampling and evaluating subsets over many iterations. A more thorough discussion on the scalability, computational complexity, and potential approaches to mitigate these challenges would make the paper more practically relevant for large-scale applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vtlEnTQRC4", "reviewer": "Reviewer_dyZD", "summary": "This paper presents a novel stochastic optimization algorithm to the *subset selection\" problem. The objective of this problem is to minimize a loss on the subsets of an ambient set, and this combinatorial problem is NP-hard.  This paper proposed two stochastic approximations where the optimization variable instead becomes a probabilistic distribution over the subsets.\n\nCompare to previous works in stochastic subset selection problem, this paper only makes little assumption over the loss function, where it is only defined over the discrete subsets and thus not differentiable over some continuous space. \n\nTo approach this setting, the authors defined a stochastic gradient, where the derivative is taken over the parameters determining the probability distributions over the subsets. They further show that this stochastic gradient is 1) unbiased, 2) bounded, and 3) relatively convex (an extended notion of strong convexity). This properties allows the author to apply a stochastic mirror descent algorithm and utilize existing analysis in the literature [1] to prove convergence of their algorithm.\n\nFinally, the authors provided a practical application of their method for the transfer learning problem and provided experimental results.\n\n[1] Zhang S, He N. On the convergence rate of stochastic mirror descent for nonsmooth nonconvex optimization. arXiv preprint arXiv:1806.04781. 2018.", "questions": "- Regarding the stochastic relaxation from $(P)$ to $(P_k)$ and $(P_B)$, do you know the optimality gap of this relaxation? If this has already been studied in the literature, it would nice to have a brief discussion over those results.\n- Just to make sure, would Lemma B.1 imply $L^* = L^*_k = L^*_B$? Also, I don't think this lemma is necessary for the proof of Theorem 5.1/5.2?\n- What is Appendix E for?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, this paper is well-written and offers a neat solution to a difficult problem. While there are not any mind-blowing mathematical techniques and the prior literature [1] does quite a bit of heavy lifting, the proposed method has sufficient contribution to warrant a publication.\n\nThe biggest strength of this paper is the writing. The authors does a good job at guiding the reader through the various steps of their method and their claims are backed-up by technical results. There are no \"fluff\" in the paper in the sense that all results contribute meaningfully to the overall picture.\n\nI checked most of the proofs and they all seem to be correct up to constants or typos. I did not check Lemma 7.1 or 7.2, but I presume some variation of those claims must be true. The other proofs all look in good shape (in a first pass), and the techniques generally follow established literature so I don't expect there is any hidden \"surprises\".\n\nMy overall assessment is that paper is not ground-breaking but does the job. It solves a relevant problem with a well-constructed method. While the lack of mathematical novelty stops it from being a great paper, it should solidly pass the threshold for acceptance.", "weaknesses": "While the overall quality of the paper is good, there are a few places where the author's intentions could be better clarified,\n\n- The formulation of the mirror descent step in line 4 of Algorithm 1 is non-standard and I suggest the author to add a few lines of math to show how one can arrive at this line from a more standard form of mirror descent.\n\n- Statements of Theorem 5.1 and 5.2 are rather difficult to parse. Initially, it wasn't immediately clear to me if smaller optimality gap $\\tau$ implies smaller $c^*$. The author should some discussion to guide the readers on how to better interpret the equation on on 291, and a plot could be nice here. Also, in the proof, line 802 suffers a similar issue and more detailed steps should be provided.\n\n- Section 6 claimed the stochastic gradient is inspired by the REINFORCE RL algorithm. But I cannot see this connection at all. I think it would be better if the authors just directly offer the intuition through a few lines of math instead of wrapping the logic around a gray box.\n\n- The experimental setup with transfer learning may be unfamiliar for many readers and a more gentle introduction to this problem is necessary. I personally had to Google this because the paragraph at the bottom of page 8 is on the shorter side. Also, I am still confused by the role of HPO. The authors should more explicitly state the objectives of these experiments and describe the terminologies for a broader audience.\n\n- The numerical results in Table 1 does not show a significant advantage for the author's approach. But I don't view this as essential to the overall quality of the paper.\n\nMinor errors/typos:\n- in Theorem 7.3, should be $(u-l)^2$.\n- line 281, should be relaxation in $(P_k^c)$.\n- line 673: should be Lemma B.1.1\n- line 837, repeated $\\in$ symbol", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)