Learning Set Functions with Implicit Differentiation

Gözde Özcan 1 Chengzhi Shi 1 Stratis Ioannidis 1

Abstract

As MLE is intractable in this setting, Ou et al. [1] propose
performing variational inference instead. In turn, they showthat approximating the distribution of oracle selections requires solving a fixed-point equation per sample. However,
fixed-point iterations may diverge in practice. In addition,
Ou et al. implement these iterations by stacking up neural network layers; this makes backpropagation expensive,
limiting their experiments to only a handful of iterations.

Ou et al. [1] introduce the problem of learning set
functions from data generated by a so-called optimal subset oracle. Their approach approximates
the underlying utility function with an energybased model. This approximation yields iterations of fixed-point updates during mean-field
variational inference. However, as the numberof iterations increases, automatic differentiation
becomes computationally prohibitive due to thesize of the Jacobians that are stacked during backpropagation. We address this challenge with implicit differentiation and examine the convergenceconditions for the fixed-point iterations. We empirically demonstrate the efficiency of our methodon subset selection applications including productrecommendation and anomaly detection tasks.

In this work, we establish a condition under which the fixedpoint iterations proposed by Ou et al. [1] are guaranteed toconverge. We also propose a more effective gradient computation utilizing the recent advances in implicit differentiation[13]–[15], instead of unrolling the fixed-point iterations via
automatic differentiation [16]. This corresponds to differentiating after infinite fixed point iterations, while remaining
tractable; we experimentally show that this improves thepredictive performance of the inferred models.

2. Problem Setup

1. Introduction

In the setting introduced by Ou et al. [1], the aim is to learn
set functions from a dataset generated by a so-called optimal
subset oracle. The dataset D consists of sample pairs of theform (S ∗ , V ), where V ⊆ Ω is a set of options, i.e., items
from a universe Ω, and S ∗ is the optimal subset of V , asselected by an oracle. The goal is to learn a set function
Fθ : 2Ω × 2Ω , → R, parameterized by θ ∈ Rd , modelingthe utility of the oracle, so that

Many interesting applications operate with set-valued outputs and/or inputs. Examples include product recommendation [2], [3], compound selection [4], set matching [5],set retrieval [6], point cloud processing [7], [8], and set
anomaly detection [9], to name a few. Motivated by this,several recent works [10]–[12] apply neural networks tolearn set functions from input/function value pairs, assuming access to a dataset generated by a function value oracle(see App. A for a review of related literature).

S ∗ = arg max Fθ (S, V ),

(1)

S⊆V

Recently, Ou et al. [1] proposed an approximate maximum
likelihood estimation framework under the supervision ofa so-called optimal subset oracle. In contrast to traditional
function value oracles, a label produced by an optimal subset oracle is the subset that maximizes an (implicit) utility
set function, in the face of several alternatives. The goal ofinference is to learn, in a parametric form, this utility function, under which observed oracle selections are optimal.

for all pairs (S ∗ , V ) ∈ D. As a motivating example, consider the case of product recommendations. Given a groundset V of possible products to recommend, a recommenderselects an optimal subset S ∗ ⊆ V and suggests these to auser. In this setting, the function Fθ (S, V ) captures, e.g., therecommender objective, the utility of the user, etc. Having
access to a dataset of pairs, the goal is to learn Fθ , effectively reverse-engineering the objective of the recommender
engine, inferring the user’s preferences, etc.

1

Northeastern University, Department of Electrical and Computer Engineering, Boston, MA 02115, USA. Correspondence to:Gözde Özcan <gozcan@ece.neu.edu>.

2.1. MLE with Energy-Based Modeling

Published at the 2 nd Differentiable Almost Everything Workshop
at the 41 st International Conference on Machine Learning, Vienna, Austria. July 2024. Copyright 2024 by the author(s).

Ou et al. [1] propose an approximate maximum likelihood
estimation (MLE) by modeling oracle behavior via a Boltz1

Learning Set Functions with Implicit Differentiationmann energy (i.e., soft-max) model [17]–[20]. They assume
that the oracle selection is probabilistic, and the probability
that S is selected given options V is given by:pθ (S | V ) = P

exp (Fθ (S, V )).
′S ′ ⊆V exp (Fθ (S , V ))2.3. ELBO Maximization
To compute ψ ∗ , Ou et al. [1] show that minimizing the KL
divergence in Eq. (5) via maximizing the corresponding
evidence lower bound (ELBO) reduces to solving a fixed
point equation. In particular, Eq. (5) is equivalent to thefollowing ELBO maximization [22], [23]:(2)This is equivalent to Eq. (1), presuming that the utility Fθ (·)
is distorted by Gumbel noise [21]. Then, given a dataset
D = {(Si∗ , Vi , )}N
i=1 , MLE amounts to:arg maxθPN∗
i=1 [log pθ (Si | Vi )] .

max F̃ (ψ, θ) + H(q(S, V, ψ)),ψwhere H(·) is the entropy and F̃ : [0, 1]|V | × Rd → R is theso-called multilinear extension of Fθ (S, V ) [24], given by:XYYF̃ (ψ, θ) =ψiFθ (S, V )(1 − ψi ). (7)

(3)

S⊆V

Notice that multiplying Fθ with a constant c > 0 makes nodifference in the behavior of the optimal subset oracle in
Eq. (1): the oracle would return the same decision underarbitrary re-scaling. However, using c · Fθ (·) in the energybased model of Eq. (2) corresponds to setting a temperatureparameter c in the Boltzmann distribution [17], [21], interpolating between the deterministic selection (c → ∞) in
Eq. (1) and the uniform distribution (c → 0).1

(8)

where the function σ : R|V | → R|V | is defined as σ(x) =|V |σ(xi ) i=1 and σ : R → R is the sigmoid function, i.e.,
σ(x) = (1 + exp (−x))−1 . The detailed derivation of thiscondition can be found in App. B.1.
Observing that the stationary condition in Eq. (8) is a fixed
point equation, Ou et al. propose solving it via the following
fixed-point iterations. Given θ ∈ Rd ,
ψ (k) ← σ(∇ψ F̃ (ψ (k−1) , θ)),(9)where k ∈ N, starting from ψ (0) ∈ [0, 1]|V | . Exact computation of the multilinear relaxation defined in Eq. (7) requires an exponential number of terms in the size of V .
However, it is possible to efficiently estimate both the multilinear relaxation and its gradient ∇ψ F̃ (ψ, θ) via Monte
Carlo sampling (see App. B.2 for details).

L(ψ ∗ , θ) = EP(V, S) [− log q(S, V , ψ ∗ )]


N
X

1 X X≈−log ψj∗ −log 1 − ψj∗  ,
N i=1
∗
∗

2.4. DiffMF and Variants
Putting everything together yields the DiffMF algorithm
introduced by Ou et al. [1]. In particular, they implement
the fixed-point iterative update steps in Eq. (9) by executing
a fixed number of iterations K, given θ. In their implementation, this amounts to stacking up K layers, each involving
an estimate of the gradient of the multilinear relaxation via
sampling, and thereby multiple copies of a neural network
representing Fθ (·). Subsequently, this extended network isentered in the loss (4), which is minimized w.r.t. θ via SGD.

