Set Transformer

The paper turns self-attention into a reusable architecture for unordered sets.

Metadata

Field	Value
Paper	Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks
Authors	Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, Yee Whye Teh
Year	2018 preprint; 2019 conference
Venue	ICML 2019
arXiv	1810.00825
PMLR	v97/lee19d
Status	full note started

One-Line Takeaway

Deep Sets gives the canonical invariant set form; Set Transformer adds attention so each element can interact with other elements before the final invariant readout.

Question

Many ML inputs are naturally sets:

• a point cloud is a set of points;
• a molecule can be treated as a set of atoms or atom-pair features;
• a bag of instances has no intrinsic order;
• a cluster assignment problem consumes an unordered collection;
• a context for meta-learning can be a set of observations.

For a set

$$X=\{x_1,x_2,\dots ,x_n\},x_i\in {\mathbb{R}}^{d_x},$$

the model should not change its meaning when the input elements are permuted. If $P$ is a permutation matrix, then an invariant set-to-vector model should satisfy

$$f(PX)=f(X),$$

and an equivariant set-to-set model should satisfy

$$g(PX)=Pg(X).$$

Deep Sets showed that invariant functions can be represented in the form

$$f(X)=\rho \left(\sum _{x\in X}\varphi (x)\right).$$

That form is clean, but the interaction between two elements $x_i$ and $x_j$ must be mediated through the pooled summary. The paper asks:

Can we build a set architecture that preserves permutation symmetry while allowing rich element-element interactions?

Main Claim

Attention is a natural operation for sets because attention weights can be computed from content rather than from fixed positions. If no positional order is injected, self-attention over elements is permutation equivariant. A final learned attention pooling step can then produce permutation-invariant outputs.

The architecture pattern is:

$$X\stackrel{\text{set encoder}}{\to }H\stackrel{\text{pooling by multi-head attention}}{\to }Z\stackrel{\text{readout}}{\to }y.$$

In compact notation:

$$f(X)=\rho \left({\mathrm{PMA}}_k\left(\mathrm{Encoder}(X)\right)\right),$$

where:

• $X\in {\mathbb{R}}^{n\times d_x}$ is an unordered set represented as a matrix;
• $n$ is the number of set elements;
• $\mathrm{Encoder}$ is a stack of attention blocks;
• ${\mathrm{PMA}}_k$ uses $k$ learned seed vectors to pool into $k$ output slots;
• $\rho$ maps pooled slots to the task output.

The important distinction from sequence Transformers is that the rows of $X$ are set elements, not ordered tokens. Without positional encodings, the model is symmetric to row permutations.

Architecture Contract

Part	Input	Output	Symmetry
MAB	query set $X$, key/value set $Y$	updated query set	equivariant in query order
SAB	set $X$	contextualized set $H$	permutation equivariant
ISAB	set $X$ and learned inducing points $I$	contextualized set $H$	permutation equivariant
PMA	set $X$ and learned seeds $S$	fixed number of pooled slots	permutation invariant to input order

The model is therefore useful when:

• the input is unordered;
• element interactions matter;
• the output is a set-level label, value, distribution parameter, or small set of slots;
• the task does not require geometric equivariance beyond permutation symmetry.

Multi-Head Attention Block

The paper uses a Multihead Attention Block, usually written as MAB. Let

$$X\in {\mathbb{R}}^{n_q\times d},Y\in {\mathbb{R}}^{n_k\times d}.$$

For a single attention head, define:

$$Q=X{W}^Q,K=Y{W}^K,V=Y{W}^V,$$

where

$$W^Q,W^K,W^V\in {\mathbb{R}}^{d\times d_h}.$$

Scaled dot-product attention is:

$$\mathrm{Attn}(X,Y)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_h}}\right)V.$$

For $h$ heads:

$$\mathrm{Multihead}(X,Y)=\mathrm{Concat}(O_1,\dots ,O_h)W^O,$$

where

$$O_j=\mathrm{softmax}\left(\frac{X{W}_j^Q(Y{W}_j^K{)}^{\top }}{\sqrt{d_h}}\right)Y{W}_j^V.$$

The block then wraps this attention operation with residual, normalization, and a row-wise feed-forward layer. A common simplified form is:

$$H=\mathrm{LayerNorm}\left(X+\mathrm{Multihead}(X,Y)\right),$$ $$\mathrm{MAB}(X,Y)=\mathrm{LayerNorm}\left(H+\mathrm{rFF}(H)\right).$$

Here $\mathrm{rFF}$ means the same feed-forward network is applied independently to each row:

