Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, Yee Whye Teh

TL;DR

The paper introduces Set Transformer, a permutation-invariant neural framework built from attention blocks (SAB, ISAB) and pooling (PMA) to model interactions within variable-size sets. It demonstrates a universal function-approximation property for permutation-invariant tasks and proposes an inductive-point-based attention variant to reduce quadratic complexity. Through extensive experiments across max-aggregation, counting, Gaussian mixture clustering, anomaly detection, and point-cloud classification, the approach achieves strong performance and scalability on diverse set-structured problems. The supplementary results provide detailed architectural, data-generation, and training protocols, validating both accuracy and computational efficiency across small- and large-scale tasks. These insights highlight Set Transformer as a versatile tool for learning on unordered, variable-sized data.

Abstract

Many machine learning tasks such as multiple instance learning, 3D shape recognition, and few-shot image classification are defined on sets of instances. Since solutions to such problems do not depend on the order of elements of the set, models used to address them should be permutation invariant. We present an attention-based neural network module, the Set Transformer, specifically designed to model interactions among elements in the input set. The model consists of an encoder and a decoder, both of which rely on attention mechanisms. In an effort to reduce computational complexity, we introduce an attention scheme inspired by inducing point methods from sparse Gaussian process literature. It reduces the computation time of self-attention from quadratic to linear in the number of elements in the set. We show that our model is theoretically attractive and we evaluate it on a range of tasks, demonstrating the state-of-the-art performance compared to recent methods for set-structured data.

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

TL;DR

The paper introduces Set Transformer, a permutation-invariant neural framework built from attention blocks (SAB, ISAB) and pooling (PMA) to model interactions within variable-size sets. It demonstrates a universal function-approximation property for permutation-invariant tasks and proposes an inductive-point-based attention variant to reduce quadratic complexity. Through extensive experiments across max-aggregation, counting, Gaussian mixture clustering, anomaly detection, and point-cloud classification, the approach achieves strong performance and scalability on diverse set-structured problems. The supplementary results provide detailed architectural, data-generation, and training protocols, validating both accuracy and computational efficiency across small- and large-scale tasks. These insights highlight Set Transformer as a versatile tool for learning on unordered, variable-sized data.

Abstract

Many machine learning tasks such as multiple instance learning, 3D shape recognition, and few-shot image classification are defined on sets of instances. Since solutions to such problems do not depend on the order of elements of the set, models used to address them should be permutation invariant. We present an attention-based neural network module, the Set Transformer, specifically designed to model interactions among elements in the input set. The model consists of an encoder and a decoder, both of which rely on attention mechanisms. In an effort to reduce computational complexity, we introduce an attention scheme inspired by inducing point methods from sparse Gaussian process literature. It reduces the computation time of self-attention from quadratic to linear in the number of elements in the set. We show that our model is theoretically attractive and we evaluate it on a range of tasks, demonstrating the state-of-the-art performance compared to recent methods for set-structured data.

Paper Structure

This paper contains 12 sections, 5 theorems, 3 equations, 1 figure, 10 tables.

Table of Contents

  1. Proofs
  2. Experiment Details
  3. Max Regression
  4. Counting Unique Characters
  5. Solving maximum likelihood problems for mixture of Gaussians
  6. Details for 2D synthetic mixtures of Gaussians experiment
  7. 2D Synthetic Mixtures of Gaussians Experiment on Large-scale Data
  8. Details for CIFAR-100 amortized clutering experiment
  9. Set Anomaly Detection
  10. Point Cloud Classification
  11. Additional Experiments
  12. Runtime of SAB and ISAB

Key Result

Lemma 1

The mean operator $\mathrm{mean}(\{x_1, \ldots, x_n \}) = \frac{1}{n} \sum_{i=1}^n x_i$ is a special case of dot-product attention with softmax.

Figures (1)

  • 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.

Theorems & Definitions (10)

  • Lemma 1
  • proof
  • Lemma 2
  • proof
  • Lemma 3
  • proof
  • Theorem 1
  • proof
  • Proposition 1
  • proof