(4)

where ψ ∗ ∈ [0, 1]|V | is the probability vector obtained bysolving:ψ ∗ = arg min KL(q(S, V, ψ) || pθ (S | V )),i∈V \Sψ − σ(∇ψ F̃ (ψ, θ)) = 0,

Learning θ by MLE is challenging precisely due to the exponential number of terms in the denominator of Eq. (2). Instead, Ou et al. [1] construct an alternative optimization objective via mean-field variational inference as follows. First,
they introduce a mean field variationalof theQ
Q approximation
density pθ given by q(S, V, ψ) = i∈S ψi i∈V \S (1−ψi ),parameterized by ψ, where ψ is the probability vector representing the probability of each element s ∈ V being in
the optimal subset S ∗ . Then, estimation via variational
inference amounts to:j∈Vi \Sii∈S

Ou et al. show that a stationary point maximizing the ELBO
in Eq. (6) must satisfy:

2.2. Variational Approximation of Energy-Based Models

j∈Si

(6)

(5)

ψ

where KL(·||·) is the Kullback-Leibler divergence, andpθ (S) is the energy-based model defined in Eq. (2). In
turn, this is found through the ELBO maximization process
we discuss in the next section.

They also introduce two variants of this algorithm, regressing ψ (0) as a function of the item features via an extra recognition network, assuming the latter are independent (terming
inference in this setting as EquiVSetind ) or correlated by aGaussian copula [25], [26] (termed EquiVSetcopula ). Compared to DiffMF, both translate to additional initial layers/steps per epoch.

1

From a Bayesian point of view, multiplying Fθ (·) with c > 0
yields the posterior predictive distribution under an uninformative
Dirichlet conjugate prior per set with parameter α = ec [17].

2Learning Set Functions with Implicit Differentiation

with the size of the ground set, we can find a unique solution
to Eq. (8), no matter where we start the iterations in Eq. (9).

2.5. Challenges
The above approach by Ou et al. [1], and constituent algorithm variants, have two drawbacks. First, the fixed-point
iterative updates given in Eq. (9) are not guaranteed to converge to an optimal solution. We indeed observed divergence experimentaly in practice. Without convergence anduniqueness guarantees, the quality of the output, ψ (K) , isheavily dependent on the selection of the starting point,ψ (0) . Moreover, as these iterations correspond to stacking
up layers, each containing multiple copies of Fθ (·) due tosampling, backpropagation is computationally prohibitive
both in terms of time as well as space complexity. In fact,
poor performance due to lack of convergence, as well as
computational considerations, led Ou et al. to set the number of iterations K ≤ 5 (even K = 1) in their experiments.
We address both of these challenges in the next section.

3.2. Efficient Differentiation through Implicit Layers
Our second contribution is to disentangle gradient computation from stacking layers together, by using the implicit
function theorem [34]. This allows us to use the recent workon deep equilibrium models (DEQs) [13], [14].
Define ψ ∗ (·) to be the map θ 7→ ψ ∗ (θ) induced by Eq. (8);equivalently, given θ, ψ ∗ (θ) is the (unique by Thm. 3.2)
limit point of iterations (9). Observe that, by the chain rule:∇θ L(ψ ∗ (θ), θ) = ∇ψ∗ L(ψ ∗ (θ), θ) · ∂θ ψ ∗ (θ).

(10)

The term that is difficult to compute here via backpropagation, that required stacking in Ou et al. [1], is theJacobian ∂θ ψ ∗ (θ), as we do not have the map ψ ∗ (·) in aclosed form. Nevertheless, we can use the implicit function
theorem (see Thm. C.4 in App. C) to compute this quantity.

3. Our Approach
We (a) establish conditions under which iterations of Eq. (8)converge to a unique solution, by utilizing Banach fixedpoint theorem and (b) establish a way to efficiently compute
the gradient of the loss at the fixed point by using the implicit
function theorem. Our results pave the way to utilize recent
tools developed in the context of implicit differentiation[13], [14], [27] to the setting of Ou et al. [1].

Indeed, to simplify the notation for clarity, we define afunction G : [0, 1]|V | × Rd → [0, 1]|V | , whereG(ψ(θ), θ) ≜ σ(∇ψ F̃ (ψ, θ)) − ψ
and rewrite Eq. (8) as G(ψ(θ), θ) = 0. Using the implicit
function theorem, given in App. C, we obtain−∂ψ G(ψ ∗ (θ), θ) ∂θ ψ ∗ (θ) = ∂θ G(ψ ∗ (θ), θ) .
{z}{z
} | {z } ||3.1. Convergence Condition for the Fixed-Point
Fixed-points can be attracting, repelling, or neutral [28],[29]. We characterize the condition under which the convergence is guaranteed in the following assumption.
Assumption 3.1. Consider the multilinear relaxation F̃ :[0, 1]|V | × Rd → R of Fθ (·), as defined in Eq. (7). For allθ ∈ Rd , supψ∈[0,1] |F̃ (ψ, θ)| < |V1 | .

A∈R|V |×|V |J∈R|V |×d

(11)

B∈R|V |×dThis yields the following way of computing the Jacobian
via implicit differentiation:
Theorem 3.3. Computing ∂θ ψ ∗ (θ) is the equivalent of solving a linear system of equations, i.e., ∂θ ψ ∗ (θ) = A−1 B,
A = I − Σ′ (∇ψ F̃ (ψ, θ)) · ∇2ψ F̃ (ψ, θ) , andAs discussed in Sec. 2, scaling Fθ (S, V ) by a positive scalaramounts to setting the temperature of a Boltzmann distribution. Moreover, neural networks are often Lipschitzregularized for bounded inputs and weights [30]–[32].
Therefore, for any such Lipschitz neural network, we cansatisfy Asm. 3.1 by appropriately setting the temperature parameter of the EBM in Eq. (2). Most importantly, satisfying
this condition guarantees convergence:
Theorem 3.2. Assume a set function Fθ : 2V → R satisfies
Asm. 3.1. Then, the fixed-point given in Eq. (8) has a uniquesolution ψ ∗ ∈ [0, 1]|V | where ψ ∗ = σ(∇ψ F̃ (ψ ∗ , θ)).
Moreover, starting with an arbitrary point ψ (0) ∈ [0, 1]|V | ,ψ ∗ can be found via the fixed-point iterative sequence described in Eq. (9) where limk→∞ ψ (k) = ψ ∗ .

(12)
B = Σ′ (∇ψ F̃ (ψ, θ)) · ∂θ ∇ψ F̃ (ψ, θ) ,|V | where Σ′ (x) = diag σ ′ (xi ) i=1 , and σ ′ (x) = (1 +exp (−x))−2 · exp (−x).

The proof is in App. E. Eq. (11) shows that the Jacobian∂θ ψ ∗ (θ), can be expressed in terms of Jacobians of G at
the solution point. This means implicit differentiation only
needs the final fixed point value, whereas automatic differentiation via the approach by Ou et al. [1] required all theiterates (see also [14]). In practice, we use JAXopt [27] for
its out-of-the-box implicit differentiation support.

