[论文解读] 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

一句话总结

本文提出了 Deep Sets++ 和 Set Transformer++，作为 Deep Sets 与 Set Transformer 的增强版本，通过引入集合归一化（set norm, sn）和干净路径等变跳跃连接，以克服深层架构中的梯度消失/爆炸问题。这些模型在多种任务上实现了最先进性能，包括点云分类与血细胞比容预测，且是首个在不进行任务特定调优的情况下，仍能保持高深度下稳定性和准确性的深度排列不变神经网络。

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.

研究动机与目标

• 解决 Deep Sets 与 Set Transformer 在加深时出现的不稳定性问题，该问题会导致梯度消失或爆炸。
• 克服层归一化在实值集合任务中导致的性能下降，因其会引入对标量变换的非期望不变性。
• 开发一种通用的、深度排列不变的架构，使其在无需任务特定工程的情况下，可在多样化任务中保持高性能。
• 引入一个新的真实世界数据集 Flow-RBC，用于在临床单细胞预测任务中对排列不变模型进行基准测试。
• 建立针对集合设计的合理残差连接与归一化方案，以支持在高深度下的可靠训练。

提出的方法

• 提出‘干净路径原则’用于等变残差连接，确保恒等映射被保留，使梯度在深层网络中稳定流动。
• 引入集合归一化（sn），一种归一化层，其对每个集合在最小维度数上进行标准化，保留有用信息的同时稳定训练过程。
• 通过将干净路径跳跃连接与集合归一化整合到原始架构中，设计出 Deep Sets++ 与 Set Transformer++，同时保持排列不变性。
• 应用集合归一化通过计算元素间的均值与标准差来标准化集合，避免层归一化带来的过度平滑与不变性问题。
• 采用残差连接方案，其中跳跃路径严格为恒等映射，以确保梯度稳定性并提升深层网络的优化性能。
• 在包括点云分类与血细胞比容预测在内的多样化任务上进行端到端训练，以验证模型在高深度下的性能与鲁棒性。

实验结果

研究问题

• RQ1 Deep Sets 与 Set Transformer 是否能有效加深，而不会遭受梯度消失或爆炸问题？
• RQ2 层归一化是否因对标量变换的非期望不变性，而在实值集合预测任务中导致性能下降？
• RQ3 新的归一化层——集合归一化——是否能提升深层排列不变网络的训练稳定性和性能？
• RQ4 干净路径等变跳跃连接是否能增强深层架构中的梯度流动与模型性能？
• RQ5 通用型深层排列不变模型是否能在点云分类等基准任务上超越任务特定变体？

主要发现

• Deep Sets++ 与 Set Transformer++ 在包括点云分类与血细胞比容预测在内的多样化任务中，性能优于或与原始模型相当。
• 在 50 层深度下，Deep Sets++ 与 Set Transformer++ 的性能优于原始的 Deep Sets 与 Set Transformer，后者在加深时性能明显下降。
• 在点云分类任务中，DS++ 与 ST++ 的深层版本性能超越了原始 Deep Sets 与 Set Transformer 论文中提出的任务特定架构。
• 在 Flow-RBC 数据集上，DS++ 与 ST++ 均优于临床基线（25.85 MSE），而原始 Deep Sets 在 50 层时未能实现超越。
• 集合归一化有效防止了层归一化在实值集合输入中引入的非期望不变性，尤其在非标量不变任务中显著提升了性能。
• 尽管参数量约为三倍，Set Transformer++ 并未显著快于 Deep Sets++；然而，Deep Sets++ 在相同硬件上训练速度约为其两倍，表明其具有更高的效率。