Probabilistic ML - 24 - Attention

[Music] All right, it's quarter past, so let's get started. Even though it's a bit up sheet symfony, I guess by now uh the group gets smaller and smaller every time uh I start a lecture. Welcome to lecture number 24, which is the last content lecture of the term. This is the feedback from last time. I well I mean there's a big red ball over here but I also kind of succeeded in pushing you towards this corner so um maybe for lecture number 23 that's okay. Here's the um detailed feedback for that lecture mostly I so which was by the way about variational inference mean field inference. I uh want to focus in particular on this question here which um so let me read it out. Basically my problem is this person writes with a lot of what we do in class is that many modeling assumptions seem motivated purely by computational tractability and beyond that they often feel arbitrary like exponential family likelihoods conjugate prior gausian approximations. Why is there a reason to believe that these assumptions hold? um they don't always feel well justified in terms of why they work and why the commitive error from them wouldn't hurt model performance when you show code at the end of the lecture it feels a bit like magic and it's hard to tell when such m methods will actually work in practice so I think first of all this is a really good comment I think when you reach that point that's an intermediate step on the way to accepting our human limitations but maybe you haven't quite made the connection yet when you think about this because what I'm what I was hoping to try to convey and maybe haven't succeeded in doing is that all scientific inference is of that form. Why do we think Brownian motion is a is a gausian distribution or a gausian process? Because it's computationally convenient, not because the particles and gases actually behave like brownie in motion. Why do we think that our world is a uklidian vector space? Because it's a convenient approximation. Why do we have I don't know why do we use partial differential equations like the NIO Stokes equations to describe the weather? Because it's a convenient computational approximation. Every single form of quantitative science is an approximation and they are always made for computational reasons. Historically, the laws of nature that you've learned in high school were designed to work on a piece of paper because the people who came up with them didn't have a computer. Now, we have slightly more complicated tools, maybe not just slightly because we have very powerful computers available, but the challenge is still the very same. So, every single model we build is wrong. And yet they still are really helpful at least some of them. Now the one sentence that I was quite surprised by is this last one. It's hard to tell when such methods will actually work in practice. So whoever wrote this I think that like I think there's one of two things might be the background for this comment. either you think computer science and machine learning doesn't work in practice. Okay, I think that was maybe a valid concern 10 years ago. I think it's difficult to still hold that opinion today with all the tools that are now flooding the market. I think it's increasingly unlikely that we'll find in hindsight that machine learning doesn't work. Or the other thing that might might be happening here is that this person thinks or hasn't realized that the tools that I've been presenting are what actually drives machine learning. Maybe you think this probabilistic machine learning class is somehow disconnected and everything about machine learning is deep learning and it's only true deep learning if it's from the deep learning regions in the south of well then the center of California probably right um but that's not true right so the models that we've been discussing for the entire term are the foundation of the algorithm s that we now call AI. But the stuff that I've showed you admittedly was maybe a bit more simplistic than the tools that you now have available. So, okay. Right. So, I didn't tell you how to build GPT. I didn't tell you how to build a new weather forecasting model on your laptop. And the reason for that is that I can't do that in 90 minutes. And you can't do it on your laptop. you need massive amounts of compute available and in our field I think that's currently sort of perceived as some kind of problem that industrial advancements industrial tools are increasingly further and further removed from the kind of tools that we can build in a teaching lecture but I think that's actually just a sign of a field maturing actually maybe it was much more problematic that 15 years ago when people gave lectures like this the tools that they presented on their laptop were pretty much the tools that people used for machine learning out in the wild. The fact that we now have companies that have thousands of employees and make a huge amount of money building actual products is just a good sign because the field has matured. But that also means that I can't teach you anymore how to build clothes right here in a in a lecture hall in 90 minutes. But that's just completely normal for a field that has industrial applications. So, so to make that this this sort of uh analog metaphor again, imagine this or maybe not even a met a metaphor, just a comparison. Imagine that this were a lecture on mechanical engineering and we were talking about the electrical motor and how it works. Then I might have a lecture assistant here, a technician who brings in some like you know like in the physics lectures up on the mortgage brings in some uh sort of contraption that's supposed to show how an electrical motor works and it's going to be some sort of wild thing made from wood and metal with a like some some coil that is hand wound and a very simple rotor and starter and then when we connect it to some like through some big blue and red wires to some uh um power source force it starts sort of cranky rotating and then we would say yeah this is how an electrical motor works. No, no one would would think oh this is what I'm building what I'm going to put into my electrical car right because you realize that the companies that make actual electrical motors have a huge amount of intrinsic knowledge that is specialized just for this application and that makes them competitive. So just in the same way these AI companies are now increasingly having internal like you know secret knowledge that allows them to build these products that you can actually use on your laptop and uh typically you don't actually use them on your laptop you just get access to some cloud server right um but it's of course clear that I can't do that can't can't like close the gap to these products anymore in the in the lecture hall so for you maybe that's also an opportunity to reflect on what you're actually trying to learn here what we what we're trying to what I'm trying to get across and what I hope that you're learning here are some principle uh some some principles some some uh uh mathematical concepts some abstractions that you can actually use well to build some small tools for yourself maybe even to start in a new niche that hasn't been cornered yet by a company to build your own specialist knowledge like some of the companies everywhere in the world including here um or that maybe qualify you when you actually start working for a company to, you know, further their technical knowledge. But what I can't offer here is to do 90-minute lectures developing entire products. That's just not going to happen and it's never going to happen again. That might have been true 15 years ago, but we're way past that and that's a good thing. So, with that sermon out of the way, let me get back to the content that I want to do in this last content lecture. I first want to wrap up a little bit this variational inference story and want to use this as a segue again on a conceptual level into some structures that are like I would say still at the quite at the forefront of deep learning attention and maybe how we can approach thinking about it from with some of the lessons learned from probabistic inference over the last decades in the field. But before we get to that, let me first wrap up or make the connection to variational inference. So here's a recap from last week's lecture. Last week Thursday, I spoke about the what used to be thought of as the most powerful, the most elegant form of variational inference, which I'm going to call free form variational inference, which is a word that I've borrowed from Kevin Murphy's book. I'm not sure that it's canon, but I think it describes quite well what um what is meant. So the setting to remind you was we have some generative model. So we've come up with one of these generative models like some exponential family some gausian joint model some mixture model that involves discrete and continuous variables to describe some data set where x are the data. Zed are some variables that we use to describe like that we put into the structure of our model to make it more expressive give it structure that might be useful and theta are some parameters of the model. So zed are variables that we want to have a posterior over. X are data that are given to us and theta are parameters that we're willing to fit but we don't actually want to estimate them like in a probabilistic fashion. We just believe that that would be too much to to ask. So um ideally we'd like to comput a postivity over zed. So that's a distribution for zed given x and theta. But let's assume that we don't quite know how to compute that because it's a bit too challenging. So now what we could do is we'll we'll try and find an approximation for this posterior. Let's call it Q. And we're going to do this by only imposing that we want Q to factoriize in some way across some set of latent variables where in general this doesn't mean that every Z I is going to be a scalar. It could just be a subset of the variable set. So we impose this factorization but other than that we're not going to impose any functional form. We're not going to say I want Q to be a Gausian or a Bernoli or a durishlay or some other exponential family or whatever. Just some probability distribution. Then we did this derivation on Thursday that of to discover that um well actually that was sort of the Tuesday already, right? And we we realized that one way to or to to describe finding an approximation for this posterior is to say that we'd like to minimize the KL divergence between our approximation and this distribution which is equivalent to maximizing the gap between the KL divergence and the evidence which is also a lower bound on the evidence the evidence lower bound the elbow. So we want to maximize this expression which happens to be this uh well difference of two expectations. One is an expectation under the approximation over the log joint and the other one is basically the entropy um so including the minus so we add the entropy of our approximation um which is by the very interesting interesting uh sort of structure as well. maximize the entropy of our Q of zed and its fit to the joint distribution. And now if we impose this factorization, we find that the distribution that maximizes this elbow is up to constants in its logarithm given by the expected value of this log joint under all the other approximations. That's a function of our Z I. Oh sorry is that J. And now we can look at this function in I mean of course here we don't see anything but in a concrete case where this actually is some function that we can write down we can stare at this on a piece of paper and see check if we can find some structure that reminds us of like well that not just reminds us but that happens to be the actual functional form of a log posterior probability distribution up to normalization which we can then fix with this constant. um over a variable zed. And I did an example, an extensive example with a gausian mixture model, actually a basian gausian mixture model that also has uncertainty over the mixture weights and the mixture components parameters, means and coariances for a gausian mixture model. And um it turned out that it was actually possible to just by writing down an equation finding the right form. And then we iterate over all of the subsets zj that are in here. And then um that leads to an algorithm that iteratively updates these approximations and in doing so maximizes the evidence lower bound and finds um minimum k divergence approximation. So I already told you this war story that this used to be how parts of machine learning worked and there were actual products by the way being built this way. I um uh during my PhD had some connection with Microsoft and back then sort of there the tool that they would always showcase at conferences was Xbox Live just this tool that like made that did matchmaking between uh online players something like 200 million users in the world and it was actually built on the basis of an algorithm like this minor variations um one some of my mentors at Microsoft were the people who built these tools so I got to see a little bit of the code I can still not allowed to tell you about it but I can tell you that it was basically something like this right a very carefully designed piece of code where someone some smart people had then sat down and found ways to make all of this scalable to 200 million people. So maybe one easy insight is that you could imagine right this was like a joint distribution over the the skills of 200 million people. If you do this in general, you're going to have some huge distribution that has a lot of degrees of freedom. So they impose some factorization assumptions and said everyone has their own skill and you just get this mean field contribution from all of the other players, everyone you've played with, right? So tools like this actually existed and they involved people standing in front of whiteboards for long hours drawing directed graphical models and then writing down extensive variational bounds trying to maximize them. And then deep learning happened. So people came around with these learned representations. People like Yan Lun founded new conferences because they couldn't get their papers published at New Rifts anymore. And um everyone kind of hopped onto the gradient descent train because it just seemed so easy. So there were these new tools that did automatic differentiation. So you could write down a big architecture that nobody understood anymore. um some feed forward perceptron or a confnet or some other new architectures that came around and somehow it just worked if you sort of like squinted enough and hired someone who knew how to get their stoastic gradient descent set up properly. And so the people who used to come from this field with these very structured approximations well okay I'm going to be a little bit facitious but yeah they basically um sort of got worried that they were losing losing the the plot and started thinking about what like whether we can also inherit some of or like leverage some of these much more user friendly let's say approaches gradient descent in variational approximations and this gave rise is to um something that Kevin Murphy calls fixed form variational inference. Um and it basically works as follows. So instead we it's actually maybe sort of the pedestrian approach to variational inference and I admit that that's a bit opinionated. Um so we're going to say we do the same thing as before but instead of doing this fanciful mean field approximation where we say ah like let's just impose factorization and then try to find the correct form for the best approximation. We actually are using a much more um sort of sledgehammer approach and just say e let's just say that the variational bound has to be a gausian. So our approximation q of zed has to be a gausian or maybe it has to be a banoli distribution if the variables we're interested in are discrete variables that are zero one or uh come from a categorical distribution. So um those are parametric distributions, right? So now we have a probability distribution over zed that has uh parameters. It just has a mean and a coariance in the case of a gausian. In the case of a categorical distribution, it just has these k for k categories variables um that are probability distributions or numbers between zero and one that sum to one. Um and now let's see if that makes our life a lot easier. And it does. Yeah. So now um the our very our elbow or our evidence lower bound still in principle has this form. But if we now know that um Q of set is some parametric distribution then one thing we can do is we can we can um well maybe if we're lucky we can just appro we can just actually compute this expectation in closed form but often you can't even with gausian approximations. So what you then might do is already like a next step of simplification you just replace this analytic integral so that this expectation of course is hiding an integral right over the distribution Q with some Monte Carlo approximation where Monte Carlo here doesn't mean mark of J Monte Carlo it's just direct samples from Q because since we're choosing Q to be this um um simple parametric distribution We can uh draw samples directly from these distributions, right? It's easy to draw samples from a gausian, easy to draw samples from a discrete distribution. Basically trivial. So there is no fancy mark of Monte Carlo going on. We just directly sample and we can directly compute such an expectation. In fact, in such settings, people often decided to just draw a single sample. So capital S was just one or is just one. So you're just evaluating this at a particular random point. Um and now this expression is an explicit function and you can implement in Jax and you can compute a gradient and then you can just follow the gradient actually the gradient of what? So there's two different things going on. Remember that we have this distribution here um that we're ob distribution here we go which involves latent variable Z and parameters theta. We want to maybe optimize both right we want to optimize a distribution over Z for parameters theta. So the derivative with respect to theta is now easy because it um uh might directly just um um you know it's just this is just a piece of code that you implement. It has parameters theta you compute gradients with automatic differentiation. You just follow the gradient. Um a minor complication is about the parameters for this distribution set. So these for example mean and coariances or maybe these pies. um because they directly affect our samples, right? So it seems like if we optimize the mean and the variance of a gausian, we would have to redraw the samples in every optimization step and then we'd get a very stochastic object that is difficult to optimize. So people came up with a cool trick called the reparameterization trick um which is initially motivated by uh the the the gausian where the structure is particularly obvious. So you remember that um you can write and we did this in the Monte Carlo lecture. You can write a sample from a gausian as a standard gausian sample. So a sample just a second with mean zero variance one and then multiply with the square root of the coariance and add the mean that gives you a sample from a gausian. So now what you can do if you want to optimize from mu and sigma is you draw the epsilon once and you keep it constant across the optimization. it never changes again. That gives you your Monte Carlo object that you can compute the sample this Monte Carlo estimate from. And now we just optimize mu and sigma. This makes the problem in a sense non-stoastic. Even though we initialized it with some random numbers epsilon, we keep them fixed afterwards. So nothing ever changes stoastically ever again and we just have a continuous optimization problem for mu and sigma. That was the question. ization that's used in autoenccoders. >> Yes. So the question was is this the same trick that is used in a variational autoenccoders? It is. So variational autoenccoders are just another way of parameterizing a variational bound for a probability distribution in a using some nonlinear transformation where the distribution Q here is effectively something that comes from a deep neural network. Um uh and you sort of reparameterize these in multiple ways. Yeah. So this is the trick for gausians. Actually there's a like and then people came up with reparameterization tricks for other distributions as well. Maybe the most um surprising one is the one for categorical distributions um which seem like the simplest possible thing right the categorical distribution is such a simple object but it's fundamentally something very discreet right so if you draw from a categorical distribution you get iid draws from the categories. So the samples kind of jump around in an unstructured fashion between the categories. So some people came up with a really cool trick. Here is this paper. Young Edal, one of the co-authors was from Tubingan um in 2017, right? So we're reasonably close in a way less than 10 years ago. um which is that it turns out that if you draw the um IID random numbers from a Goombal distribution um the Goombble distribution is another exponential family that was originally invented for the modeling of extreme values maxima then um you can you can reparameterize these gumbal iid draws in this way. So pi are the parameters of our categorical distribution. We take the logarithm of those. We add those gumball random variables. We divide by something like a temperature a parameter theta towel sorry. Then take the softmax of the whole thing. Then for small values of toao those are actual iid draws from the categorical distribution with parameter pi. So if you take to zero this becomes a non-ifferiable distribution. But if you take tor to something very small like 10us 3 then this thing is obviously differentiable with respect to pi you can keep the gloomal random variables constant and you get a categorical distribution that can be um where the random numbers can be optimized with respect to their parameters and that's actually how such categorical random numbers are now implemented. If you read the dock string for the categorical random numbers in Jacks, you'll find that you use the Goomba trick to um make differentiable uh elbows. So we can use such reparameterization tricks and these sort of um sampling tricks to do gradient descent under variational bounds. And um so for fun I implemented this um on Friday briefly. So I simplified my life a little bit. I didn't do the full Beijian thing. I just kept uh like basically did like a re-implementation of our gausian mixture model problem with um just means and coariances as parameters rather than um as parameters of a variational distribution. and uh and a a a mixture distribution pi with probabilities for the mixture components pi um where for which I even put in a prior a dishly prior because that allows us to model sparity a little bit and um so I did this I'm not going to go through the code I mean the basic idea is that I've implemented some elbow that takes in the parameters of a model in this case they are um the uh the the log sorry the logit probabilities of the cluster members membership assignments are uh which I call row these are the logits the gausian parameters mean and coariance for the coariances I I store actually the square root of the covariance because it has to be a parameterization of the positive definite cone of of uh coariance matrices right I can't take a non-definite coariance matrix um the mean and uh pi where is pi here for which I take the logits again and implement the entire elbow explicitly in this form and that gives us something we can directly optimize so the thing returns an elbow of which I can compute a gradient um actually I compute the gradient of the negative elbow because then I can feed it into an optimizer that takes a descent step. Um, so atom for example and then I use a standard deep learning library like optax to just do gradient descent and the output looks like this. So what you see here is a gausian mixture model with three components being learned on this simple toy data set that we've looked at a few times already. This sort of old faithful geyser and I've decided to use three because two is a bit boring. What you can see is also the mixture components that are being learned. You can imagine which one's which. This one is this one. This one is this one. And this tiny one is this one. Uh by fixing the parameters of the prior, I can also make this small one go away. And then it learns a representation that sort of gets realigned a little bit. The interesting thing though is that training this model actually takes a non-trivial amount of work. So it took me about an hour to get this to train. I had to switch from SGD to ADOM because SGD got stuck all the time. Overall the training process still takes quite some time. So this red curve here that you can see which is the loss over time takes 15,000 steps of atom which on this machine with jitted code takes like 20 seconds or so. Okay, it's still fast enough. I mean, it's also a tiny little data set, but the EM loop takes about I don't know like 0.5 seconds or so. It just goes and it's there basically. Um, so this I think the main reason I'm showing you this is that I think there's a very high level cautionary tale in there that implementing this model is arguably much easier. You don't have to think much about the structure. I'm just writing down that I want the model to be a sum of gausian a weighted sum of gausian. Okay, let's do gradient descent on that. So the first thing that happens when I do that is I realize that I have to be careful about the parameterization of the model. Oh, I need like a square root of the coariance of a gaus and that was actually not straightforward to get to work. Also, you need to get like jacks to do this right for you. There's a lot of vmaps going on here. That's very annoying. Um, and then once you've implemented it, you still need to run an optimizer that's actually not particularly efficient. So it really takes a lot of time to converge, even though 20 seconds isn't a lot, but it's much much slower than using a structured algorithm like here. So what we're doing is we're basically we're paying a price in compute for our lack of understanding of the problem. And there are these structured algorithms that are very like specific for particular types of problems which run much much much faster like EM or its variational generalizations to Beijian mixture models. But getting those right requires tedious derivations literally on a piece of paper that to this day we cannot really automate properly at least not in a general form. There are some tools that allow some optimiz some optimization but they're quite limited. So maybe that's where we were at. I don't know 2011 12 or so 13 I don't know. So for example I highlighted some paper here. Um there's a paper by Tamara Broadick and several other people. she was um quite influential at that time in this these developments of um uh these uh streaming variational approximations. She had several papers on um uh varants of variational inference that work on streaming data that came under the humorous names like MAD base and SAT base. Actually they were called SDA base but the joke was obvious. Um and um so around that time maybe 2013 145 yeah people really started building these more unstructured gradient descent trained models also Maxwelling with his variational autoenccoders became like quite successful around this time and in many ways these tools provided functionality that you can't quite get from completely structured variational bounds also because they were able to learn representations but in the process. Maybe in hindsight I would say we lost a little bit of the structure that can actually lead to drastic computational gains. And from my perspective, and I admit that this is a bit like squinting at the field with and with a lot of simplifications, we're now continuously rediscovering some of the structure that used to make these models very like efficient in convergence in new forms that again let us regain some computational efficiency. And the example, the story that I want to tell for this today since it's the last content lecture of term is the one of attention. So to get to that point, I have to take a few detours for a few minutes and to get you into the right mindset. So we're going to start from this picture. We've just learned a gausian mixture model with gradient descent. And now we're going to make a few changes. So the first thing we're going to do is to realize that mixture models aren't just for gausians anymore. We could use other forms of distributions as the base measure for a mixture model. So this is what I mean by that. Um a mixture model the name for a mixture model for me is a model of this form. So we we are observing data X which is supposed to be generated from a generative model with some latent variable set which has the property that so first of all of course we assume that every datim comes IID from this model that's our standard assumption for pretty much like 95% of machine learning but then we additionally assume that every single DATM comes with probability pi K from a component number K where each component is identified by a distribution P let's call it PK or P of theta K for the DATM XN. So to make a DATM number n we first draw a categorical variable with probability distribution pi to decide which component it comes from and then we draw from the component. So, so far our component distribution has been a gausian. These component distributions are also called the base measure of the mixture model. Gausians of course are good models for real valued observations. So if the xns are real numbers maybe gausian are a good base measure. But um if the XN are binary random variables, for example, if you model language where you have the identity of the token, then you might have might want some other mixture uh uh sorry, one other base measure of your mixture model. And maybe a natural one is a Bernoli distribution. So if these are um of binary form or also if they're binary strings if they're identified by binary strings our tokens then we could say let's um to let's say our xn are is a binary random string of length d then we could write a base probability for such an xn so an element of the base measures as a product over the the um entries in our binary string. So the string is of length D times independent probabilities for every single bit in our string where the bits are switched on with probability mu K D in cluster K at position D and switched off with probability 1 minus mu KD. So there's a lot of indexes flowing around here. What this says is that we have we observe x's each x is a string of length d. Every element of the string is a binary object. So it's either zero or one. And our assumption is that every entry in this binary string is drawn independently from the other ones with a binary probability given by mu k d where k identifies the cluster d the position in our string. Okay, good. Some people are noting. Um, of course you could think of more complicated models, but that's going to be the simple one that we use. And we could pick some other base measure. If you don't like this one, you can pick another one and then it's just going to be another um mixture model. In fact, this idea of using mixtures to generalize on a class of distributions is quite common to the point where my colleague Antonio Vgari at who's in Edinburgh is making making jokes about this. You can click on this is actually not an underline, it's a link you can click on it. He takes you to a YouTube video where he introduces some uh interesting mixture structures from a theoretical perspective and makes this joke that he regularly does that he's proposing a law that every time someone comes up with a good probability distribution for something interesting within a few years someone will come up with a mixture model version of it and that's pretty I think pretty true you can find lots of examples of this or he has also chronicled some examples of these so mixture models can take other distrib distributions as B as base measures. Let's see what happens if we replace in our Gaussian mixture model the Gaussian base measure with Beni and then we'll get a mixture of Benoli distributions. I'll do that over two slides. These are the stuff that we had for the EM version of Gausian mixture models. So our assumption was that we have data that comes from a mixture of gausian. So this is the logarithm of the equation that was on the previous slide. We have these Z and K which are the like beautiful elegant extra variables we have introduced that indicate which data belongs to which cluster. They they select an entry from this sum and they say entry number data point number xn has been drawn by first deciding to use cluster k by drawing from the probability distribution parameterized by pi k and then drawing from a gausian that has mean mu k and variance sigma k. To do em we then do two things. We first compute the in the E step the expected value of the latent variables. When we do that, we set an approximate distribution over the latent variables to the true posterior. Um we found that we can do this for these discrete variables by applying base theorem which involves the prior probability for the zetss and then the likelihood. The likelihood involves these gausians. That's why it's orange. And then we do base theorem. And we sort of realize, I've pointed this out a few times already, that using base theorem for categorical variables is like a softmax. Base theorem is a softmax. It's a softmax on the logarithm of these expressions. So the log of pi, so this stuff up here, right? Log of pi k plus log of gausian. Those are the logits that go into the softmax. Here it is. Those are called the responsibilities. That's the E step of our Gaussian mixture model EM algorithm. And then um we in the M step maximize this um evidence lower bound by keeping this postivia over zed fixed for a moment and maximizing with respect to the parameters. The parameters here are pi and mu and sigma. Um and um so now what happens if we replace this with the uh with with a with a banoli mixture base measure in our mixture model. So I'm just going to go one slide forward and I'm replacing everywhere there was something orange gausian with a banoli nothing else. So before there were gausian now there are bundle distributions everything else stays exactly the same. So now we have in our mixture components this product over individual pixels or entries and a binary string whatever you want to call these discrete objects that make up a string um with their probabilities to be switched on which I call mu kd by the way mu sounds like mean that's because it is a mean right and a discrete distribution that's the mean activation rate um uh times the probability to be switched off. We still do an EM step where in E we um compute a soft max over lock probabilities. The the lock probabilities now look like this. Um the pi is still in there. There's nothing has changed here. And then we maximize the evidence lower bound which now contains the expected value under our posterior of this expression. And there used to be a loion here. Now there is this log benoli. Okay. And for a loop I can write out what it lobbity is because it's you know it's just a logarithm of this thing. So it's easy to write down. So what would happen if we now maximize this elbow with respect to uh the these mean probabilities discrete probability switched on with respect to mu. Well I take this elbow and I compute the derivative with respect to mu which is what I also did for the gausian case. And you can do it some nasty equations with sums that you have to rearrange. Basically what you do here is what you write this down. You take the derivative. That's easy, right? The derivative of the log is one over huh the thing for one minus the log is minus one over the thing. Uh now you get some equations that like some of them involve mu, some of them involve x. Annoyingly the mu is in the in the numerator h sorry the denominator. So you have to expand these fractions so that you get some terms with mus that you can rearrange. Make sure that you notice that some terms don't contain a mu and some do. You rearrange, you get this expression for mu. So that also kind of makes sense. What is our mean probability for entry number D in our binary string to be switched on? Well, it's the weighted sum of the responsibilities for the clusters um among the clusters for each datim over the number of times the empirical frequency with which entry number D in string number n was switched on divided by the total number. So it's really just an empirical frequency basically. And then you can implement this. Um here is my uh approximations. You by the way you can also estimate pi of course but it's the same thing as before. So there's no change to the gausian mixture model. It's the same line. And here is my uh simple while loop em algorithm for fitting the comp the parameters mu of a benoli mixture model. The pies which are the cluster probabilities and the responsibilities are in k. And again it's a while loop with a while equal to true expression. And it's still fine because we know that every step of this algorithm increases the evidence lower bound. So it has to converge um and also makes the in the E step it makes the bound tight in the M step it maximizes the bound so it has to converge to some um uh fixed point and then it'll break and now what I've done as a final thing before the break is I've applied this mixture model to mnest because again it's just a 90-minute lecture so what kind of data set am I going to use I'm going to use mnest not because mnest is hard it's easy actually running this algorithm on emnest takes 10 seconds on my machine and now we have a mixture model for images that contain digits. So maybe to be clear in the final few moments before we go into the break what I've done is I've taken the emnest data set I've binarized those uh data points right so emnest has grayscale images to make my life easy I've just binarized them I just set a threshold at I don't know 0.5 or something set everything above to one everything below to zero and then I treated these pixels or I reshaped those images into strings of length what 768 or whatever however long mnest is and then ran this algorithm from the previous slide on this. And of course, I decided not to use 10 clusters because that would be too easy. But instead, I used 14 clusters. Just told the algorithm, find 14 clusters. And here is the clusters that you see. So as you can tell well I mean maybe this algorithm has found some structure in this data set that is maybe reminiscent of what's going on here right there is it looks like digits and there are some clusters that are maybe to our eyes not particularly good representations of of digits like for example this cluster number three here or Um maybe also this cluster number nine which isn't quite clear whether it's an eight or something else actually. Maybe it's a five and an eight mixed together but it's good enough for what we're going to do next. Yes. >> Ah so I randomly initialized in this code. I should probably have uploaded this code to Jupyter Hub. Um I I randomly initialized the means the mu k to um actually not just one particular datim but like three randomly chosen images per cluster mean so that initially they look a little bit like the sum of three three digits put together and then I ran this piece of code. Um and the pi case I randomly initialized I think or maybe I even set them to all equal to one over k. It doesn't really matter much. Um so you can clearly that's enough right? You don't need to initialize the r and k because they are directly updated. So we just need pi and mu and then I ran this until it converged and this thing takes on this data set I think it does something like 70 steps or so. So it's a few iterations but it's not like 15,000 gradient descent steps. Okay. So this was a preparation. I mean it's also a bit of a side note. You don't have to use gausian if you do a mixture model. You can use bendis. But now we have a data set that we can play with that I'm going to use to motivate attention after the break. So if there are no questions, that doesn't seem to be the case, then I'll see you in five minutes. So now I should say um like quick disclaimer what I'm going to do is for the last 45 content minutes of this of this entire course I um will do an attempt and it's really just an attempt to talk about well transformers and aspects of transformers. I know that I mean most of you will have taken some introductory machine learning class either in tubing or somewhere else. You've probably heard some introduction to what transformers are. Maybe you've heard them about them in the the deep learning class. Everyone feels like they have to talk about transformers because they drive you know LLMs. And um this is always based on like whenever someone tries to put together a a lecture on transformers, they end up looking at this paper which is maybe one of the most celebrated papers of machine learning in the last decade. Um and that is a really wellchosen title as well. um built by these um now famous eight people at uh at well back then Google and then they went somewhere else partly because Google didn't quite or couldn't quite keep them. Um and I should say that everyone in my impression everyone who's trying to teach about transformers is really struggling to do it and they always end up putting this picture on their slide as I have now done too. This picture comes actually from a blog post which is linked here as well. But what it is is just three figures from this paper put together. So here is the first one. It's the big one. That's the transformer architecture. Then there's some zoomins that involve dot product attention on multi head attention. So clearly attention is important as the title of the paper suggests. Yeah. And you need it to build this thing apparently. So when this paper came out and people realized that it was a powerful idea that you can build LLMs with it that can't work, right? And we now all have these co-pilots. Everyone stared at this image. What's going on here? And I'm guessing that's probably what you do as well. When this came up, I remember that there were like journal clubs here in Tubing as well. Everyone was staring at this image. I was like, what? And then they tried to pick it apart. So that's what we're trying to do as well. And I I really I mean I seriously think to this day the community has not really found a clear canonical way to explain what exactly is going on in these models. And I think the reason for this is that there are so many things happening at once and as you can see in this picture. So what I'm going to do is I'm going to zoom in twice as the Lilian Fang in this blog post has also doneoop into the core apparently structure of um transformers which is attention the attention block and then I'll try and get my own intuition for what's actually happening in these blocks. I'm not claiming that that's the uniquely correct or only intuition but I mean you've taken deep learning classes in ML. You've heard about transformers elsewhere already. So you already had some story. This is going to be my probabilistic machine learning approach to them. Transformers then mix these ideas from attention in multiple ways. Um they also use structures that are called self attention, multi head attention. And then there is the stuff going on down here which is um the embedding of the tokens from uh um text which I'm basically not going to be talking about. Um ah something interesting just happened here on my machine that you couldn't see. Good. Um um I'm going to make a note to that at the end very briefly, but for us it's not going to be particularly important. I also think that it's not the interesting bit. It's just what you need to do to get language correctly represented. But this is not a language modeling class, so whatever. Um and I should say I think I'm not alone with struggling to get this right. If you read textbooks on transformers and uh attention, you can tell that everyone was struggling to figure out what to talk about. For example, Kevin Murphy, who is maybe written the most up-to-date te textbook on machine learning at the moment. So like all credit to him in his book from 2024, if you look at it, it's freely available online. There is a whole chapter on attention. And I would say I can tell him flip-flopping about trying to figure out what he's talking about. So in on page 518 he says oh let's think of a deep learning architecture where phi is some nonlinearity. So like you have like the building block of a deep architecture is one where you have an input to the layer x and then you multiply by some weights. Just imagine that if you had weights that depend on the input that's called attention. And he writes that for five lines and it's like what is a very generic structure and it doesn't actually look like what you see. And then he goes, "Actually, maybe no. Maybe maybe attention is something like that people describe as a soft dictionary lookup." Maybe you've heard about that term before. Has someone Yes, some people are nodding. Soft dictionary lookup. Uhhuh. So, attention is a is a function that looks a bit like this. Um, which we'll come back to where you think of a database, right? Databases have have uh are are like pieces of code that you can give a query to. I say they say I have something. Can you tell me more about this? That's my query. And then what the what the dictionary does is it goes through the list of its keys. You can think of a Python dictionary as well and it finds the key that matches your query and then returns the entry that is stored under this key. That entry is called the value, right? So attention is somehow a soft version of that where we try to match the query to all the possible keys. They match in some non-binary way. Then we do a soft max over those queries and get give back a weighted sum of the values. Uhhuh. Interesting. So we're going to use that structure. And then some people have decided that there's a basic form, a standard form of attention that goes beyond this. So this a here is just some abstract function called the attention score. But there is a special version of that that's sometimes described as the standard way of doing attention which you could also find on Wikipedia. for example, by the way, the Wikipedia article on attention again is people struggling to describe something they haven't fully understood. Um, so which is called scaled.tproduct attention or QKV attention for query key value. Um, which has this form. So here the attention score happens to be an inner product. So we get in some query that somehow gets maybe linearly translated into some capital Q and then we have a key that also gets somehow translated into a capital K. We take an inner product between these two matrices and we somehow scale by the square root of the dimensionality of these matrices. Take the softmax multiply by some representations of the value and now everything here is dotproducts. So these are matrices. This is a matrix that can multiply with this and now we get multi outputs here. The first thing to realize here is that this square root over DK is not really important. It's just something you do to get gradient descent to work well under the assumption that the Q's and the Ks are IID random variables with mean zero and standard deviation one. The dividing by this makes the variance unit. So it's somehow invariant to the dimensionality of the objects. and then gradient descent is going to work. Well, let's forget about that for a second. Interest the interesting bit is this inner product here and a soft max and the V. Let's see if we can approach from a very different direction and end up with a structure that looks like this. And then we'll see if we can zoom out again towards the transformer architecture. And the hint is in this form here. You can see that there is a sum and a soft max and then values. Can someone guess already what the map is going to be to the stuff we did before the break? Yes. >> Yes. So in the mixture models we also had a soft max. These were the responsibilities over something and then there were base measures that the softmax got multiplied with and there was a sum in front and that gave us predictions. So let's see if we can do that. So I'm going to build a setup. My task is not going to be a large language model. My task is a predictive problem but it's a little bit like a large language model. Right? So remember what a large language model does is it takes a token uh sorry strings of tokens a sequence of tokens. It's a sequence model right? You give it sentences that you've typed. It looks at your basically your most recent tokens that you've produced with some trickery going on positional encoding and so on. And then it predicts the next tokens. It tells you what comes next in your text pretty much right. That's the idea. So I've made my own little uh token prediction problem which is much much more simple. I've taken emnest and then I've paired pairs of digits. So I've taken always a uniform random sample without replacement from emnest for a um even number. So 0 2 4 6 or 8. And then put right next to it an odd number 1 3 5 7 9. Okay. they because I've drawn them from emnest they come from different writers they have different handwriting and so on they don't match to each other at all but it's always the case that on the left there is an even number and on the right there is its successor next number so next to a four is a five next to an eight is a nine next to a two is a three okay good so now the task is going to be I give you the context the left side of the image your task is to predict what comes next so you have to predict the images the pixels in this image. So let's see if we can do this. Right? So the this is basically a regression problem or a yeah a structural prediction problem. I give you this. This is X. You tell me Y which is the blue thing on the right. So the first thing I did to see how hard this is is I tried to do this with a feed forward perceptron in an encoder decoder style architecture. So I made a little uh piece of code that I forgot to upload which is just a simple feed forward perceptron one. Here it is. It's basically just this function. I mean feed forward person is a big word. It's just this function um which is a two-layer model. So it takes in the input X. X is an image of size 28x 28. We are reshaped into a 768 pixels long vector. It's binary. And then I do a linear transformation or actually an fine transformation. So I multiply with a weight matrix. I add a bias. Then I take a tanh nonlinearity. Tanh because I wanted to make this model as similar as possible to an attention mechanism. A tan h is a sigmoid, right? Softmax is also a bit like a sigmoid. So I use tanh. Um then I multiply again by a linear weight, add a bias, and then take the logistic because I'm predicting binary random variables. So the logistic is the right loss function for the output. Um and then I did so this is what this architecture looks like right we take X binary thing map through the linear transformation add a bias we get a hidden representation we map out again into something that's shaped like an image and predict the out image and now I do gradient descent on this loss um I did this up here you see loss goes down right make a prediction the prediction looks like this not good which is good for me because it like makes my story work. This is not a good prediction. Why is that? Well, it's because these weights here have to learn a very complicated map from 768 pixels to 768 pixels. And they don't know anything about the structure of this data set. They don't know that there are 10 digits in there or that they somehow have something to do with how the pixels activate in some shape. they just have to predict what the distribution of the pixels is and they just are not particularly good at doing so. So the prediction is pretty bad. Um and that's a good base case for us. So this is a bit like you know doing language modeling if you don't have any structure no attention whatsoever. Um and um notice that it's maybe useful to keep in mind I cunningly chose this model class to be to contain an interesting number of parameters. So I'm mapping from the 10 from the 784 sorry 84 it is of course pixels to 10 and then from 10 out to 784 again. So the total number of parameters is um both of these weight matrices are of size 7,840. So 10 * 784 plus the biases. In one case there is a 10 entries long bias and in one case it's 784 long. So the total number of parameters is 16,474. Now I'm going to build a different model. I'm I said before the break gradient descent on unstructured models is not particularly efficient. Sometimes it's much better to do EM. Why don't we train a mixture model with EM? That's two things, right? Using a mixture model is a powerful formalism to describe structure in a probability distribution. It's one of these models that we just impose, right? It's like just like your question on the feedback. It's just this thing that I come up with some latent variables that I think describe what's going on in terms of some very simple distributions benoli binary random variables. So what I'm going to say is that my model actually is a mixture model of beni but because I have this predictive structure my benoli mixture model actually has two parts. So I'm going to say the data we see actually let me go back to the original data set. This data I say comes from a mixture model of Bernoli where each Bernoli distribution is a joint distrib as well it's a product over independent probabilities for all the pixels in such an image. So that's two * 784 pixels for the left side and the right side, the orange side and the blue side. Yes. Good. Some people are nodding. So here is my mixture model. I'm saying the probability for the pixels of these images, the orange and the blue ones are a sum over weights in the mixture. I'm going to say there are 10 mixture components. So notice that there are actually only five clusters, not 10. But I'd give the model a little bit more freedom. Five because right there are five even numbers. um times a probability to be in that cluster which is initially set to the unit uh like one over K one/10th probability and then the mixture base measure which consists of a big product from for D from 1 to 2 thou 2 * 784. So these are the individual pixels in the image which I immediately block into two structures an orange and a blue structure where each pixel has its own probability to be switched on or switched off. Those are the parameters of my model. So there are 2 * 784 parameters in this model plus the 10 parameters for this pi k vector. So that's less parameters than my my two layer perceptron that I've trained before. And um now I've trained this model with EM using the exactly the algorithm that I showed you before the break. And this algorithm has learned this mixture model up here. So here are the 10 components of my um mixture model. And you can see that this is I mean it's averaging over all like like elements in in the learned clusters. But maybe you agree with me that these are reasonably interesting clusters, right? So each of these clusters looks like pairs of even numbers and odd numbers. So because this model doesn't know that there are equally sized clusters and five of them, it sort of just by random initialization ended up putting a bit of bit too much focus on this 01 cluster. Here we go. It used three of these instead of two. And in exchange, it only has one of the two clusters, which otherwise you could have two for, but whatever, right? Maybe it's just because the ones have more variability, and so it decided to use three of its clusters for those. Ah, whatever. It's fine. We're going to live with that. So now, let's see if probability theory tells us how to make a prediction for the blue bit given the orange bit. So let's just think about how you would do that. How do you make a prediction with probability theory? Yes, >> base theorem. Thank you very much. So let's write down base theorem. So the probability well how do we make a prediction in a mixture model? Actually we use base theorem or we have to use it twice. We first have to make compute a posterior for the clusters right for the zs the latent variables of which there are 10 for each image. And then once we have a probability distribution over the cluster membership, we can then make a prediction for the right side of the image. So let's first compute the for serial distribution given an image over the cluster membership. But in our setting, we will only be given the left hand side of the image. The right hand side is not visible. So it's actually a latent variable that we don't know yet. So we're going to marginalize over it. That's what probability theory tells us to do, right? We're going to use base theorem and we'll sum out over all the exponentially many like 2 to the 768 86 possible values of the right hand side of the image. So I do that here. So base theorem says the probability of Z equal to K is prior probability for Z equal to K which is just given by pi K in my model times the probability to observe the left hand side of the image if it comes from cluster K times the probability for the right hand side image and now we have to marginalize so there's a big sum over all the possible combinations binary combinations of the right hand side image pixels divide by the evidence. So the evidence is a sum over all possible values of zed and the left hand side and the right hand side image which we don't see. So we have to again marginalize over the right hand side image. And now we notice that the same term shows up in the numerator and the denominator. So these exponentially many terms in this big blue product and some they just vanish. They just cancel out. Nice. So because we don't know them and we don't know them in any of the class clusters, they just gone. So we're left with just the orange bit. Okay. So now we remember, huh? So this is what we're computing here is a categorical distribution under base theorem. Categorical posteriors under base theorem are a softmax. So we can write them as a softmax over well they are actually v softmax. Basically every time a softmax shows up, it's probably someone doing base theorem. So we write base theorem as a softmax over the log probabilities, the logits. So they are the logarithm of the prior probability plus the logarithm of the base measure. Remember the base measure is this product over the individual pixels um with mean um parameters. So we put them in here. So the logarithm of that is going to be a sum over the d pixels. So there are 786 of these pixels over x that's the value of the binary val v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v variable in the pixel number d in the input image times the logarithm of its probability under the mixture model plus 1 - x * log of 1 - mu. Now we see that X in red here shows up in only some of these terms and then there are other terms which don't depend on X at all and we can separate them out because it's a Bernoli distribution. The Bernoli is the exponential family for which the sufficient statistics are the inputs themselves. So just X. So we rearrange the sum into terms that involve X and terms that don't involve X. Those are just some numbers. There are just 10 of them, right? So it's the logarithm of pi k. There are 10 possible values for that. And then it's the sum over all the d's for log of 1 minus mu kd. It's just a bunch of parameters in our model. We sum over them. So now we are left with 10 numbers called bk b for bias. And then there are these d uh entries over which we do a sum. and we do a sum over x pixel number d times the log odds between mu and 1 minus mu. So that's just a new variable that is a deterministic transformation of our model parameters mu. So I'm just going to call them w kd weights kd. It's just a new name for this parameter in my model. And now what you see is that what we're doing here is a softmax over a linear transformation actually an aphine transformation of the input. Huh? So what this thing does is it produces a prediction first over which component we're in. And now what we can do is with this discrete distribution with 10 entries which is like the E step in well it's not actually EM because we're not learning the model we've decided someone has already given us the model right if we have the W's then we have the model already then we can make now after this E step a direct prediction for the right hand side namely we just sum over these probabilities stored in this softmax X times the mean predictions on the right hand side for this model. So that's what I do here. So the prediction for the output the right hand side of the image given the left hand side of the image in orange and our model parameters is uh sum over the posterior distributions for our latent variable set under the model which we just computed. That's the softmax times our predictive distribution for the right hand side. It's just a probability distribution, right? A beni distribution over all the pixels. We can draw from it. We can also just evaluate its mean. So renaming things this is like computing a responsibility for each cluster times multiplying with the probability for the right hand side which we we've just saw on the previous slide that this responsibility is a soft max over a linear transformation of the input with some weights which are given by the mixture um component base measure parameters up to some logarithmic transformation times the right hand side prediction of the image. So what happens here is we take the input which we could call a query. So that's the left hand side image that I put in and then we multiply we do a dotproduct with the weights. Those weights are these log odds log of mu over 1 minus mu. Um and over here I've plotted these weights. So I've re I mean each each uh row k of this matrix w k contains 786 probabilities. Right. So these W's are the logarithm of those probabilities divided by one minus those probabilities which I plot here. And you can see that well they look like basically I mean they're up to some shifts in how I make the plots right it's just an im show. So it automatically scales the numbers basically look like digits like even digits like the left hand side of the digits. So what happens inside of this um dot product here is that we take the input image which is like a reshaped version of a of a digit image of an even number and then we compare it pixel by pixel to those keys these rows of the matrix W which is clearly a form of this soft dictionary lookup right we are comparing what the image looks like compared to all the keys. We overlay them basically m check how well they match. Then we do a soft max over that up to a shift which just takes care that everything is sort of well centered and then we multiply with the predictions which are these mean predictions for every individual case that they map to and these are the values. So here on the left at the bottom I show you the biases as well. They're not particularly interesting. Um, and so I built a little little app which I also have here on Jupiter. Here we go. Um, so what you see here on the left is an input image, a particular one. In this case, it's clearly a two. And then in these small uh boxes you see all the keys which also their biases plotted above. Um, and then so in gray you see the keys and you just about see overlaid the query. So that's the image. Here's this two. Each of those gives us a match. We take a soft max of that. That gives us this distribution here that those are called the attention weights because clearly this is the same algorithmic structure. And here in this case the match is very clear like this model says it has to be this one. And then it picks with probability one oops from the corresponding predictive distributions which is going to be this one. And it looks like a three. So now I can move around. I can try oops something else. Here is a second test image. Now it looks like a zero. We get a map a match um to uh those last two keys. Right? So this query fits particularly well to these last two keys. And um as you can see over here, right? And the prediction looks like a one like a weighted sum of ones. Here is a four which maps clearly to the cluster number zero and two. So this query matches these two keys roughly equally well. We get a five prediction. There's also a case of a bad fit. Um actually let me see. I had an example of one that didn't work particularly well. Uh where is one that's not so great? This one. Yeah. So here's an example of one that is badly matched. This six, I think it's a six. Doesn't look much like a six. It's a bit rotated. So you can see that it overlaps well with the cluster that's for the four just because it looks a little bit like a rotated four as well, annoyingly. And um also to cluster number five, which is the six. And so it sort of makes a weighted prediction of a five and a seven. Okay. So what is my point here? I'm saying one way to approach this idea of attention is that well it sort of also is talked about as a soft dictionary lookup that um associates uh each query that compares every qu one query with a set of keys returns a weighted sum over the match between keys and queries and then weighs the output the values according to these matches. Um this softmax structure is the internal like that's the thing that really drives this fundamentally and you can think of this as emerging probabilistically from a mixture of predictive distributions for the the given the context. We first compute a conditional distribution for clusters and then the clusters provide predictive distributions for what follows next. where the individual um components if you like take responsibility for the input and each of them then having taken like considered their responsibility make a prediction for um what might happen next. So here I'm using the word responsibility which comes from mixture models and DM to apply it to attention which comes from the sort of the the deep learning architecture. And the reason this works well is that the softmax the attent in the attention mechanism or the responsibility in a mixture model allows the model to tune down a prediction if it doesn't fit well to the context. Right? If you if you see an image that doesn't look like what you're good at, you just shut up. That's basically what happening in in these models, right? While in a general feed forward model, everyone has to make a contribution all the time. So there's no way to specialize onto one particular input at least not naturally. And here through the softmax this e emerges naturally which we can interpret as a probabilistic functionality. So now okay that's just attention that's my way of thinking about attention. Now I going to say attention is just mixture modeling where we have predictive a mixture of our predictive distributions. So how do we get from attention to the whole transformer story? So there it gets a little bit more complicated and I think that in that hidden in that is actually a bit of a story as well. So I'm going to show you this and I'm going to try and summarize in a few minutes and see if I can get this across. that like advanced warning what's going to happen now is that I think the connection between what structures happen in or actually get used and transformers and how we can think about them probabilistically gets ever more strained as we try to leverage them. But let's see how far we can get. So the first thing is multi head attention. So in this fancy full cartoon from the attention is all you all you need paper is um this is sort of the next step that says once you know how attention works you can now do multiple attention heads together. So here what they do is they say you you like given a set of value key and query um what we're doing is we're feeding it through multiple linear heads as we now call them in deep learning. then um do attention for each single one of them separately and then concatenate them afterwards waiting them in a linear fashion. So for us the structure I mean this is the structure for multi head attention. You put in one set of key queries and values and then you weigh them all differently to compute different keys different queries different values where these weightings these W's are all learned separately um and then you combine them uh well just all together by putting a linear computing a linear combination of them that can also be learned. So such a structure would emerge in our mixture model analogy if we would consider several possible explanations for what's going on. So we could have like a mixture of mixture if you like a higher order combination of mixture models where um we compute first for every individual mixture model its predictive distribution just like we did on the previous slide. So that's the softmax thing. But now in this softmax there is different uh distributions different base measures for the keys. So different ways of thinking about our data set in different sort of languages according to different sufficient statistics according to different base measures in the model each of which is associated with a different predictive distribution. And that's actually also if you read the textbooks the idea that people now assign to multi head attention they say well it's like you come up with different models. So maybe picking up again on your feedback that I showed on the in the beginning of the lecture. It's like saying I don't like the gausian as a model. I also don't like the bani. Let's find 20 other let's find 20 other different models different base measures that I may be even able to learn um and represent the data in terms of these different mixture models. And then we combine them linearly which amounts to again learning a weighted combination of these different predictions based on how well I trust them basically or how much how often I've seen them work well in the data set. Huh. So the next interesting question is then and this is something that I think is tends to immediately draw attention when people talk about no pun intended about uh transformers and that's this self attention business. So something that like really like people think long and hard about is the fact that in transformers the um keys, queries and values are constructed by feeding actually let me go back to the original image. Where is it? by feeding the inputs into all three parts of the attention head. So this this uh threepronged arrow here is supposed to say that you take the input and you make it a key, a query and a value at the same time by weighing it differently. So this the picture for this is this one. This is from the blog post at Google when this paper came out when we were excited about it. they have this whole long spiel about so basically this is the math right we're saying the key into our attention mechanism is going to be the input x times some matrix that we learned the query is also the input x times some matrix w and the values are also x times some weight matrices so why is this good ah so there is a um an explanation that uh is centers around this picture which is from this Google block which you can like click on up here that says by representing the input or feeding the input into all three layers we can let the model learn contextaware embeddings in this attention way. So here's an example of two sentences. Um the the bottom is the sentence that you're supposed to look at. The animal didn't cross the street because it was too tired and the animal didn't cross the street because it was too wide. Now the trick here of course is this there's this this core reference resolution going on that the um this um um adjective at the end tells us what the it refers to. In one case the it refers to the street which is too wide and in one case it refers to the animal which is too tired. Uh now um in this soft dictionary lookup the word it which is from the input is supposed to attend to once the street and one the animal because it sort of learns that due to the context wide and tired it has to go one way or the other. So where would we get such a structure from? And here admittedly gets a little bit maybe strained. Um I nevertheless want to point out as the last thing I want to show to you what actually happens when we train our Bernoli mixture model um for our mixtures of Bernoli with the EM algorithm. So here I've um shown you again the code for training a mixture model of Benis that um I showed you before. It's the same code. Nothing changed. It's actually the exact same file being loaded here into the latte code. And I want you now to stare at where the X's show up. So there's an X here and an X here. And then there is an X here again when we fit the mean. So these are the two steps of the algorithm. The E step and the MTEP. Remember the EEP is the thing that computes the soft max, the bit that goes into the attention weights and it contains the X's. And then there is the um bit that computes the means. So the means are in the predictive distributions and in the weights of the softmax. So they are the keys, right? So the these are the the weights in the attention uh mechanism. This is the query in this context. These are the weights and they are also the predictive distribution. They're also the values. The left hand side of the image is the is the is the key. The right hand side image is the value. So if I rearrange this. So again we notice that we can write this update for this softmax here as some kind of dotproduct with a weight matrix. Right? There's a weight matrix here with entries w n k which are given by the logarithm of mu over 1 - mu and then a bias. Um and uh this update here for the means clearly is also a multiplication between a weight matrix. Now the weight matrix is the attention heads right and K. And there is a literally an at in here. It's just a matrix uh product. So I can rewrite this update in a notation that explicitly forms matrices called um attention weights M and then well previous slides I called them a weight matrix W but now because we have this attention mechanism metaphor already we can also call it the keys right. Um so now we sort of use the notation from this attention mechanism with a key matrix K and attention weights M M being given by the soft max over um um query time key. Then um we see that in this em loop we actually use the x's both to set the attention the attention weights as a query and also to update the keys because they directly use this map m. So if I wanted to I tried to write this down in the sort of same cartoon style that the attention is all you need paper uses here. So what happens in the EM loop is that we take the input as both the values and the queries where actually we use different parts of the inputs as queries and keys in our example sorry as queries and values in our example the left hand side and the right hand side but it's just a masking so you could call that masked attention if you like um and we put in some keys then we do an attention mechanism right so we multiply the keys and the queries do a softmax over it multiply by the values and then the only thing we do is take a logarithm because that's the notation we use for EM because we use log probability distributions and then we run a loop around it. So now the output of this thing will come in in the in the next iteration and it'll effectively we use effects from the input as well to get new keys. So in EM when we are learning a mixture model in the examples I showed you before there was no learning it was just once you have a mixture model you naturally get an attention mechanism. If we are learning a mixture model, we are doing something that is I would say reminiscent of self attention. Namely, we use a data set X as both the queries and the values and then feed them into the keys repeatedly finding maps between matches between the keys and the queries to make the keys look a lot like the structure in the queries. Okay, so admittedly a bit belabored and it's not the exact same thing as um self attention, but it concludes my story of how we can use probabilistic structure to gain maybe like a like an X-ray view of these elaborate structures that now have emerged in deep learning. And maybe that's a way that lights sort of a dim path forward at the end of this course on probabilistic machine learning on how to continue to think about existing evolving models in a probabilistic fashion. So what I said today and let me like do that as the final sentence was first we started with variational out variational inference. I pointed out that one can be there is a there's a very tempting path forward to by giving up about the structure that gave us algorithms like EM invented a 100 years ago by people like Hugo Dynos to build very efficient computational tools that converge very quickly and use structure in our model to perform well but require tedious handwritten notation and derivations. So a few years ago people started doing this with variational bounds moving to pure gradient descent and in the process we occasionally not always but occasionally lost some important structure that made the algorithms converge fast. There's a similar thing in modeling for sequences in deep learning where like more recently structures like the attention mechanism have helped lead to architectures that can converge much more efficiently to good representations of language and they now power large language models as the maybe currently most prominent use case for AI. We've seen that one can think about the internal structure that according to the original paper is all you need to build a powerful large language model namely attention and motivate it from a probabilistic perspective as the structure that emerges from a mixture structure in a probabilistic generative model namely a soft dictionary lookup according to base theorem which happens to be the softmax and then I try to make generalizations to other model aspects of transformers to multi head attention to self attention. I could also have talked about embeddings and positional encoding but it's not so interesting. Um and the question now is where the connection breaks and where it stays interesting. So I don't think most people think currently about attention from a mixture component perspective. But remember that last Thursday I told you all the cool things you can do with a probabilistic perspective on mixture models. You can for example learn mixture models that learn to drop particular components that learn how many components your model has. We can't do that in transformers at the moment. Maybe there is a way to leverage our probabilistic ideas to improve architectures like transformers so that they for example learn to shrink and grow as necessary. But maybe the connection also breaks. Maybe the last thing I showed you about self attention isn't actually useful and it doesn't lead us anywhere. I don't know at the moment, but I think it's interesting to keep thinking about it. And as we move forward now as a field, we see more and more of these increasingly complicated structures emerge. Transformers, I think, are a great example for this. They're so complicated that when they like they've been around for what five year, no, seven years now and people are still struggling to really get a hold on what's happening there. We just know that they empirically work quite well and they've already been overloaded with all this new notation, new ideas, positional encoding, long uh context windows and so on and so on. And I will I think it's going to be very interesting to see over the next few years whether the this connection to basic ideas for more simple parts of machine learning will get closer again whether we discovered that they were just a detour for something quite basic or whether they are just engineering details that are so tricky that we can't generalize beyond them that's I think a good point to end with the content of this lecture I'll be here on Thursday with a quick summary and also to answer your inevitable questions about the exam. Um, until then, thanks a lot that you were here for the whole term. See you on Thursday.