4. Experiments

The proof can be found in App. D, and relies on the Banach
fixed-point theorem [33]. Theorem 3.2 implies that as longas F̃ (ψ, θ) is bounded and this bound is inversely correlated

We evaluate our proposed method on two real-world, two
synthetic datasets for set anomaly set detection and product3

Learning Set Functions with Implicit Differentiation

Product Recommendation (Amazon)

AD

Datasets
CelebA
GaussianMoons
apparelbathbeddingcarseats
diaperfeedingfurnituregear
healthmedia
safetytoys

EquiVSetind
Test JCTime (s)53.16 ± 0.68 1822.33 ± 208.36
90.87 ± 0.16
20.76 ± 1.89
56.47 ± 0.52
17.98 ± 0.97
66.29 ± 1.66
37.13 ± 12.32
68.00 ± 0.69
32.42 ± 5.22
65.36 ± 0.29
40.73 ± 4.24
20.64 ± 1.15
13.31 ± 5.43
72.23 ± 1.75
52.11 ± 17.47
71.36 ± 1.13
83.91 ± 25.05
14.60 ± 1.92
5.64 ± 1.48
64.06 ± 1.07
42.58 ± 3.29
63.35 ± 1.14
34.34 ± 4.36
56.21 ± 1.17
46.36 ± 6.71
14.51 ± 0.57
4.95 ± 1.04
63.08 ± 0.68
32.32 ± 4.43

EquiVSetcopula
Test JCTime (s)54.31 ± 1.06 1889.18 ± 163.30
90.98 ± 0.04
27.24 ± 5.92
58.46 ± 0.20
17.70 ± 0.59
76.28 ± 0.61
71.23 ± 6.31
75.82 ± 1.45
46.78 ± 7.114
75.81 ± 0.95
58.77 ± 11.75
20.30 ± 1.26
12.06 ± 5.87
82.52 ± 0.69
82.23 ± 21.69
81.75 ± 1.44
116.54 ± 41.32
15.87 ± 1.75
7.69 ± 3.62
72.34 ± 1.36
61.06 ± 15.40
70.99 ± 0.93
60.25 ± 10.77
55.29 ± 1.78
52.13 ± 15.07
13.94 ± 1.77
4.43 ± 0.30
67.97 ± 0.98
48.00 ± 6.45

Test JC
51.98 ± 0.64
90.96 ± 0.05
58.54 ± 0.10
69.65 ± 0.73
71.75 ± 0.33
66.32 ± 1.00
20.46 ± 1.08
81.10 ± 0.78
77.34 ± 0.35
15.78 ± 2.01
59.44 ± 1.19
58.69 ± 0.49
51.45 ± 1.23
13.77 ± 0.45
63.19 ± 1.42

DiffMFTime (s)1837.20 ± 221.18
59.63 ± 18.64
44.39 ± 6.42
106.66 ± 19.01
92.89 ± 3.37
94.47 ± 12.73
19.31 ± 10.48
173.69 ± 31.34
239.76 ± 57.10
6.65 ± 1.06
95.18 ± 13.29
73.97 ± 11.68
55.07 ± 11.51
5.27 ± 0.76
78.08 ± 18.49

iDiffMF2
Test JCTime (s)51.17 ± 0.74
2075.74 ± 225.35
90.81 ± 0.08
37.76 ± 0.73
58.47 ± 0.17
36.12 ± 1.95
69.66 ± 1.62
83.55 ± 21.35
76.59 ± 0.82
62.74 ± 9.24
75.99 ± 0.88
67.55 ± 16.24
21.25 ± 0.94
39.84 ± 12.28
82.12 ± 1.42
99.11 ± 36.21
80.30 ± 0.72
109.27 ± 24.42
21.97 ± 3.46
42.62 ± 9.42
58.47 ± 2.39
42.61 ± 8.56
72.08 ± 1.22
78.65 ± 13.01
57.88 ± 0.74
68.61 ± 7.66
30.36 ± 1.46
43.57 ± 3.95
67.49 ± 1.08
73.69 ± 26.15

iDiffMF∗
Test JCTime (s)51.10 ± 0.83
2206.81 ± 455.08
91.05 ± 0.02
43.22 ± 5.50
58.56 ± 0.26
35.88 ± 1.39
69.31 ± 1.77
75.07 ± 20.58
76.78 ± 0.77
63.52 ± 11.36
75.99 ± 0.88
65.44 ± 14.89
21.23 ± 0.91
34.78 ± 6.55
82.11 ± 1.65
101.38 ± 56.45
80.30 ± 0.72
108.00 ± 22.30
21.97 ± 3.46
38.52 ± 9.15
58.47 ± 2.39
41.35 ± 8.39
72.46 ± 0.87
82.37 ± 12.84
57.74 ± 0.74
64.60 ± 5.04
30.36 ± 1.46
40.60 ± 4.20
67.10 ± 1.00
52.22 ± 10.02

Table 1. Test Jaccard coefficient and training time for set anomaly detection (AD) and product recommendation tasks, across all fivealgorithms. iDiffMF2 and iDiffMF∗ correspond to our algorithm with Frobenius and nuclear norm rescaling, respectively. Bold andunderline indicate the best and second best performance results, respectively.

where Sˆi∗ is the prediction of the model.

recommendation tasks (see App. F). We closely follow theexperimental setup of Ou et al. [1]. We report the mean andstandard variation for all experiment results based on 5 runs
with different seeds.

4.2. Results
We report the performance of our proposed iDiffMF methods against the baselines on Tab. 1. For the majorityof the cases, iDiffMF∗ achieves either the best or thesecond-best JC score. While the next best competitor,
EquiVSetcopula , performs the best on some of the datasets
its performance is sub-optimal on the remaining datasets.
For the Amazon furniture and safety datasets, iDiffMF variants give significantly better results than EquiVSetcopula ,even though EquiVSetcopula is faster. This is probably because EquiVSetcopula converges to a local optimum andfinishes training earlier. It is also important to highlightthat we evaluate iDiffMF using JAX+Flax [36]–[39] while
we use PyTorch to evaluate the baselines. Therefore, thedifferences in running time can also be explained with theframework differences. Finally, even though iDiffMF executes fixed-point iterations until convergence, as opposed to
K = 1 or 5 in remaining methods [1], the average runningtimes are comparable across datasets.

4.1. Algorithms
DiffMF [1]: Differentiable mean field variational inference
algorithm described in Section 2.3 where, as per Ou et al.,the number of iterations is set to K = 5 for all datasets.
EquiVSetind [1]: Equivariant variational inference algorithm proposed by Ou et al. [1]. It is a variation of the
DiffMF algorithm where the parameter ψ is predicted by anadditional recognition network as a function of the features.
As per Ou et al., we set K = 1 for all datasets.
EquiVSetcopula [1]: A correlation-aware version of the
EquiVSetind algorithm where the relations among the input
elements are modeled by a Gaussian copula. As per Ou etal., we set K = 1 for all datasets.
iDiffMF: Our proposed implicit alternative to the DiffMF
algorithm where we solve the fixed-point condition in
Eq. (8) with a very low tolerance threshold, instead of running the fixed-point iterations in Eq. (9) for only a fixed
number of times. Although DNNs are bounded, the exact
computation of their Lipschitz constant is NP-hard [31]. In
our implementation, we satisfy the condition in Asm. 3.1 bynormalizing the gradient of the multilinear relaxation by its
norm and scaling it with the size of the ground set, i.e., we
multiply ∇ψ F̃ (ψ, θ) with 2/(|V |∥∇ψ F̃ (ψ, θ)∥). For theL2 (∥ · ∥2 ) and nuclear (∥ · ∥∗ ) norms, we observe that thisheuristic works well in practice and propose two variations
as iDiffMF2 and iDiffMF∗ , respectively.

