Lightning Talk: Bayesian Neural Networks With Variational Inference in PyTorch - Lars Heyen

This talk discusses the integration of AI and uncertainty quantification in weather forecasting. The speaker, a post-doctoral researcher at the Weather Institute of Technology, highlights the importance of providing uncertainty measures alongside predictions, such as distinguishing between a sunny forecast and the possibility of rain. He introduces variational inference as a solution, explaining how it replaces traditional neural network weights with distributions, allowing the output to include confidence intervals. The speaker presents Torch Blue, a software library designed for Bayesian learning and uncertainty estimation, demonstrating its ease of use in PyTorch. He emphasizes the benefits of using Torch Blue, including vectorized sampling and compatibility with distributed training, making it an effective tool for high-stakes decision-making beyond weather predictions.

Hi everyone. Thanks for coming. My name is Lars and I'm a post-doctoral researcher at the weather Institute of Technology. And I work mainly on two things. The first is uncertainty quantification. And the second is actually AI in weather forecasting. if you weren't aware that's in the recent years become a very big thing AI in weather forecasting to the point where some of the weather predictions that you

might have on your phones could already be AI generated. And even despite all this technical uh progress that we've made sometimes the predictions are still wrong. The your weather forecast is telling you that it's sunny while you're standing in the rain and well, that's not optimal. So, what maybe might have helped you if uh is if the model had not just told you that the most likely

event is that it's going to be sunny, but maybe it had told you that that was just by a slim margin and there was still like a 49% chance of rain. In that case, you could have prepared and still brought an umbrella. There's just a tiny problem. Um most typical neural networks don't really allow you to do that in detail because typically what you have is uh

you give your network some input. You mix that up with some weights in your preferred form linear layers, transformers, whatever. Uh those weights you learn and what you get out is one prediction, one number or one set of numbers that is just a point prediction. So, there's in that sense not really room for uncertainty, no concept of uncertainty if you have just one And that is where

variational inference comes in And the the key idea of variational inference is quite simple. You replace your weights with distributions, and you learn the mean and standard deviation or some parameterization in general of those distributions. And what that does for you is that now your output is also a distribution, and that of course has then some notion of uncertainty or confidence intervals that you can use to

make a better informed decision. So, uh there there is quite a bit of math behind this. you need to make a number of choices regarding distributions. You at some point need to be aware of the approximations that you're making, and uh in short, this is all very easy to get wrong. And those can be subtle mistakes that then in the end uh hard to trace back. And

uh that is why the team that I'm part of has made those mistakes for you corrected them. So, we built uh software library called Torch Blue. So, Torch for Bayesian learning and uncertainty estimation. And um let's just look at a very simple example. So, this is very much broken down just uh very basic code in Torch, right? You define in some way your model, you define a

loss function, and you put that into a training loop together with your training data or data loader, and and then you just let have fun, and in the end you hopefully have a good model. So, if you now want to use Torch Blue to do this with variational inference, you change or add just four lines of code, and one of those is the import statement, so I

won't really count that. so what we have as features is really just um uh for this purpose, you can just use an auto convert function that essentially does this replacement step of you for you, right? So you had just weights as trainable uh parameters before and that makes it so you now train parameters of a distribution for each weight. Then you need to that's a bit we

have specific select your task type. For the most part, uh what's written here would will just work and you use need to select a different loss function because now you're comparing no longer one number or one set of numbers to your labels, but instead a distribution and that's a bit more complicated, so we've got that handled for you in the back end. Okay, so what does this

give us? Well, I going back to weather stuff, if we now had a temperature forecast for the well, coming days and I guess yesterday as well. previously would we would have just had these temperatures that would have appeared on our apps. Now we also have some confidence um that tell us that um maybe I mean we can trust the forecast for for for Wednesday still, but afterwards

it gets a bit iffy and um it might get a bit colder than the model predicted on average. Uh and the way we get those error bars is essentially by now our weights are We sample from those several times, so we have several realizations of our model. We feed the same input through and that gives us an output ensemble. The problem with that is just it's kind

of expensive. You have to to your um essentially forward step through the model several times. And this is where partially the performance optimizations that we've also mostly hidden from you come in. Um, in this for this example that's uh, our use of essentially vectorized sampling just using uh, PyTorch's vmap functionality to really make use of the hardware that you have, right? Instead of doing this sequentially, it's

very natural to just do this in parallel because the different samples don't really interact during the forward step. And this is just one of uh, the advantages that we have with Torch Blue. Um, of course we have this auto conversion function and that uh, also works with more complex models than just a single linear layer that I showed. But if you prefer to build up your model

step by step, uh, from individual layers, you can also do that in the inter- interface is designed in such a way that it's almost identical to uh, baseline PyTorch even down to the to the layer uh, names. Uh, so if for example you just want to add like two Bayesian layers at the end of your model to get a bit of an uncertainty prediction, it would still

be doable very easily. Uh, it is uh, compatible with Torch's uh, distributed data parallel more or less out of the box. So, if you're using that, that's not a problem. You can still continue to use Um, we have not just the vectorized sampling, we also have some more tricks to to save on compute or make uh, better use of your hardware. And if you want to get

a bit fancy on the math side and uh, have some specific distributions that you would like to put into your model because it just makes sense. For example, you're predicting a temperature and temperatures can't below can't go below absolute zero. So, maybe you want a distribution that's not symmetric, but positive only, right? Stuff like that is also very easy to implement yourself in framework. And of course,

now I've talked about weather a lot, but in general, I think that variation inference is really something you should consider if the decisions you're making with your networks are critical. Say, electrical load forecasting for network management, for forecasting for finance. There's a lot of money hanging on that. So, if you want to make more informed decisions, I recommend that you try out giving your model to give

your model the ability to just say, I don't know. I think you should try to watch Blue for If you have any questions, feel free to ask me. If you have any comments, concerns, yeah, either more or less right now or catch me later during the conference. And with that, I thank you. I have a question. Yes. How do you leverage this in recommendation setting? Good question.

So, um I don't know if what if it was audible in the background. So, the question was how to leverage that in recommendation setting. So, um what what is is a recommendation setting more because I'm I'm not an expert on that. I'm mostly in weather. Is that more would you say that's more similar to a classification problem or like element-wise classification or more similar to a regression

problem? So, classification is um so, essentially assigning probabilities to a finite set of um of outcomes. In that case, um technically, your standard neural networks already can give you somewhat of a probability distribution, um which is like the typical image classification gives you class probabilities. so, um the standard way it's applied in formulas like that is just you use this as a method of generating ensembles and

then you average your class probabilities in some way over these ensembles. Um that's one way to do it uh that I've seen for this type of problem, um personal opinion, I don't like that you're throwing away like the spread of your ensemble uh for that typically, but yeah. But, that is how I would uh make use of it in that context. Thank you for the question. Further

questions? Do you have any >> Have not thought about that before to be honest. But um can we like talk afterwards? Okay, then thank you again and have a nice evening.