[论文解读] Set Representation Learning with Generalized Sliced-Wasserstein Embeddings

Navid Naderializadeh, Soheil Kolouri+3 more|arXiv (Cornell University)|Mar 5, 2021

Domain Adaptation and Few-Shot Learning参考文献 42被引用 5

一句话总结

本文提出广义切片-Wasserstein嵌入（GSWE），一种可几何解释的集合表示学习框架，将集合元素视为概率测度的样本，并利用广义切片-Wasserstein距离的精确欧几里得嵌入来学习不变表示。与基于注意力机制的方法（如Set Transformers）相比，GSWE在参数更少的情况下，在监督和无监督集合学习任务中均实现了最先进性能。

ABSTRACT

An increasing number of machine learning tasks deal with learning representations from set-structured data. Solutions to these problems involve the composition of permutation-equivariant modules (e.g., self-attention, or individual processing via feed-forward neural networks) and permutation-invariant modules (e.g., global average pooling, or pooling by multi-head attention). In this paper, we propose a geometrically-interpretable framework for learning representations from set-structured data, which is rooted in the optimal mass transportation problem. In particular, we treat elements of a set as samples from a probability measure and propose an exact Euclidean embedding for Generalized Sliced Wasserstein (GSW) distances to learn from set-structured data effectively. We evaluate our proposed framework on multiple supervised and unsupervised set learning tasks and demonstrate its superiority over state-of-the-art set representation learning approaches.

研究动机与目标

• 开发一种可几何解释的框架，用于从无序集合结构数据中学习表示。
• 通过利用最优传输理论，解决现有池化机制在集合学习中的局限性。
• 通过广义切片-Wasserstein距离的精确欧几里得嵌入，实现高效且参数高效的集合嵌入。
• 在监督和无监督集合学习基准上，相比最先进方法，实现更优性能。
• 通过在池化模块中避免多头注意力机制，降低模型复杂度，同时保持或提升准确率。

提出的方法

• 该方法将每个集合建模为从其元素导出的概率测度，将这些元素视为来自某一分布的样本。
• 采用广义切片-Wasserstein（GSW）距离作为集合与参考集合之间概率测度的几何距离度量。
• 为GSW距离推导出精确的欧几里得嵌入，从而实现在固定维空间中的高效且可微分的优化。
• 参考集合元素为可微分参数，在训练过程中端到端学习，使模型能够自适应地学习有意义的参考分布。
• 主干网络通过置换等变函数（如前馈网络）处理每个集合元素，随后使用学习到的参考集合进行基于GSW的池化。
• 该框架支持单个和多个参考集合，实现灵活且可扩展的集合表示学习。

实验结果

研究问题

• RQ1 基于最优传输的可几何解释的集合嵌入方法，是否能优于现有池化机制在集合表示学习中的表现？
• RQ2 在准确率和参数效率方面，GSWE与基于注意力的模型（如Set Transformers）相比表现如何？
• RQ3 为GSW距离使用精确欧几里得嵌入是否能提升在多样化集合学习任务中的泛化能力和鲁棒性？
• RQ4 与RepSet等方法相比，GSWE对参考集合数量和参考集合基数的敏感性如何？
• RQ5 GSWE是否能在监督和无监督集合学习中均实现最先进性能，且参数量显著少于基于注意力的基线模型？

主要发现

• 在Set-Circles和Set-Boxes数据集上，GSWE在最近邻检索准确率方面优于Set Transformer（PMA），尤其当参考集合包含多个元素时表现更优。
• 在监督文本分类任务中，GSWE在所有四个数据集（twitter、bbcsport、classic、amazon）上均达到与PMA相当或更高的峰值5折交叉验证准确率，且参数量更少。
• GSWE在不同数量的参考集合下均保持稳健性能，而RepSet在参考集合较少时性能显著下降，凸显GSWE的稳定性和高效性。
• 当参考集合元素多于一个时，GSWE相较于PMA的性能提升最为显著，表明GSWE能更好地捕捉集合中的结构关系。
• GSWE在监督和无监督集合学习任务中均实现了最先进性能，证明了其有效性与可扩展性。
• GSWE参数量的减少（源于无需多头注意力机制）促使其在效率和性能上优于基于注意力的基线模型。