[论文解读] Unsupervised Learning of Group Invariant and Equivariant Representations

Robin Winter, Marco Bertolini+3 more|arXiv (Cornell University)|Feb 15, 2022

Medical Imaging and Analysis被引用 7

一句话总结

该论文提出了一种无监督自编码器框架，可对任意群 $G$ 的群不变性和群等变表征进行解耦，通过可学习的群作用对齐重建结果。该方法在学习多种数据类型（包括点云、分子和3D形状）的不变表征方面达到了最先进性能，且架构约束极少，在下游任务中表现出强大的泛化能力。

ABSTRACT

Equivariant neural networks, whose hidden features transform according to representations of a group G acting on the data, exhibit training efficiency and an improved generalisation performance. In this work, we extend group invariant and equivariant representation learning to the field of unsupervised deep learning. We propose a general learning strategy based on an encoder-decoder framework in which the latent representation is separated in an invariant term and an equivariant group action component. The key idea is that the network learns to encode and decode data to and from a group-invariant representation by additionally learning to predict the appropriate group action to align input and output pose to solve the reconstruction task. We derive the necessary conditions on the equivariant encoder, and we present a construction valid for any G, both discrete and continuous. We describe explicitly our construction for rotations, translations and permutations. We test the validity and the robustness of our approach in a variety of experiments with diverse data types employing different network architectures.

研究动机与目标

• 开发一种无监督表示学习方法，显式分离数据在群作用下的不变分量与等变分量。
• 解决在数据坐标参数化中未显式体现对称性的情况下，无监督学习群不变表征的挑战。
• 构建一个适用于离散与连续群（包括旋转、平移和排列）的通用框架。
• 通过在自编码器中引入群对称性作为归纳偏置，提升泛化能力和训练效率。
• 使该框架能够应用于分子和3D点云等真实世界数据，这些数据具有高维输入空间。

提出的方法

• 该框架采用编码器-解码器自编码器架构，其中潜在码 $z$ 被分解为群不变分量和群作用分量。
• 群作用函数 $\psi$ 预测适当的群变换 $g \in G$，以对齐重建数据 $\delta(\eta(x), g)$ 与输入 $x$，从而在存在对称性的情况下实现重建。
• 编码器被约束为在 $G$ 下等变，确保 $\eta(g \cdot x) = \rho_G(g) \cdot \eta(x)$，其中 $\rho_G$ 是 $G$ 的表示。
• 解码器通过将预测的群作用 $g$ 应用于标准重建 $\hat{x}$ 来重建输入，使用 $\delta(z, g) = \rho_X(g) \cdot \hat{x}$。
• 该方法具有通用性，适用于任意群 $G$，包括 $SO(3)$、$\mathbb{R}^3$ 和 $S_N$，并对每类群提供了显式构造。
• 该框架兼容标准自编码器与变分自编码器，支持对不变表征进行生成建模。

实验结果

研究问题

• RQ1 我们能否在不依赖数据增强或对称性感知监督的情况下，无监督地学习群不变性和等变表征？
• RQ2 编码器与群作用函数需满足哪些必要的数学条件，才能确保不变分量与等变分量的正确解耦？
• RQ3 如何构建一个适用于任意群 $G$（包括连续与离散群）的通用框架？
• RQ4 该框架能否提升表示学习任务（如形状分类与分子构象重建）中的泛化能力与下游性能？
• RQ5 与标准自编码器相比，该解耦表征在潜在空间的鲁棒性与结构方面表现如何？

主要发现

• 在QM9数据集上，该方法对分子构象的重建均方根误差为 $0.15 \pm 0.07~\text{\AA}$，在5000个未见测试样本上实现了原子类型识别的完美准确率。
• 在ShapeNet数据集上，基于不变嵌入训练的KNN分类器达到81%准确率，而基于非不变嵌入的分类器仅达63%，表明其泛化能力显著更优。
• TSNE可视化显示，不变自编码器生成了结构良好、类别分明的潜在空间，而非不变模型则因方向差异导致聚类杂乱无序。
• 该框架成功学习了在旋转与平移下Tetris方块的解耦表征，潜在空间聚类与形状类别高度匹配。
• 该方法在多种架构与数据类型（包括3D点云、分子构象和2D图像）上均表现出一致的性能提升，具有良好的泛化能力。
• 该框架与变分自编码器完全兼容，支持对不变表征的概率建模，具有生成应用的潜力。