Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks

1. Introduction

• Set-input tasks include multiple instance learning, 3D shape recognition, sequence ordering, set operations, and few-shot classification using image support sets.Examples include multiple instance learning (Dietterich et al., 1997; Maron & Lozano-P´erez, 1998) and few-shot image classification (Finn et al., 2017; Snell et al., 2017; Lee & Choi, 2018).Source: §1 · p. 1
• Set-input models should be permutation invariant and handle sets of any size, requirements that classical feed-forward networks violate and RNNs meet only with order sensitivity.Source: §1 · p. 1
• Set pooling independently embeds each element and aggregates embeddings, satisfying both requirements and universally approximating set functions (Edwards & Storkey (2017); Zaheer et al. (2017)).Source: §1 · p. 1
• Because pooling processes elements independently, it discards interactions; in amortized clustering, its quantization cannot depend on set contents, limiting solution quality and contributing to under-fitting.Source: §1 · pp. 1–2
• Self-attention captures pairwise or higher-order set interactions, while an inducing-point method reduces computation from O(n^2) to O(nm) and attention-based aggregation supports dependent multiple outputs.Here, m is a fixed hyperparameter; meta-clustering is an example where each cluster center depends on the others.Source: §1 · p. 2
• The Set Transformer achieves state-of-the-art performance on most of the set-input tasks evaluated in the paper.Source: §1 · p. 2

2.1. Pooling Architecture for Sets

• Pooling over element embeddings gives a simple permutation-invariant model, while most set architectures separate elementwise encoding from aggregating and decoding the encoded features.The paper identifies φ as the encoder and ρ(pool(·)) as the decoder.Source: §2.1 · p. 2
• Sum pooling with continuous ρ and φ can represent every permutation-invariant function, and invariance persists when the encoder uses permutation-equivariant layers.The paper attributes these results to Zaheer et al. (2017).Source: §2.1 · p. 2

2.2. Attention

• Attention weights each value according to the dot-product similarity between its key and a query, producing a weighted sum of values.The attention function maps query vectors to outputs using key-value pairs.Source: §2.2 · p. 2
• Multi-head attention projects queries, keys, and values into h subspaces, applies attention separately in each, then linearly transforms the concatenated outputs.Multi-head attention was originally introduced in Vaswani et al. (2017).Source: §2.2 · pp. 2–3

3. Set Transformer

• Unlike fixed mean pooling, the Set Transformer uses parameterized aggregation that adapts to the task, while attention layers operate throughout the encoder and decoder.Source: §3 · p. 3

3.1. Permutation Equivariant (Induced) Set Attention Blocks

• SAB self-attention captures pairwise interactions that can be composed into higher-order features; ISAB uses trainable inducing points to reduce attention cost from O(n^2) to O(nm).ISAB first lets the inducing points attend to X, then lets X attend to those transformed points; SAB and ISAB are permutation equivariant.Source: §3.1 · pp. 3–4

3.2. Pooling by Multihead Attention

• PMA aggregates encoded features by attention from k learnable seed vectors, usually using k=1; when k outputs are correlated, an SAB models their interactions afterward.Attention can weight instances unequally, which is useful when a target depends strongly on a particular instance.Source: §3.2 · p. 4

3.3. Overall Architecture

• The encoder stacks SABs or ISABs and the decoder aggregates features before a feed-forward network; ℓ ISAB stacks cost O(ℓnm), versus O(ℓn^2) for SAB stacks.PMA with more than one seed should be followed by SABs to model correlations among outputs.Source: §3.3 · p. 4

3.4. Analysis

• The Set Transformer is permutation invariant and a universal approximator of permutation-invariant functions.Source: §3.4 · p. 4

Pooling architectures for permutation invariant mappings

• Pooling architectures serve diverse set tasks; Zaheer et al. (2017) discuss their structure and partial universality proof, Wagstaff et al. (2019) analyze their limitations, and Bloem-Reddy & Teh (2019) link them to probabilistic exchangeability.Source: §4 · pp. 4–5

Attention-based approaches for sets

• Vinyals et al. (2016), Yang et al. (2018), and Ilse et al. (2018) use attention-weighted pooling; Set Transformer adds multihead aggregation and post-pooling self-attention for correlations among outputs.PMA with one seed and single-head attention roughly corresponds to those earlier approaches; Mishra et al. (2018) use attention for sequence-based meta-learning, while Kim et al. (2019) apply self-attention to query sets for conditional regression.Source: §4 · p. 5

