[論文レビュー] 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

ひとこと要約

本稿では、条件付き変分オートエンコーダーの原理を組み込んだ階層的潜在分布拡張による確率的トランスフォーマー（Probabilistic 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.

研究の動機と目的

• RNA二次構造や分子性質といった曇りやノイズ、多モードな生物学的データを扱う際、決定論的トランスフォーマーの限界を克服すること。
• トランスフォーマーの潜在空間における確率的推論を可能にし、ソフトマックス出力確率を超えた不確実性の表現を可能にすること。
• 階層的変分推論を用いて潜在的なデータ分布を学習することで、系列データタスクにおける生成モデルの性能を向上させること。
• 実世界の生物学的データを用いた判別的（RNA折りたたみ）および生成的（分子設計）設定の両方で、モデルの有効性を示すこと。

提案手法

• トランスフォーマーのアテンションおよびフィードフォワードブロックの直後に挿入される、階層的潜在分布モデリングを可能にする新規な確率的レイヤーを導入。
• 階層的U-Netにインspiredされた階層的構造を持つ条件付き変分オートエンコーダー（cVAE）フレームワークを採用し、複数レベルの不確実性表現を可能に。
• 一般化された下界（ELBO）と制約付き最適化（GECO）を用いた学習法を採用し、変分推論の改善を図るためのオンラインハイパーパramータアニーリング技術を統合。
• 推論時に潜在空間でのサンプリングを実施し、計算グラフ内でのサンプルベクトルの最適化を可能にし、複雑な分布のモデル化を実現。
• 自己アテンションによるグローバルな受容 field を維持しつつ、既存のトランスフォーマー部品および学習パイプラインと互換性を保つ。
• 標準的なトランスフォーマーと同一の残差接続およびレイヤーナーミャライゼーションスキームを採用し、アーキテクチャの安定性を確保。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーのアーキテクチャに階層的潜在分布を導入することで、系列データにおける複雑で多モードなデータ分布を効果的に学習・回復できるか？
• RQ2 データ自体に構造的曇りや測定ノイズが内在する場合、Probabilistic TransformerはRNA折りたたみタスクでどの程度の性能を示すか？
• RQ3 このモデルは、そのような分子の潜在分布を暗黙的に学習することで、複数の望ましい化学的性質を同時に満たす分子をどの程度効果的に生成できるか？
• RQ4 標準的なトランスフォーマーと比較して、確率的潜在空間の統合が曇りやノイズのある生物学的データに対する一般化性能およびロバストネスを向上させるか？
• RQ5 提案されたGECO用のオンラインハイパーパramータアニーリング技術は、変分オートエンコーディング設定における学習安定性および性能にどのように影響を与えるか？

主な発見

• Probabilistic Transformerは、特に曇りの強いデータセットで標準的なトランスフォーマーが苦戦する状況においても、RNA折りたたみベンチマークで最先端の性能を達成した。
• 合成された系列データ分布タスクにおいて、モデルは潜在的な多モードデータ分布を効果的に学習・回復し、分布モデリングの能力を示した。
• 性質に基づく分子生成タスクにおいて、MolGPTよりもはるかに多くの固有の分子（例：TPSA、LogP、SAS）をターゲット性質に適合させる分子を生成した。1つのターゲットセットに対して最大430.8個の固有分子を生成した。
• モデルは、多様で化学的に妥当な分子を生成し、さまざまなターゲット性質の組み合わせについての例示予測でも高い性質精度を示した。
• アブレーションスタディの結果、確率的レイヤーが性能向上に寄与していることが確認され、提案されたカッパアニーリング技術が学習安定性とELBO最適化を改善した。
• 判別的および生成的両設定で優れた結果を達成しており、不確実性のモデル化と多様で高品質な出力の生成という二重の能力を確認した。