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

Lee, Lee, Kim, Kosiorek, Choi and Teh, Oxford, POSTECH, AITRICS and Kakao (2018; ICML 2019) — ARXIV-1810.00825

Key takeaways

• Two blocks, both learned arrays attended to or from. The multihead attention block (MAB) is a Transformer encoder block without positional encoding or dropout, so it is permutation-equivariant. The ISAB adds m trainable inducing points I: first H = MAB(I, X), the inducing points attending to the n input elements, then MAB(X, H), the inputs attending back. Each block costs O(nm) and returns a set of size n, “analogous to low-rank projection or autoencoder models”. The idea is credited to inducing-point methods in sparse Gaussian processes and the Nyström method. PMA pools with k learned seed vectors attending to the encoded set (k = 1 for one output, k = 4 for four clusters), followed by self-attention among the k outputs.
• The model is permutation-invariant and a universal approximator of permutation-invariant functions (Propositions 1–2), proved in the supplement.
• Where attention helps, measured over five runs each. Counting unique Omniglot characters: SAB + PMA reaches 0.6037 accuracy against 0.4382 for sum-pooled DeepSets-style rFF + Pooling and 0.4471 for dot-product attention pooling (Table 2). Amortized clustering of mixtures of Gaussians: on CIFAR-100 features, adjusted Rand index 0.9210 for ISAB (16) + PMA, 0.9015 for SAB + PMA, 0.7612 for rFF + PMA and 0.5593 for rFF + Pooling (Table 3). The paper reads this as the attention-based decoder (PMA) mattering more than the attention-based encoder. Set anomaly detection on CelebA: AUROC 0.5941 against 0.5643.
• Inducing points cost little, and sometimes help. With only 16 inducing points, ISAB (16) + PMA beat full self-attention on both clustering datasets, which the authors attribute to regularisation. On ModelNet40 point clouds, where full attention is too slow at 1,000 and 5,000 points, ISAB (16) + Pooling reached 0.8915 and 0.9040, while ISAB (16) + PMA was best only at 100 points (0.8454) and fell behind at larger sizes (Table 4). DeepSets’ own engineered point-cloud result, with augmentation and a different optimizer, was 0.90 at 5,000.
• What is not shown. No wall-clock or FLOP comparison between SAB and ISAB, and no depth study: the encoders it sets out (eqs. 13–14) are two blocks deep. The tasks are small sets (6–10 images for counting, up to 5,000 points) and synthetic or meta-learning problems.

Standing in the anthology

Filed as the paper the Perceiver (LIT-747) calls “most closely related to our work”. The Perceiver’s appendix sets out the difference, and it is the reason this note is not the origin of SOTA-436. ISAB maps the input to a small array and immediately back to the input size, so a stack of ISABs still costs O(LMN): every layer pays for the full input. PMA maps to an array whose size is fixed by the task, one seed for classification. The Perceiver feeds the input once into a PMA-like cross-attention whose output is large (512) and independent of the task, then runs a 48-block Transformer on that array alone, which is what takes the input size out of depth. Same primitive, composed differently: the Set Transformer is the source of learned queries cross-attending to a large input, and the Perceiver of the recommendation built from them. The Perceiver does not compare against it.

It is the record’s first set-input architecture; the record holds no DeepSets note, the pooling baseline it is measured against. Permutation invariance otherwise appears only in LIT-251, which is about the symmetry of hidden units in weight space, a different object. PMA is also the pattern a learned-query resampler uses to pool a variable number of visual tokens to a fixed few; the record holds no paper that does so.

Unread — no NOTE.