FSPool: Learning Set Representations with Featurewise Sort Pooling

TL;DR

FSPool introduces a sorting-based pooling method for set representations that improves auto-encoder reconstructions and enhances accuracy and convergence in set encoding tasks.

Contribution

The paper proposes FSPool, a novel permutation-equivariant pooling method based on feature sorting, addressing the responsibility problem in set prediction models.

Findings

FSPool yields better reconstructions on toy and MNIST datasets.

Replacing pooling with FSPool improves accuracy and convergence speed.

Auto-encoders with FSPool outperform traditional methods.

Abstract

Traditional set prediction models can struggle with simple datasets due to an issue we call the responsibility problem. We introduce a pooling method for sets of feature vectors based on sorting features across elements of the set. This can be used to construct a permutation-equivariant auto-encoder that avoids this responsibility problem. On a toy dataset of polygons and a set version of MNIST, we show that such an auto-encoder produces considerably better reconstructions and representations. Replacing the pooling function in existing set encoders with FSPool improves accuracy and convergence speed on a variety of datasets.