Learning Neural Set Functions Under the Optimal Subset Oracle
Anonymous Author(s)
Abstract
Learning set functions becomes increasingly more important in many applicationslike product recommendation and compound selection in AI-aided drug discovery.
The majority of existing works study methodologies of set function learning underthe function value oracle, which, however, requires expensive supervision signals.
This renders it impractical for applications with only weak supervisions under the
Optimal Subset (OS) oracle, the study of which is surprisingly overlooked.
In thiswork, we present a principled yet practical maximum likelihood learning framework, termed as EquiVSet , 1 that simultaneously meets the following desiderata oflearning set functions under the OS oracle: i) permutation invariance of the set massfunction being modeled; ii) permission of varying ground set; iii) minimum prior;and iv) scalability.
The main components of our framework involve: an energybased treatment of the set mass function, DeepSet-style architectures to handlepermutation invariance, mean-field variational inference, and its amortized variants.
Thanks to the elegant combination of these advanced architectures, empirical studies on three real-world applications (including Amazon product recommendation,set anomaly detection and compound selection for virtual screening) demonstratethat EquiVSet outperforms the baselines by a large margin.
Introduction
Many real-world applications involve prediction of set-value outputs, such as recommender systemswhich output a set of products to customers, anomaly detection that predicts the outliers from themajority of data (Zhang et al., 2020) and compound selection for virtual screening in drug discoveryaims at extracting the most effective compounds from a given compound database (Gimeno et al.
All of these applications implicitly learn a set function (Rezatofighi et al., 2017; Zaheer et al.
,2017) that measures the utility of a given set input, such that the most desirable set output has thehighest (or lowest w.l.o.g ) utility value.
More formally, consider a recommender system: given a set of product candidates V , it is expectedto recommend a subset of products S ∗ ⊆
V to the user, which would satisfy the user most, i.e. ,offering the maximum utility to the user.
We assume the underlying process of determining S ∗ canbe modelled by a utility function F θ ( S ; V ) parameterized by θ , and the following criteria:
[Equation]
There are mainly two settings for learning the utility function.
The first one, namely function value(FV) oracle, targets at learning F θ ( S ; V ) to fit the utility explicitly, under the supervision of datain the form of { ( S i , f i ) } for a fixed ground set V , where f i is the true utility function value of the subset S i .
However, training in this way is prohibitively expensive, since one needs to constructlarge amounts of supervision signals for a specific ground set V (Balcan & Harvey, 2018).
Here weconsider an alternative setting, which learns F θ ( S ; V ) in an implicit way.
More formally, with thedata in form of { ( V i , S ∗ i ) } Ni =1 , where S ∗ i is the optimal subset (OS) corresponding to V i , our goal isto estimate θ such that for all possible ( V i , S ∗ i ) , it satisfies equation (1).
The OS oracle is arguablymore practical than the FV oracle, which alleviates the need for explicitly labeling utility values for alarge amount of subsets.
Though being critical for practical success, related study on set utility function learning under the
OS supervision oracle is surprisingly lacked.
The most relevant work is the probabilistic greedymodel (PGM), which solves optimization (1) with a greedy maximization algorithm (Tschiatscheket al., 2018).
Specifically, PGM interprets the maximization algorithm as to construct differentiabledistributions over sequences of items in an auto-regressive manner.
However, such construction ofdistributions is problematic for defining distributions on sets due to the dependency on the samplingorder.
Therefore, they alleviate this issue by enumerating all possible permutations of the samplingsequence (detailed discussion is given in Appendix A).
Such enumerations scale poorly due to thecombinatorial cost O ( | V | !) , which hinders PGM’s applicability to real-world applications.
To learn set functions under the OS oracle, we advocate the maximum likelihood paradigm (Stigler,1986).
Specifically, this learning problem can be viewed from a probabilistic perspective
[Equation]
where the constraint admits the learned set function to obey the objective defined in (1).
Given limiteddata { ( V i , S ∗ i ) } Ni =1 sampled from the underlying data distribution P ( V, S ) , one would maximize theempirical log likelihood: (cid:80) Ni =
[log p θ ( S ∗ i | V i )] .
The most important step is to construct a properset distribution p θ ( S | V ) whose probability mass monotonically grows with the utility function
F θ ( S ; V ) and satisfy the following additive requirements: (i) permutation invariance : the probabilitymass should not change under any permutation of the elements in S ;
(ii) varying ground set : thefunction should be able to process input sets of variable size; iii) minimum prior : we shouldmake no assumptions of the set probability, i.e., with maximum entropy, which is equivalent to theuninformative prior (Jeffreys, 1946); and iv) scalibility : the learning algorithm should be scalable tolarge-scale datasets and run in polynomial time.
In this paper, we propose Equi variant V ariational inference for Set function learning ( EquiVSet ),a new method for learning set functions under the OS oracle, which satisfies all the requirements.
Specifically, we use an energy-based model (EBM) to construct the set mass function.
EBMs aremaximum entropy distributions, which satisfies the minimum prior requirement.
Moreover, bymodeling the energy function with DeepSet-style architectures (Zaheer et al., 2017; Lee et al., 2019),the two requirements, i.e., permutation invariance and varying ground set are naturally satisfied.
Unfortunately, the flexibility of EBMs exacerbates the difficulties of learning and inference, since theinputs of set are discrete and lie in an exponentially-large space.
To remedy this issue, we developan approximate maximum likelihood approach which estimates the marginals via the mean-fieldvariational inference, resulting in an efficient training manner under the supervision of OS oracles.
In order to ensure scalability , an amortized inference network with permutation equivariance isproposed, which allows the model to be trained on large-scale datasets.
Although it may be seen as combining existing components in approximate inference, the proposedframework addresses a surprisingly overlooked problem in the set function learning communitiesusing an intuitive yet effective method.
Our main contributions are summarized below:
• We formulate set functions learning problems under the OS supervision oracle using the
maximum likelihood principle;
• We present an elegant framework based on EBMs which satisfies the four desirable require-
ments and is efficient both at training and inference stages;
Real-world experiments demonstrate effectiveness of the proposed OS learning framework
Energy-Based Modeling for Set Function Learning
The first step to solve problem (2) is to construct a proper set mass function p θ ( S | V ) monotonicallygrowing with the utility function F θ ( S ; V ) .
There exits countless ways to construct such a probabilitymass function, such as the sequential modeling in PGM (Tschiatschek et al., 2018, Section 4).
Herewe resort to the energy-based treatment:
[Equation]
where the utility function F θ ( S ; V ) stands for the negative energy, with higher utility representinglower energy.
The energy-based treatment is attractive, partially due to its maximum entropy ( i.e. ,minimum prior) property.
That is, it assumes nothing about what is unknown, which is known asthe “noninformative prior” principle in Bayesian modeling (Jeffreys, 1946).
This basic principleis, however, violated by the set mass function defined in PGM.
We refer detailed motivation of theenergy-based modeling to Appendix B.1.
In addition to the minimum prior , the energy-based treatment also enables the set mass functionp θ ( S | V ) to meet the other two requirements, i.e. permutation invariance and varying ground set , bydeliberately designing a suitable set function F θ ( S ; V ) .
However, modeling such a proper function isnontrivial, since classical feed-forward neural networks (e.g., the ones designed for submodular setfunctions (Bilmes & Bai, 2017)) violate both two criteria, which restricts their applicability to theproblems involving a set of objects.
Fortunately, Zaheer et al.
(2017) sidestep this issue by introducinga novel architecture, namely DeepSet.
They theoretically prove the following Proposition.
Proposition 1.
All permutation invariant set functions can be decomposed in the form f ( S )
=
ρ (cid:0)(cid:80) s ∈ S κ ( s ) (cid:1) , for suitable transformations κ and ρ .
By combining the energy-based model in (3) with this architecture, we could construct a validset mass function to meet two important criteria: permutation invaraince and varying ground set .
However, the flexibility of EBMs exacerbates the difficulties of learning and inference, since thepartition function Z is typically intractable and the input of sets is undesirably discrete.
Approximate Maximum Likelihood Learning with OS Supervision Oracle
In this section, we explore an effective framework for learning set functions under the supervision ofoptimal subset oracles.
We start with discussing the principles for learning parameter θ , followed bydiscussing the detailed inference method for discrete EBMs.
3.1 Training Discrete EBMs Under the Guidance of Variational Approximation
For discrete data, e.g., set, learning the parameter θ in (3) via maximum likelihood is notoriouslydifficult.
Although one could apply techniques, such as ratio matching (Lyu, 2012), noise contrastiveestimation (Tschiatschek et al., 2016), and contrastive divergence (Carreira-Perpinan & Hinton, 2005),they generally suffer from instability on high dimensional data, especially when facing very largeground set in real-world applications.
Instead of directly maximizing the log likelihood, we consideran alternative optimization objective that is computationally preferable.
Specifically, we first fit avariational approximation to the EBM by solving
[Equation]
4) where D ( ·||· ) is a discrepancy measure between two distributions, p θ ( S ) 3 is the EBM defined in (3),and q ( S ; ψ ) denotes the mean-field variational distribution with the parameter ψ ∈ [0 , 1] | V | standingfor the odds that each item s ∈ V shall be selected in the optimal subset S ∗ . Note that the optimalparameter ψ ∗ of (4) can be viewed as a function of θ .
In this regard, we can optimize the parameter the condition V for brevity.
In some specific context, it would be helpful to regard subset S as a binary vector, i.e. , S :
{+= { 0 , 1 } | V | with the i -th element equal to 1 meaning i ∈ S and 
+}
meaning i ̸∈ S .
1: for k ← 1 , . . .
, K do 2: for i ← 1 , . . .
, | V | in parallel do 3: sample m subsets
S n ∼ q
S ; ( ψ ( k − 1)
| ψ ( k − 1) i ← 0))
4: update variational parameter ψ ( k ) i ← σ ( 1 m m (cid:80) n =1 [ F θ ( S n + i ) − F θ ( S n )]) 5: end for 6: end for
[Figure] [Figure]
θ by minimizing the following cross entropy loss, 4 which is well-known to be implementing themaximum likelihood estimation (Goodfellow et al., 2016) w.r.t.
the surrogate distribution q ( S ; ψ ∗ ) ,
L ( θ ; ψ ∗ ) =
E [Equation]
[Equation] [Equation] [Equation] [Equation]
1 N
N (cid:88) i =  [Equation] (cid:88)
j [Equation] S ∗ i log ψ ∗ j − (cid:88)
j ∈ V i \ S ∗ i log(1 − ψ ∗ j ) 
 .
