A Variational Bayesian Method for Sequence Models

Dario Coscia

International School of Advanced Studies,

University of Amsterdam

Scientific Discovery: from Experiment to Simulation

1

Galileo Galilei

- inventor of scientific method 15^th/16^th century

“Scientific method process involves observations, forming a hypotheses, making predictions, conducting an experiment and finally analyzing the results"

Scientific Discovery: from Simulation to Emulation

2

Neural Emulators: AI driven Emulators

3

CFD Neural Emulator

CFD Simulator

Physics Informed Neural Emulator

Coscia, Dario, et al. "Physics-Informed Neural networks for Advanced modeling." 

Neural Emulator are reaching standard simulator performances at much lower inference cost and higher generalizability

Sequence Models: A Brief Overview

4

Radford, Alec. "Improving language understanding by generative pre-training."

Van Den Oord, Aaron, et al. "WaveNet: A generative model for raw audio."

Keller, T. Anderson, et al. "Traveling waves encode the recent past and enhance sequence learning."

Sequence models are a class of machine learning models designed to process and generate data that is inherently sequential, e.g. text or audio

Sequence Models: A Brief Overview

5

Sequence models predict sequences of data by modelling the probability distribution of each element given previous elements, capturing temporal dependencies and uncertainty

Autoregressive Neural Emulators in AI4Science

6

Autoregressive models can be used to build Neural Emulators to solve PDEs or perform weather forecast by iteratively applying the emulator on its output given an initial state

Bodnar, Cristian, et al. "Aurora: A foundation model of the atmosphere"

Recurrent Neural Emulators in AI4Science

7

Recurrent Neural Networks (RNNs) can be used to discover novel Molecules by generating chemical structures as sequence of tokens, Chemical Language Models (CLMs)

Özçelik, Rıza, et al. "Chemical language modeling with structured state space sequence models."

​

Can we Trust Sequence Models?

8

Large Language Models Hallucination

Diverging Solutions between Emulator and PDE Solver

time

diverging

Uncertainty Quantification

9

Uncertainty Quantification aims to represent and measure the model’s confidence in its predictions

Aleatoric Uncertainty

Epistemic Uncertainty

Bayesian Modelling and Uncertainty Decomposition

10

Bayes’ theorem updates the prior information of model parameters based on the collected evidence and likelihood

The predictions are obtained by model averaging and latent’s marginalization

Uncertainties can be decomposed

Bayesian Deep Learning – UQ and Search

11

Bayesian Models by construction model epistemic uncertainty by assigning probabilistic weights and update them using Bayes’ rule

Bayesian methods are key also for search, e.g. molecular design/ language models alignment

1. Korbak, Tomasz, et al. "RL with KL penalties is better viewed as Bayesian inference."

2. Papamarkou, Theodore, et al. "Position: Bayesian deep learning is needed in the age of large-scale AI."

Variational Bayes for Bayesian Deep Learning

12

Variational Inference optimizes the parameters of some parameterized posterior, e.g. Neural Networks, such that it is a close approximation of the true intractable posterior

Variational Dropout and Local Reparametrization Trick: Scaling to Large Networks

13

Variational Dropout generalizes the concept of Dropout to learnable dropout rates

The Local Reparameterization trick translates uncertainty about global parameters into local noise that is independent across datapoints in the minibatch

Sampling activation allows scaling

Bayesianize Sequence Models: BARNN

14

BARNN is a practical, scalable, calibrated and accurate methodology to turn any autoregressive or recurrent model into its Bayesian version

Joint states-weight Sequence Model

15

Jointly sample in an alternating way weights and states

Variational Bayes Training

16

Optimize via Variational Inference the Evidence Lower BOund

Temporal Variational Dropout and Temporal VAMP

17

The posterior is obtained by extending Variational Dropout to variable in time dropout rates

The best prior maximazing the Temporal ELBO is given by the aggregated posterior in time

The optimizable weights are decoupled

=

scalability

The KL is independent on Model Weights!

Experiments

18

1. Does the tVAMP prior outperforms ”static” common priors?

​

• Can the BARNN model produce well-calibrated and sharp uncertainty estimates while maintaining the same level of accuracy as a non-Bayesian model?

​

• Is BARNN generally applicable and scalable?

Time Series

PDEs

Molecules

AI4Science Applications!

BARNN for TimeSeries: tVAMP prior

19

BARNN outperforms standard MLP autoregression and Ensemble Dropout in both prediction and confidence estimation

​

The tVAMP prior empirically outperforms the commonly used log-uniform prior

BARNN in AI4Science: PDEs and Molecules

20

Choosing a gaussian state update with small variance is equivalent to update with a standard Autoregressive PDE Solvers

BARNN solves PDEs with Calibrated Uncertainties

21

SOTA results on solving PDEs

Having a small root mean-square error (RMSE) between CFD solution and Neural Emulator solution does not mean to have calibrated uncertainties! BARNN achieves both

BARNN dynamically adapts uncertainties over time

22

MAP provides convergent metrics when uncertainties are not needed

Confidence intervals are tighter in regions of small uncertainties and larger otherwise in the Kuramoto Sivashinsky PDE

BARNN in AI4Science: Molecules

23

Molecules can be represented as text using SMILES and a Language Model can be used to generate new molecules

with,

BARNN excels in molecule generation and long-range dependency modelling

24

The molecules generated by BARNN are statistically richer

BARNN generates more complex molecules (many rings) showing great long-dependency modelling abilities

BARNN excels in molecular properties prediction

25

The molecules generated by BARNN better match the training dataset properties, showing higher density matching ability

Conclusions

26

BARNN is a Bayesian Sequential model. It performs UQ and achieves state-of-the-art predictive results while being scalable to large nets.

AI4Science highlights the necessity of Uncertainty Quantification, without UQ surrogate models can not be safely applied

Extending these approaches to LLMs will be impactful for chatbots, allowing them to recognize when they lack knowledge

References

27

Backup Slides

28

Predictive Distribution and Uncertainty Estimates

29

The predictive distribution is obtained by marginalizing over the weights for each time-steps

Uncertainties can be decomposed at each time-steps

Monte Carlo sampling can be used to obtain moments estimates

KL divergence in Temporal ELBO

30

Given the posterior, the tVAMP prior is obtained by:

The KL is independent in each time-steps on the static weights: