Learning Set-equivariant Functions with SWARM Mappings
This addresses a foundational problem in machine learning for building permutation-invariant models, with incremental improvements over prior set-equivariant approaches.
The authors tackled the problem of learning set-equivariant functions, which map sets of entities to sets of the same cardinality regardless of order, by proposing a new neural network architecture called SWARM mapping based on a gated recurrent network, and demonstrated that it achieves state-of-the-art results compared to existing methods like the Set Transformer.
In this work we propose a new neural network architecture that efficiently implements and learns general purpose set-equivariant functions. Such a function f maps a set of entities x = {x1, . . . , xn} from one domain to a set of same cardinality y = f (x) = {y1, . . . , yn} in another domain regardless of the ordering of the entities. The architecture is based on a gated recurrent network which is iteratively applied to all entities individually and at the same time syncs with the progression of the whole population. In reminiscence to this pattern, which can be frequently observed in nature, we call our approach SWARM mapping. Set-equivariant and generally permutation invariant functions are important building blocks for many state of the art machine learning approaches. Even in applications where the permutation invariance is not of primary interest, as to be seen in the recent success of attention based transformer models (Vaswani et. al. 2017). Accordingly, we demonstrate the power and usefulness of SWARM mappings in different applications. We compare the performance of our approach with another recently proposed set-equivariant function, the Set Transformer (Lee et.al. 2018) and we demonstrate that models solely based on SWARM layers gives state of the art results.