5. Conclusion
We improve upon an existing learning set functions with
an optimal subset oracle setting by characterizing the convergence condition of the fixed point iterations resulting
during MLE approximation and by using implicit differentiation over automatic differentiation. Our results performbetter than or comparable to the baselines for the majorityof the cases without the need of an additional recognition
network.

We train all models with an Adam optimizer [35] with alearning rate η = 10−4 and a batch size of 128.
Metrics: We measure the performance of the algorithms
by their mean Jaccard Coefficient (JC) score. The mean
PN |S ∗ ∩Sˆ∗ |JC over all optimal oracle subset pairs is N1 i=1 |Si∗ ∪Sˆi∗ | ,i

6. Acknowledgements
We gratefully acknowledge support from the NSF grantCCF-1750539.

i

4

Learning Set Functions with Implicit Differentiation

References
[1]

[2]

[3]
[4]

[5]

[6]

[7]

[8]

[9]
[10]
[11]

[12]

[13]
[14]

[15]
[16]

[17]

Z. Ou, T. Xu, Q. Su, Y. Li, P. Zhao, and Y. Bian, “Learningneural set functions under the optimal subset oracle,” Advances in Neural Information Processing Systems, vol. 35,
pp. 35 021–35 034, 2022.
H. Bonab, M. Aliannejadi, A. Vardasbi, E. Kanoulas, and
J. Allan, “Cross-market product recommendation,” in Proceedings of the 30th ACM International Conference on
Information & Knowledge Management, 2021, pp. 110–
119.
J. B. Schafer, J. Konstan, and J. Riedl, “Recommendersystems in e-commerce,” in Proceedings of the 1st ACM
conference on Electronic commerce, 1999, pp. 158–166.
X. Ning, M. Walters, and G. Karypis, “Improved machinelearning models for predicting selective compounds,” inProceedings of the 2nd ACM Conference on Bioinformatics,
Computational Biology and Biomedicine, 2011, pp. 106–
115.
Y. Saito, T. Nakamura, H. Hachiya, and K. Fukumizu, “Exchangeable deep neural networks for set-to-set matching
and learning,” in European Conference on Computer Vision,
Springer, 2020, pp. 626–646.
Q. Feng, Y. Zhou, and R. Lan, “Pairwise linear regression
classification for image set retrieval,” in Proceedings of theIEEE conference on computer vision and pattern recognition, 2016, pp. 4865–4872.
H. Zhao, L. Jiang, C.-W. Fu, and J. Jia, “Pointweb: Enhancing local neighborhood features for point cloud processing,”
in Proceedings of the IEEE/CVF Conference on Computer
Vision and Pattern Recognition (CVPR), 2019.
A. Gionis, D. Gunopulos, and N. Koudas, “Efficient andtumble similar set retrieval,” in Proceedings of the 2001
ACM SIGMOD international conference on Managementof data, 2001, pp. 247–258.
M. Mašková, M. Zorek, T. Pevnỳ, and V. Šmídl, “Deepanomaly detection on set data: Survey and comparison,”Pattern Recognition, p. 110 381, 2024.
M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R.
Salakhutdinov, and A. J. Smola, “Deep sets,” Advances inneural information processing systems, vol. 30, 2017.
J. Lee, Y. Lee, J. Kim, A. Kosiorek, S. Choi, and Y. W.
Teh, “Set transformer: A framework for attention-based
permutation-invariant neural networks,” in International
conference on machine learning, PMLR, 2019, pp. 3744–
3753.
S. Tschiatschek, A. Sahin, and A. Krause, “Differentiablesubmodular maximization,” in Proceedings of the 27th
International Joint Conference on Artificial Intelligence,
2018, pp. 2731–2738.
S. Bai, J. Z. Kolter, and V. Koltun, “Deep Equilibrium Models,” Advances in neural information processing systems,
vol. 32, 2019.
Z. Kolter, D. Duvenaud, and M. Johnson, “Deep Implicit
Layers - Neural ODEs, Deep Equilibirum Models, andBeyond,” NeurIPS, 2020. [Online]. Available: https ://implicit-layers-tutorial.org.
Z. Huang, S. Bai, and J. Z. Kolter, “(Implicit)2 : Implicit
layers for implicit representations,” NeurIPS, 2021.
A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang, Z.
DeVito, Z. Lin, A. Desmaison, L. Antiga, and A. Lerer,“Automatic differentiation in pytorch,” 2017.

[18]

[19]

[20]
[21]

[22]
[23]

[24]

[25]
[26]
[27]

[28]
[29]

[30]

[31]

[32]
[33]
[34]
[35]

5

K. P. Murphy, Machine learning: a probabilistic perspective. MIT press, 2012.
A. Mnih and G. Hinton, “Learning nonlinear constraints
with contrastive backpropagation,” in Proceedings. 2005
IEEE International Joint Conference on Neural Networks,
2005., IEEE, vol. 2, 2005, pp. 1302–1307.
G. Hinton, S. Osindero, M. Welling, and Y.-W. Teh, “Unsupervised discovery of nonlinear structure using contrastive backpropagation,” Cognitive science, vol. 30, no. 4,
pp. 725–731, 2006.
Y. LeCun, S. Chopra, and R. Hadsell, “A tutorial on energybased learning,” 2006.
A. Kirsch, S. Farquhar, P. Atighehchian, A. Jesson, F.
Branchaud-Charron, and Y. Gal, “Stochastic batch acquisition: A simple baseline for deep active learning,” Transactions on Machine Learning Research, 2023.
D. P. Kingma and M. Welling, “Auto-encoding variational
bayes,” arXiv preprint arXiv:1312.6114, 2013.
D. M. Blei, A. Kucukelbir, and J. D. McAuliffe, “Variational inference: A review for statisticians,” Journal of theAmerican statistical Association, vol. 112, no. 518, pp. 859–
877, 2017.
G. Calinescu, C. Chekuri, M. Pal, and J. Vondrák, “Maximizing a monotone submodular function subject to a matroid constraint,” SIAM Journal on Computing, vol. 40,
no. 6, pp. 1740–1766, 2011.
A. Sklar, “Random variables, joint distribution functions,and copulas,” Kybernetika, vol. 9, no. 6, pp. 449–460, 1973.
R. B. Nelsen, An Introduction to Copulas, second. New
York, NY, USA: Springer, 2006.
M. Blondel, Q. Berthet, M. Cuturi, R. Frostig, S. Hoyer,
F. Llinares-López, F. Pedregosa, and J.-P. Vert, “Efficientand modular implicit differentiation,” Advances in neural
information processing systems, vol. 35, pp. 5230–5242,
2022.
B. Davies, Exploring chaos: Theory and experiment. CRC
Press, 2018.
A. Rechnitzer, Fixed points - summary [lecture notes].
dynamical systems and chaos — 620341. 2003. [Online].
Available: https : / / personal . math . ubc . ca /~andrewr/620341/pdfs/fp_sum.pdf.
C. Szegedy, W. Zaremba, I. Sutskever, J. Bruna, D. Erhan, I.
Goodfellow, and R. Fergus, “Intriguing properties of neuralnetworks,” in 2nd International Conference on Learning
Representations, ICLR 2014, 2014.
A. Virmaux and K. Scaman, “Lipschitz regularity of deepneural networks: Analysis and efficient estimation,” Advances in Neural Information Processing Systems, vol. 31,
2018.
H. Gouk, E. Frank, B. Pfahringer, and M. J. Cree, “Regularisation of neural networks by enforcing lipschitz continuity,”
Machine Learning, vol. 110, pp. 393–416, 2021.
S. Banach, “Sur les opérations dans les ensembles abstraits
et leur application aux équations intégrales,” Fundamenta
mathematicae, vol. 3, no. 1, pp. 133–181, 1922.
S. G. Krantz and H. R. Parks, The implicit function theorem: history, theory, and applications. Springer Science &Business Media, 2002.
D. P. Kingma and J. Ba, “Adam: A Method for Stochastic
Optimization,” in ICLR, 2015.

