Set Transformer: A Framework for Attention-based

Permutation-Invariant Neural Networks

Juho Lee 1 2 Yoonho Lee 3 Jungtaek Kim 4 Adam R. Kosiorek 1 5 Seungjin Choi 4 Yee Whye Teh 1

Abstract 1997; Maron & Lozano-Pérez, 1998) is an example of such

a set-input problem, where a set of instances is given as an
arXiv:1810.00825v3 [[Link]] 26 May 2019

Many machine learning tasks such as multiple input and the corresponding target is a label for the entire
instance learning, 3D shape recognition and few- set. Other problems such as 3D shape recognition (Wu et al.,
shot image classification are defined on sets of in- 2015; Shi et al., 2015; Su et al., 2015; Charles et al., 2017),
stances. Since solutions to such problems do not sequence ordering (Vinyals et al., 2016), and various set op-
depend on the order of elements of the set, mod- erations (Muandet et al., 2012; Oliva et al., 2013; Edwards &
els used to address them should be permutation Storkey, 2017; Zaheer et al., 2017) can also be viewed as the
invariant. We present an attention-based neural set-input problems. Moreover, many meta-learning (Thrun
network module, the Set Transformer, specifically & Pratt, 1998; Schmidhuber, 1987) problems which learn
designed to model interactions among elements using different, but related tasks may also be treated as set-
in the input set. The model consists of an encoder input tasks where an input set corresponds to the training
and a decoder, both of which rely on attention dataset of a single task. For example, few-shot image clas-
mechanisms. In an effort to reduce computational sification (Finn et al., 2017; Snell et al., 2017; Lee & Choi,
complexity, we introduce an attention scheme in- 2018) operates by building a classifier using a support set
spired by inducing point methods from sparse of images, which is evaluated with query images.
Gaussian process literature. It reduces computa-
tion time of self-attention from quadratic to linear A model for set-input problems should satisfy two critical
in the number of elements in the set. We show requirements. First, it should be permutation invariant —
that our model is theoretically attractive and we the output of the model should not change under any permu-
evaluate it on a range of tasks, demonstrating in- tation of the elements in the input set. Second, such a model
creased performance compared to recent methods should be able to process input sets of any size. While these
for set-structured data. requirements stem from the definition of a set, they are not
easily satisfied in neural-network-based models: classical
feed-forward neural networks violate both requirements,
1. Introduction and RNNs are sensitive to input order.
Recently, Edwards & Storkey (2017) and Zaheer et al.
Learning representations has proven to be an essential prob-
(2017) propose neural network architectures which meet
lem for deep learning and its many success stories. The
both criteria, which we call set pooling methods. In this
majority of problems tackled by deep learning are instance-
model, each element in a set is first independently fed into
based and take the form of mapping a fixed-dimensional
a feed-forward neural network that takes fixed-size inputs.
input tensor to its corresponding target value (Krizhevsky
Resulting feature-space embeddings are then aggregated
et al., 2012; Graves et al., 2013).
using a pooling operation (mean, sum, max or similar).
For some applications, we are required to process set- The final output is obtained by further non-linear processing
structured data. Multiple instance learning (Dietterich et al., of the aggregated embedding. This remarkably simple ar-
1 chitecture satisfies both aforementioned requirements, and
Department of Statistics, University of Oxford, United King-
dom 2 AITRICS, Republic of Korea 3 Kakao Corporation, Repub- more importantly, is proven to be a universal approximator
lic of Korea 4 Department of Computer Science and Engineering, for any set function (Zaheer et al., 2017). Thanks to this
POSTECH, Republic of Korea 5 Oxford Robotics Institute, Univer- property, it is possible to learn a complex mapping between
sity of Oxford, United Kingdom. Correspondence to: Juho Lee input sets and their target outputs in a black-box fashion,
<[Link]@[Link]>. much like with feed-forward or recurrent neural networks.
Proceedings of the 36 th International Conference on Machine Even though this set pooling approach is theoretically at-
Learning, Long Beach, California, PMLR 97, 2019. Copyright tractive, it remains unclear whether we can approximate
2019 by the author(s).
 Set Transformer

complex mappings well using only instance-based feature 2. Background

extractors and simple pooling operations. Since every el-
ement in a set is processed independently in a set pooling 2.1. Pooling Architecture for Sets
operation, some information regarding interactions between Problems involving a set of objects have the permutation
elements has to be necessarily discarded. This can make invariance property: the target value for a given set is the
some problems unnecessarily difficult to solve. same regardless of the order of objects in the set. A sim-
Consider the problem of amortized clustering, where we ple example of a permutation invariant model is a network
would like to learn a parametric mapping from an input that performs pooling over embeddings extracted from the
set of points to the centers of clusters of points inside the elements of a set. More formally,
set. Even for a toy dataset in 2D space, this is not an easy net({x1 , . . . , xn }) = ρ(pool({φ(x1 ), . . . , φ(xn )})). (1)
problem. The main difficulty is that the parametric mapping
must assign each point to its corresponding cluster while Zaheer et al. (2017) have proven that all permutation in-
modelling the explaining away pattern such that the resulting variant functions can be represented as (1) when pool is
clusters do not attempt to explain overlapping subsets of the sum operator and ρ, φ any continuous functions, thus
the input set. Due to this innate difficulty, clustering is justifying the use of this architecture for set-input problems.
typically solved via iterative algorithms that refine randomly
Note that we can deconstruct (1) into two parts: an encoder
initialized clusters until convergence. Even though a neural
(φ) which independently acts on each element of a set of n
network with a set poling operation can approximate such an
items, and a decoder (ρ(pool(·))) which aggregates these
amortized mapping by learning to quantize space, a crucial
encoded features and produces our desired output. Most
shortcoming is that this quantization cannot depend on the
network architectures for set-structured data follow this
contents of the set. This limits the quality of the solution
encoder-decoder structure.
and also may make optimization of such a model more
difficult; we show empirically in Section 5 that such pooling Zaheer et al. (2017) additionally observed that the model
architectures suffer from under-fitting. remains permutation invariant even if the encoder is a stack
of permutation-equivariant layers:
In this paper, we propose a novel set-input deep neural
network architecture called the Set Transformer, (cf. Trans- Definition 1. Let Sn be the set of all permutations of indices
former, (Vaswani et al., 2017)). The novelty of the Set {1, . . . , n}. A function f : X n → Y n is permutation equiv-
Transformer is in three important design choices: ariant iff for any permutation π ∈ Sn , f (πx) = πf (x).

An example of a permutation-equivariant layer is

1. We use a self-attention mechanism to process every fi (x; {x1 , . . . , xn }) = σi (λx + γpool({x1 , . . . , xn }))
element in an input set, which allows our approach to (2)
naturally encode pairwise- or higher-order interactions
between elements in the set. where pool is the pooling operation, λ, γ are learnable scalar
variables, and σ(·) is a nonlinear activation function.

2. We propose a method to reduce the O(n2 ) computation 2.2. Attention

time of full self-attention (e.g. the Transformer) to
O(nm) where m is a fixed hyperparameter, allowing Assume we have n query vectors (corresponding to a set
our method to scale to large input sets. with n elements) each with dimension dq : Q ∈ Rn×dq .
An attention function Att(Q, K, V ) is a function that
maps queries Q to outputs using nv key-value pairs K ∈
3. We use a self-attention mechanism to aggregate fea- Rnv ×dq , V ∈ Rnv ×dv .
tures, which is especially beneficial when the prob-

Att(Q, K, V ; ω) = ω QK > V. (3)
lem requires multiple outputs which depend on each
other, such as the problem of meta-clustering, where The pairwise dot product QK > ∈ Rn×nv measures how
the meaning of each cluster center heavily depends its similar each pair of query and key vectors is, with weights
location relative to the other clusters. computed with an activation function ω. The output
ω(QK > )V is a weighted sum of V where a value gets
more weight if its corresponding key has larger dot product
We apply the Set Transformer to several set-input problems
with the query.
and empirically demonstrate the importance and effective-
ness of these design choices, and show that we can achieve Multi-head attention, originally introduced in Vaswani
the state-of-the-art performances for the most of the tasks. et al. (2017), is an extension of the previous attention
 Set Transformer

(a) Our model (b) MAB (c) SAB (d) ISAB

Figure 1. Diagrams of our attention-based set operations.

scheme. Instead of computing a single attention func- neural network blocks with their own parameters, and not
tion, this method first projects Q, K, V onto h different fixed functions.
dMq , dq , dv -dimensional vectors, respectively. An atten-
M M
Given matrices X, Y ∈ Rn×d which represent two sets of
tion function (Att(·; ωj )) is applied to each of these h pro-
d-dimensional vectors, we define the Multihead Attention
jections. The output is a linear transformation of the con-
Block (MAB) with parameters ω as follows:
catenation of all attention outputs:
MAB(X, Y ) = LayerNorm(H + rFF(H)), (6)
Multihead(Q, K, V ; λ, ω) = concat(O1 , · · · , Oh )W O , where H = LayerNorm(X + Multihead(X, Y, Y ; ω)),
(4) (7)
where Oj = Att(QWjQ , KWjK , V WjV ; ωj ) (5) rFF is any row-wise feedforward layer (i.e., it pro-
cesses each instance independently and identically), and
Note that Multihead(·, ·, ·; λ) has learnable parameters LayerNorm is layer normalization (Ba et al., 2016). The
M
λ = {WjQ , WjK , WjV }hj=1 , where WjQ , WjK ∈ Rdq ×dq , MAB is an adaptation of the encoder block of the Trans-
former (Vaswani et al., 2017) without positional encoding
M M
WjV ∈ Rdv ×dv , W O ∈ Rhdv ×d . A typical choice for the
dimension hyperparameters is dM M and dropout. Using the MAB, we define the Set Attention
q = dq /h, dv = dv /h,
d = dq . For brevity, we set dq = dv = d, dM M Block (SAB) as
q = dv = d/h
throughout the rest of the paper. Unless otherwise
√ specified, SAB(X) := MAB(X, X). (8)
we use a scaled softmax ωj (·) = softmax(·/ d), which
our experiments were worked robustly in most settings. In other words, an SAB takes a set and performs self-
attention between the elements in the set, resulting in a set
of equal size. Since the output of SAB contains information
3. Set Transformer about pairwise interactions among the elements in the input
In this section, we motivate and describe the Set Trans- set X, we can stack multiple SABs to encode higher order
former: an attention-based neural network that is designed interactions. Note that while the SAB (8) involves a multi-
to process sets of data. Similar to other architectures, a Set head attention operation (7), where Q = K = V = X, it
Transformer consists of an encoder followed by a decoder could reduce to applying a residual block on X. In practice,
(cf. Section 2.1), but a distinguishing feature is that each it learns more complicated functions due to linear projec-
layer in the encoder and decoder attends to their inputs to tions of X inside attention heads, (3) and (5).
produce activations. Additionally, instead of a fixed pooling A potential problem with using SABs for set-structured
operation such as mean, our aggregating function pool(·) data is the quadratic time complexity O(n2 ), which may be
is parameterized and can thus adapt to the problem at hand. too expensive for large sets (n  1). We thus introduce
the Induced Set Attention Block (ISAB), which bypasses
3.1. Permutation Equivariant (Induced) Set Attention this problem. Along with the set X ∈ Rn×d , additionally
Blocks define m d-dimensional vectors I ∈ Rm×d , which we call
inducing points. Inducing points I are part of the ISAB
We begin by defining our attention-based set operations,
itself, and they are trainable parameters which we train
which we call SAB and ISAB. While existing pooling meth-
along with other parameters of the network. An ISAB with
ods for sets obtain instance features independently of other
m inducing points I is defined as:
instances, we use self-attention to concurrently encode the
whole set. This gives the Set Transformer the ability to com- ISABm (X) = MAB(X, H) ∈ Rn×d , (9)
pute pairwise as well as higher-order interactions among m×d
where H = MAB(I, X) ∈ R . (10)
instances during the encoding process. For this purpose,
we adapt the multihead attention mechanism used in Trans- The ISAB first transforms I into H by attending to the
former. We emphasize that all blocks introduced here are input set. The set of transformed inducing points H, which
 Set Transformer

contains information about the input set X, is again attended 3.3. Overall Architecture
to by the input set X to finally produce a set of n elements.
Using the ingredients explained above, we describe how we
This is analogous to low-rank projection or autoencoder
would construct a set transformer consists of an encoder and
models, where inputs (X) are first projected onto a low-
a decoder. The encoder Encoder : X 7→ Z ∈ Rn×d is a
dimensional object (H) and then reconstructed to produce
stack of SABs or ISABs, for example:
outputs. The difference is that the goal of these methods is
reconstruction whereas ISAB aims to obtain good features Encoder(X) = SAB(SAB(X)) (13)
for the final task. We expect the learned inducing points
to encode some global structure which helps explain the Encoder(X) = ISABm (ISABm (X)). (14)
inputs X. For example, in the amortized clustering problem
We point out again that the time complexity for ` stacks
on a 2D plane, the inducing points could be appropriately
of SABs and ISABs are O(`n2 ) and O(`nm), respectively.
distributed points on the 2D plane so that the encoder can
This can result in much lower processing times when using
compare elements in the query dataset indirectly through
ISAB (as compared to SAB), while still maintaining high
their proximity to these grid points.
representational power. After the encoder transforms data
Note that in (9) and (10), attention was computed between X ∈ Rn×dx into features Z ∈ Rn×d , the decoder aggre-
a set of size m and a set of size n. Therefore, the time gates them into a single or a set of vectors which is fed into
complexity of ISABm (X; λ) is O(nm) where m is a (typ- a feed-forward network to get final outputs. Note that PMA
ically small) hyperparameter — an improvement over the with k > 1 seed vectors should be followed by SABs to
quadratic complexity of the SAB. We also emphasize that model the correlation between k outputs.
both of our set operations (SAB and ISAB) are permutation
equivariant (definition in Section 2.1): Decoder(Z; λ) = rFF(SAB(PMAk (Z))) ∈ Rk×d (15)
Property 1. Both SAB(X) and ISABm (X) are permuta- where PMAk (Z) = MAB(S, rFF(Z)) ∈ Rk×d , (16)
tion equivariant.
3.4. Analysis
3.2. Pooling by Multihead Attention
Since the blocks used to construct the encoder (i.e., SAB,
A common aggregation scheme in permutation invariant ISAB) are permutation equivariant, the mapping of the en-
networks is a dimension-wise average or maximum of the coder X → Z is permutation equivariant as well. Combined
feature vectors (cf. Section 1). We instead propose to aggre- with the fact that the PMA in the decoder is a permutation
gate features by applying multihead attention on a learnable invariant transformation, we have the following:
set of k seed vectors S ∈ Rk×d . Let Z ∈ Rn×d be the set of
features constructed from an encoder. Pooling by Multihead Proposition 1. The Set Transformer is permutation invari-
Attention (PMA) with k seed vectors is defined as ant.

PMAk (Z) = MAB(S, rFF(Z)). (11) Being able to approximate any function is a desirable prop-
erty, especially for black-box models such as deep neural
Note that the output of PMAk is a set of k items. We use networks. Building on previous results about the universal
one seed vector (k = 1) in most cases, but for problems such approximation of permutation invariant functions, we prove
as amortized clustering which requires k correlated outputs, the universality of Set Transformers:
the natural thing to do is to use k seed vectors. To further
Proposition 2. The Set Transformer is a universal approxi-
model the interactions among the k outputs, we apply an
mator of permutation invariant functions.
SAB afterwards:
H = SAB(PMAk (Z)). (12) Proof. See supplementary material.
We later empirically show that such self-attention after pool-
ing helps in modeling explaining-away (e.g., among clusters 4. Related Works
in an amortized clustering problem).
Pooling architectures for permutation invariant map-
Intuitively, feature aggregation using attention should be pings Pooling architectures for sets have been used in
beneficial because the influence of each instance on the various problems such as 3D shape recognition (Shi et al.,
target is not necessarily equal. For example, consider a 2015; Su et al., 2015), discovering causality (Lopez-Paz
problem where the target value is the maximum value of a et al., 2017), learning the statistics of a set (Edwards &
set of real numbers. Since the target can be recovered using Storkey, 2017), few-shot image classification (Snell et al.,
only a single instance (the largest), finding and attending to 2017), and conditional regression and classification (Gar-
that instance during aggregation will be advantageous. nelo et al., 2018). Zaheer et al. (2017) discuss the structure
 Set Transformer

in general and provides a partial proof of the universality

Table 1. Mean absolute errors on the max regression task.
of the pooling architecture, and Wagstaff et al. (2019) fur-
ther discuss the limitation of pooling architectures. Bloem- Architecture MAE
Reddy & Teh (2019) provides a link between probabilistic rFF + Pooling (mean) 2.133 ± 0.190
exchangeability and pooling architectures. rFF + Pooling (sum) 1.902 ± 0.137
rFF + Pooling (max) 0.1355 ± 0.0074
Attention-based approaches for sets Several recent
works have highlighted the competency of attention mecha- SAB + PMA (ours) 0.2085 ± 0.0127
nisms in modeling sets. Vinyals et al. (2016) pool elements
in a set by a weighted average with weights computed using
an attention mechanism. Yang et al. (2018) propose AttSets
for multi-view 3D reconstruction, where dot-product atten- ments five times and report performance metrics evaluated
tion is applied to compute the weights used to pool the on corresponding test datasets. Along with baselines, we
encoded features via weighted sums. Similarly, Ilse et al. compared various architectures arising from the combina-
(2018) use attention-based weighted sum-pooling for multi- tion of the choices of having attention in encoders and de-
ple instance learning. Compared to these approaches, ours coders. Unless specified otherwise, “simple pooling” means
use multihead attention in aggregation, and more impor- average pooling.
tantly, we propose to apply self-attention after pooling to
model correlation among multiple outputs. PMA with k = 1 • rFF + Pooling (Zaheer et al., 2017): rFF layers in
seed vector and single-head attention roughly corresponds encoder and simple pooling + rFF layers in decoder.
to these previous approaches. Although not permutation • rFFp-mean/rFFp-max + Pooling (Zaheer et al., 2017):
invariant, Mishra et al. (2018) has attention as one of its rFF layers with permutation equivariant variants in
core components to meta-learn to solve various tasks using encoder (Zaheer et al., 2017, (4)) and simple pooling +
sequences of inputs. Kim et al. (2019) proposed attention- rFF layers in decoder.
based conditional regression, where self-attention is applied • rFF + Dotprod (Yang et al., 2018; Ilse et al., 2018):
to the query sets. rFF layers in encoder and dot product attention based
Modeling interactions between elements in sets An im- weighted sum pooling + rFF layers in decoder.
portant reason to use the Transformer is to explicitly model • SAB (ISAB) + Pooling (ours): Stack of SABs (ISABs)
higher-order interactions among the elements in a set. San- in encoder and simple pooling + rFF layers in decoder.
toro et al. (2017) propose the relational network, a simple • rFF + PMA (ours): rFF layers in encoder and PMA
architecture that sum-pools all pairwise interactions of el- (followed by stack of SABs) in decoder.
ements in a given set, but not higher-order interactions.
• SAB (ISAB) + PMA (ours): Stack of SABs (ISABs)
Similarly to our work, Ma et al. (2018) use the Transformer
in encoder and PMA (followed by stack of SABs) in
to model interactions between the objects in a video. They
decoder.
use mean-pooling to obtain aggregated features which they
fed into an LSTM.
5.1. Toy Problem: Maximum Value Regression
Inducing point methods The idea of letting trainable vec-
tors I directly interact with data points is loosely based on To demonstrate the advantage of attention-based set aggre-
the inducing point methods used in sparse Gaussian pro- gation over simple pooling operations, we consider a toy
cesses (Snelson & Ghahramani, 2005) and the Nyström problem: regression to the maximum value of a given set.
method for matrix decomposition (Fowlkes et al., 2004). m Given a set of real numbers {x1 , . . . , xn }, the goal is to
trainable inducing points can also be seen as m independent return max(x1 , · · · , xn ). Given prediction p, we use the
memory cells accessed with an attention mechanism. The mean absolute error |p − max(x1 , · · · , xn )| as the loss func-
differential neural dictionary (Pritzel et al., 2017) stores pre- tion. We constructed simple pooling architectures with three
vious experience as key-value pairs and uses this to process different pooling operations: max, mean, and sum. We
queries. One can view the ISAB is the inversion of this idea, report loss values after training in Table 1. Mean- and sum-
where queries I are stored and the input features are used as pooling architectures result in a high mean absolute error
key-value pairs. (MAE). The model with max-pooling can predict the output
perfectly by learning its encoder to be an identity function,
and thus achieves the highest performance. Notably, the
5. Experiments Set Transformer achieves performance comparable to the
To evaluate the Set Transformer, we apply it to a suite of max-pooling model, which underlines the importance of
tasks involving sets of data points. We repeat all experi- additional flexibility granted by attention mechanisms — it
can learn to find and attend to the maximum element.
 Set Transformer

60
0.
55
Accuracy
0.
ISAB(n)+PMA
Figure 2. Counting unique characters: this is a randomly sampled

50
SAB+PMA

0.
set of 20 images from the Omniglot dataset. There are 14 different SAB + Pooling
rFF + PMA
characters inside this set.

45
rFF + Pooling

0.
1 2 3 4 5 6 7 8 9 10 11
Number of Inducing Points (n)
Table 2. Accuracy on the unique character counting task.

Architecture Accuracy Figure 3. Accuracy of ISABn + PMA on the unique character

counting task. x-axis is n and y-axis is accuracy.
rFF + Pooling 0.4382 ± 0.0072
rFFp-mean + Pooling 0.4617 ± 0.0076
rFFp-max + Pooling 0.4359 ± 0.0077 5.3. Amortized Clustering with Mixture of Gaussians
rFF + Dotprod 0.4471 ± 0.0076
rFF + PMA (ours) 0.4572 ± 0.0076 We applied the set-input networks to the task of maxi-
SAB + Pooling (ours) 0.5659 ± 0.0077 mum likelihood of mixture of Gaussians (MoGs). The
SAB + PMA (ours) 0.6037 ± 0.0075 log-likelihood of a dataset X = {x1 , . . . , xn } generated
from an MoG with k components is
n
X k
X
log p(X; θ) = log πj N (xi ; µj , diag(σj2 )). (17)
5.2. Counting Unique Characters i=1 j=1

In order to test the ability of modelling interactions between The goal is to learn the optimal parameters θ∗ (X) =
objects in a set, we introduce a new task of counting unique arg maxθ log p(X; θ). The typical approach to this prob-
elements in an input set. We use the Omniglot (Lake et al., lem is to run an iterative algorithm such as Expectation-
Maximisation (EM) until convergence. Instead, we aim
2015) dataset, which consists of 1,623 different handwritten to learn a generic meta-algorithm that directly maps the
characters from various alphabets, where each character is input set X to θ∗ (X). One can also view this as amor-
represented by 20 different images. tized maximum likelihood learning. Specifically, given a
dataset X, we train a neural network to output parameters
We split all characters (and corresponding images) into train, f (X; λ) = {π(X), {µj (X), σj (X)}kj=1 } which maximize
validation, and test sets and only train using images from the
 
train character classes. We generate input sets by sampling |X|
X k
X
between 6 and 10 images and we train the model to predict EX  log πj (X)N (xi ; µj (X), diag(σj2 (X))) . (18)
the number of different characters inside the set. We used i=1 j=1

a Poisson regression model to predict this number, with

We structured f (·; λ) as a set-input neural network and
the rate λ given as the output of a neural network. We
learned its parameters λ using stochastic gradient ascent,
maximized the log likelihood of this model using stochastic
where we approximate gradients using minibatches of
gradient ascent.
datasets.
We evaluated model performance using sets of images sam-
We tested Set Transformers along with other set-input net-
pled from the test set of characters. Table 2 reports accuracy,
works on two datasets. We used four seed vectors for the
measured as the frequency at which the mode of the Poisson
PMA (S ∈ R4×d ) so that each seed vector generates the
distribution chosen by the network is equal to the number
parameters of a cluster.
of characters inside the input set.
Synthetic 2D mixtures of Gaussians: Each dataset con-
We additionally performed experiments to see how the num-
tains n ∈ [100, 500] points on a 2D plane, each sampled
ber of incuding points affects performance. We trained
from one of four Gaussians.
ISABn + PMA on this task while varying the number of in-
ducing points (n). Accuracies are shown in Figure 3, where CIFAR-100: Each dataset contains n ∈ [100, 500] images
other architectures are shown as horizontal lines for compar- sampled from four random classes in the CIFAR-100 dataset.
ison. Note first that even the accuracy of ISAB1 + PMA Each image is represented by a 512-dim vector obtained
surpasses that of both rFF + Pooling and rFF + PMA, and from a pretrained VGG network (Simonyan & Zisserman,
that performance tends to increase as we increase n. 2014).
 Set Transformer

Table 3. Meta clustering results. The number inside parenthesis indicates the number of inducing points used in ISABs of encoders. We
show average likelihood per data for the synthetic dataset and the adjusted rand index (ARI) for the CIFAR-100 experiment. LL1/data,
ARI1 are the evaluation metrics after a single EM update step. The oracle for the synthetic dataset is the log likelihood of the actual
parameters used to generate the set, and the CIFAR oracle was computed by running EM until convergence.

Synthetic CIFAR-100
Architecture LL0/data LL1/data ARI0 ARI1
Oracle -1.4726 0.9150
rFF + Pooling -2.0006 ± 0.0123 -1.6186 ± 0.0042 0.5593 ± 0.0149 0.5693 ± 0.0171
rFFp-mean + Pooling -1.7606 ± 0.0213 -1.5191 ± 0.0026 0.5673 ± 0.0053 0.5798 ± 0.0058
rFFp-max + Pooling -1.7692 ± 0.0130 -1.5103 ± 0.0035 0.5369 ± 0.0154 0.5536 ± 0.0186
rFF + Dotprod -1.8549 ± 0.0128 -1.5621 ± 0.0046 0.5666 ± 0.0221 0.5763 ± 0.0212
SAB + Pooling (ours) -1.6772 ± 0.0066 -1.5070 ± 0.0115 0.5831 ± 0.0341 0.5943 ± 0.0337
ISAB (16) + Pooling (ours) -1.6955 ± 0.0730 -1.4742 ± 0.0158 0.5672 ± 0.0124 0.5805 ± 0.0122
rFF + PMA (ours) -1.6680 ± 0.0040 -1.5409 ± 0.0037 0.7612 ± 0.0237 0.7670 ± 0.0231
SAB + PMA (ours) -1.5145 ± 0.0046 -1.4619 ± 0.0048 0.9015 ± 0.0097 0.9024 ± 0.0097
ISAB (16) + PMA (ours) -1.5009 ± 0.0068 -1.4530 ± 0.0037 0.9210 ± 0.0055 0.9223 ± 0.0056

Figure 4. Clustering results for 10 test datasets, along with centers and covariance matrices. rFF+Pooling (top-left), SAB+Pooling
(top-right), rFF+PMA (bottom-left), Set Transformer (bottom-right). Best viewed magnified in color.

We report the performance of the oracle along with the set- tive processes, network architectures, and training schemes
input neural networks in Table 3. We additionally report along with additional experiments with various numbers of
scores of all models after a single EM update. Overall, inducing points in the supplementary material.
the Set Transformer found accurate parameters and even
outperformed the oracles after a single EM update. This 5.4. Set Anomaly Detection
may be due to the relatively small size of the input sets;
some clusters have fewer than 10 points. In this regime, We evaluate our methods on the task of meta-anomaly de-
sample statistics can differ substantially from population tection within a set using the CelebA dataset. The dataset
statistics, which limits the performance of the oracle while consists of 202,599 images with the total of 40 attributes.
the Set Transformer can adapt accordingly. Notably, the We randomly sample 1,000 sets of images. For every set,
Set Transformer with only 16 inducing points showed the we select two attributes at random and construct the set
best performance, even outperforming the full Set Trans- by selecting seven images containing both attributes and
former. We believe this is due to the knowledge transfer one image with neither. The goal of this task is to find the
and regularization via inducing points, helping the network image that does not belong to the set. We give a detailed
to learn global structures. Our results also imply that the description of the experimental setup in the supplementary
improvement from using the PMA is more significant than material. We report the area under receiver operating char-
that of the SAB, supporting our claim of the importance acteristic curve (AUROC) and area under precision-recall
of attention-based decoders. We provide detailed genera- curve (AUPR) in Table 5. Set Transformers outperformed
all other methods by a significant margin.
 Set Transformer

Table 4. Test accuracy for the point cloud classification task using 100, 1000, 5000 points.

Architecture 100 pts 1000 pts 5000 pts

rFF + Pooling (Zaheer et al., 2017) - 0.83 ± 0.01 -
rFFp-max + Pooling (Zaheer et al., 2017) 0.82 ± 0.02 0.87 ± 0.01 0.90 ± 0.003
rFF + Pooling 0.7951 ± 0.0166 0.8551 ± 0.0142 0.8933 ± 0.0156
rFF + PMA (ours) 0.8076 ± 0.0160 0.8534 ± 0.0152 0.8628 ± 0.0136
ISAB (16) + Pooling (ours) 0.8273 ± 0.0159 0.8915 ± 0.0144 0.9040 ± 0.0173
ISAB (16) + PMA (ours) 0.8454 ± 0.0144 0.8662 ± 0.0149 0.8779 ± 0.0122

they augmented data (scaling, rotation) and used a differ-

ent optimizer (Adamax) and learning rate schedule. Set
Transformers were superior when given small sets, but were
outperformed by ISAB (16) + Pooling on larger sets. First
note that classification is harder when given fewer points.
We think Set Transformers were outperformed in the prob-
lems with large sets because such sets already had sufficient
information for classification, diminishing the need to model
complex interactions among points. We point out that PMA
Figure 5. Sampled datasets. Each row is a dataset, consisting of outperformed simple pooling in all other experiments.
7 normal images and 1 anomaly (red box). In each subsampled
dataset, a normal image has two attributes (rightmost column) 6. Conclusion
which anomalies do not.
In this paper, we introduced the Set Transformer, an
attention-based set-input neural network architecture. Our
Table 5. Meta set anomaly results. Each architecture is evaluated
proposed method uses attention mechanisms for both en-
using average of test AUROC and test AUPR.
coding and aggregating features, and we have empirically
Architecture Test AUROC Test AUPR validated that both of them are necessary for modelling
Random guess 0.5 0.125 complicated interactions among elements of a set. We also
rFF + Pooling 0.5643 ± 0.0139 0.4126 ± 0.0108 proposed an inducing point method for self-attention, which
rFFp-mean + Pooling 0.5687 ± 0.0061 0.4125 ± 0.0127 makes our approach scalable to large sets. We also showed
rFFp-max + Pooling 0.5717 ± 0.0117 0.4135 ± 0.0162 useful theoretical properties of our model, including the fact
rFF + Dotprod 0.5671 ± 0.0139 0.4155 ± 0.0115
that it is a universal approximator for permutation invariant
SAB + Pooling (ours) 0.5757 ± 0.0143 0.4189 ± 0.0167 functions. An interesting future work would be to apply
rFF + PMA (ours) 0.5756 ± 0.0130 0.4227 ± 0.0127 Set Transformers to meta-learning problems. In particular,
SAB + PMA (ours) 0.5941 ± 0.0170 0.4386 ± 0.0089
using Set Transformers to meta-learn posterior inference in
Bayesian models seems like a promising line of research.
Another exciting extension of our work would be to model
5.5. Point Cloud Classification the uncertainty in set functions by injecting noise variables
We evaluated Set Transformers on a classification task using into Set Transformers in a principled way.
the ModelNet40 (Chang et al., 2015) dataset1 , which con-
tains three-dimensional objects in 40 different categories.
Each object is represented as a point cloud, which we treat
as a set of n vectors in R3 . We performed experiments with Acknowledgments JL and YWT’s research leading to
input sets of size n ∈ {100, 1000, 5000}. Because of the these results has received funding from the European Re-
large set sizes, MABs are prohibitively time-consuming due search Council under the European Union’s Seventh Frame-
to their O(n2 ) time complexity. work Programme (FP7/2007-2013) ERC grant agreement no.
617071. JL has also received funding from EPSRC under
Table 4 shows classification accuracies. We point out that
grant EP/P026753/1. JL acknowledges support from IITP
Zaheer et al. (2017) used significantly more engineering
grant funded by the Korea government(MSIT) (No.2017-
for the 5000 point experiment. For this experiment only,
0-01779, XAI) and Samsung Research Funding & Incuba-
1
The point-cloud dataset used in this experiment was obtained tion Center of Samsung Electronics under Project Number
directly from the authors of Zaheer et al. (2017). SRFC-IT1702-15.
 Set Transformer

References Krizhevsky, A., Sutskever, I., and Hinton, G. E. ImageNet

classification with deep convolutional neural networks.
Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization.
In Advances in Neural Information Processing Systems
arXiv e-prints, arXiv:1607.06450, 2016.
(NeurIPS), 2012.
Bloem-Reddy, B. and Teh, Y.-W. Probabilistic sym-
Lake, B. M., Salakhutdinov, R., and Tenenbaum, J. B.
metry and invariant neural networks. arXiv e-prints,
Human-level concept learning through probabilistic pro-
arXiv:1901.06082, 2019.
gram induction. Science, 350(6266):1332–1338, 2015.
Chang, A. X., Funkhouser, T., Guibas, L., Hanrahan, P.,
Huang, Q., Li, Z., Savarese, S., Savva, M., Song, S., Lee, Y. and Choi, S. Gradient-based meta-learning with
Su, H., Xiao, J., Yi, L., and Yu, F. ShapeNet: An learned layerwise metric and subspace. In Proceedings
information-rich 3D model repository. arXiv e-prints, of the International Conference on Machine Learning
arXiv:1512.03012, 2015. (ICML), 2018.

Charles, R. Q., Su, H., Kaichun, M., and Guibas, L. J. Point- Lopez-Paz, D., Nishihara, R., Chintala, S., Schölkopf, B.,
Net: Deep learning on point sets for 3D classification and and Bottou, L. Discovering causal signals in images. In
segmentation. In Proceedings of the IEEE Conference on Proceedings of the IEEE Conference on Computer Vision
Computer Vision and Pattern Recognition (CVPR), 2017. and Pattern Recognition (CVPR), 2017.

Dietterich, T. G., Lathrop Richard, H., and Lozano-Pérez, T. Ma, C.-Y., Kadav, A., Melvin, I., Kira, Z., AlRegib, G., and
Solving the multiple instance problem with axis-parallel Peter Graf, H. Attend and interact: higher-order object
rectangles. Artificial intelligence, 89(1-2):31–71, 1997. interactions for video understanding. In Proceedings of
the IEEE Conference on Computer Vision and Pattern
Edwards, H. and Storkey, A. Towards a neural statistician. In Recognition (CVPR), 2018.
Proceedings of the International Conference on Learning
Representations (ICLR), 2017. Maron, O. and Lozano-Pérez, T. A framework for multiple-
instance learning. In Advances in Neural Information
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- Processing Systems (NeurIPS), 1998.
learning for fast adaptation of deep networks. In Pro-
ceedings of the International Conference on Machine Mishra, N., Rohaninejad, M., Chen, X., and Abbeel, P.
Learning (ICML), 2017. A simple neural attentive meta-learner. In Proceedings
of the International Conference on Machine Learning
Fowlkes, C., Belongie, S., Chung, F., and Malik, J. Spectral (ICML), 2018.
grouping using the Nyström method. IEEE Transactions
on Pattern Analysis and Machine Intelligence, 25(2):215– Muandet, K., Fukumizu, K., Dinuzzo, F., and Schölkopf,
225, 2004. B. Learning from distributions via support measure ma-
chines. In Advances in Neural Information Processing
Garnelo, M., Rosenbaum, D., Maddison, C. J., Ramalho, Systems (NeurIPS), 2012.
T., Saxton, D., Shanahan, M., Teh, Y. W., Rezende, D. J.,
and Eslami, S. M. A. Conditional neural processes. In Oliva, J., Póczos, B., and Schneider, J. Distribution to dis-
Proceedings of the International Conference on Machine tribution regression. In Proceedings of the International
Learning (ICML), 2018. Conference on Machine Learning (ICML), 2013.

Graves, A., Mohamed, A.-r., and Hinton, G. E. Speech Pritzel, A., Uria, B., Srinivasan, S., Puigdomenech, A.,
recognition with deep recurrent neural networks. In Pro- Vinyals, O., Hassabis, D., Wierstra, D., and Blundell,
ceedings of the IEEE International Conference on Acous- C. Neural episodic control. In Proceedings of the In-
tics, Speech, and Signal Processing (ICASSP), 2013. ternational Conference on Machine Learning (ICML),
2017.
Ilse, M., Tomczak, J. M., and Welling, M. Attention-based
deep multiple instance learning. In Proceedings of the Santoro, A., Raposo, D., Barret, D. G. T., Malinowski, M.,
International Conference on Machine Learning (ICML), Pascanu, R., and Battaglia, P. A simple neural network
2018. module for relational reasoning. In Advances in Neural
Information Processing Systems (NeurIPS), 2017.
Kim, H., Mnih, A., Schwarz, J., Garnelo, M., Eslami, A.,
Rosenbaum, D., Vinyals, O., and Teh, Y. W. Attentive Schmidhuber, J. Evolutionary Principles in Self-Referential
neural processes. In Proceedings of International Confer- Learning. PhD thesis, Technical University of Munich,
ence on Learning Representations, 2019. 1987.
 Set Transformer

Shi, B., Bai, S., Zhou, Z., and Bai, X. DeepPano:

deep panoramic representation for 3-D shape recogni-
tion. IEEE Signal Processing Letters, 22(12):2339–2343,
2015.
Simonyan, K. and Zisserman, A. Very deep convolutional
networks for large-scale image recognition. arXiv e-
prints, arXiv:1409.1556, 2014.

Snell, J., Swersky, K., and Zemel, R. Prototypical networks

for few-shot learning. In Advances in Neural Information
Processing Systems (NeurIPS), 2017.
Snelson, E. and Ghahramani, Z. Sparse Gaussian processes
using pseudo-inputs. In Advances in Neural Information
Processing Systems (NeurIPS), 2005.
Su, H., Maji, S., Kalogerakis, E., and Learned-Miller, E.
Multi-view convolutional neural networks for 3D shape
recognition. In Proceedings of the IEEE International
Conference on Computer Vision (ICCV), 2015.

Thrun, S. and Pratt, L. Learning to Learn. Kluwer Academic

Publishers, 1998.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones,
L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Atten-
tion is all you need. In Advances in Neural Information
Processing Systems (NeurIPS), 2017.
Vinyals, O., Bengio, S., and Kudlur, M. Order matters:
sequence to sequence for sets. In Proceedings of the
International Conference on Learning Representations
(ICLR), 2016.
Wagstaff, E., Fuchs, F. B., Engelcke, M., Posner, I., and
Osborne, M. On the limitations of representing functions
on sets. arXiv:1901.09006, 2019.
Wu, Z., Song, S., Khosla, A., Yu, F., Zhang, L., Tang, X.,
and Xiao, J. 3D ShapeNets: a deep representation for vol-
umetric shapes. In Proceedings of the IEEE Conference
on Computer Vision and Pattern Recognition (CVPR),
2015.
Yang, B., Wang, S., Markham, A., and Trigoni, N. Atten-
tional aggregation of deep feature sets for multi-view 3D
reconstruction. arXiv e-prints, arXiv:1808.00758, 2018.
Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B.,
Salakhutdinov, R. R., and Smola, A. J. Deep sets.
In Advances in Neural Information Processing Systems
(NeurIPS), 2017.
 Supplementary Material for Set Transformer

Juho Lee 1 2 Yoonho Lee 3 Jungtaek Kim 4 Adam R. Kosiorek 1 5 Seungjin Choi 4 Yee Whye Teh 1

1. Proofs
arXiv:1810.00825v3 [[Link]] 26 May 2019

Pn
Lemma 1. The mean operator mean({x1 , . . . , xn }) = 1
n i=1 xi is a special case of dot-product attention with softmax.

Proof. Let s = 0 ∈ Rd and X ∈ Rn×d .

  n
sX > 1X
Att(s, X, X; softmax) = softmax √ X= xi
d n i=1

Lemma 2. The decoder of a Set Transformer, given enough nodes, can express any element-wise function of the form
Pn p
p
1
1
n i=1 zi .

Proof. We first note that we can view the decoder as the composition of functions

Decoder(Z) = rFF(H) (1)

where H = rFF(MAB(Z, rFF(Z))) (2)

We focus on H in (2). Since feed-forward networks are universal function approximators at the limit of infinite nodes, let
1
the feed-forward layers in front and back of the MAB encode the element-wise functions z → z p and z → z p , respectively.
We let h = d, so the number of heads is the same as the dimensionality of the inputs, and each head is one-dimensional. Let
the projection matrices in multi-head attention (WjQ , WjK , WjV ) represent projections onto the jth dimension and the output
matrix (W O ) the identity matrix. Since the mean operator is a special case of dot-product attention, by simple composition,
we see that an MAB can express any dimension-wise function of the form

n
! p1
1X p
Mp (z1 , · · · , zn ) = z . (3)
n i=1 i

Pn
Lemma 3. A PMA, given enough nodes, can express sum pooling ( i=1 zi ).

Proof. We prove this by construction.

Set the seed s to a zero vector and let ω(·) = 1 + f (·), where f is any activation function such that f (0) = 0. The identiy,
sigmoid, or relu functions are suitable choices for f . The output of the multihead attention is then simply a sum of the
values, which is Z in this case.

We additionally have the following universality theorem for pooling architectures:

Theorem 1. Models of the form rFF(sum(rFF(·))) are universal function approximators in the space of permutation
invariant functions.

Proof. See Appendix A of ?.

 Supplementary Material for Set Transformer

By Lemma 3, we know that decoder(Z) can express any function of the form rFF(sum(Z)). Using this fact along with
Theorem 1, we can prove the universality of Set Transformers:
Proposition 1. The Set Transformer is a universal function approximator in the space of permutation invariant functions.

Proof. By setting the matrix W O to a zero matrix in every SAB and ISAB, we can ignore all pairwise interaction terms
in the encoder. Therefore, the encoder(X) can express any instance-wise feed-forward network (Z = rFF(X)). Directly
invoking Theorem 1 concludes this proof.

While this proof required us to ignore the pairwise interaction terms inside the SABs and ISABs to prove that Set
Transformers are universal function approximators, our experiments indicated that self-attention in the encoder was crucial
for good performance.

2. Experiment Details
In all implementations, we omit the feed-forward layer in the beginning of the decoder (rFF(Z)) because the end of the
previous block contains a feed-forward layer. All MABs (inside SAB, ISAB and PMA) use fully-connected layers with
ReLU activations for rFF layers.
In the architecture descriptions, FC(d, f ) denotes the fully-connected layer with d units and activation function f . SAB(d, h)
denotes the SAB with d units and h heads. ISABm (d, h) denotes the ISAB with d units, h heads and m inducing points.
PMAk (d, h) denotes the PMA with d units, h heads and k vectors. All MABs used in SAB and PMA uses FC layers with
ReLU activations for FF layers.

2.1. Max Regression

Given a set of real numbers {x1 , . . . , xn }, the goal of this task is to return the maximum value in the set max(x1 , · · · , xn ).
We construct training data as follows. We first sample a dataset size n uniformly from the set of integers {1, · · · , 10}. We
then sample real numbers xi independently from the interval [0, 100]. Given the network’s prediction p, we use the actual
maximum value max(x1 , · · · , xn ) to compute the mean absolute error |p − max(x1 , · · · , xn )|. We don’t explicitly consider
splits of train and test data, since we sample a new set {x1 , . . . , xn } at each time step.

Table 1. Detailed architectures used in the max regression experiments.

Encoder Decoder
FF SAB Pooling PMA
FC(64, ReLU) SAB(64, 4) mean, sum, max PMA1 (64, 4)
FC(64, ReLU) SAB(64, 4) FC(64, ReLU) FC(1, −)
FC(64, ReLU) FC(1, −)
FC(64, −)

We show the detailed architectures used for the experiments in Table 1. We trained all networks using the Adam optimizer (?)
with a constant learning rate of 10−3 and a batch size of 128 for 20,000 batches, after which loss converged for all
architectures.

2.2. Counting Unique Characters

The task generation procedure is as follows. We first sample a set size n uniformly from the set of integers {6, . . . , 10}.
We then sample the number of characters c uniformly from {1, . . . , n}. We sample c characters from the training set of
characters, and randomly sample instances of each character so that the total number of instances sums to n and each set of
characters has at least one instance in the resulting set.
We show the detailed architectures used for the experiments in Table 3. For both architectures, the resulting 1-dimensional
output is passed through a softplus activation to produce the Poisson parameter γ. The role of softplus is to ensure that γ is
always positive.
 Supplementary Material for Set Transformer

Table 2. Detailed results for the unique character counting experiment.

Architecture Accuracy
rFF + Pooling 0.4366 ± 0.0071
rFF + PMA 0.4617 ± 0.0073
rFFp-mean + Pooling 0.4617 ± 0.0076
rFFp-max + Pooling 0.4359 ± 0.0077
rFF + Dotprod 0.4471 ± 0.0076
SAB + Pooling 0.5659 ± 0.0067
SAB + Dotprod 0.5888 ± 0.0072
SAB + PMA (1) 0.6037 ± 0.0072
SAB + PMA (2) 0.5806 ± 0.0075
SAB + PMA (4) 0.5945 ± 0.0072
SAB + PMA (8) 0.6001 ± 0.0078

Table 3. Detailed architectures used in the unique character counting experiments.

Encoder Decoder
rFF SAB Pooling PMA
Conv(64, 3, 2, BN, ReLU) Conv(64, 3, 2, BN, ReLU) mean PMA1 (8, 8)
Conv(64, 3, 2, BN, ReLU) Conv(64, 3, 2, BN, ReLU) FC(64, ReLU) FC(1, softplus)
Conv(64, 3, 2, BN, ReLU) Conv(64, 3, 2, BN, ReLU) FC(1, softplus)
Conv(64, 3, 2, BN, ReLU) Conv(64, 3, 2, BN, ReLU)
FC(64, ReLU) SAB(64, 4)
FC(64, ReLU) SAB(64, 4)
FC(64, ReLU)
FC(64, −)

The loss function we optimize, as previously mentioned, is the log likelihood log p(x|γ) = x log(γ) − γ − log(x!). We
chose this loss function over mean squared error or mean absolute error because it seemed like the more logical choice when
trying to make a real number match a target integer. Early experiments showed that directly optimizing for mean absolute
error had roughly the same result as optimizing γ in this way and measuring |γ − x|. We train using the Adam optimizer
with a constant learning rate of 10−4 for 200,000 batches each with batch size 32.

2.3. Solving maximum likelihood problems for mixture of Gaussians

2.3.1. D ETAILS FOR 2D SYNTHETIC MIXTURES OF G AUSSIANS EXPERIMENT
We generated the datasets according to the following generative process.

1. Generate the number of data points, n ∼ Unif(100, 500).

2. Generate k centers.

µj,d ∼ Unif(−4, 4), j = 1, . . . , 4, d = 1, 2. (4)

3. Generate cluster labels.

π ∼ Dir([1, 1]> ), zi ∼ Categorical(π), i = 1, . . . , n. (5)

4. Generate data from spherical Gaussian.

xi ∼ N (µzi , (0.3)2 I). (6)

 Supplementary Material for Set Transformer

Table 4 summarizes the architectures used for the experiments. For all architectures, at each training step, we generate 10
random datasets according to the above generative process, and updated the parameters via Adam optimizer with initial
learning rate 10−3 . We trained all the algorithms for 50k steps, and decayed the learning rate to 10−4 after 35k steps.
Table 5 summarizes the detailed results with various number of inducing points in the ISAB. Figure ?? shows the actual
clustering results based on the predicted parameters.

Table 4. Detailed architectures used in 2D synthetic experiments.

Encoder Decoder
rFF SAB ISAB Pooling PMA
FC(128, ReLU) SAB(128, 4) ISABm (128, 4) mean PMA4 (128, 4)
FC(128, ReLU) SAB(128, 4) ISABm (128, 4) FC(128, ReLU) SAB(128, 4)
FC(128, ReLU) FC(128, ReLU) FC(4 · (1 + 2 · 2), −)
FC(128, ReLU) FC(128, ReLU)
FC(4 · (1 + 2 · 2), −)

Table 5. Average log-likelihood/data (LL0/data) and average log-likelihood/data after single EM iteration (LL1/data) the clustering
experiment. The number inside parenthesis indicates the number of inducing points used in the SABs of encoder. For all PMAs, four seed
vectors were used.

Architecture LL0/data LL1/data

Oracle -1.4726
rFF + Pooling -2.0006 ± 0.0123 -1.6186 ± 0.0042
rFFp-mean + Pooling -1.7606 ± 0.0213 -1.5191 ± 0.0026
rFFp-max + Pooling -1.7692 ± 0.0130 -1.5103 ± 0.0035
rFF+Dotprod -1.8549 ± 0.0128 -1.5621 ± 0.0046
SAB + Pooling -1.6772 ± 0.0066 -1.5070 ± 0.0115
ISAB (16) + Pooling -1.6955 ± 0.0730 -1.4742 ± 0.0158
ISAB (32) + Pooling -1.6353 ± 0.0182 -1.4681 ± 0.0038
ISAB (64) + Pooling -1.6349 ± 0.0429 -1.4664 ± 0.0080
rFF + PMA -1.6680 ± 0.0040 -1.5409 ± 0.0037
SAB + PMA -1.5145 ± 0.0046 -1.4619 ± 0.0048
ISAB (16) + PMA -1.5009 ± 0.0068 -1.4530 ± 0.0037
ISAB (32) + PMA -1.4963 ± 0.0064 -1.4524 ± 0.0044
ISAB (64) + PMA -1.5042 ± 0.0158 -1.4535 ± 0.0053

2.3.2. 2D S YNTHETIC M IXTURES OF G AUSSIANS E XPERIMENT ON L ARGE - SCALE DATA

To show the scalability of the set transformer, we conducted additional experiments on large-scale 2D synthetic clustering
dataset. We generated the synthetic data as before, except that we sample the number of data points n Unif(1000, 5000)
and set k = 6. We report the clustering accuracy of a subset of comparing methods in Table 6. The set transformer with only
32 inducing points works extremely well, demonstrating its scalability and efficiency.

2.3.3. D ETAILS FOR CIFAR-100 AMORTIZED CLUTERING EXPERIMENT

We pretrained VGG net (?) with CIFAR-100, and obtained the test accuracy 68.54%. Then, we extracted feature vectors of
50k training images of CIFAR-100 from the 512-dimensional hidden layers of the VGG net (the layer just before the last
layer). Given these feature vectors, the generative process of datasets is as follows.

1. Generate the number of data points, n ∼ Unif(100, 500).

2. Uniformly sample four classes among 100 classes.
3. Uniformly sample n data points among four sampled classes.
 Supplementary Material for Set Transformer

Table 6. Average log-likelihood/data (LL0/data) and average log-likelihood/data after single EM iteration (LL1/data) the clustering
experiment on large-scale data. The number inside parenthesis indicates the number of inducing points used in the SABs of encoder. For
all PMAs, six seed vectors were used.

Architecture LL0/data LL1/data

Oracle -1.8202
rFF + Pooling -2.5195 ± 0.0105 -2.0709 ± 0.0062
rFFp-mean + Pooling -2.3126 ± 0.0154 -1.9749 ± 0.0062
rFF + PMA (6) -2.0515 ± 0.0067 -1.9424 ± 0.0047
SAB (32) + PMA (6) -1.8928 ± 0.0076 -1.8549 ± 0.0024

Table 7. Detailed architectures used in CIFAR-100 meta clustering experiments.

Encoder Decoder
rFF SAB ISAB rFF PMA
FC(256, ReLU) SAB(256, 4) ISABm (256, 4) mean PMA4 (128, 4)
FC(256, ReLU) SAB(256, 4) ISABm (256, 4) FC(256, ReLU) SAB(256, 4)
FC(256, ReLU) SAB(256, 4) ISABm (256, 4) FC(256, ReLU) SAB(256, 4)
FC(256, ReLU) FC(256, ReLU)) FC(4 · (1 + 2 · 512), −)
FC(256, ReLU) FC(256, ReLU)
FC(256, −) FC(256, ReLU)
FC(4 · (1 + 2 · 512), −)

Table 7 summarizes the architectures used for the experiments. For all architectures, at each training step, we generate 10
random datasets according to the above generative process, and updated the parameters via Adam optimizer with initial
learning rate 10−4 . We trained all the algorithms for 50k steps, and decayed the learning rate to 10−5 after 35k steps.
Table 8 summarizes the detailed results with various number of inducing points in the ISAB.

2.4. Set Anomaly Detection

Table 9 describes the architecture for meta set anomaly experiments. We trained all models via Adam optimizer with
learning rate 10−4 and exponential decay of learning rate for 1,000 iterations. 1,000 datasets subsampled from CelebA
dataset (see Figure ??) are used to train and test all the methods. We split 800 training datasets and 200 test datasets for the
subsampled datasets.

2.5. Point Cloud Classification

We used the ModelNet40 dataset for our point cloud classification experiments. This dataset consists of a three-dimensional
representation of 9,843 training and 2,468 test data which each belong to one of 40 object classes. As input to our
architectures, we produce point clouds with n = 100, 1000, 5000 points each (each point is represented by (x, y, z)
coordinates). For generalization, we randomly rotate and scale each set during training.
We show results our architectures in Table 10 and additional experiments which used n = 100, 5000 points in Table ??. We
trained using the Adam optimizer with an initial learning rate of 10−3 which we decayed by a factor of 0.3 every 20,000
steps. For the experiment with 5,000 points (Table ??), we increased the dimension of the attention blocks (ISAB16 (512, 4)
instead of ISAB16 (128, 4)) and also decayed the weights by a factor of 10−7 . We also only used one ISAB block in the
encoder because using two lead to overfitting in this setting.

3. Additional Experiments
3.1. Runtime of SAB and ISAB
We measured the runtime of SAB and ISAB on a simple benchmark (Figure 1). We used a single GPU (Tesla P40) for this
experiment. The input data was a constant (zero) tensor of n three-dimensional vectors. We report the number of seconds it
 Supplementary Material for Set Transformer

Table 8. Average clustering accuracies measured by Adjusted Rand Index (ARI) for CIFAR100 clustering experiments. The number
inside parenthesis indicates the number of inducing points used in the SABs of encoder. For all PMAs, four seed vectors were used.

Architecture ARI0 ARI1

Oracle 0.9151
rFF + Pooling 0.5593 ± 0.0149 0.5693 ± 0.0171
rFFp-mean + Pooling 0.5673 ± 0.0053 0.5798 ± 0.0058
rFFp-max + Pooling 0.5369 ± 0.0154 0.5536 ± 0.0186
rFF+Dotprod 0.5666 ± 0.0221 0.5763 ± 0.0212
SAB + Pooling 0.5831 ± 0.0341 0.5943 ± 0.0337
ISAB (16) + Pooling 0.5672 ± 0.0124 0.5805 ± 0.0122
ISAB (32) + Pooling 0.5587 ± 0.0104 0.5700 ± 0.0134
ISAB (64) + Pooling 0.5586 ± 0.0205 0.5708 ± 0.0183
rFF + PMA 0.7612 ± 0.0237 0.7670 ± 0.0231
SAB + PMA 0.9015 ± 0.0097 0.9024 ± 0.0097
ISAB (16) + PMA 0.9210 ± 0.0055 0.9223 ± 0.0056
ISAB (32) + PMA 0.9103 ± 0.0061 0.9119 ± 0.0052
ISAB (64) + PMA 0.9141 ± 0.0040 0.9153 ± 0.0041

Table 9. Detailed architectures used in CelebA meta set anomaly experiments. Conv(d, k, s, r, f ) is a convolutional layer with d output
channels, k kernel size, s stride size, r regularization method, and activation function f . If d is a list, each element in the list is distributed.
FC(d, f, r) denotes a fully-connected layer with d units, activation function f and r regularization method. If d is a list, each element in
the list is distributed. SAB(d, h) denotes the SAB with d units and h heads. PMA(d, h, nseed ) denotes the PMA with d units, h heads
and nseed vectors. All MABs used in SAB and PMA uses FC layers with ReLU activations for rFF layers.

Encoder Decoder
rFF SAB Pooling PMA
Conv([32, 64, 128], 3, 2, Dropout, ReLU) mean PMA4 (128, 4)
FC([1024, 512, 256], −, Dropout) FC(128, ReLU, −) SAB(128, 4)
FC(256, −, −) FC(128, ReLU, −) FC(256 · 8, −, −)
FC([128, 128, 128], ReLU, −) SAB(128, 4) FC(128, ReLU, −)
FC([128, 128, 128], ReLU, −) SAB(128, 4) FC(256 · 8, −, −)
FC(128, ReLU, −) SAB(128, 4)
FC(128, −, −) SAB(128, 4)

took to process 10,000 sets of each size. The maximum set size we report for SAB is 2,000 because the computation graph
of bigger sets could not fit on our GPU. The specific attention blocks used are ISAB4 (64, 8) and SAB(64, 8).
 Supplementary Material for Set Transformer

Table 10. Detailed architectures used in the point cloud classification experiments.

Encoder Decoder
rFF ISAB Pooling PMA
FC(256, ReLU) ISAB(256, 4) max Dropout(0.5)
FC(256, ReLU) ISAB(256, 4) Dropout(0.5) PMA1 (256, 4)
FC(256, ReLU) FC(256, ReLU) Dropout(0.5)
FC(256, −) Dropout(0.5) FC(40, −)
FC(40, −)

Figure 1. Runtime of a single SAB/ISAB block on dummy data. x axis is the size of the input set and y axis is time (seconds). Note that
the x-axis is log-scale.