Set Norm and Equivariant Skip Connections: Putting the Deep in Deep Sets

Lily H Zhang, Veronica Tozzo, John M. Higgins, Rajesh Ranganath

Keywords: deep learning, permutation invariance, normalization, residual connections

Abstract: Permutation invariant neural networks are a promising tool for predictive modeling of set data. We show, however, that existing architectures struggle to perform well when they are deep. In this work, we address this issue for the two most widely used permutation invariant networks, Deep Sets and its transformer analogue Set Transformer. We take inspiration from previous efforts to scale neural network architectures by incorporating normalization layers and skip connections that work for sets. First, we motivate and develop set norm, a normalization tailored for sets. Then, we employ equivariant residual connections and introduce the ``clean path principle'' for their placement. With these changes, our many-layer Deep Sets++ and Set Transformer++ models reach comparable or better performance than their original counterparts on a diverse suite of tasks, from point cloud classification to regression on sets of images. We additionally introduce Flow-RBC, a new single-cell dataset and real-world application of permutation invariant prediction. On this task, our new models outperform existing methods as well as a clinical baseline. We open-source our data and code here: link-omitted-for-anonymity.

One-sentence Summary: We build deep permutation invariant models Deep Sets++ and Set Transformer++ using normalization layers and residual connections tailored for sets, and we introduce a new single-cell dataset for permutation invariant prediction.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2206.11925/code)