[Paper Review] Set Norm and Equivariant Skip Connections: Putting the Deep in Deep Sets

Lily H. Zhang, Veronica Tozzo+2 more|arXiv (Cornell University)|Jun 23, 2022

Explainable Artificial Intelligence (XAI)4 citations

TL;DR

This paper introduces Deep Sets++ and Set Transformer++, enhanced versions of Deep Sets and Set Transformer with set norm (sn) and clean-path equivariant skip connections to overcome vanishing/exploding gradients in deep architectures. These models achieve state-of-the-art performance on diverse tasks, including point cloud classification and hematocrit prediction, and are the first deep permutation-invariant networks to maintain stability and accuracy at high depths without task-specific tuning.

ABSTRACT

Permutation invariant neural networks are a promising tool for making predictions from sets. However, we show that existing permutation invariant architectures, Deep Sets and Set Transformer, can suffer from vanishing or exploding gradients when they are deep. Additionally, layer norm, the normalization of choice in Set Transformer, can hurt performance by removing information useful for prediction. To address these issues, we introduce the clean path principle for equivariant residual connections and develop set norm, a normalization tailored for sets. With these, we build Deep Sets++ and Set Transformer++, models that reach high depths with comparable or better performance than their original counterparts on a diverse suite of tasks. We additionally introduce Flow-RBC, a new single-cell dataset and real-world application of permutation invariant prediction. We open-source our data and code here: https://github.com/rajesh-lab/deep_permutation_invariant.

Motivation & Objective

• Address the instability of Deep Sets and Set Transformer when deepened, which leads to vanishing or exploding gradients.
• Overcome the performance degradation caused by layer normalization in real-valued set tasks, where it induces unwanted invariance to scalar transformations.
• Develop a general-purpose, deep permutation-invariant architecture that maintains high performance across diverse tasks without task-specific engineering.
• Introduce a new real-world dataset, Flow-RBC, for benchmarking permutation-invariant models in clinical single-cell prediction.
• Establish a principled design for residual connections and normalization tailored for sets, enabling reliable training at high depths.

Proposed method

• Propose the 'clean path principle' for equivariant residual connections, ensuring identity mappings are preserved and gradients flow stably through deep networks.
• Introduce set norm (sn), a normalization layer that standardizes each set over the minimal number of dimensions, preserving useful information while stabilizing training.
• Design Deep Sets++ and Set Transformer++ by integrating clean-path skip connections and set norm into the original architectures, maintaining permutation invariance.
• Apply set norm to standardize sets by computing mean and standard deviation across elements, avoiding the oversmoothing and invariance issues of layer norm.
• Use a residual connection scheme where the skip path is strictly identity-mapped, ensuring gradient stability and improved optimization in deep networks.
• Train models end-to-end on diverse tasks, including point cloud classification and hematocrit prediction, to validate performance and robustness at high depths.

Experimental results

Research questions

• RQ1 Can Deep Sets and Set Transformer be effectively deepened without suffering from vanishing or exploding gradients?
• RQ2 Does layer normalization degrade performance in real-valued set prediction tasks due to unwanted invariance to scalar transformations?
• RQ3 Can a new normalization layer—set norm—improve training stability and performance in deep permutation-invariant networks?
• RQ4 Do clean-path equivariant skip connections enhance gradient flow and model performance in deep architectures?
• RQ5 Can general-purpose deep permutation-invariant models outperform task-specific variants on benchmark tasks like point cloud classification?

Key findings

• Deep Sets++ and Set Transformer++ achieve superior or comparable performance to their original counterparts across a diverse suite of tasks, including point cloud classification and hematocrit prediction.
• At 50 layers, Deep Sets++ and Set Transformer++ outperform the original Deep Sets and Set Transformer, which degrade in performance when deepened.
• On point cloud classification, the deep versions of DS++ and ST++ surpass the performance of task-specific architectures proposed in the original Deep Sets and Set Transformer papers.
• On the Flow-RBC dataset, both DS++ and ST++ outperform the clinical baseline (25.85 MSE), while the original Deep Sets at 50 layers fails to do so.
• Set norm effectively prevents the unwanted invariance induced by layer norm, particularly in real-valued set inputs, leading to improved performance on non-scalar-invariant tasks.
• Despite having approximately three times more parameters, Set Transformer++ is not significantly faster than Deep Sets++; however, Deep Sets++ trains roughly twice as fast on the same hardware, indicating better efficiency.