[Paper Review] 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 mechanisms4 citations

TL;DR

This paper introduces the Probabilistic Transformer, a hierarchical latent distribution enhancement to the Transformer architecture that models ambiguity and uncertainty in data by incorporating conditional variational auto-encoding principles. It achieves state-of-the-art performance in RNA folding on ambiguous datasets and outperforms existing methods in property-based molecule generation by implicitly learning complex data distributions.

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.

Motivation & Objective

• To address the limitations of deterministic Transformers in modeling ambiguous, noisy, or multi-modal biological data such as RNA secondary structures and molecular properties.
• To enable probabilistic reasoning in the latent space of Transformers, allowing for uncertainty representation beyond softmax output probabilities.
• To improve generative modeling in sequential data tasks by learning underlying data distributions through hierarchical variational inference.
• To demonstrate the model's effectiveness in both discriminative (RNA folding) and generative (molecule design) settings with real-world biological data.

Proposed method

• Introduces a novel probabilistic layer inserted after attention and feed-forward blocks in the Transformer, enabling hierarchical latent distribution modeling.
• Employs a conditional variational auto-encoder (cVAE) framework with hierarchical structure inspired by the hierarchical U-Net, allowing for multi-level uncertainty representation.
• Uses the generalized evidence lower bound (ELBO) with constrained optimization (GECO) for training, incorporating an online hyperparameter annealing technique for improved variational inference.
• Performs sampling in the latent space during inference, enabling refinement of sampled vectors within the computational graph to model complex distributions.
• Preserves the global receptive field via self-attention while remaining compatible with existing Transformer components and training pipelines.
• Applies a residual connection and layer normalization scheme consistent with standard Transformers, ensuring architectural stability.

Experimental results

Research questions

• RQ1 Can a hierarchical latent distribution in the Transformer architecture effectively learn and recover complex, multi-modal data distributions in sequential data?
• RQ2 How does the Probabilistic Transformer perform on RNA folding tasks when data contains inherent structural ambiguity and measurement noise?
• RQ3 To what extent can the model generate molecules with multiple desired chemical properties by implicitly learning the underlying distribution of such molecules?
• RQ4 Does the integration of probabilistic latent space improve generalization and robustness on ambiguous or noisy biological data compared to standard Transformers?
• RQ5 How does the proposed online hyperparameter annealing technique for GECO affect training stability and performance in variational auto-encoding settings?

Key findings

• The Probabilistic Transformer achieves state-of-the-art performance on RNA folding benchmarks, particularly excelling on highly ambiguous datasets where standard Transformers struggle.
• On the synthetic sequential distribution task, the model successfully learns and recovers the underlying multi-modal data distribution, demonstrating its capacity for distribution modeling.
• In property-based molecule generation, the Probabilistic Transformer generates significantly more unique molecules meeting target properties (e.g., TPSA, LogP, SAS) than MolGPT, with up to 430.8 unique molecules for a single target set.
• The model generates diverse, chemically valid molecules with high property accuracy, as shown in example predictions for various target property combinations.
• Ablation studies confirm that the probabilistic layers contribute meaningfully to performance, and the proposed kappa annealing technique improves training stability and ELBO optimization.
• The model achieves superior results in both discriminative and generative settings, confirming its dual capability in modeling uncertainty and generating diverse, high-quality outputs.