Neural Set Function Extensions: Learning with Discrete Functions in High Dimensions

Integrating functions on discrete domains into neural networks is key to developing their capability to reason about discrete objects. But, discrete domains are (1) not naturally amenable to gradient-based optimization, and (2) incompatible with deep learning architectures that rely on representations in high-dimensional vector spaces. In this work, we address both difficulties for set functions, which capture many important discrete problems. First, we develop a framework for extending set functions onto low-dimensional continuous domains, where many extensions are naturally defined. Our framework subsumes many well-known extensions as special cases. Second, to avoid undesirable low-dimensional neural network bottlenecks, we convert low-dimensional extensions into representations in high-dimensional spaces, taking inspiration from the success of semidefinite programs for combinatorial optimization. Empirically, we observe benefits of our extensions for unsupervised neural combinatorial optimization, in particular with high-dimensional representations.


Introduction
While neural networks are highly effective at solving tasks grounded in basic perception (Chen et al., 2020;Vaswani et al., 2017), discrete algorithmic and combinatorial tasks such as partitioning graphs, and finding optimal routes or shortest paths have proven more challenging. This is, in part, due to the difficulty of integrating discrete operations into neural network architectures (Battaglia et al., 2018;Bengio et al., 2021;Cappart et al., 2021a). One immediate difficulty with functions on discrete spaces is that they are not amenable to standard gradient-based training. Another is that discrete functions are typically expressed in terms of scalar (e.g., Boolean) variables for each item (e.g., node, edge to be selected), in contrast to the high-dimensional and continuous nature of neural networks' internal representations. A natural approach to addressing these challenges is to carefully choose a function on a continuous domain that extends the discrete function, and can be used as a drop-in replacement.
There are several important desiderata that such an extension should satisfy in order to be suited to neural network training. First, an extension should be valid, i.e., agree with the discrete function on discrete points. It should also be amenable to gradient-based optimization, and should avoid introducing spurious minima. Beyond these requirements, there is one additional critical consideration. In both machine learning and optimization, it has been observed that high-dimensional representations can make problems "easier". For instance, neural networks rely on high-dimensional internal representations for representational power and to allow information to flow through gradients, and performance suffers considerably when undesirable low-dimensional bottlenecks are introduced into network architectures (Belkin et al., 2019;Veličković & Blundell, 2021). In optimization, lifting to higher-dimensional spaces can make the problem more well-behaved (Goemans & Williamson, 1995;Shawe-Taylor et al., 2004;Du et al., 2018). Therefore, extending discrete functions to highdimensional domains may be critical to the effectiveness of the resulting learning process, yet remains largely an open problem.
With those considerations in mind, we propose a framework for constructing extensions of discrete set functions onto high-dimensional continuous spaces. The core idea is to view a continuous point x in space as an expectation over a distribution (that depends on x) supported on a few carefully chosen discrete points, to retain tractability. To evaluate the discrete function at x, we compute the expected value of the set function over this distribution. The method resulting from a principled formalization of this idea is computationally efficient and addresses the key challenges of building continuous extensions. Namely, our extensions allow gradient-based optimization and address the dimensionality concerns, allowing any function on sets to be used as a computation step in a neural network.
First, to enable gradient computations, we present a method based on a linear programming (LP) relaxation for constructing extensions on continuous domains where exact gradients can be computed using standard automatic differentiation software (Abadi et al., 2016;Bastien et al., 2012;Paszke et al., 2019). Our approach allows task-specific considerations (e.g., a cardinalilty constraint) to be built into the extension design. While our initial LP formulation handles gradients, and is a natural formulation for explicitly building extensions, it replaces discrete Booleans with scalars in the unit interval [0, 1], and hence does not yet address potential dimensionality bottlenecks. Second, to enable higher-dimensional representations, we take inspiration from classical SDP relaxations, such as the celebrated Goemans-Williamson maximum cut algorithm (Goemans & Williamson, 1995), which recast low-dimensional problems in high-dimensions. Specifically, our key contribution is to develop an SDP analog of our original LP formulation, and show how to lift LP-based extensions into a corresponding high-dimensional SDP-based extensions. Our general procedure for lifting low-dimensional representations into higher dimensions aligns with the neural algorithmic reasoning blueprint (Veličković & Blundell, 2021), and suggests that classical techniques such as SDPs may be effective tools for combining deep learning with algorithmic processes more generally.

Problem Setup
Consider a ground set [n] = {1, . . . , n} and an arbitrary function f : 2 [n] → R ∪ {∞} defined on subsets of [n]. For instance, f could determine if a set of nodes or edges in a graph has some structural property, such as being a path, tree, clique, or independent set (Bello et al., 2016;Cappart et al., 2021a). Our aim is to build neural networks that use such discrete functions f as an intermediate layer or loss. In order to produce a model that is trainable using standard auto-differentiation software, we consider a continuous domain X onto which we would like to extend f , with sets embedded into X via an injective map e : 2 [n] → X . For instance, when X = [0, 1] n we may take e(S) = 1 S , the Boolean vector whose ith entry is 1 if i ∈ S, and 0 otherwise. Our approach is to design an extension F : X → R of f and consider the neural network NN 2 • F • NN 1 (if f is used as a loss, NN 2 is simply the identity). To ensure that the extension is valid and amenable to automatic differentiation, we require that 1) it agrees with f on all discrete points: F(e(S)) = f (S) for all S ⊆ [n] with f (S) < ∞, and 2) F is continuous.
There is a rich existing literature on extensions of functions on discrete domains, particularly in the context of discrete optimization (Lovász, 1983;Grötschel et al., 1981;Calinescu et al., 2011;Vondrák, 2008;Bach, 2019;Obozinski & Bach, 2012;Tawarmalani & Sahinidis, 2002). These works provide promising tools to reach our goal of neural network training. Building on these, our method is the first to use semi-definite programming (SDP) to combine neural networks with set functions. There are, however, different considerations in the neural network setting as compared to optimization. The optimization literature often focuses on a class of set functions and aims to build extensions with desirable optimization properties, particularly convexity. We do not focus on convexity, aiming instead to develop a formalism that is as flexible as possible. Doing so maximizes the applicability of our method, and allows extensions adapted to task-specific desiderata (see Section 3.1).

