[论文解读] BasisVAE: Translation-invariant feature-level clustering with Variational Autoencoders

Kaspar Märtens, Christopher Yau|arXiv (Cornell University)|Mar 6, 2020

Natural Language Processing Techniques参考文献 20被引用 8

一句话总结

BasisVAE 提出了一种联合概率框架，将变分自编码器（VAEs）与解码器中的贝叶斯混合先验相结合，以同时实现非线性降维和特征级别的聚类。通过将特征建模为学习到的基函数的线性组合，并引入具有聚类特异性的缩放和偏移参数，该方法实现了对高维表格数据的可解释性、平移不变聚类——在单细胞基因表达数据上表现尤为出色，其在揭示具有生物学意义的基因表达模式方面优于标准聚类方法和 VAE。

ABSTRACT

Variational Autoencoders (VAEs) provide a flexible and scalable framework for non-linear dimensionality reduction. However, in application domains such as genomics where data sets are typically tabular and high-dimensional, a black-box approach to dimensionality reduction does not provide sufficient insights. Common data analysis workflows additionally use clustering techniques to identify groups of similar features. This usually leads to a two-stage process, however, it would be desirable to construct a joint modelling framework for simultaneous dimensionality reduction and clustering of features. In this paper, we propose to achieve this through the BasisVAE: a combination of the VAE and a probabilistic clustering prior, which lets us learn a one-hot basis function representation as part of the decoder network. Furthermore, for scenarios where not all features are aligned, we develop an extension to handle translation-invariant basis functions. We show how a collapsed variational inference scheme leads to scalable and efficient inference for BasisVAE, demonstrated on various toy examples as well as on single-cell gene expression data.

研究动机与目标

• 为解决标准 VAE 在提供高维表格数据（如基因组学）中可解释的、聚类级别的洞察方面的局限性。
• 将降维与特征聚类统一到单一概率模型中，避免使用临时的两阶段工作流程。
• 通过显式建模具有相似形状但不同偏移的特征的平移参数（δ），实现平移不变聚类。
• 为识别具有共享功能动态的特征组提供一种可扩展且可解释的框架，尤其适用于单细胞组学数据。
• 使研究人员不仅能发现形状相似性，还能通过推断参数（λ 和 δ）揭示特征之间的相对时序或相位偏移。

提出的方法

• 引入一种改进的 VAE 解码器，采用基于类别潜变量的分层贝叶斯聚类先验，以将特征分配至聚类。
• 将每个特征的响应建模为基函数的线性组合，其中包含聚类特异性的缩放参数（λ）和位移参数（δ）。
• 采用折叠变分推断方案，即使在过度指定聚类数量 K 的情况下，也能实现高效且稀疏的推断。
• 在解码器中使用零膨胀负二项分布似然，以建模非负、零膨胀的单细胞基因表达数据。
• 通过允许聚类特异性位移参数（δ）扩展模型，以实现平移不变性，从而将具有相似形状但不同相位偏移的特征分组。
• 将模型应用于真实世界数据，包括合成基因表达数据和单细胞精子发生数据，以展示对潜在结构和特征分组的联合发现。

实验结果

研究问题

• RQ1 能否在 VAE 框架内构建一个联合的降维与特征聚类模型，以提升高维表格数据中的可解释性？
• RQ2 如何在深度生成框架中显式建模平移不变性，以将具有相似功能形状但不同相位偏移的特征分组？
• RQ3 在 VAE 解码器中引入概率聚类先验是否能相比后处理聚类方法，带来更准确且稀疏的聚类分配？
• RQ4 BasisVAE 在不依赖伪时间或聚类结构先验知识的情况下，能在多大程度上恢复单细胞数据中已知的生物基因表达模式？
• RQ5 推断出的缩放参数（λ）和位移参数（δ）是否能为基因表达的相对时序或动态提供具有生物学意义的见解？

主要发现

• BasisVAE 在合成基因表达数据中成功识别出五个不同的聚类，其中具有平移不变性的变体正确地将具有相同形状但不同相位的特征分组。
• 在单细胞精子发生数据中，具有平移不变性的 BasisVAE 将 Tex101、Ly6k、Sdc4、Tubb5 和 Ccnd2 等基因分配至同一聚类，揭示了尽管存在偏移，其表达动态仍具共享性。
• 模型推断出的 δ 参数使得能够解释同一聚类中基因之间的相对时序差异，例如某基因相对于另一基因的表达延迟。
• 具有尺度不变性的 BasisVAE 将形状相同但偏移不同的基因分配至不同聚类，凸显了平移不变性在准确生物分组中的重要性。
• 折叠变分推断实现了高效且稀疏的推断，即使在过度指定聚类数（K=50）的情况下，仍能获得有意义且稳定的聚类分配。
• BasisVAE 在揭示合成数据和真实单细胞数据中真实潜在结构方面优于 k-means 和基于相关性的相似性度量，如通过特征和潜在变量的重排所证实。