This is also known as the marginal-based loss (Domke, 2013), which trains probabilistic modelsby evaluating them using the marginals approximated by an inference algorithm.
Despite notexactly bounding the log-likelihood of (3), this objective, as pointed out by Domke (2013), benefitsfrom taking the approximation errors of inference algorithm into account while learning.
However,minimizing (5) requires the variational parameter ψ ∗ being differentiable w.r.t.
θ .
Inspired bythe differentiable variational approximation to the Markov Random Fields (Krähenbühl & Koltun,2013; Zheng et al., 2015; Dai et al., 2016), below, we extend this method to the deep energy-basedformulation, which admits an end-to-end training paradigm with the back-propagation algorithm.
3.2 Differentiable Mean Field Variational Inference
To solve the optimization problem (4), we need to specify the variational distribution q ( S ; ψ ) andthe divergence measure D ( ·||· ) , such that the optimum marginal ψ ∗ is differentiable w.r.t.
themodel parameter θ .
A natural choice is to restrain q ( S ; ψ ) to be fully factorizable, which leads toa mean-field approximation of p θ ( S ) .
The simplest form of q ( S ; ψ ) would be a | V | independent
Bernoulli distribution, i.e. , q ( S ; ψ ) = (cid:81) i ∈
S ψ i (cid:81) i ̸∈ S (1 − ψ i ) , ψ ∈ [0 , 1] | V | .
Further restricting thediscrepancy measure D ( q || p ) to be the Kullback-Leibler divergence, we recover the well-knownmean-field variational inference method.
It turns out that minimizing the KL divergence amounts tomaximizing the evidence lower bound (ELBO)
[Equation]
where f F θ mt ( ψ ) is the multilinear extension of F θ ( S ) (Calinescu et al., 2007), which is defined as
[Equation]
To maximize the ELBO in (6), one can apply the fixed point iteration algorithm.
Specifically,for coordinate ψ i , the partial derivative of the multilinear extension is ∇ ψ i f F θ mt ( ψ ) , and for theentropy term, it is ∇ ψ i H ( q )
= log 1 − ψ i ψ i .
Thus, the stationary condition of maximizing ELBO isψ i =
σ ( ∇ ψ i f F θ mt ( ψ ))
, i = 1 , . .
, | V | , where σ is the sigmoid function, which means ψ i should beupdated as ψ i ← σ ( ∇ ψ
i f F θ mt ( ψ )) .
This analysis leads to the traditional mean field iteration, whichupdates each coordinate one by one (detailed derivation in Appendix B.2).
In this paper, we suggest
[Equation]
to update ψ in a batch manner, which is more efficient in practice.
More specifically, we summerizethe mean field approximation as the following fixed-point iterative update steps
[Equation]
We denote the above iterative steps as a function termed as MFVI( ψ , V, K ) , which takes initialvairational parameter ψ , ground set V , and number of iteration steps K as input, and outputs theparameter ψ ∗ after K steps.
Note that, MFVI( ψ , V, K ) is differentiable w.r.t.
the parameter θ , sinceeach fixed-point iterative update step is differentiable.
Thereby, one could learn θ by minimizingthe cross entropy loss in (5).
However, the computation complexity raises from the derivative ofmultilinear extension f F θ mt ( ψ ) defined in (7), which sums up all the possible subsets in the space of size
| V | .
Fortunately, the gradient ∇ ψ f F θ mt can be estimated efficiently via Monte Carlo approximationmethods, since the following equation holds.
[Equation]
(11) in which we use S + i to denote the set union S ∪{ i } .
Detailed derivation is provided in Appendix B.3.
According to (11), we can estimate the partial derivative ∇ ψ i f F θ mt via Monte Carlo approximation:i) sample m subsets S n , n = 1 , . . .
, m from the surrogate distribution q ( S ; ( ψ | ψ i ← 0)) ; ii)approximate the expectation by the average 1 m (cid:80) nk =
[ F θ ( S n + i ) − F θ ( S n )] .
After training, the OSfor a given ground set can be sampled via
S =
topN( ψ ∗ ) , (12) where the function topN( x ) returns the indexes of top N values in the vector x , and ψ ∗ =
MFVI( ψ , V, K ) denotes the optimal variational parameter after K -step mean-field iteration.
Weterm this method as Diff erentiable M ean F ield ( DiffMF ) and summarize the complete training andinference process in Algorithm 2 and 1, respectively.
Amortizing Inference with Equivariant Neural Networks
Although DiffMF can learn set function F θ in an effective way, it undesirably has two notoriousissues: i) the computation is in general prohibitively expensive, since DiffMF involves a typicallyexpensive sampling loop per data point; ii) some information regarding interactions between elementsis discarded, since DiffMF assumes a fully fatorizable variational distribution.
In this section, wefirst propose to amortize the inference process with an additional recognition neural network, andthen extend it to considering correlation for more accurate approximations.
4.1 Equivariant Amortized Variational Inference
To enable training the proposed model on a large-scale dataset, we propose to amortize the approximate inference process with an additional recognition neural network which outputs parameter ψ forthe variational distribution q ϕ ( S ; ψ ) , 5 where ϕ denotes the parameter of neural networks.
A properrecognition network involving set objects shall satisfy permutation equivariance .
Definition 1.
A function f : X d → Y d is called permutation equivalent when upon permutation ofthe input instances permutes the output labels, i.e., for any permutation π : f ( π ([ x 1 , . . .
, x d ]))
=
π ( f ([ x 1 , . . .
, x d ])) .
Zaheer et al.
(2017) propose to formulate the permutation equivariant architecture as :
[Equation]
(13) where s i denotes the i th element in the set S , λ, γ are learnable scalar variables, and ρ, κ areany proper transformations.
Note that the output value of f i : 2 V → [0 , 1] is relative to the i thcoordinate, but not the order of the elements in S .
Thus the equivariant recognition network, denotedas ψ = EquiNet( V ; ϕ ) : 2 V → [0 , 1] | V | , can be defined as EquiNet i :
=
f i , which takes the groundset V as input and outputs the distribution parameter ψ for q ϕ ( S ; ψ ) .
EquiNet ( 𝑉 ; 𝜙 ) 𝑉 𝜙 (
[Figure]
) 𝜓 (") 𝜃
MFVI ( 𝜓 (") , 𝑉 , 𝐾 ) 𝜓 $ ( 𝜃 )
MLE on 𝑞 % ( 𝑆 ; 𝜓 $ ( 𝜃 ))
[Figure]
4.2 Correlation-aware Inference with Gaussian Copula
Due to the mean-field assumption, the proposed variational distribution cannot model the interactionsamong elements in the input set.
We address this issue by introducing Gaussian copula (Nelsen,2007), which is a cumulative distribution function (CDF) of random variables ( u 1 , . . .
, u | V | ) overthe unit cube [0 , 1] | V | , with u i ∼ Uniform(0 , 1) .
More formally, given a covariance matrix Σ , the
Gaussian copula C Σ with parameter Σ is defined as
[Equation]
where Φ Σ stands for the joint CDF of a Gaussian distribution with zero mean and covariance matrixΣ , and Φ − 1 is the inverse CDF of standard Gaussian.
With the location parameter ψ output by
EquiNet( V ; ϕ ) , we can induce correlation into the Bernoulli distribution via the following way: i)
sample an auxiliary noise g ∼ N ( 0 , Σ ) ; ii) apply element-wise Gaussian CDF u
= Φ diag( Σ ) ( g ) ;
iii) obtain binary sample via s =
I ( ψ ≤ u ) , 6 where ψ ≤ u means ∀
i
, ψ i ≤ u i , I ( · ) is the indicatorfunction, and diag( Σ ) returns the diagonal matrix of Σ .
In practice, the covariance matrix Σ couldbe generated by another neural network with the input ground set.
We refer the discussion on it to
Appendix C, and demonstrate how to efficiently construct and sample from a non-diagonal Gaussiandistribution, while retaining a permutation equivariant sampling process.
To learn the parameters of the variational distribution, one can maximize the ELBO objective in (6).
However, the ELBO has no differentiable closed-form expression w.r.t.
ϕ .
7 To remedy this, we relaxthe binary variable s to a continuous one by applying the Gumbel-Softmax trick (Jang et al.
, 2016),resulting in an end-to-end training process with backpropagation.
4.3 Details of Training and Inference
Our model consists of two components: the EBM p θ ( S ) and the variational distribution q ϕ ( S ; ψ ) .
As shown in Fig.
2, these two components are trained in a cooperative learning scheme (Xie et al.,2018).
Specifically, we train the variational distribution q ϕ with fixed θ firstly by maximizing the
ELBO in (6).
To train the energy model p θ , we first initialize the variational parameter ψ (0) with theoutput of equivariant recognition network EquiNet( V ; ϕ ) .
This enables us to get a more accuratevariational approximate, since q ϕ has modeled the correlation among the elements in the set.
Noticethat ψ (0) does not depend on θ directly.
To learn θ , we take one further step of mean-field iteration
MFVI( ψ (0) , V, 1) , which flows the gradient through θ and enables to optimize θ using the crossentropy loss in (5) ( i.e., if we skip step 3 in Alg. 3, and feed ψ (0) to step 4, the gradient would notflow through θ ).
However, if we take multiple steps, it inclines to converge to the local optima that isthe same as the original mean-field iteration.
As a result, the benefit of correlation-aware inferenceprovided by the Gaussian copula would be diminished.
Detailed analysis is provided in Appendix F.6.
The training procedure is summarized in Alg.
3 (the complete version is given in Appendix D).
For inference in the test time, given a ground set V , we initialize the variational parameter via ψ (0) =
EquiNet( V ) , then run one step mean-field iteration ψ ∗ ← MFVI( ψ (0) , V, 1) .
Finally, the corresponding OS is obtained by applying the rounding method in (12).
We term our method as Equi variant
V ariational Inference for Set Function Learning ( EquiVSet ), and respectively use EquiVSet ind and
EquiVSet copula to represent two variants with independent and copula variational posterior.
Related Work
Set function learning.
There is a growing literature on learning set functions with deep neuralnetworks.
Zaheer et al.
(2017) designed the DeepSet architecture to create permutation invariantand equivariant function for set prediction.
Lee et al.
(2019) enhanced model ability of DeepSet byemploying transformer layer to introduce correlation among instances of set, and Horn et al.
(2020)extended this framework for time series.
It is noteworthy that they all learn set functions under thefunction value oralce and can be employed as the backbone of the utility function F θ ( S ; V ) in ourmodel.
Dolhansky & Bilmes (2016); Bilmes & Bai (2017); Ghadimi & Beigy (2020) have alsodesigned deep architectures for submodular set functions, however, these designs can not handle thevarying ground set requirement.
There are papers studying the learnability of specific set functions(e.g.
, submodular functions and subadditive functions) in a distributional learning setting (Balcanet al., 2012; Badanidiyuru et al., 2012; Balcan & Harvey, 2018) under the function value oracle,they mainly provide sample complexity with inapproximability results under the probably mostlyapproximately correct (PMAC) learning model.
Other methods relevant to our setting are TSPN(Kosiorek et al., 2020) and DESP (Zhang et al., 2020).
However they both focused on generating setobjects under a given condition.
While we aim at predicting under the optimal subset oracle.
Energy-based modeling.
Energy based learning (LeCun et al., 2006) is a classical framework tomodel the underlying distribution over data.
Since it makes no assumption of data, energy-basedmodels are extremely flexible and have been applied to wide ranges of domains, such as datageneration (Nijkamp et al., 2019), out-of-distribution detection (Liu et al., 2020), game-theoreticvaluation algorithms (Bian et al., 2022) and biological structure prediction (Shi et al., 2021).
Learning
EBMs can be done by applying some principled methods, like contrastive divergence (Hinton, 2002),score matching (Hyvärinen & Dayan, 2005), and ratio matching (Lyu, 2012).
For inference, it iswidely exploited gradient-based MCMC methods (Welling & Teh, 2011; Grathwohl et al., 2021).
Inthis paper, we train discrete EBMs under the supervision of OS oracle by running mean-field inference.
This new training paradigm may have some implicit connections with conventional methods, likenoise contrastive estimation.
We leave it for future works.
Amortized and Copula variational inference.
Instead of approximating separate variables foreach data point, amortized variational inference (VI) (Kingma & Welling, 2013) assumes that thevariational parameters can be predicted by a parameterized function of the data (Zhang et al., 2018).
The idea of amortized VI has been widely applied in deep probabilistic models (Hoffman et al., 2013;
Garnelo et al., 2018).
Although this procedure would introduce an amortization gap (Cremer et al.,2018), which refers to the suboptimality of variational parameters, amortized VI enables significantspeedups and combines probabilistic modeling with the representational power of deep learning.
Copula is the other method to improve the representational power for VI.
Tran et al.
(2015) usedcopula to augment the mean-field VI for better posterior approximation.
Suh & Choi (2016) adopted
Gaussian copula in VI to model the dependency structure of observed data.
However, none of themcan be directly applied to our setting involving discrete latent variables.
Empirical Studies
We evaluate the proposed methods on various tasks: product recommendation, set anomaly detection,and compound selection.
All experiments are repeated five times with different random seeds andtheir means and standard deviations are reported.
The model architectures and training details aredeferred to Appendix E. Additional experiments on synthetic datasets and varying ground set are givenin Appendices F.1 and F.2, respectively.
Ablation studies on hyper-parameter choices (e.g. MFVIiteration steps, number of MC samples, and rank of perturbation) are provided in Appendix F.6.
Evaluations.
We evaluate the methods using the mean Jaccard coefficient (MJC).
Specifically, foreach sample ( V, S ∗ ) , denoting the corresponding model predict as S ′ , the Jaccard coefficient isdefined as JC( S, S ′ )
= | S ′ ∩ S
| | S ′ ∪ S
| .
Then the MJC metric can be computed by averaging over all samplesin the test set: MJC = 1 |D t | (cid:80) ( V,S ∗ ) ∈D
t JC
( S ∗ , S ′ ) .
Baselines.
We compare our solution variants, i.e., DiffMF , EquiVSet ind , and EquiVSet copula to thefollowing three baselines:
- Random:
The expected performance of random guess, which estimates how difficult the task is.
[Table]
- PGM (Tschiatschek et al.
(2018), see Appendix A):
The probabilistic greedy model, which ispermutation invariant but computationally prohibitive.- DeepSet (NoSetFn) (Zaheer et al., 2017): The deepset architecture, satisfying permutation invariant,is the backbone of our models.
Its adapted version: 2 V → [0 , 1] | V | , which serves as the amortizednetworkds in EquiVSet , could work as a baseline since its output stands for the probability of whichinstance should be selected.
We train it with cross entropy loss and sample the subset via topNrounding in (12).
The term “NoSetFn” is used to emphasize that this baseline does not learn a setfunction explicitly, although it can be adapted to our empirical studies.
Product Recommendation.
In this experiment, we use the Amazon baby registry dataset (Gillenwater et al., 2014), which contains numerous subsets of products selected by different customers.
Amazon characterizes each product in a baby registry as belonging to a specific category, such as“toys” and “furniture”.
Each product is characterized by a short textual description and we representit as a 768 dimensional vector using the pre-trained BERT model (Devlin et al., 2018).
For each category, we generate samples ( V, S ∗ ) as follows.
Firstly, we filter out those subsets selectedby customers whose size is equal to 1 or larger than 30 .
Then we split the remaining subset collection
S into training, validation and test folds with a 1 : 1 : 1 ratio.
Finally for each OS oracle S ∗ ∈ S , werandomly sample additional 30 − | S ∗ | products from the same category to construct V \ S ∗ .
In thisway, we construct one data point ( V, S ∗ ) for each customer, which reflects this real world scenario:
V contains 30 products displayed to the customer, and the customer is interested in checking | S ∗ | ofthem.
Note that this curation process is different from that of (Tschiatschek et al., 2018, Section 5.3),which is deviated from the real world scenario (Detailed discussion in Appendix E.5.).
The performance of all the models on different categories are shown in Table 1.
Evidently, ourmodels perform favorably to the baselines.
Compared with PGM, which learns the set functionvia a probabilistic greedy algorithm, we can observe that our models, which model the the setfunctions with energy-based treatments, achieves better results on all settings.
Although DeepSetis also permutation invariant, our model still outperforms it by a substantial margin, indicating thesuperiority of learning the set function explicitly.
Set Anomaly Detection.
In this experiment, we evaluate our methods on two image datasets: thedouble MNIST (Sun, 2019) and the CelebA (Liu et al., 2015b).
For each dataset, we randomly splitthe training, validation, and test set to the size of 10 , 000 , 1 , 000 , and 1 , 000 , respectively.
Double MNIST:
The dataset consists of 1000 images for each digit ranging from 00 to 99 .
For eachsample ( V, S ∗ ) , we randomly sample n ∈ { 2 , . . .
, 5 } images with the same digit to construct the OSoracle S ∗ , and then select 20 − | S ∗ | images with different digits to construct the set V \ S ∗ .
CelebA:
The CelebA dataset contains 202 , 599 images with 40 attributes.
We select two attributes at randomand construct the set with the size of 8 .
For each ground set V , we randomly select n ∈ { 2 , 3 }images as the OS oracle S ∗ , in which neither of the two attributes is present.
See Fig.
3 and Fig. 4 in
Appendix E.6 for illustrations of sampled data.
From Table 2, we see that the variants of our model consistently outperform bot baseline methodsstrongly.
Furthermore, we observe that by introducing the correlation to the variational distribution,significant performance gains can be obtained, demonstrating the benefits of relaxing the independentassumption by using Gaussian copula.
Compound Selection in AI-aided Drug Discovery.
A critical step in drug discovery is to selectcompounds with high biological activity (Wallach et al., 2015;
Li et al., 2021; Ji et al., 2022), diversityand satisfactory ADME (absorption, distribution, metabolism, and excretion) properties (Gimenoet al., 2019).
As a result, virtual screening is typically a hierarchical filtering process with several
[Table] [Table]
necessary filters, e.g., first choosing the highly active compounds, then selecting diverse subsets fromthem, and finally excluding compounds that are bad for ADME.
We finally arrive at a compoundsubset after a series of these steps.
Given the OS supervision signals, we can learn to conductthis complicated selection process in an end to end manner.
As a result, it will eliminate the needfor intermediate supervision signals, which can be very expensive or impossible to obtain due topharmacy’s personal protection policy.
For example, measuring the bioactivity and ADME propertiesof a compound has to be done in wet labs, and pharmaceutical companies are reluctant to disclosethe data.
Here we simulate the OS oracle of compound selection by applying the two filters : highbioactivity and diversity filters, based on the following two datasets.
PDBBind (Liu et al., 2015a):
This dataset consists of experimentally measured binding affinitiesfor bio-molecular complexes.
We construct our dataset using the “refined” subsets therein, whichcontains 179 protein-ligand complexes.
BindingDB 8 : It is a public database of measured bindingaffinities, which consists of 52 , 273 drug-targets with small, drug-like molecules.
Instead of providingcomplexes, here only the target amino acid sequence and compound SMILES string are provided.
We apply the same filtering process to construct samples ( V, S ∗ ) for these two datasets.
Specifically,we first randomly select a number of compounds to construct the ground set V , whose size is 30 and300 for PDBBind and BindingDB, respectively.
Then 13 compounds with the highest bioactivity arefiltered out, accompanied by a distance matrix measured by the corresponding fingerprint similarityof molecules.
To ensure diversity, the OS oracle S ∗ is generated by the centers of clusters which arepresented by applying the affinity propagation algorithm.
We finally obtain the training, validation,and test set with the size of 1,000, 100, and 100, respectively, for both two datasets.
Detaileddescription is provided in Appendix E.7.
From Table 3, one can see that our methods magnificently outperform the random guess.
Thisindicates that the proposed EquiVSet framework has great potential for drug discovery to facilitatethe virtual screening task by modeling the complicated hierarchical selection process.
Besides,improvements of EquiVSet can be further observed by comparing with DeepSet, which simplyequips the deepset architecture with cross entropy loss, illustrating the superiority of explicit setfunction learning and energy-based modeling.
Although comparable results could be achieved by
PGM with sequential modeling, which satisfies permutation invariance and differentiability, ourmodels still outperform it.
This is partially because our models additionally maintain the other threedesiderata of learning set functions, i.e., varying ground set, minimum prior, and scalability.
We alsoconduct a fairly simple task in Appendix F.5, in which only the bioactivity filter is considered.
Tosimulate the full selection process, we leave it as important future work due to limited labels.
Discussion and Conclusions
We proposed a simple yet effective framework for set function learning under the OS oracle.
Byformulating the set probability with energy-based treatments, the resulting model enjoys the virtuesof permutation invariance , varying ground set , and minimum prior .
A scalable training and inferencealgorithm is further proposed by applying maximum log likelihood principle with the surrogate ofmean-field inference.
Real-world applications confirm the effectiveness of our approaches.
Limitations & Future Works.
The training objective in (5) does not bound the log-likelihood of
EBMs.
This limitation can be mitigated by applying advanced discrete EBMs trainer, e.g., GFlowNet(Zhang et al., 2022).
We leave it as a future work.
In addition, the proposed framework has thepotential to facilitate learning to select subsets for other applications (Iyer et al., 2021), includingactive learning (Kothawade et al., 2021), targeted selection of subsets, selection of subsets forrobustness (Killamsetty et al., 2020), and selection of subsets for fairness.
Though we considerlearning generic neural set functions in this work, it is beneficial to consider building useful priorsinto the neural set function architectures, such as set functions with the diminishing returns prior(Bilmes & Bai, 2017) and the bounded curvature/submodularity ratio prior (Bian et al., 2017).
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes]
• Did you include the license to the code and datasets?
[No] The code and the data are
proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes] See Sec. 7.(c)
Did you discuss any potential negative societal impacts of your work?
[No] That is notexplicitly pertinent to any negative social impact.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes](b) Did you include complete proofs
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix E.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Appendix E.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix E.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] See Sec.
(b) Did you mention the license of the assets?
[Yes] The assets we use are public.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
Our code is available at a anonymous link.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A] All datasets we evaluate on are open source datasets.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
Appendix
for “Learning Neural Set Functions Under
the Optimal Subset Oracle”
Contents
A Details of the Probabilistic Greedy Model
B Derivations
B.1 Derivations of the Maximum Entropy Distribution . . . . . . . . . . . . . . . . . . . . . . .
B.2 Derivations of the Fixed Point Iteration . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.3 Derivations of the Gradient of Multilinear Extension . . . . . . . . . . . . . . . . . . . . . .
C Low-Rank Perturbation for the Covariance Matrix
D Detailed Pseudo Code of EquiVSet Algorithms
E Experimental Details[-
-]
E.1 The Architecture of EquiVSet . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.2 Implementation Details . .
E.3 Baselines . .
E.4 Assumptions on the Underlying Data Generative Distribution . . . . . . . . . . . . . . . . . E.5 Detailed Experiment Settings for Product Recommendation . . . . . . . . . . . . . . . . . . E.6 Detailed Experiment Settings for Set Anomaly Detection . . . . . . . . . . . . . . . . . . . .
E.7 Detailed Experiment Settings for Compound Selection . . . . . . . . . . . . . . . . . . . . .
F Additional Experiments[-
-]
F.
Synthetic Experiments . .
A Details of the Probabilistic Greedy Model
The probabilistic greedy model (PGM) solves optimization (1) with a differentiable extension ofgreedy maximization
algorithm (Tschiatschek et al., 2018).
Specifically, denote the first j chosenelements as S j = { s 1 , . . .
, s j } ⊆ V , PGM samples the ( j + 1) th element from the candidate set
V \ S j with the probability proportional to exp( F θ ( s j +1 + S j ) /γ ) , which raises the probability ofthe selected elements in the sequence π = { s 1 , s 2 , . . .
, s k } asp
[Equation]
|
[Equation]
[Equation] (cid:89) [Equation] (14) where γ is a temperature parameter, S 0 = ∅ , and s + S :
= S ∪ { s } .
Note that, the computation ofp θ ( π | V ) depends on the order of sequence π , which would make the learned parameter θ sensitive tothe sampling order.
To alleviate this problem, Tschiatschek et al.
(2018) finally construct the set mass function by enumerating all possible permutations
[Equation]
(15) where Π S is the permutation space generated from S .
After training, the OS oracle S can be sampledvia sequential decoupling p ( s i +1 | S i )
∝ exp( F θ ( s i +
+ S i ) /γ ) .
However, maximizing the loglikelihood of (15) is prohibitively expensive and unscalable due to the exponential time complexityof enumerating all permutations.
Although one can apply Monte Carlo approximation to avoid that,i.e., approximating log p θ ( S | V )
= log (cid:80) π ∈ Π S p θ ( π | V ) with log p θ ( π | V ) , π ∼ Π S , such a simpleestimator is biased, resulting in a permutation variant model.
B DerivationsB.1 Derivations of the Maximum Entropy Distribution
The first step to solve problem (2) is to construct a proper set mass function p θ ( S | V ) monotonicallygrowing with the utility function F θ ( S ; V ) .
There exits countless ways to construct such a probabilitymass function, such as p θ ( S | V )
∝ F θ ( S ; V ) and the set mass function defined in PGM, i.e., Eq.
Here, one would care about what the most appropriate set mass function should be?
Generally weprefer the model to assume nothing about what is unknown.
More formally, we should choose themost “uniform" distribution, which maximizes the Shannon entropy H ( p ) =
− (cid:80) S ⊆
V p ( S ) log p ( S ) .
This principle is known as “noninformative prior" (Jeffreys, 1946), which has been widely appliedin many physical systems (Jaynes, 1957a,b).
It turns out that the energy-based model is the onlydistribution with maximum entropy.
More specifically, the following theorem holds:
Theorem 1.
Let P µ := { p ( S ) : E p [ F ( S )]
= µ } be a set of distributions satisfying the expectationconstraint E p [ F ( S )]
= µ , and p λ have density
[Equation]
If E p λ [ F ( S )]
= µ , then p λ maximizes the entropy H ( p ) over P µ ; moreover, the distribution p λ isunique.
Proof.
The derivation below is adapted from Jaynes (1957a) in the context of set function learning,for completeness.
We rewrite the maximum entropy problem in the form of
[Equation]
−
[Equation]
Introducing Lagrange multipliers α ( S ) > 0 for the constraint p ( S ) > 0 , β ∈ R for the normalizationconstraint that (cid:80) S ⊆
V p ( S )
= 1 , λ for the constraint that E p [ F ( S )]
= µ , and , we obtain thefollowing Lagrangian:
[Equation]
[Equation]
Now we take derivatives and obtain
[Equation]
Since this function is convex in p , the minimizing p can be find by setting this equal to zero
[Equation]
[Equation]
Note that in this setting we always have p ( S ) > 0 .
By complementary slackness, the constraintp ( S ) > 0 is unnecessary and we have α ( S ) =
To satisfy the constraint (cid:80) S ⊆
V p ( S )
= 1 , we takeβ = 1 − + log Z = − 1 + log (cid:80) S ⊆
V exp( λF ( S )) .
Then the optimal mass p has the form
[Equation]
So we reach the form of p ( S ) we would like to have.
Next we show the distribution p λ is unique.
Assume there exists any other distribution p ∈ P µ , suchthat p = argmax p H ( p ) .
In this case, we have
[Equation]
As KL ( p || p λ ) ≥
unless p = p λ , we have shown that p λ is the unique distribution maximizing theentropy, as desired.
Discussion.
Theorem 1 shows that EBM is the maximum entropy distribution, which verifies theassertion that energy-based treatments of set function enjoy the minimum prior property.
It should benoted that the model proposed by Tschiatschek et al.
(2018) violates this requirement.
They usedsequencial modeling to construct p ( S ) (see (14) and (15)).
Although this approach simplifies thesampling process, it introduces undesirable inductive bias.
B.2 Derivations of the Fixed Point Iteration
In this section, we give the detailed derivation for the fixed point iteration (FPI) of MFVI:
[Equation]
First, recall that we want to maximize the ELBO:
[Equation]
The formula (20) is obtained by setting the partial derivative w.r.t.
coordinate i of ELBO to be 0 :
[Equation]
which implies
[Equation]
This is exactly the formula of FPI used in the mean-field variational inference algorithm.
Note thatthe FPI actually corresponds to the gradient ascent with an adaptive step size vector α as
[Equation]
where σ ( x ) = (1 + exp( − x )) − 1 denotes the sigmoid function.
To verify this, we haveψ ( k +1) i = ψ ( k ) i + α ( k ) i (cid:16) ∇ ψ ( k ) i
f F θ mt ( ψ ( k ) )
+ ∇ ψ ( k ) i H ( q ( S ; ψ ( k ) ))
(cid:17) = ψ ( k ) i + σ ( ∇ ψ i f F θ mt ( ψ ( k ) )) − ψ ( k ) i ∇ ψ ( k ) i
f F θ mt ( ψ ( k ) )
+ log[(1 − ψ ( k ) i ) /ψ
( k ) i ] (cid:32) ∇ ψ ( k ) i
f F θ mt ( ψ ( k ) )
+ log (1 − ψ ( k ) i ) ψ ( k ) i (cid:33) =
(1 + exp −∇ ψ ( k ) i f F θ mt ( ψ ( k ) ))) − 1 . The connection to gradient ascent further confirms the soundness of our FPI algorithm.
B.3 Derivations of the Gradient of Multilinear Extension
In this section, we prove that the gradient of multilinear extension can be estimated using Monte
Carlo sampling.
Specifically we have
[Equation]
Discussion.
The Monte Carlo (MC) approximation of ∇ ψ i f F θ mt is unbiased.
Thereby, although exactlycalculating (22) has exponential time complexity, we can apply MC sampling to approximate it in apolynomial time, resulting a scalable training algorithm.
It is worth to note that the MC approximationused in PGM (see (15)) is biased.
That is they approximate log p θ ( S | V )
= log (cid:80) π ∈ Π S p θ ( π | V ) withlog p θ ( π | V ) , π ∼ Π S .
Although such a biased approximation can be computed in polynomial time,they undesirably introduce permutation variance.
C Low-Rank Perturbation for the Covariance Matrix
In the construction of Gaussian copula C Σ , we require a positive semi-definite matrix Σ ∈ R | V |×| V | ,whose elements are generally modeled as the output of neural networks.
Thereby, if the size of groundsize V is large, the number of neural network outputs will be prohibitively large.
Meanwhile, based onthe definition of set, covariance matrix Σ is further required to satisfy permutation equivariance .
Toremedy this issue, we propose to employ a more efficient strategy, namely Lower-Rank Perturbation ,which restricts the covariance matrix to the form
[Equation]
(23) where Σ ∈ R | V |×| V | + is a diagonal matrix with positive entries and P
= [ p 1 , p 2 , . .
. , p v ] is a lowerrank perturbation matrix with p i ∈ R | V | and v ≪ | V | .
In this way, the number of neural networkoutputs can be dramatically reduced from | V | 2 to v | V | .
Another benefit of constructing Σ in this wayis that, it is convenient to employ the DeepSet architecture in (13) to output D and p i for i = 1 , . .
, v ,such that they are permutation equivariant , and the resulting covariance matrix Σ
= D + P P Tis also permutation equivariant .
Moreover, the lower-rank perturbation trick permits us to avoidusing Cholesky decomposition to sample a Gaussian noise with covariance Σ , which is prohibitivelyexpensive.
Specifically, the Gaussian noise g ∼ N ( 0 , Σ ) can be reparameterized as
[Equation]
(24) where
ϵ 1 ∼ N ( 0 , I | V | ) and ϵ 2 ∼ N ( 0 , I v ) .
In this way, the sampling complexity can be reducedfrom O ( | V | 3 ) to O ( v 2 | V | ) .
D Detailed Pseudo Code of EquiVSet Algorithms
We provide the pseudo-code for EquiVSet in Alg. 4.
The training procedure consists of two steps: i)train q ϕ with fixed θ ; ii) train p θ under the guidance of q ϕ .
Specifically, to train q ϕ , we first fix theparameter θ of the set function and then optimize ϕ by maximizing the ELBO in (6).
To train p θ , wefirst initialize the variational parameter ψ via EquiNet and then run K steps mean-field iteration tomake ψ dependent with θ .
Finally, the parameter θ can be optimized by minimizing the cross entropyin (5).
Note that, we set K as 1 in our experiments.
Input : { V i , S ∗ i } Ni =1 : η
Output : Optimal parameters (
1: θ, ϕ ← Initialize parameter 2: repeat 3:
Sample training data point ( V, S ∗ ) ∼ { V i , S ∗ i } Ni = 4:
Obtain variational parameter ψ via EquiNet ψ ← EquiNet( V ; ϕ ) 5:
Sample m subsets via CopulaBernoulli( V, v, τ ) or IndBernoulli( V, τ ) S n ∼ q ( S ; ψ ) , n = 1 , 2 , . . .
, m 6:
Update the parameter ϕ by maximizing ELBO in (6) ϕ ← ϕ + η ∇ ϕ (cid:32)m m (cid:80) n =1 F θ ( S n ) + | V | (cid:80) i =1 [ ψ i log ψ i + (1 − ψ i ) log(1 − ψ i )] (cid:33) 7:
Initialize parameter ψ via EquiNet ψ (0) ← EquiNet(V; ϕ ) 8: for k ← 1 , . .
, K do 9: for i ← 1 , . . .
, | V | in parallel do 10:
Sample m subsets via CopulaBernoulli( V, v, τ ) or IndBernoulli( V, τ )
S n ∼ q
S ; ψ (0)
| ψ (0)
i ← 0) , n = 1 , 2 , . . .
, m 11:
Update the variational parameter ψ ψ ∗ i ← σ (cid:18) 1 m m (cid:80) n =
[ F θ ( S n + i ) − F θ ( S n )] (cid:19) 12: end for 13: end for 14:
Update the parameter θ by minimizing the cross entropy loss in (5) θ ← θ − η
[Table] [Table]
[Table]
E Experimental DetailsE.1 The Architecture of EquiVSet
In this section, we provide a detail architecture description of EquiVSet copula .
EquiVSet copulaconsists of two different components that are implemented as neural networks: (i) the set functionwhich is permutation invaraint and (ii) the recognition network which is permutation equivariant.
Weemploy the DeepSet architecture to implement these two components, with the detailed architecturesare given in Table 4.
[Table]
In Table 4, InitLayer( S, d ) denotes the set transformation function, which encodes the set objectsinto vector representations.
FC( d, h, f ) denotes the fully-connected layer with activation functionf .
diag( v ) is a diagonal matrix with the elements of diagonal being vector v .
[ p ] v = [ p 1 , . .
. , p v ]denotes a matrix with p representing a column perturbation vector.
Note that we also propose twovariant methods, i.e., DiffMF and EquiVSet ind .
For DiffMF , we apply the same architecture of thethe set function in Table 4.
We also exploit the same architecture for EquiVSet ind , but discardingthe copula components, i.e., D and P .
In all experiments, we implement our models following thesame architecture with the difference being that we apply various InitLayer to different datasets. The
[Table] architectures of InitLayer for different datasets are depicted below.
Synthetic datasets.
The synthetic datasets consist of the Tow-Moons and Gaussian-Mixture datasets.
Each instance of the set is a two-dimensional vector, which represents the corresponding Cartesiancoordinates.
In this dataset, the InitLayer is a one-layer feed-forward neural network FC(2 , 256 , − ) .
Amazon Baby Registry.
The Amazon baby registry dataset consists of a set of products that arecharacterized by a short textual description.
We transform them into vector representations using thepre-trained BERT module (Devlin et al., 2018).
Thereby, each instance of the set is a 768 dimensionalfeature vector.
The InitLayer is modelled as FC(768 , 256 , − ) .
Double MNIST.
The double MNIST dataset consists of different digit images ranging from 00 to99 .
Each image has the shape of (64 , 64) and we reshape it into (4096 , ) .
Therefore, the InitLayer isdesigned as FC(4096 , 256 , − ) .
CelebA.
The CelebA dataset contains 202 , 599 number of face images.
Each image is in the shape of(3 , 64 , 64) .
We employ convolutional neural networks as InitLayer .
Specifically, the architecture of
InitLayer is
ModuleList([Conv(32 , 3 , 2 , ReLU) , Conv(64 , 4 , 2 , ReLU) ,
Conv(128 , 5 , 2 , ReLU) , MaxPooling , FC(128 , 256 , − )]) , where Conv( d, k, s, f ) is a convolutional layer with d output channels, k kernel size, s stride size,and activation function f .
PDBBind.
The PDBBind database consists of experimentally measured binding affinities forbiomolecular complexes (Liu et al., 2015a).
It provides detailed 3D Cartesian coordinates of both ligands and their target proteins derived from experimental ( e.g. , X-ray crystallography) measurements.
The atomic convolutional network (ACNN) (Gomes et al., 2017) provides meaningful vector featuresfor complexes by constructing nearest neighbor graphs based on the 3D coordinates of atoms andpredicting binding free energies.
In this work, we apply the output of last second layer of the ACNNmodel followed by feed-forward neural networks to obtain the representations of complexes.
Moreformally, the InitLayer is defined as
ModuleList([ACNN
[: − 1] , FC(1922 , 2048 , ReLU) , FC(2048 , 256 , − )]) , where ACNN[: − 1] denotes the ACNN module without the last prediction layer, whose outputdimensionality is 1922 .
BindingDB.
The BindingBD dataset containts 52 , 273 drug-target pairs.
We exploit the DeepDTAmodel (Öztürk et al., 2018) to encode drug-target pairs as vector representations.
Specifically, the
DeepDTA model first represents the drug compound and target protein as sequences of one-hotvectors and encodes them as feature vectors using convolutional neural networks.
The detailedarchitecture of InitLayer used in this dataset is demonstrated in Table 5.
[Table]
E.2 Implementation Details
Here we provide a detailed description of the hyperparameters setup for our model EquiVSet and itsvariants.
EquiVSet contains four important hyperprameters: the number of Monte Carlo samplingm and mean-field iteration steps K in Alg. 1, and the rank of lower-rank perturbation v in (23).
We set m = 5 , v = 5 throughout the experiments.
For the mean-field iteration steps K , we set itas 5 for the variant model DiffMF , and 1 for EquiVSet ind and EquiVSet copula .
It is noted that thehyperparameters above are empirically set, and we have detail sensitivity analysis in Appendix F.6.
The proposed models are trained using the Adam optimizer (Kingma & Ba, 2014) with a fixed learningrate 1 e − 4 and weight decay rate 1 e − 5 .
We choose the batch size from { 4 , 8 , 16 , 32 , 64 , 128 } , sincethe model sizes for various datasets are different and we choose the largest batch size to enable it canbe trained on a single Tesla V100-SXM2-32GB GPU.
We apply the early stopping strategy to train the models, including the baselines and our models.
That is if the performances are not improved in continuous 6 epochs, we early stop the trainingprocess.
Each dataset is trained for maximum 100 epochs.
After each epoch, we validate the modeland save the model with the best performance on the validation set.
After training, we evaluate theperformance of saved models on the test set.
We repeat all experiments 5 times with different randomseeds and the average performance metrics and their standard deviations are reported as the finalperformances.
E.3 Baselines
Throughout the experiments, we compared our models with three conventional approaches: randomguess, probabilistic greedy model (PGM) (Tschiatschek et al., 2018) and DeepSet (Zaheer et al.,2017).
Further descriptions of the benchmarks and implementation details are as follows
- Random: We report the expected value of the Jaccard coefficient (JC) of random guess.
This baselineprovides an estimate of how difficult the task is.
Specifically, given a data point ( V, S ∗ ) , it can becomputed as E ( JC ( V, S
[Equation] [Equation]
- PGM Tschiatschek et al.
[Equation] 2018): PGM is the most relevant method to the set functions learningunder the OS oracle, that solves optimization (1) using greedy maximization algorithm with the virtues of differentiability and permutation invariance.
We employ the same architecture defined in
Table to model the set function F θ ( S ) in (14).
The temperature parameter γ is empirically set as
We
use Monte Carlo sampling to estimate (15).
That is we randomly sample one permutationπ ∼ Π S and use p θ ( π | V ) to approximate p θ ( S | V ) .
The model is trained using the Adam optimizer,with batch size choosing from { 4 , 8 , 16 , 32 , 64 , 128 } , fixed learning rate 1 e − 4 , and fixed weightdecay rate 1 e − 5 .
[Table]
- DeepSet (NoSetFn)
(Zaheer et al., 2017): DeepSet is a neural-network-based architecture thatsatisfies permutation invariance and varying ground sets.
Although the DeepSet architecture can beemployed here to sample the optimal subset oracle, it does not learn the set functions explicitly.
Weexploit the same architecture of F θ ( S ) in Table 4, but drop the SumPooling operator to ensure thedimensionality of output is | V | .
This baseline is trained by minimizing the objective in (5) using the
Adam optimizer with batch size choosing from { 4 , 8 , 16 , 32 , 64 , 128 } , fixed learning rate 1 e − 4 , andfixed weight decay rate 1 e − 5 .
E.4 Assumptions on the Underlying Data Generative Distribution
In this section, we discuss the assumptions made about the data distribution for better understandingthe set functions learning under optimal subset oracle.
Generally speaking, for any scenario withthe output being a subset S of the given ground set V of the input, the proposed approach couldbe applied to predict the subset S of the given ground set V .
The only loose assumption is thatthe optimal subset oracle S ∗ of a given ground set V is generated by some underlying distributionformulated via a utility function that maximizes the utility value of OS oracle (see (1) in the mainpaper).
We further assume the utility function could be parameterized by a deep neural network,thanks to the universal approximation theorem (Leshno et al., 1993).
This assumption is very weak and generally makes sense in practice. We also apply this assumptionto the datasets used in the experiments. Specifically, in the product recommendation (Appendix E.5), V is the set of recommended products, and S ∗ is the one the customer buys.
Undoubtedly, theunderlying generative distribution, or say the utility function is specified by the selection process ofcustomers.
In the set anomaly detection (Appendix E.6), given a ground set V , S ∗ is generated as theone containing anomaly data points.
Therefore, the utility function in this setting is formulated as theanomaly pattern.
Moreover, in the compound selection Appendix E.7, we applied high bioactivity anddiversity filters to select compounds.
In this case, the utility function is determined by the bioactivityand diversity of the group of compounds.
E.5 Detailed Experiment Settings for Product Recommendation
Detailed Descriptions of the Amazon Baby Registry Dataset.
The Amazon baby registry data(Gillenwater et al., 2014) consists of baby registry data collected from Amazon and is split intoseveral datasets according to product categories, such as toys, furniture, etc.
For each category, whichcan be considered as the product database, Amazon provides multiple sets of products selected bydifferent customers.
Thereby, these subsets of products can be viewed as OS oracles.
To ensure thateach ground set V only contains one OS oracle S ∗ , we construct the sample ( V, S ∗ ) as follows.
Foreach subset of products selected by an anonymous user, we filter it out if its size is equal to 1 or largerthan 30 .
For each OS oracle S ∗ in the remaining subsets, we randomly sample 30 − | S ∗ | products inthe same category to construct V \ S ∗ .
We summarize the statistics of the categories in Table 6.
Comparing with the Setting of (Tschiatschek et al., 2018, Section 5.3).
In (Tschiatschek et al.
,2018, Section 5.3), Tschiatschek et al.
(2018) consider an alternative setting which is different from [Figure]
ours. Specifically, they construct the ground set V as all the products in a category, and view theselected subsets of all the customers as the corresponding optimal subsets.
That is why they have thedata points in the form of { V, ( S ∗ 1 , . . .
, S ∗ N ) } .
This is a bit problematic since it is deviated from thereal world scenario: naturally the chosen subset S ∗ i shall depend on both V and the i -th customer’spersonal preference.
However, the customer is fully anonymized, so no information can be extractedfrom this dataset.
In order to be aligned with the real world scenario, we curate the dataset in the following way, inorder to make data samples in the OS supervision oracle with the data in the form of { V i , S ∗ i } .
For each category, we generate samples ( V, S ∗ ) as follows.
Firstly, we filter out those subsets selectedby customers whose size is equal to 1 or larger than 30 .
Then we split the remaining subset collection
S into training, validation and test folds with a 1 : 1 : 1 ratio.
Finally for each OS oracle S ∗ ∈ S , werandomly sample additional 30 − | S ∗ | products from the same category to construct V \ S ∗ .
In this way, we construct one data point ( V, S ∗ ) for each customer, which reflects this real worldscenario:
V contains 30 products displayed to the customer, and the customer is interested in checking| S ∗ | of them.
This is also consistent with real world recommender system, as users can only browsea small number of products at a time since the screen size of the device is limited, and the user haslimited attention.
E.6 Detailed Experiment Settings for Set Anomaly Detection
In this experiment, we evaluate our methods on two real-world datasets:
Double MNIST:
The dataset consists of 1000 images for each digit ranging from 00 to 99 .
For eachsample ( V, S ∗ ) , we randomly sample n ∈ { 2 , . . .
, 5 } images with the same digit to construct the
OS oracle S ∗ , and then select 20 − | S ∗ | images with different digits to construct the set V \ S ∗ .
Anexample is shown in Fig.
CelebA:
The CelebA dataset contains 202 , 599 images with 40 attributes.
As shown in Fig.
4, we select two attributes at random and construct the set with the size of 8 .
Foreach ground set V , we randomly select n ∈ { 2 , 3 } images as the OS oracle S ∗ , in which neither ofthe two attributes is present. [Table]
E.7 Detailed Experiment Settings for Compound Selection
Alg.
7 shows the corresponding data generation process of simulating the OS oracle for compoundselection.
In this algorithm, random _ choose( C , n ) means randomly choosing n compounds fromthe database C
( i.e., PDBBind or BindingDB), and topK _ bioactivity( V, m ) represents selectingthe top- m compounds with highest biological activity from the ground set V .
These two operatorscombine together to form the bioactivity filter, in which we set ( n, m ) as (30 , 10) , and (300 , 100)for PDBBind and BindingDB, respectively.
To further apply the diversity filter, we use the RDKit 9tools to compute the similarity between each molecule pair based on their topological fingerprints.
This operator corresponds to the line 3 of Alg.
7, in which cal _ fingerprint _ similarity( S ) returns thesimilarity matrix M ∈ R | S |×| S | of the set of compounds S .
Since rows (or columns) of the similaritymatrix can be regarded as the features of the corresponding molecules, the molecules are clusteredbased on these similarity features by applying the affinity propagation algorithm.
The OS oracle
S ∗ is finally represented by the center of each cluster.
Note that, each compound consists of twosmall molecules, i.e., the protein-ligand molecules in PDBBind, and the drug-target molecules in
BindingDB.
We use the protein and drug molecules to compute the fingerprint similarity for PDBBindand BindingDB, respectively.
F Additional ExperimentsF.1 Synthetic Experiments
We demonstrate the effectiveness of our models on learning set functions with two synthetic datasets:the two-moons dataset with additional noise of variance σ 2 = 0 .
1 , and mixture of Gaussians12 N ( µ 0 , Σ )
+ 12 N ( µ 1 , Σ ) , with µ 0 =
[ 1 √ 2 , 1 √ 2 ] T , µ 1 = − µ 0 , Σ
= 14 I 2 .
Take the Gaussianmixture as an example, the data generation procedure ( S ∗ , V ) ∼ P ( S ∗ , V ) is as follow: i) samplea component: b ∼ Bernoulli ( 1 2 ) ; ii) sample 10 points from N ( µ b , Σ ) to construct S ∗ ; iii) sample90 points for V \ S ∗ from N ( µ 1 − b , Σ ) .
We collect 1 , 000 samples for training, validation, and test,respectively.
A qualitative result of the EquiVSet copula is shown in Fig.
5, where the green dots represent correctmodel predictions, the red crosses are incorrect model predictions, and the yellow triangles representthe data points in the subset oracle S ∗ that are missed by the model.
One can see that the mostconfusing points are located at the intersection of two components.
We also illustrate the quantitativeresults in Fig.
As expected, our methods provide significantly better performance over othermethods, with averaged 59 .
16% and 100 .
07% improvements compared to PGM.
on the Two-Moonsand Gaussian-Mixture datasets, respectively.
F.2 Experiments on Varying Ground Set
Thanks to the virtues of DeepSet, our models are able to process input sets of variable sizes, which istermed as varying ground set property.
To examine the impact of ground set sizes, we care about [Figure] [Figure] the following two questions: i) how well the model performs on different sizes of ground set duringthe test time; and ii) how well does the model train on ground sets of different sizes?
To answerthese two questions, we conduct experiments on the synthetic datasets using the proposed model
EquiVSet copula .
Set Size Transferability Analysis We first experiment to understand the pattern of set size transferability.
In this experiment, we train the model using fixed sizes of the ground set but test the trainedmodel on different sizes.
We present two scenarios: train on a small size but test on a large one, andtrain on a large size but test on a small one.
For the former one, we fix the size of OS oracle S ∗ tobe 10 , and train the model with ground set V of size 100 .
After training, we test it using varyingsizes of ground set in the range of { 200 , 400 , 600 , 800 , 1000 } .
For the latter one, we fix the size of
OS oracle S ∗ to be 10 , and train the model with ground set V of size 1000 .
After training, we testit using varying sizes of ground set in the range of { 100 , 200 , 400 , 600 , 800 } .
The former and latterexperiments are conducted on the Two-Moons and Gaussian-Mixture datasets, respectively, with theresults shown in Fig.
As can be seen, the performance would be slightly reduced if tested on adifferent size.
Moreover, increasing the difference would enlarge the reduction.
Selection Ratio Analysis To answer the second question, we fix the size of OS oracle S ∗ to be 10 ,and experiment with different selection ratios | S ∗ | | V | in the range of { 0 . 5 , 0 . 2 , 0 . 1 , 0 . 05 , 0 . 01 } .
Unlikethe set size transferability analysis, in this experiment, the selection ratios are the same during trainingand testing.
Fig.
8 shows the performance of different ratios on two synthetic datasets.
We observethat increasing the ratio would deteriorate the model performance.
This phenomenon makes intuitivesense, since sampling subset from a large collection is more difficult.
Moreover, the model performsworst when the ratio is equal to 0 .
This is partly because S ∗ and V \ S ∗ are randomly sampled fromone of two components.
When | S ∗ | = | V \ S ∗ | , the model struggles to identify the optimal subset.
F.3 Comparison with Set Transformer
Set Transformer (Lee et al., 2019), which satisfies permutation invariant, is a well-known architectureused to model interactions among elements in the input set. Similar to DeepSet, Set Transformercould be adapted to serve as a baseline. Specifically, the architecture of the SetTransformer (NoSetFn)
[Figure]
[Table]
baseline is
[Table]
ModuleList(InitLayer( V, 256) , SAB(256 , 500 , 2 , − ) , SAB(500 , 1 , 2 , sigmoid)) , where SAB( d, h, m, f ) 10 denotes the set attention block (Lee et al., 2019) with d dimensional setinput, h dimensional set output, m multi-head attentions, and activation function f .
We train theadapted Set Transformer model: 2 V → [0 , 1] | V | with cross entropy loss and sample the subset viatopN rounding in (12).
It is noteworthy that, like DeepSet (NoSetFn), SetTransformer (NoSetFn)deos not learn a set function explicitly, although it can be adapted as a baseline and can be viewed asmerely modelling the amortized network in our EquiVSet framework.
For fair comparison, we also replace the DeepSet backbone with Set Transformer in EquiVSet .
Specifically, the InitLayer in Table 4 is replaced with
ModuleList(InitLayer( S, 256) , SAB(256 , 500 , 2 , − ) , SAB(500 , 256 , 2 , − )) .
Experiments are conducted on product recommendations, with the results shown in Table 7.
It showsthat the proposed approaches with the Set Transformer backbone outperform the Set Transformer(NoSetFn) comprehensively.
One could also compare the results of Table 1 in the paper.
It canbe seen that the proposed EquiVSet (with DeepSet backbone) also performs better than the Set
Transformer baseline.
Moreover, EquiVSet (with DeepSet backbone) outperforms EquiVSet (with
Set Transformer backbone) consistently, indicating that EquiVSet has great potential to be improvedwith more advanced architecture.
F.4 Experiments on Set Anomaly Detection with
-MNIST and CIFAR-10
In this experiment, we further perform set anomaly detection on the other two datasets: F-MNIST(Xiao et al., 2017) and CIFAR-10 (Krizhevsky et al., 2009).
Both two datasets contain images with10 different labels.
For each dataset, we randomly sample n ∈ { 2 , 3 } images as the OS oracle S ∗ ,and then select 8 − | S ∗ | images with different labels to construct the set V \ S ∗ .
We finally obtainthe training, validation, and test set with the size of 10 , 000 , 1 , 000 , 1 , 000 , respectively, for both twodatasets.
Illustrations of sampled data are shown in Fig.
[Figure]
$LUSODQH &DW 'RJ
[Figure]
[Table]
[Table]
The results are shown in Fig.
We see that the variants of our model consistently outperform baselinemethods strongly.
Moreover, DiffMF seems to perform better than EquiVSet ind and EquiVSet copulain set anomaly detection (similar results can be found in Table 2).
However, this is not a consistentphenomenon.
It seems that in most scenarios, e.g., product recommendation, compound selection,and synthetic dataset, EquiVSet performs better than DiffMF .
F.5 Experiments on Compound Selection with Only the Bioactivity Filter
To further evaluate the potential of EquiVSet for drug discovery, we consider an alternative settinghere.
In contrast to the task in Sec. 6, which aims at selecting the most active compounds whilepreserving diversity, the task defined here only focuses on selecting the compounds with the highestbioactivity, which results a relatively simple selection process.
The following is a detailed description.
PDBBind: To construct a data point ( V, S ∗ ) , we randomly sample 30 complexes as the ground set Vfrom the PDBBind database, and S ∗ is generated by the five most active complexes in V .
Finally, weobtain the training, validation, and test set with the size of 1 , 000 , 100 , 100 , respectively.
BindingDB:
We construct the ground set V by randomly sampling 300 drug-targets from the BindingDB databaseand generate S ∗ with the 15 most active drug-target pairs.
We finally obtain the training, validation,and test set with the size of 10 , 00 , 1 , 00 , and 1 , 00 , respectively.
Table 9 shows that our methods outperform the baselines.
Meanwhile, the baselines also showsatisfactory results.
That is because identifying the most active compounds is a relatively simpletask, especially for the PDBBind dataset with complex structures.
More specifically, the modelcould predict the activity value of complexes precisely without considering the interactions betweenelements in the set, since the structure of complexes has provided sufficient information for this task.
It is worth noting that the models in this task perform better than that in Sec. 6 partly because aone-layer filter (i.e. bioactivity) represents an easier way to replicate the OS oracle than a two-layerfilter (i.e. bioactivity and diversity).
Nevertheless, both experimental results in Sec.
6 and heredemonstrate the effectiveness of EquiVSet for facilitating the complicated compound selectionprocess.
F Sensitivity Analysis of Hyperparameters
The proposed model EquiVSet copula has four important hyperparameters: the number of Monte Carlosampling m and mean-field iteration step K in Alg. 1, and the rank of lower-rank perturbation v in(23).
In this section, we discuss the impact of these hyperparameters through a sensitivity analysis onthe Amazon product datasets.
[Figure] Impact of the Mean Field Iteration Step Since iteration step K controls the convergence ofmean-field iterative algorithms, this hyperparameter is highly relevant to the final performanceof EquiVSet copula .
We experiment with different K on the Amazon product dataset.
The resultsare shown in the first row of Fig.
We notice that increasing K would degenerate the model’sperformance.
This seems to be embarrassingly surprising at first glance, since a large stride Kencourages convergence with guarantee, resulting in a more robust training process.
It is worth to benoted that in this method, we apply an amortized variational distribution to initialize the parametersfor mean-field iterative algorithms.
Since the amortized variational distribution is modeled with
Gaussian copula, it can effectively capture the correlation among elements in the set, such thatobtaining a better local optimal.
However, if the iterative step K is large, the model inclines toconverge to the local optimal that is the same as the original mean-field iteration.
As a result, thebenefit of correlation-aware inference provided by the Gaussian copula would be diminished.
Thisexplains why the iterative step K cannot be set too large.
The number of Monte Carlo (MC) sampling m playsan important role in the proposed method.
It is widely known that increasing number of sampleswould reduce the variance of MC sampling.
Therefore, using larger m would result in a betterapproximation of the gradient of multilinear extension ∇ ψ f F θ mt and thus better performance.
Thishypothesis is validated by the empirical results show in the second row of Fig.
It can be seen thatas the sample number increases, the performance rises steadily at first and then gradually convergesinto a certain level.
Undoubtedly, a large number would increase the computational complexity.
Inthis regard, we uniformly set it as 5 in all experiments.
Impact of the Lower-rank Perturbation Lower-rank perturbed covariance matrix enables theproposed method to model the correlation information of elements in the set.
To investigate itsimpacts, we evaluate the performance of EquiVSet copula under different values of rank v .
The resultsare demonstrated in the third row of Fig.
Notably, the proposed model with v = 0 is equivalent to
EquiVSet ind .
It can be seen that as the number of ranks increases, the performances also increase,indicating the hypothesis that employing the variational distribution with correlations can increasethe model’s representational capacity and thereby results in a better approximation in turn.
It is worthnoting that the most significant performance improvement is observed between the models with v
= 0and v = 1 , and then as the value of v continues to increase, the improvement becomes relativelysmall.
This indicates that it is feasible to set the v to a relatively small value to save computationalresources while retaining competitive performance.