$$\mathrm{rFF}(H{)}_i=\mathrm{FFN}(H_i).$$

This row-wise sharing is part of the permutation behavior.

Self-Attention Block

The Self-Attention Block is simply:

$$\mathrm{SAB}(X)=\mathrm{MAB}(X,X).$$

Each element attends to every other element. If $X$ is permuted by $P$, then the attention matrix is permuted consistently:

$$Q^{\prime }=PX{W}^Q=PQ,$$ $$K^{\prime }=PX{W}^K=PK,$$ $$V^{\prime }=PX{W}^V=PV.$$

The logits become:

$$Q^{\prime }(K^{\prime }{)}^{\top }=PQ{K}^{\top }{P}^{\top }.$$

Row-wise softmax preserves this permutation structure:

$$\mathrm{softmax}(PA{P}^{\top })=P\mathrm{softmax}(A)P^{\top }.$$

Therefore:

$$\mathrm{Attn}(PX,PX)=P\mathrm{Attn}(X,X).$$

The output rows move with the input rows. This is permutation equivariance.

Induced Self-Attention Block

Full self-attention over $n$ set elements costs:

$$O(n^{2}d)$$

per layer up to constants and head count. For large sets, this can dominate compute and memory.

Set Transformer introduces Induced Self-Attention Block:

$${\mathrm{ISAB}}_m(X)=\mathrm{MAB}\left(X,\mathrm{MAB}(I,X)\right),$$

where:

• $I\in {\mathbb{R}}^{m\times d}$ is a learned set of inducing points;
• $m$ is much smaller than $n$ in the efficient regime;
• $\mathrm{MAB}(I,X)$ lets inducing points attend to the input set;
• $\mathrm{MAB}(X,\cdot )$ lets input elements attend back to the induced summary.

The two attention passes cost roughly:

$$O(mnd)+O(nmd)=O(nmd).$$

When $m\ll n$:

$$O(nmd)\ll O(n^{2}d).$$

The inducing points act like a learned information bottleneck. They are not input elements. They are trainable latent slots that summarize interactions in a lower-rank way.

Pooling by Multi-Head Attention

Simple set pooling often uses:

$$z=\sum _{i=1}^n{h}_i$$

or

$$z=\frac{1}{n}\sum _{i=1}^n{h}_i.$$

Set Transformer instead uses Pooling by Multi-Head Attention:

$${\mathrm{PMA}}_k(X)=\mathrm{MAB}(S,X),$$

where:

• $S\in {\mathbb{R}}^{k\times d}$ is a set of learned seed vectors;
• $k$ is the number of output slots;
• $X$ supplies keys and values;
• $S$ supplies queries.

For one output slot, $k=1$, PMA is a learned set-level readout:

$$z={\mathrm{PMA}}_1(X)\in {\mathbb{R}}^{1\times d}.$$

For multiple related outputs, $k>1$:

$$Z={\mathrm{PMA}}_k(X)\in {\mathbb{R}}^{k\times d}.$$

This matters for tasks like clustering, where the output is not just one scalar label but several structured slots.

Why PMA Is Invariant

Let the input set encoding be $H$. PMA computes attention from fixed learned seeds $S$ to $H$:

$${\mathrm{PMA}}_k(H)=\mathrm{MAB}(S,H).$$

If $H$ is permuted:

$$H^{\prime }=PH.$$

Then:

$$K^{\prime }=PH{W}^K=PK,$$ $$V^{\prime }=PH{W}^V=PV.$$

The seed query is unchanged:

$$Q=S{W}^Q.$$

The attention logits become:

$$Q(K^{\prime }{)}^{\top }=Q{K}^{\top }{P}^{\top }.$$

The attention weights are permuted over keys, and the values are permuted in the same way. The weighted sum is unchanged:

$$\mathrm{softmax}(Q{K}^{\top }{P}^{\top })PV=\mathrm{softmax}(Q{K}^{\top })V.$$

Thus PMA is invariant to input row order.

Encoder-Decoder Forms

The paper frames Set Transformer as an encoder-decoder family.

SAB Encoder

The direct, expressive version is:

$$\mathrm{Encoder}(X)=\mathrm{SAB}\left(\mathrm{SAB}(X)\right).$$

Then:

$$f(X)=\rho \left({\mathrm{PMA}}_k\left(\mathrm{Encoder}(X)\right)\right).$$

This gives full pairwise attention at every layer.

ISAB Encoder

The efficient version is:

