D IFFERENTIABLE E XPECTATION -M AXIMIZATION FOR S ET R EPRESENTATION L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
We tackle the set2vec problem, the task of extracting a vector representation from an input set comprised of a variable number of feature vectors.
Although recent approaches based on self attention such as (Set)Transformers were very successful due to the capability of capturing complex interaction between set elements, the computational overhead is the well-known downside.
The inducing-point attention and the latest optimal transport kernel embedding (OTKE) are promising remedies that attain comparable or better performance with reduced computational cost, by incorporating a fixed number of learnable queries in attention.
In this paper we approach the set2vec problem from a completely different perspective.
The elements of an input set are considered as i.i.d.
samples from a mixture distribution, and we define our set embedding feed-forward network as the maximum-a-posterior (MAP) estimate of the mixture which is approximately attained by a few ExpectationMaximization (EM) steps.
The whole MAP-EM steps are differentiable operations with a fixed number of mixture parameters, allowing efficient auto-diff backpropagation for any given downstream task.
Furthermore, the proposed mixture set data fitting framework allows unsupervised set representation learning naturally via marginal likelihood maximization aka the empirical Bayes.
Interestingly, we also find that OTKE can be seen as a special case of our framework, specifically a single-step EM with extra balanced assignment constraints on the E-step.
Compared to OTKE, our approach provides more flexible set embedding as well as prior-induced model regularization.
We evaluate our approach on various tasks demonstrating improved performance over the state-of-the-arts.
I NTRODUCTION
Effectively learning succinct and salient feature representation for complex structured data is a central part of recent deep learning.
In particular the set structured data, arising in many important scientific tasks including bioinformatics and NLP tasks, poses several challenges to deal with varying numbers of elements and the constraints of permutation invariance, in contrast with conventional instance-based representation learning (Zaheer et al., 2017; Edwards & Storkey, 2017; Lee et al., 2019a; Skianis et al., 2020).
Perhaps the key to successful set representation learning is to capture the interaction between elements effectively, implemented as self attention (Bahdanau et al., 2015) and (set)transformer models (Vaswani et al., 2017; Lee et al., 2019a).
Despite their superb performance in numerous application problems, one well-known drawback of the self-attention approaches is high computational (time and memory) overhead.
Among others, the inducing-point attention (Lee et al., 2019a) and the latest optimal transport kernel embedding (OTKE) (Mialon et al., 2021) are promising remedies that attain comparable or better performance with reduced computational cost.
This is enabled by incorporating a fixed number of trainable reference vectors as queries in attention, and especially OTKE, with the attention scheme based on optimal transport, showed state-of-the-art results on several tasks in bioinformatics and NLP.
Although OTKE is motivated from self attention, in this paper we show that OTKE can be derived from a completely different perspective: it is the maximum likelihood estimate of a Gaussian mixture model obtained by a single Expectation-Maximization (EM) step (Dempster et al., 1977) with some constraints on the E-step, which aims to model set elements as i.i.d.
data samples.
From this perspective we propose several interesting extension/generalization such as: i) we perform multiple
EM steps instead of a single one, ii) unlike treating the references (model parameters) just as an initial iterate of the EM, we also place a prior distribution on the references to incorporate the prior effect to the output embedding in a more principled Bayesian manner.
With this generalization, our set embedding function becomes a (differentiable) feed-forward network defined as a multi-step EM for the maximum-a-posterior (MAP) estimate of a Gaussian mixture.
Furthermore, the proposed mixture set data fitting framework allows unsupervised set representation learning naturally via marginal likelihood maximization aka the empirical Bayes.
On several set/sequence representation learning problems including biological sequence classification problems and NLP tasks, our approach shows considerable improvement over the state-of-the-arts including OTKE.
Our main contributions are summarized below:
We derive OTKE by a novel EM-based mixture fitting perspective, which allows us to generalize it to a more flexible and principled set embedding function that can offer more modeling options, prior-induced model regularization, and unsupervised learning via empirical Bayes.
The whole MAP-EM steps are differentiable operations with a fixed number of mixture parameters, enabling efficient auto-diff back-propagation for any given downstream task.
Our approach shows improved results over state-of-the-arts on various set embedding tasks in bioinformatics and NLP that require effective interaction modeling among set elements.
B ACKGROUND ON OTKE
OTKE (Mialon et al., 2021) is a computationally efficient trainable model for representing a variablelength set, by aggregating a variable number of features/elements in a set into a fixed-size vector.
The input to the model is a set S
= { x 1 , . . .
, x n } ( n varies from instance to instance), where each x i ∈ R
d is an input feature vector which may be an output of some feature extractor network, thus involving some trainable parameters in it.
For instance, in the OTKE paper, they often consider a parametric Nyström approximation of the RKHS ϕ ( x ) (Williams & Seeger, 2001).
We abuse the notation throughout the paper, denoting by x both a raw input element and an extracted feature ϕ ( x ) , where the distinction will be clear from context.
The OTKE model’s set embedding function, emb : S → R D where S ∋ S , is defined as follows:
[Equation]
where cat () indicates concatenation of vectors in column, and hence the output is of dimension
D = d · p .
The trainable parameters of the embedding function are the p reference vectors
Z = { z 1 , . .
. , z p } with z j ∈ R d , where they serve as queries to determine the weights Q in the attention form (1).
Instead of following the popular dot-product strategy (Vaswani et al., 2017), the weight matrix Q is the (unique) solution of the optimal transport problem (Villani, 2008; Cuturi, 2013) between S and Z with the cost matrix C defined as the negative kernel C ij = − k
( x i , z j ) , that is,
[Equation]
where the objective is augmented with the entropic term (impact ϵ ), H ( Q )
=
− (cid:80) ij Q ij log Q ij , allowing the efficient Sinkhorn-Knopp (SK) algorithm to be applied (Cuturi, 2013).
This OT-based weighting is shown to be more effective than the dot-product in the biological application domain (Mialon et al., 2021).
The key difference from dot-product is the balanced assignment constraints (cid:80)
i Q ij = 1 /p for all j = 1 , . .
. , p , and the main motivation for this is to have all references z j contribute equally to describing the set S .
This may be a reasonable strategy in certain scenarios to prevent a small number of references from affecting the attention dominantly.
Note that (2) can be solved by a few fixed point matrix scaling iterations (SK algorithm):
The SK algorithm finds the optimal solution as Q =
Diag ( u ) A Diag ( v ) , where A ij = exp( − C ij /ϵ ) and the vectors u ∈ R n + and v ∈ R p + are the fixed points of u i = 1 n / ( Av ) i , v j = 1 p / ( A ⊤ u ) j for i = 1 , . .
. , n , j = 1 , . .
p .
The fixed point iteration usually converges quickly after a few iterations.
Since the whole SK iterations constitute the feed-forward pass for the final output (1), hence OT is differentiable.
Note that as the transport cost in OT is defined as the negative kernel, it becomes simply the dotproduct in the RKHS, i.e., k ( x, z ) =
x ⊤ z (with the RKHS embedding notation, ϕ ( x ) ⊤ z ), where z lives in the same RKHS.
OTKE can be extended to a multi-head network, for which we run the above
OTKE procedure multiple times independently, each with different parameters Z , then concatenate the outputs.
OTKE has several notable advantages over the self attention.
First, it is computationally efficient, scalable to very large dataset.
Unlike the quadratic cost of the self attention, OTKE uses a fixed number ( p ) of references serving as queries.
Secondly, it allows unsupervised learning where the model parameters Z can be learned without target (class) labels via the Wasserstein (or k-means alternatively) clustering of set elements as per the OT problem.
This unsupervised learning can also serve as pre-training for the potentially unknown target downstream tasks.
We also note that regarding the first benefit, the formula (1) is also very similar to the inducing-point attention of
SetTransformer (Lee et al., 2019a) (known as the PMA layer) except for the OT-based weighting scheme.
In Appendix B.2 we provide a proof showing that PMA is indeed a special case of OTKE.
E XPECTATION M AXIMIZATION V IEW OF OTKE
While OTKE is motivated and derived from the attention mechanism, in this section we derive the
OTKE embedding formulas (1-2) from a completely different perspective, namely the EM-based maximum likelihood learning of the Gaussian mixture for set data.
Specifically we consider a mixture of p Gaussians with equal mixing proportions and fixed shared spherical covariances, that is,
[Equation]
where c denotes the component latent variable, and the learnable parameters are only the means θ = { z 1 , . .
. , z p } .
By regarding the set elements x i ’s as i.i.d.
data samples, the log-likelihood of the set S
= { x 1 , . . .
, x n } under the model is log p ( S | θ ) =
(cid:80) i log p ( x i | θ ) .
Now we perform one E-step in the EM toward maximizing the log-likelihood, which is derived by the following Jensen lower bound:
[Equation]
where q ( j | i ) is the variational distribution and the bound is tight if q ( j | i )
=
p ( c i =
j | x i ) .
Although the latter can be derived in a closed form, we aim to formulate an optimization problem for q ( j | i ) .
We use a joint variational form instead, i.e., q ij :=
q ( j | i ) · q ( i ) with q ( i ) :
= 1 /n .
Then maximizing the lower bound in (4) with respect to the ( n × p ) matrix q = ( q ij ) is equivalent to solving:
[Equation]
, where
[Equation]
By having the quadratic kernel k ( x, z ) = − 12 || x − z
||
which is (conditionally) positive definite (Schölkopf & Smola, 2002), we see that (5) reduces to (2) except for the first constraints (cid:80) i q ij = 1 /p (the other constraints (cid:80) j q ij = 1 /n are automatically met by construction).
Note that the quadratic kernel approximates the dot-product (in the RKHS) if the norms of the points are roughly constant.
Hence the conventional E-step of the EM algorithm coincides with the OT-based attention weight computation (2) without the balanced assignment constraints (cid:80)
i q ij = 1 /p .
Conversely, if we consider to maximize the lower bound (4) with the extra balanced assignment constraints on the variational distribution q = ( q ij ) , the corresponding E-step exactly reduces to the OT solution (2).
Once the E-step is done (i.e., q is found), we perform the M-step, that is, maximizing the lower bound in (4) with respect to θ while q is fixed.
This admits a closed-form update equation,
[Equation]
Note that if we imposed the balanced assignment constraints, the denominator of (6) is constant ( 1 /p ), and by concatenating z newj ’s, it exactly reduces to the OTKE formula (1) (up to a constant factor).
But, even without the constraints, the M-step in (6) suggests to normalize (balance) the scales of the transformed vectors (cid:80) i q
ij x i across j = 1 , . . .
, p , according to the reference marginals (cid:80) i q ij (how many points are assigned to z j ).
This establishes a very close connection between OTKE and the EM-updated mean parameters of the Gaussian mixture.
Note that our derivation showed that OTKE is recovered by just one EM step, and there is possibility of extending it by taking multiple EM steps.
Furthermore, this opens up other interesting generalization of OTKE through the (EM) mixture set data fitting perspective.
We summarize below reasonable and interesting generalization/extension of OTKE in four folds:
Instead of fixing mixing proportions and (shared) covariances in the Gaussian mixture as in (3) we can also incorporate them as parameters of the model in conjunction with the means.
Instead of just one EM step, we can take multiple ( k ) EM steps with k as a hyperparameter.
It becomes an option whether the E-step is solved with the balanced assignment constraints (solving the OT problem) or not (regular E-step).
This enables more flexible modeling. 4.
OTKE can be seen as an (single-step) EM which starts from the references z j ’s as an initial iterate.
Not just using them as an initial iterate, we can also place a prior distribution on the Gaussian mixture parameters, that is, p ( θ ) which serves as learnable model parameters.
This allows us to define the embedding as the posterior distribution emb ( S )
=
p ( θ | S ) .
These extensions of OTKE are incorporated into a single framework, which becomes a novel set embedding function viewed as a differentiable MAP Expectation-Maximization for a Gaussian mixture.
D IFFERENTIABLE
[Equation]
We consider a Gaussian mixture model with non-fixed mixing proportions and covariances:
[Equation]
where the covariances V j are typically chosen as diagonal ( d × d ) matrices 1 to reduce the computational overhead.
Given the mixture θ , the set elements are considered i.i.d., p ( S | θ ) =
(cid:81) ni
=1 p ( x i | θ ) .
Now we impose a prior distribution on θ where the Dirichlet-Normal-Inverse-Wishart ( Dir-NIW for short) would be a fairly natural choice:
[Equation]
This defines a graphical model (whose diagram shown in Fig.
1) where Λ generates mixture parameters θ ’s, one for each set S , and each θ generates x 1 , . . .
, x n ∈ S .
We emphasize that the reason we consider the prior is that we can express the process of how the model parameters impact on the final output in a more principled manner, beyond just having them as an initial iterate.
For simplicity, the prior parameters η, λ, ν are shared across the components.
That is, Dir ( π ; η ) ∝ (cid:81)
j π η − 1 j where η is scalar, and NIW ( z j , V j ) =
N ( z j ; µ j , V j /λ ) ·
IW ( V j ; Σ j , ν ) for all j = 1 , . . .
, p , where λ, ν are scalar, µ j ∈ R d , and Σ j are positive diagonal matrices.
Then we define our set embedding function emb ( S )
→ R D as the MAP estimate of the Gaussian mixture θ , denoted by θ MAP ( S ) :
= arg max θ p ( θ | S, Λ) , which is approximated by the k -step EM update.
And the prior parameters Λ will constitute the model parameters of the embedding function emb ( S ) to be trained.
The EM update equations for the MAP can be derived in closed forms.
More specifically, from the following lower-bound of the log-posterior (“ = c ” means equal up to constant)
[Equation]
the E-step maximizes the lower bound with respect to q (with θ fixed), which yields
[Equation]
The former formula in (10) is the regular E-step without the balanced assignment constraints, and the latter is the solution with the constraints where SK ( C, ϵ ) indicates the solution of the OT problem via the SK algorithm (c.f., (2) and (5)).
Note that we multiply the latter by n to have conditional distributions q ( j | i ) .
The M-step is the log-prior regularized weighted log-likelihood maximization with respect to θ (with q fixed), which admits a closed-form solution (Appendix B.1 for derivations)
[Equation] [Equation]
where all operations including squaring are element-wise.
The τ values originate from the prior, defined as: τ η = η − 1 , τ λ = λ , τ 1 = 1 , and τ ν = ν
+ d + 2 .
Note that τ ’s represent the strength of the prior impact: turning them off to 0 completely ignores the prior and leads to the maximum likelihood M-step update.
Instead of dealing with them differently, we simplify it by collapsing them into a single hyperparameter τ (= τ η = τ λ
= τ 1 = τ ν ) .
Then the simplified M-step becomes: π
[Equation]
Now we can control τ to impose the prior impact in the set embedding: smaller (larger) τ for lower (higher, resp.)
prior impact.
We treat τ as the hyperparameter (e.g., chosen by validation) while Λ = { µ, Σ } constitutes the learnable prior parameters.
As an initial iterate θ 0 of the EM, we use the mode of prior p ( θ | Λ) , which equals
[Equation]
Note that θ 0 is a function of Λ .
We deal with multiple ( k ) EM steps, θ
→ θ 1 → · · · → θ k , each following E-step (10) and M-step (12), and the final θ k is a function of Λ .
Overall, from Λ to the final iterate θ k constitutes a feed-forward pass, and it is composed of all differentiable operations.
The final output emb ( S ; Λ) = θ k ≈ θ MAP ( S ; Λ) is then transformed to a vector representation by concatenating the parameters in θ k in a fixed order (e.g., cat ( π k 1 , z k 1 , V k 1 , π k 2 , z k 2 , V k 2 , . . .
, π kp , z kp , V kp ) ).
The feed-forward network of emb ( S ; Λ) is visualized in Fig.
2 (Left).
Multi-head representation.
We can also extent the model to have multi-head representation, which is simply done by replicating the process H times with different prior parameters Λ 1 , ..., Λ H .
See
Fig.
2 (Right).
We have H outputs from the MAP-EM set embedding networks, denoted by θ ( h ) for h = 1 , . . .
, H .
Then the final output is a pooling from these θ ( h ) ’s, and we propose three different options for this multi-head pooling (MHP): i) PC (parameter-concatenation) simply concatenates θ ( h ) ’s in a fixed order, ii) SB (select-best) computes the marginal data likelihood log p ( S | Λ h ) for each h and select/return the one with the largest score only, and iii) SB2 outputs the additional one-hot encoding of the selected mixture’s ID in addition to SB’s output.
Note that PC makes sense in that we have the supports of the input set covered by all H mixtures together, while SB strategy is motivated from the conventional practice in EM learning where we run EM algorithm multiple times with differential initial iterates and take the one with the highest likelihood score.
SB2 additionally links the output to the specific mixture parameters selected.
Unsupervised learning.
DIEM can be typically used in supervised learning, where the output of DIEM ( S ) is fed into subsequent layers (including the classification head) in a prediction network, and updated in an end-to-end learning manner. Beyond this, DIEM can be trained with only sets { S 1 , . . .
, S m } without target labels.
This is naturally done by marginal data likelihood maximization for mixture learning, where we maximize (cid:80) mi =1 log p ( S i ) with respect to the prior parameters Λ .
We can perform EM for optimizing the objective (in this case, the off-line EM until convergence with the entire training sets, without the auto-diff backprop mode), or alternatively the k-means clustering.
The latter leads to the unsupervised learning algorithm similar to that of OTKE (Mialon et al., 2021).
R ELATED W ORK
Set representation learning.
Prior to the attention-based approaches, the neural statistician (Edwards & Storkey, 2017) aims to build a model that learns the statistics of an input set, while DeepSet (Zaheer et al., 2017) has the element-wise nonlinear transformation followed by a simple mean or max pooling to aggregate features.
Although the universality of this network architecture was partially proved, the limitation was found (Wagstaff et al., 2019), and network structures more suitable for exchangeability and permutation invariance were suggested (Garnelo et al., 2018; Bloem-Reddy & Teh, 2019).
Attention-based methods.
Self attention was shown to be very successful for modeling set-structured data (Vinyals et al., 2016;
Yang et al., 2018; Ilse et al., 2018; Kim et al., 2019).
While there were several sophisticated attempts to reduce the quadratic computational cost of self attention (Wang et al., 2020; Kitaev et al., 2020), the inducing-point approaches that introduce learnable parameters as queries are especially suitable for set inputs (Lee et al., 2019a; Pritzel et al., 2017; Skianis et al., 2020; Mialon et al., 2021).
The attention mechanism was analyzed rigorously in (Tsai et al., 2019) providing unifying kernel interpretation.
Applying multiple EM steps in our DIEM is architecturally similar to the recent Perceiver model (Jaegle et al., 2021) with multiple latent-attention layers.
Mixture-based representation.
There was an attempt to define the similarity between two images (image as a set of fixed feature vectors, e.g., Fisher vectors (Perronnin et al., 2010)) using mixture models about a decade ago (Liu & Perronnin, 2008).
Similar to ours, they represent each image as a
MAP estimated mixture density, and the similarity between two images is measured as a divergence between the corresponding mixture densities.
However, instead of treating the pipeline of whole EM steps as a differentiable object and learning the prior parameters as in our DIEM, they merely apply a single MAP-EM step to obtain a fixed density for each image, where the initial mixture for EM is also fixed from global offline learning with entire training images.
E VALUATION
We evaluate our DIEM model empirically on two different types of tasks: i) counting and clustering problems (Sec. 6.1, 6.2) to verify the model’s capability of learning general set representations by modeling interaction between set elements, ii) large-scale biological sequence classification and NLP tasks (Sec. 6.3, 6.4, 6.5) to test the performance of the proposed model on real-world problems in both supervised and unsupervised settings.
The hyperparameters in our DIEM include: p (the mixture order), H (the number of heads), k (the number of EM steps), τ (prior impact), and the multi-head pooling strategy (either of PC, SB, or SB2).
We report the results of the best combinations that are selected by cross validation.
The empirical study on the impact of these hyperparameters is summarized in Appendix C.4.
The other option is the mixture modeling: we learn both means and covariances for counting, clustering, and SCOP 1.
datasets, while covariances are fixed for SST-2 and DeepSEA.
6. OMNIGLOT U NIQUE C HARACTER C OUNTING
This task, originally devised in (Lee et al., 2019a), is useful to verify capability of modeling interaction between set elements.
From the OMNIGLOT dataset (Lake et al., 2015) which contains 1,
[Table]
different handwritten characters with 20 examples each, we build a set by randomly selecting n images, and the task is to predict the number of unique characters in the set.
Following the data construction protocol from (Lee et al., 2019a), we split the 1,623 characters into training/validation/test splits so that the test dataset contains only unseen characters.
Specifically, we randomly choose the number of images n uniformly from { N min , . . .
, N max } , then choose the number of unique characters c randomly from { c min , . .
. , n } .
Then we randomly sample c different classes (characters), and sample n images from those classes while ensuring that all classes have at least one image.
We form two experimental settings: i) Small set ( N min = 6 , N max = 10 , c min = 1 ) originally used in (Lee et al., 2019a) and ii) (a more challenging)
Large set ( N min = 11 , N max = 30 , c min = 5 ).
For the models, we adopt a similar feature extractor architecture as (Lee et al., 2019a): first apply four
Conv-BN-ReLU layers to each (element) image to have feature representation ϕ ( x ) , then perform the set embedding emb ( S ) whose output is fed into a fully connected layer to return the Poisson parameter λ .
The final loss function is the negative Poisson log-likelihood − log p ( c | λ ) .
The batch size is 32 (sets) for both datasets.
Results.
We report the 0/1 prediction accuracy averaged over 10 random runs in Table 1.
OTKE, after validation, uses ( p = 100 , H = 1) for the small set and ( p = 50 , H = 2) for the large set.
Our DIEM takes ( p = 20 , H = 3 , k = 2 , τ = 10 − 3 , PC) for the small set and ( p = 50 , H = 2 , k = 2 , τ = 10 − 3 ,
PC) for the large set.
For both datasets, our DIEM outperforms the competing approaches by large margin.
Unlike OTKE, we have prior regularization and multiple EM steps which appear to be important to improve the performance.
6. A MORTIZED C LUSTERING
The task of amortized clustering is to learn a function f ( S )
→ Θ where S = { x 1 , . . .
, x n } is a (data) set of n points, and Θ = { α, m, C } is the parameters of the mixture of Gaussians, p ( x ; Θ)
= (cid:80)
j α j N ( x ; m j , C j ) , best fit to the data.
We follow the setup similar to (Lee et al., 2019a) where instead of providing the supervised data of pairs 2 { ( S, Θ) } , we are only given the set data S , and the loss function is the negative log-likelihood − log p ( S | Θ) .
Following (Lee et al., 2019a), we form two datasets: i)
Synthetic dataset where n 2D data points are sampled from a randomly generated mixture of four Gaussians, and ii) CIFAR-100 images where we first sample four classes out of 100, and sample n images from the classes to form a set S .
In both cases n ∈ { 100 , . . . , 500 } .
For
CIFAR-100, each image is represented as a 512-dim vector from a VGG network pre-trained on the training set (Simonyan & Zisserman, 2014).
For our DIEM (and OTKE), we first apply a simple linear layer transformation to have features ϕ ( x ) .
Then the output of the set embedding emb ( S ) is fed into a fully connected layer to have the Gaussian mixture parameters Θ .
For the competing DeepSet and SetTransformer models, we faithfully follow the network structures in (Lee et al., 2019a).
The details are summarized in Appendix C.2.
Results.
The results are reported in Table 2 where the test log-likelihood scores and the adjusted rand index (ARI) are shown for the synthetic and CIFAR-100 datasets, respectively.
As references, we also report the oracle performances: (Synthetic) the true mixture used to generate the data, and (CIFAR-100) the full offline EM-trained mixture model.
Although SetTransformer greatly outperforms DeepSet, our DIEM exhibits further improvement over SetTransformer by large margin.
Relatively poor performance of OTKE may be due to only one-step EM update, compared to our
DIEM that used k = 3 EM steps.
For instance, DIEM also showed degraded performance when k = 1 is used (Fig. 10 in Appendix C.4).
Another reason of success of DIEM is the prior-induced regularization ( τ = 0 . 01 ) whereas OTKE has no such regularization (Fig. 11 in Appendix C.4).
[Table]
[Table]
6. P ROTEIN F OLD C LASSIFICATION T ASK ON SCOP 1.
Protein fold classification is the well-known task in bioinformatics where the goal is to predict the fold class for a given protein sequence.
We use the preprocessed data of the SCOP version 1.75 and 2.06 from (Hou et al., 2019; Mialon et al., 2021), which consists of 19,245 sequences (14,699/2, training/validation from SCOP 1.75 and 2,533 test from SCOP 2.06).
Each input protein is a sequence of amino acids whose length ranges from tens to thousands, and each amino acid is represented as a 45-dimensional vector based on the PSSM, secondary structure, and more.
It is a multi-classification task to predict one of the 1,195 different folds to which a protein sequence belongs.
For the models, we adopt the overall network architecture from (Chen et al., 2019a; Mialon et al., 2021): each input sequence is treated as a set S , each set element ϕ ( x ) is a learnable Gaussian kernel mapping on 10-mers via Nyström RKHS approximation with L anchor points (we use L =for unsupervised learning, and L = 512 / 128 for supervised learning), and the output of the set embedding emb ( S ) is fed into the final softmax linear classification layer.
For the competing methods, we replace the set embedding emb ( S ) by: the global mean pooling layer for CKN (Chen et al., 2019a), the OTKE layer (Mialon et al., 2021), and our DIEM layer.
In the unsupervised learning, we sequentially train-then-fix the first two layers without class labels: train ϕ ( x ) via k-means, fix it, do unsupervised training of emb ( S ) , and fix it, after which the last classification layer is trained with class labels.
We use the k-means clustering for unsupervised learning.
In the supervised learning case, the whole network is trained end-to-end.
Results.
Classification accuracy (top 1/5/10) for unsupervised and supervised learning are shown in
Table 3.
In our DIEM, we fix the covariances V as identity and perform k = 1 OT E-step for the unsupervised learning, while we learn priors for V and take k = 2 regular E-steps for the supervised learning.
The details of the hyperparameters are summarized in Table 12 of Appendix C.3.
The result shows that our approach outperforms all state-of-the-arts for both learning settings.
Compared to
OTKE, the increased accuracy of our DIEM can be mainly attributed to the prior-induced model regularization ( τ = 10 − 3 for unsupervised learning) and multiple EM steps ( k = 2 for supervised learning).
See also Fig. 12 in the Appendix for the hyperparameter impact analysis.
6. NLP T ASK : S ENTIMENT C LASSIFICATION ON SST-
The dataset (Socher et al., 2013) consists of 70,042 movie reviews with positive/negative binary sentiment.
We follow the experimental protocol from (Mialon et al., 2021) where the original 67,
[Table]
training data are split into 80% / 20% training/validation sets, and the 872 validation reviews form a test set.
This is because test evaluation on the original test set requires online submission in the
GLUE (Wang et al., 2019)
leaderboard.
Similar to (Mialon et al., 2021), we adopt the 768-dimensional word vectors from the pre-trained BERT model (Devlin et al., 2019; Wolf et al., 2019), and use the
Gaussian RKHS mapping ϕ ( x ) approximated by the Nyström method with 2048 (unsupervised) or 64 (supervised) filters.
Overall network architectures and learning options are also similar to (Mialon et al., 2021).
Results.
Table 4 summarizes the results.
For the unsupervised learning, the proposed DIEM attains the best accuracy, even outperforming the pre-trained BERT model with the popular [CLS] embedding or mean pooling.
Unlike OTKE that takes a large number of parameters ( p = 300 , H = 1 ), our
DIEM has far fewer parameters ( p = 3 , H = 2 , k = 1 with PC multi-head pooling).
When we trained
OTKE with the same complexity architecture, it only attained accuracy 86 .
6 ± 0 .
The improved accuracy of DIEM over OTKE is mainly due to the prior-induced regularization (we use τ = 10 − 3 ) and the regular E-step (with the OT E-step, the performance is on par, 86 .
3 ± 0 . 8 ).
For the supervised learning, both DIEM and OTKE has p = 30 , H = 4 , and the performance difference mainly comes from the regular E-step taken by DIEM ( τ = 10 − 6 , SB).
Although DIEM performs comparably well, it falls short of the fine-tuned BERT model.
As alluded in (Mialon et al., 2021), the set embedding approaches may not be very attractive in this case due to short sentences (small cardinality sets).
6. C HROMATIN P ROFILE D ETECTION ON D EEP SEA
[Table]
Finally, we test our DIEM on the large-scale DeepSEA dataset (Zhou & Troyanskaya, 2015) with about five million genomic sequences.
It is a multi-label classification task to predict 919 chromatin profiles.
As in (Mialon et al., 2021), the overall network architecture follows the 1D Conv layer from DeepSEA, while the pooling layer after the Conv is replaced by the DIEM layer.
The RKHS feature mapping is not used (identity mapping), and the Gaussian positional encoding 3 was adopted.
The supervised learning results are shown in Table 5, where we have marginal improvement over DeepSEA, and the same performance as OTKE using the ( p = 64 , H = 1 , k = 2 , τ = 1 ) hyperparameter choice.
C ONCLUSION
In this paper we proposed a novel differentiable EM model for set representation learning.
Our model is built from the perspective of fitting a Gaussian mixture model to the set data that are viewed as i.i.d.
samples, which offers more flexibility and prior-induced model regularization in a principled
Bayesian manner.
The proposed model is also shown to generalize the recent set embedding models based on optimal transport and attention, leading to a computationally efficient model with superb performance on tasks in bioinformatics and NLP.
The mixture fitting perspective can potentially solve other important problems in set representation learning.
For instance, deciding the number of components/references can be translated into the mixture order selection problem, and tackled by well-known information criteria methods.
This will be left as interesting future research.
Appendix
We provide some formal derivations, experimental details and additional results that are skipped in the main paper.
A C OMPLEXITY AND R UNNING T IME
In this section we analyze the complexity, especially the number of parameters and time complexity for a forward pass, of our DIEM.
We also compare the wall clock running time for OTKE (Mialon et al., 2021) and our DIEM with different numbers of EM steps. A. C OMPLEXITY Notation. First, we recall the notation: The input set S
= { x 1 , . . .
, x n } contains n d -dimensional feature vectors x i ∈ R d .
We consider the OTKE model with p references { z j } pj =1 (each z j ∈ R d ) as parameters.
To be comparable, DIEM has p Normal-Inverse-Wishart (NIW) parameters { µ j , Σ j } pj = ( µ j ∈ R d , Σ j is a ( d × d ) diagonal matrix), each of which defines a prior for each component
N ( z j , V j ) in the Gaussian mixture.
For simplicity, we assume single-head representations (i.e.,
H = 1 ) for both models, since multi-head ( H heads) representation increases all compute and memory linearly by H .
Numbers of parameters.
The numbers of parameters are: pd (OTKE) vs. 2 pd (DIEM), Hence
DIEM requires only twice as many parameters as OTKE.
Time complexity.
First, one forward pass in OTKE consists of: [O1] computing the cost/kernel matrix C which is ( n × p ) with C ij = − k
( x i , z j ) , [O2] solving the OT problem (2) by the SinkhornKnopp (SK) matrix scaling algorithm, and [O3] performing attention-weighted sums (1) to have the final embedding.
The least amount of time for [O1] is O ( npd ) to inspect each element of x i and z j once (e.g., linear kernels and squared-distance based kernels require O ( npd ) ).
Denoting the number of SK iterations by M , [O2] takes O ( npM ) , while [O3] requires O ( npd ) .
Thus, the time complexity for one forward pass in OTKE is O
( np ( M + d )) .
For DIEM with k EM steps, [D1] each E-step (10) amounts to computing the Gaussian likelihoods, and additionally running the SK algorithm if the OT E-step is used.
[D2] Each M-step (12) mainly consists of the weighted sums.
[D1] takes O ( npd ) for the regular E-step, and O ( np ( M + d )) for the
OT E-step, while [D2] requires only O ( npd ) .
Hence, the time complexity for one forward pass in our DIEM is O ( knp ( M + d )) , only k times as much as OTKE.
We summarize it in Table 6.
The actual wall clock running times are compared in the next section.
A. (F ORWARD P ASS ) R UNNING T IME
To measure the forward pass time for OTKE and our DIEM, we consider first the large-scale SCOP 1.75 dataset: the input feature dimension d = 512 (from Gaussian RKHS) and the set cardinality n = 1091 .
To see the impact of the set cardinality, we inject random noise features to the sets to make twice and four times larger sets ( n = 2182 and n = 4364 ).
For fair comparison, both OTKE and DIEM use one head ( H = 1 ) with the number of references (or mixture components) p = 100 .
The batch size is 128.
We run all models on the same machine, Core i7 3.50GHz CPU and 128 GB
RAM with a single GPU (RTX 2080 Ti).
Fig.
3 compares the average per-batch forward pass time for OTKE and DIEM with different EM steps ( k = 1 to 4 ).
Although DIEM has an option to choose the E-step type (either regular E-step or OT E-step), we take all OT E-steps to have fair comparison with OTKE.
(For the running time comparison between these two E-step types, please refer to Fig.
As shown in Fig.
3, increasing the number of EM steps in our DIEM incurs only a small amount of computational overhead.
The impact of the set cardinality n on the running time exhibits similar trends for both OTKE and DIEM.
In Fig.
4, we compare running time of OTKE, SetTransformer (Lee et al., 2019a), and DIEM (with k = 3 EM steps).
Here we use a minimal network structure for SetTransformer: one SAB block for the encoder, one PMA and one SAB block for the decoder.
For larger sets ( × 2 and × 4 ) we reduce
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
the batch size for SetTransformer to fit into the GPU memory.
We need not reduce the batch size for
OTKE and DIEM as they fit into the GPU memory for all three set cardinalities.
The result shows that OTKE and DIEM are scalable to large sets whereas SetTransformer takes significantly longer time, which is well aligned with the fact that self-attention suffers from quadratic computational cost.
In Fig.
5, we compare the running time between OT E-steps and regular E-steps in our DIEM.
The maximum number of Sinkhorn-Knopp iterations for solving the OT problem is set to 10.
While the regular E-step is consistently faster than the OT E-step, the increased overhead of the OT E-step seems to be minor compared to the overall running time.
For smaller datasets, we consider the OMNIGLOT counting (large set) and CIFAR100 clustering datasets.
The running time of OTKE and DIEM on these two datasets is depicted in Fig.
For the
OMNIGLOT large set, we set p = 50 , H = 2 , d = 64 for both models where n ∈ { 11 , . .
. , 30 } and batch size is 32 sets.
For CIFAR-100 clustering, we test with p = 4 , H = 5 , d = 512 for both models where n ∈ { 100 , . . .
, 500 } and batch size is 10 sets.
Our DIEM takes regular E steps.
We vary the number of EM steps k in DIEM from 1 to 5.
The additional forward pass time incurred by extra EM steps in DIEM
looks permissible, the overall running time of DIEM being constant factor comparable to that of OTKE.
In Appendix B.1, we provide detailed derivations for MAP EM update equations (10-11) in the main paper.
In Appendix B.2, we formally show that the inducing-point attention module of
SetTransformer (Lee et al., 2019a) known as PMA, is actually a special case of OTKE (Mialon et al., 2021). [Figure] [Figure] [Figure]
B. MAP EM D ERIVATIONS
E-step is optimization of the lower bound in (9) with respect to q ( j | i ) with θ fixed.
Hence it is identical to the maximum likelihood E-step for the lower bound of (4).
Now we derive the M-step, maximizing the lower bound in (9) with respect to θ with q ( j | i ) fixed.
That is,
[Equation]
The log-prior term of the Dir-NIW density can be expanded as (“ = c ” refers to equality up to constant ):
[Equation]
The weighted log-likelihood term in (14) can be written as:
[Equation]
Combining (15) and (16) [Equation] we take derivatives with respect to θ and solve for θ as follows:
[Equation]
where the equality between (19) and (20) can be easily verified from the optimal z ′ j in (18).
B. PMA OF S ET T RANSFORMER IS A S PECIAL C ASE OF OTKE
The PMA (inducing-point attention) module of SetTransformer (Lee et al., 2019a) is an attention block that incorporates a fixed number of learnable parameters as queries, and plays an important role in SetTransformer to have a fixed size vector representation for a variable-length input set.
Formally,
PMA module has learnable p vectors Z = { z 1 , . .
. , z p } , and for a given set S
= { x 1 , . . .
, x n } , it performs the (dot-product) attention operation with each z j as a query and S as keys and values.
Hence each z j is transformed as follows:
[Equation]
where W q and W are the weight matrices for queries and keys/values, respectively.
Since z j itself is a learnable parameter vector, it is redundant to introduce W q , and we can subsume it under z j , which leads to
[Equation]
Next, we build a special OTKE module by having the RKHS feature mapping ϕ ( x i ) =
Wx i , which corresponds to the linear kernel k ( x i , z j ) =
z ⊤ j ϕ ( x i ) =
z ⊤ j Wx i .
Furthermore, we consider the OT [Table] [Equation] problem between Z and S without the second constraints (cid:80) pj =1 Q ij = 1 /n, ∀ i .
That is,
[Equation]
with C ij = − k
( x i , z j ) = − z ⊤ j Wx i .
We can solve (23) using the Lagrange multiplier, [Equation]
[Equation]
We set the derivative to 0 and solve for Q ij as follows
[Equation]
for some constant α j .
Applying the constraints (cid:80)
ni =1 Q ij = 1 /p determines α j and Q ij as:
[Equation]
By choosing the entropic regularizer trade-off ϵ = √ d and applying the OTKE embedding formula z j → √ p (cid:80) i Q ij ϕ ( x i ) , we get OTKE
[Equation]
27) which is equivalent to PMA (22) up to a constant factor.
Hence PMA of SetTransformer can be viewed as an OTKE module without the equal element assignment constraints (cid:80) pj =1 Q ij = 1 /n, ∀ i .
C E XPERIMENTAL D ETAILS AND A DDITIONAL R ESULTS C. OMNIGLOT U NIQUE C HARACTER C OUNTING
(Original) Small set.
This experimental setup is the same as that of (Lee et al., 2019a).
The sets of
OMNIGLOT images are generated by: set cardinality n ∼ { 6 , 7 , 8 , 9 , 10 } and the number of unique characters k ∼ { 1 , . . . , n } .
The batch is composed of 32 sets.
The results are summarized in Table 7.
Large set.
This large set experiment is built from the sets generated by: set cardinality n ∼ { 11 , 12 , . .
. , 30 } and the number of unique characters k ∼ { 5 , . .
. , n } .
The batch is composed ofsets.
The results are shown in Table 8. [Table] [Table]
Network architectures and learning options.
The feature vectors ϕ ( x ) of images (set elements) are formed by applying four Conv(64,3,2,BN,ReLU) layers.
For OTKE and our DIEM the output of the set embedding emb ( S ) is fed into the fully-connected layers FC(64,ReLU) – FC(1,softplus) to return the (scalar) Poisson parameter.
For SetTransformer ϕ ( x ) is fed into SAB(64,4) – SAB(64,4) –
PMA(1,8,8) – FC(1,softplus) layers where SAB( d , h ) means the set attention block with d/h units/- heads, and PMA( k , d , h ) is the Pooling Multihead Attention layer with k vectors and d/h units/heads.
When the inducing-point attention is used, SAB(64,4) blocks are replaced by ISAB( m ,64,4) that has m inducing points ( m = 32 or 64) with 64/4 units/heads.
For DeepSet, ϕ ( x ) is fed into FC(64,ReLU) – FC(64,-) – Mean-Pool – FC(64,ReLU) – FC(1,softplus) layers.
For all models, we use the Adam optimizer (Kingma & Ba, 2015) with learning rate 10 − 4 and batch size = 32 sets, until 200K iterations.
C. A MORTIZED C LUSTERING
Synthetic 2D data clustering.
The results are summarized in Table 9.
Synthetic small set data clustering.
We test the models on sets with smaller cardinality: N ∼ { 50 , . . .
, 100 } instead of the previous N ∼ { 100 , . . .
The results are shown in Table 10.
CIFAR-100 clustering.
Results are shown in Table 11.
Network architectures and learning options: For OTKE and our DIEM, we first apply a linear layer FC( dh ,-) to the din -dimensional input vectors to have features ϕ ( x ) , where dh = 128 , din =for the synthetic dataset and the small synthetic dataset, and dh = 256 , din = 512 for the CIFAR100.
The output of the set embedding emb ( S ) is then fed into FC( dh ,ReLU) –
FC( dout ,-) where dout = 4 + 4 · din + 4 for α , m , and C , respectively, in the order-4 spherical-covariance Gaussian mixture parameter output Θ = { α, m, C } .
We also apply softmax and softplus layers for α and C , respectively.
For SetTransformer, the din -dimensional inputs are fed into SAB( dh ,4) – SAB( dh ,4) – PMA(4, dh ,4) – SAB( dh ,4) – SAB( dh ,4) – FC( dout ,-) layers, whereas for the inducing-point [Table] are shown (the higher the better). In DIEM, k = 2(00) (or
[Table]
[Table] 6 ± 0 . DIEM ( p = 50 ,H = 1 ,k = 2(00) ,τ = 10 − 6 ,L = 512 ) (Ours) 89 .
6 ± 0 .
4 ± 0 .
6 ± 0 .attention, the first two SAB( dh ,4) blocks are replaced by ISAB( m , dh ,4).
For DeepSet, the layers are composed of k d FC( dh ,ReLU) layers, Mean-Pool, and ( k d − 1) FC( dh ,ReLU), followed by
FC( dout ,-), where k d = 4 for synthetic and k d = 6 for CIFAR-100.
For all models, we use the
Adam optimizer with learning rate 10 − 3 ( 10 − 4 for CIFAR-100) and batch size = 10 sets, until 50K iterations.
C. P ROTEIN F OLD C LASSIFICATION T ASK ON SCOP 1.
Unsupervised learning.
The best performing DIEM model is found with ( p = 100 , H = 1 , k = 1 , τ = 10 − 3 ) following the OT E-step.
See also Fig.
12 for the results of different prior strength τ values.
Supervised learning.
The results for several different hyperparameter choices are shown in Table 12.
Network architectures and learning options.
We faithfully follow (Mialon et al., 2021) for the overall network architectures and learning options.
For instance, the Adam optimizer is used with batch size 128 for 100 epochs, and the learning rate is initially 0.01 and halved if there is no decrease in validation loss for 5 consecutive epochs.
C. E MPIRICAL S TUDY ON H YPERPARAMETER I MPACTS
We empirically analyze the impact of various hyperparameters used in our DIEM model.
They are summarized as follows.
OMNIGLOT (small).
We first study the impact of prior strength τ and multi-head pooling strategy.
While fixing the number of EM steps as k = 2 , we vary τ ∈ { 10 − 6 , 10 − 3 , 10 0 } and the multi-head pooling strategy ∈ { PC , SB , SB2 } .
The results are shown in Fig.
We have the best accuracy when τ is the middle value 10 − 3 compared to the extreme ones ( 10 − 6 lowest prior impact and 10 0 highest impact).
Regarding the multi-head pooling strategy, PC (parameter concatenation) performs the best for small τ , while SB2 (select best with the selected mixture ID one hot encoding) is slightly better than PC for τ = 10 0 .
Next we test the impact of the number of EM steps k and the prior impact τ .
For this we fix the multi-head pooling strategy as SB2 or PC.
As shown in Fig.
8, having multiple k = 2 EM steps consistently outperform a single EM step.
Unlike k = 2 , the single EM step models attained the best performance with the strongest prior impact τ = 10 0 .
Synthetic clustering.
We also test the impact of multi-head pooling strategy and prior strength τ on this synthetic clustering dataset for our DIEM model ( p = 4 , H = 5 , k = 3 ).
As shown in Fig.
SB2 strategy outperforms the other two strategies for all τ values.
For SB2, the impact of τ appears to be minor, while for the other two strategies, we have opposite behavior; larger τ performs better for PC, and smaller τ improves the performance for SB.
Fig.
10 visualizes the impact of the number of EM steps, where the model is with ( p = 4 , H = 5 , τ = 0 .
01 , SB2).
It clearly shows that having multiple EM steps significantly improve the performance over the single EM step.
There is only a little improvement from k = 2 to k = 3 . [Figure] [Figure] [Figure]
SB
SB [Figure] [Figure] CIFAR-100 clustering.
We vary the prior impact hyperparameter τ from 10 − 6 to 10 0 , and the results are depicted in Fig.
It indicates that choosing a moderate prior trade-off ( 10 − 2 ) yields the best accuracy.
SCOP 1.75 protein fold classification.
We test on the protein fold classification dataset the impact of the prior strength τ .
For the unsupervised learning setting, we vary τ
and the results are shown in Fig.
We have similar trend as CIFAR-100 clustering, where the moderate strength τ = 10 −performs the best.
Finally, we compare the types of E-steps performed for the supervised learning setup.
For the DIEM model ( p = 50 , H = 1 , k = 2 , τ = 10 − 6 , L = 128 ), the k = 2 EM steps are run in two different ways: the regular E-steps and the OT E-steps (i.e., without and with imposing the balanced assignment constraints, respectively).
The results are shown in Fig. 13, and the regular
E-steps are more effective than OT E-steps for this particular dataset. [Figure] [Figure]