Preprint.
A DVANCING O PTIMAL S UBSET O RACLE VIA L EARNING R ELAXATION OF N EURAL S ET F UNCTIONS Yongquan Shi1 , Zijing Ou2 , Shiping Wang1 , Yatao Bian3 Fuzhou University, 2 Imperial College London, 3 National University of Singapore [email protected] 1
arXiv:2607.11555v1 [cs.LG] 13 Jul 2026
A BSTRACT Learning neural set functions is pivotal to a wide range of important applications, including compound selection in AI-driven drug discovery and product recommendation. Recent work has introduced optimal subset oracles to implicitly learn set functions under practical weakly supervised settings, where model parameters are optimized through mean-field variational inference. However, these frameworks rely on Monte Carlo sampling to estimate gradients of the evidence lower bound when updating the variational distribution. Repeated sampling across iterations incurs substantial computational overhead, while the resulting stochasticity can destabilize the optimization trajectory. In this work, we reinterpret the evidence lower bound as a continuous relaxation of the set function and learn a surrogate objective that replaces sampling-based ELBO gradient estimation during variational optimization. The learned surrogate provides stable and efficient gradients throughout the continuous domain, thereby reducing computational overhead and accelerating inference. Furthermore, we establish an approximation guarantee for the proposed framework under submodular maximization and characterize its connection to variational free energy. Experiments on a variety of real-world tasks demonstrate consistent improvements over existing baselines.
1
I NTRODUCTION
Set-value prediction has a wide range of applications in real-world scenarios and plays a crucial role in many tasks. For example, recommendation systems select products that are likely to interest a user (Coppolillo et al., 2024), anomaly detection identifies outliers from the majority of observations (Zhang et al., 2020), and AI-driven drug discovery prioritizes promising compounds from large candidate databases (Gimeno et al., 2019). These tasks require explicitly or implicitly learning a set function (Rezatofighi et al., 2017; Zaheer et al., 2017) that assigns a utility value to each candidate subset, with more desirable subsets receiving higher values. More formally, our objective is to select an optimal subset S ∗ from a given large ground set V , such that it attains the highest utility value among all candidate subsets according to a set function Fθ (S; V ) parameterized by θ. This process can be understood as optimizing the following criteria: S ∗ = arg max Fθ (S; V ).
(1)
S∈2V N
A direct approach is to learn Fθ (S; V ) in a supervised manner from tuples (Vi , Si , Ui )i=1 , where Ui denotes the utility of the candidate subset Si , a setting commonly referred to as a function-value (FV) oracle (Balcan & Harvey, 2018). However, this training paradigm is often prohibitively expensive, as it requires collecting utility annotations for a sufficiently large and diverse set of candidate subsets (Ou et al., 2022). To overcome this limitation, an alternative approach to optimizing objective (1) is to implicitly learn the set function from a probabilistic perspective (Tschiatschek et al., 2018). This approach estimates the parameter θ in a supervised manner using pairs {(Vi , Si∗ )}N i=1 , where Si∗ denotes the optimal subset corresponding to Vi , serving as an optimal subset (OS) oracle. With limited data sampled from the underlying distribution P(S, V ), the OS oracle learns latent patterns PN by maximizing the empirical log-likelihood i=1 log pθ (Si∗ |Vi ) over the observed data. Compared with the FV oracle, the OS oracle is generally more practical, as it avoids explicit utility labeling for candidate subsets, thereby reducing annotation costs in practice. 1
Preprint.
Figure 1: (left) Optimization trajectories over the same number of steps on a black-box objective: Monte Carlo sampling (blue) introduces stochastic noise, while surrogate-based optimization (red) converges faster and more smoothly. (middle) Comparison of different baselines on the CelebA dataset in terms of performance, per-iteration training time, and model size. (right) Time required by different baselines to reach the same validation performance on the CelebA dataset. Here, “− R” denotes variants that replace Monte Carlo sampling with a learned set function relaxation. Although the OS oracle is conceptually appealing, the distribution p(S|V ) is generally intractable in practice. Consequently, most existing OS oracles adopt a variational inference framework (Blei et al., 2017), in which a variational distribution qψ (S|V ) is fitted to approximate the target distribution. Here, ψ ∈ [0, 1]|V | parameterizes |V | independent Bernoulli variables, with ψi representing the probability that element i is included in the predicted subset S. Within this framework, optimizing the variational distribution requires maximizing the evidence lower bound (ELBO), which is equivalent to minimizing the divergence between the two distributions. Unfortunately, exact evaluation of the ELBO gradient is computationally prohibitive, as it involves an exponential summation over all possible subsets S. When the ground set V becomes moderately large (e.g., |V | ≥ 30), enumerating the entire 2|V | -subset space becomes intractable. Existing approaches attempt to address this challenge through the Monte Carlo (MC) samplingbased approximation (Xie et al., 2024a;b; Özcan et al., 2025); however, this strategy comes with limitations. Accompanied by the unrolled multi-step gradient ascent, extensive sampling incurs substantial computational overhead. Meanwhile, the induced stochasticity destabilizes the optimization trajectory, making the model prone to suboptimal solutions. In practice, the ELBO can be viewed as a relaxation of the underlying set function. Motivated by this, our core idea is to learn a differentiable surrogate objective that replaces sampling-based ELBO gradient estimation during variational optimization. In this paper, we propose to learn a Relaxation of Set functions (ReSet) within optimal subset oracles, thereby replacing MC–based approximations with a flexible and learnable paradigm. By leveraging the learned relaxation, ReSet substantially reduces optimization overhead and enables faster updates of the variational distribution. We provide illustrative comparisons in Figure 1 to demonstrate this. This optimization process can be interpreted as a greedy strategy for continuous submodular maximization, and we derive an approximation ratio that provides a theoretical lower-bound guarantee. Furthermore, we explore the intrinsic relationship between the ELBO and variational free energy, revealing that our method offers advantages beyond merely optimizing a fixed objective by adaptively learning the trade-off between expected energy and entropy. Our main contributions can be summarized as follows: • We propose to learn a set function relaxation that replaces sampling-based ELBO gradient estimation, enabling efficient and stable learning of neural set functions. • We guarantee that the proposed method converges to a stationary point and derive a constant approximation ratio for the solution, ensuring theoretical soundness. • We revisit the variational inference process through the lens of free energy, demonstrating that our method is able to learn an adaptive trade-off between expected energy and entropy. • Experimental results indicate that the proposed method consistently outperforms baseline methods across diverse real-world tasks. 2
Preprint.
2
P RELIMINARIES
We begin by introducing the optimization framework for optimal subset oracles. From a probabilistic perspective, neural set function learning under the OS oracle can be formulated as a maximum likelihood problem (Stigler, 1990): arg max EP(V,S) [log pθ (S|V )], θ
s. t. pθ (S|V ) ∝ exp(Fθ (S; V )), ∀S ∈ 2V ,
(2)
where pθ (S|V ) denotes the target set distribution we aim to identify. By treating the utility function Fθ (S; V ) as the negative energy, the set distribution can be formulated via an energy-based model: pθ (S|V ) =
exp(Fθ (S; V )) , Z
Z=
X
exp(Fθ (S ′ ; V )).
(3)
S ′ ⊆V
Such formulation follows the principle of maximum entropy (Guiasu & Shenitzer, 1985), thereby embodying the least prior assumptions. This endows the model with the ability to maintain the least possible bias toward unknown information, consistent with the notion of non-informative priors in Bayesian modeling (Jeffreys, 1946). Although this approach is conceptually attractive, learning the parameters θ is notoriously difficult. Following Ou et al. (2022), the training of energy-based model is cast in a variational inference framework, where a variational distribution q(ψ) is introduced as an approximation to the set mass function pθ (S|V ), the fitting target can be expressed as1 ψ ∗ = arg min D[q(ψ)||pθ (S|V )],
(4)
ψ
where D(·||·) is a discrepancy measure between two distributions, ψ ∈ [0, 1]|V | represents the probability that each element s ∈ V is included in the optimal subset S ∗ . By achieving this objective, the variational distribution q(ψ) can be regarded as a function of θ, which allows the parameter θ to be effectively learned in a straightforward manner through the cross-entropy loss: L(θ; ψ ∗ ) = EP(V,S) [− log q(ψ ∗ )] X X ≈ − log ψi∗ − log(1 − ψi∗ ). i∈S ∗
(5)
i∈V \S ∗
In this way, the entire framework can be trained in a cooperative learning fashion (Xie et al., 2018): with θ fixed, the variational distribution q(ψ) is optimized to approximate the energy model pθ (S|V ) via (4); then, θ is updated through gradient flow on the cross-entropy loss in (5).
3
M ETHOD
Although Section 2 outlines a feasible framework for learning set functions under energy-based modeling, the resulting optimization remains nontrivial. In this section, we show how this challenge can be effectively addressed through a straightforward end-to-end learnable paradigm. 3.1
P ROBLEM S ETUP
In order to optimize objective (4), we need to specify both the divergence measure D(·||·) and the form of the variational distribution q(ψ), while ensuring that ψ ∗ remains differentiable with respect to θ. A common choice is to define Q the variational distribution q(ψ) as a product of |V | Q independent Bernoulli variables, i.e., q(ψ) = i∈S ψi i∈V \S (1 − ψi ), which implies that q(ψ) is fully factorized and can be regarded as a mean-field approximation of pθ (S|V ). By further choosing the Kullback-Leibler divergence as the measure, the objective (4) can be reformulated as maximizing the evidence lower bound (ELBO): Fθ min KL(q(ψ)||pθ (S|V )) ⇐⇒ max fmt (ψ) + H(q(ψ)), ψ ψ | {z } ELBO
1
We use the shorthand q(ψ) to denote qψ (S|V ) here.
3
(6)
Preprint.
Fθ (ψ) denotes the multilinear extension of Fθ (S) (Calinescu et al., 2007), defined as where fmt X Y Y Fθ (ψ) := fmt Fθ (S) ψi (1 − ψi ), i∈S
S⊆V
(7)
i∈S /
and H(·) is the entropy of a given distribution. To maximize the ELBO in (6), a straightforward Fθ (ψ))2 , approach is to apply fixed-point iteration, where the stationary condition is ψi = σ(∇ψi fmt here σ(·) denotes the sigmoid function. In this way, ψ can be updated via the following procedure: ψ (0) ← Initialize in [0, 1]|V | ,
(8)
Fθ (ψ (k−1) )))−1 , ψ (k) ← (1 + exp(−∇ψ(k−1) fmt
(9)
∗
ψ ←ψ
(K)
,
(10)
which was referred to as differentiable mean-field variational inference by Ou et al. (2022). Fθ (ψ) entails The crux of the matter is that computing the derivative of the multilinear extension ∇fmt summing over all possible subsets, which is computationally intractable. Existing methods (Xie et al., 2024a;b; Özcan et al., 2025) approximate using Monte Carlo (MC) sampling3 : Fθ ∇ψi fmt (ψ) ≈
M i 1 Xh Fθ (S (m) + i) − Fθ (S (m) ) , M m=1
(11)
where S (m) is the m-th independent sample from q(S; (ψ|ψi ← 0)), and S (m) +i denotes the union of the set S (m) ∪ i. However, the reliance on extensive sampling incurs substantial computational overhead, and the induced stochasticity can further destabilize the optimization trajectory. 3.2
L EARNING A DAPTIVE S ET F UNCTION R ELAXATION
In practice, when optimizing the ELBO with respect to θ while holding ψ fixed, the entropy term H(·) in Eq. (6) depends only on ψ and is therefore constant with respect to θ. It can thus be omitted without affecting the optimization or the gradients with respect to θ. Consequently, we can further treat the ELBO as a multilinear extension of the set function Fθ (S; V ). Motivated by this observation, our core idea is to learn a differentiable surrogate objective gθ : [0, 1]|V | → R+ as a relaxation of the set function, replacing sampling-based ELBO gradient estimation during variational optimization.4 The surrogate enables stable gradient computation throughout the continuous domain without requiring repeated MC estimation at each optimization step. Notably, constructing this surrogate function is easy to implement in practice. According to Karalias et al. (2022), the extensions of set function, including gθ (·) here, can be modeled using simple neural networks. Consequently, the variational parameter ψ can be effectively optimized through gradient ascent ψ (k+1) = ψ (k) + α∇ψ(k) gθ (ψ (k) ),
(12)
where α is the step size. Subsequently, the parameters θ can be updated through the marginal-based loss (Domke, 2013) defined in Equation (5). Note that, in this case, the loss is propagated through the entire optimization process, which involves computing second-order derivatives (i.e., gradients of gradients). These derivatives can be efficiently evaluated using Hessian–vector products, which scale linearly with the model size (Dagréou et al., 2024), similar to standard first-order backpropagation in feed-forward architectures. A brief description of our method is given in Algorithm 1, with more comprehensive procedure is detailed in Appendix D.2. In fact, this framework can be viewed as a continuous submodular maximization problem: max gθ (ψ),
s.t.
ψ
ψ ∈ X,
where the constraint set X is a box constraint, forming a standard compact convex set n o X = ψ ∈ R|V | | 0 ≤ ψi ≤ 1, ∀i ∈ {1, . . . , |V |} . 2
(13)
(14)
The detailed derivation of this stationary condition is provided in Appendix B.1. Further details on gradient approximation via Monte Carlo sampling are provided in Appendix B.2. 4 Note that we do not introduce an additional neural network here. The original OS framework already uses a network to approximate Fθ for MC-based gradient estimation, while our method directly learns the relaxation. 3
4
Preprint.
Algorithm 1 Variational Optimization with ReSet 1: Initialize the variational parameter via the auxiliary recognition network
ψ (0) ← RecNetϕ (V ) 2: for k = 0 to K − 1 do 3: Compute the ELBO surrogate gθ (ψ (k) ) with Eq. (24) 4: Perform gradient ascent on ψ using Eq. (22) ψ (k+1) ← ψ (k) + α∇ψ gθ (ψ (k) ) + ϵk 5: end for 6: Set ψ ∗ ← ψ (K)
The gradient ∇gθ (·) of the differentiable surrogate objective provides continuous surrogate marginal-gain signals, indicating which elements (probability components) should be increased to improve the objective. Through projected gradient ascent in (12), the model selects a locally optimal direction at each iteration, which globally emulating the continuous greedy algorithm (Calinescu et al., 2007). The learned surrogate function is typically a weakly DR-submodular function, as enforcing exact submodularity is difficult with neural parameterizations. In this case, the proposed framework admits the following guarantees (Pedramfar et al., 2024; Hassani et al., 2017): Theorem 3.1. Suppose the relaxation surrogate function gθ (·) : [0, 1]|V | → R+ satisfies γ-weakly submodularity.5 Then the following result holds: 1) For any stationary point ψ ∗ ∈ [0, 1]|V | , we have gθ (ψ ∗ ) ≥ γe−γ gθ (ψOPT ),
(15)
where gθ (ψOPT ) = maxψ gθ (ψ) is the globally optimum. 2) If gθ (·) is L-smooth6 , gradient ascent with step size α < L1 converges to a stationary point. Above theorem indicates that our learning framework is theoretically sound. By choosing smooth activation functions such as sigmoid, we ensure that the optimization process will converge to a stationary point. More importantly, the stationary point to which it converges has a constant approximation ratio with a strict guarantee (at least γe−γ of the global optimal value). As the final solution ψ ∗ obtained is typically fractional, we can obtain a high-quality discrete solution S ∗ using rounding techniques, which can be interpreted as Pipage rounding (Calinescu et al., 2007), ensuring that the expected utility matches the fractional solution: Fθ ES ∗ ∼R(ψ∗ ) [Fθ (S ∗ )] = fmt (ψ ∗ ).
(16)
After sufficient training, the learned surrogate and the underlying multilinear objective are expected to yield aligned assessments of solution quality.7 Therefore, we can further derive a lower-bound guarantee for the resulting discrete solution: ES ∗ ∼R(ψ∗ ) [Fθ (S ∗ )] ≳ γe−γ Fθ (SOPT ). 3.3
(17)
C ONNECTION TO VARIATIONAL F REE E NERGY
The energy (utility) function in (3) evaluates the plausibility of a specific state, which inherently reflects a localized perspective. In contrast, the Helmholtz free energy provides a global measure, which is defined as F = −kB T log Z, (18) where kB is the Boltzmann constant, and T denotes the temperature. It is worth noting that Z here is the partition function from Equation (3). It aggregates all possible states of the system, thereby encapsulating the model’s energy distribution across the entire data space. Within the variational 5 This assumption is not overly restrictive, as γ ∈ [0, 1] continuously characterizes the degree of submodularity: γ = 1 recovers exact DR-submodularity, whereas γ = 0 imposes no submodularity requirement. 6 L-smoothness is a standard regularity assumption widely adopted in first-order optimization and continuous DR-submodular maximization. 7 See Appendix A.3 for a detailed formalization and derivation.
5
Preprint.
framework, for a given ground set V and energy function Eθ (·; V ), the free energy in (18) can be reformulated as X Fθ (ψ) = ES∼q(ψ) [Eθ (S; V )] + kB T [ψi log ψi + (1 − ψi ) log(1 − ψi )] . (19) i∈V
Here, the first term denotes the expected energy under the mean-field variational distribution q(ψ), while the second term accounts for its entropy. Here, kB is the Boltzmann constant, T denotes the temperature, and their product kB T controls the trade-off between expected energy and entropy. Proposition 3.2 (Variational Free Energy). Consider a probabilistic model p(x, z) defined by an energy function E(x, z), together with a variational distribution q(z|x) that approximates the true posterior p(z|x), the variational free energy associated admits a decomposition into an expected energy term and an entropy term: F = ⟨E⟩q − Hq , (20) where ⟨E⟩q = Eq(z|x) [− log pθ (x, z)] denotes the expected energy under the variational distribution q(z|x), and Hq = Eq(z|x) [− log q(z|x)] corresponds to the entropy. Notably, the set utility function naturally plays the role of negative energy here, i.e., Fθ (·) = −Eθ (·), which implies the expected energy term in (19) can be interpreted as the negative counterpart of the Fθ multilinear extension fmt in (6). Consequently, by setting kB T = 1, the variational free energy becomes equivalent to the negative ELBO. Corollary 3.3. Consider a probabilistic model with joint distribution p(x, z) and a variational distribution q(z|x) approximating the true posterior p(z|x). When setting kB T = 1, the variational free energy is equal to the negative evidence lower bound, i.e., F = −ELBO. Therefore, maximizing ELBO in (6) is equivalent to minimizing the variational free energy Fθ (ψ): Eθ max ELBO ⇐⇒ min fmt (ψ) − H(q(ψ)) = Fθ (ψ). ψ
ψ
(21)
Conceptually, this equivalence indicates that variational inference can be interpreted as an energy minimization process: the model seeks balance between fitting the observed data and maintaining uncertainty. By minimizing Fθ (ψ), the model assigns low energy to real samples while encouraging the variational distribution to maintain high entropy, thereby preventing mode collapse. More importantly, this implies that our method offers advantages beyond the objective in (6). In contrast to directly optimizing the ELBO at a fixed balancing relation (temperature), we can directly fit the free energy in (19) to learn an adaptive temperature T , thereby allowing the model to identify an adaptive trade-off between expected energy and entropy.
4
N ETWORK A RCHITECTURE
Although we have established theoretical guarantees above, realizing a smooth and well-behaved energy landscape remains notoriously difficult in practice (Gladstone et al., 2025), as real-world problems often involve high-dimensional objective surfaces. To address this, we leverage Langevin Dynamics (Du & Mordatch, 2019) by injecting a stochastic noise term, thereby promoting exploration of the energy landscape: ψ (k+1) = ψ (k) + α∇ψ gθ (ψ (k) ) + ϵk , ϵk ∼ N (0, τ ), (22) where τ is the magnitude of the noise ϵ. Besides this, to ensure that the optimization trajectory of the predictive distribution can adapt to diverse problem settings, we randomize the gradient ascent step size α, which improves generalization. Taken together, these improvements equip the model with a stronger capacity to navigate complex and rugged objective landscapes in real-world tasks. In addition, to further accelerate inference and improve scalability, we introduce an auxiliary recognition network8 that amortizes inference by providing a reasonable initialization of the variational distribution q(ψ) for a given set V before optimizing the objective in (6): ψ (0) = RecNetϕ (V ). (23) Furthermore, modeling the surrogate gθ requires permutation invariance, which cannot be achieved by classic feed-forward networks. We adopt the DeepSets architecture (Zaheer et al., 2017), which overcomes this limitation through the following proposition: 8
Following the setup of Ou et al. (2022); see Appendix B.3 for details.
6
Preprint.
Proposition 4.1. All permutation invariant set functions can be decomposed in the form f (S) = P ρ si ∈S κ(si ) , for suitable transformations κ and ρ. Following the DeepSets framework, we further construct permutation-invariant sufficient representations9 for both the superset V and the subset S by employing the soft P assignment ψ. These representations can be approximated by empirical estimations R(si ) = si ∈S δ(si ) Bloem-Reddy & Teh (2020a), where δ(si ) represents an atom of unit mass located at si , such as one-hot embeddings. Hence, the neural network construction can be outlined as follows: X |V | |V | X gθ (ψ) = σ θ1 ζ1 (ψi si ) + θ2 ζ2 ((1 − ψi )si ) . (24) i=1
i=1
Here, the feed-forward modules θ1 and θ2 are accompanied by a non-linear activation layer denoted by σ(·). Due to the limitations of the mean-field assumption, the variational distribution q(ψ) is insufficient for capturing interactions among elements within the input set. To remedy this, we design the architecture of function ζ as the multi-head attention (Vaswani et al., 2017), enabling the capture of interactions among elements. The overall architecture of the proposed network is illustrated in Appendix D.1. Leveraging this network architecture, we show in the following section that our proposed framework empirically outperforms the baselines.
5
E MPIRICAL S TUDIES
We evaluate the proposed method on a diverse range of tasks, including product recommendation, set anomaly detection, compound selection, and synthetic experiments. All experiments are conducted with five different random seeds, and we report the corresponding mean performance along with standard deviation. Detailed descriptions of the model architectures and training configurations are provided in Appendix E. Additional evaluations including ablation studies, hyperparameter sensitivity analyses, and experiments on varying ground set sizes are provided in Appendix F. Evaluation Metric. The goal of optimal subset selection is to identify an optimal subset S ∗ from a given ground set V . In experiments, we adopt the mean Jaccard coefficient (MJC) as the evaluation metric to assess the performance of different methods. For each ground set, the Jaccard coefficient |S ′ ∩S ∗ | ′ is defined as JC(S ′ , S ∗ ) = |S ′ ∪S ∗ | , which evaluates the similarity between the predicted subset S ∗ and the ground-truth optimal subset S . Subsequently, the MJC is obtained by averaging the Jaccard coefficients across all test instances. Baselines. Our method is designed as a plug-in module that can be integrated into existing optimal subset oracle frameworks. The resulting variants are denoted by the suffix “-R”. We evaluate these variants against four standard baselines: Random, Probabilistic Greedy Model (PGM) (Tschiatschek et al., 2018), DeepSet (Zaheer et al., 2017) and Set Transformer (Lee et al., 2019); as well as two optimal subset oracles: EquiVSet (Ou et al., 2022) and INSET (Xie et al., 2024a). Detailed descriptions of these baselines are provided in Appendix E.3. Table 1: Results on the synthetic datasets. Method
Two Moons
Gaussian Mixture
Random PGM DeepSet Set Transformer
5.5 36.0 ± 2.0 47.2 ± 0.3 57.4 ± 0.2
5.5 43.8 ± 0.9 44.6 ± 0.2 90.5 ± 0.2
EquiVSet EquiVSet-R
58.7 ± 0.2 99.2 ± 0.3
90.9 ± 0.2 91.0 ± 0.1
INSET INSET-R
59.0 ± 0.3 99.7 ± 0.1
90.9 ± 0.2 90.9 ± 0.1
(a) EquiVSet
(b) EquiVSet-R
Figure 2: Visualization on Two-Moons.
Synthetic Experiments. We first evaluate the methods on two synthetic datasets: the two-moons dataset (Pedregosa et al., 2011) with additional noise of variance σ 2 = 1, and the mixture of Gaus9 Following Xie et al. (2024a), we incorporate background information from V into the set function; see Appendix C for details.
7
Preprint.
sians represented as 21 N (µ0 , Σ) + 12 N (µ1 , Σ), where µ0 = [ √12 , √12 ]⊤ , µ1 = −µ0 , Σ = 14 I. Taking the Gaussian Mixture dataset as an example, the data are generated according to the following procedure: i) Initially, we sample an index b ∼ Bernoulli( 12 ); ii) we then sample 10 points from the Gaussian distribution N (µb , Σ) to construct the ground-truth subset S ∗ ; iii) Finally, we sample additional 90 points from N (µ1−b , Σ) to construct V \S ∗ . For both datasets, we generate 1,000 samples for training, validation, and test, respectively. From the quantitative results in Table 1, it can be observed that the proposed method achieves a significant improvement of nearly 70% over the baseline on the more challenging Two Moons dataset. In addition, we provide visualization results of the model predictions after training for 50 epochs in Figure 2. The green dots represent correct model predictions, the red crosses are incorrect model predictions, and the yellow triangles represent data points in subset oracle S ∗ that are missed by the model. One can see that EquiVSet tends to misclassify more points at the intersection, while also missing some in-distribution points. In comparison, EquiVSet-R learns a more cohesive and clearer distribution boundary, while correctly identifying all the difficult samples in the overlapping region. Table 2: Experimental results (%) on product recommendation across categories in the Amazon dataset, where “Set-T” denotes Set Transformer. Categories
Random
PGM
DeepSet
Set-T
EquiVSet
EquiVSet-R
INSET
INSET-R
Toys Furniture Gear Carseats Bath Health Diaper Bedding Safety Feeding Apparel Media
8.3 6.5 7.7 6.6 7.6 7.6 8.4 7.9 6.5 9.3 9.0 9.4
44.1 ± 0.4 17.5 ± 0.7 47.1 ± 0.4 23.0 ± 1.0 56.4 ± 0.8 44.9 ± 0.2 58.0 ± 0.9 48.0 ± 0.6 25.0 ± 0.6 56.0 ± 0.8 53.3 ± 0.5 44.1 ± 0.9
42.9 ± 0.5 17.6 ± 0.7 38.1 ± 0.2 21.0 ± 1.0 42.4 ± 0.6 44.8 ± 0.4 45.7 ± 0.5 48.2 ± 0.8 22.1 ± 0.4 43.0 ± 0.2 50.7 ± 0.4 42.0 ± 1.0
62.6 ± 2.0 17.6 ± 0.8 64.7 ± 0.6 22.0 ± 1.0 71.6 ± 0.5 69.0 ± 1.0 78.9 ± 0.5 76.0 ± 2.0 23.4 ± 0.9 75.3 ± 0.6 68.0 ± 2.0 53.0 ± 2.0
68.0 ± 2.0 17.2 ± 0.9 70.0 ± 0.2 21.0 ± 1.0 75.7 ± 0.9 70.0 ± 2.0 83.0 ± 1.0 77.0 ± 1.0 25.0 ± 3.0 81.0 ± 0.7 75.0 ± 1.0 57.0 ± 1.0
77.9 ± 1.0 35.9 ± 1.3 78.7 ± 2.6 26.9 ± 1.0 84.7 ± 0.9 81.4 ± 0.6 90.4 ± 0.6 86.8 ± 1.2 32.1 ± 9.4 87.2 ± 1.4 82.2 ± 0.9 67.0 ± 1.0
76.9 ± 0.5 16.9 ± 5.0 80.8 ± 1.2 23.1 ± 3.4 86.2 ± 0.5 81.2 ± 0.5 88.0 ± 0.7 85.7 ± 1.0 23.8 ± 1.5 88.2 ± 1.0 83.7 ± 0.3 62.0 ± 2.3
79.0 ± 0.6 36.5 ± 1.5 82.1 ± 1.1 30.2 ± 3.9 86.7 ± 0.6 82.2 ± 0.2 90.4 ± 0.5 87.6 ± 1.7 38.3 ± 4.0 88.6 ± 0.5 83.3 ± 1.3 68.1 ± 1.1
Product Recommendation. We use the Amazon baby registry dataset (Gillenwater et al., 2014) as a benchmark for the product recommendation task, which contains multiple categories of sub-products selected by different customers. Each product in the dataset is registered under a specific category, such as “Diaper” or “Feeding”, and is represented by a brief textual description. In practice, each product is represented by a 768-dimensional embedding obtained from a pre-trained BERT model (Lee & Toutanova, 2018). Following Ou et al. (2022), we partition the data into training, validation, and test sets using a 1:1:1 ratio. The objective is to identify the subset of products that a customer finds most relevant within a candidate set of size 30. The performance of all methods across different categories is reported in Table 2. Overall, our approach consistently surpasses the two optimal subset oracle baselines. Notably, the performance improvements are more evident on challenging categories, such as Furniture, Carseats, Safety, and Media. Set Anomaly Detection. Experiments on set anomaly detection are conducted using four widely adopted benchmarks: Double MNIST (Sun, 2019), F-MNIST (Xiao et al., 2017), CelebA (Liu et al., 2015b), and CIFAR-10 (Krizhevsky & Hinton, 2009).10 According to the protocol of Zaheer et al. (2017); Ou et al. (2022), we split each dataset into 10,000 training, 1,000 validation, and 1,000 test set instances. Taking the Double MNIST dataset as an example, we form the OS oracle S ∗ by sampling n ∈ {2, 3, 4} images of the same class, with the complement V \S ∗ comprising 20 − |S ∗ | images sampled from non-target classes. From Table 3, it is evident that our proposed variants consistently outperform the baseline methods with notable improvements across all datasets.11 For instance, on the Double MNIST dataset, our proposed variant EquiVSet-R achieves an improvement of approximately 83.9% compared to the original framework. 10
Due to the absence of public code and the relative ambiguity in the protocol descriptions of Ou et al. (2022), we provide a reconstructed and well-defined experimental standard for Double MNIST, F-MNIST, and CIFAR-10; further discussions are provided in Appendix E.5. 11 Results for PGM are unavailable as the official implementation has not been released.
8
Preprint.
Table 3: Experimental results (%) on Anomaly Detection and Compound Selection tasks. Double MNIST
Anomaly Detection CelebA F-MNIST
CIFAR10
Compound Selection PDBBind BindingDB
Random PGM DeepSet Set Transformer
9.3 18.5 ± 1.0 22.0 ± 2.6
21.9 48.1 ± 0.6 44.0 ± 0.5 52.7 ± 0.8
20.9 59.7 ± 1.2 69.8 ± 2.9
20.1 29.0 ± 0.3 27.0 ± 1.7
9.9 91.0 ± 1.0 90.1 ± 1.1 91.9 ± 1.5
0.9 69.0 ± 2.0 71.0 ± 2.0 71.5 ± 1.0
EquiVSet EquiVSet-R
38.6 ± 2.5 71.0 ± 1.6
54.9 ± 0.5 58.8 ± 0.6
58.4 ± 1.2 72.6 ± 2.1
30.4 ± 0.7 32.4 ± 2.2
92.4 ± 1.1 93.0 ± 2.6
72.1 ± 0.9 78.3 ± 0.8
INSET INSET-R
46.2 ± 1.1 73.2 ± 1.8
58.0 ± 1.2 59.2 ± 0.2
66.0 ± 2.8 75.9 ± 1.4
32.0 ± 0.9 35.8 ± 1.3
93.5 ± 0.8 96.0 ± 0.6
73.4 ± 1.0 77.8 ± 1.2
Compound Selection in AI-aided Drug Discovery. In the drug discovery process, screening compounds with high bioactivity (Wallach et al., 2015; Li et al., 2021; Ji et al., 2023), structural diversity, and favorable ADME (absorption, distribution, metabolism, and excretion) properties (Gimeno et al., 2019) is a critical step. Virtual screening pipelines typically proceed through multiple sequential filtering stages. For example, researchers first select compounds with high bioactivity, then extract diverse subsets, and finally remove candidates with unfavorable ADME properties. After these steps, a final subset of compounds is obtained. Due to privacy protection policies in pharmaceutical practice, intermediate supervision signals are often difficult to access or prohibitively expensive to obtain. As a result, models are expected to learn the composite screening process in an end-to-end manner, which is well aligned with the characteristics of the OS oracle. Nevertheless, full end-to-end prediction is often more complex and requires additional domain expertise. Here, we focus on the high-bioactivity selection stage, following the experimental setting of Xie et al. (2024a). In the experiments, we evaluate the methods on PDBBind (Liu et al., 2015a) and BindingDB (Liu et al., 2007) datasets using the same high-bioactivity filter. As shown in Table 3, our method consistently improves upon both OS-oracle baselines across the two datasets. The gains are particularly pronounced on the more challenging BindingDB benchmark.
6
R ELATED W ORK
Set function learning is commonly studied through function-value (FV) oracles, which learn to predict the utility of a given set (Wendler et al., 2019; 2021; De & Chakrabarti, 2022). Representative architectures include DeepSets (Zaheer et al., 2017), Set Transformer (Lee et al., 2019), and neural models for submodular functions (Bilmes & Bai, 2017; Dolhansky & Bilmes, 2016; Balcan & Harvey, 2018). However, these methods typically require dense utility supervision, limiting their practicality. Optimal subset (OS) oracles address this issue by learning latent set functions from optimal subsets under an energy-based variational inference framework (Ou et al., 2022), with subsequent extensions incorporating contextual information, hierarchical attention, and implicit differentiation (Xie et al., 2024a;b; Özcan et al., 2025). Nevertheless, existing OS methods use Monte Carlo sampling for gradient estimation, leading to inefficient training. Our work builds on this line by learning a relaxation of neural set functions, while also connecting OS oracle learning with energy minimization and weakly DR-submodular maximization. Further discussion is provided in Appendix G.
7
C ONCLUSION
In this paper, we propose ReSet, which learns a relaxation of neural set functions to advance optimal subset oracles. The central idea of our approach is to replace Monte Carlo sampling over the fixed landscape with a learnable paradigm. By constructing a surrogate of the evidence lower bound, our method enables efficient and stable gradient estimation across the continuous domain, thereby accelerating and improving the variational inference process. We further establish approximation guarantees for the resulting solutions, characterize their convergence to stationary points, and show that the proposed method can adaptively learn the trade-off between expected utility and entropy. As a general plug-in module, our method can be integrated into various optimal subset oracle frameworks, with extensive experiments demonstrating its effectiveness and scalability. 9
Preprint.
R EFERENCES Maria-Florina Balcan and Nicholas JA Harvey. Submodular functions: Learnability, structure, and optimization. SIAM Journal on Computing, 47(3):703–754, 2018. Gantavya Bhatt, Arnav Das, and Jeff Bilmes. Deep submodular peripteral networks, 2024. Jeffrey Bilmes and Wenruo Bai. Deep submodular functions. arXiv preprint arXiv:1701.08939, 2017. David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):859–877, 2017. Benjamin Bloem-Reddy and Yee Whye Teh. Probabilistic symmetries and invariant neural networks. Journal of Machine Learning Research, 21(90):1–61, 2020a. Benjamin Bloem-Reddy and Yee Whye Teh. Probabilistic symmetries and invariant neural networks, 2020b. Gruia Calinescu, Chandra Chekuri, Martin Pál, and Jan Vondrák. Maximizing a submodular set function subject to a matroid constraint. In Integer Programming and Combinatorial Optimization, pp. 182–196, 2007. Miguel A Carreira-Perpinan and Geoffrey Hinton. On contrastive divergence learning. In International workshop on artificial intelligence and statistics, pp. 33–40, 2005. Erica Coppolillo, Giuseppe Manco, and Aristides Gionis. Relevance meets diversity: A user-centric framework for knowledge exploration through recommendations. In Proceedings of the ACM SIGKDD conference on knowledge discovery and data mining, pp. 490–501, 2024. Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. How to compute hessianvector products? In The Third Blogpost Track at ICLR, 2024. Abir De and Soumen Chakrabarti. Neural estimation of submodular functions with applications to differentiable subset selection. In Advances in Neural Information Processing Systems, pp. 19537–19552, 2022. Josip Djolonga and Andreas Krause. Differentiable learning of submodular models. In Advances in Neural Information Processing Systems, 2017. Brian W Dolhansky and Jeff A Bilmes. Deep submodular functions: Definitions and learning. In Advances in Neural Information Processing Systems, volume 29, 2016. Justin Domke. Learning graphical model parameters with approximate marginal inference. IEEE transactions on pattern analysis and machine intelligence, 35(10):2454–2467, 2013. Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in neural information processing systems, 32, 2019. Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Learning iterative reasoning through energy minimization. In International Conference on Machine Learning, pp. 5570–5582, 2022. Yilun Du, Jiayuan Mao, and Joshua B Tenenbaum. Learning iterative reasoning through energy diffusion. In International Conference on Machine Learning, pp. 11764–11776, 2024. Jennifer A Gillenwater, Alex Kulesza, Emily Fox, and Ben Taskar. Expectation-maximization for learning determinantal point processes. Advances in Neural Information Processing Systems, 27, 2014. Aleix Gimeno, Marı́a José Ojeda-Montes, Sarah Tomás-Hernández, Adrià Cereto-Massagué, Raúl Beltrán-Debón, Miquel Mulero, Gerard Pujadas, and Santiago Garcia-Vallvé. The light and dark sides of virtual screening: what is there to know? International journal of molecular sciences, 20 (6):1375, 2019. 10
Preprint.
Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Peixuan Han, Hyeonjeong Ha, Aman Chadha, Yilun Du, Heng Ji, Jundong Li, and Tariq Iqbal. Energy-based transformers are scalable learners and thinkers. arXiv preprint arXiv:2507.02092, 2025. Joseph Gomes, Bharath Ramsundar, Evan N Feinberg, and Vijay S Pande. Atomic convolutional networks for predicting protein-ligand binding affinity. arXiv preprint arXiv:1703.10603, 2017. Silviu Guiasu and Abe Shenitzer. The principle of maximum entropy. The mathematical intelligencer, 7(1):42–48, 1985. Hamed Hassani, Mahdi Soltanolkotabi, and Amin Karbasi. Gradient methods for submodular maximization. In Advances in Neural Information Processing Systems, volume 30, 2017. Jiajun He, Yuanqi Du, Francisco Vargas, Yuanqing Wang, Carla P Gomes, José Miguel HernándezLobato, and Eric Vanden-Eijnden. Feat: Free energy estimators with adaptive transport. arXiv preprint arXiv:2504.11516, 2025. Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005. Harold Jeffreys. An invariant form for the prior probability in estimation problems. Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences, 186(1007):453–461, 1946. Yuanfeng Ji, Lu Zhang, Jiaxiang Wu, Bingzhe Wu, Lanqing Li, Long-Kai Huang, Tingyang Xu, Yu Rong, Jie Ren, Ding Xue, et al. Drugood: Out-of-distribution dataset curator and benchmark for ai-aided drug discovery–a focus on affinity prediction problems with noise annotations. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 8023–8031, 2023. Nikolaos Karalias, Joshua Robinson, Andreas Loukas, and Stefanie Jegelka. Neural set function extensions: Learning with discrete functions in high dimensions. In Advances in Neural Information Processing Systems, volume 35, pp. 15338–15352, 2022. Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. Suraj Kothawade, Jiten Girdhar, Chandrashekhar Lavania, and Rishabh Iyer. Deep submodular networks for extractive data summarization, 2020. A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Handbook of Systemic Autoimmune Diseases, 1(4), 2009. Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, et al. A tutorial on energybased learning. Predicting structured data, 1(0), 2006. JDMCK Lee and K Toutanova. Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 3(8):4171–4186, 2018. Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariant neural networks. In International conference on machine learning, pp. 3744–3753, 2019. Shuangli Li, Jingbo Zhou, Tong Xu, Liang Huang, Fan Wang, Haoyi Xiong, Weili Huang, Dejing Dou, and Hui Xiong. Structure-aware interactive graph neural networks for the prediction of protein-ligand binding affinity. In Proceedings of the 27th ACM SIGKDD conference on knowledge discovery & data mining, pp. 975–985, 2021. Tiqing Liu, Yuhmei Lin, Xin Wen, Robert N Jorissen, and Michael K Gilson. Bindingdb: a webaccessible database of experimentally determined protein–ligand binding affinities. Nucleic acids research, 35(suppl 1):D198–D201, 2007. Zhihai Liu, Yan Li, Li Han, Jie Li, Jie Liu, Zhixiong Zhao, Wei Nie, Yuchen Liu, and Renxiao Wang. Pdb-wide collection of binding data: current status of the pdbbind database. Bioinformatics, 31 (3):405–412, 2015a. 11
Preprint.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pp. 3730–3738, 2015b. Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983, 2016. Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya. On learning sets of symmetric elements, 2020. Zijing Ou, Tingyang Xu, Qinliang Su, Yingzhen Li, Peilin Zhao, and Yatao Bian. Learning neural set functions under the optimal subset oracle. In Advances in Neural Information Processing Systems, pp. 35021–35034, 2022. Gözde Özcan, Chengzhi Shi, and Stratis Ioannidis. Learning set functions with implicit differentiation. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 19777–19785, 2025. Hakime Öztürk, Arzucan Özgür, and Elif Ozkirimli. Deepdta: deep drug–target binding affinity prediction. Bioinformatics, 34(17):i821–i829, 2018. Mohammad Pedramfar, Christopher John Quinn, and Vaneet Aggarwal. A unified approach for maximizing continuous dr-submodular functions, 2024. Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit-learn: Machine learning in python. the Journal of machine Learning research, 12:2825–2830, 2011. S Hamid Rezatofighi, Vijay Kumar Bg, Anton Milan, Ehsan Abbasnejad, Anthony Dick, and Ian Reid. Deepsetnet: Predicting sets with deep neural networks. In International Conference on Computer Vision, pp. 5257–5266, 2017. Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in artificial intelligence, pp. 574–584, 2020. Stephen M Stigler. The history of statistics: The measurement of uncertainty before 1900. Harvard University Press, 1990. Shao-Hua Sun. Multi-digit mnist for few-shot learning. GitHub repository, 2019. Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In Proceedings of the 25th international conference on Machine learning, pp. 1064– 1071, 2008. Sebastian Tschiatschek, Aytunc Sahin, and Andreas Krause. Differentiable submodular maximization. In Proceedings of the International Joint Conference on Artificial Intelligence, pp. 2731– 2738, 2018. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, volume 30, 2017. Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661–1674, 2011. Izhar Wallach, Michael Dzamba, and Abraham Heifets. Atomnet: a deep convolutional neural network for bioactivity prediction in structure-based drug discovery. arXiv preprint arXiv:1510.02855, 2015. Runqian Wang and Yilun Du. Equilibrium matching: Generative modeling with implicit energybased models. arXiv preprint arXiv:2510.02300, 2025. Chris Wendler, Markus Püschel, and Dan Alistarh. Powerset convolutional neural networks. In Advances in Neural Information Processing Systems, 2019. 12
Preprint.
Chris Wendler, Andisheh Amrollahi, Bastian Seifert, Andreas Krause, and Markus Püschel. Learning set functions that are sparse in non-orthogonal fourier bases. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 10283–10292, 2021. Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017. Binghui Xie, Yatao Bian, Yongqiang Chen, Peilin Zhao, Bo Han, Wei Meng, James Cheng, et al. Enhancing neural subset selection: Integrating background information into set representations. In International Conference on Learning Representations, 2024a. Binghui Xie, Yixuan Wang, Yongqiang Chen, Kaiwen Zhou, Yu Li, Wei Meng, and James Cheng. Horse: hierarchical representation for large-scale neural subset selection. In Advances in neural information processing systems, pp. 4852–4877, 2024b. Jianwen Xie, Yang Lu, Ruiqi Gao, and Ying Nian Wu. Cooperative learning of energy-based model and latent variable model via mcmc teaching. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018. Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017. David W Zhang, Gertjan J Burghouts, and Cees GM Snoek. Set prediction without imposing structure as conditional density estimation. In International Conference on Learning Representations, 2020.
13
Preprint.
Appendix for “Advancing Optimal Subset Oracle via Learning Relaxation of Neural Set Functions” C ONTENTS A Proofs and Derivations
15
A.1 Proof of Theorem 3.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
15
A.2 Proof of Corollary 3.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
18
A.3 Derivation of the Discrete Approximation Guarantee . . . . . . . . . . . . . . . . . . . . . .
19
B Details of Differentiable Mean Field Variational Inference in OS Oracle
19
B.1 Derivations of the Fixed Point Iteration . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
19
B.2 Details of Monte Carlo Gradient Estimation . . . . . . . . . . . . . . . . . . . . . . . . . . .
20
B.3 Accelerating Inference with Equivariant Neural Networks . . . . . . . . . . . . . . . . . . .
20
C Incorporate Sufficient Invariant Statistics of Background Information
20
D Detailed Procedure of ReSet
21
D.1 Overall Framework . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
21
D.2 Detailed Pseudo Code . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
22
E Experimental Details
23
E.1 The Architecture of ReSet . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
23
E.2 Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
24
E.3 Baselines . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
24
E.4 Detailed Experimental Settings for Product Recommendation . . . . . . . . . . . . . . . . .
25
E.5 Detailed Experimental Settings for Set Anomaly Detection . . . . . . . . . . . . . . . . . . .
25
E.6 Detailed Experimental Settings for Compound Selection . . . . . . . . . . . . . . . . . . . .
27
F Additional Experiments
28
F.1
Ablation Study . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
28
F.2
Computational Efficiency . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
28
F.3
Comparison with Low-Variance Gradient Estimators . . . . . . . . . . . . . . . . . . . . . .
29
F.4
Sensitivity Analysis of Hyperparameters . . . . . . . . . . . . . . . . . . . . . . . . . . . .
29
F.5
Set Size Transferability Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
30
G Related Work
31
G.1 Set Function Learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
31
G.2 Energy-based Modeling . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
31
G.3 Submodular Function Learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
31
14
Preprint.
A
P ROOFS AND D ERIVATIONS
A.1
P ROOF OF T HEOREM 3.1
Theorem 3.1. Suppose the relaxation surrogate function gθ (·) : [0, 1]|V | → R+ satisfies γ-weakly submodularity. Then the following result holds: 1) For any stationary point ψ ∗ ∈ [0, 1]|V | , we have gθ (ψ ∗ ) ≥ γe−γ gθ (ψOPT ),
(15)
where gθ (ψOPT ) = maxψ gθ (ψ) is the globally optimum. 2) If gθ (·) is L-smooth , gradient ascent with step size α < L1 converges to a stationary point. Proof. Recall that we consider the continuous relaxation X = [0, 1]|V | .
max gθ (ψ), ψ∈X
Here, gθ (·) is a differentiable surrogate of the multilinear extension induced by the discrete set function Fθ (·). Under the assumptions of the theorem, gθ : [0, 1]|V | → R+ is L-smooth and γweakly DR-submodular for some γ ∈ (0, 1]. Specifically, for any x, y ∈ [0, 1]|V | satisfying x ≤ y, ∇gθ (x) ≥ γ∇gθ (y),
(25)
where all vector inequalities are understood coordinate-wise. Let ψOPT ∈ arg max gθ (ψ) ψ∈X
denote a global maximizer. For brevity, let d ≜ |V |, η ≜ 1/T , and define Gt ≜ gθ ψ (t) ,
GOPT ≜ gθ (ψOPT ) .
We first verify that the iterates generated by Eq. (12) remain feasible. The claim holds at t = 0 because ψ (0) = 0. Moreover, since 0 ≤ v (t) ≤ 1 − ψ (t) and η ∈ (0, 1], we have 0 ≤ ψ (t+1) = ψ (t) + ηv (t) ≤ (1 − η)ψ (t) + η1 ≤ 1. Therefore, ψ (t) ∈ [0, 1]d for every t. We next establish two consequences of γ-weak DR-submodularity that will be used below. For any x, y ∈ [0, 1]d satisfying x ≤ y, Eq. (25) implies ⟨∇gθ (x), y − x⟩ ≥ γ (gθ (y) − gθ (x)) .
(26)
Indeed, letting z ≜ y − x ≥ 0 gives Z 1 gθ (y) − gθ (x) =
⟨∇gθ (x + sz), z⟩ ds 0
≤
1 ⟨∇gθ (x), z⟩ , γ
(27)
where the inequality follows from x ≤ x + sz, z ≥ 0, and Eq. (25). In addition, nonnegativity and γ-weak DR-submodularity imply that, for any x, y ∈ [0, 1]d , gθ (x ∨ y) ≥ (1 − γ∥x∥∞ ) gθ (y). 15
(28)
Preprint.
To establish this inequality, let r ≜ ∥x∥∞ . The result is immediate when r = 0. Otherwise, define z ≜ x ∨ y − y ≥ 0. For every s ∈ [0, 1/r], both y + rsz and y + sz belong to [0, 1]d , and y + rsz ≤ y + sz. It follows that Z 1/r ⟨∇gθ (y + rsz), z⟩ ds gθ (x ∨ y) − gθ (y) = r 0
Z 1/r ≥ rγ
⟨∇gθ (y + sz), z⟩ ds h i z = rγ gθ y + − gθ (y) r ≥ −rγgθ (y), 0
(29)
where the last inequality follows from gθ (·) ≥ 0. Rearranging proves Eq. (28). We now apply these inequalities to bound the progress at each iteration. Consider the comparison direction d(t) ≜ ψOPT ∨ ψ (t) − ψ (t) = ψOPT − ψ (t) . (30) +
Since 0 ≤ d(t) ≤ 1 − ψ (t) , the vector d(t) is feasible for the linear optimization problem in Eq. (12). The optimality of v (t) , together with Eqs. (26) and (28), therefore yields D E D E ∇gθ (ψ (t) ), v (t) ≥ ∇gθ (ψ (t) ), d(t) h i ≥ γ gθ ψ (t) ∨ ψOPT − Gt h i ≥ γ 1 − γ∥ψ (t) ∥∞ GOPT − Gt . (31) For every coordinate i, the direction constraint and γ ≤ 1 imply (t)
vi
(t)
≤ 1 − ψi
(t)
≤ 1 − γψi .
Consequently, (t+1)
1 − γψi
(t)
(t)
= 1 − γψi − ηγvi (t) ≥ (1 − ηγ) 1 − γψi .
(32)
Since ψ (0) = 0, induction gives (t)
1 − γψi
≥ (1 − ηγ)t
for every coordinate i. Hence, 1 − γ∥ψ (t) ∥∞ ≥ (1 − ηγ)t . Combining Eqs. (31) and (33), we obtain D E ∇gθ (ψ (t) ), v (t) ≥ γ (1 − ηγ)t GOPT − Gt .
(33)
(34)
By L-smoothness and Eq. (34), D E Lη 2 ∥v (t) ∥22 Gt+1 ≥ Gt + η ∇gθ (ψ (t) ), v (t) − 2 Lη 2 d ≥ (1 − ηγ)Gt + ηγ(1 − ηγ)t GOPT − , 2 where the last inequality also uses ∥v (t) ∥22 ≤ d. 16
(35)
Preprint.
Letting a ≜ 1 − ηγ and recursively applying Eq. (35) yields T −1
T
GT ≥ a G0 + T ηγa
≥ T ηγaT −1 GOPT −
T −1 Lη 2 d X T −1−t a GOPT − 2 t=0
T Lη 2 d , 2
(36)
PT −1
aT −1−t ≤ T . Substituting η = 1/T gives Ld γ T −1 GOPT − . GT ≥ γ 1 − T 2T
where we used G0 ≥ 0, a ∈ [0, 1], and
t=0
(37)
For γ ∈ (0, 1] and T ≥ 2, applying log(1 − u) ≥ −u/(1 − u) for u ∈ [0, 1) gives γ(T − 1) γ ≥− (T − 1) log 1 − ≥ −γ, T T −γ and hence
γ T −1 ≥ e−γ . T Substituting this inequality into Eq. (37) and recalling that d = |V |, we conclude that L|V | gθ ψ (T ) ≥ γe−γ gθ (ψOPT ) − . 2T
1−
(38)
Thus, choosing T ≥ L|V |/(2ε) gives the stated γe−γ approximation guarantee up to an additive error ε. Having established the objective-value guarantee, we now show that the optimization process also enjoys a standard convergence guarantee when instantiated with projected gradient ascent. Specifically, consider 1 ψ (k+1) = ΠX ψ (k) + α∇gθ (ψ (k) ) , (39) 0<α< , L where ΠX (·) denotes the Euclidean projection onto X . The first-order optimality condition of the projection gives D E ψ (k) + α∇gθ (ψ (k) ) − ψ (k+1) , ψ − ψ (k+1) ≤ 0, ∀ψ ∈ X. (40) Setting ψ = ψ (k) in Eq. (40) yields D E 2 1 ψ (k+1) − ψ (k) . ∇gθ (ψ (k) ), ψ (k+1) − ψ (k) ≥ α 2
(41)
Combining Eq. (41) with L-smoothness gives D E gθ (ψ (k+1) ) ≥ gθ (ψ (k) ) + ∇gθ (ψ (k) ), ψ (k+1) − ψ (k) 2 L ψ (k+1) − ψ (k) 2 2 2 1 L (k) ≥ gθ (ψ ) + − ψ (k+1) − ψ (k) α 2 2 2 1 ≥ gθ (ψ (k) ) + ψ (k+1) − ψ (k) , 2α 2
−
(42)
where the last inequality follows from α < 1/L. Thus, the objective values are monotonically nondecreasing. Since gθ is continuous and X is compact, gθ is bounded from above on X . Summing Eq. (42) over k therefore gives ∞ X
ψ (k+1) − ψ (k)
2
< ∞, 2
k=0
17
(43)
Preprint.
and consequently lim
k→∞
ψ (k+1) − ψ (k)
= 0.
(44)
2
To express this result in terms of first-order stationarity, define the projected gradient mapping Gα (ψ) ≜
1 [ΠX (ψ + α∇gθ (ψ)) − ψ] . α
(45)
By Eqs. (39) and (44), lim
k→∞
Gα (ψ (k) )
= 0.
(46)
2
Because X is compact, the sequence {ψ (k) } admits at least one accumulation point. Let ψ̄ be any such point. The continuity of the projected gradient mapping and Eq. (46) imply Gα (ψ̄) = 0. Equivalently, ∇gθ (ψ̄), ψ − ψ̄ ≤ 0,
∀ψ ∈ X.
(47)
Therefore, projected gradient ascent converges to first-order stationarity in the sense that its projected gradient mapping vanishes and every accumulation point is stationary. Together with Eq. (38), this establishes both the approximation and convergence guarantees stated in the theorem. A.2
P ROOF OF C OROLLARY 3.3
Corollary 3.3. Consider a probabilistic model with joint distribution p(x, z) and a variational distribution q(z|x) approximating the true posterior p(z|x). When setting kB T = 1, the variational free energy is equal to the negative evidence lower bound, i.e., F = −ELBO. Proof. From the variational characterization of the Helmholtz free energy established above, when kB T = 1, the variational free energy admits the form F = ⟨E⟩q − Hq = Eq(z|x) [Eθ (x, z)] + Eq(z|x) [log q(z|x)] .
(48)
Recall that the evidence lower bound (ELBO) is defined as ELBO = Eq(z|x) [log pθ (x, z)] − Eq(z|x) [log q(z|x)] .
(49)
Since the joint distribution pθ (x, z) is defined through the energy function as log pθ (x, z) = −Eθ (x, z) − log Zθ , substituting into Eq. (49) yields ELBO = −Eq(z|x) [Eθ (x, z)] − log Zθ − Eq(z|x) [log q(z|x)] . Noting that the Helmholtz free energy satisfies F = − log Zθ when kB T = 1, we obtain ELBO = −F − Eq(z|x) [Eθ (x, z)] + Eq(z|x) [log q(z|x)] . Rearranging terms and using Eq. (48) immediately gives F = −ELBO. In our setting, the set utility function naturally plays the role of negative energy, i.e., Fθ (·) = −Eθ (·). Under the ground-set formulation, the expected energy term in (19) therefore corresponds to the Fθ negative multilinear extension −fmt (ψ). As a consequence, maximizing the ELBO in Eq. (6) is equivalent to minimizing the variational free energy: Eθ max ELBO ⇐⇒ min fmt (ψ) − H(q(ψ)) = min Fθ (ψ). (50) ψ
ψ
ψ
18
Preprint.
A.3
D ERIVATION OF THE D ISCRETE A PPROXIMATION G UARANTEE
In the idealized regime of sufficient training, we assume that the learned surrogate and the original multilinear objective induce aligned normalized assessments of solution quality, such that Fθ fmt (ψ ∗ ) gθ (ψ ∗ ) ≤ εsol , − Fθ (SOPT ) gθ (ψgOPT )
(51)
where εsol ≥ 0 denotes the residual solution-level discrepancy, which is expected to decrease as training progresses. By Theorem 3.1, we have gθ (ψ ∗ ) ≥ γe−γ . gθ (ψgOPT ) Combining this result with the lower bound implied by Eq. (51) yields Fθ fmt (ψ ∗ ) ≥ γe−γ − εsol Fθ (SOPT ).
(52)
We next obtain a discrete solution by rounding S ∗ ∼ R(ψ ∗ ). By the definition of the multilinear extension, Fθ ES ∗ ∼R(ψ∗ ) [Fθ (S ∗ )] = fmt (ψ ∗ ). Combining this identity with Eq. (52) gives ES ∗ ∼R(ψ∗ ) [Fθ (S ∗ )] ≥ γe−γ − εsol Fθ (SOPT ). Therefore, as εsol becomes negligible under sufficient training, the resulting discrete solution admits the following approximate lower-bound guarantee: ES ∗ ∼R(ψ∗ ) [Fθ (S ∗ )] ≳ γe−γ Fθ (SOPT ).
(53)
B
D ETAILS OF D IFFERENTIABLE M EAN F IELD VARIATIONAL I NFERENCE IN OS O RACLE
B.1
D ERIVATIONS OF THE F IXED P OINT I TERATION
Following Ou et al. (2022), we present a detailed derivation of the fixed-point iteration (FPI) associated with Eq. (9): −1 (k+1) Fθ ψi ← 1 + exp −∇ψ(k) fmt (ψ (k) ) . (54) i
We begin by recalling that the objective is to maximize the evidence lower bound (ELBO), max
ψ∈[0,1]|V |
X
Fθ (S)
S⊆V
|
Y
ψi
i∈S
Y
|V | X
ψi log ψi + (1 − ψi ) log(1 − ψi ) .
(55)
i=1
i∈S /
{z
F
(1 − ψi ) − }
|
fmtθ (ψ)
{z
−H(q(S;ψ))
}
The update in Eq. (54) is obtained by enforcing stationarity of the ELBO with respect to coordinate ψi . Specifically, setting the partial derivative to zero yields Fθ Fθ ∇ψi fmt (ψ) + ∇ψi H(q(S; ψ)) = ∇ψi fmt (ψ) + log
1 − ψi = 0. ψi
Solving for ψi leads to Fθ ψi = 1 + exp −∇ψi fmt (ψ)
−1
,
which coincides with the fixed-point update employed in mean-field variational inference. 19
Preprint.
B.2
D ETAILS OF M ONTE C ARLO G RADIENT E STIMATION
According to Eq. (54), the variational parameters can be updated via fixed-point iteration. Nevertheless, a major computational challenge arises from evaluating the gradient of the multilinear extension Fθ fmt (ψ) defined in Eq. (7), as it involves summation over all 2|V | possible subsets. To address this Fθ issue, Ou et al. (2022) showed that the gradient ∇ψ fmt admits an expectation form, which enables efficient estimation through Monte Carlo sampling. Specifically, the partial derivative with respect to ψi can be derived as X Y Y Fθ ∇ψi fmt = ∇ψi Fθ (S) ψi (1 − ψi ) S⊆V
i∈S
i̸∈S
= Eq(S;(ψ|ψi ←1)) [Fθ (S)] − Eq(S;(ψ|ψi ←0)) [Fθ (S)] X Y Y X = Fθ (S) ψj (1 − ψj ′ ) − =
j ′ ̸∈S
S⊆V,i∈S
j∈S\{i}
X
[Fθ (S + i) − Fθ (S)]
S⊆V \{i}
Y j∈S
Fθ (S)
S⊆V \{i}
Y
ψj
Y j∈S
Y
ψj
(1 − ψj ′ )
j ′ ∈S,j ′ ̸=i
(1 − ψj ′ )
j ′ ∈V \S\{i}
= Eq(S;(ψ|ψi ←0)) [Fθ (S + i) − Fθ (S)] .
(56)
Fθ Here, the S + i is denote the set union S ∪ {i}. Based on Eq. (56), the gradient ∇ψi fmt can be approximated using Monte Carlo estimation: i) sample m subsets Sn , n = 1, . . . , m from thePsurrogate distribution q(S; (ψ|ψi ← 0)); ii) approximate the expectation by the average n 1 k=1 [Fθ (Sn + i) − Fθ (Sn )]. This procedure provides a tractable and unbiased estimator of m the gradient required for the fixed-point updates.
B.3
ACCELERATING I NFERENCE WITH E QUIVARIANT N EURAL N ETWORKS
Although Monte Carlo–based gradient estimation makes the optimization of the variational parameters ψ feasible, it is accompanied by a notorious drawback: the computational cost is often prohibitively high. In particular, the DiffMF procedure in Eq. (8) typically involves an inner sampling loop for each data instance, which substantially increases the overall training overhead. To improve scalability on large-scale datasets, prior works (Ou et al., 2022; Xie et al., 2024a) propose to amortize the approximate inference process by introducing an auxiliary recognition network. This network directly predicts the variational parameters ψ of the distribution qϕ (S; ψ), where ϕ denotes the learnable parameters of the recognition model. Crucially, since the inputs are set-valued, the recognition network must satisfy the property of permutation equivariance, as formalized in Proposition 4.1. Specifically, each output coordinate is associated with a particular element of the set, and is invariant to the ordering of elements in S. In our framework, optimizing each data point similarly requires multiple gradient ascent steps as in Eq. (12). Therefore, we likewise employ an auxiliary recognition network to accelerate the inference procedure. Concretely, the equivariant recognition network is defined as ψ = RecNet(V ; ϕ) : 2V → [0, 1]|V | , where each component RecNet i := fi takes the ground set V as input and outputs the corresponding variational parameter, collectively specifying the distribution qϕ (S; ψ).
C
I NCORPORATE S UFFICIENT I NVARIANT S TATISTICS OF BACKGROUND I NFORMATION
Within our framework, we adopt the architecture proposed by Xie et al. (2024a) and implement its relaxed variant, INSET-R. The central idea of INSET is to incorporate contextual information from the ground set V into the set function, yielding a formulation of the form F (V, S), in contrast to EquiVSet, which models the set function solely as F (S). This mechanism is referred to by Xie et al. (2024a) as the fusion of invariant sufficient representations of the subset S and the ground set V . P Specifically, Bloem-Reddy & Teh (2020b) show that the empirical measure MS (S) = si ∈S δ(si ) can be served as a valid invariant sufficient representation, where δ(si ) denotes a unit-mass atom located at si , for example instantiated via one-hot encodings. Furthermore, building on the result of Zaheer et al. (2017), this empirical measure can be efficiently approximated using the aggrega20
Preprint.
P tion form ρ s∈S ϕ(s) . This formulation provides a practical and effective means of realizing invariant sufficient representations in neural architectures. Proposition C.1. If f is a valid permutation invariant function on S, it can be approximated arbi P trarily close in the form of f (S) = ρ s∈S ϕ(s) , for suitable transformations ϕ and ρ. In practice, an encoder ϕ(·) is employed to produce element-wise embeddings, while the mapping ρ(·) may be instantiated by a variety of feedforward architectures, including fully connected layers equipped with nonlinear activations. In a related line of work, Maron et al. (2020) establish that universal of invariant sufficient representations can be achieved via the aggregation P approximation P P form S∈V s∈S ϕ(s), which simplifies to the equivalent expression xj ∈V ϕ(xj ). Therefore, the neural network architecture for set functions injected with superset background information can be expressed as ! !! ni n X X θ(S, V ) = σ θ1 ϕ(xi ) + θ2 ϕ(xj ) . (57) i=1
i=1
Here, the feed-forward modules θ1 and θ2 are accompanied by a non-linear activation layer denoted by σ. As a comparison, we provide the set function architectures for four different optimal subset oracles in Table 4. Table 4: Comparison of set function architectures across four different optimal subset oracles. EquiVSet
EquiVSet-R
INSET
INSET-R
FCS (256, 500, ReLU) FCS (500, 1, −) -
MHAS (256, 4) FCS (256, 512, ReLU) FCS (512, 1, −) -
FCV (256, 500, ReLU) FCV (500, 1, −) FCS (256, 500, ReLU) FCS (500, 1, −) -
MHAV (256, 4) MHAS (256, 4) FCV (256, 512, ReLU) FCV (512, 1, −) FCS (256, 512, ReLU) FCS (512, 1, −)
Framework
𝝍 (𝝍)
Ground Set
MHA Encoder
Subset Info
V
𝝍𝝍(𝝍 (𝝍) )
𝝍 − 𝝍 (𝝍)
MHA Encoder
Background Info
𝝍 (𝝍)
RecNet 𝝍
ELBO Surrogate
Update
𝝍𝝍𝝍(𝝍 (𝝍) ) 𝝍𝝍 (𝝍)
× 𝝍 steps
Output: 𝝍 (𝝍 )
Figure 3: Overall framework of the proposed network architecture.
D
D ETAILED P ROCEDURE OF R E S ET
D.1
OVERALL F RAMEWORK
5
Figure 3 illustrates the overall architecture of ReSet. Given a ground set V , an auxiliary recognition network RecNetϕ first produces an initial variational parameter ψ (0) , which amortizes the inference process. At each optimization step, the current soft assignment ψ (k) and its complement 1 − ψ (k) are used to construct subset-aware and background-aware representations of the input set. These two branches are encoded by multi-head attention modules to capture element-wise dependencies, and are then combined to produce the ELBO surrogate gθ (ψ (k) ). The gradient of this surrogate 21
Preprint.
with respect to ψ (k) is used to iteratively update the variational parameter for K steps, yielding the final prediction ψ (K) . This design enables efficient differentiable optimization without Monte Carlo gradient estimation. D.2
D ETAILED P SEUDO C ODE
We provide the pseudo-code for the training and inference procedures of ReSet in Algorithms 2 and 3, respectively. The training process consists of three key steps: i) generating the initial variational parameters ψ (0) via RecNetϕ ; ii) iteratively updating ψ (k) for K steps by maximizing the ELBO using the surrogate function gθ (·) via gradient ascent; and iii) computing the cross-entropy loss based on the optimized ψ (K) in Eq. (5) and updating the network parameters through backpropagation. During inference, only steps i) and ii) are executed. Notably, we incorporate Langevin dynamics exclusively during training to facilitate exploration, while opting for a deterministic approach during inference to ensure stability and finer convergence. In our experiments, the number of iterative steps K is set to 1. Algorithm 2 ReSet (Training) Require: {Vi , Si∗ }N i=1 : training dataset; η: learning rate; K: number of iterations; α: step size for gradient ascent Ensure: Optimal parameters θ∗ , ϕ∗ 1: θ, ϕ ← Initialize parameter 2: while θ, ϕ not converged do 3: Sample training data point (V, S ∗ ) ∼ {Vi , Si∗ }N i=1 4: Obtain the initial variational parameter ψ (0) via the auxiliary recognition network ψ (0) ← RecNetϕ (V ) 5: for k = 0 to K − 1 do 6: Calculate the ELBO surrogate gθ (ψ (k) ) using Eq. (24) 7: Update the parameter ψ by gradient ascent in Eq. (22) ψ (k+1) ← ψ (k) + α∇ψ gθ (ψ (k) ) + ϵk 8: end for 9: Set ψ ∗ ← ψ (K) 10: Update θ and ϕ by minimizing the cross-entropy loss in Eq. (5) P P (θ, ϕ) ← (θ, ϕ) − η∇(θ,ϕ) − i∈S ∗ log ψi∗ − i∈V \S ∗ log(1 − ψi∗ ) 11: end while 12: Set θ ∗ , ϕ∗ ← θ, ϕ 13: return θ ∗ , ϕ∗
Algorithm 3 ReSet (Inference) Require: V : test ground set; θ∗ , ϕ∗ : trained parameters; K: number of iterations; α: step size for gradient ascent Ensure: S ∗ : Predicted optimal subset 1: Obtain the initial variational parameter via the recognition network 2: ψ (0) ← RecNetϕ∗ (V ) 3: for k = 0 to K − 1 do 4: Calculate the ELBO surrogate gθ∗ (ψ (k) ) using Eq. (24) 5: Update ψ by gradient ascent: 6: ψ (k+1) ← ψ (k) + α∇ψ(k) gθ∗ (ψ (k) ) 7: end for 8: Set ψ ∗ ← ψ (K) 9: Discretize the continuous solution 10: S ∗ ← Rounding(ψ ∗ ) 11: return S ∗
22
Preprint.
Table 5: Detailed architectures of ReSet. Set (Surrogate) Function
Recognition Network
InitLayer(V, 256) MHA(256, 4) LayerNorm(256) SumPooling FC(256, 1, −)
InitLayer(V, 256) FC(256, 512, ReLU) FC(512, 1, Sigmoid) -
E
E XPERIMENTAL D ETAILS
E.1
T HE A RCHITECTURE OF R E S ET
Similar to the framework of Ou et al. (2022); Xie et al. (2024a), ReSet consists of two neural components: (i) a set function that is permutation invariant, and (ii) a recognition network that is permutation equivariant. Both components are implemented using the DeepSets architecture, with detailed architectural specifications provided in Table 5. Specifically, InitLayer(din , dout ) encodes the set objects into vector representations, FC(din , dout , f ) is a fully-connected layer with activation function f , MHA(d, h) denotes a multi-head attention layer with h heads. Synthetic datasets. We consider two synthetic benchmarks, namely Two-Moons and Gaussian Mixture. Each set element is represented by a two-dimensional vector encoding its Cartesian coordinates. For both datasets, the InitLayer is instantiated as a single-layer feed-forward network with architecture FC(2, 256, −). Amazon Baby Registry. The Amazon Baby Registry dataset consists of sets of products, where each product is associated with a short textual description. We encode these descriptions using a pre-trained BERT model (Lee & Toutanova, 2018), resulting in a 768-dimensional embedding for each element. The corresponding InitLayer is implemented as FC(768, 256, −). Double MNIST. The Double MNIST dataset contains digit images spanning all combinations from 00 to 99. Each image has a resolution of (64, 64) and is reshaped into a vector of dimension 4096. Accordingly, we employ a fully connected initialization layer with architecture FC(4096, 256, −). F-MNIST. The Fashion-MNIST dataset consists of grayscale images of fashion items from 10 categories. Each image has a spatial resolution of (28, 28) and is flattened into a 784-dimensional vector. Similar to Double MNIST, the InitLayer is implemented as a feed-forward neural network with architecture FC(784, 256, −). CelebA. The CelebA dataset comprises 202,599 face images, each of size (3, 64, 64). For this dataset, the InitLayer is realized using a convolutional neural network. Specifically, its architecture is given by ModuleList([Conv(32, 3, 2, ReLU), Conv(64, 4, 2, ReLU), Conv(128, 5, 2, ReLU), MaxPooling, FC(128, 256, −)]), where Conv(d, k, s, f ) denotes a convolutional layer parameterized by d output channels, kernel size k, stride s, and activation function f . CIFAR-10. The CIFAR-10 dataset consists of 60, 000 color images of size 32 × 32 pixels. For this dataset, the feature extraction network is implemented as a deep convolutional neural network consisting of three cascaded functional blocks. Specifically, the architecture is formulated as: ModuleList ([Block(32), Block(64), Block(128), AvgPool, FC(128, 256, −)]) , where each Block(d) is a sub-network designed to capture hierarchical spatial features, defined as the sequence: [Conv(d, 3, 2, ReLU), BN, Conv(d, 3, 1, ReLU), BN]. To maintain spatial resolution during convolution, zero-padding is applied such that p = ⌊k/2⌋. Each BN denotes a Batch Normalization layer applied after the convolution. It should be noted that the code related to this part was not publicly released in Ou et al. (2022), so we constructed the InitLayer architecture ourselves. 23
Preprint.
PDBBind. The PDBBind database provides experimentally determined binding affinities for biomolecular complexes (Liu et al., 2015a), together with high-resolution three-dimensional Cartesian coordinates of ligands and their corresponding target proteins, typically obtained via techniques such as X-ray crystallography. To extract informative representations, we leverage the atomic convolutional neural network (ACNN) (Gomes et al., 2017), which constructs nearest-neighbor graphs based on atomic spatial coordinates. In our implementation, the representation is obtained from the penultimate layer of ACNN, followed by additional feed-forward layers. Formally, the InitLayer is defined as ModuleList([ACNN[: −1], FC(1922, 2048, ReLU), FC(2048, 256, −)]), where ACNN[: −1] denotes the ACNN architecture without its final prediction layer, yielding an output dimension of 1,922. BindingDB. The BindingDB dataset consists of 52,273 drug–target interaction pairs. We adopt the DeepDTA framework (Öztürk et al., 2018) to encode each drug–target pair into a fixed-dimensional vector. DeepDTA represents drug compounds and target proteins as one-hot encoded sequences, which are subsequently processed by convolutional neural networks. The full specification of the InitLayer used for this dataset is reported in Table 6. Table 6: Detailed architectures of InitLayer in the BindingDB dataset. Drug
Target
Conv(32, 4, 1, ReLU) Conv(32, 4, 1, ReLU) Conv(64, 6, 1, ReLU) Conv(64, 8, 1, ReLU) Conv(96, 8, 1, ReLU) Conv(96, 12, 1, ReLU) MaxPooling MaxPooling FC(96, 256, ReLU) FC(96, 256, ReLU) Concat FC(512, 256, −)
E.2
I MPLEMENTATION D ETAILS
We provide the training details and hyperparameter settings for ReSet. The proposed model is trained using the Adam optimizer (Kingma & Ba, 2017) with a fixed weight decay of 1e − 5 and a batch size of 128. During training, we employ a cosine annealing learning rate scheduler (Loshchilov & Hutter, 2016) together with an early stopping strategy. The initial learning rate is selected via grid search over the range [1e − 5, 1e − 3] using half-decade intervals, with the minimum learning rate fixed at one-tenth of the chosen initial value. Training is conducted for up to 500 epochs, and early stopping is triggered if no performance improvement is observed for 20 consecutive epochs. For datasets that are more difficult to converge, we allow a longer patience. After training, the saved models are evaluated on the test set. All experiments are repeated 5 times with different random seeds, and the mean performance along with their standard deviations are reported. Within the ReSet framework, the gradient ascent step size α and the noise magnitude ϵ in the Langevin dynamics constitute two critical hyperparameters governing model performance. To enhance the adaptivity of the learning process, we allow both parameters to be updated during training. Specifically, the initial value of α is selected empirically. In contrast, although ϵ is treated as a learnable parameter, it is constrained to remain at a small scale of 1e − 5. This strategy aims to ensure numerical stability during training while enabling precise control over the determinacy of the model. E.3
BASELINES
Throughout our experiments, we evaluate the proposed models against four conventional approaches: Random Guess, Probabilistic Greedy Model (PGM) (Tschiatschek et al., 2018), DeepSet (Zaheer et al., 2017), and Set Transformer (Lee et al., 2019). Additionally, we compare our method with two state-of-the-art optimal subset oracles: EquiVSet (Ou et al., 2022) and INSET (Xie et al., 2024a). Detailed descriptions of these benchmarks are provided below. 24
Preprint.
• Random. We report the expected Jaccard coefficient obtained by random guessing. This baseline serves as a reference for assessing the intrinsic difficulty of the task. • PGM (Tschiatschek et al., 2018). The probabilistic greedy model tackles Problem (2) by introducing a differentiable relaxation of the greedy maximization procedure, which relies on enumerating all possible element permutations. For a detailed exposition, we refer readers to Appendix A of Ou et al. (2022). • DeepSet (Zaheer et al., 2017). By learning a permutation-invariant mapping from 2|V | to [0, 1]|V | , the DeepSet architecture provides a generic backbone for modeling set functions in optimal subset oracle frameworks (Ou et al., 2022; Xie et al., 2024a;b). We therefore include it as a baseline in our evaluation. To adapt DeepSets for point-wise tasks, we shift from an invariant to a permutation equivariant design. Instead of global reduction, we broadcast the mean-pooled global descriptor to each element, allowing for an output of size |V | that preserves the identity of individual members. • Set Transformer (Lee et al., 2019). Extending the DeepSet framework, Set Transformer incorporates self-attention mechanisms to model pairwise interactions among set elements. Similar to DeepSet, it can be adapted for subset selection tasks and is included as a competitive baseline. Similar to our adaptation of DeepSet, we bypass the original global pooling (PMA) and invariant decoder, utilizing only the encoder blocks. This preserves the permutation equivariance of the self-attention mechanism, allowing for element-wise prediction while maintaining an output dimensionality of |V |. • EquiVSet (Ou et al., 2022). EquiVSet aims to learn an optimal subset oracle without explicitly specifying the underlying set function. Instead, the set function is modeled as an energy-based formulation and approximated through a variational distribution. Optimization via mean-field variational inference requires iterative unrolling with Monte Carlo sampling to estimate gradients of the multilinear extension, which introduces substantial computational overhead. • INSET (Xie et al., 2024a). Building upon EquiVSet, INSET augments the set function Fθ (·) by injecting background information from the ground set V . By explicitly accounting for the symmetry of the ground set, this design introduces a stronger inductive bias and more effectively exploits contextual relationships within the global set. E.4
D ETAILED E XPERIMENTAL S ETTINGS FOR P RODUCT R ECOMMENDATION
The Amazon Baby Registry dataset (Gillenwater et al., 2014) is collected from Amazon and organized into multiple subsets based on product categories, such as toys and furniture. Within each category, which serves as a product universe, Amazon records multiple product subsets curated by different customers. These customer-selected subsets naturally correspond to optimal subset (OS) oracles. To ensure that each ground set V is associated with a single OS oracle S ∗ , we construct data samples in the form of (V, S ∗ ) as follows. We first discard customer-selected subsets whose cardinality is equal to 1 or exceeds 30. For each remaining OS oracle S ∗ , we randomly sample 30 − |S ∗ | additional products from the same category to form the complement set V \ S ∗ . Statistics of the resulting datasets across different categories are reported in Table 7. E.5
D ETAILED E XPERIMENTAL S ETTINGS FOR S ET A NOMALY D ETECTION
CelebA: The CelebA dataset consists of 202,599 images annotated with 40 binary attributes. As illustrated in Fig. 4, we randomly choose two attributes and form sets of fixed cardinality 8. For each ground set V , an OS oracle subset S ∗ is constructed by randomly sampling n ∈ {2, 3} images such that neither of the selected attributes is exhibited. Based on this procedure, we generate training, validation, and test splits containing 10,000, 1,000, and 1,000 samples, respectively. Comparing with the Setting of Ou et al. (2022). Although we aim to align our experimental setup as closely as possible with that of Ou et al. (2022), the experimental protocols for the anomaly detection task on Double MNIST (Sun, 2019), F-MNIST (Xiao et al., 2017), and CIFAR-10 (Krizhevsky & Hinton, 2009) are not publicly available. Moreover, the corresponding descriptions in the original work are relatively ambiguous. For clarity and reproducibility, we therefore re-establish and standardize the experimental settings for these three datasets. Specifically, for Double MNIST, we 25
Preprint.
Table 7: The statistics of Amazon product dataset. Categories
#products
|D|
|V |
P
|S ∗ |
E[|S ∗ |]
minS ∗ |S ∗ |
maxS ∗ |S ∗ |
Toys Furniture Gear Carseats Bath Health Diaper Bedding Safety Feeding Apparel Media
62 32 100 34 100 62 100 100 36 100 100 58
2,421 280 4,277 483 3,195 2,995 6,108 4,524 267 8,202 4,675 1,485
30 30 30 30 30 30 30 30 30 30 30 30
9,924 892 16,288 1,576 12,147 11,053 25,333 17,509 846 37,901 21,176 6,723
4.09 3.18 3.80 3.26 3.80 3.69 4.14 3.87 3.16 4.62 4.52 4.52
3 3 3 3 3 3 3 3 3 3 3 3
14 6 10 6 11 9 15 12 5 23 21 19
Bangs & No_Beard
Big_Lips & Male
Smiling & Brown_Hair
Black_Hair & Big_Nose
Figure 4: Example data points from the CelebA dataset adapted from Ou et al. (2022). Each row corresponds to one sample and contains |S | anomalous images (highlighted in red) along with 8−|S | normal images. Within each sample, normal images exhibit both selected attributes (shown in the rightmost column), whereas anomalous images exhibit neither attribute. generate a total of 1, 000 images distributed across all digit pairs from 00 to 99.12 For F-MNIST and CIFAR-10, we use the official datasets obtained via online sources. During sample construction, we randomly select n ∈ {2, 3, 4} images from the same class to form the optimal subset oracle S ∗ for all three datasets. For the Double MNIST dataset, V \S ∗ is constructed by selecting 20 − |S ∗ | images, whereas for F-MNIST and CIFAR-10, it is formed by drawing 10 − |S ∗ | images from different classes. Notably, the images in V \S ∗ are drawn from distinct classes, with at most one image per class. Finally, we generate samples on the fly, producing 10, 000, 1, 000, and 1, 000 instances for the training, validation, and test sets, respectively. Visual illustrations are provided in Figures 5 and 6.
Figure 5: A sampled data for the Double MNIST dataset, which consists of |S ∗ | images with the same digit (red box, 76 in this case) and 20 − |S ∗ | images with different digits. 12 The source code for dataset generation is publicly available at https://github.com/ shaohua0116/MultiDigitMNIST.
26
Preprint.
Shirt
Bird
Sandal
Horse
Bag
Frog
Coat
Truck
Dress
Deer
Figure 6: Illustrative samples from the F-MNIST (left) and CIFAR-10 (right) datasets. Each row represents a single sample composed of |S ∗ | images sharing the same class label (highlighted in red) and 10 − |S ∗ | additional images drawn from distinct classes. The corresponding labels are shown in the rightmost column.
E.6
D ETAILED E XPERIMENTAL S ETTINGS FOR C OMPOUND S ELECTION
In the drug discovery pipeline, identifying compounds that exhibit strong bioactivity (Wallach et al., 2015; Li et al., 2021; Ji et al., 2023), sufficient structural diversity, and favorable ADME (absorption, distribution, metabolism, and excretion) profiles (Gimeno et al., 2019) constitutes a fundamental screening stage. In practice, virtual screening is commonly implemented as a sequence of filtering operations. A typical workflow first prioritizes candidates based on predicted bioactivity, subsequently enforces diversity constraints, and finally excludes compounds with undesirable ADME characteristics, yielding a refined candidate subset. However, owing to privacy considerations and cost constraints in pharmaceutical development, intermediate supervision signals corresponding to these individual filtering stages are often inaccessible or prohibitively expensive to obtain. Consequently, learning the entire screening procedure in an end-to-end fashion becomes desirable, a setting that naturally aligns with the formulation of optimal subset (OS) oracles. Despite this alignment, fully end-to-end modeling remains challenging and often demands substantial domain knowledge. In this work, we therefore adopt a simplified yet representative setting by focusing on the prediction of high-bioactivity compounds, following the experimental protocol of Xie et al. (2024a). Experiments are conducted on the PDBBind (Liu et al., 2015a) and BindingDB (Liu et al., 2007) datasets using a unified bioactivity-based filtering criterion. Algorithm 4 illustrates the data generation procedure used to simulate the OS oracle for the compound selection task. Specifically, random choose(C, n) denotes uniformly sampling n compounds from a database C (either PDBBind or BindingDB) to construct the ground set V , while topK bioactivity(V, m) selects the top-m compounds exhibiting the highest biological activity from V . These two operations together form a bioactivity-based filtering process, through which the OS oracle S ∗ is obtained. In our experiments, we set (n, m) = (30, 10) for PDBBind and (300, 100) for BindingDB, respectively. This procedure yields a data point (V, S ∗ ), where V represents a candidate compound set and S ∗ corresponds to the subset of compounds with the highest bioactivity.
Algorithm 4 OS Oracle Construction Procedure for the Compound Selection Task Input: C: compound database; n: cardinality of the ground set; m: number of highly active compounds Output: A data sample (V, S ∗ ) 1: Ground Set Formation: 2: Uniformly sample n compounds from the database C 3: V ← random choose(C, n) 4: Bioactivity-based Selection (OS Oracle): 5: Select the top-m compounds with the highest bioactivity scores from V 6: S ∗ ← topK bioactivity(V, m) 7: return (V, S ∗ )
27
Preprint.
F
A DDITIONAL E XPERIMENTS
F.1
A BLATION S TUDY
Our framework diverges from the original OS Oracle baseline primarily through two key components: the integration of multi-head self-attention and the incorporation of Langevin dynamics. To address potential concerns that the performance gains might stem solely from these modules—despite our model’s significantly lower parameter count—we conduct an ablation study across several model variants. Specifically, we augment the baseline with an identical self-attention layer (denoted as “-MHA”) and evaluate a version of our proposed method without Langevin dynamics (denoted as “w/o L”). Table 8 presents the ablation results. It is evident that our proposed method consistently outperforms both the original baseline and its MHA-enhanced variant. Notably, in several cases, the performance gap remains marginal upon removing Langevin dynamics; in certain instances, the deterministic version even surpasses the full model. This suggests that the benefit of Langevin dynamics is dataset-dependent. When the objective landscape is relatively smooth, removing the stochastic term can yield comparable or slightly better performance by reducing noise in the optimization trajectory. In more challenging landscapes, Langevin dynamics can improve exploration and help stabilize training. Thus, its usefulness depends on problem complexity and can be determined empirically. Nevertheless, the overall results provide strong empirical evidence for the efficacy of our proposed framework. Table 8: Performance comparison (%) of different OS Oracle variants in the ablation study. Variants
F.2
Toys
Furniture
Carseats
Diaper
Safety
Media
EquiVSet EquiVSet-MHA
68.0 ± 2.0 75.7 ± 1.4
17.2 ± 0.9 15.1 ± 1.8
21.0 ± 1.0 20.6 ± 3.2
83.0 ± 1.0 88.5 ± 0.7
25.0 ± 3.0 13.9 ± 0.6
57.0 ± 1.0 62.9 ± 2.8
EquiVSet-R(w/o L) EquiVSet-R
78.7 ± 0.7 77.9 ± 1.0
33.8 ± 4.8 35.9 ± 1.3
26.5 ± 2.4 26.9 ± 1.0
89.8 ± 1.0 90.4 ± 0.6
37.6 ± 5.3 32.1 ± 9.4
64.6 ± 4.0 67.0 ± 1.0
INSET INSET-MHA
76.9 ± 0.5 77.8 ± 1.3
16.9 ± 5.0 17.4 ± 1.5
23.1 ± 3.4 23.5 ± 2.0
88.0 ± 0.7 90.2 ± 1.7
23.8 ± 1.5 14.0 ± 0.5
62.0 ± 2.3 65.5 ± 2.4
INSET-R(w/o L) INSET-R
79.3 ± 0.7 79.0 ± 0.6
36.9 ± 1.0 36.5 ± 1.5
29.7 ± 2.8 30.2 ± 3.9
89.9 ± 0.5 90.4 ± 0.5
37.4 ± 5.0 38.3 ± 4.0
67.8 ± 1.3 68.1 ± 1.1
C OMPUTATIONAL E FFICIENCY
Table 9 compares the inference time and memory usage of different methods under increasing sample and set sizes. Across all settings, replacing Monte Carlo sampling with the proposed learned relaxation consistently improves computational efficiency. For EquiVSet, the relaxation variant achieves up to 2.6× faster inference and reduces memory usage by up to 1.6×. The improvement is more pronounced for INSET, where INSET-R obtains up to 3.8× speedup and 2.5× lower memory usage. These gains become especially important as N and |V | grow, suggesting that the proposed relaxation effectively alleviates the computational and memory overhead caused by sampling-based gradient estimation. Table 9: Computational efficiency on the Two Moons dataset under varying sample and set sizes. The average inference time per epoch (ms) and peak memory usage (MB) are reported. Method (#Params)
N =1000, V =100 Time Memory
N =2000, V =200 Time Memory
N =3000, V =300 Time Memory
EquiVSet (0.76M) EquiVSet-R (0.26M)
178 68
936 569
686 370
1942 1352
1575 1021
3043 2468
INSET (1.14M) INSET-R (0.52M)
479 125
2554 1021
1892 732
5368 2483
4420 2025
8505 4396
28
Preprint.
F.3
C OMPARISON WITH L OW-VARIANCE G RADIENT E STIMATORS
Table 10 compares our method with low-variance gradient estimators by increasing the number of Monte Carlo samples. Although using more samples slightly improves performance, it also substantially increases inference time. In contrast, our method achieves the best performance for both EquiVSet and INSET while requiring significantly less inference time. This suggests that directly learning the relaxation provides a more favorable performance-efficiency trade-off than reducing gradient variance through additional sampling. Table 10: Comparison with low-variance gradient estimators on the Toys dataset. Performance (%) and average inference time (ms) are reported. Number of Samples
F.4
10
20
30
40
50
Ours
EquiVSet
MJC Time
71.6±0.7 98
74.0±0.6 160
74.2±0.5 223
75.3±0.2 290
75.9±0.6 360
77.9±1.0 34
INSET
MJC Time
76.8±1.0 165
77.2±0.5 300
76.9±1.3 423
76.9±0.6 547
77.2±0.4 669
79.0±0.6 54
S ENSITIVITY A NALYSIS OF H YPERPARAMETERS
Impact of the Gradient Ascent Step Size α. In ReSet, the step size α for gradient ascent in Eq. (22) serves as a critical hyperparameter governing both convergence velocity and numerical stability, directly impacting the final performance. An excessively large step size may lead to divergence, while a disproportionately small one risks trapping the optimization in suboptimal local minima and unnecessarily prolonging inference time. We evaluate the performance of our two proposed OS Oracle variants across different datasets under various α configurations, with results illustrated in Figures 7 and 8. Furthermore, in our practical implementation, we set α as a learnable parameter to achieve an adaptive step size. This approach alleviates the burden of manual hyperparameter tuning for unseen datasets and enables the model to dynamically calibrate its optimization dynamics for diverse tasks. The performance of this self-adaptive configuration is denoted as “Adaptive” in the corresponding figures. The empirical results indicate that a fixed step size near 0.5 generally yields favorable outcomes. Notably, the adaptive strategy consistently outperforms the fixed-parameter configurations in most scenarios, demonstrating its superior capability in navigating the optimization landscape.
Figure 7: Sensitivity analysis of EquiVSet-R performance under different values of α, where “Adaptive” denotes the variant with α as a learnable parameter. 29
Preprint.
Figure 8: Sensitivity analysis of INSET-R performance under different values of α, where “Adaptive” denotes the variant with α as a learnable parameter.
F.5
S ET S IZE T RANSFERABILITY A NALYSIS
Following Ou et al. (2022), we experiment to understand the pattern of set size transfer ability. In this experiment, we train the models using fixed sizes of the ground set but test the trained model on different sizes. We fix the size of OS oracle S ∗ to be 10, and train the model with ground set V of size 100. After training, we test it using varying sizes of ground set in the range of {200, 400, 600, 800, 1000}. The experiments with two OS Oracle variants are conducted on the Two-Moons and Gaussian-Mixture datasets, respectively, with the results shown in Figures 9 and 10. As illustrated in the results, both proposed OS Oracle variants exhibit a observable degree of set-size transferability. When trained on a fixed ground set size and evaluated on increasingly larger sets, the models experience a gradual performance degradation, which aligns with theoretical expectations regarding the increased complexity of the search space. This decline is relatively marginal on the Gaussian Mixture dataset but becomes more pronounced on the Two Moons dataset. We attribute this discrepancy to the intrinsic task complexity; the non-linear, intertwined geometry of the TwoMoons distribution likely poses a greater challenge for zero-shot size transfer compared to the more clustered Gaussian-Mixture distribution.
Figure 9: Synthetic results of EquiVSet-R under the set-size transferability setting are reported. The blue bars correspond to scenarios where the ground set size used at test time matches that employed during training, whereas the yellow bars indicate evaluations conducted with mismatched ground set sizes at inference.
30
Preprint.
Figure 10: Synthetic results of INSET-R under the set-size transferability setting are reported. The blue bars correspond to scenarios where the ground set size used at test time matches that employed during training, whereas the yellow bars indicate evaluations conducted with mismatched ground set sizes at inference.
G
R ELATED W ORK
G.1
S ET F UNCTION L EARNING
Many existing studies focus on learning an unknown function to predict the value of a given set (Wendler et al., 2019; 2021; De & Chakrabarti, 2022), commonly referred to as a function-value (FV) oracle. Zaheer et al. (2017) introduced the DeepSets architecture to model permutationinvariant and permutation-equivariant functions for set-valued prediction. Building on DeepSets, Lee et al. (2019) incorporated the Transformer architecture to capture relational dependencies among elements within a set. Additionally, several works have explored deep architectures for submodular set functions (Bilmes & Bai, 2017; Dolhansky & Bilmes, 2016; Balcan & Harvey, 2018). It is important to note that these methods typically require a substantial amount of supervision, which poses challenges for their practical deployment. Optimal Subset Oracle. Ou et al. (2022) propose to learn latent set functions using optimal subsets rather than explicit function values as supervision, formulating the problem under an energybased modeling framework with parameters updated via variational inference. Building upon this framework, Xie et al. (2024a;b) further incorporate contextual information from the superset and hierarchical attention mechanisms. Özcan et al. (2025) further improve the fixed-point iterations in variational inference through implicit differentiation. However, these methods require unrolled Monte Carlo sampling for gradient estimation during iterative optimization, rendering the training process inefficient. G.2
E NERGY- BASED M ODELING
Energy-based learning (LeCun et al., 2006) provides a classical framework for modeling underlying data distributions, in which model parameters are learned via principled techniques such as contrastive divergence (Carreira-Perpinan & Hinton, 2005; Tieleman, 2008; Du & Mordatch, 2019) and score matching (Hyvärinen & Dayan, 2005; Vincent, 2011; Song et al., 2020). In the framework of Ou et al. (2022), energy-based model are employed to parameterize the set mass function. Building upon this formulation, we reinterpret the associated variational inference procedure from the perspective of free energy estimation (He et al., 2025) and energy minimization (Gladstone et al., 2025). This viewpoint aligns closely with a broad line of work that casts learning and inference as energy minimization problems (Du et al., 2022; 2024; Wang & Du, 2025). G.3
S UBMODULAR F UNCTION L EARNING
In practical applications, learning set functions typically necessitates the incorporation of structural priors. The underlying objectives are often assumed to be submodular, i.e., exhibiting the diminishing returns property. Unlike prior work that predominantly relies on function value oracles for 31
Preprint.
parameter estimation (Dolhansky & Bilmes, 2016; Bilmes & Bai, 2017; Djolonga & Krause, 2017; Kothawade et al., 2020; De & Chakrabarti, 2022; Bhatt et al., 2024), this work learns latent set functions characterized by more general weak submodular properties through a more practical and computationally efficient optimal subset oracle.
32