Conditional Set Generation with Transformers

TL;DR

This paper introduces the Transformer Set Prediction Network (TSPN), a permutation-equivariant model for set prediction that improves over previous methods like DSPN in set quality and size accuracy, tested on image and object detection tasks.

Contribution

We propose TSPN, a novel transformer-based permutation-equivariant set generator that outperforms DSPN in set quality and size accuracy.

Findings

TSPN outperforms DSPN in set quality.

TSPN achieves higher accuracy in predicted set sizes.

TSPN performs well on MNIST and CLEVR datasets.

Abstract

A set is an unordered collection of unique elements--and yet many machine learning models that generate sets impose an implicit or explicit ordering. Since model performance can depend on the choice of order, any particular ordering can lead to sub-optimal results. An alternative solution is to use a permutation-equivariant set generator, which does not specify an order-ing. An example of such a generator is the DeepSet Prediction Network (DSPN). We introduce the Transformer Set Prediction Network (TSPN), a flexible permutation-equivariant model for set prediction based on the transformer, that builds upon and outperforms DSPN in the quality of predicted set elements and in the accuracy of their predicted sizes. We test our model on MNIST-as-point-clouds (SET-MNIST) for point-cloud generation and on CLEVR for object detection.