# Probabilistic Transformer: Modelling Ambiguities and Distributions for RNA Folding and Molecule Design

arXiv: https://arxiv.org/abs/2205.13927 · HTML: https://mlsanity.com/paper/2205.13927/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 129 references that are in the map.
- First posted on arXiv 27 May 2022.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Learning Deep Transformer Models for Machine Translation](https://mlsanity.com/paper/1906.01787/) (2019)
- [GuacaMol: Benchmarking Models for De Novo Molecular Design](https://mlsanity.com/paper/1811.09621/) (2018)
- [Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models](https://mlsanity.com/paper/1811.12823/) (2018)
- [Molecular De Novo Design through Deep Reinforcement Learning](https://mlsanity.com/paper/1704.07555/) (2017)
- [Molecular Transformer - A Model for Uncertainty-Calibrated Chemical Reaction Prediction](https://mlsanity.com/paper/1811.02633/) (2018)
- [Constrained Graph Variational Autoencoders for Molecule Design](https://mlsanity.com/paper/1805.09076/) (2018)
- [Deep Reinforcement Learning for De-Novo Drug Design](https://mlsanity.com/paper/1711.10907/) (2017)
- [Objective-Reinforced Generative Adversarial Networks (ORGAN) for Sequence Generation Models](https://mlsanity.com/paper/1705.10843/) (2017)
- [Taming VAEs](https://mlsanity.com/paper/1810.00597/) (2018)
- [A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities](https://mlsanity.com/paper/1905.13077/) (2019)
- [Transformer-based Conditional Variational Autoencoder for Controllable Story Generation](https://mlsanity.com/paper/2101.00828/) (2021)
- [Variational Transformers for Diverse Response Generation](https://mlsanity.com/paper/2003.12738/) (2020)
- [RNA Secondary Structure Prediction By Learning Unrolled Algorithms](https://mlsanity.com/paper/2002.05810/) (2020)
- [Transformer Uncertainty Estimation with Hierarchical Stochastic Attention](https://mlsanity.com/paper/2112.13776/) (2021)
- [Molecular Generative Model Based On Adversarially Regularized Autoencoder](https://mlsanity.com/paper/1912.05617/) (2019)

## Built on by (most cited first, 7 of 7)

- [Beyond Sequence: Impact of Geometric Context for RNA Property Prediction](https://mlsanity.com/paper/2410.11933/) (2024)
- [Improving Deep Learning Optimization through Constrained Parameter Regularization](https://mlsanity.com/paper/2311.09058/) (2023)
- [Scalable Deep Learning for RNA Secondary Structure Prediction](https://mlsanity.com/paper/2307.10073/) (2023)
- [Deciphering RNA Secondary Structure Prediction: A Probabilistic K-Rook Matching Perspective](https://mlsanity.com/paper/2212.14041/) (2022)
- [Rethinking Performance Measures of RNA Secondary Structure Problems](https://mlsanity.com/paper/2401.05351/) (2023)
- [M3-20M: A Large-Scale Multi-Modal Molecule Dataset for AI-driven Drug Design and Discovery](https://mlsanity.com/paper/2412.06847/) (2024)
- [Memory and Bandwidth are All You Need for Fully Sharded Data Parallel](https://mlsanity.com/paper/2504.03655/) (2025)

## Abstract (arXiv)

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.