Learning Set Functions with Implicit Differentiation
[36]

[37]
[38]

[39]
[40]
[41]

[42]

[43]

[44]

[45]

[46]

[47]

[48]

[49]

[50]

J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C.
Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, JAX: Composable transformations of Python+NumPy programs, version 0.3.13, 2018. [Online]. Available: http://github.
com/google/jax.
R. Frostig, M. J. Johnson, and C. Leary, “Compiling machine learning programs via high-level tracing,” Systems
for Machine Learning, vol. 4, no. 9, 2018.
J. Heek, A. Levskaya, A. Oliver, M. Ritter, B. Rondepierre,A. Steiner, and M. van Zee, Flax: A neural network library and ecosystem for JAX, version 0.8.5, 2023. [Online].
Available: http://github.com/google/flax.
P. Lippe, UvA Deep Learning Tutorials, https : / /uvadlc - notebooks . readthedocs . io / en /latest/, 2024.
M. Kimura, R. Shimizu, Y. Hirakawa, R. Goto, and Y. Saito,“On permutation-invariant neural networks,” arXiv preprint
arXiv:2403.17410, 2024.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones,
A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is
all you need,” Advances in neural information processing
systems, vol. 30, 2017.
L. Zhang, V. Tozzo, J. Higgins, and R. Ranganath, “Setnorm and equivariant skip connections: Putting the deep indeep sets,” in International Conference on Machine Learning, PMLR, 2022, pp. 26 559–26 574.
G. Giannone and O. Winther, “Scha-vae: Hierarchical context aggregation for few-shot generation,” in International
Conference on Machine Learning, PMLR, 2022, pp. 7550–
7569.
M. Michalkiewicz, J. K. Pontes, D. Jack, M. Baktashmotlagh, and A. Eriksson, “Implicit surface representations as layers in neural networks,” in Proceedings of theIEEE/CVF International Conference on Computer Vision(ICCV), 2019.
V. Sitzmann, J. Martel, A. Bergman, D. Lindell, and G.
Wetzstein, “Implicit neural representations with periodic
activation functions,” Advances in neural information processing systems, vol. 33, pp. 7462–7473, 2020.
Y. Strümpler, J. Postels, R. Yang, L. V. Gool, and F.
Tombari, “Implicit neural representations for image compression,” in European Conference on Computer Vision,
Springer, 2022, pp. 74–91.
S. W. Fung, H. Heaton, Q. Li, D. McKenzie, S. Osher,and W. Yin, “Jfb: Jacobian-free backpropagation for implicit networks,” in Proceedings of the AAAI Conferenceon Artificial Intelligence, vol. 36, 2022, pp. 6648–6656.
M. Retchin, B. Amos, S. Brunton, and S. Song, “Koopman
constrained policy optimization: A koopman operator theoretic method for differentiable optimal control in robotics,”
in ICML 2023 Workshop on Differentiable Almost Everything: Differentiable Relaxations, Algorithms, Operators,
and Simulators, 2023.
E. Dupont, A. Golinski, M. Alizadeh, Y. W. Teh, and A.
Doucet, “Coin: Compression with implicit neural representations,” in Neural Compression: From Information Theory
to Applications–Workshop@ ICLR 2021, 2021.
D. Grattarola and P. Vandergheynst, “Generalised implicit
neural representations,” Advances in Neural Information
Processing Systems, vol. 35, pp. 30 446–30 458, 2022.

[51]
[52]
[53]

[54]

[55]

[56]

[57]
[58]

[59]
[60]

[61]

6

E. Winston and J. Z. Kolter, “Monotone operator equilibrium networks,” Advances in neural information processing
systems, vol. 33, pp. 10 718–10 728, 2020.
R. T. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud, “Neural ordinary differential equations,” Advancesin neural information processing systems, vol. 31, 2018.
X. Li, T.-K. L. Wong, R. T. Chen, and D. Duvenaud, “Scalable gradients for stochastic differential equations,” in International Conference on Artificial Intelligence and Statistics,
PMLR, 2020, pp. 3870–3882.
B. Amos and J. Z. Kolter, “Optnet: Differentiable optimization as a layer in neural networks,” in International
Conference on Machine Learning, PMLR, 2017, pp. 136–
145.
T. Apostol, Mathematical Analysis (Addison-Wesleyseries in mathematics). Addison-Wesley, 1974, ISBN:
9780201002881. [Online]. Available: https://books.
google.com/books?id=Le5QAAAAMAAJ.
W. Rudin, Principles of Mathematical Analysis (International series in pure and applied mathematics). McGrawHill, 1976, ISBN: 9780070856134. [Online]. Available:
https : / / books . google . com / books ? id =kwqzPAAACAAJ.
J. V. Burke, “Nonlinear optimization,” Lecture Notes, Math,vol. 408, p. 80, 2014.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B.
Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss,
V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M.
Brucher, M. Perrot, and E. Duchesnay, “Scikit-learn: Machine learning in Python,” Journal of Machine Learning
Research, vol. 12, pp. 2825–2830, 2011.
Z. Liu, P. Luo, X. Wang, and X. Tang, “Deep learning
face attributes in the wild,” in Proceedings of International
Conference on Computer Vision (ICCV), 2015.
J. A. Gillenwater, A. Kulesza, E. Fox, and B. Taskar,“Expectation-maximization for learning determinantal point
processes,” Advances in Neural Information Processing
Systems, vol. 27, 2014.
J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert:
Pre-training of deep bidirectional transformers for language
understanding,” in North American Chapter of the Association for Computational Linguistics, 2019. [Online]. Available: https : / / api . semanticscholar . org /CorpusID:52967399.

