PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:1810.00825v3cs.LGstat.ML

TL;DR

Set-structured learning requires permutation-invariant models that can represent interactions among elements, while simple pooling may discard such information. The paper introduces the attention-based Set Transformer, using attention for encoding and aggregation and inducing points for scalable self-attention. It reports state-of-the-art performance across most evaluated set-input tasks and establishes theoretical properties including universality for permutation-invariant functions.

  • Problem

    Set-input tasks require permutation invariance, variable-size input handling, and interactions among elements that simple pooling may discard.

  • Method

    The Set Transformer uses attention-based encoding and aggregation, with inducing points reducing self-attention complexity for large sets.

  • Results

    The Set Transformer achieves state-of-the-art performance on most evaluated set-input tasks.

  • Takeaways & Limitations

    Attention mechanisms model complicated interactions among set elements, while the inducing-point operation makes the approach scalable to large sets.

Abstract

from arXiv · show

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.

1. Introduction

Set-input tasks require models that are permutation invariant and can handle sets of varying size, but simple pooling may discard interactions among elements. The Set Transformer addresses this with attention-based encoding and aggregation, including an inducing-point scheme for scalable self-attention, and achieves state-of-the-art performance across evaluated tasks.

  • Motivation: Set-input problems require permutation-invariant outputs and support for input sets of any size.Classical feed-forward networks violate both requirements, while RNNs are sensitive to input order.
  • Limitations of pooling: Set pooling methods satisfy these requirements and are universal approximators, but independently processing elements can discard information about their interactions.This information loss may make some mappings unnecessarily difficult to solve.
  • Proposed architecture: The Set Transformer uses self-attention to encode pairwise and higher-order interactions among elements in the input set.Its attention-based aggregation is also designed for problems with multiple mutually dependent outputs, such as meta-clustering.
  • Scalability: O(nm) replaces O(n^2) self-attention computation, where m is a fixed hyperparameter, enabling scaling to large input sets.The method is introduced to reduce the computation time of full self-attention.
  • Empirical evaluation: The Set Transformer achieves state-of-the-art performance on most evaluated set-input tasks while empirically validating its design choices.The paper applies the model to several set-input problems.

2. Background

Set models commonly use an encoder-decoder structure with elementwise encoding followed by pooling, while attention provides a mechanism for modeling relationships among set elements. The background introduces permutation-equivariant attention operations and multihead attention, including their computational and parameter conventions.

  • Pooling architecture: Set-input models must preserve permutation invariance: reordering a set must not change its target value.A standard pooling architecture applies an elementwise encoder and then aggregates the resulting embeddings.
  • Pooling architecture: Sum pooling with continuous functions can represent all permutation-invariant functions, providing a universal-approximation justification for this architecture.This result is attributed to Zaheer et al. (2017).
  • Pooling architecture: The pooling architecture decomposes into an encoder that independently processes each element and a decoder that aggregates encoded features.Most set-structured-data architectures follow this encoder-decoder structure.
  • Equivariance: Permutation-equivariant layers transform outputs consistently with input reordering, and the described layer combines an element with a pooled representation of the set.The layer uses learnable scalar variables and a nonlinear activation function.
  • Attention: Attention maps queries to outputs using key-value pairs, with dot products determining weights for a weighted sum of values.Multihead attention applies attention across multiple projected representations and linearly transforms their concatenation.
  • Attention: Multihead attention concatenates the outputs of h attention heads and applies an output projection.The formulation includes learnable projection parameters and standard dimension choices such as dM_v = d_v/h and d = d_q.

3. Set Transformer

The Set Transformer uses trainable attention-based encoder and decoder blocks to model interactions within sets while preserving permutation invariance. Inducing points reduce encoder complexity, and learnable seed vectors provide adaptive aggregation for single or correlated outputs.

  • Encoder: Self-attention concurrently encodes the whole set, capturing pairwise and higher-order interactions that independent pooling discards.Stacked SABs encode higher-order interactions from pairwise relationships.
  • Encoder: ISAB replaces quadratic self-attention with O(nm) computation using m trainable inducing points, where m is typically small.The inducing points first attend to the input and are then attended to by the input set.
  • Decoder: PMA aggregates encoded features through multihead attention from k learnable seed vectors instead of fixed mean or maximum pooling.The number of seeds determines whether the decoder produces one output or a set of outputs.
  • Decoder: For tasks requiring correlated outputs, such as amortized clustering, multiple seeds are followed by SAB to model interactions among outputs.The paper reports that this post-pooling self-attention helps model explaining-away among clusters.
  • Permutation properties: SAB and ISAB are permutation equivariant, while PMA is permutation invariant, making the complete Set Transformer permutation invariant.The encoder maps sets equivariantly, and the decoder’s PMA produces invariant outputs.
  • Theory: The paper proves that Set Transformers are universal approximators of permutation invariant functions.The proof is provided in the supplementary material.