$$\mathrm{Encoder}(X)={\mathrm{ISAB}}_m\left({\mathrm{ISAB}}_m(X)\right).$$

Then:

$$f(X)=\rho \left({\mathrm{PMA}}_k\left(\mathrm{Encoder}(X)\right)\right).$$

This trades exact full interaction for induced latent interaction.

Comparison to Deep Sets

Property	Deep Sets	Set Transformer
Core form	$\rho ({\sum }_i\varphi (x_i))$	attention encoder + PMA
Element interaction before pooling	indirect through sum	direct through attention
Invariance mechanism	symmetric pooling	equivariant attention plus invariant pooling
Complexity	often $O(n)$ after row-wise encoding	$O(n^2)$ for SAB or $O(nm)$ for ISAB
Strength	simple, universal form, stable baseline	richer interaction modeling
Weakness	pooling bottleneck	heavier compute and more hyperparameters

Deep Sets is often the right first baseline. Set Transformer is worth using when predictions depend on relations among elements, such as nearest-neighbor structure, clustering, pairwise compatibility, or context-dependent importance.

Comparison to Sequence Transformer

Question	Sequence Transformer	Set Transformer
Input assumption	ordered tokens	unordered elements
Positional encoding	central to order awareness	omitted unless extra structure is intended
Attention role	token mixing plus order-aware context	element interaction under permutation symmetry
Output	sequence, next token, class token, span	set-level value, slots, set prediction
Wrong use case	unordered set treated as arbitrary sequence	ordered sequence where order matters

The absence of positional encoding is not a missing feature here. It is the architectural contract.

Complexity

For $n$ elements, hidden width $d$, and $h$ heads:

Module	Attention pattern	Approximate attention cost
SAB	all elements attend to all elements	$O(n^{2}d)$
ISAB	$m$ inducing points mediate attention	$O(nmd)$
PMA	$k$ seeds attend to $n$ elements	$O(knd)$

Memory follows the attention matrix sizes:

$$\mathrm{SAB}:O(n^2),$$ $$\mathrm{ISAB}:O(nm),$$ $$\mathrm{PMA}:O(kn).$$

The architecture therefore has three important size knobs:

• $n$: number of input elements;
• $m$: number of inducing points;
• $k$: number of output slots.

What the Inducing Points Mean

The inducing points are easy to misread. They are not cluster centers in the input space by definition. They are learned latent query vectors used to route information.

One can think of:

$$U=\mathrm{MAB}(I,X)$$

as a learned summary with $m$ slots. Then:

$$H=\mathrm{MAB}(X,U)$$

uses that summary to update each input element. If $m$ is too small, the bottleneck can underfit. If $m$ is too large, the computation approaches full attention while adding more parameters.

Evidence Reading

The paper evaluates the architecture on synthetic and practical set tasks. The main evidence is not that Set Transformer is always best, but that attention-based interaction improves over simple pooled set functions in tasks where interactions matter.

Claim	Evidence Type	What It Supports	Caveat
Attention helps set modeling	comparison with Deep Sets-style baselines	element interaction before pooling is useful	task-dependent
PMA improves readout	pooling comparisons	learned pooling can beat fixed pooling	seed count and decoder design matter
ISAB improves efficiency	induced attention design and experiments	lower cost than full self-attention	approximation quality depends on $m$
Multi-output slots are useful	clustering-style outputs	set outputs can be modeled with multiple seeds	slot semantics may be unstable

Benchmark Card

Field	Value
Main task family	set-structured learning
Input unit	unordered elements
Output unit	set label, distribution parameter, or multiple output slots
Main baseline	Deep Sets and task-specific set models
Central mechanism	self-attention without positional ordering
Efficiency mechanism	inducing-point attention
Main architectural risk	using attention when simpler symmetric pooling is enough
Not the claim	solving geometric equivariance, graph edge semantics, or ordered sequence modeling

How to Decide Between Deep Sets and Set Transformer

Use a Deep Sets-style model first when:

• each element contributes independently after feature extraction;
• the set size is large and compute is tight;
• interpretability of a simple pooled summary matters;
• the target depends mostly on aggregate statistics.

Use Set Transformer when:

• pairwise or higher-order relations matter;
• the importance of an element depends on other elements;
• clustering or multiple output slots are needed;
• the set size is moderate or ISAB can control cost;
• attention maps are useful diagnostic artifacts.

Molecular and Structural Modeling Reading

Set Transformer is relevant to molecular and structural modeling, but it should not be treated as a full replacement for graph or geometric models.

Possible fit:

• unordered atom feature sets;
• residue feature sets before explicit geometry is used;
• ligand conformer bags;
• multiple docking pose candidates;
• pocket residue sets;
• sets of local environment descriptors.

Potential mismatch:

• chemical bonds are not just unordered co-occurrence;
• 3D coordinates require rotation/translation-aware handling;
• pairwise distances and angles may need explicit geometric features;
• graph topology may be better handled by Graph neural networks;
• equivariant tasks may require Geometric deep learning.

A useful hybrid pattern is:

$$\text{local graph/geometric encoder}\to \text{set of object embeddings}\to \text{Set Transformer readout}.$$

For example, each conformer can be encoded by a structure model, and a Set Transformer can aggregate a set of conformer embeddings.

Implementation Notes

Masking

Real batches often contain variable-size sets. If a batch tensor has shape:

$$X\in {\mathbb{R}}^{B\times n_{\max }\times d},$$

then padding masks are required. The attention logits for padded keys should be set to a large negative value before softmax:

$$A_{ij}=\{\begin{array}{ll}\frac{q_i^{\top }{k}_j}{\sqrt{d_h}},& \text{valid key }j,\\ -\infty ,& \text{padded key }j.\end{array}$$

Without masking, the model can learn artifacts from padding count.

No Accidental Position Features

If the input order is arbitrary, do not add:

• index embeddings;
• order-dependent augmentation;
• sorting by label-derived quantities;
• positional encodings that are not meaningful features.

If elements are sorted by a domain feature, the model may silently become order-dependent through that feature pipeline.

Normalization

Layer normalization is usually compatible with set symmetry because it is applied row-wise over features. Batch normalization can be more delicate if statistics depend on padding, set size, or batching patterns.

Pooling Slot Count

The seed count $k$ is an architectural choice:

$$S\in {\mathbb{R}}^{k\times d}.$$

Use:

• $k=1$ for one set-level embedding;
• $k>1$ for multiple output slots, mixture components, cluster-like outputs, or structured predictions.

Common Misreadings

”Attention makes the model invariant.”

Self-attention over a set is permutation equivariant, not invariant. The invariant part comes from the pooling/readout stage:

$$\mathrm{SAB}(PX)=P\mathrm{SAB}(X),$$

but:

$$\mathrm{PMA}(\mathrm{SAB}(PX))=\mathrm{PMA}(\mathrm{SAB}(X)).$$

“Set Transformer is just Transformer without position embeddings.”

That is only partly true. The important additions are PMA and ISAB, which adapt attention to set readout and large-set scaling.

”Inducing points are interpretable prototypes.”

They may become semantically suggestive, but the architecture only defines them as learned latent queries. Interpretability is not guaranteed.

”Permutation invariance is always desirable for molecules.”

Only some molecular representations are sets. A molecule also has graph structure, atom types, bond types, stereochemistry, conformers, and 3D geometry. Treating it as a plain set can discard useful structure.

Later-Paper Checklist

When reading a later set, graph, point-cloud, or multimodal architecture paper, ask:

• Does the model need permutation invariance, equivariance, or neither?
• Is the symmetry exact by construction or learned approximately?
• Where do element interactions happen before pooling?
• Is pooling fixed, learned, attention-based, or task-specific?
• Does complexity scale as $O(n)$, $O(nm)$, or $O(n^2)$?
• Are there latent slots, inducing points, memory tokens, or bottleneck tokens?
• Does the task require geometric equivariance beyond set symmetry?
• Are results compared to Deep Sets and attention-based set baselines?

Why It Matters

Set Transformer is an important architecture paper because it clarifies how attention can be used outside ordered sequences. It connects three ideas that now appear repeatedly in modern models:

• permutation-aware modeling;
• latent bottleneck tokens;
• learned pooling/readout.

This makes it a bridge from Deep Sets to later latent-array architectures such as Perceiver IO, and to object-centric or slot-based modeling patterns.

Limitations

• SAB is quadratic in set size.
• ISAB introduces the inducing-point hyperparameter $m$.
• Attention does not encode graph edges unless they are added through features or masks.
• Plain set symmetry is insufficient for many geometric tasks.
• Output slots can be permutation-sensitive among themselves unless the task and loss handle slot matching.
• Empirical gains can depend on task size, feature quality, and baseline strength.

Connections

• Set Transformer
• Deep Sets
• Attention
• Pooling and readout
• Inductive bias
• Graph neural networks
• Geometric deep learning
• Deep Sets
• Perceiver IO
• Architecture papers