Learning Set Functions with Implicit Differentiation

Supplementary Material
A. Related Work
Learning Set Functions. There is a line of work that tries to extend the capabilities of neural networks for functions
on discrete domains, i.e., set functions. Set functions are characterized by their permutation invariance. We refer
the reader to [40] for a detailed overview of permutation-invariant architectures. One such seminal architecture, Deep
Sets [10], characterizes a general family of functions for permutation invariant objectives. Following the effectiveness oftransformers [41], set transformers [11] process input sets with an attention-based neural network architecture. Zhang etal. [42] improves over both of these methods by introducing a novel set normalization layer. Ou et al. [1] study the problemof learning set functions under the supervision of optimal subset oracles with a mean-field variational approximation-based
method. We build upon their approach by using the implicit differentiation toolbox in the context of set function learning.
There are existing models that use variational auto-encoders (VAE) [22] for set outputs: SetVAE adopts the aforementioned
set transformers [11] into VAEs for point cloud generation tasks and SCHA-VAE [43] applies the same idea for few-shot
image generation. These are both generative models and not applicable to subset selection tasks.
Implicit Differentiation. Implicit representations [15], [44]–[46], layers [13]–[15], and differentiation [27], [47], [48] havegained a lot of traction in the recent years. Implicit representations [49], [50] refer to the case when high-frequency, often
discretely measured data is modeled with a continuous function, whereas implicit layers is a term that characterizes modelarchitectures. Huang et al. [15] study the complementary benefits of combining implicit representations with implicit layers.
Deep equilibrium models (DEQ) [13], [51], an example of implicit layers, directly compute the fixed-point resulting fromstacking up weight-tied forward sequence models and also directly differentiate through the stacked fixed-point equations
via implicit differentiation. A recent tutorial [14] highlights the advantages of implicit differentiation by including otherexamples from neural ordinary differential equations [52], [53] and differentiable optimization problems [54]. Blondel etal. [27] provide a JAX-based, modular automatic implicit differentiation tool in order to aid practitioners with implementation.
For feed-forward networks, Fung et al. [47] eliminate the Jacobian that results from implicit differentiation. It is important
to note that they also use Banach fixed-point theorem [33] to prove the convergence of the fixed-point iterations. However,they directly assume the mapping in focus is a contraction where we specify the condition under which the mapping in focus
is a contraction (see Asm. 3.1).

B. Derivations for Section 2
B.1. Derivation of the Fixed-Point
Rewriting the ELBO by plugging in the definition of entropy,F̃ (ψ, θ) + H(q(S, ψ)) = F̃ (ψ, θ) −|V |X

[ψi log ψi + (1 − ψi ) log(1 − ψi )] .

i=1

Taking the partial derivative of this expression with respect to the ith coordinate and setting it to zero, yields
ψi∂ F̃ (ψ, θ)
− log
= 0,
∂ψi
1 − ψiexp∂ F̃ (ψ, θ)
ψi
=
,
∂ψi
1 − ψi!
∂ F̃ (ψ, θ)∂ F̃ (ψ, θ)exp− ψi exp= ψi ,∂ψi∂ψi!
!!
∂ F̃ (ψ, θ)∂ F̃ (ψ, θ)= ψi 1 + exp
,
exp
∂ψi
∂ψi


!
(ψ,θ)
exp ∂ F̃∂ψ
1
∂ F̃ (ψ, θ)
i
=
 =σ


ψi =
,
∂ψi
1 + exp ∂ F̃ (ψ,θ)
1 + exp − ∂ F̃ (ψ,θ)
∂ψi

∂ψi

7

(13)

Learning Set Functions with Implicit Differentiation

where σ(x) = (1 + exp (−x))−1 is the sigmoid function.
B.2. Gradient Computation via Sampling
Lemma B.1. Given a set function F : 2V → R and a vector of probabilities ψ ∈ [0, 1]|V | where ψi = P[i ∈ S], thegradient of the multilinear relaxation of the set function F (S) is∂ F̃ (ψ)= ES∼ψ|ψi ←0 [F (S + i) − F (S)].
∂ψi
Proof.
X
Y Y∂ F̃ (ψ)ψi=
F (S)(1 − ψi ),
∂ψi
i∈S

S⊆V

i∈S
/

= ES∼ψ|ψi ←1 [F (S)] − ES∼ψ|ψi ←0 [F (S)],
X
Y
Y
X
=
F (S)
ψj
(1 − ψj ) −S⊆V, i∈S=

X

j∈S\{i}j ∈S/[F (S + i) − F (S)]

Y

ψj

j∈S

S⊆V \{i}Y

F (S)S⊆V \{i}

Y

j∈S

ψj

Y

(1 − ψj ),j ∈S,/ j̸=i(1 − ψj ),j ∈S/= ES∼ψ|ψi ←0 [F (S + i) − F (S)].

Corollary B.2. Knowing Lemma B.1, the gradient of the multilinear relaxation F̃ (ψ, θ), is defined as follows∂ F̃ (ψ, θ)= F̃ ([ψ]+i , θ) − F̃ ([ψ]−i , θ).
∂ψi

(14)

Eq. (14) can be computed by producing random samples of S.

C. Technical Preliminaries
Theorem C.1. (Multivariate Mean Value Theorem [55]–[57]) If f : Rn → Rm is continuously differentiable, then forevery x, y ∈ Rn , there exists a z ∈ [x, y], such that∥f (x) − f (y)∥2 ≤ sup ∥∂f (z)∥F ∥x − y∥2 ,z∈[x,y]where ∥ · ∥2 is the L2 norm and ∥ · ∥F is the Frobenius norm.
Definition C.2. A mapping T : X → X is called a contraction on X if there exists a constant ϵ ∈ [0, 1) such that for all
x, y ∈ X,∥T (x) − T (y)∥2 ≤ ϵ∥x − y∥2 .
Theorem C.3. (Banach’s Fixed Point Theorem [33]) Let T : X → X be a contraction on X. Then T has a unique fixed
point x∗ ∈ X where T (x∗ ) = x∗ .
Theorem C.4. (Implicit Function Theorem [27], [34]) Given a continuously differentiable function G : Rn × Rd → Rn , animplicitly defined function x∗ : Rd → Rn of θ ∈ Rd , and an optimal solution x∗ (θ); letG(x∗ (θ), θ) = 0.

(15)

For (x0 , θ 0 ) satisfying G(x0 , θ 0 ) = 0, if the Jacobian ∂x G evaluated at (x0 , θ 0 ) is a square invertible matrix, then there
exists a function x∗ (·) defined on a neighborhood of θ 0 such that x∗ (θ 0 ) = x0 . Furthermore, for all θ in this neighborhood,
we have that G(x∗ (θ), θ) = 0 and ∂x∗ (θ) exists. According to the chain rule, the Jacobian ∂x∗ (θ) satisfies∂x G(x∗ (θ), θ)∂x∗ (θ) + ∂θ G(x∗ (θ), θ) = 0.

(16)

∗

