Learning Functions on Multiple Sets using Multi-Set Transformers

TL;DR

This paper introduces a versatile deep learning architecture called Multi-Set Transformers for learning functions on multiple permutation-invariant sets, demonstrating its universality and superior performance on various tasks including statistical distance estimation.

Contribution

The paper presents a novel multi-set transformer architecture that is a universal approximator for functions on multiple sets and extends to sets of any dimension with equivariance.

Findings

Outperforms existing methods on counting, alignment, and distinguishability tasks.

Provides more accurate estimates of KL divergence and mutual information.

Demonstrates universality and generalization capabilities of the architecture.

Abstract

We propose a general deep architecture for learning functions on multiple permutation-invariant sets. We also show how to generalize this architecture to sets of elements of any dimension by dimension equivariance. We demonstrate that our architecture is a universal approximator of these functions, and show superior results to existing methods on a variety of tasks including counting tasks, alignment tasks, distinguishability tasks and statistical distance measurements. This last task is quite important in Machine Learning. Although our approach is quite general, we demonstrate that it can generate approximate estimates of KL divergence and mutual information that are more accurate than previous techniques that are specifically designed to approximate those statistical distances.