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

1 INTRODUCTION Many real-world tasks involve predictions on sets or multisets , from point cloud classification ( Guo et al. , 2020 ; Wu et al. , 2015 ; Qi et al. , 2017a ) to the prediction of health outcomes from single-cell data ( Regev et al. , 2017 ; Lähnemann et al. , 2020 ; Liu et al. , 2021 ; Yuan et al. , 2017 ) . In the context of health prediction among other applications , there is the need for reliable methods that can be applied by application-based practitioners without the additional requirement of engineering models for every application task . An important property of models applied to input sets is permutation invariance : for any permutation of the instances in the input set , the model prediction stays the same . Deep Sets ( Zaheer et al. , 2017 ) and Set Transformer ( Lee et al. , 2019 ) are two general-purpose deep permutation invariant models . Such models have been proven to be universal approximators of permutation invariant functions under the right conditions ( Zaheer et al. , 2017 ; Wagstaff et al. , 2019 ) . In practice , however , the architectures are tailored to specific tasks to achieve good performance . In this work , we explore a more general approach to achieving good performance : making networks deeper , a strategy which has yielded benefit across many tasks ( He et al. , 2016b ; Wang et al. , 2019b ) . We first note that naively stacking more layers on both the Deep Sets and Set Transformer architectures hurts performance . Consider Figure 1 . Deep Sets 50 layers ( panel ( a ) ) has significantly worse performance than Deep Sets 3 layers due to the problem of vanishing gradients . Residual connections are often used to combat vanishing gradients , but even the Set Transformer architecture which consists of residual connections has worse performance at a higher depth ( panel ( b ) ) . Another strategy often used in combination with residual connections in deep architectures is the use of normalization layers , but the use of layer norm does not improve or deteriorates performance of both architectures at high depths , one with residual connections ( Set Transformer ) and one without ( Deep Sets ) . In fact , the Set Transformer paper ( Lee et al. , 2019 ) mentions layer norm as a component of their architecture but turns it off in actual implementation . In this work , we tackle the challenge of designing deep permutation invariant networks with the aim of providing general purpose architectures . • Based on real-world applications , we formalize desiderata for normalization layers when inputs are sets and systematically analyze options based on a unified framework . We introduce set norm , aimed at optimizing the desiderata we introduce . Set norm is a general , easy-to-implement normalization that can be applied on any architecture . We empirically demonstrate its benefit with ablation studies ( Section 3 , Section 5 ) . • We consider settings for residual connections which prevent vanishing gradients and maintain permutation equivariance . Through extensive experiments , we show that the best setting is that of equivariant residual connections ( ERC ) coupled the “ clean residual path ” principle ( Section 4 ) . To the best of our knowledge , we are the first to analyze this principle across architectural designs and to suggest a new benefit of certain residual connections : alleviation of information loss ( Section 5 ) . • We develop Deep Sets++ ( DS++ ) and Set Transformer++ ( ST++ ) , leveraging set norm and ERC . We demonstrate that both models yield comparable or improved performance over their existing counterparts across a variety of tasks ( Section 5 ) . Among other results , these new architectures yield better accuracy on point cloud classification than the task-specific architectures proposed in the original Deep Sets and Set Transformer papers . The fact that set norm and equivariant residual connections yield improvements on two different architectures suggests that they may also be able to benefit other permutation invariant architectures in the future . • Given growing interest in single-cell analysis ( Regev et al. , 2017 ) and the need for methods which can extract signal from single-cell data ( Lähnemann et al. , 2020 ) , we introduce a novel clinical single-cell dataset for prediction on sets called Flow-RBC . The dataset consists of red blood cell ( RBC ) measurements and hematocrit levels ( i.e . the fraction of blood volume occupied by RBCs ) for 100,000+ patients . The size and presence of a prediction target ( hematocrit ) makes this dataset unique , even among single-cell datasets in established repositories like the Human Cell Atlas ( Regev et al. , 2017 ) . Flow-RBC not only provides a new real-world benchmark for methodological innovation but also a potential source of information for biomedical science ( Section 5 ) . 2 BACKGROUND . Let S be the number of samples in a set . A function f : XS ! Y is permutation invariant if any permutation ⇡ of the input set results in the same output : f ( ⇡x ) = f ( x ) . A function : XS ! YS permutation equivariant if , for any permutation ⇡ , the outputs are permuted accordingly : ( ⇡x ) = ⇡ ( x ) . A sum/max-decomposable function f : XS ! Y is one which can be expressed via an equivariant encoder : XS ! ZS that maps each input sample in X to a latent space Z , a sum or max aggregation over the sample dimension of the encoder output , and an unconstrained decoder ⇢ : Z ! Y : f ( x ) = ⇢ ⇣X ( x ) ⌘ . ( 1 ) An architecture is a universal approximator of permutation invariant functions if and only if it is sum/max-decomposable with sufficient conditions on the latent space dimension ( Wagstaff et al. , 2019 ) . Both the Deep Sets and Set Transformer architectures can be written in the form of Equation ( 1 ) . Concretely , Deep Sets architecture consists of an equivariant encoder made up of feedforward layers , a sum or max aggregation , and a decoder also made up of feedforward layers ( Zaheer et al. , 2017 ) . The Set Transformer encoder consists of multi-attention blocks called inducing point attention blocks ( Lee et al. , 2019 ) , a variant of the original transformer block ( Vaswani et al. , 2017 ) modified to handle large set sizes . The aggregation is learned via attention , and the decoder consists of transformer blocks with self-attention . Both architectures are permutation invariant since they consist of a permutation equivariant encoder and a permutation invariant aggregation over the outputs of the encoder . In order to maintain permutation invariance of the architectures upon modification , one needs to ensure that changes to the encoder preserve permutation equivariance , and changes to the aggregation maintain permutation invariance . In the next two sections , we discuss how to incorporate normalization layers and residual connections while maintaining permutation invariance overall . 3 NORMALIZATION FOR SET DATA . Incorporating normalization layers into the permutation equivariant encoders requires careful consideration , as not all normalization layers are appropriate to use . We study normalization layers which consist of two operations , standardization and transformation . This setting captures most common normalizations ( Ioffe & Szegedy , 2015 ; Ba et al. , 2016 ; Ulyanov et al. , 2016 ) . Let a 2 RB⇥S⇥D be the activation before the normalization operation , where B is the size of the batch , S is the number of samples in a set ( the tensor is zero-padded to the largest sample size ) , and D is the feature dimension . First , the activations are standardized based on a setting L which defines which dimensions utilize separate statistics . For instance , L = { B , S } denotes that each set in a batch and each sample in a set gets its own mean and standard deviation for standardization , which are repeated over the D dimension so that µL ( a ) , L ( a ) 2 RB⇥S⇥D match the original activation tensor for elementwise subtraction and division . A standardization operation can be defined as : ¯aL = a µL ( a ) L ( a ) . ( 2 ) Next , the standardized activations are transformed through learned parameters which differ only over a setting of dimensions M . For instance , M = { D } denotes that each feature is transformed by a different scale and bias , which are shared across the sets in the batch and samples in the sets . Let ! M , ! M 2 RB⇥S⇥D denote the learned parameters and represent elementwise multiplication , any transformation operation can be defined as : ˆaM = ¯a ! M + ! M . ( 3 ) We can now summarize a wide array of normalization possibilities simply via choice of L and M . 3.1 NORMALIZATION DESIDERATA . We formally introduce the desiderata for a general-purpose normalization layer for prediction on sets . In particular : ( i ) The normalization must be permutation equivariant . ( ii ) The computation of the statistics for standardization ( µL ( a ) , L ( a ) ) should not be affected by the presence of sets of varying sizes . ( iii ) The normalization should not deteriorate performance due to removal of information that is useful for prediction . Item ( ii ) is meant to address the concern that real-world set inputs are often not all the same size , e.g . not all point cloud objects are made up of the same number of points , and not all patients will have the same number of cells in a blood sample . If the presence of different set sizes forces a choice for an application-specific implementation of a normalization layer , then such a normalization is not general-purpose . Item ( iii ) refers to the fact that the standardization operation in a normalization layer incurs a loss of information , as activations which were once of varying means and variances across certain dimensions are now forced to be zero mean and unit variance . For real-valued sets in particular , this can greatly limit predictive performance , e.g . if the removed means and variances would help in the output prediction . As an example , the mean cell volume of a blood sample can be indicative of patient health , and a normalization layer that makes it impossible to distinguish between patients ’ mean cell volumes will restrict performance across a wide variety of health outcomes predictions . 3.2 THE ONLY SUITABLE TRANSFORMATION FOR SETS IS ON THE FEATURE DIMENSION .. Only normalization layers adopting a per feature transformation , i.e . M = { D } , meet the permutation equivariance desiderata ( i ) , as per sample transformations break permutation equivariance and per set transformations cause the prediction output to change when an element ’ s position in the batch changes ( see Appendix B for proofs ) . Proposition 1 . Per feature transformation is the only setting which maintains permutation equivariance and prediction behavior which is agnostic to batch position . Thus , to systematically compare normalizations , we only need to consider the settings for L . 3.3 EXISTING NORMALIZATIONS DO NOT MEET ALL DESIDERATA .. We now consider possible settings for standardization : { } , { B } , { S } , { D } , { B , S } , { B , D } , { S , D } . We first note that L = { S } and L = { S , D } destroy permutation equivariance ( see Appendix B , for proof ) . Of the remaining settings , we note that two already exist in the literature . The first , which we call feature norm , is used in the permutation invariant PointNet and PointNet++ architectures for point cloud classification ( Qi et al. , 2017a ; b ) and performs a per feature standardization followed by a per feature transformation , i.e . L = { D } , M = { D } ( see Figure 2 ( i ) ) . The second is layer norm , used in Set Transformer , which applies standardization individually to each set and each sample and transformation separately per feature , i.e . L = { B , S } , M = { D } ( see Figure 2 ( ii ) ) . Both these normalization layers have drawbacks . Feature norm calculates statistics over the batch and consequently may exhibit batch effects : sets in the batch are no longer independent , and performance can be heavily impacted by batch size ( Hoffer et al. , 2017 ) . Additionally , like in batch norm , in feature norm sets are standardized differently at training and test time.1 Moreover , in the context of sets of varying sizes , each possible implementation of feature norm ( e.g . computing statistics over the batch , averaging statistics computed on each set separately ) will lead to each set ( or samples in the set ) having a different impact in the computation of the statistics . Therefore , the use of feature norm requires careful consideration , especially in applications having as input real-valued sets with wide ranges of both set sizes and first moments . A similar property holds for normalizations which calculate statistics over the batch ( B 62 L ) such as L = { } . Layer norm , on the other hand , removes information which may be particularly useful for set prediction tasks . By performing standardization on each sample in each set separately , layer norm maps two samples whose activations differ in only a scale and bias to the same output . The inability to distinguish such samples after layer norm can hurt prediction , as we see in Section 5.3 . Proposition 2 . Applying layer norm in its most common placement ( after linear projection , before non-linearity ) removes mean and variance information from each sample . The use of layer norm can present a representation issue similar to the oversmoothing problem in Graph Convolutional Networks ( Li et al. , 2018 ) , where nodes ( samples ) get mapped to the same value as the number of layers in the network increase and neighborhoods of nodes become similar in representation . 1Feature norm is implemented as batch norm on a transposed tensor . Thus , test batches are standardized by a running mean and variance while training batches are standardized by their own batch statistics .

Summary: The paper aims to improve training of deep neural networks applied to sets. The authors investigate permutation equivariant normalizations and residual connections in DeepSet and SetTransformer models, and propose a modification of layer norm which calculates standardization statistics over both features and all set elements. They also propose a modification to DeepSet and SetTransformer models inspired from pre-norm ResNet and pre-norm Transformer. The modifications are experimentally tested on a pointcloud classification dataset, synthetic variance regression, and a new dataset for set prediction contributed by the authors.