Therefore, computing ∂x (θ) becomes the equivalent of solving the following linear system of equations−∂x G(x∗ (θ), θ) ∂x∗ (θ) = ∂θ G(x∗ (θ), θ) .
|{z
} | {z } |{z}A∈Rn×nJ∈Rn×d8

B∈Rn×d(17)

Learning Set Functions with Implicit Differentiation

D. Proof of Theorem 3.2
Before stating our proof we need to state the following corollary:
Corollary D.1. Knowing Corollary B.2, we can write the Hessian of the multilinear relaxation as 
∂ 2 F̃ (ψ, θ) 
= F̃ ([ψ]+i,+j , θ)] − F̃ ([ψ]−i,+j , θ)] − F̃ ([ψ]+i,−j , θ)] − F̃ ([ψ]−i,−j , θ)] ,∂ψi ∂ψj(18)= F̃ ([ψ]+i,+j , θ)] − F̃ ([ψ]−i,+j , θ) − F̃ ([ψ]+i,−j , θ) + F̃ ([ψ]−i,−j , θ),2

if i ̸= j, otherwise ∂∂ψF̃ (ψ,θ)∂ψ = 0.
ij

We proof the following lemma using this corollary:
Lemma D.2. Elements of the Hessian given in Eq. (18) are bounded with 4 supψ∈[0,1] F̃ (ψ, θ) , i.e.,2supψi ,ψj ∈[0,1] ∂∂ψF̃ (ψ,θ)≤ 4 supψ∈[0,1] F̃ (ψ, θ) .
∂ψij

Proof. By Corollary D.1, we have 
∂ 2 F̃ (ψ, θ)= F̃ ([ψ]+i,+j , θ)] − F̃ ([ψ]−i,+j , θ)] − F̃ ([ψ]+i,−j , θ)] − F̃ ([ψ]−i,−j , θ)] .
∂ψi ∂ψj
Using the triangular inequality twice, we get∂ 2 F̃ (ψ, θ)≤ F̃ ([ψ]+i,+j , θ)] − F̃ ([ψ]−i,+j , θ)] + F̃ ([ψ]+i,−j , θ)] − F̃ ([ψ]−i,−j , θ)] ,∂ψi ∂ψj≤ F̃ ([ψ]+i,+j , θ)] + F̃ ([ψ]−i,+j , θ)] + F̃ ([ψ]+i,−j , θ)] + F̃ ([ψ]−i,−j , θ)] ,≤ 4 supF̃ (ψ, θ) .

ψ∈[0,1]Equipped with this lemma, we are ready to proof Thm. 3.2:

Proof. For simplicity, define a mapping Tθ : [0, 1]|V | → [0, 1]|V | where Tθ (ψ) = σ(∇ψ F̃ (ψ, θ)). Given F̃ is apolynomial w.r.t. ψ and the sigmoid is a smooth function, Tθ is continuously differentiable w.r.t. ψ in [0, 1]|V | . By themultivariate equivalent of the mean-value theorem (see Theorem C.1 in Appendix C), for every x, y ∈ |0, 1||V | , there existsa ψ ∈ |0, 1||V | , such that∥Tθ (x) − Tθ (y)∥2 ≤ sup ∥∂Tθ (ψ)∥F ∥x − y∥2 ,ψ∈[0,1]∥σ(∇ψ F̃ (x, θ)) − σ(∇ψ F̃ (y, θ))∥2 ≤ sup ∥∂ψ σ(∇ψ F̃ (ψ, θ))∥F ∥x − y∥2 .

(19)ψ∈[0,1]ih ∂ F̃ (ψ,θ) ∂ 2 F̃ (ψ,θ)
From Eq. (23), we know that ∂ψ σ(∇ψ F̃ (ψ, θ)) = σ ′
∂ψ
∂ψ ∂ψ
i

exp (−x). Then,

9

i

j

1≤i,j≤|V |where σ ′ (x) = (1 + exp (−x))−2 ·Learning Set Functions with Implicit Differentiationv
u |V | |V |uX X
sup ∥∂ψ σ(∇ψ F̃ (ψ, θ))∥F = sup tσ′ψ∈[0,1]

ψ∈[0,1]i∂ F̃ (ψ, θ)
∂ ψi

j

v
u |V | |V |
uX X
≤t

!

2

∂ 2 F̃ (ψ, θ),∂ ψi ∂ ψj!
!2∂ F̃ (ψ, θ) ∂ 2 F̃ (ψ, θ)
sup
,
∂ ψi
∂ ψi ∂ψjψi ,ψj ∈[0,1]ji!
∂ 2 F̃ (ψ, θ)
∂
F̃
(ψ,
θ)
,
= |V | supσ′∂ ψi∂ ψi ∂ ψjψi ,ψj ∈[0,1]!
∂ 2 F̃ (ψ, θ)∂ F̃ (ψ, θ)′.
≤ |V | sup σsup∂ ψi∂ ψi ∂ψjψi ∈[0,1]
ψi ,ψj ∈[0,1]

Since arg maxx∈R σ ′ (x) = 0, we know that supψi ∈[0,1] σ ′

σ′

2



∂ F̃ (ψ,θ)∂ψ i≤ 14 . Moreover, Lemma D.2 gives us a bound for≤ 4 supψ∈[0,1] F̃ (ψ, θ) . As a result,the elements of the Hessian matrix where supψi ,ψj ∈[0,1] ∂∂ψF̃ (ψ,θ)∂ψ
ijsup ∥∂ψ σ(∇ψ F̃ (ψ, θ))∥F ≤ |V | supψ∈[0,1](20)F̃ (ψ, θ) .

ψ∈[0,1]

According to Assumption 3.1, we know that supψ∈[0,1] F̃ (ψ, θ) < |V1 | . Therefore,(21)

sup ∥∂ψ σ(∇ψ F̃ (ψ, θ))∥F < 1.
ψ∈[0,1]

Plugging this in Eq. (19) above, we get(22)∥σ(∇ψ F̃ (x, θ)) − σ(∇ψ F̃ (y, θ))∥2 < ∥x − y∥2 .

This means Eq. (8) is a contraction on [0, 1]|V | (see Definition C.2 in Appendix C). Thus, according to Banach fixed-point
theorem (see Theorem C.3 in Appendix C) the iterations given in Eq. (9) are bound to converge to a unique solution.

E. Proof of Theorem 3.3
Proof. For n = |V |,A = I − ∂ψ σ(∇ψ F̃ (ψ, θ)) = I − ∂σ (∇ F̃ (ψ,θ))
..
.

...
..
.

∂σn (∇ψ F̃ (ψ,θ))
∂ψ1

...

1


=I −





=I −






=I −


ψ

∂ψ1

σ′



σ′



σ′



∂ F̃ (ψ,θ)
∂ψ1



..
.

∂ F̃ (ψ,θ)∂ψn∂ F̃ (ψ,θ)
∂ψ1

..
.

0



h

∂σ(∇ψ F̃ (ψ,θ))∂ψ1∂σ1 (∇ψ F̃ (ψ,θ))
∂ψn



...
..
.

σ′

∂ F̃ (ψ,θ)∂ψn ∂ψ1

...

σ′

...
..
.

0
..
 .



σ′