Scalar Set Function Extensions
We start by presenting a general framework for extending set functions onto X = [0, 1] n , where a set S ⊆ [n] is viewed as the Boolean indicator vector e(S) = 1 S ∈ {0, 1} n whose ith entry is 1 if i ∈ S and 0 otherwise. We call extensions onto [0, 1] n scalar since each item i is represented by a single scalar value-the ith coordinate of x ∈ X . These scalar extensions will become the core building blocks in developing high-dimensional extensions in Section 4.
A classical approach to extending discrete functions on sets represented as Boolean indicator vectors 1 S is by computing the convex-envelope, i.e., the point-wise supremum over linear functions that lower bound f (Falk & Hoffman, 1976;Bach, 2019). Doing so yields a convex function whose value at a point x ∈ [0, 1] n is the solution of the following linear program (LP): (primal LP) The set P f of all feasible solutions (z, b) is known as the (canonical) polyhedron of f (Obozinski & Bach, 2012) and can be seen to be non-empty by taking the coordinates of z to be sufficiently small (possibly negative). Variants of this optimization program are frequently encountered in the theory of matroids and submodular functions (Edmonds, 2003) where P f is commonly known as the submodular polyhedron (see Appendix A for an extended discussion). By strong duality, we may solve the primal LP by instead solving its dual: (dual LP) whose optimal value is the same as the primal LP. The dual LP is always feasible (see e.g., the Lovász extension in Section 3.1). However, F does not necessarily agree with f on discrete points in general, unless the function is convex-extensible (Murota, 1998).
To address this important missing piece, we relax our goal from solving the dual LP to instead seeking a feasible solution to the dual LP that is an extension of f . Since the dual LP is defined for a fixed x, a feasible solution must be a function y S = p x (S) of x. If p x were to be continuous and a.e. differentiable in x then the value S p x (S)f (S) attained by the dual LP would also be continuous and a.e. differentiable in x since gradients flow through the coefficients y S = p x (S), while f (S) is treated as a constant in x. This leads us to the following definition: Definition (Scalar SFE). A scalar SFE F of f is defined at a point x ∈ [0, 1] n by coefficients p x (S) such that y S = p x (S) is a feasible solution to the dual LP. The extension value is given by and we require the following properties to hold for all S ⊆ [n]: 1) p x (S) is a continuous function of x and 2) F(1 S ) = f (S) for all S ⊆ [n].
Efficient evaluation of F requires that p x (S) is supported on a small collection of carefully chosen sets S. This choice is a key inductive bias of the extension, and Section 3.1 gives many examples with only O(n) non-zero coefficients. Examples include well-known extensions, such as the Lovász extension, as well as a number of novel extensions, illustrating the versatility of the SFE framework.
Thanks to the constraint S y S = 1 in the dual LP, scalar SFEs have a natural probabilistic interpretation. An SFE is defined by a probability distribution p x such that fractional points x can be written as an expectation E S∼px [1 S ] = x over discrete points using p x . The extension itself can be viewed as arising from exchanging f and the expectation operation: . This interpretation is summarized in Figure 1.
Scalar SFEs also enjoy the property of not introducing any spurious minima. That is, the minima of F coincide with the minima of f up to convex combinations. This property is especially important when training models of the form f • NN 1 (i.e., f is a loss function) since F will guide the network NN 1 towards the same solutions as f .
Proposition 1 (Scalar SFEs have no bad minima). If F is a scalar SFE of f then: Obtaining set solutions. Given an architecture F • NN 1 and input problem instance G, we often wish to produce sets as outputs at inference time. To do this, we simply compute x = NN 1 (G), and select the set S in supp S {p x (S)} with the smallest value f (S). This can be done efficiently if, as is typically the case, the cardinality of supp S {p x (S)} is small.

Constructing Scalar Set Function Extensions
A key characteristic of scalar SFEs is that there are many potential extensions of any given f . In this section, we provide examples of scalar SFEs, illustrating the capacity of the SFE framework for building knowledge about f into the extension. See Appendix C for all proofs and further discussion.
Lovász extension. Re-indexing the coordinates of x so that x 1 ≥ x 2 . . . ≥ x n , we define p x to be supported on the sets S 1 ⊆ S 2 ⊆ · · · ⊆ S n with S i = {1, 2, . . . , i} for i = 1, 2, . . . , n. The coefficient are defined as y Si = p x (S i ) := x i − x i+1 and p x (S) = 0 for all other sets. The resulting Lovász extension-known as the Choquet integral in decision theory (Choquet, 1954;Marichal, 2000)-is a key tool in combinatorial optimization due to a seminal result: the Lovász extension is convex if and only if f is submodular (Lovász, 1983), implying that submodular minimization can be solved in polynomial-time (Grötschel et al., 1981).
can be encoded in an n × n matrix S ∈ {0, 1} n×n whose ith column is 1 Si . In this notation, the dual LP constraint S⊆[n] y S 1 S = x can be written as Sp = x, where the ith coordinate of p defines p x (S i ). The bounded cardinality extension generalizes the Lovász extension to focus only on sets of cardinality at most k ≤ n. Again, re-index x so that x 1 ≥ x 2 . . . ≥ x n . Use the first k sets S 1 ⊆ S 2 ⊆ · · · ⊆ S k , where S i = {1, 2, . . . , i}, to populate the first k columns of matrix S. We add further n − k sets: S k+i = {j + i | j ∈ S k } for i = 1, . . . , n − k, to fill the rest of S. Finally, p x (S i ) can be analytically calculated from p = S −1 x, where S is invertible since it is a Toeplitz banded upper triangular matrix.
Permutations and involutory extensions. We use the same S, p notation. Let S be an elementary permutation matrix. Then it is involutory, i.e., SS = I, and we may easily determine p = Sx given S and x. Note that p x (S i ) = p i must be non-negative since x and S are non-negative entry-wise. Finally, restricting x to the n-dimensional Simplex guarantees that p 1 ≤ 1, which ensures p x is a probability distribution (any remaining mass is placed on the empty set). The extension property can be guaranteed on singleton sets as long as the chosen permutation admits a fixed point at the argmax of x. Any elementary permutation matrix S with such a fixed point yields a valid SFE.
Singleton extension. Consider a set function f for which f (S) = ∞ unless S has cardinality one. To ensure F is finite valued, p x must be supported only on the sets S i = {i}, i = 1, . . . , n. Assuming x is sorted so that . It is shown in Appendix C that this defines a scalar SFE, except for the dual LP feasibility. However, when using F as a loss function, minimization drives x towards the minima min x F(x) which are dual feasible. So dual infeasibility is benign in this instance and we approach the feasible set from the outside.
Multilinear extension. The multilinear extension, widely used in combinatorial optimization (Calinescu et al., 2011), is supported on all sets with coefficients p x (S) = i∈S x i i / ∈S (1 − x i ), the product distribution. In general, evaluating the multilinear extension exactly requires 2 n calls to f , but for several interesting set functions, e.g., graph cut, set cover, and facility location, it can be computed efficiently in O(n 2 ) time (Iyer et al., 2014).

Neural Set Function Extensions
This section builds on the scalar SFE framework-where each item i in the ground set [n] is represented by a single scalar-to develop extensions that use high-dimensional embeddings to avoid introducing low-dimensional bottlenecks into neural network architectures. The core motivation that lifting problems into higher dimensions can make them easier is not unique to deep learning. For instance, it also underlies kernel methods (Shawe-Taylor et al., 2004) and the lift-and-project method for integer programming (Lovász & Schrijver, 1991).
Our method takes inspiration from prior successes of semi-definite programming for combinatorial optimization (Goemans & Williamson, 1995) by extending onto X = S n + , the set of n × n positive semi-definite (PSD) matrices. With this domain, each item is represented by a vector, not a scalar.