Modeling interactions between elements in sets

• Santoro et al. (2017) sum-pool pairwise interactions but not higher-order ones, while Ma et al. (2018) use a Transformer for video-object interactions and mean-pool features into an LSTM.Source: §4 · p. 5

Inducing point methods

• Set Transformer’s trainable inducing vectors draw on sparse Gaussian-process inducing points and the Nyström method; m vectors can act as attention-accessed memory cells.The differential neural dictionary stores prior experience as key-value pairs, whereas ISAB reverses this arrangement by storing queries I and using input features as key-value pairs.Source: §4 · p. 5

5.1. Toy Problem: Maximum Value Regression

• On maximum-value regression, Set Transformer performance was comparable to max pooling, while mean- and sum-pooling architectures had high mean absolute error.Max pooling predicted the target perfectly; the authors note that attention can learn to find and attend to the maximum element.Source: Table 1 · p. 5, §5.1 · p. 5

5.2. Counting Unique Characters

• ISAB1 + PMA exceeded both rFF+Pooling and rFF+PMA accuracy, and performance generally increased with the number of inducing points.The task predicts the number of unique characters in sets of 6–10 Omniglot images sampled from held-out character classes.Source: Figure 3 · p. 6, §5.2 · p. 6

5.3. Amortized Clustering with Mixture of Gaussians

• Set Transformers found accurate mixture parameters, outperformed oracles after one EM update, and the 16-inducing-point variant performed best, even over the full model.The task learns to map an input dataset directly to mixture-of-Gaussians parameters rather than running EM to convergence; experiments used synthetic 2D mixtures and CIFAR-100 datasets.Source: §5.3 · pp. 6–7, Table 3 · p. 7

5.5. Point Cloud Classification

• Set Transformers were superior on small point sets, but ISAB (16) + Pooling outperformed them on larger sets.For the 5,000-point experiment, Zaheer et al. (2017) used more engineering, including data augmentation, a different optimizer, and a different learning-rate schedule.Source: Table 4 · p. 8, §5.5 · p. 8

5.4. Set Anomaly Detection

• Set Transformers outperformed all other methods by a significant margin on meta-anomaly detection, evaluated using test AUROC and AUPR.Each CelebA set contained seven images with two randomly selected attributes and one image with neither; the task was to identify the outlier.Source: §5.4 · p. 7, Table 5 · p. 8

• Attention mechanisms for both encoding and feature aggregation were empirically necessary to model complicated interactions among set elements.Source: §6 · p. 8
• The proposed inducing-point method for self-attention makes the Set Transformer scalable to large sets.Source: §6 · p. 8
• The Set Transformer is a universal approximator for permutation-invariant functions.Source: §6 · p. 8
• Future directions include using Set Transformers for meta-learning, especially Bayesian posterior inference, and principled noise injection to model uncertainty in set functions.Source: §6 · p. 8

Proofs

• The decoder, given enough nodes, can express the element-wise functions specified by Lemma 2.The proof constructs this capability through the decoder's attention and feed-forward layers.Source: §1 · p. 11
• PMA, given enough nodes, can express sum pooling.The construction uses a zero seed and a suitable activation so multihead attention outputs the sum of the values.Source: §1 · p. 11
• The architecture rFF(sum(rFF(·))) is a universal function approximator for permutation-invariant functions.This pooling universality theorem is used to establish the Set Transformer's universality.Source: §1 · pp. 11–12
• The Set Transformer is a universal function approximator for permutation-invariant functions.The proof makes the encoder instance-wise by zeroing attention output matrices, then invokes the universality result for pooling architectures.Source: §1 · p. 12
• The universality argument ignores encoder pairwise interactions, although experiments found self-attention crucial for good performance.Setting each SAB and ISAB output matrix to zero reduces the encoder to an instance-wise feed-forward network.Source: §1 · p. 12

2. Experiment Details

• Architecture notation defines FC(d, f), SAB(d, h), ISABm(d, h), and PMAk(d, h); MAB feed-forward layers use fully connected ReLU layers, and the decoder’s initial feed-forward layer is omitted.The initial decoder layer is omitted because the preceding block already ends in a feed-forward layer.Source: §2 · p. 12