4. Related Works

Related work covers pooling architectures, attention-based aggregation, explicit interaction modeling, and inducing-point methods for set-structured data.

  • Pooling architectures: Pooling architectures independently embed set elements and aggregate them with operations such as mean, sum, or max.These methods satisfy permutation invariance and support inputs of varying size, but independent processing can discard interactions.
  • Attention-based approaches: Attention-based set methods compute weighted pooling, including weighted averages or weighted sums based on learned attention scores.Examples include attention mechanisms for multi-view 3D reconstruction and multiple instance learning.
  • Interaction modeling: Prior interaction models include relational networks that sum-pool pairwise interactions, while Transformer-based work models object interactions in videos.The relational-network example does not model higher-order interactions.
  • Inducing point methods: Set Transformer inducing points are related to sparse Gaussian-process inducing points and Nyström matrix decomposition methods.The paper also characterizes the trainable inducing points as independent memory cells accessed through attention.

5. Experiments

The experiments evaluate Set Transformer variants and baselines on maximum regression, unique-character counting, amortized clustering, anomaly detection, and point-cloud classification. Across these tasks, attention-based aggregation and inducing-point architectures show strong performance, while results vary with set size and model configuration.

  • Maximum Value Regression: Set Transformer matches max-pooling performance on maximum-value regression, while mean- and sum-pooling produce high mean absolute error.Attention can learn to identify and attend to the maximum element.
  • Counting Unique Characters: ISAB1 + PMA surpasses both rFF+Pooling and rFF+PMA on unique-character counting, with accuracy tending to increase as the inducing-point count n grows.The task predicts the number of distinct characters in Omniglot image sets.
  • Amortized Clustering: Set Transformers find accurate mixture-of-Gaussian parameters on synthetic and CIFAR-100 datasets, with the 16-inducing-point variant achieving the best reported performance.The evaluation uses average likelihood per data for synthetic mixtures and adjusted rand index for CIFAR-100, including scores after one EM update.
  • Set Anomaly Detection: Set Transformers outperform all other methods by a significant margin on CelebA meta-anomaly detection.The task identifies the one image lacking two attributes shared by the other seven images; performance is reported with AUROC and AUPR.
  • Point Cloud Classification: Set Transformers are superior on small point clouds, whereas ISAB (16) + Pooling outperforms them on larger sets; PMA outperforms simple pooling in the other experiments.The comparison reports test accuracy using 100, 1000, and 5000 points.

6. Conclusion

The paper introduces Set Transformer as an attention-based architecture for set inputs, combining attention-based encoding and aggregation. It reports empirical validation of attention mechanisms, scalability through inducing points, and universality for permutation invariant functions.

  • Set Transformer uses attention mechanisms for both encoding and aggregating features.
  • Experiments indicate that both encoder self-attention and attention-based aggregation are necessary for modelling complicated set interactions.
  • The inducing point method for self-attention makes the approach scalable to large sets.
  • The model is a universal function approximator for permutation invariant functions.
  • The paper identifies meta-learning posterior inference and principled uncertainty modelling as future research directions.

Supplementary Material for Set Transformer

The supplementary material passage identifies the paper’s authors.

  • The listed authors are Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh.
  • The author list includes six named researchers.
  • The passage presents the authors as contributors to Set Transformer.

1. Proofs