Lifting Set Function Extensions to Higher Dimensions
We embed sets into S n + via the map e(S) = 1 S 1 S . To define extensions on this matrix domain, we translate the linear programming approach of Section 3 into an analogous SDP formulation: where we switch from lower case letters to upper case since we are now using matrices. Next, we show that this choice of primal SDP is a natural analog of the original LP that provides the right correspondences between vectors and matrices by proving that primal LP feasible solutions correspond to primal SDP feasible solutions with the same objective value (see Appendix A for a discussion on the SDP and its dual). To state the result, note that the embedding e(S) = 1 S 1 S is a particular case of the correspondence x ∈ [0, 1] n → √ x √ x .
Proposition 2. (Containment of LP in SDP) For any x ∈ [0, 1] n , define X = √ x √ x with the square-root taken entry-wise. Then, for any (z, b) ∈ R n + × R that is primal LP feasible, the pair (Z, b) where Z = diag(z), is primal SDP feasible and the objective values agree: Tr(X Z) = z x.
Proposition 2 establishes that the primal SDP feasible set is a spectrahedral lift of the positive primal LP feasible set, i.e., feasible solutions of the primal LP lead to feasible solutions of the primal SDP. As with scalar SFEs, to define neural SFEs we consider the dual SDP: We demonstrate that for suitable X this SDP has feasible solutions via an explicit construction in Section 4.2. This leads us to define a neural SFE which, as with scalar SFEs, is given by a feasible solution to the dual SDP that satisfies the extension property whose coefficients are continuous in X: Definition (Neural SFE). A neural set function extension of f at a point X ∈ S n + is defined as where y S,T = p X (S, T ) is a feasible solution to the dual SDP and for all S, T ⊆ [n]: 1) p X (S, T ) is continuous at X and 2) it is valid, i.e., F(1 S 1 S ) = f (S) for all S ⊆ [n].

Constructing Neural Set Function Extensions
We constructed a number of explicit examples of scalar SFEs in Section 3.1. For neural SFEs we employ a different strategy. Instead of providing individual examples of neural SFEs, we develop a single recipe for converting any scalar SFE into a corresponding neural SFE. Doing so allows us to build on the variety of scalar SFEs and provides an additional connection between scalar and neural SFEs. In Section 5 we show the empirical superiority of neural SFEs over their scalar counterparts.
Our construction is given in the following proposition: Then, p X defines a neural SFE F at X.
See Appendix D for proof. The choice of decomposition will give rise to different extensions. Here, we instantiate our neural extensions using the eigendecomposition of X. Since eigenvectors may not belong to [0, 1] n we reparameterize by first applying a sigmoid function before computing the scalar extension distribution p x . In practice we found that neural SFEs work just as well even without this sigmoid function-i.e., allowing scalar SFEs to be evaluated outside of [0, 1] n . The continuity of the neural SFE F when using the eigendecomposition follows from a variant of the Davis-Kahan theorem (Yu et al., 2015), which requires the additional assumption that the eigenvalues of x are distinct. For efficiency, in practice we do not use all n eigenvectors, and use only the k with largest eigenvalue. This is justified by Figure 3, which shows that in practical applications X often has a rapidly decaying spectrum.
Evaluating a neural SFE requires an accessible closed-form expression, the precise form of which depends on the underlying scalar SFE. Further, from the definition of Neural SFEs we see that if a scalar SFE is supported on sets with a property that is closed under intersection (e.g., bounded cardinality), then the supporting sets of the corresponding neural SFE will also inherit that property. This implies that the neural counterparts of the Lovász, bounded cardinality Lovász, and singleton/permutation extensions have the same support as their scalar counterparts. An immediate corollary is that we can easily compute the neural counterpart of the Lovász extension which has a simple closed form: Corollary 1. For X ∈ S n + consider the eigendecomposition X = n i=1 λ i x i x i . Let p xi be as in the Lovász extension: p xi (S ij ) = σ(x i,j ) − σ(x i,j+1 ), where σ is the sigmoid function, and x i is sorted so x i,1 ≥ . . . ≥ x i,n and S ij = {1, . . . , j}, with p xi (S) = 0 for all other sets. Then, the neural Lovász extension is: Complexity and obtaining sets as solutions. In general, the neural SFE relies on all pairwise intersections S ∩ T of the scalar SFE sets, requiring O(m 2 ) evaluations of f when the scalar SFE is supported on m sets. However, when the scalar SFE is supported on a family of sets that is closed under intersection-e.g., the Lovász and singleton extensions-the corresponding neural SFE requires only O(m) function evaluations. Discrete solutions can be obtained efficiently by returning the best set out of all scalar SFEs p xi .

Experiments
We experiment with SFEs as loss functions in neural network pipelines on discrete objectives arising in combinatorial and vision tasks. For combinatorial optimization, SFEs network training with a continuous version of the objective without supervision. For supervised image classification, they allow us to directly relax the training error instead of optimizing a proxy like cross entropy.

Unsupervised Neural Combinatorial Optimization
We begin by evaluating the suitability of neural SFEs for unsupervised learning of neural solvers for combinatorial optimization problems on graphs. We use the ENZYMES, PROTEINS, IMDB, MUTAG, and COLLAB datasets from the TUDatasets benchmark (Morris et al., 2020), using a 60/30/10 split for train/test/val. We test on two problems: finding maximum cliques, and maximum independent sets. We compare with three neural network based methods. We compare to two common approaches for backpropogating through discrete functions: the REINFORCE algorithm (Williams, 1992), and the Straight-Through estimator (Bengio et al., 2013). The third is the recently proposed probabilistic penalty relaxation (Karalias & Loukas, 2020) for combinatorial optimization objectives. All methods use the same GNN backbone, comprising a single GAT layer (Veličković et al., 2018) followed by multiple gated graph convolution layers Li et al. (2015).
In all cases, given an input graph G = (V, E) with |V | = n nodes, a GNN produces an embedding for each node: X ∈ R n×d . For scalar SFEs d = 1, while for neural SFEs we consider XX in order to produce an n × n PSD matrix, which is passed as input to the SFE F. The set function f used is problem dependent, which we discuss below. Finally, see Appendix F for training and hyper-parameter optimization details, and Appendix E for details on data, hardware, and software.
MaxClique problem is to find the largest set S that is a clique: i.e., f (S) = |S| · 1{S a clique}.
The goal is to find the largest S in the graph that is independent, i.e., f (S) = |S| · 1{S an ind. set}. MIS differs significantly from MaxClique due to its high heterophily.
Results. Table 1 displays the mean and standard deviation of the approximation ratio f (S)/f (S * ) of the solver solution S and an optimal S * on the test set graphs. The neural Lovaśz extension outperforms its scalar counterpart in 8 out of 10 cases, often by significant margins, for instance improving a score of 0.778 on PROTEINS MaxClique to 0.926. The neural SFE proved effective at boosting poor scalar SFE performance, e.g., 0.311 on ENZYMES MIS, to the competitive performance of 0.775. Neural Lovaśz outperformed or equalled and straight-through in 9 out of 10 cases, and the method of Karalias & Loukas (2020) in 6 out of 10.

