BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Authors: Dario Coscia
Max Welling
Nicola Demo
Gianluigi Rozza
Venue: ICML 2025 poster (2025)
Keywords: Bayesian Modelling
Variational Inference
Autoregressive Neural Networks
Recurrent Neural Networks
Uncertainty Quantification
Scientific Machine Learning
Primary area: probabilistic_methods->bayesian_models_and_methods
TL;DR: BARNN is a Bayesian framework that converts any autoregressive or recurrent neural network into a Bayesian version, improving performance and uncertainty quantification with minimal model changes

Abstract: Autoregressive and recurrent networks have achieved remarkable progress across various fields, from weather forecasting to molecular generation and Large Language Models. Despite their strong predictive capabilities, these models lack a rigorous framework for addressing uncertainty, which is key in scientific applications such as PDE solving, molecular generation and machine l earning Force Fields.
To address this shortcoming we present BARNN: a variational Bayesian Autoregressive and Recurrent Neural Network. BARNNs aim to provide a principled way to turn any autoregressive or recurrent model into its Bayesian version. BARNN is based on the variational dropout method, allowing to apply it to large recurrent neural networks as well. We also introduce a temporal version of the
“Variational Mixtures of Posteriors” prior (tVAMP-prior) to make Bayesian inference efficient and well-calibrated. Extensive experiments on PDE modelling and molecular generation demonstrate that BARNN not only achieves comparable or superior accuracy compared to existing methods, but also excels in uncertainty quantification and modelling long-range dependencies.

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Dario Coscia 12
Max Welling?
Nicola Demo 3 Gianluigi Rozza
Abstract
on, challenging their application outside their training do- main (Papamarkou al., 2024). Overcoming this behaviour is critical, not only in scientific applications where phenom- ena can show complex data distribution shifts away from the training data, but also in deep over-parametrized mod- cls. For example, in weather forecasting, rapidly changing climate patterns can cause significant deviations from the training data and lead unpredictable results, while over- parametrized LLMs often provide incorrect answers with high confidence, highlighting issues with model calibra- tion (Jiang et al., 2021; Xiao et al., 2022; Yang et al., 2022). This problem is also prevalent across many other domains, and Bayesian approaches present a promising direction for improvement (Papamarkou et al., 2024).
Autoregressive and recurrent networks have achieved remarkable progress across various fields, from weather forecasting to molecular gen- eration and Large Language Models. Despite their strong predictive capabilities, these mod- els lack a rigorous framework for addressing un- certainty, which is key in scientific applications such as PDE solving, molecular generation and Machine Learning Force Fields. To address this shortcoming present BARNN: a variational Bayesian Autoregressive and Recurrent Neural Network. BARNNs aim to provide principled way to turn any autoregressive or recurrent model into its Bayesian version. BARNN is based on the variational dropout method, allowing to apply it to large recurrent neural networks as well. We also introduce temporal version of the "Vari- ational Mixtures of Posteriors" prior (tVAMP- prior) to make Bayesian inference efficient and well-calibrated. Extensive experiments on PDE modelling and molecular generation demonstrate that BARNN not only achieves comparable or superior accuracy compared to existing methods, but also excels in uncertainty quantification and modelling long-range dependencies.
This work aims to close the gap between autoregressive/ recurrent models and Bayesian methods, presenting a fully Bayesian, scalable, calibrated and accurate autoregressive/ recurrent model, named BARNN: Bayesian Autoregressive and Recurrent Neural Network. BARNN provides a princi- pled way to turn any autoregressive/ recurrent model into its Bayesian version. BARNN the network weights are evolved jointly with the observable states (e.g. language tokens, PDE states, etc.), creating joint probabilistic model (see Figure 1). A new state is generated by sampling from the distribution conditioned on previous states and the cur- rent network weights. At the same time, the latter are drawn from a variational posterior distribution that is condi- tioned on the previous states. We derive variational lower bound for efficiently training BARNN that is related the VAE (Kingma & Welling, 201 2014) (ELBO) objective. To make Bayesian Inference over a large number of network weights computationally efficient, we propose an extension of Vari- ational Dropout (Kingma et al., 2015; Gal & Ghahramani, 2016), which also opens the door to explore quantization and network efficiency (Louizos Welling, 2017: Van Baalen et al., 2020) for overparametrrized autoregressive models, Finally, we introduce temporal version of the "Variational Mixtures of Posteriors" (Tomczak & Welling, 2018) prior (CVAMP-prior) to make Bayesian inference efficient and well-calibrated.
1. Introduction
Autoregressive and recurrent models have demonstrated impressive advancements different disciplines, from weather prediction (Bodnar et al., 2024; Lam et al., 2022),
molecules generation (Shi al., 2020; Simm al., 2020) and Large Language Models (LLMss) (Bengio et al., 2000;
Vaswani, 2017; Radford et al., 2019). However, while au- toregressive models show strong predictive abilities, they are also prone to overfitting to the specific tasks they are trained
Mathematics Area, International School of Advanced Studies, Italy IInformatics Institute. University of Amsterdam, The Nether- lands 'FAST Computing Srl, Italy. Correspondence to: Dario Coscia Mlario.coscia@ee sissa.it>.
Building on the recent successful application of autoregres- sive and recurrent models to Scientific Machine Learning tasks (Pfaff al., 2021: Lippe et al., 2024; Segler al., 2018) we apply our methodology to uncertainty quantifica-
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
p(w(11)
p(w.)
p(w(-1)
@ 31-2)
((w;
y1-1))
g(w yc)
wi
W
wt
W
My
m
M
P( (1-2,
p(dt Be-1
p( Ur, (b--1)
P(V-11 11-22
p(y: ())1,
p( y.,
(a)
states
weights
(b)
evolution
evolution
Figure I: BARNN generative model (a) and inference model (b). Solid lines represent the generative process, whereas dotted lines indicate the variational approximation used for inference. The network weights w, evolve jointly with the states yr. forming a coupled probabilistic model. During inference, the weights are sampled from the variational posterior allowing for improved parameter estimation based on observed data. The draw depicts the one-step causal evolution of the states.
tion (UQ) for Neural PDE Solvers, and molecules genera- tion with Language models for drug discovery. Experiments demonstrate that BARNN not only offers greater accuracy but also provides calibrated and sharp uncertainty estimates, and excels in modelling long-range molecular dependencies compared to related methods. To the best of our knowledge, BARNN is the first approach that transforms any autore- gressive or recurrent model into Bayesian version with minimal modifications, ensuring greater accuracy compared to its not Bayesian counterpart and providing structured way to quantify uncertainties.
cantly from the training data distribution (Papamarkou et al., 2024). However, so far in the literature, appears that quan- tifying the uncertainty of these model parameters, namely epistemic unceriainty, is far from being the norm, and prin- cipled ways to do it are still missing (Aichberger et al., 2024).
2.2. Bayesian Modelling and Uncertainty Quantification
Bayes' theorem (Bayes, 1763) offers systematic approach for updating beliefs based on new evidence, profoundly in- fluencing a wide array of scientific disciplines. In Deep Learning (DL), Bayesian methods have extensively been applied (Hinton Van Camp, 1993; Korattikara et al., 2015; Graves, 2011; I; Kingma & Welling, 2014). They provide probabilistic treatment of the network parameters, enable the use of domain knowledge through priors and overcome hyper-parameter tuning issues through the use of hyper- priors (Papamarkou et al., 2024). Bayesian Models have been historically applied to model epistemic uncertainty, i.c. uncertainty in network parameters. For example in (Gal & Ghahramani, 2016) the Dropout (Srivastava al., 2014) method linked to Gaussian Processes (Rasmussen & Williams, 2005) showing how to obtain uncertainties in neu- ral networks; while in (Kingma et al., 2015) the authors connect Gaussian Dropout objectives to Stochastic Gradient Variational Bayesian Inference, allowing to learn dropout co- efficients and showing better uncertainty. This last method is the most related work ours, but differently from our approach, the weights do not evolve in time, leading to less
2. Background and Related Work
2.1. Autoregressive and Recurrent Networks
Autoregressive and Recurrent models represent a joint prob- ability distribution as a product of conditional distributions, leveraging the probability product rule (Bengio al., 2000; Uria al., 2016). In recent years, autoregressive and current models have been successfully applied in various to different tasks, such as text (Bengio et al., 2000; Vaswani, 2017), graphs (Li et al., 2018; Liu et al., 2018), audio (Dieleman et al., 2016), and, more recently, PDE mod- elling (Brandstetter et al., 2022a;b) and molecules genera- tion (Segler al., 2018; Özşelik al., 2024; Schmidinger al., 2024). Despite their vast applicability, those mod- els are known to overfit, leading to unreliable predictions, especially in over- parameterized regimes, such as modern LLMs, or when the testing data distribution shifts signifi-
2

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
accurate uncertainties (see the experiment section 4.2).
(section 3.1). We show how to optimize the model by deriv- ing novel variational lower bound objective that strongly connects to the VAEs (Kingma & Welling, 2014) framework (section 3.2). Finally, we present scalable variational pos- terior and prior for efficient and scalable weight-parameter sampling (section 3.3).
2.3. Neural PDE Solvers
recent fast-growing field of research is that of surrogate modelling, where DL models, called Neural PDE Solvers, used to learn solutions to complex physical phenom- ena (Li et al., 2020; Bhattacharya et al., 2021; Rozza al., 2022; Pichi et al., 2024; Coscia et al., 2024; Li et al., 2024). In particular, Autoregressive Neural PDE Solvers (Brand- stetter et al., 2022b; Sanchez-Gonzalee et al., 2020) gained significant attention since they can be used to infer solutions of temporal PDEs orders of magnitude faster compared to standard PDE solvers, and generate longer stable solutions compared to standard (not-autoregressive) Neural Operators. However, Autoregressive Neural PDE Solvers accumulate late errors in each autoregressive step (Brandstetter et al., 2022b), slightly shifting the data distribution over long rollouts. This yields inaccurate solutions for very long rollouts, and meth- ods to quantify the uncertainty have been developed: PDE Refiner (Lippe et al., 2024) obtains the PDE solution by refining a latent variable starting from unstructured random Gaussian noise, similarly to denoising diffusion models (Ho et al., 2020); while in GraphCast (Lam et al., 2022) the au- thors introduced& a method called Input Perturbation, which adds small random Gaussian noise to the initial condition, and unroll the autoregressive model on several samples to obtain uncertainties.
3.1. The State- Weith- Model
Autoregressive models represent a joint probability distribu- tion as a product of factorized distributions over states:
p(yo: P.y...... yr) II ply: 11-1....., yo),
where new state yr is sampled from a distribution condi- tioned all the previous states yy-11.... yy. and T is the state-trajectory length. In particular, deep autoregressive models (Schmidhuber et al., 1997) approximate the factor- ized distribution p(ye yt-1..... 3o0 with neural network with optimizable deterministic weights w. We construct a straightforward Bayesian extension of this framework by jointly modelling states y, and weights wi via j distri- bution p(yo, W1, y1, w2; y2,...,, WT, yr) - p(yo.T, Wi:T) states and weights, accounting for the variability in time for both'. The full joint distribution is given by:
T p wiT) II p(y: yt-1..... yo. wi (pwp). (2)
Hence, a new state y/ is obtained by sampling from the distribution conditioned on previous states and the current weights wt while the latter are sampled from a prior distri- bution p(w)), see Figure 1.
2.4. RNN for Molecule Generation
Drug design involves discovering new molecules that can effectively bind to specific biomolccular targets. In re- cent years, generative DL has emerged as a powerful tool for molecular generation (Shi et al., 2020; Eijkel- boom et al., 2024). A particularly promising approach is the use of Chemical Language Models (CLMs), where molecules are represented as molecular strings. In this area, RNNs have shown significant potential (Segler et al., 2018; Özçelik et al., 2024; Schmidinger et al., 2024). In the experiments section, we will demonstrate how classical RNNs, when combined with BARNN, can enhance perfor- mance particullll by improving the model's ability to capture long-range dependencies, generating statistically more robust molecules, and achieving better learning of molecular properties
3.2. The Temporal Variational Lower Bound
Learning the model in eq. (2) requires maximising the log- likelihood log p(y2() Unfortunately, directly optimizing log p(yo:) by integrating eq. (2) over the weights, or by expectation maximization is intractable; thus we propose to optimize pogo(r)) by variational inference (Kingma & Welling, 2014). The variational posterior over network weights given the states (o(W | yo:r) is parametrized by (different immeindependent) weights Φ and reads:
T
Gop (W])T (o:r) = II 2 (wr y yo)
(3)
Given the variational posterior above, in Appendix A.I we derive the following variational lower bound to be max- imised over the variational parameters:
3. Methods
Bayesian Autoregressive and Recurrent Neural Network (BARNN) is a practical and easy way to turn any autore- gressive or recurrent model into its Bayesian version. The BARNN framework creates a joint distribution over the observable states (e.g. language tokens, PDE states, etc.) and model weights, with both alternatingly evolving in time
L( (E.-U((17) [uuingg ply wi))
(4)
DKL (wo )lr=))]lp(w:)]]
'We indicate a sequence (ax, ax+1, ak- ab+t) with akakez VI > k E N.
3

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
The equation above has two nice properties. First, it resem- bles the VAE objective (Kingma & Welling, 2014) creating connection between Bayesian networks and latent vari- able models. Second, BARNN incorporates into the ELBO the timestep dependency, allowing for adjustable weights in time which we will show provide sharper, better cali- brated uncertainties as well the ability to better capture long range dependencies in the sequences. Finally, in the limit of time-independent peaked distribution with prior and pos- terior perfectly matching, the BARNN loss simplifies to the standard log- logllikeihhod optimization used in standard autoregressive and recurrent models (see Appendix A.5 for derivation). This suggests interpreting eq. (4) as a form of Bayesian weight regularization during autoregressive model training. Finally, once the model trained, the predictive distribution and uncertainty estimates can be easily com- puted by Monte Carlo sampling (see Appendix A.4).
where w/ are the weights time / for layer 1 of the au- toregressive recurrent network. Second, assume go(w)) y0.--1) follows normal distribution as follows:
go(w) yoot-1) =
(6)
w/ = afss'(1 , _E~NN(,
where the square Ω25 is done component-wise, mean- ing weights in each layer are not correlated, and of are scalar positive dropout variational coefficients depending on yo:r-1 (and possibly t depending on the specific ap- plications). In practice, a network encoder Eg is used to output the vector of dropouts coefficients for all layers a [a...... at]] E (yo.--)) with &b the encoder weights and the recaaareparametiizaatoonn trick is applied (Kingma et al., 2015) component-wise for linear layers. Thus, the variational parameter to optimize are Φ = (92, {b), with Ω = (Stt,...,,S* S2 static weights, and ab encoder weights.
3.3. Variational Dropout Approximation
When examining the generative model in eq. (2) the lower bound in eq. (4), a key challenge arises: how can we efficiently sample network weights? Directly sampling network weights is impractical for large networks, making alternative approaches necessary. Variational Dropout (VD) (Kingma et al., 2015; Molchanov et al., 2017) offers a solu- tion by reinterpreting traditional Dropout (Srivastava et al., 2014) randomly zeros out network weights during training- as form of Bayesian regularization, and use the local reparametrization trick for sampling only the activa- tions resulting in computational efficiency. We re-interpret VD for sampling dynamic weights w, during training and inference. Our goal is to do a reparametrization of the weights wi = fS2 axe c), with Ω dettrmitttriiiiitti net- work weights, on deterministic time-dependent variational weights, p(E) some random noise, and f differen- tiable function. This parametrization is very flexible: (i) it allows to apply the rotaaal-pprammerizattooo trick (Kingma & Welling, 2014); (ii) it uses fixed network weights Z en- abling to scale the Bayesian methodology to large networks; (iii) it incorporates time and previous state dependency by perturbing 12 with deterministic time-dependent variational weights 0g and adds stochasticity with random noise e, making the methodology easily implementable with little change to the main autoregressive or recurrent model. We present specific method for defining this reparameteriza- tion while leaving other potential approaches and extensions for future work.
Aggregated Variational Posterior in Time Prior To find suitable prior for the model, first observe that if the KI KL-divergence in eq. (4) is independent of Ω then max- imising the variational bound L with respect to Ω for fixed Or, is equivalent to maximise the expected log-likelihood log p(y | 52 sb, 30.c-1)², where we made explicit the weight dependency. Thus, we look for a prior that allows us to have a KL-term independent of f2.
Following the ideas of (Tomezak & Welling, 2018), can find that the best prior for the lower bound in eq. (4) is given by (proof in Appendix A.2):
p(ws)) :
f
f ply0:t-1))(u I 10:--1) dyoat-1-
We call this prior Temporal Variational Mixture of Posterior (IVAMP). This result is very similar to VAMP obtained in (Tomezak & Welling, 2018) but with a few differences: (i) VAMP is prior on the latent variables, while tVAMP is a prior on the Bayesian Network weights; (ii) tVAMP is time-dependent, i.e. for each time t the best prior is given by the aggregated posterior, while in (Tomczak & Welling, 2018) time dependency was not considered. Assuming the prior factorize similarly to the posterior, by carrying out the computations (see Appendix A.2), we obtain the following
Variational Posterior: First, assume that weights wi fac- torize over layers E (I....,L) L) for the posterior 4s, imply-
ing:
L
(wr yo:s-1)= = II go(w vo:t-1), 1=1
(5)
"This is the same optimising the network weights S2 in standard deterministic networks.

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
prior indicating with & the batch index:
3.5. Bayesian Recurrent Neural Networks
p(w)) = NN (7/55)"
Given a sequence yo: ........ yT standard Recurrent Neu- ral Network (RNN) (Bengio et al., 2000; Schmidhuber et al., 1997) computes eq. (1) by introducing hidden variables (he Ji that store' the input information up to the time t. Given these variables, the conditional probability is mod- elled as:
1 B: N
Σ a -
(8)
N
(at(vbn-1)) :/ N Σ
VI,6
Ply Ve-1,-.., yo) = o a(h-11) h, = -NN(N h1:
(1I)
Finally, for the given prior and posterior defined above, the KL-divergence indeed independent of Z (proof in Appendix A.3) and reads:
where a is the softmax function, and NN is a neural network with deterministic weights w and a specific gate mecha- nism depending on the RNN structure. Importantly, the variable h, contains historical information up to time I, i.e. knowledge of yt-11....;3 y00
DKL (ap(ws 10.--1) (lx=-)))p((w..))
1.
12'
c
3!
a2
of 1 - 2 In
(9)
+
7
2
Extending RNNs to Bayesian RNNs is straightforward with BARNN. In particular, the states distribution becomes:
where [22| is the dimension of the weights in layer I.
p(ye Vt-1. Yo w,) = ohh-11) he = NN ht-1;
(12)
3.4. Bayesian Neural PDE Solvers
Autoregressive Neural PDE Solvers map PDE states yt to future states Yt-1, given a specific initial state yo (Brand- stetter et al., 2022b). Those solvers are deterministic and compute the solution without providing estimates of uncer- tainty. BARNN can be used to convert any Autoregressive Neural PDE Solver into its Bayesian version by adopting few steps. First, we model the joint weight-state distribution as:
while the posterior distribution over the weights is obtained by conditioning on the hidden states:
qep yi-1..... yo) = qup (we 9t-1 h(-2).
(13)
The probabilistic weights we are obtained again with the transformation (S2 Ore; E) applied layerwise as explained in eq. (6), with the encoder Es taking yc-1, hi-2 input. Finally, learning is done by maximizing eq. (4), which, given the specific state distribution, becomes the minimiza- tion of the cross entropy loss (Bengio et al., 2000) commonly used in causal language modelling or next token prediction, plus the negative KL divergence in eq. (9). The training algorithm is reported in Appendix D.
T p(yo:T, Wi:)) = II plyel yt-1, (p)(ur).
(10)
which is a special case of eq. (2) when markovianity of first order on the states is assumed. Then, we as- sume a Gaussian state distribution p(y. | 3)-1,w,) = N (NNO 1; w;) diag(a?)), where the mean at time is obtained by applying the Autoregressive Neural PDE Solver to the state yt-1, while the standard deviation is not learned. Specifically, NO indicates any Neural PDE Solver archi- tecture, with w, its probabilistic weights, showing that the framework is independent of the specific architecture used. The probabilistic weights are obtained with the transforma- tion (S2 at,, E) applied layerwise as explained in eq. (6), with the encoder Ex taking only one state as input (due to markovian updates). Learning done by maximizing eq. (4), which, given the specific state distribution, becomes the minimization of the standard one-step MSE loss (Brandstet- ter et al., 2022a) commonly used to train Autoregressive Neural PDE Solvers, plus a Bayesian weight regularization term given by the negative KL divergence in eq. (9). The training algorithm is reported in Appendix
4. Experiments
We demonstrate the effectiveness of BARNN by applying it to different tasks. To begin, we validate the tVAMP prior on synthetic time series dataset and demonstrate its superior- ity over the widely used log-uniform prior. Then, we show that BARNN combined with Neural PDE Solvers (Bar-Sinai et al., 2019; Brandstetter et al., 2022b; Li et al., 2020) can solve PDEs and quantify the related uncertainty. Finally, we apply BARNN to RNNs for molecule unconditional gen- eration (Segler et al., 2018; Özşelik et al., 2024) using the SMILES syntax and show stronger generation capabilities compared to existing methods. We make our code publicly available at http:://giiuub.co/darrroooosssss
Wiit the convention that h-1 =
5

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Burgers
KS
KdV
1.0
0
0.5
0.0
-5
0
4
6
20
40
60
20
40
60
X
X
X
Figure 2: Prediction and uncertainty intervals for different PDEs at last time-step. The figure depicts the CFD solution (red dotted line) alongside the BARNN mean prediction (solid blue line) with +3 std for uncertainty quantification (shaded blue area).
4.1. Time Series
1.0
1--1
Burgers
1=-1
16-1
We evaluate the performance of the proposed (VAMP prior on a synthetic time series forecasting task. The dataset is constructed by generating sequences of sinusoidal signals with varying frequencies and phases:
0.8
2.0
RMSE
0.6
0.4
0.2
0.5
0.0
It-1 3r
To 0
(14)
sin (11Te + Bi),
j=1
T1N
with a; U0.55, 1.5), 8 3n], 21. 100). This setup serves controlled environment to test the effect of prior choice model uncertainty and forecasting accuracy. We generate 1024 trajectories for training, and 100 for testing. All models are trained using the same ar- chitecture, a 2-layer neural network of 64 units and Relu activation. The mean-square-error (MSE) loss minimized for 1500 epochs using the Adam (Kingma Ba, 2014) optimizer with learning rate 10-1 and weight decay 10-8
200
200
200
Unrolling
Uncelling
Uncolling step
Figure 3: RMSE and NLL for different Neural Solvers predicting solutions different PDEs.
4.2. PDE Modelling
IVAMP Outperforms Log-Uniform Prior Synthetic Forecasting Task Table presents an ablation study com- paring different prior choices for BARNN the synthetic time series dataset. We also compare BARNN against standard non- -ay-Bayeiian MLP, and Monte Carlo Dropout (Dropout) (Gal & Ghahramani, 2016) with different levels of dropout probability. The BARNN model using the pro- posed (VAMP prior achieves the best performance across all metrics, outperforming both the log-uniform prior and standard baselines such dropout and deterministic MLPs. Notably, tVAMP improves over the log-uning uniform prior, the negative log- likeliidldoodd from -0.092 to -0.166 indicating better-calibrated predictive uncertainty, while also achieving the lowest RMSE and ECE. These results demonstrate the effectiveness of the tVAMP prior in capturing uncertainty and enhancing forecasting accuracy time series models.
PDE modelling consists of finding solutions to partial dif- ferential equations, generalising to unseen equations within given family. We test BARNN on three famous PDEs, namely Burgers, Kuramoto- Serumoto--ivashinsky (KS) and Korteweg de Vries (KdV) (Evans, 2022; Brandstetter et al., 2022a), and compare its performance against multiple UQ tech- niques for Neural PDE Solvers: Monte Carlo Dropout (Dropout) (Gal & Ghahramani, 2016), Variational Dropout with Empirical Bayes (ARD) (Kharitonov et al., 2018), which a variation of Variational Dropout, and the recently proposed Input Perturbation (Lam al., 2022) (Perturb), PDE Refiner (Lippe et al., 2024) (Refiner). A more detailed explanation of dataset generation, metrics, neural operator architectures, baselines and hyperparameters can be found in Appendix B, along with additional results.
6

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Table I:
Ablation for different BARNN priors for the sinusoidal time-series forecasting dataset. The results report RMSE, NLL, and ECE statistics. Static method reports the RMSE obtained the initial state is not propagated, while MLP the base (non- con-Bayesian) architecture. BARNN uses the same MI.P architecture, and ablated on different priors.
Model
Prior
MSE (L)
NLL (L)
ECE (1)
Static
0.490=0.000
MLP
0.081 =0.011
Dropout (p=0.5)
0.072=0.001
0.593=0.161
0.084 #0.010
Dropout (p=0.2)
0.048 0.048=0.004
-0.075+0.004
0.068 2.0009
BARNN
log-uniform
0.045=0.005
-0.092+0.064
0.050+0.016
BARNN
tVAMP
0.043
-0.166-0.019
0.049
Table NLL and ECE for Burgers, KS and KdV PDEs for different Neural Solvers. The mean and std are computed for four different random weights initialization seeds. Solvers are unrolled for 320 steps, break is reported when the solver metric diverged.
0.0170
Burgers
R05SU
0,0180
MAP estimate
KS
BMSEE
estimate
Model
NLL (L)
Burgers
KS
KdV
20
100
BARNN
-2.51 1.0.9
0..5
0.94 0.94-11.25 1.23
ARD
5.00=7.11
10.7#11.6
break
RMSE
0.0G85
KdV
MAP estimate
0,0590
Dropout
0.06=1.59
- -0.48+0.65
0.40+1157
Perturb
2.01 00.08
-0.65-41.07
0.332
100
Refiner
break
break
break
number of essemble memebers
Figure 4: Variation of RMSE for increasing ensemble mem- bers in BARNN.
Model
ECE (4)
Burgers
KS
KdV
BARNN
0.05 10.03
4.040
0.10
highlights that point-wise RMSE alone is insufficient for evaluating Neural PDE Solvers, especially in UQ scenarios, 0.14
ARD
0.21
0.17:00.01
0.10
Dropout
0.17=0.15
0.10#0.01
as low RMSE values can still mask incorrect uncertainty estimates from the solver.
Perturb 0.16
0.05
0.11 0.105
Refiner
0.30=0.08
0.23
0.34
Confidence interval adaptahility Figure 2 illustrates the BARNN predictions alongside the 99.7% confidence in- tervals across different regions of the domain. Notably, BARNN dynamically adjusts the width of these confidence intervals based on the uncertainty present in each region. For example, in areas of high uncertainty, such as with the KS equation, the model appropriately broadens the intervals, reflecting its lower confidence in those predictions. Con- versely, regions where the model closely aligns with the true solution, in the case of the KdV equation, the con- fidence intervals narrow considerably. This demonstrates BARNN's ability to flexibly adapt to varying levels of uncer- tainty, offering robust and reliable uncertainty quantification that aligns with the behaviour of the underlying PDEs.
BARNN solves PDEs with calibrated uncertainties We unroll the trained Neural Solvers for 320 steps and com- pute the negative log-likelihood (NLL) and the expected calibration error (ECE) of the predictions. NLL measures balance between the accuracy of model prediction and the sharpness of the model's standard deviation, while the ECE measures how well the estimated probabilities match the observed probabilities. Table 2 reports the results, show- ing that BARNN can accurately solve PDEs with lower NLL and ECE compared to all existing methods. We found that, despite the models exhibiting similar root mean square errors (RMSE). their negative logllikeihhood varied sig- nificantly (Figure 3), driven by underestimated ensemble variance, which in turn caused a sharp increase in the NLL (as seen in models like Refiner and ARD). This finding
Test number of ensemble members in BARNN A de- sirable property of ensemble methods is that for suffi-

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
space needed to be explored, and long-range dependencies within molecular structures. We evaluate BARNN the molecule generation task, by representing molecule with the SMILES syntax and pre-train a recurrent neural network (LSTM gate-mechanism) (Schmidhuber et al., 1997) on the ChEMBL dataset from (Özçelik et al., 2024), which con- tains approximately 1.9 M SMILES. We test the BARNN LSTM model against standard LSTM (Schmidhuber et al., 1997), and the SMILES LSTM (Segler al., 2018) which differ from the standard by applying Dropout to the LSTM input and hidden variables. Additional details on the dataset, metrics, hyperparameters, and baselines can be found in Appendix C.
1e3
BARNN
Number of Errors
2.0
SMILES LSTM
ISTM
1.5
1.0
0.5
0.0
Misc. Branching Valence
Rings
Total
Error type
Source of error for invalid generated SMILES.
BARNN correctly learn the SMILES syntax BARNN perfectly learns the SMILES syntax generating molecules with high validity, diversity, novelty and uniqueness com- pared to the baselines (Table 3). For validity, diversity and uniqueness BARNN archives 1% higher values than SMILES LSTM, and 2% than LSTM, while diversity is only slightly improved (0.1%) compared to the baselines.
Number of Molecules
10*
10a
10a
Long range dependencies in SMILES We sampled 50K SMILES from the language models and for the invalid molecules examined the source of errors (Figure 5a). In- terestingly, most of the errors language models commit are due to ring assignment, e.g. opening but not closing the ring. For ring assignments, BARNN outperforms the baselines, improving by 30% over SMILES LSTM and 50% over LSTM. In SMILES notation, ring closure bonds are represented by assigning matching digits to connected atoms, indicating where the ring opens and closes. A good language model must identify when the ring opens and remember to close it (with the same digit), long-range dependency. This is analyzed by counting the numerical markers in each SMILES string and assessing validity (Fig- ure 5b). As expected, SMILES strings with up to 4 rings show errors, but validity decreases the number of rings increases. BARNN outperforms baselines, generat- ing more valid molecules, particularly for molecules with many rings, demonstrating superior handling of long-range dependencies.
10a
SMILES Numerical Markers Count
(b) Analysis of the impact of multiple rings in a SMILES string. Shaded colours indicate number of total SMILEE (valid + invalid), while solid colours indicate only valid ones.
Figure 5: Error analysis for SMILES, mean and std shown for four different random weight initialization seeds.
ciently high number of ensemble members, the statistical moments (e.g. mean, std) converge. Figure 4 shows that with only ~ 30 ensemble members BARNN provides con- vergent RMSE, emphasizing that accurate moments esti- mates can be obtained with just few additional stochastic forward passes. In Appendix B.3 we also report ECE and NLL, which follow similar trends. Interestingly, if uncer- tainties are not required for specific application, using the maximum a-posteriori estimate (MAP) for the network weights only requires a single forward pass to achieve a convergent RMSE, as depicted in Figure 4. This makes the approach as fast as standard non-ensemble methods, which can be advantageous in practical applications.
BARNN perfectly learns molecular properties We test the ability of BARNN to learn the molecular properties of the training data by computing the Wasserstein distance W2 between model-predicted properties and training prop- erties. Following (Segler al., 2018), we compute Molecu- lar weight (MW), H-bond donors (HBDs), H-bond accep- tors (HBAs), Rotable Bonds, LogP, and TPSA. Overall, the BARNN approach obtains the lowest Wasserstein distance for all molecular properties (Table 4), demonstrating an ex- cellent ability to match the training data distribution. In particular, for molecular weight, which is the hardest prop-
4.3. Molecule Generation
Molecular generation involves creating new molecules with desired properties, which is difficult due to the vast chemical
8

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Table 3: Designing drug-like molecules de novo with BARNN, mean and std shown for four different random weight initialization seeds.
Model
Validity (1)
Diversity (1)
Novelty (t)
Uniqueness (1)
BARNN
95. 05.09:t0.34
88 63.00.01
95.41 95.41=0024 20.24
95. 0..06
SMILES LSTM
94. 94.6040..2
88.57400.9
94.25 14.15
94.58-0.22
LSTM
93.02#0.33
88.57=0.02
92.81 #0.34
92.14=0.15
Table 4:
Wasserstein distance between models' and training dataset's ' molecular properties distribution, mean and std shown for four different random weight initialization seeds.
Model
W2 (1)
MW
HBDs
HBAs
Rotable Bonds
LogP
TPSA
BARNN
2.53+11003
0.02 #0.01
0.05 t0.03
0.13 0.008
0.06 #0.01
0.90 0.338
SMILES LSTM
10.4+3.71
0.03+0002
0.15+00.06
0.14+0.05
0.11
+0.97
LSTM
12.3
0.03+0.01
0.20#0.07
0.13 0.07
0.12#0.08
1.84+1.04
erty to predict, BARNN obtains a substantial decrease in W2 compared the baselines. The results on molecular prop- erties in Table 4, along with those on statistical properties in Table 3, demonstrate that BARNN not only more accu- rately reflects the dataset properties but also generates, average, more statistically plausible molecules, significantly advancing the quality of molecular generation.
more effective priors that incorporate prior knowledge could improve the model's performance. Another potential area of research is investigating local pruning at each autoregres- sive step. particularly for large over pere-aaaetteeieed networks such as LLMs. Finally, deeper examination of the roles of epistemic and aleatory uncertainty over time represents an important direction for further study.
5. Conclusions
A Acknowledgments
We have introduced BARNN, scalable, calibrated and ac- curate methodology to turn any autoregressive recurrent model to its Bayesian version, bridging the gap between autoregressive// recurrent methods and Bayesian inference. BARNN creates a joint probabilistic model by evolving network weights along with states. We propose novel vari- ational lower bound for efficient training, extending Varia- tional Dropout to compute dynamic network weights based on previous states, and introduce temporal VAMP-prior to enhance calibration. We demonstrate BARNN's applica- tion in Autoregressive Neural PDE Solvers and molecule generation using RNNs. Experiments demonstrated that BARNN surpasses existing Neural PDE Solvers in accu- racy and provides calibrated uncertainties. BARNN is also sharp, requiring only few ensemble estimates and, un- certainties not needed, only one forward pass using the weights maximum a-posteriori. For molecules, BARNN generated molecules with higher validity, diversity, novelty and uniqueness compared to its non-Bayesian counterparts, and it also excels in learning long-range and molecular prop- erties. BARNN a very general and flexible framework that can be easily applied other domains of DL, such as text, audio or time series. Several promising extensions for BARNN could also be explored. For example, developing
D. Coscia, N. Demo and G. Rozza acknowledge the sup- port provided by PRIN "FaReX Full and Reduced or- der modeling of coupled systems: focus non-matching methods and automatic learning" project, the European Re- search Council Executive Agency by the Consolidator Grant project AROMA-CFD "Advanced Reduced Order Meth- ods with Applications in Computational Fluid Dynamics" GA 681447, H2020-ERC CoG 2015 AROMA-CFI -FDD PI G. Rozza, and the CINECA award under the ISCRA ini- tiative, for high-performance computing resources and sup- port availability. N. Demo and G. Rozza conducted this work within the research activities of the consortium iN- EST (Interconnected North-East Innovation Ecosystem), Piano Nazionale di Ripresa Resilienza (PNRR) Mis- sione 4 Componente 2, Investimento 1.5 D.D. 1058 23/06/2022, ECS00000043, supported by the European Union's NexIGenerationEU program.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
References
Eijkelboom, F., Bartosh, G., Naesseth, C. A., Welling, M., and de Meent, J.-W. Variational Flow Matching for Graph Generation. arXiv preprint 2rXi::24000048453 2024.
Aichberger, L., Schweighofer, K., and Hochreiter, S. Re- thinking uncertainty estimation in natural language gen- eration. arXiv preprint arXiv:2412.15170 2024.
Evans, L. C. Partial Differential Equations. In American Mathematical Society, 2022.
Bar-Sinai, Y., Hoyer, S., Hickey, J., and Brenner, M. P. Learning Data-Driven Discretizations for Partial Differ- ential Equations. National Academy of Sciences, 116(31); 15344-15349, 2019.
Falcon, W. and The PyTorch Lightning team. PyTorch Lightning. PyTorch Lighming, March 2019. doi:
URL betpss://githen com com/Lightning- wtwpiiiiinnA//gggttttatttae lightning.
10.5281/7emodo.388888
Bayes, T. An Essay towards solving Problem in the Doctrine of Chances. By the late Rev. Mr. Bayes, FRS communicated by Mr. Price, in a letter to John Canton, A.M.F.R.S. Philosophical transactions of the Royal Soci- ety of London, 1763.
Gal, Y. and Ghahramani, Z. Dropout as a Bayesian Ap- proximation: Representing Model Uncertainty in Deep Learning. International Conference on Machine Learn- ing. 2016.
Bengio, Y., Ducharme, R., and Vincent, P. A Neural Proba- bilistic Language Model. Advances in neural information processing systems, 2000.
Graves, A. Practical Variational Inference for Neural Net- works. Advances in Neural Information Processing Sys- lems, 24, 2011.
Bhattacharya, K., Hosseini, B., Kovachki, N. B., and Stu- art, A. M. Model Reduction and Neural Networks for Parametric PDEs. The SMAI journal of computational mathematics, 7:121-157, 2021.
Hinton, G. E. and Van Camp, D. Keeping the Neural Net- works Simple by Minimizing the Description Length of the Weights. In Conference on Computational Learning Theory, 1993.
Bodnar, C., Bruinsma, W. P., Lucic, A., Stanley, M., Brand- stetter, J., Garvan, P., Riechert, M., Weyn, J., Dong, H., Vaughan, A., et al. Aurora: A Foundation Model of the Atmosphere. arXiv preprint arXiv:2405. 130633 2024.
Ho, J., Jain, A., and Abbeel, P. Denoising Diffusion Proha- bilistic Models. Advances in Neural Information Process- ing Systems, 2020.
Jiang, Z., Araki, J., Ding, H., and Neubig, G. How Can We Know When Language Models Know? On the Cal- ibration of Language Models for Question Answering. Transactions of the Association for Computational Lin- guistics, 2021.
Brandstetter, J., Welling, M., and Worrall, D. E. Lie Point Symmetry Data Augmentation for Neural PDE Solvers. International Conference on Machine Learn- ing, 2022a.
Brandstetter, J., Worrall, D. E., and Welling, M. Message Passing Neural PDE Solvers. In International Conference on Learning Representations, 2022b.
Kharitonov, V., Molchanov, D., and Vetrov, D. Varia- tional Dropout via Empirical Bayes. In arXiv preprint arXiv:1811.0083: 1.00596, 2018.
Coscia, D., Ivagnes, A., Demo, N., and Rozza, G. Physics- informed neural networks for advanced modeling. Jour- nal of Open Source Software, 8(87):5352, 2023.
Kingma, D. P. and Ba, J. L. Adam: A method for stochas- tic optimization. International Conference on Learning Representations, 2014.
Coscia, D., Demo, N., and Rozza, G. Generative Adversarial Reduced Order Modelling. Scientific Reports, 2024.
Kingma, D. P. and Welling, M. Auto-Encoding Variational Bayes. International Conference on Learning Repre- sentations, 2014.
Depeweg, S., Hernandez-Looatoo J.-M., Doshi Velez, F., and Udluft, S. Decomposition of Uncertainty in Bayesian Deep Learning for Efficient and Risk-Sensitte Learning. In International Conference on Muchine Learning. 2018.
Kingma, D. P., Salimans, T., and Welling, M. Variational Dropout and the Local Reparameterization Trick. In Ad- vances in Neural Information Processing Systems, 2015.
Dieleman, S., Zen, H., Simonyan, K., Vinyals, O., Graves, A., Kalchbrenner, N., Senior, A., Kavukeuoglu, K., et al. Wavenet: Generative Model for Raw Audio. arXiv preprint arXiv: /609 00499, 2016.
Kohl, G., Chen, L., and Thuerey, N. Benchmarking Regres- sive Conditional Diffusion Models for Turbulent Flow Simulation. In ICML 2024 Al for Science Workshop, 2024.
10

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Korattikara, A., Rathod, V., Murphy, K., and Welling, M. Bayesian Dark Knowledge. In International Conference on Neural Information Processing Systems, 2015.
Deep Learning is Needed in the Age of Large-Scale AI. In International Conference Machine Learning, 2024.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z.., Gimelshein, N., Antiga, L., al. Pytorch: An imperative style, nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Lam, R., Sanchez-Gonzalees A., Willson, M., Wirnsberger, P., Fortunato, M., Alet, F., Ravuri, S., Ewalds, T., Eaton- Rosen, Z., Hu, W., et al. GraphCast: Learning Skill- ful Medium-Ranee Global Weather Forecasting, arXiv preprint arXiv:2212.12779, 2022.
Pfaff, T., Fortunato, M., Sanchez-Gonzalez, A., and Battaglia, P. Learning mesh-based simulation with graph networks. In International Conference on Learning Rep- resentations, 2021.
Lam, R., Sanchez-Gonzales A., Willson, M., Wirnsberger, P., Fortunato, M., Alet, F., Ravuri, S., Ewalds, T., Eaton- Rosen, Z., Hu, W., et al. Learning Skillful Medium Range Global Weather Forecasting, In American Association for the Advancement of Science, 2023.
Pichi, F., Moya, B., and Hesthaven, J. S. A Graph Convo- lutional Autoencoder Approach to Model Order Reduc- tion for Parametrized PDEs. Journal of Computational Physics, pp. 12762,, 2024.
Li, Y., Vinyals, O., Dyer, C.,, Pascanu, R., and Battaglia, P. Learning Deep Generative Models of Graphs. arXiv preprint arXiv: 033343324, 2018.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language Models are Unsupervised Multitask Learners. OpenAl blog, 2019.
Li, 7.,, Kovachki, B., Arizzadenesheli, K., Bhattacharya, K., Stuart, A., Anandkumar, A., et al. Fourier Neural Operator for Parametric Partial Differential Equations. In International Conference on Learning Representations, 2020.
Rasmussen, C. E. and Williams, C. K. Gaussian Processes for Machine Learning. The MIT Press, 2005.
Rozza, G., Stabile, G., and Ballarin, F. Advanced Reduced Order Methods and Applications in Computational Fluid Dynamics. SIAM, 2022.
Li, Z., Zheng, H., Kovachki, N., Jin, D., Chen, H., Liu, B., Azizzadenesheli, K., and Anandkumar, A. Physies- Informed Neural Operator for Learning Partial Differ- ential Equations. ACM/JMS Journal of Data Science, 2024.
Sanchez-Goozalez,, A., Godwin, J., Pfaff, T., Ying, R., Leskovec, J., and Battaglia, P. Learning to Simulate Complex Physics with Graph Networks. International Conference on Machine Learning, 2020.
Lippe, P., Veeling, B., Perdikaris, P., Turner, R., and Brand- stetter, J. PDE-Refiner: Achieving Accurate Long Roll- outs with Neural PDE Solvers. In Advunces in Neural Information Processing Systems, 2024.
Schmidhuber, J., Hochreiter, S., et al. Long Short-term Memory, Neural Comput, 98): 1997.
Liu, Q., Allamanis, M., Brockschmidt, M., and Gaunt, A. Constrained Graph Variational Autoencoders for Molecule Design. Advances in Neural Information Pro- cessing Systems, 31. 2018.
Schmidinger, N., Schneckenreiter, L., Seidl, P., Schimunek, J., Hoedt, P.-J., Brandstetter, J., Mayr, A., Luukkonen, S., Hochreiter, S., and Klambauer, G. Bio-xIstm: Gen- erative modeling, representation and in-context learning of biological and chemical sequences. arXiv preprint arXiv:241 1.04165, 2024.
Louizos, C. and Welling, M. Multiplicative Normalizing Flows for Variational Bayesian Neural Networks. In International Conference Machine Learning, 2017.
Segler, M. H., Kogej. T., Tyrchan, C., and Waller, M. P. Gen- erating Focused Molecule Libraries for Drug Discovery with Recurrent Neural Networks. ACS Central Science, 2018.
Molchanov, D., Ashukha, A., and Vetrov, D. Variational Dropout Sparsifies Deep Neural Networks. In Interna- tional Conference on Machine Learning, 2017.
Shi, C., Xu, M., Zhu, Z., Zhang, W., Zhang, M., and Tang. J. GraphAF: Flow-based Autoregressive Model for Molec- ular Graph Generation. In International Conference on Learning Representations, 2020.
Özęelik, R., de Ruiter, S., Criscuolo, E., and Grisoni, F. Chemical Language Modeling with Structured State Space Sequence Models. Nature Communications, 15(1): 6176, 2024.
Simm, G., Pinsler, R., and Hernández-Lobato, J. M. Re- inforcement Learning for Molecular Design Guided by Quantum Mechanics. In International Conference on Machine Learning, 2020.
Papamarkou, T., Skoularidou, M., Palla, K., Aitchison, L., Arbel, J., Dunson, D., Filippone, M., Fortuin, V., Hennig, P., Hernández-Lobato, J. M., et al. Position: Bayesian

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Dropout: A Simple Way to Prevent Neural Networks from Overfitting. In The Journal of Machine Learning Research, 2014.
Tomezak, J. and Welling, M. VAB with a VampPrior. In International Conference Artificial Intelligence and Statistics, 2018.
Uria, B., Côté, M.-A., Gregor, K., Murray, I., and Larochelle, H. Neural Autoregressive Distribution Esti- mation. Journal of Machine Learning Research, 2016.
Van Baalen, M., Louizos, C., Nagel, M., Amjad, R. A., Wang, Y., Blankevoort, T., and Welling, M. Bayesian Bits: Unifying Quantization and Pruning. Advances in neural information processing systems, 2020.
Van der Maaten, L. and Hinton, G. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Vaswani, A. Attention Is All You Need. Advances in Neural Information Processing Systems, 2017.
Xiao, Y., Liang, P. Bhatt, U., Neiswanger, W. Salakhut- dinov, R., and Morency, L.-P. Uncertainty Quantifica- tion with Pre-trained Language Models: A Large-Scale Empirical Analysis. In Findings of the Association for Computational Linguisties: EMNLP 2022, 2022.
Yang, A. X., Robeyns, M., Wang, X., and Aitchison, L. Bayesian Low-rank Adaptation for Large Language Mod- els. In International Conference on Learning Representa- tions, 2022.
Zdrazil, B., Felix, E., Hunter, F., Manners, E. J., Blackshaw, J. Corbett, S., de Veij, M., Joannidis, H., Lopez, D. M., Mosquera, J. F,, et al. The ChEMBL Database in 2023: Drug Discovery Platform Spanning Multiple Bioactivity Data Types and Time Periods. Nucleic Acids Research, 52:D1 180-D1192, 2024.
12

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
SUPPLEMENTARY MATERIALS BARNN: BAYESIAN AUTOREGRESSIVE AND RECURRENT NEURAL NETWORK
TABLE OF CONTENTS
A Proofs and Derivations
14
A.1 Variational Lower Bound proof
14
A.2 Variational Mixture of Posterior in Time proof...
14
A.3 KL-divergence proof
16
A.4 Predictive Distribution and Uncertainty Estimates
16
A.5 Connection not Bayesian Autoregressive and Recurrent Networks
17
B BARNN PDEs Application
18
B.1 Data Generation and Metrics
18
B.2 Model Architectures, Hyperparameters, and Computational Costs
19
B.3 Additional Results..
20
BARNN Molecules Application
23
C.1 Data Generation and Metries
23
C.2 Model Architectures, Hyperparameters, and Computational Costs
23
C,3 Additional Results..
24
D Algorithms and Software
25
D.I Pseudocodes
26
D.2 Software
26
13

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
A. Proofs and Derivations
This Appendix Section devoted the formal proofs introduced in the main text. In A. we prove the variational lower bound, while the --AMPP proof found in A.2. Finally, A.3 contains the derivation of KL-divergence between prior and posterior presented in 3.3.
A.1. Variational Lower Bound proof
We aim to optimize the log- likelihood for the model:
T
p(yo.T, w1:r) II
plys 10--1, wi w.)))..)
(15)
By defining the variational posterior:
go (wir Bo:T) II gowwr y0:c-1);
(16)
obtain the following evidence lower bound:
log p(Bo.r) = log / p p(yo.T, WItT) dwit
lo (wir
yo:r) dwier
log
p(Bo:T: Wir))
9 (IIT yo:r)
p(yO:T, wI:T)
2 (ww. (w (oor)
log 44 (W1:1 your)
T
p(y Mloee wi (p(s))
-Ewir
log
4 (w,
10o
1=1
(17)
log p(3/1 You--1, wi)p(wr) q (w: you--1)
Fu Swo (we (mm
(wwrgglwwllbee. plog(ly your-1, w[]] DEY. lao(wr you ([w(()))
Σ ($)
Interestingly, by being Bayesian, ended up having sum of evidence lower bounds, which tells us that maximise the model likelihood need to maximise the cumulative evidence lower bound. For performing scalable training approximate the found evidence lower bound as usually done in causal language modelling or next token prediction by:
((Φ) Essupny (Cumulative ($, t)].
(18)
A.2. Variational Mixture of Posterior in Time proof
We want to find the best prior that maximises the evidence lower bound:
log p(yorT, whr))
L(Φ) E torrrsaat w r(BOTT)
go (wier yorr)
14

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
then:
p" (wi))) = arg ((wer)
(())p(yO:)) dyorr
/ plyoor)a( yo.r) logp(wir)) dyo.rdwi.
arg ma
/
V pl pluoor))(wi (orr) dyo.T log (ogglwrr)dn
arg maxp(
(19)
[ p(yoor)a(ure) Borr) dyor p(WI:))
arg max paxp)(xr)) p(WIT)
H
- =
pl ((wrrr yorr) dyo:r
where H is the cross entropy and has a maximum when the two distributions match. In particular, for each time é
p¹ (w()))== / (Wii:) dwpi p(yo:r)a(wr:r yorr) dyoer dua = f p(y0:--1)9(w) yo:r-1) dyout-1
This result is very similar to the one obtained in (Tomczak & Welling, 2018) but with few differences: (i) (Tomezak & Welling, 2018) is prior on the latent variables, here is prior on the Bayesian Network weights; (ii) the presented prior is time-dependent, i.e. for each time / the best prior is given by the aggregated posterior, while in (Tomczak & Welling, 2018) time dependency was not considered.
We approximate the aggregated distribution assuming factorizes similarly to the variational posterior:
L p" (wi) x where, we can compute B/. and of functions of at:
L IIII / p(uo:r-))W(aaddd (a Jdyo..1 IIII N (B (360) 1=1
(20)
1=1 8:30
B Buuly Wwirr
- - p(soo-1) [f Whin (afst (alki Mydwhsi dyoa-1
(21)
= Shist J 1ole Styp
Σ ((
N
and,
(2 Varulin Whit
= -f pleooc-)) [f - af )2N (N (a/st ( Rddulss dyo---1
(22)
N
= =
(ofle0n--i)) zo (2 Σ
(o
N
k=1
In summary,
N 3
Σ
afly2115t) of (af(v=1,,))) Vl,t.
(23)
15

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
A.3. KL--divergence proof
We want compute DE (gol y0.c-))lp(wee] Due factorization on the layers, the KL divergence reads: ]
But Lof (w (wlloo. [golus your-1)
log ( p(w))
DeL. (up 10:--1) [[p(w:)] =
1=1
(24)
Σ Dx (golwss you ())-))))((w)))])
We now expand the KL divergence DEL [golwst [lxp(]) for each layer 1. Let whyz representing the matrix entries of wi then:
golluiie -oit-l)) plwtyje)
DKL [golw)) you- [0xt-))lp(w{]] = Σ
Eny ~as (w 1o..
log
(25)
Rech
We model q (wig: you-11 ) = Vaaa (a/ ))) and p(w1)x)) = NN (1A )), rates (see in the main text). Then eq. (25) becomes:
where at, By, 7 are the dropout
(a 3/
(a
(a
DeL [up(w) you- [Pp((())
+
- - In
(7
(1/
(3/0))))
$1.fi
(26)
a
B
+
1
2
Apcii
In the last step, we used the fact that the dropout rates are positive. Combining eq. (26) with eq. (24) concludes the proof. Notice that this KI. divergence is independent on 22,5 V.
A.4. Predictive Distribution and Uncertainty Estimates
This subsection shows how to use BARNN for inference, and uncertainty quantification. In the inference case, we are interested in sampling from previous states yo...., 1 the causally consecutive one yr, while in UQ we aim to model epistemic (model weight uncertainty) and/or aleatory uncertainty (data's inherent randomness).
The model predictive distribution p(ye yt-1...., yo) is obtained by marginalizing over the weights Wi sampled from the posterior distribution:
Predictive Distribution
p 1e.1, yo) Eawrng [p(ye Yo:t-1,
wi)g (1n=-1,we)gool
yo:--1)]
(27)
In practical terms. this involves executing multiple stochastic forward passes through the neural networks and then averaging the resulting outputs (Gal & Ghahramani, 2016; Louizos & Welling, 2017).
Uncertainty Estimates The predictive distribution allows for a clear variance decomposition using the law of total variance (Depeweg et al., 2018):
of
2 (E y: yo:r-1, + Ee,, (a² yo:r-1, Wr)
(28)
epistemic uncertainty
aleatoric uncertainty
Yo:-1, the mean and variance of the distribution with wi go(we 10. and Elye yo:- 2 plyel Vi-1...., "o. w,) respectively.
Monte Carlo Estimates Samples from eq. (27), and uncertainties from eq. (28) can be computed by Monte Carlo estimates. Assume, for fixed initial state yo. to unroll an autoregressive or recurrent network for a full trajectory i-times,
16

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
with ......, D), each time sampling different wy. Let Peg of the mean and variance of plye w for the i-sample of w/ Then, the following unbiased Monte Carlo estimates can be used:
o D
Σ 1
ph
1
(29)
epistemic uncertainty
aleatoric uncertainty
A.5. Connection not Bayesian Autoregressive and Recurrent Networks
In this subsection we want to derive connection between the variational lower bound in eq. (4), and the standard log-likelihood commonly used for training autoregressive networks. We start by rewriting eq. (4):
((Φ) EE-ue(1.11 [Bwrrgalw, (oo-1) p(y. 10 w.)] DE [go [yoc-1))l((we]]]
(30)
Now we make the following assumptions:
1a.1) w. w 11 = (1.
(a.2) q(wl yo:t-1) p(w)= =6( 22) Vt (1,
Assumption (a. ensures that weights do not change over time (static weights). Assumption (a.2) fixes the weights to single value Ω (deterministic weights).
Under assumption (a.2) the KL-divergence is zero, and with assumption (a.1) the expectation reduces to single value; thus we are left with the following:
L = [[., [log (2)-)..()))
(31)
which is the standard log-likelihan ligllikel optimized in autoregressive networks.
17

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
B. BARNN PDEs Application
This section shows how BARNN can used to build Bayesian Neural PDE Solvers, provides all the details for reproducing the experiments, and shows additional results.
B.1. Data Generation and Metries
Data Generation: We focus on PDEs modelling evolution equations, although our method can be applied to vast range of timedeependent differential equations. Specifically we consider three famous types of PDEs, commonly used as benchmark for Neural PDE Solvers (Brandstetter al., 2022a;b; Bar-Sinai et al., 2019):
Burgers
dau + ud u vd
Kuramoto Sivashinsky
deu + vd, u + Ogzu Owazz u
(32)
Korteweg de Vries
0,u + v0, + Uxrzu=
where v 0 is the viscosity coefficient. We follow the data generation setup of (Brandstetter et al., 2022a), applying periodic boundary conditions, and sampling the initial conditions u² from distribution over truncated Fourier series (x) 10 k= A& sinnsnnnntt (2 & + 7h), where {Ak, Mj, 4 h random coefficients in (Brandstetter et al., 2022a) The space-time discretization parameters are reported in Table 5.
Table 5: PDE parameters setup. The discretization time and space is indicated by n, and ng respectively, tmax represents the simulation physical time, and I, is the domain length.
PDE
Thg
T2z
mma
L
Burgers
0.1
0.25
14
2x
Train
KS
0.1
0.25
14
64
Kdv
0.1
0.25
64
Burgers
0.1
0.25
32
2m
Test
KS
0.1
0.25
32
64
Kdv
0.1 0.25
32
64
Metrics: To evaluate the performance of the probabilistic solvers, we focus on different metrics:
1. Root Mean Square Error (RMSE) (Brandstetter et al., 2022b): measures the match of the ensemble prediction means and true values
2. Negative Log-Likelihood (NLL): represents the trade-off between low standard deviations and error terms, where the latter are between ensemble prediction means and true values
3. Expected Calibration Error (ECE): measures how well the estimated probabilities match the observed probabilities
Let (4, de the ensemble mean and variance predicted by the probabilistic Neural Solver for different times t, Q the Gaussian quantile function of p. then:
RMSE
Σ
[.((x)) 2x(x)
[u (x) Me (
NLL 21 2,
log(2TT(?
(33)
07
ECE Ep-U(0,1)) 1pp-10 ||p-Pobll
E s @p(x(
Pobs fg
2,
18

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
B.2. Model Architectures, Hyperparameters, and Computational Costs
All models were trained for 7000 epochs using Adam with 5 10 4 of learning rate and 10-8 of weight decay for regularization. Computations were performed on a single Quadro RTX 4000 GPU with 8-GB of memory and requires approximately one day to train all models. Unless otherwise stated, the mean and std of the predictions computed using 100 ensemble members for all models, and models are unrolled for 320 steps in inference.
We perform all experiments with a Fourier Neural Operator (Li et al., 2020) which is a standard baseline for Neural PDE Solvers. The Neural Operator is composed of 8 layers of size 64 with 32 modes and swish activation (Li et al., 2020; Brandstetter et al., 2022a) for total of approximately 1 million trainable parameters. We found that using batch normalization worsens all the uncertainty estimates. We used one linear layer to lift the input and one project the output. For dropout models (Dropout and ARD Dropout), found that applying dropout only the linear layer of the Fourier layer and on the projection layer gave the best uncertainties, we used this also for BARNN for fair comparison. In the following, we report the specific hyperparameters for each model.
Dropout We used a dropout rate of 0.2, applying it only to the Fourier linear layers. We also tested other dropout rates 0.5, 0.8, but they worsened the uncertainty. In Table we report an ablation study for different dropout rates, due to the computational burden we run the study only for one seed.
Table 6: Results for different PDEs and Dropout rates on the validation set for 140 steps unrolling. Results report RMSE, NLL, and ECE statistics.
PDE
Dropout Rate
RMSE (4)
NLL (4)
ECE (4)
0.2
0.014
-3.070
0.022
Burgers
0.5
0.030
-1.779
0.077
0.8
0.055
-1.562
0.030
0.2
0.044
-0.416
0.120
KS
0.5
0.066
-1.247
0.065
0.8
0.183
0.911
0.127
0.2
0.017
-3.350
0.135
KdV
0.5
0.040
-2.280
0.078
0.8
0.080
-1.310
0.097
Input Perturbation We employed the standard FNO discussed above and perturbed the input with Gaussian noise during training (Lam al., 2023), ut gemirr u4 + 8 max /u,cc, with (0, II) and 5 0.01. In Table 7 report ablation study for different perturbation scales, due to the computational burden we run the study only for one seed.
PDE Refiner We implemented our version of PDE Refiner, following the pseudocode provided in the paper (Lippe et al., 2024). We found the model to be very sensitive to the choice of the refinement steps R and minimum variance Omin hyperparameters, in accordance with previous studies (Kohl et al., 2024). For the PDE experiments we used R = 2, miin = 2- 10-6, Table we report an ablation study for different perturbation scales, due to the computational burden we run the study only for one seed.
The model does not have hyperparameters, we follow (Kharitonov et al., 2018) for the implementation. We used global dropout rate to have fair comparison with standard Dropout, and we applied ARD-Dropout only the Fourier layers.
ARD-Dropout
BARNN Similarly to Dropout and ARD-Dropout, we applied the dropout mask only on the Fourier fully connected layers. The dropout coefficients an are given by a specific encoder Es. In our implementation. the encoder takes as input the state ya and the time variable t. We map the state channel to a fixed dimension of 16 using one linear layer. Consequently, we map the time variable using a sinusoidal positional encoding with period T 1000 to scale and shift the state. The resulting state embedding mapped to the output shape (number of variational layers) with four other linear layers of
19

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Table 7: Results for different PDEs and Input Perturbations scales å on the validation set for 140 steps unrolling. Results report RMSE, NLL, and ECE statistics.
PDE
Perturbation
RMSE (4)
NLL (4)
ECE (4)
0.001
0.017
1.583
0.366
Burgers
0.01
0.014
-2.085
0.181
0.1
0.093
0.218
0.193
0.001
0.015
-1.048
0,188
KS
0.01
0.041
-1.942
0.110
0.1
0.404
0,413
0.078
0.001
0.014
-0.641
0.231
KdV
0.01
0.027
-2.032
0.111
0.1
0.155
0.115
0.189
Table &: Results for different PDEs and Refiner Hyperparameters on the validation set for 140 steps unrolling. Results report RMSE, NLL, and ECE statistics.
PDE
muin
R
RMSE (4)
NLL (L)
ECE (L)
9-10-6
2
0.019
44.18
0,216
2.10
0.106
184.4
0.547
2.10-6
0.021
16.64
0.196
2-10-7
2
0.014
86.24
0.248
Burgers
2-10-7
3
0.026
240.4
0.239
2-10-7
4
0.023
100.3
0.367
2-10-6 6
0.027
6.089
0.167
2.10-6
0.081
76.76
0.232
2-10-6
0.041
8,456
0.179
2.10-7
2
0.024
9.330
0.189
KS
2-10-7
3
0.025
6.977
0.183
2.10-
4
0.037
6.250
0.175
1.10-6
2
0.011
10.17
0.402
2.10-6
3
0.031
72.10
0.343
2.10-6
4
0.022
39.90
0.285
2-10-7
2
0.043
86.74
0.356
KdV
2.10-7
3
0.042
192.5
0.456
10
0.042
92.33
0.321
size 16 interleaved with swish activation and final sigmoid activation return the dropout probabilities The dropout rates can be obtained by computing P:/(1 P:). The built encoder contains less than 50K parameters, thus a tiny network compared to modern Neural Solver networks, not significantly affecting the training time.
B.3. Additional Results
Comparison of Model Accuracies Table presents the RMSE results for the analysed PDEs across various models, The based methods (BARNN, ARD, and Dropout) consistently show the best performance, indicating that Bayesian approaches can be advantageous not only for uncertainty quantification but also for pointwise accuracy. The differences in RMSE are relatively small among the top-performing models, suggesting that they generally perform similar
20

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Burgers
KS
KdV
0.12
0.10
NILL
ECE
0.08
1
0.06
0.04
20
40
60
80
100
20
40
60
80
100
number of ensemble memebers
number of ensemble memebers
Figure 6: Caption
Table 9: RMSE for Burgers, KS and Kdv PDEs for different Neural Solvers. The mean and std are computed for four different random weights initialization seeds. Solvers are unrolled for 320 steps.
Model
RMSE (4)
Burgers
KS
KdV
BARNN
0.019
0.217-00.047
0.061 10.026
ARD
0.020
0.161 1.0052
0.000000000
Dropout
0.019=00000
0.172200.02
0.053#00.02
Perturb 0.022
0.222=00000
0.0000
Refiner 0.042=00.224
0.407=0.402
0.775
level. However, as discussed in Section 4.2, RMSE alone is insufficient for evaluating neural PDE solvers, particularly in uncertainty quantification scenarios, as low RMSE values can obscure inaccurate uncertainty estimates from the solver (see Figure 3).
Test number of ensemble members in BARNN We extend the analysis presented in Section 4.2 by examining the convergence behaviour with varying numbers of ensemble members. Figure 6 illustrates how the negative log-likelihood (NLL) and expected calibration error (ECE) evolve as the ensemble size increases. Similar to the trend observed in RMSE (refer to Figure 4), the NLL and ECE metrics stabilize after approximately ensemble members. This indicates that, beyond this point, adding more members yields diminishing returns terms of performance improvement. These results highlight that BARNN is capable of generating precise and well-calibrated predictions, offering sharp uncertainty estimates without requiring large ensemble. The saturation of these statistics reinforces the efficiency of BARNN, making it robust method for uncertainty quantification with minimal computational overhead.
Solution rollouts Figures 7, 8, and showcase examples of 1-dimensional rollouts over 320 steps, generated using 100 ensemble members with BARNN. As anticipated, the error accumulates over time, which is characteristic of the autoregressive nature of the PDE solver. However, the model's variance also increases in tandem, highlighting its ability to identify regions with higher errors and demonstrating impressive adaptability. Despite the growing errors, the overall solution remains closely aligned with the predictive mean, reflecting strong pointwise accuracy. Mcanwhile, the variance effectively captures the rise errors, providing reliable indication of uncertainty and reinforcing the model's adaptability in challenging regions of the solution space. This dual capability of maintaining accuracy while adapting to uncertainty underscores robustness of the BARNN approach.
21

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Ground Truth y
Mcan μ
2
19.3
12.9
-1
6.40
-2
-3
0.00
Variance O²
Error ly - -41?
0.008
25.7
19.3
0..006
12.9
0.004
6.40
0.002
0.00
0.000
0.00
0.90
1.80
2.70
3.60
4.50
5.40
0.00
0.90
1.80
2.70
3.60
4.50
5.40
Figure 7: Exemplary of I-dimensional Burgers rollout.
Ground Truth
Mean μ
19.2
12.8
-1
6.40
-2
0.00
-3
Variance 2
Error ly 24
25.7
19.2
12.8
6.40
0.00
0.00
9.00
18.2
27.2
36.5
45.5
54.8
0.00
9.00
18.2
27.2
36.5
45.5 54.8
Figure 8: Exemplary of I-dimensional
rollout.
22

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Ground Truth y
Mean μ
25.7
19.3
12.9
6.40
0.00
Variance o²
Error ly u1
0.6
25.7
19.3
12.9
0.2
6.40
0.00
9.00
18.2
27.2
36.5
45.5
54.8
0.00
9.00
18.2
27.2 36.5
45.5
54.8
Figure 9: Exemplary of I-dimensional Korteweg de Vries rollout.
C. BARNN Molecules Application
This section shows how BARNN can be used to build Bayesian Recurrent Neural Networks, provides all the details for reproducing the experiments, and shows additional results.
C.1. Data Generation and Metries
The experiments focus on unconditional molecule generation using the SMILES syntax, a notation used to represent the structure of a molecule as line of text. We use the data set provided (Özşelik et al., 2024), which contains a collection of 1.9 M SMILES extracted from the ChEMBL data set (Zdrazil et al., 2024). Before training, the SMILES strings were tokenized using regular expression, containing all elements and special SMILES characters, e.g., numbers, brackets, and more. Therefore, each atom, or special SMILES character, is represented by single token.
Data Generation:
Metrics:
To evaluate the performance of the language models, we focus on different metrics:
1. Validity: the percentage generated SMILES corresponding to chemically valid molecules
2. Diversity: the percentage of unique Mureko scaffolds among generated molecules.
3. Novelty: the percentage of generated molecules absent from the training dataset
4. Uniqueness: the percentage structurally-uniquu molecules among the generated designs
Moreover, we compute the Wasserstein distance between the molecular properties distribution for the language model and the ChEMBL dataset to ensure molecular properties are matched. Similarly to (Segler et al., 2018) we analyse the molecular weight, H-bond donors and acceptors, rotatable bonds, LogP, and TPSA.
C.2. Model Architectures, Hyperparameters, and Computational Costs
All models were trained for 12 epochs using Adam with 2- 10-4 learning rate and 10 1 of weight decay for regularization. The training was distributed across four Tesla P100 GPUs with 16-GB of memory each, and we used a batch size of 256 for memory requirements. Approximately hours are needed to train one model.
23

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
1.
BASNK
3.3
6.
Desstty
0.033
20
02
0.000
00
9.0
0-
Molecuaar
500
14 H-boad
H H-bond
20
0.15-
0.23
0.010
10
5.15
0.004
-00-
2.03
0.002
2,03
0.000
-100
Rotable Bonds
25
40
Log?
TPSA
220
400
Figure 10: t-SNE representation of different molecular prop- erty deseriptors for BARNN and chEMBL molecules. The distributions of both sets overlap.
Figure 11: Distribution density for each molecular property descriptor for BARNN and chEMBL molecules. The distri- butions for different descriptors overlap significantly.
We perform all experiments with an LSTM gate mechanism (Schmidhuber et al., 1997) for recurrent networks. The tokens were mapped to a one-hot encoding of the size the vocabulary, processed by the language model, and mapped back to the vocabulary size by a linear layer. We used three LSTM-type layers with the hidden dimension of 1024 accounting for 25 M parameters model overall. For the dropout L.TTM we used 0.2 as dropout coefficient, following indications from previous works (Segler al., 2018). For BARNN the encoder composed of two fully connected (linear) layers of size 128 interleaved by LeakyReLU activation which we found to be the best, and final sigmoid activation to return the dropout probabilities Pt- The dropout rates can be obtained by computing ar Pt/(1 Pt) elementwise. Once models were trained, 50K molecules were sampled to compute the statistics reported in the experiment section.
C.3. Additoonal Results
In this section, we provide additional results for molecule generation using BARNN.
Molecular properties predictions We compute six molecular property descriptors-namety Molecular Weight (MW), H-bond donors (HBDs), H-bond acceptors (HBAs), Rotable Bonds, LogP, and Topological Polar Surface Area (TPSA)-for for both BARNN- -eeneregeee molecules and the chEMBL dataset. To evaluate whether the molecular properties of the BARNN- generated molecules align with those of the chEMBL data, we employ t-SNE (1-Distributed Stochastic Neighbor Embedding) (Van der Maaten Hinton, 2008), a dimensionality reduction technique, to visualize the two-dimensional latent space (Figure 10). The t-SNE plot reveals that the distributions of the two sets significantly overlap, suggesting that BARNN effectively captures the underlying molecular properties. For more detailed analysis, we perform Kernel Density Estimation (KDE) for each molecular property descriptor and present the resulting distributions for both BARNN and chEMBL in Figure 11. The KDEs from both datasets almost perfectly overlap, even for distributions with multiple peaks (c.g., H-bond donors and acceptors), indicating that BARNN not only captures the overall distribution but also the single characteristics of the molecular properties with high fidelity.
Robustness to temperature change Temperature is a crucial hyperparameter recurrent neural networks that influences the randomness of predictions by scaling the logits before applying the softmax function. As T 0, the model tends to select the most likely element according the Conditional Language Model (CLM) prediction, leading to more deterministic outputs. Conversely, the temperature increases, the selection becomes more random, higher temperatures flatten the logits. In our analysis, we evaluate the robustness of BARNN to temperature variations during the sampling process, specifically examining the validity and uniqueness of the generated sequences, as shown in Figure 13. Across all models, we observe a general trend: both validity and uniqueness decline as the temperature increases. At lower temperatures, BARNN outperforms the baseline models, maintaining higher validity and uniqueness. However, at very high temperatures, BARNN and the base LTM show comparable performance, with both models achieving similar levels of validity and uniqueness.
24

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
BARNN SMILES LSTM
BARNN
0.9
0.9
SMILES LSTM
LSTM
LSTM
Validity
0.8
Uniquenes
0.8
0.7
0.7
0.6
0.6
1.0
1.2
1.4
1.6
1.0
1.2
1.4
1.6
Temperature
Temperature
Figure 12: t-SNE representation of different molecular prop- crty descriptors for BARNN and chEMBL molecules. The distributions of both sets overlap.
Figure 13: Distribution density for each molecular property descriptor for BARNN and chEMBL molecules. The distri- butions for different descriptors overlap significantly.
Figure 14: Samples of generated molecules with BARNN.
D. Algorithms and Software
This section reports the pseudocode algorithms for training and performing inference using a BARNN model and the software used.
25

---

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
D.I. Pseudocodes
Algorithm 1 Local reparametrization trick: Given H minibatch of activations for layer / 1...., reparametrization trick to compute the linear layer output. It requires the dropout rate coefficient a for layer 1, and the Neural Operator static weights S2 for layer I. For the first layer, activation corresponds the state input.
......,, L use the local
M < a H Q4 (Compute the ] V + (a 2 [Compute the variance, power is taken elementwise]
E ~ (0, I) (Sample Gaussian state)
return M + V E
Algorithm 2 Training Bayesian Neural PDE Solvers: For given batched input data trajectory, Neural Solver and variational posterior encoder, we draw random timepoint i, get our input data trajectory. compute the dropout rates with the posterior encoder, perform forward pass using the recrameppaaa trick, and finally, perform the supervised learning task with the according labels plus the KL. regularization. (vb.... y are the data trajectories, NO the neural operator architecture with static weights, and Es is the variational posterior encoder.
while not converged do
- [[1 [Draw uniformly a random ending point}
ag < (y)")) (Compute dropout rates for each NO layer)
y1 < NO S2, a.) [Apply local reparametrization trick to linear layers)
C(S2, vb) < 1 ye1l2 D (ar) (Compute the loss using DKL from eq. (9))
Optimize Ωt zb by descending the gradient of L
end while
Algorithm Training Bayesian Recurrent Neural Networks: For given batched input and output data taken from the sequence we compute the dropout rates with the posterior encoder, perform forward pass using the local-reparametrization trick, and finally, perform the supervised learning task with the according labels plus the KL regularization. {ye.... y k= are the sequences tokenized, RNN is the recurrent architecture with Ω static weights, o, is the RNN probability output at step t, and Es is the variational posterior encoder.
while not converged do h < 0 [Initialize to zero hidden ) for t E [0...., - do
a + (3)) (Compute dropout rates for each RNN layer)
a < F (h) (Compute dropout rates for each RNN layer)
o h RNN(Y: h; 12, a², att) (Apply RNN forward with local reparametrizationn trick)
end for
C(S vb) + CrossEntropyLoss(( y1.T) (az( (Compute the loss)
Optimize Ω ab by descending the gradient of f
end while
D.2. Software
We perform the PDEs experiment using PINA (Coscia et al., 2023) software, which Python library based on PyTorch (Paszke al., 2019) and PyTorch Lightning (Falcon & The PyTorch Lightning team, 2019) used for Scientific Machine Learning and includes Neural PDE Solvers, Physics Informed Networks and more. For the Molecules experiments we used PyTorch Lightning (Falcon & The PyTorch Lightning team, 2019) and RDKit for postprocessing analysis of the molecules.
26


Reviews:
[{"review_id": "qkUgiZLace", "reviewer": "Reviewer_4bFX", "summary": "This paper introduces BARNN, a framework designed to turn any autoregressive or recurrent deep learning model into a Bayesian version with minimal modifications. The authors propose jointly modeling both the states (e.g., tokens, PDE solutions, molecule strings) and the model’s weights as they evolve in time. By extending existing approaches to variational Bayesian inference, this work derives a novel temporal version of a variational lower bound and leverages a time-dependent extension of Variational Dropout. The work demonstrates BARNN in two main applications: (1) Neural PDE solvers for uncertainty quantification, and (2) RNN-based molecule generation, with experiments showing stronger calibration of uncertainties, improved long-range dependency handling, and better alignment with observed statistical properties of the underlying data.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are delighted that the reviewer finds our approach *elegant and general, allowing flexible application across domains* and the tVAMP prior and time-dependent dropout mechanics *novel and easy to implement*. We appreciate that the reviewer considers that our *claims are generally supported by well-chosen experiments*. Finally, we thank the reviewer for the detailed comments and suggestions, which we address point by point below:\n\n**From a high-level view, the presented proofs appear logically consistent. However, a more detailed or step-by-step check of the proofs would help confirm there are no hidden assumptions that might be restrictive.**\n- We think the reviewer understood correctly the core of the proofs. For the temporal variational lower bound proof (Appendix A.1), given the specifics of the generative model distribution and the posterior distribution, no assumptions (e.g. independence or Markovian-memoryless properties) are made for deriving the lower bound. For the temporal VAMP prior proof (Appendix A.2), the approximation comes in eq. (19) while the general formulation of eq. (18) does not make any assumption. \n\n**Weaknesses:**\n**The method’s computational overhead for extremely large models (like multi-billion-parameter LLMs) is not fully explored**\n- We acknowledge the fact that we did not test, due to computational resources, the methodology on large-scale models. However, we want to highlight that our methodology can scale to large models due to the decoupling between posterior and main model weights. Indeed, the temporal dropout coefficients are obtained from a posterior model with shared weights, which contains order of magnitude lower number of parameters, compared to the main network which is left untouched in the parameter size.\n\n**Weaknesses:**\n**Real-world PDEs often involve high-dimensional grids (2D, 3D); the current experiments may not fully reflect potential complexities in multi-dimensional domains.**\n- We thank the reviewer for the suggestion, which we also find interesting to investigate. We believe our results will also scale to 2D and 3D experiments as we decouple the dropout rates from the main model weights. In our paper, we did not apply to 2D and 3D experiments for computational resources limitations (as GPU memory was limited).\n\n**Weaknesses:**\n**More thorough ablations could be done on alternative priors besides the new “tVAMP prior**\n- We thank the reviewer for this suggestion, which is shared with reviewer WqtE. To show how tVAMP excels over standard prior and fixed dropout rates, we conducted a test on a synthetic time-series dataset. This test was handmade to show how the prior affects the network modelling (i) for varying amplitudes and frequencies in the states, (ii) understand the effect of time adaptivitly of the prior. We also add a simpler static dropout model for reference in the hope of answering the question raised in Other Comments or Suggestions. The results report RMSE, NLL, and ECE statistics. *Static* method reports the RMSE obtained if the initial state is not propagated, while *MLP* is the base (non-bayesian) architecture (2 layer of with 64 with relu activation). BARNN uses the same MLP architecture, and it is ablate on different priors.\n\nDataset:\n\\begin{align}\n\\begin{cases}\n    x_t &= x_{t-1} + \\frac{3\\pi}{100}, \\quad x_0=0\\\\\\\\\n    y_t &=\\frac{1}{5}\\sum_{j=1}^5\\sin(\\alpha_i x_t +\\beta_i), \\quad \\alpha_i\\sim U[0.5, 1.5],\\, \\beta_i\\sim U[0, 3\\pi],\\, t\\in\\{1, 2, \\dots, 100\\}\n\\end{cases}\n\\end{align}\n\nTable:\n| **Method**        | **Prior**      | **MSE** (↓)         | **NLL** (↓)           | **ECE** (↓)           |\n|-------------------|----------------|---------------------|-----------------------|-----------------------|\n| Static            | -              | 0.490 ± 0.000       | -                     | -                     |\n| MLP               | -              | 0.081 ± 0.011       | -                     | -                     |\n| Dropout (p=0.5)   | -              | 0.072 ± 0.004       | 0.593 ± 0.461         | 0.084 ± 0.010         |\n| Dropout (p=0.2)   | -              | 0.048 ± 0.004       | -0.075 ± 0.004        | 0.068 ± 0.009         |\n| BARNN             | log-uniform    | 0.045 ± 0.003       | -0.092 ± 0.064        | 0.050 ± 0.016         |\n| BARNN             | tVAMP          | **0.043 ± 0.001**   | **-0.166 ± 0.019**    | **0.049 ± 0.008**     |\n\n**Why did you pick a tVAMP prior specifically, instead of simpler approximations?....**\n- We thank the reveiwer for the interesting question. We picked the tVAMP prior because during preliminary experiments we found that using a standard log-uniform prior (commonly used for VD) did not give us good uncertainties, and we attribute it to the fact that the prior pushed to heavily spared network weights.  This is also confirmed by the time-series experiment results or the PDE (Table 1) results, where log-uniform and ARD priors obtain suboptimal results.", "claims_and_evidence": "1. BARNN can be applied to “any” autoregressive or recurrent model with minimal architectural changes. The paper shows that only a small set of modifications is needed: introducing time-dependent “dropout coefficients” and a new prior derived from the aggregated posterior. They illustrate this with PDE models and an LSTM-based molecule generator, both using standard training objectives (MSE for PDE forecasting, cross-entropy for next-token prediction).\n\n2. BARNN provides better-calibrated and sharper uncertainties than existing methods. In experiments with PDE solvers, BARNN outperforms Monte Carlo Dropout, Input Perturbation, and other baselines on negative log-likelihood and expected calibration error.\n\n3. BARNN enhances long-range dependency modeling. In the molecule-generation experiments, the authors highlight ring-closure errors as a challenging long-range dependency in SMILES. BARNN yields fewer ring-closure mistakes compared to baseline LSTMs, suggesting improved capacity to track tokens that appear far apart in the sequence.\n\n4. BARNN-derived ensembles require fewer samples for stable statistical estimates and can revert to a single forward pass via MAP if uncertainty is not needed. The paper demonstrates that with about 30 ensemble members, the PDE solver’s predicted mean and variance converge. If uncertainty is not required, using the MAP estimate makes it as fast as a single forward pass, matching deterministic baselines on RMSE.\n\nOverall, the claims are generally supported by well-chosen experiments that compare BARNN with multiple baselines. However, some details (e.g., how BARNN might scale to extremely large language models) are not deeply addressed, and real-world performance in large-scale domains is left for future work.", "methods_and_evaluation_criteria": "The evaluation criteria appear sound for the chosen tasks. For additional confidence, the paper checks multiple seeds, compares ensemble sizes, and tests different PDEs. Overall, the methodology and metrics are consistent with accepted best practices in both scientific PDE modeling and generative modeling in chemistry.", "theoretical_claims": "1. The key theoretical contribution is a novel variational lower bound that explicitly models time-varying network weights and states, together with a new prior inspired by VAMP.\n2. The authors provide a derivation that connects BARNN’s objective to the standard VAE framework.\n3. The correctness of the proofs hinges on: factorization of the generative model over states and weights, the application of the local reparameterization trick for weight sampling, the derivation of the “best prior” in the sense of aggregated variational posteriors.\n\nFrom a high-level view, the presented proofs appear logically consistent. However, a more detailed or step-by-step check of the proofs would help confirm there are no hidden assumptions (e.g., independence assumptions) that might be restrictive.", "experimental_designs_or_analyses": "- The PDE examples use well-known benchmark equations (Burgers, Kuramoto–Sivashinsky, Korteweg–de Vries) and evaluate multi-step rollouts (320 steps), which is a reasonable stress test.\n- For molecule generation, a large ChEMBL-based dataset is used, covering real-world chemical diversity. Standard metrics in de novo drug design—validity, novelty, uniqueness, distribution matching—are computed.\n- The comparison includes widely used baseline methods for uncertainty quantification (Dropout, Input Perturbation, etc.) and for molecule generation (standard RNNs).\n- One limitation is that the PDE results rely on single-dimensional or low-dimensional PDE domains (1D PDE experiments). Similarly, the molecule generation uses a single tokenization approach (SMILES). Nonetheless, these design choices are typical entry points in the literature.", "relation_to_broader_scientific_literature": "This paper broadly builds on two themes: (1) Bayesian methods for deep networks (Variational Dropout, VAMP priors, etc.) and (2) autoregressive structures in PDEs, language modeling, or molecule generation.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n- The approach is elegant and general, allowing flexible application across domains.\n- Empirical performance is strong, particularly regarding calibration and the ability to generate valid molecules.\n- The derived prior (tVAMP) and time-dependent dropout mechanics feel both novel and easy to implement.\n\nWeaknesses:\n- The method’s computational overhead for extremely large models (like multi-billion-parameter LLMs) is not fully explored.\n- Real-world PDEs often involve high-dimensional grids (2D, 3D); the current experiments may not fully reflect potential complexities in multi-dimensional domains.\n- More thorough ablations could be done on alternative priors besides the new “tVAMP prior.”", "other_comments_or_suggestions": "It could be helpful to include an additional ablation that compares time-dependent dropout with a simpler “static” dropout factor, clarifying precisely how much performance gain is due to the time dimension vs. simply having an advanced prior.", "questions_for_authors": "Why did you pick a tVAMP prior specifically, instead of simpler approximations (e.g., a fixed log-uniform prior for dropout)? Did you try them, and how did they perform?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mjSpNsdQuY", "reviewer": "Reviewer_SBfz", "summary": "This article proposes a new Bayesian recurrent neural network, mainly by introducing variational dropout into recurrent neural networks. Experiments can prove the model's ability to quantify uncertainty.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and comments, and we appreciate that the reviewer finds the paper *well written* with the experiments *sufficient to demonstrate the validity of the model*. We address the reviewer's comments point by point below in the hope of clarifying some misunderstandings, improving the paper and eventually raise their score:\n\n**Some work attempts to probabilistic model attention in order to obtain good uncertainty estimation ability**\n- We thank the reviewer for the references; we find  Bayesian Attention Modules very relevant and will cite them in the related works. A key difference between Bayesian Attention Modules and our methodology is that our method explicitly models uncertainties in time using a joint state-weight model and applies to any autoregressive or recurrent model, while Bayesian Attention Modules are a technology specifically tailored to Transformers models, and the Bayesian weights are static.\n\n**Weakness**\n**The proposed method can be seen as an application of variational dropout to an autoregressive model but does not provide any new insight.**\n- We thank the reviewer for their time in reading the paper and finding parallelisms with Variational Dropout (VD). We understand that our method shares similarities with VD, from which we inherit the scalability to large networks. However, in variational dropout methods (i) the weights do not evolve in time, and (ii) the prior is empirical, while we provide a time-dependent new prior, which is the best for the objective in eq. (4). As reviewer 4bFX highlighted, *the derived prior (tVAMP) and time-dependent dropout mechanics feel both novel and easy to implement*, and the reviewer WqtE believes that our methodology *is of pragmatic interest due to its broad applicability in uncertainty-aware sequence modeling*. Finally, we also found that static weights lead to less accurate uncertainty metrics and adaptivity. For example, see Table 1, where ARD Dropout or Dropout shows lower accuracy in terms of NLL and ECE compared to our methodology.\n\n**Autoregressive models based on attention mechanisms are not discussed.**\n- We thank the reviewer for pointing out attention-based models, which we agree we did not mention exhaustively. We tried to cover a wide variety of references on autoregressive models, not focusing only on a specific model. Nevertheless, we did mention the Vaswani, A. et al. Attention Is All You Need paper for transformer reference, and the Radford, A., et al. Language Models are Unsupervised Multitask Learners for moder decoder-only language models. If the reviewer thinks we missed other relevant papers which will improve the manuscript, we will be happy to add them if pointed to.\n\n**Questions For Authors:**\n**Is there any difference between probabilistic modeling of autoregressive models and normal neural network models, like convolutional neural networks?**\n- We thank the reviewer for the very interesting question. We did find that standard techniques used to model uncertainties in non-autoregressive models tend to obtain less accurate uncertainty metrics when applied to autoregressive solvers (see, for instance, Table 1. where Ensemble Dropout does not obtain as good performances as our methodology). This is mainly due to error accumulation and long-term dependencies, which non-autoregressive models do not need to deal with. \n\n**Why not use Transformer, a more general and powerful autoaggressive neural network?**\n- We thank the reviewer for the question, which we are glad to discuss. In our experiment, we focused in the AI4Science field since we believe uncertainty quantification will be key there. In the AI4Science field, we choose PDEs and Molecule generation as tasks since they are very different and use different models. For PDEs, autoregressive models are mostly based on multiple-step predictions using a Neural Operator. For molecule generation, RNNs were first applied, and just more recently, Transformers have been used, and it is not clear that transformers are more powerful in this specific task. In summary, with the paper, we wanted to show a general methodology for Bayesian sequence modelling,  hoping to inspire potential applications to future work given the simplicity to implement and the generality of the methodology.\n\n**Why not experiment on text to better illustrate the effectiveness of the model on large data sets?**\n- We agree with the reviewer that the NLP task is an interesting avenue of research. Even though not an NLP task, we did employ text for the molecule generation task where the SMILES syntax is learned by the model to generate new molecules. For a more specific NLP task (e.g. language translation), a complication for us was how to define uncertainty properly, which seems to be an ongoing open research. Indeed, we choose our experiment tailored to support each claim in the paper, as reviewer WqtE also highlighted (*all the claim are supported by clear and convincing evidence*).", "claims_and_evidence": "This paper builds a variational Bayesian autoregressive and recurrent neural network. BARNNs aim to provide a principled way to turn any autoregressive or recurrent model into its Bayesian version. BARNN is based on the variational dropout method, allowing us to apply it to large recurrent neural networks as well.  I think is correct .", "methods_and_evaluation_criteria": "Quantifying uncertainty in autoregressive models is a crucial and key issue.", "theoretical_claims": "A posteriori inference technique based on variational autoencoder, I think, is correct", "experimental_designs_or_analyses": "The experiments can verify the method is effective and are taken on different tasks.", "relation_to_broader_scientific_literature": "Some work attempts to probabilistic model attention in order to obtain good uncertainty estimation ability. \n\n[1] Bayesian Attention Modules; Xinjie Fan, Shujian Zhang, Bo Chen, and Mingyuan Zhou; NeurIPS 2020: Advances in Neural Information Processing Systems, Dec. 2020. \n[2] Deng, Yuntian, Yoon Kim, Justin Chiu, Demi Guo, and Alexander Rush. \"Latent alignment and variational attention.\" Advances in neural information processing systems 31 (2018).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strongth \n1. Uncertainty modeling is a key problem in neural networks, and I think the paper as a whole is workable, and the results look good.\n\n2. The experiments are sufficient to demonstrate the validity of the model.\n\n3. The paper is well wrightten.\n\nWeakness\n1. The proposed method can be seen as an application of variational dropout to an autoregressive model but does not provide any new insight. \n\n2. Autoregressive models based on attention mechanisms are not discussed.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Is there any difference between probabilistic modeling of autoregressive models and normal neural network models, like convolutional neural networks?\n\n2. Why not use Transformer, a more general and powerful autoagressive neural network?\n\n3. Why not experiment on text to better illustrate the effectiveness of the model on large data sets?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TCo4mESDN2", "reviewer": "Reviewer_fti8", "summary": "This paper addresses the problem of uncertainty quantification with autoregressive and recurrent neural neworks. Variational Bayes method is applied to infer the posterior distribution of network parameters, and further techniques are developed using variational dropout methods. Applications are made to PDE solving and molecular generation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's time and effort in making constructive comments and providing us with additional references. We are glad that the reviewer finds our experiments *nice* and acknowledges that *the choices of applications are real-world driven*. We answer below the questions raised by the reviewer, and we also ask if there are other concerns about the paper which not lead to a higher score in the initial review round; in that case, we will be glad to discuss.\n\n**For Eq.(6), why is the normal distribution assumed to have variance being the square of its mean?**\n- Thank you for pointing out this crucial step in our methodology. We choose to model the variance as the square of the mean, as also done in vanilla Variational Dropout, for efficiency proposes. This parametrization allows us to not use extra model parameters, permitting scaling to very large neural networks. Also, this approach allows us not to change the original (non-Bayesian) model and add the BARNN methodology on top of it, which, as  4bFX highlighted, is “easy to implement”\n\n**I do have a question regarding he confidence intervals presented in Figure 2. It is said in the paper that the presented CIs are 99.7%, which are still very narrowly concentrating near the mean prediction/posterior predictive mean. To me, it appears that the variational inference approach might have signficantly underestimated the uncertainty.**\n- We thank the reviewer for the comment on the CIs. We agree with the reviewer that in the picture, the confidence interval seems concentrated near the mean prediction/posterior predictive mean. However, this is mainly a plotting issue because we used a very thick line width to have three nice displayed figures side by side. To convince you, please look at the ECE and NLL in Table 1. Underestimating the variance (as in the case of Refiner, for example) would lead to a very high ECE and exploding NLL (which is not our case).\n\n**The design of the factorization format of variational posterior in Eq.(3) reminds of the line of work on modular Bayes inference, where we deliberately cut off certain dependencies between parameter blocks and parts of observations, in order to improve computational efficiency / reduce potential influence of model misspecification due to incorrect inclusion of certain dependencies.**\n- We thank the reviewer for pointing out this reference, which we believe is relevant. Indeed, we cut off non-causal dependencies in eq (3) (i.e. the current weights are only dependent on the past). We will look more in depth into modular bayes inference and add a connection to eq. (3) in the final manuscript. If there are other concerns or suggestion, we are happy to discuss.", "claims_and_evidence": "The motivation is clear and the use of variational inference for infering autoregressive and recurrent neural network parameters is a reasonable choice.\n\nEq.(3) on the factorization form of variational posterior appears to a good choice, which also associates with the idea of modular Bayes. More to be discussed later.\n\nFor Eq.(6), why is the normal distribution assumed to have variance being the square of its mean?", "methods_and_evaluation_criteria": "Evaluations are reasonable, ranging over both metrics on predicitve accuracy and metrics on uncertainty quantification.", "theoretical_claims": "The derivations of ELBO and related formulas look correct.", "experimental_designs_or_analyses": "The experimental designs look nice and the choices of applications are real-world driven.\n\nI do have a question regarding he confidence intervals presented in Figure 2. It is said in the paper that the presented CIs are 99.7%, which are still very narrowly concentrating near the mean prediction/posterior predictive mean. To me, it appears that the variational inference approach might have signficantly underestimated the uncertainty.", "relation_to_broader_scientific_literature": "The design of the factorization format of variational posterior in Eq.(3) reminds of the line of work on modular Bayes inference, where we deliberately cut off certain dependencies between parameter blocks and parts of observations, in order to improve computational efficiency / reduce potential influence of model misspecification due to incorrect inclusion of certain dependencies.", "essential_references_not_discussed": "Here is an example reference on modular Bayes: \n\nBayarri, M. J., J. O. Berger, and F. Liu. \"Modularization in Bayesian analysis, with emphasis on analysis of computer models.\" Bayesian Anal. 4(1): 119-150", "other_strengths_and_weaknesses": "See above.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "See Claims And Evidence & Experimental Designs Or Analyses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VnBKBKHBVS", "reviewer": "Reviewer_WqtE", "summary": "This paper introduces a Bayesian version of an RNN by modeling the time-dependent weights as a reparameterized random variable, where the posterior is learned in an amortized way using a time-dependent ELBO. A specific variational posterior parameterization is provided, where each layer’s weights are parameterized via multiplicative Gaussian dropout with a time-dependent scaling factor as parameters. These dropout scaling factors evolve overtime and are generated by an encoder network that conditions on past observations. Instead of a fixed isotropic Gaussian prior, the model employs a Temporal Variational Mixture of Posteriors (tVAMP) prior, which aggregates past variational posteriors to improve uncertainty modeling and calibration. The framework is applied to Neural PDE solvers and molecular language models (unconditional SMILES generation), demonstrating better uncertainty quantification (lower NLL) in PDE modeling and higher validity in unconditional molecular generation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for their time and valuable feedback. In particular, we are delighted that the reviewer believes our work *is of pragmatic interest due to its broad applicability in uncertainty-aware sequence modeling*, and that their appreciate the experiment section. We answer the questions raised below point by point:\n\n**Weak claim(promise):\nautoregressive models are also prone to overfitting to the specific tasks they are trained on, challenging their application outside their training domain. -> Not explicitly demonstrated**\n- We thank the reviewer for analysing in depth our results section, we appreciate the time and effort. We agree that we did not stress in our experiment analysis autoregressive model overfitting. However, we want to highlight that our Table 2 results show the overfitting behaviour of autoregressive models. In particular, SMILES LSTM (it has dropout on the recurrent layers) and BARNN can be thought of as regularized versions of the vanilla LSTM. The latter, compared to the regularised models, is more prone to reproduce molecules already seen during training, shown by lower novelty and uniqueness, which is an indication of overfitting to the training data distribution.\n\n**Essential References Not Discussed:\nI think literature [1] , which also introduces a variational Bayesian scheme for RNNs, is highly relevant and should be discussed.**\n- We thank the reviewer for pointing out the methodology to us, which shares the similarity of modelling uncertainties in RNNs, and we will add in section 2. However, differently from our methodology, the approach is only suited to RNNs, while our approach can be applied to various autoregressive or recurrent models. Furthermore, but most importantly, the model in [1] can hardly scale to large networks mainly due to two reasons (i) weights are sampled directly, in contrast we use the local-reparametrization allowing scaling to large networks (ii) the number of parameters doubles because they model mean and standard deviation directly, while we use a variance proportional to the mean.\n\n**Weaknesses:\nHard to position in context: The incorporation of Bayesian methods into RNNs has a long history, with key prior work (see \"Essential References Not Discussed\") that is not acknowledged or discussed.**\n- See above in “Essential References Not Discussed” answer.\n\n**Weaknesses: Combined with the presentation style, this gives the impression of overclaiming novelty, as if applying Bayesian inference to RNNs is itself a novel contribution.**\n- We thank you for the comment, and we do not want to claim that applying Bayesian inference to RNNs is our novel contribution. However, to the best of our knowledge, we are the first approach that can transform any autoregressive or recurrent model into a Bayesian one with minimal modification, obtaining similar or superior performances to non-Bayesian counterparts, scale to large networks and at the same time provide UQ. \n\n**Weaknesses: Ablation study of the prior: The tVAMP prior is a key contribution, yet the paper does not include an ablation study comparing it to a simpler isotropic Gaussian prior.**\n- We thank the reviewer for the suggestion, which is shared with reviewer 4bFX. We conducted a test on a synthetic time-series dataset to show how tVAMP excels over standard log uniform prior and fixed dropout rates. A simple isotropic Gaussian would not lead to a KL which is independent on the likelihood model parameters, therefore losing the regularization property of VD. Please see in reviwer's 4bFX answer the result table and explanation (not inserted here for words limit).\n\n**Weaknesses: Motivation for using RNNs over stronger sequence models: While this may not be a direct weakness, it is unclear why the paper focuses on Bayesian RNNs when stronger sequence models (e.g., S4) exist.**\n- We thank the reviewer for the comment, also suggesting a possible application to BARNN for SSMs. We wanted to focus on RNNs because the model in eq. (2) (3) is heavily recurrent. Of course, extending to SSMs and Transformers models would be interesting, but a specific ancoder $E$ is needed, and further research needs to be conducted. In summary, with the paper, we wanted to show a general methodology for Bayesian sequence modelling, hoping to inspire potential applications to future work.\n\n**Eq.~8: should it be $E_{\\phi}(\\mathbf{y}_{0:t-1}^k)$?**\n- Yes, we write it more compactly as $\\alpha^l_t(\\mathbf{y}_{0:t-1}^k)$, we will add a note on the paper, thank you for spotting it!\n\n**Intuition Behind Long-Term Dependency Claims**\n- We thank you for the interesting question, and we claim that this is due to the time-adaptivity of the model weights. Indeed, both SMILES LSTM (with Dropout) and standard LSTM obtain similar performances, worse than BARNN, which only has temporal adaptivity differently from them.\n\n**Practical Setting of $N$ in Eq.~(8)**\n- Correct! $N$ is set to the batch size. We will update the paper to be clearer.", "claims_and_evidence": "Overall, all the claim are supported by clear and convincing evidence. \n\n\nStrong claim:\n- BARNN offers greater accuracy -> Table 1\n- BARNN calibrated and sharp uncertainty estimates -> Table 1, Fig. 2\n- BARNN excels in modelling long-range molecular dependencies compared to related methods. -> Fig. 5 \n- BARNN is the first approach that transforms any autoregressive or recurrent model into its Bayesian version with minimal modifications. -> see related work below\n\nWeak claim(promise):\n- autoregressive models are also prone to overfitting to the specific tasks they are trained on, challenging their application outside their training domain. -> Not explicitly demonstrated", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense to me, for detailed elaboration please see below \"Experimental Designs Or Analyses\".", "theoretical_claims": "This paper does not have any theoretical claim. \n\nThe ELBO derivation and the marginalized prior are plausible to me.", "experimental_designs_or_analyses": "- Neural PDE solver benchmark:\n\n    - The use of autoregressive models for solving time-dependent PDEs is well-motivated, as these models naturally capture temporal dependencies, and evaluating uncertainty estimation is important.\n    - The uncertainty calibration metrics (NLL/ECE) are appropriate and commonly used in probabilistic modeling.  \n\n- SMILES molecular generation benchmark:    \n    - Autoregressive models are widely used in molecular generation, making SMILES-based generation a standard benchmark.\nBayesian uncertainty is particularly valuable for molecular design and exploration, where uncertainty-aware sampling improves model reliability.  \n    - The Wasserstein distance, t-SNE visualization (Fig. 10), and property histograms (Fig. 11) faithfully capture different aspects of distribution similarity—global shape (Wasserstein), structural overlap (t-SNE), and individual property distributions (histograms).", "relation_to_broader_scientific_literature": "This work is related to Bayesian deep learning and Bayesian autoregressive models. The proposed Bayesian RNN structure, which uses a multiplicative scaling factor scheme, provides a generalizable framework for transforming autoregressive models and RNNs into their Bayesian counterparts. This builds on prior work in variational Bayesian methods [1] and VAMP priors [2], extending them to time-dependent settings. The method is of pragmatic interest due to its broad applicability in uncertainty-aware sequence modeling.\n\n\n[1] https://proceedings.neurips.cc/paper/2015/hash/bc7316929fe1545bf0b98d114ee3ecb8-Abstract.html  \n[2] https://arxiv.org/abs/1705.07120", "essential_references_not_discussed": "I think literature [1] , which also introduces a variational Bayesian scheme for RNNs, is highly relevant and should be discussed.\n\n[1] https://arxiv.org/pdf/1704.02798", "other_strengths_and_weaknesses": "Strength\n- **A unified framework for Bayesian RNNs and autoregressive models**  \nThe paper provides a complete pipeline, making it potentially applicable to various autoregressive tasks.\n\nWeaknesses\n- **Hard to position in context**\n    - The incorporation of Bayesian methods into RNNs has a long history, with key prior work (see \"Essential References Not Discussed\") that is not acknowledged or discussed.\n    - Combined with the presentation style, this gives the impression of overclaiming novelty, as if applying Bayesian inference to RNNs is itself a novel contribution.\n- **Ablation study of the prior**:  The tVAMP prior is a key contribution, yet the paper does not include an ablation study comparing it to a simpler isotropic Gaussian prior.\n- **Motivation for using RNNs over stronger sequence models**: While this may not be a direct weakness, it is unclear why the paper focuses on Bayesian RNNs when stronger sequence models (e.g., S4) exist.", "other_comments_or_suggestions": "-  Eq.~8: should it be $E_{\\phi}(\\boldsymbol y_{0:t-1}^k)$ (i.e. a $E_{\\phi}(\\cdot)$ is missing)?", "questions_for_authors": "- **Intuition Behind Long-Term Dependency Claims** \n    - BARNN is claimed to help capture long-term dependencies. In the SMILES generation task, do the authors have any intuition on why Bayesian modeling specifically mitigates the ring-closing issue compared to deterministic baselines?\n- **Practical Setting of \n$N$ in Eq.~(8)**: \n    - is the $N$ in Eq.\\~(8) pragmatically set by Batch size of training data or the whole data size?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper provides a framework to turn any autoregressive or recurrent neural network into a Bayesian version by incorporating a novel variational framing.  Specifically the authors define a temporal variational lower bound that they then sample from using a variant of variational dropout.   As the prior, the authors adapt the VAMP prior to a temporal setting (tVAMP), The authors show that their BARNN method outperforms  the non-Bayesian variant across a variety of experiments including PDE solvers, Molecule generation and some nice 1D uncertainty plots.  

The scores were towards accept (4, 3, 3, 2) with an average leaning accept.  The reviewers seemed to find the paper well written and structured.  The method seemed sound, elegant and experiments are convincing.  In terms of weaknesses, the reviewers suggested that the method may not scale to very-large autoregressive models, e.g. billions of parameter models and one reviewer found that the method seemed like a somewhat straightforward application of variational dropout to temporal models.  The lowest review scored seemed to question novelty and seemed to want the authors to compare to some existing literature.  In general, however, the other reviewers found the contribution to be solid and the experiments convincing.  There were desires for some additional ablations, for example to validate the contribution of the tVAMP prior.

Overall this seems like a strong paper that addresses a significant problem of interest to the community (improving uncertainty of autoregressive models).  The method is quite general, not overly difficult, and empirically convincing.

Note there are some typos in the paper, e.g. multiple times "factories" was written instead of factorize.
Decision: Accept (poster)