The proofs establish theoretical expressivity for Set Transformer components and the full architecture. They show that pooling architectures and Set Transformers can approximate permutation invariant functions, while experiments find encoder self-attention important for performance.

  • PMA with enough nodes can express sum pooling.
  • The construction uses a zero seed and a suitable activation to make multihead attention output a sum of the values.
  • Models of the form rFF(sum(rFF(·))) are universal function approximators for permutation invariant functions.
  • The Set Transformer is a universal function approximator for permutation invariant functions.
  • The universality proof obtains instance-wise feed-forward processing by setting W_O to a zero matrix in every SAB and ISAB.
  • Experiments indicate that encoder self-attention remains crucial for good performance despite being ignored in the universality proof.

2. Experiment Details

The experiments use specified Set Transformer building blocks and training procedures across max regression and unique character counting tasks. Architectural notation distinguishes self-attention, inducing-point self-attention, and pooling modules.

  • Architecture notation: SAB(d, h), ISAB_m(d, h), and PMA_k(d, h) denote modules with units, heads, inducing points, and pooling vectors as applicable.
  • Implementation: All MABs inside SAB and PMA use fully connected feed-forward layers with ReLU activations.
  • Max regression: Max regression samples set sizes uniformly from 1 through 10 and real-valued elements independently from [0, 100].
  • Max regression: The max regression objective is mean absolute error against the actual maximum value.
  • Training: Networks are trained with Adam at learning rate 10^-3, batch size 128, for 20,000 batches.
  • Unique character counting: Unique character counting samples set sizes from 6 through 10 and character counts from 1 through the sampled set size.
  • Unique character counting: The unique character counting output passes through softplus to produce a positive Poisson parameter γ.
  • Unique character counting: Training uses the Poisson log likelihood rather than mean squared or mean absolute error for integer targets.

2.3. Solving maximum likelihood problems for mixture of Gaussians

The experiments evaluate set-based architectures on synthetic and CIFAR-100 clustering tasks, including large-scale synthetic datasets and varying numbers of inducing points. Training uses generated datasets and compares clustering likelihoods or accuracies across architectures.

  • Synthetic clustering: Synthetic datasets vary the number of points, generate cluster centers and labels, and sample observations from spherical Gaussians.The standard process samples n from Unif(100, 500), generates k centers, draws labels from a categorical distribution, and generates data from spherical Gaussians.
  • Synthetic clustering: The synthetic clustering experiments train architectures for 50k steps with Adam and compare likelihood before and after one EM iteration.The learning rate starts at 10^-3, decays to 10^-4 after 35k steps, and results are summarized for different numbers of inducing points.
  • Large-scale clustering: Large-scale synthetic clustering samples 1,000–5,000 points with six clusters, and the Set Transformer with 32 inducing points works extremely well.The authors describe this result as demonstrating scalability and efficiency.
  • CIFAR-100 clustering: CIFAR-100 architectures are trained for 50k steps with Adam and evaluated across different numbers of ISAB inducing points.The learning rate begins at 10^-4 and decays to 10^-5 after 35k steps.

2.4. Set Anomaly Detection

The experiments include meta set anomaly detection on CelebA-derived datasets and point-cloud classification on ModelNet40. The point-cloud setting tests multiple set sizes with rotation and scale augmentation.

  • Set anomaly detection: CelebA meta set anomaly experiments use 1,000 subsampled datasets, split into 800 training and 200 test datasets.All models are trained with Adam at learning rate 10^-4 and exponential decay for 1,000 iterations.
  • Point-cloud classification: ModelNet40 contains 9,843 training and 2,468 test point clouds spanning 40 object classes.Each point cloud is represented by three-dimensional coordinates.
  • Point-cloud classification: Point-cloud experiments evaluate sets with 100, 1,000, and 5,000 points, using random rotation and scale during training for generalization.The 5,000-point experiment uses wider attention blocks and one ISAB block in the encoder.

3. Additional Experiments

Additional experiments examine runtime scaling of SAB and ISAB blocks and provide architecture details for clustering, anomaly detection, and point-cloud tasks. The runtime benchmark specifically contrasts feasible set sizes under the two attention schemes.

  • Runtime: Runtime measurements process 10,000 constant sets for each input size on a single Tesla P40 GPU.The benchmark uses zero tensors containing n three-dimensional vectors and evaluates ISAB4(64, 8) and SAB(64, 8).
  • Runtime: SAB runtime is reported only up to set size 2,000 because larger computation graphs could not fit on the GPU.The comparison is visualized with input-set size on the x axis and time in seconds on the y axis.
Loading 1810.00825v3…