Constraint Satisfaction Problems
Constraint satisfaction problems ask if there exists a set satisfying a given set of conditions (Kumar, 1992;Cappart et al., 2021b). In this section, we apply SFEs to the k-clique problem: given a graph, determine if it contains a clique of size k or more. We test on the ENZYMES and PROTEINS datasets. Since satisfiability is a binary classification problem we evaluate using F1 score. Results. Figure 2 shows that by specifically searching over sets of size k using the cardinality constrained Lovász extension from Section 3.1, we significantly improve performance compared to the Lovász extension, and REINFORCE. This illustrates the value of SFEs in allowing task-dependent considerations (in this case a cardinality constraint) to be built into extension design. During training the performance of a classifier h is typically assessed using the training error

Training Error as a Classification Objective
Since training error itself is non-differentiable, it is standard to train h to optimize a differentiable surrogate such as the cross-entropy loss. Here we offer an alternative training method by continuously extending the non-differentiable mappingŷ → 1{y i =ŷ}. This map is a set function defined on single item sets, so we use the singleton extension (definition in Section 3.1). Our goal is to demonstrate that the resulting differentiable loss function closely tracks the training error, and can be used to minimize it. We do not focus on test time generalization. Figure 6 shows the results. The singleton extension loss (left plot) closely tracks the true training error at the same numerical scale, unlike other common loss functions (see Appendix G for setup details). While we leave further consideration to future work, training error extensions may be useful for model calibration (Kennedy & O'Hagan, 2001) and uncertainty estimation (Abdar et al., 2021).

Ablations
Number of Eigenvectors. Figure 3 compares the runtime and performance of neural SFEs using only the top-k eigenvectors from the eigendecomposition X = n i=1 λ i x i x i with k ∈ {1, 2, 3, 4, 5, 6} on the maximum clique problem. For both ENZYMES and PROTEINS, performance increases with k-easily outperforming scalar SFEs and REINFORCE-until saturation around k = 4, while runtime grows linearly with k. Histograms of the eigenvalues produced by trained networks show a rapid decay in the spectrum, suggesting that the smaller eigenvalues have little effect on F.
Comparison to Naive High-Dimensional Extension. We compare neural SFEs to a naive highdimensional alternative which, given an n×d matrix X simply computes a scalar SFE on each column independently and sums them up. This naive function design is not an extension, and the dependence on the d dimensions is linearly separable, in contrast to the complex non-linear interactions between columns of X in neural SFEs. Figure 4 shows that this naive extension, whilst improving over one-dimensional extensions, performs considerably worse than neural SFEs.

Related Work
Neural combinatorial optimization Our experimental setup largely follows recent work on unsupervised neural combinatorial optimization (Karalias & Loukas, 2020;Schuetz et al., 2022;Xu et al., 2020;Toenshoff et al., 2021;Amizadeh et al., 2018), where continuous relaxations of discrete objectives are utilized. In that context, it is important to take into account the key conceptual and methodological differences of our approach. For instance, in the unsupervised Erdős goes neural (EGN) framework from Karalias & Loukas (2020), the probabilistic relaxation and the proposed choice of distribution can be viewed as instantiating a multilinear extension. As explained earlier, this extension is costly in the general case (since f must be evaluated 2 n times, and summed) but can be computed efficiently in closed form in certain cases. On the other hand, our extension framework offers multiple options for efficiently computable extensions without imposing any further conditions on the set function. For example, one could efficiently (linear time in n) compute the scalar and neural Lovász extensions of any set function with only black-box access to the function. This renders our framework more broadly applicable. Furthermore, EGN incorporates the problem constraints additively in the loss function. In contrast to that, our extension framework does not require any commitment to a specific formulation in order to obtain a differentiable loss. This provides more flexibility in modelling the problem, as we can combine the cost function and the constraints in various other ways (e.g., multiplicatively). For general background on neural combinatorial optimization, we refer the reader to the surveys (Bengio et al., 2021;Cappart et al., 2021a;Mazyavkina et al., 2021).  (Tawarmalani & Sahinidis, 2002) to barycentric approximation and interpolation schemes in computer graphics (Guessab, 2013;Hormann, 2014). Convex extensions have also been used for combinatorial penalties with structured sparsity (Obozinski & Bach, 2012, 2016, and general minimization algorithms for set functions (El Halabi & Jegelka, 2020).
Stochastic gradient estimation. SFEs produce gradients for f requiring only black-box access.
There is a wide literature on sampling-based approaches to gradient estimation, for instance the REINFORCE algorithm (Williams, 1992) (i.e., score function estimator). However, sampling introduces noise which can cause unstable training and convergence issues, prompting significant study of variance reducing control variates (Gu et al., 2017;Liu et al., 2018;Grathwohl et al., 2018;Wu et al., 2018;Cheng et al., 2020). SFEs can avoid sampling (and noise) all-together, as our extensions are differentiable and can be computed deterministically. A closely related, yet distinct, task is to produce gradients through sampling operations, which introduce non-differentiable nodes in neural network computation graphs. The Straight-Through Estimator (Bengio et al., 2013), arguably the simplest solution, treats sampling as the identity map in the backward pass, yielding biased gradient estimates. The Gumbel-Softmax trick (Maddison et al., 2017;Jang et al., 2017), provides an alternative method to sample from categorical distributions (also benefiting from variance reduction (Paulus et al., 2020a)). The trick can be seen through the lens of the more general Perturb-and-MAP framework that treats sampling as a perturbed optimization program. This framework has since been used to generalize the trick to more complex distributions (Paulus et al., 2020b) and to differentiate through the parameters of exponential families for learning and combinatorial tasks (Niepert et al., 2021). Broadly, these techniques relax a discrete distribution into a continuous one by utilizing a noise distribution and assuming access to a continuous loss function. SFEs are complementary to this setup, addressing the problem of designing continuous extensions. . Whilst sharing the algorithmic alignment motivation of SFEs, the convex programming connection is mostly cosmetic: these works directly embed solvers into network architectures, while SFEs use convex programs as an analytical tool, without requiring solver access.

Conclusion
We introduced Neural Set Function Extensions, a framework that enables evaluating set functions on continuous and high dimensional representations. We showed how to construct such extensions and demonstrated their viability in a range of tasks including combinatorial optimization and image classification. Notably, neural extensions deliver good results and improve over their scalar counterparts, further affirming the benefits of problem-solving in high dimensions. We provide anonymized code. Open source code will be released after review. (d) Did you discuss whether and how consent was obtained from people whose data you're using/curating?

A Optimization programs: extended discussion
In this section, we provide an extended discussion of the key components of our LP and SDP formulations and the relationships between them. Apart from supplying derivations, another goal of this section is to illustrate that there is in fact flexibility in the exact choice of formulation for the LP (and consequently the SDP). We provide details on possible variations as part of this discussion as a guide to users who may wish to adapt the SFE framework.
A.1 LP formulation: Derivation of the dual.
First, recall that our primal LP is defined as The dual is In order to standardize the derivation, we first convert the primal maximization problem into minimization (this will be undone at the end of the derivation). We have The Lagrangian is The optimal solution p * to the primal problem is then where d * is the optimal solution to the dual. From the Lagrangian, y S = 1.
Our proposed dual formulation is then obtained by switching from maximization to minimization and negating the objective. It can also be verified that by taking the dual of our dual, the primal is recovered (see El Halabi (2018, Def. 20) for the derivation).
A.2 Connections to submodularity, related linear programs, and possible alternatives.
Our LP formulation depends on a linear program known to correspond to the convex closure (Murota, 1998, Eq. 3.57) (convex envelope) of a discrete function. Some readers may recognize the formal similarities of this formulation with the one used to define the Lovász extension (Bilmes, 2022). Namely, for x ∈ R n we can define the Lovász Extension as where the feasible set, known as the base polytope of a submodular function, is defined as B f = {z ∈ R n : z 1 S ≤ f (S) S ⊂ [n], and z 1 S = f (S) when S = [n]}. Base polytopes are also known as generalized permutahedra and have rich connections to the theory of matroids, since matroid polytopes belong to the class of generalized permutahedra Ardila et al. (2010).
An alternative option is to consider x ∈ R n + , then the Lovász extension is given by where P f is the submodular polyhedron as defined in our original primal LP. The subtle differences between those formulations lead to differences in the respective dual formulations. In principle, those formulations can be just as easily used to define set function extensions. Overall, there are three key considerations when defining a suitable LP: • The constraints of the primal.
• The domain of the primal variables z, b and the cost x.
• The properties of the function being extended.
Below, we describe a few illustrative example cases for different choices of the above: • Adding the constraint z 1 S = f (S) when S = [n] leads to y [n] ∈ R n for the dual. This implies that the coefficients cannot be interpreted as probabilities in general which is what provides the guarantee that the extension will not introduce any spurious minima. S⊆[n] y S = 1 is just an affine hull constraint in that case. • For b = 0, the constraint S⊆[n] y S = 1 is not imposed in the dual and the probabilistic interpretation of the extension cannot be guaranteed. Examples that do not rely on this constraint include the homogeneous convex envelope (El Halabi et al., 2018) and the Lovász extension as presented above. However, even for b = 0, from the definition of the Lovász extension it is easy to see that it retains the probabilistic interpretation when x ∈ [0, 1]. • Consider a feasible set defined by P f R n + and let x ∈ R n + . If the function f is submodular, non-decreasing and normalized so that f (∅) = 0 (e.g., the rank function of a matroid), then the feasible set is called polymatroid and f is a polymatroid function. Again, in that case the Lovász extension achieves the optimal objective value (Schrijver et al., 2003, Eq. 44.32). In that case, the constraint S⊆[n] y S 1 S = x of the dual is relaxed to S⊆[n] y S 1 S ≥ x. This feasible set of the dual will allow for more flexible definitions of an extension but it comes at the cost of generality. For instance, for a submodular function that is not non-decreasing, one cannot obtain the Lovász extension as a feasible solution to the primal LP, and the solutions to this LP will not be the convex envelope in general.

A.3 SDP formulation: The geometric intuition of extensions and deriving the dual.
In order to motivate the SDP formulation, first we have to identify the essential ingredients of the LP formulation. First, the constraint S⊆[n] y S 1 S = x captures the simple idea that each continuous point is expressed as a combination of discrete ones, each representing a different set, which is at the core of our extensions. Then, ensuring that the continuous point lies in the convex hull of those discrete points confers additional benefits w.r.t. optimization and offers a probabilistic perspective.
Consider the following example. The Lovász extension identifies each continuous point in the hypercube with a simplex. Then the continuous point is viewed as an expectation over a distribution supported on the simplex corners. The value of the set function at a continuous point is then the expected value of the function over those corners under the same distribution, i.e., E S∼px [1 S ] = x leads to E S∼px [f (S)] = F(x). As long as the distribution p x can be differentiated w.r.t x, we obtain an extension that can be used with gradient-based optimization. It is clear that the construction depends on being able to identify a small convex set of discrete vectors that can express the continuous one.
This can be formulated in higher dimensions, particularly in the space of PSD matrices. A natural way to represent sets in high dimensions is through rank one matrices that are outer products of the indicator vectors of the sets, i.e., 1 S 1 S is the matrix representation of S similar to how 1 S is the vector representation. Hence, in the space of matrices, our goal will be again to identify a set of discrete matrices that represents sets that can express a matrix of continuous values.
The above considerations set the stage for a transition from linear programming to semidefinite programming, where the feasible sets are spectrahedra. Our SDP formulation attempts to capture the intuition described in the previous paragraphs while also maintaining formal connections to the LP by showing that feasible LP regions correspond to feasible SDP regions by simply projecting the LP regions on the space of diagonal matrices (see Proposition 2).
Derivation of the dual. Recall that our primal SDP is defined as We will show that the dual is 1 2 y S,T (1 S 1 T + 1 T 1 S ) and
As before, we convert the primal to a minimization problem: First, we will standardize the formulation by converting the inequality constraints into equality constraints. This can be achieved by adding a positive slack variable d S,T to each constraint such that In matrix notation this is done by introducing the positive diagonal slack matrix D to the decision variable Z, and extending the symmetric matrices in each constraint where diag(e S,T ) is a diagonal matrix where all diagonal entries are zero except at the diagonal entry corresponding to the constraint on S, T which has a 1. Using this reformulation, we obtain an equivalent SDP in standard form: Next, we form the Lagrangian which features a decision variable y S,T for each inequality, and a dual matrix variable Λ. We have For the solution to the primal p * , we have For our Lagrangian we have the dual function Thus, the dual function min y S,T − 1 = 0.
The first two conditions imply the linear matrix inequality (LMI) From the definition of A S,T we know that its additional diagonal entries will correspond to the variables y S,T . Combined with the conditions above, we arrive at the constraints of the dual y S,T ≥ 0, y S,T = 1.

This leads us to the dual formulation
y S,T = 1.
Then, we can obtain our original dual by switching to minimization and negating the objective.

B Scalar Set Function Extensions Have No Bad Minima
In this section we re-state and prove the results from Section 3. The first result concerns the minima of F, showing that the minimum value is the same as that of f , and no additional minima are added (besides convex combinations of discrete minimizers). These properties are especially desirable when using an extension F as a loss function (see Section 5) since it is important that F drive the neural network NN 1 towards producing discrete 1 S outputs. Proposition 4 (Scalar SFEs have no bad minima). If F is a scalar SFE of f then: where the last equality simply uses the fact that S⊆[n] p x (S) = 1. This proves the first claim.
To prove the second claim, suppose that x minimizes F(x) over x ∈ X . This implies that the inequality in the above derivation must be tight, which is true if and only if For a given S, this implies that either p x (S) = 0 or f (S) = min S⊆[n] f (S). Since x = E px [1 S ] = S⊆[n] p x (S) · 1 S = S:px(S)>0 p x (S) · 1 S . This is precisely a convex combination of points 1 S for which f (S) = min S⊆[n] f (S). Since F is a convex combination of exactly this set of points 1 S , we have the second claim.

C Examples of Vector Set Function Extensions
This section re-defines the vector SFEs given in Section 3.1, and prove that they satisfy the definition of an SFEs. One of the conditions we must check is that F is continuous. A sufficient condition for continuity (and almost everywhere differentiability) that we shall use for a number of constructions is to show that F is Lipschitz. A very simple computation shows that it suffices to show that x ∈ X → p x (S) is Lipschitz continuous.
Lemma 1. If the mapping x ∈ [0, 1] n → p x (S) is Lipschitz continuous and f (S) is finite for all S in the support of p x , then F is also Lipschitz continuous. In particular, F is continuous and almost everywhere differentiable.
Proof. The Lipschitz continuity of F(x) follows directly from definition: where L is the maximum Lipschitz constant of x → p x (S) over any S in the support of p x , and k is the maximal cardinality of the support of any p x .
In general k can be trivially bounded by 2 n , so F is always Lipschitz. However in may cases the cardinality of the support of any p x is much smaller than 2 n , leading too a smaller Lipschitz constant. For instance, k = n in the case of the Lovász extension.

C.1 Lovász extension.
Recall the definition: x is sorted so that x 1 ≥ x 2 ≥ . . . ≥ x d . Then the Lovász extension corresponds to taking S i = {1, . . . , i}, and letting p x (S i ) = x i − x i+1 , the non-negative increments of x (where recall we take x n+1 = 0). All other sets have zero probability. For convenience, we introduce the shorthand notation Any remaining probability mass is assigned to the empty set: p x (∅) = 1 − x 1 , which contributes nothing to the extension F since f (∅) = 0 by assumption. All that remains is to check that For a given k ∈ [n], note that the only sets S i with non-zero kth coordinate are S 1 , . . . , S k , and in all cases (1 Si ) k = 1. So the kth coordinate is precisely Extension. Consider an arbitrary S ⊆ [n]. Since we assume x = 1 S is sorted, it has the form 1 S = (1, 1, . . . , 1 k times , 0, 0, . . . 0) . Therefore, for each j < k we have a j = x j − x j+1 = 1 − 1 = 0 and for each j > k we have a j = x j − x j+1 = 0 − 0 = 0. The only non-zero probability is where the the final equality follows since by definition S k corresponds exactly to the vector (1, 1, . . . , 1 k times , 0, 0, . . . 0) = 1 S and so S k = S.
Continuity. The Lovász is a well-known extension, whose properties have been carefully studied. In particular it is well known to be a Lipschitz function Bach (2019). However, for completeness we provide a simple proof here nonetheless.
Lemma 2. Let p x be as defined for the Lovász extension. Then Proof. First note that p x is piecewise linear, with one piece per possible ordering x 1 ≥ x 2 ≥ . . . ≥ x n (so n! pieces in total). Within the interior of each piece p x is linear, and therefore Lipschitz. So in order to prove global Lipschitzness, it suffices to show that p x is continuous at the boundaries between pieces (the Lipschitz constant is then the maximum of the Lipschitz constants for each linear piece).

Now consider a point
Consider the perturbed point x δ = x − δe i with δ > 0, and e i denoting the ith standard basis vector. To prove continuity of p x it suffices to show that for any S ∈ Ω we have p x δ (S) → p x (S) as δ → 0 + .
There are two sets in the support of p x whose probabilities are different under p x δ , namely: S i = {1, . . . , i} and S i+1 = {1, . . . , i, i + 1}. Similarly, there are two sets in the support of p x δ whose probabilities are different under p x , namely: S i = {1, . . . , i−1, i+1} and S i+1 = {1, . . . , i, i+1} = S i+1 . So it suffices to show the convergence p x δ (S) → p x (S) for these four S. Consider first S i : where the final equality uses the fact that x i = x i+1 . Next consider S i+1 = S i+1 : Finally, we consider S i : completing the proof.

C.2 Bounded cardinality Lovaśz extension.
The bounded cardinality extension coefficients p x (S) are the coordinates of the vector y, where y = S −1 x and the entries (i, j) of the inverse are Equivalence to the Lovaśz extension. We want to show that the bounded cardinality extension is equivalent to the Lovaśz extension when k = n. Let T i,k = {j | (j − i) mod k = 0, for i ≤ j ≤ n, j ∈ Z + }, i.e., T i,k stores the indices where j − i is perfectly divided by k. From the analytic form of the inverse, observe that the i-th coordinate of y is p x (S i ) = j∈T i,k (x j − x j+1 ). For k = n, we have T i,n = {j | (j − i) mod n = 0} = {i}, and therefore p x (S i ) = x i − x i+1 , which are the coefficients of the Lovász extension.
Feasibility. The equation y = S −1 x guarantees that the constraint x = n i=1 y Si 1 Si is obeyed.
Recall that x is sorted in descending order like in the case of the Lovász extension. Then, it is easy to see that p is always contained in the summation for p x (S i ). Therefore, by restricting x in the probability simplex it is easy to see that To secure tight equality, we allocate the rest of the mass to the empty set, i.e., p x (∅) = 1 − n i=1 p x (S i ), which does not affect the value of the extension sicne the corresponding Boolean is the zero vector.
Extension. To prove the extension property we need to show that F(1 S ) = f (S) for all S with |S| ≤ k. Consider any such set S and recall that we have sorted 1 S with arbitrary tie breaks, such that x i = 1 for i ≤ |S| and x i = 0 otherwise. Due to the equivalence with the Lovaśz extension, the extension property is guaranteed when k = n for all possible sets. For k < n, consider the following three cases for T i,k .
• When i > |S|, T i,k = ∅ because for sorted x of cardinality at most k, we know for the coordinates that x i = x i+1 = 0. For i > k, this implies that p x (S i ) = 0.
Bringing it all together, F(1 S ) = n i=1 p x f (S i ) = p x (S)f (S) = f (S) since the sum contains only one nonzero term, the one that corresponds to i = |S|.
Continuity. Similar to the Lovaśz extension, p x in the bounded cardinality extension is piecewise linear and therefore a.e. differentiable with respect to x, where each piece corresponds to an ordering of the coordinates of x. On the other hand, unlike the Lovaśz extension, the mapping x → p x (S) is not necessarily globally Lipschitz when k < n, because it is not guaranteed to be Lipschitz continuous at the boundaries.

C.3 Singleton extension.
Feasibility. The singleton extension is not dual LP feasible. However, one of the key reasons why feasibility is important is that it implies Proposition 1, which show that optimizing F is a reasonable surrogate to f . In the case of the singleton extension, however, Proposition 1 still holds even without feasibility for f . This includes the case of the training accuracy loss, which can be viewed as minimizing the set function f ({ŷ}) = −1{y i =ŷ}.
Here we give an alternative proof of Proposition 1 for the singleton extension. Consider the same assumptions as Proposition 1 with the additional requirement that min S f (S) < 0 (this merely asserts hat S = ∅ is not a trivial solution to the minimization problem, and that the minimizer of f is unique. This is true, for example, for the training accuracy objective we consider in Section 5.
Proof of Proposition 1 for singleton extension. For x ∈ X = [0, 1] n , where the final inequality follows since min j∈[n] f (S j ) < 0. Taking x = (1, 0, 0, . . . , 0) shows that all the inequalities can be made tight, and the first statement of Proposition 1 holds. For the second statement, suppose that x ∈ X = [0, 1] n minimizes F. Then all the inequality in the preceding argument must be tight. In particular, tightness of the final inequality implies that x 1 = 1. Meanwhile, tightness of the first inequaliity implies that x i −x i+1 = 0 for all i for which f (S i ) = min j∈[n] f (S j ), and tightness of the second inequality implies that x n+1 = 0. These together imply that x = 1⊕0 n−1 where 1 is a 1 × 1 vector with entry equal to one, and 0 n−1 is an all zeros vectors of length n − 1, and ⊕ denotes concatenation. Since f (S 1 ) = min j∈[n] f (S j ) is the unique minimize we have that x = 1 S1 ∈ Hull arg min 1 S i :i∈[n] f (S i ) , completing the proof.
Extension. Consider an arbitrary i ∈ [n]. Since we assume x = 1 {i} is sorted, we are without loss of generality considering 1 {1} = (1, 0, . . . , 0, 0, . . . 0) . Therefore, we have p x (S 1 ) = x 1 − x 2 = 1 − 0 = 1 and for each j > 1 we have p x (S j ) = x j − x j+1 = 0 − 0 = 0. The only non-zero probability is p x (S 1 ), and so Continuity. The proof of continuity of the singleton extension is a simple adaptation of the proof used for the Lovaśz extension, which we omit.

C.4 Permutations and Involutory Extension.
Feasibility. It is known that every elementary permutation matrix is involutory, i.e., SS = I. Given such an elementary permutation matrix S, since S(Sx) = Sp x = x, the constraint S⊆[n] y S 1 S = x is satisfied. Furthermore, S⊆[n] y S = 1 can be secured if x is in the simplex, since the sum of the elements of a vector is invariant to permutations of the entries.
Extension. If the permutation has a fixed point at the maximum element of x, i.e., it maps the maximum element to itself, then any elementary permutation matrix with such a fixed point yields an extension on singleton vectors. Without loss of generality, let x = e 1 , where e 1 is the standard basis vector in R n . Then Se 1 = e 1 and therefore p x (e 1 ) = 1. This in turn implies F(e 1 ) = 1 · f (e 1 ). This argument can be easily applied to all singleton vectors.
Continuity. The permutation matrix S can be chosen in advance for each x in the simplex. Since p x = Sx, the probabilities are piecewise-linear and each piece is determined by the fixed point induced by the maximum element of x. Consequently, p x depends continuously on x.
have for any S, T ⊆ [n] that showing SDP feasibility. That the objective values agree is easily seen since: Next, we provide a proof for the construction of neural extensions. Recall the statement of the main result.
Proposition. Let p x induce a scalar SFE of f . For X ∈ S n + with distinct eigenvalues, consider the decomposition X Then, p X defines a neural SFE F at X.
Proof. We begin by showing through the eigendecomposition of X that the F defined by p X (S, T ) is dual SDP feasible. It is clear that S,T p X (S, T ) = 1 as long as n i=1 λ i = 1, which can be easily enforced by appropriate normalization of X. Recall from the eigendecomposition we have X = n i=1 λ i v i v i where we have fixed each v i ∈ [0, 1] n through a sigmoid. Using the scalar SFE p x we may write each v i as a convex combination v i = S p vi (S)1 S . For each i we may use this representation to re-express the outer product of v i with itself: Summing over all eigenvectors v i yields the relation X = S,T ⊆[n] p X (S, T )(1 S 1 T + 1 T 1 S ), proving dual SDP feasibility.
Next, consider an input X = 1 S 1 S . In this case, the only eigenvector is 1 S with eigenvalue λ = |S| since X1 S = 1 S (1 S 1 S ) = 1 S |S|. That is, p X (T , T ) = p 1 S (T )p 1 S (T ).
For X = 1 S 1 S , 1 S is clearly an eigenvector with eigenvalue λ = |S| because X1 S = 1 S (1 S 1 S ) = 1 S |S|. So, taking1 S = 1 S / |S| to be the normalized eigenvector of X, we have X = |S|1 S1 S = = p X (S, S)1 S 1 S for p X (S, S) = 1. Therefore, the corresponding neural SFE is All that remains is to show continuity of neural SFEs. Since the scalar SFE p x is continuous in x by assumption, all that remains is to show that the map sending X to its eigenvector with i-th largest eigenvalue is continuous. We handle sign flip invariance of eignevectors by assuming a standard choice for eigenvector signs-e.g., by flipping the sign where necessary to ensure that the first non-zero coordinate is greater than zero. The continuity of the mapping X → v i follows directly from Theorem 2 from Yu et al. (2015), which is a variant of the Davis-Kahan theorem. The result shows that the angle between the i-th eigenspaces of two matrices X and X goes to zero in the limit as X → X . We shall open source our code with MIT License, and have provided anonymized code as part of the supplementary material for reviewers.

E.2 Data Details
This paper uses five graph datasets: ENZYMES, PROTEINS, IMDB-BINART, MUTAG, and COL-LAB. All data is accessed via the standardized PyG API. In the case of COLLAB, which has 5000 samples available, we subsample the first 1000 graphs only for training efficiency. All experiments Use a train/val/test split ratio of 60/30/10, which is done in exactly one consistent way across all experiments for each dataset.

F Unsupervised Neural Combinatorial Optimization Experiments
All methods use the same GNN backbone: a combination of GAT Veličković et al. (2018)

F.1 Discrete Objectives
Maximum Clique. For the maximum clique problem, we could simply take f to compute the clique size (with the size being zero if S is not a clique). However, we found that this objective led to poor results and unstable training dynamics. So, instead, we select a discrete objective that yielded the much more stable results across datasets. It is defined for a graph G = ([n], E) as, where w is a measure of size of S and q measures the density of edges within S (i.e., distance from being a clique). The scalar c is a constant, taken to be c = 2 in all cases except REINFORCE for which c = 2 proved ineffective, so we use c = 4 instead. Specifically, w(S) = i,j∈S 1{(i, j) ∈ E} simply counts up all the edges between nodes in S, and q(S) = −2w(S)/(|S| 2 − |S|) is the ratio (with a sign flip) between the number of edges in S, and the number of undirected edges (|S| 2 −|S|)/2 there would be in a clique of size |S|. If G were directed, simply remove the factor of 2. Note that this f is minimized when S is a maximum clique.
Maximum Independent Set. Similarly for maximum independent set we use the discrete objective, where w is a measure of size of S and q measures the number of edges between nodes in S (the number should be zero for an independent set), and c = 2 as before. Specifically, we take w(S) = |S|/n, and q(s) = 2 i,j∈S 1{(i, j) ∈ E}/(|S| 2 − |S|), as before.
All Neural SFEs, unless otherwise stated, use the top k = 4 eigenvectors corresponding to the largest eigenvalues. This is an important efficiency saving step, since with k = n, i.e., using all eigenvectors, the resulting Neural Lovaśz extension requires O(n 2 ) set function evaluations, compared to O(n) for the scalar Lovaśz extension. By only using the top k we reduce the number of evaluations to O(kn). Wall clock runtime experiments given in Figure 3 show that the runtime of the Neural Lovaśz extension is around ×k its scalar counterpart, and that the performance of the neural extension gradually increases then saturates when k gets large. To minimize compute overheads we pick the smallest k at which performance saturation approximately occurs.
Instead of calling the pre-implemented PyTorch eigensolver torch.linalg.eigh, which calls LAPACK routines, we use the power method to approximate the first k eignevectors of X. This is because we found the PyTorch function to be too numerically unstable in our case. In contrast, we found the power method, which approximates eigenvectors using simple recursively defined polynomials of X, to be significantly more reliable. In all cases we run the power method for 5 iterations, which we found to be sufficient for convergence.

F.3 Baselines.
This section discusses various implementation details of the baseline methods we used. The basic training pipeline is kept identical to SFEs, unless explicitly said otherwise. Namely, we use nearly identical model architectures, identical data loading, and identical HPO parameter grids.

REINFORCE.
We compared with REINFORCE (Williams (1992)) which enables backpropagation through (discrete) black-box functions. We opt for a simple instantiation for the score estimator where p(S|θ) = i∈S p i j / ∈S (1 − p j ), i.e., each node is selected independently with probability p i = g θ (y) for i = 1, 2, . . . , n, where g θ is a neural network and y some input attributes. We maximize the expected reward, i.e., For all experiments with REINFORCE, the expected reward is computed over 250 sampled actions S which is approximately the number of function evaluations of neural SFEs in most of the datasets. Here, f is taken to be the corresponding discrete objective of each problem (as described earlier in section F.1). For maximum clique, we normalize rewards f (S) by removing the mean and dividing by the standard deviation. For the maximum independent set, the same strategy led to severe instability during training. To alleviate the issue, we introduced an additional modification to the rewards: among the sampled actions S, only the ones that achieved higher than average reward were retained and the rewards of the rest were set to 0. This led to more stable results in most datasets, with the exception of COLLAB were the trick was not sufficient.
These issues highlight the instability of the score function estimator in this kind of setting. Additionally, we experimented by including simple control variates (baselines). These were: i) a simple greedy baseline obtained by running a greedy algorithm on each input graph ii) a simple uniform distribution baseline, where actions S were sampled uniformly at random. Unfortunately, we were not able to obtain any consistent boost in either performance or stability using those techniques. Finally, to improve stability, the architectures employed with REINFORCE were slightly modified according to the problem. For example, for the independent set we additionally applied a sigmoid to the outputs of the final layer.
Erdos Goes Neural. We compare with recent work on unsupervised combinatorial optimization (Karalias & Loukas, 2020). We use the probabilistic methodology described in the paper to obtain a We omit additive constants as in practice they not affect the optimization. For the maximum independent set, we follow the methodology from the paper to derive the following loss: β was tuned through a simple line search over a few possible values in each case. Following the implementation of the original paper, we use the same simple decoding algorithm to obtain a discrete solution from the learned probabilities.
Straight Through Estimator. We also compared with the Straight-Through gradient estimator (Bengio et al., 2013). This estimator can be used to pass gradients through sampling and thresholding operations, by assuming in the backward pass that the operation is the identity. In order to obtain a working baseline with the straight-through estimator, we generate level sets according to the ranking of elements in the output vector x of the neural network. Specifically, given x ∈ [0, 1] n outputs from a neural network, we generate indicator vectors 1 S k , where S k = {j| x j ≥ x k } for k = 1, 2, . . . , n. Then our loss function was computed as where f is the corresponding discrete objective from section F.1. At inference, we select the set that achieves the best value in the objective while complying with the constraints.
Ground truths. We obtain the maximum clique size and the maximum independent set size s for each graph by expressing it as a mixed integer program and using the Gurobi solver (Gurobi Optimization, LLC, 2021).

F.4 k-Clique Constraint Satisfaction
Ground truths. As before, we obtain the maximum clique size s for each graph by expressing it as a mixed integer program and using the Gurobi solver (Gurobi Optimization, LLC, 2021). This is converted into a binary label 1{s ≥ k} indicating if there is a clique of size k or bigger.
Implementation details. The training pipeline, including HPO, is identical to the MaxClique setup. The only difference comes in the evaluation-at test time the GNN produces an embedding x, and the largest clique S in the support of p x is selected. The model prediction for the constraint satisfaction problem is then 1{|S| ≥ k}, indicating whether the GNN found a clique of size k or more. Since this problem is. binary classification problem we compute the F1-score on a validation set, and report as the final result the F1-score of that same model on the test set.

G Training error as an objective
Recall that for a K-way classifier h : X → R K withŷ(x) = arg max k=1,...,K h(x) k , we consider the training error 1 n n i=1 1{y i =ŷ(x i )} calculated over a labeled training dataset {(x i , y i )} n i=1 to be a discrete non-differentiable loss. The set function in question is y → 1{y i = y}, which we relax using the singleton method described in Section 3.1.
Training details. For all datasests we use a standard ResNet-18 backbone, with a final layer to output a vector of the correct dimension depending on the number of classes in the dataset. CIFAR10 and tinyImageNet models are trained for 200 epochs, while SVHN uses 100 (which is sufficient for convergence). We use SGD with momentum mom = 0.9 and weight decay wd = 5 × 10 −4 and a cosine learning rate schedule. We tune the learning rate for each loss via a simple grid search of the values lr ∈ {0.01, 0.05, 0.1, 0.2}. For each loss we select the learning rate with highest accuracy on a validation set, then display the training loss and accuracy for this run.

H Pseudocode: A forward pass of Scalar and Neural SFEs
To illustrate the main conceptual steps in the implementation of SFEs, we include two torch-like pseudocode examples for SFEs, one for scalar and one for neural SFEs. The key to the practical implementation of SFEs within PyTorch is that it is only necessary to define the forward pass. Gradients are then handled automatically during the backwards pass.
Observe that in both Algorithm, 1 and Algorithm 2, there are two key functions that have to be implemented: i) getSupportSets, which generates the sets on which the extension is supported. ii) getCoeffs, which generates the coefficients of each set. Those depend on the choice of the extension and have to be implemented from scratch whenever a new extension is designed. The sets of the neural extension and their coefficients can be calculated from the corresponding scalar ones, using the definition of the Neural SFE and Proposition 3.

I.1 Limitations and Future Directions
Our SFEs have proven useful for learning solvers for a number of combinatorial optimization problems. However there remain many directions for improvement. One direction of particular interest is to scale our methods to instances with very large n. This could include simply considering larger graphs, or problems with larger ground sets-e.g., selecting paths. We believe that a promising approach to this would be to develop localized extensions that are supported on sets corresponding to suitably chosen sub-graphs, which would enable us to build in additional task-specific information about the problem.

I.2 Broader Impact
Our work focuses on a core machine learning methodological goal of designing neural networks that are able to learn to simulate algorithmic behavior. This program may lead to a number of promising improvements in neural networks such as making their generalization properties more reliable (as with classical algorithms) and more interpretable decision making mechanisms. As well as injecting algorithmic properties into neural network models, our work studies the use of neural networks for solving combinatorial problems. Advances in neural network methods may lead to advances in numerical computing more widely. Numerical computing in general-and combinatorial optimization in particular-impacts a wide range of human activities, including scientific discovery and logistics planning. Because of this, the methodologies developed in this paper and any potential further developments in this line of work are intrinsically neutral with respect to ethical considerations; the main responsibility lies in their ethical application in any given scenario.