[论文解读] Probabilistic Transformer: Modelling Ambiguities and Distributions for RNA Folding and Molecule Design

Jörg K. H. Franke, Frederic Runge+1 more|arXiv (Cornell University)|May 27, 2022

RNA and protein synthesis mechanisms被引用 4

一句话总结

本文提出了概率Transformer（Probabilistic Transformer），这是一种对Transformer架构的分层潜在分布增强方法，通过引入条件变分自编码原理，对数据中的模糊性和不确定性进行建模。该方法在模糊数据集上的RNA折叠任务中达到最先进性能，并在基于属性的分子生成任务中优于现有方法，通过隐式学习复杂数据分布实现性能提升。

ABSTRACT

Our world is ambiguous and this is reflected in the data we use to train our algorithms. This is particularly true when we try to model natural processes where collected data is affected by noisy measurements and differences in measurement techniques. Sometimes, the process itself is ambiguous, such as in the case of RNA folding, where the same nucleotide sequence can fold into different structures. This suggests that a predictive model should have similar probabilistic characteristics to match the data it models. Therefore, we propose a hierarchical latent distribution to enhance one of the most successful deep learning models, the Transformer, to accommodate ambiguities and data distributions. We show the benefits of our approach (1) on a synthetic task that captures the ability to learn a hidden data distribution, (2) with state-of-the-art results in RNA folding that reveal advantages on highly ambiguous data, and (3) demonstrating its generative capabilities on property-based molecule design by implicitly learning the underlying distributions and outperforming existing work.

研究动机与目标

• 解决确定性Transformer在建模模糊、噪声或多模态生物数据（如RNA二级结构和分子性质）方面的局限性。
• 在Transformer的潜在空间中实现概率推理，允许表示超出Softmax输出概率的不确定性。
• 通过分层变分推理学习底层数据分布，提升序列数据任务中的生成建模性能。
• 在真实世界生物数据的判别性（RNA折叠）和生成性（分子设计）设置中，验证模型的有效性。

提出的方法

• 在Transformer的注意力和前馈网络模块之后引入一种新颖的概率层，实现分层潜在分布建模。
• 采用受分层U-Net启发的分层结构条件变分自编码器（cVAE）框架，支持多层级不确定性表示。
• 使用带约束优化的广义证据下界（GECO）进行训练，结合在线超参数退火技术，提升变分推理性能。
• 在推理过程中对潜在空间进行采样，允许在计算图内对采样向量进行细化，以建模复杂分布。
• 通过自注意力机制保持全局感受野，同时与现有Transformer组件和训练流程保持兼容。
• 采用与标准Transformer一致的残差连接和层归一化方案，确保架构稳定性。

实验结果

研究问题

• RQ1 Transformer架构中的分层潜在分布能否有效学习并恢复序列数据中的复杂、多模态数据分布？
• RQ2 当数据本身存在固有的结构模糊性和测量噪声时，概率Transformer在RNA折叠任务中的表现如何？
• RQ3 该模型在隐式学习目标分子底层分布的基础上，能在多大程度上生成具有多种期望化学性质的分子？
• RQ4 与标准Transformer相比，集成概率潜在空间是否能提升在模糊或噪声生物数据上的泛化能力和鲁棒性？
• RQ5 所提出的GECO在线超参数退火技术在变分自编码设置中对训练稳定性和性能有何影响？

主要发现

• 概率Transformer在RNA折叠基准测试中达到最先进性能，尤其在高度模糊的数据集上表现优异，而标准Transformer则难以应对。
• 在合成序列分布任务中，模型成功学习并恢复了底层多模态数据分布，证明了其在分布建模方面的强大能力。
• 在基于属性的分子生成任务中，概率Transformer生成的符合目标性质（如TPSA、LogP、SAS）的唯一分子数量显著多于MolGPT，单个目标集合下最多可达430.8个唯一分子。
• 模型生成的分子具有多样性且化学结构合理，且属性预测准确率高，如多种目标属性组合下的示例预测所示。
• 消融实验表明，概率层对性能有显著贡献，且所提出的kappa退火技术提升了训练稳定性和ELBO优化效果。
• 该模型在判别性和生成性设置中均取得优越结果，证实其在建模不确定性与生成多样化、高质量输出方面的双重能力。