2.1. Max Regression

• Max regression predicts max(x1, · · · , xn) for n sampled uniformly from 1–10, with each xi sampled independently from [0, 100], and is measured by mean absolute error.A new set is sampled at each time step rather than using explicit train and test splits.Source: §2.1 · p. 12
• Max-regression networks use Adam at learning rate 10^-3 with batch size 128 for 20,000 batches, when loss had converged for all architectures.The detailed architectures are shown in Table 1.Source: §2.1 · p. 12, Table 1 · p. 12

2.2. Counting Unique Characters

• Unique-character sets contain 6–10 instances and 1–n distinct characters, with at least one instance of each sampled character.Instances are sampled so their total equals n.Source: §2.2 · p. 12
• The final output of each architecture is a 1-dimensional value, and softplus ensures the Poisson parameter γ is positive.The architectures for this experiment are detailed in Table 3.Source: §2.2 · p. 12, Table 3 · p. 13
• The unique-character model is trained for 200,000 batches of size 32 using Adam at learning rate 10^-4 and log likelihood for a Poisson target.The optimized loss is log p(x|γ) = x log(γ) −γ −log(x!); early experiments found roughly the same result with direct mean absolute error optimization and evaluation as |γ −x|.Source: §2.2 · p. 13

2.3.1. DETAILS FOR 2D SYNTHETIC MIXTURES OF GAUSSIANS EXPERIMENT

• Synthetic datasets used 100–500 points, four 2D centers drawn uniformly from −4 to 4, and categorical cluster labels sampled from Dirichlet-distributed mixture weights.Training used 10 random datasets per step for 50k Adam steps, with the learning rate reduced from 10^-3 to 10^-4 after 35k; Table 5 reports LL0/data and LL1/data across inducing-point counts.Source: §2.3.1 · pp. 13–14, Table 5 · p. 14

2.3.3. DETAILS FOR CIFAR-100 AMORTIZED CLUTERING EXPERIMENT

• For CIFAR-100 clustering, the authors extracted 512-dimensional VGG features from 50,000 training images and formed datasets by sampling four classes and then n points from them.VGG test accuracy was 68.54%; training used 10 random datasets per step for 50k Adam steps, with the learning rate reduced from 10^-4 to 10^-5 after 35k.Source: §2.3.3 · pp. 14–15

2.3.2. 2D SYNTHETIC MIXTURES OF GAUSSIANS EXPERIMENT ON LARGE-SCALE DATA

• In large-scale 2D synthetic clustering with n sampled uniformly from 1,000 to 5,000 and k = 6, the Set Transformer with 32 inducing points works extremely well, demonstrating scalability and efficiency; Table 6 reports average log-likelihood per datum before and after one EM iteration.Source: §2.3.2 · p. 14, Table 6 · p. 15

2.4. Set Anomaly Detection

• Models used Adam with learning rate 10−4 and exponential decay for 1,000 iterations on 1,000 CelebA-subsampled datasets, split into 800 training and 200 test datasets.Table 9 describes the experiment architecture.Source: §2.4 · p. 15

2.5. Point Cloud Classification

• ModelNet40 provides 9,843 training and 2,468 test point clouds across 40 object classes.Each point cloud is a three-dimensional representation.Source: §2.5 · p. 15
• Inputs contain 100, 1,000, or 5,000 points represented by (x, y, z) coordinates, with random rotation and scaling during training.Source: §2.5 · p. 15
• For 5,000-point experiments, the attention blocks use ISAB16(512, 4), weights decay by 10^-7, and the encoder uses one ISAB block because two led to overfitting.The other described training settings use Adam with an initial learning rate of 10^-3, decayed by a factor of 0.3 every 20,000 steps.Source: §2.5 · p. 15

3.1. Runtime of SAB and ISAB

• The benchmark timed 10,000 zero-tensor sets of n three-dimensional vectors on one Tesla P40 GPU using ISAB4(64, 8) and SAB(64, 8).Source: §3.1 · pp. 15–16
• SAB measurements stopped at set size 2,000 because larger computation graphs could not fit on the GPU.Source: §3.1 · p. 16
• Figure 1 plots the runtime of a single SAB or ISAB block against input set size, using a log-scale x-axis.Source: Figure 1 · p. 17