∂ F̃ (ψ,θ)∂ψn= I − Σ′ (∇ψ F̃ (ψ, θ)) · ∇2ψ F̃ (ψ, θ) ,∂ F̃ (ψ,θ)
∂ψ1



..
.

∂ F̃ (ψ,θ)
∂ψn






 
10

i

1

∂ψn

∂ F̃ (ψ,θ)∂ 2 ψ1∂σ(∇ψ F̃ (ψ,θ))∂ψn  ∂ F̃ (ψ,θ) 
∂σ
∂ψ




=I −∂σn (∇ψ F̃ (ψ,θ))
..
.



...

...

∂ψ1

..
.

∂σ



∂ F̃ (ψ,θ)
∂ψn



∂ψ1

,
...
..
.
...



∂ F̃ (ψ,θ)∂ψ1 ∂ψn ∂ F̃ (ψ,θ)∂ 2 ψn


,


∂σ



∂ F̃ (ψ,θ)
∂ψ1

∂ψn



,



∂ F̃ (ψ,θ) 
..
.

∂σ



∂ψn

∂ψn

(23)

∂ F̃ (ψ,θ)∂ 2 ψ1

..
.

...
..
.

∂ F̃ (ψ,θ)∂ψ1 ∂ψn∂ F̃ (ψ,θ)∂ψn ∂ψ1...

∂ F̃ (ψ,θ)∂ 2 ψn

..
.




,


Learning Set Functions with Implicit Differentiation ′σ (x1 )
 ..
′
n
n×n
′
where the function Σ : R → R
defined as Σ (x) =  .
0(1 + exp (−x))

−2

...
..
.
...

0
..
.
σ ′ (xn )· exp (−x).

 and σ ′ : R → R is σ ′ (x) =Similarly,
B = ∂θ σ(∇ψ F̃ (ψ, θ)) =h ∂σ (∇ F̃ (ψ,θ))
..
.

...
..
.

∂σn (∇ψ F̃ (ψ,θ))
∂θ1

...

1


=






=





=


ψ

∂θ1

σ′



σ′



σ′



∂ F̃ (ψ,θ)
∂ψ1



..
.

∂ F̃ (ψ,θ)∂ψn∂ F̃ (ψ,θ)
∂ψ1

..
.

0



∂σ(∇ψ F̃ (ψ,θ))∂θ1∂σ1 (∇ψ F̃ (ψ,θ))
∂θd



∂θd

...
..
.

σ′

∂ F̃ (ψ,θ)∂ψn ∂θ1

...

σ′

...
..
.

0
..
 .



∂ F̃ (ψ,θ)∂ψ1 ∂θ1σ′∂ F̃ (ψ,θ)∂ψn= Σ′ (∇ψ F̃ (ψ, θ)) · ∂θ ∇ψ F̃ (ψ, θ) .

∂ F̃ (ψ,θ)
∂ψ1



..
.

∂ F̃ (ψ,θ)
∂ψn



i

,

∂θ1

..
.

∂σ



∂ F̃ (ψ,θ)
∂ψn

∂σ

...
..
.

1


 
=
 

∂σn (∇ψ F̃ (ψ,θ))
..
.



...

∂σ(∇ψ F̃ (ψ,θ))∂θd  ∂ F̃ (ψ,θ) 
∂σ
∂ψ

...



∂σ

∂ F̃ (ψ,θ)

(24)

...
..
.

∂ F̃ (ψ,θ)∂ψ1 ∂θd∂ F̃ (ψ,θ)∂ψn ∂θd...

∂ F̃ (ψ,θ)∂ψn ∂θdavg(|S ∗ |)
2.5
10
100
4.52
3.80
3.87
3.26
4.14
4.62
3.18
3.8
3.69
4.52
3.16
4.09

min(|S ∗ |)
2
10
10
3
3
3
3
3
3
3
3
3
3
3
3

  ∂ψ1 ∂θ1

..

.
 

∂ψn

∂θn

∂ F̃ (ψ,θ)∂ψ1 ∂θd ∂ F̃ (ψ,θ)∂ψn ∂θd





,



∂ F̃ (ψ,θ) 
..
.




,


∂ F̃ (ψ,θ)
∂ψ1

∂θn

...

∂θ1



..
.




,


AmazonF. Dataset Details
DatasetCelebA
GaussianMoons
apparelbathbeddingcarseats
diaperfeedingfurnituregear
healthmedia
safetytoys|Ω|202, 599
100
100
100
100
100
34
100
100
32
100
62
58
36
62

|D|
10000
1000
1000
4, 675
3, 1954, 524483
6, 108
8, 202280
4, 2772, 995
1, 485267
2, 421

|V |
8
100
100
30
30
30
30
30
30
30
30
30
30
30
30

max(|S ∗ |)
3
10
10
19
11
12
6
15
2
6
10
9
19
5
14

Table 2. Dataset configurations.

Moons and Gaussian. There are two classes in the synthetic datasets, whose labels are Bernoulli sampled with p = 0.5.
Based on this label, the optimal subset and ground set pairs are constructed as follows: 1) sampling 10 points within the classas S ∗ ; and 2) sampling 90 points from the other class as V \S ∗ . This process is repeated until |Dtraining | = |Dvalidation | =|Dtest | = 1000.Following the experimental procedure of Ou et al. [1], we use the SCIKIT- LEARN package [58] to generate
the Moons dataset, consisting of two interleaving moons with some noise with varianceh σ 2 = 0.1.
i For the Gaussian dataset,
we sample data from a mixture of Gaussians 12 N (µ0 , Σ) + 12 N (µ1 , Σ), where µ0 = √12 , √12 , µ1 = −µ0 , and Σ = 14 I.
CelebA. The CelebFaces Attributes dataset (CelebA) [59] is a large-scale face dataset used extensively in computer visionresearch, particularly for tasks such as face detection, face attribute recognition, etc. The dataset contains 202, 599 face
images of 10, 177 celebrities with various poses and backgrounds. Besides, each image is annotated with 40 binary attributes,
describing facial features and properties (e.g., having a mustache, wearing a hat or glasses, etc.). Following previous work,11

Learning Set Functions with Implicit Differentiation

we select 2 attributes at random and construct the set V with a size of 8 and the oracle set S ∗ with a size of 2 or 3, where
neither attribute is present (e.g. not wearing glasses and hats).
Amazon. The Amazon Baby Registry dataset [60] includes various subsets of baby registry products chosen by customers.
These products are then organized into 18 distinct categories. From these, 12 categories are selected. Each product in thedataset is described by a textual description, which has been transformed into a 768-dimensional vector using a pre-trained
BERT model [61]. For each category, the (V, S ∗ ) pairs are sampled using the following process. First, we exclude subsets
chosen by customers that contain only one item or more than 30 items. Next, we divide the remaining subsets into training,validation, and test sets equally. For each oracle subset S ∗ ∈ S, we randomly sample 30 − |S ∗ | additional products from thesame category to ensure that V contains exactly 30 products. This method constructs a data point (V, S ∗ ) for each customer,
simulating a real-world scenario where V represents 30 products shown to the customer, and S ∗ represents the subset of
products the customer is interested in.

12