Axial Attention & MetNet: A Neural Weather Model for Precipitation Forecasting
deep learning machine learning arxiv google attention mechanism attention transformer rnn recurrent weather long-range layers convolutions cnns rain physics
MetNet is a predictive neural network model for weather prediction. It uses axial attention to capture long-range dependencies. Axial attention decomposes attention layers over images into row-attention and column-attention in order to save memory and computation.

https://ai.googleblog.com/2020/03/a-neural-weather-model-for-eight-hour.html
https://arxiv.org/abs/1912.12180

Abstract:
Weather forecasting is a long standing scientific challenge with direct social and economic impact. The task is suitable for deep neural networks due to vast amounts of continuously collected data and a rich spatial and temporal structure that presents long range dependencies. We introduce MetNet, a neural network that forecasts precipitation up to 8 hours into the future at the high spatial resolution of 1 km2 and at the temporal resolution of 2 minutes with a latency in the order of seconds. MetNet takes as input radar and satellite data and forecast lead time and produces a probabilistic precipitation map. The architecture uses axial self-attention to aggregate the global context from a large input patch corresponding to a million square kilometers. We evaluate the performance of MetNet at various precipitation thresholds and find that MetNet outperforms Numerical Weather Prediction at forecasts of up to 7 to 8 hours on the scale of the continental United States.

Authors: Casper Kaae Sønderby, Lasse Espeholt, Jonathan Heek, Mostafa Dehghani, Avital Oliver,Tim Salimans, Shreya Agrawal, Jason Hickey, Nal Kalchbrenner

Links:
YouTube: https://www.youtube.com/c/yannickilcher
Twitter: https://twitter.com/ykilcher
BitChute: https://www.bitchute.com/channel/yannic-kilcher
Minds: https://www.minds.com/ykilcher
 Hi there. So what you're looking at here is a weather forecast model. Specifically the very top row is a new weather forecast model called NetNet by Google Research. So the goal of weather prediction is pretty simple. You want to know what the weather is going to be in the future. Specifically here you want to know precipitation rates. And so this is a new work that uses neural network instead of physical models in order to predict precipitation. So in the middle here you see this is the ground truth of what really happened at that particular time. You see precipitation rates in red here moving across the country. Now the bottom there is a physical model and as far as I understand it, physical models have been used so far to make weather predictions. Which basically means that you simulate these rain clouds and the movement of them across the country. And you do a physical simulation like a particle simulation type of thing and then that allows you to predict and then you run that maybe multiple times and you get an idea of the kind of distribution you're going to get. Now what MetNet does is it simply uses a neural network to predict the outcome directly. So there's no physical simulation involved. There is just a neural network that takes as input what's the situation now and maybe over a stretch of time. And then you ask it please make a prediction in eight hours or something. And then the MetNet will make that prediction and it will just output it like snap. No physical simulation needed. And you also see here that MetNet outputs things in kind of a cloud way, in a probabilistic way. In one forward pass you don't need to run it multiple times. But we'll get to that. On the bottom here you see the measurement. So the axis is F1. F1 is kind of the overlap of how well you're able to predict the precipitation. And you see here the MetNet is above the HRR baseline for most of this time. Up to 480 minutes into the future. Which is eight hours I believe. All right. So the paper is the following. It's called MetNet, a Neural Weather Model for Precipitation Forecasting. And I'm not going to read all the names here. The main corresponding authors are Caspar K. Sonderby and Nalkalke Brenner. And it's a team of Google research. So specifically they use the input of these two things here. So one is this GOS16, which is what you see here on the left. And the precipitation rates are here depicted on the right. So you want to take these things as input into your model. Now how do you do that? Of course we want to build a neural network. And this is the architecture they come up with. So on the bottom here they feed in the data. And they feed in the data in 15 minute interval from 90 minutes into the pass. So you have to imagine it like this. So there's a timeline. I'm going to use a little bit of a finer thing. So there's a timeline. And you let's say are here. This is now. And then here in the future, this is maybe one hour into the future. This is your target, right? This is you are here and you're looking out. You would like to know what's the precipitation going to be in one hour from now. What Metnet does is it takes an input. And specifically it takes the last 90 minutes before now as an input. And it samples it in frequencies of 15 minute intervals. So each one of these is going to be 15 minutes. And each 15 minutes you get like a snapshot of this entire of the input region. Now the input region, if I can jump back here to the website for a second, they show it what the input region is. The input region, if you want to predict in the middle of this small square, the input region is actually the entire 1024 square kilometers around it. So it's very big input. Though the actual region you consider is the inside 64 square kilometers. But you take in information from the big region. And the main point of the paper, I believe, is how to do that. All right, so each 15 minutes you take in a snapshot. And these are these snapshots here on the bottom. So these are, and you have to imagine in here, every 15 minutes there's a stack of these inputs. So what are these inputs? These inputs are some kind of features that you have. So there is the target time, which in this case would be this one hour here. There is the month, day and hour, which is important for weather prediction, right? So the time of year, time of day and so on. Longitude latitude is probably pretty important. Elevation map is probably pretty important. So these you can see, these are all maps. Now sometimes, and this is how you encode things in these. Since it's a neural network, you know, all of these things must be of the same dimensions here. So if you have 256 dimensions here and probably 256 dimensions here, then all of these things must be of the same dimension. And if you want to give a feature such as the target time, which in this case, let's say it's one hour, you just put here one hour, what's one hour? Let's say 60 minutes. So you just put the number 60 here, 60, 60, 60, 60, 60, 256 times and 256 times 250, 65, sorry, 265 times, no, 56. I'm confusing with German. So this is how you encode features. It's pretty primitive, but it turns out it works the best if you do it this way. All right. So you have these planes and some, as I said, are just features such as the target time, month, day and hour and so on. Elevation, I guess is a map, is like an elevation map of the region you consider. And this corresponds now to this, these 64 kilometers times 64 kilometers here. And that's exactly what these center crops are here. So this center crop thing, that now, this thing here, this plane, sorry, is these 64 by 64 region. That's this plane here. And also the, that's the precipitation and the GOES, that's this thing here. Now we also have these down sampled things, which these are the 1024 kilometers. So this here and this here, these are the 1024 square kilometer patches, but they are down sampled. So everything is down sampled, I guess, to 256 by 256 pixels. So you don't really take into account every nuance of that very big, of that very big input, but you do down sample it. So you kind of get the big picture of the outer frame and in the inner frame, you take it in a much higher resolution in order to get the details. All right. So you stack all of this up into a big tensor and then you feed it into here into a spatial down sampler, which I guess, no, I have read is a, some just a convolutional neural network, right? So this is your typical image processing pipeline. So you do this for each of these stacks, right? And then what you get out of it is a lower size representation right here. So you get these representations and then you let a temporal encoder run over it. What does a temporal encoder do? This in particular is a convolutional LSTM. And if you already know what an LSTM is, a convolutional LSTM is nothing more than an LSTM that has as intermediate layers, convolutional layers. So it's pretty suited to do, for example, videos or any sort of image processing that goes over time like this one. So the temporal encoder simply starts out here with an initial state. My pens are screwing me today. So it starts out here with an initial state and then it simply inputs each of these representations, takes them one by one, runs across time, right? And each time producing a new intermediate representation of the input until it finally reaches this here, final representation. So this thing here is a single final representation of all of this input, right? Of this entire time span of all of these stacks here. Yeah, so you can press this into a single input with first a convolutional network to downsample each time point individually and then with a recurrent neural network, an LSTM, to integrate the information over time. You end up with this single piece here. And then what you do, so you still, here you still retain kind of an image sort of thing. So this representation here, you can see it in the background. Maybe I'll get down my scribbles here. This here is still sort of an image tensor, though I guess it's a hidden representation, so you couldn't really look at it. But it still has dimensions of images. So this here is still, I think, the same or corresponding to these dimensions here. So this still has some spatial information where this might be north-south here in this axis, might be east and west, right? And then these are just the hidden channels, the channels of the hidden representations, right? So what you would like to do now is to basically encode information from the space around you. If you look at, let's look at one of these, one of the big pictures. What you would like to do in weather prediction, let's say you are right here. What's a good example, you are right here, right? Now if you want to know if this particular cloud over here is going to move to your direction, what you want to know is, for example, is there a mountain range here, right? Because then it's more probable that this cloud is going maybe to move up there. You would also want to know how this cloud here moves, right? If this cloud here moves somewhere here around, then it's probably this cloud down here might be pulled with it or something like this. So you're very much, sorry, you're there. You're very much kind of want to look out into each of the directions here and you want to incorporate kind of what's happening across the space. We're already used to kind of convolutional networks being able to do this, but in here the authors use attention to do that. So if you don't know what attention is, my most popular video is about attention and you can do attention for images. So the way that works is that you have a series of images of stacked blocks of a neural network. Let me draw this here. So you have an image here and let's say it has just four pixels, right? So you have the next layer of these four pixels, right? So you have layers of this. So the next layers of the four pixels, they all emit what are called queries and queries are just vectors. So each pixel emits a single vector. Let's say this, that, that, this, right? And each of the lower layers emits what is called a key. This, this, this, this. And now the keys and the queries are routed together based on their inner product. So these two would be routed together. This would probably be routed here. This as well. This would probably routed here. So what in effect each of the pixels of the higher layer can look at specific pixels of the lower layer. Now you can imagine this is exactly what we want here in that if there is a mountain range here and we might be interested in that. So we'd be able from our, from our point here to specifically attend to that location using, using attention, right? So the authors here build basically a stacked model of attention layers. And that's what's happening in the third part here. And this is the attention is in order to incorporate long range dependencies. As I made the example with the mountain range, this might be far away, but it might actually influence your weather very much. So the attention is to incorporate these long range dependencies. But the problem with attention is, is as you saw in the example, each of these pixels can attend to each of the pixels in the lower layer. So what you'd end up with, so each can attend to that. This can attend to each. This can attend to each. You'll see you'll end up with 16 connections. Can't even draw them. So you end up with 16 connections. In general, if you have D here, you will end up with a D squared number of things you need to calculate, right? So if this here, and now of course we have images. So generally we'll think of D by D pixels. Now we have D by D pixels and that thing squared number of things we need to calculate. This quickly gets too much. So in, for example, MNIST, you have 28 by 28 pixel images. This is 780 or 2 or something. I don't quite remember. But you'll have to calculate this squared many connections between things. This is simply impossible pretty quickly, especially if you scale up the images and then have some channels in here as well. So attention for image processing has been a bit lagging compared to natural language processing. In natural language processing, you usually have maybe 500 tokens or something. Images you have much more. So attention is much more expensive. So you can't really do it on current hardware. Now this paper uses something called axial attention. Axial attention is kind of the trick of how to make this tension happen for images. And for that I want to switch over to this paper. It's called Axial Attention in Multidimensional Transformers by some of the same authors. So Jonathan Ho and Nell Coutt Brenner, also of Google Brain and UC Berkeley, proposed this axial transformer. Now they originally proposed axial attention for autoregressive models. If you know transformers, they also started by making autoregressive models, so language modeling and so on. But we can decouple the axial attention from the autoregressivity of these models. So I'm not going to talk about autoregressive models, it's just axial attention. So what is axial attention? It's pretty simple actually. And I want to start by talking about convolutions. So what does a convolution do? Let's just take a one-dimensional image, which is pretty boring, but let's say it has these eight pixels here. So this is an image, it just has one row of eight pixels. What do I do when I run a convolutional filter across that? This is the lower layer, and now this is the next layer that is produced by a convolution. So for each of the pixels in the next layer, what I can do with the convolutional layer, I can look at its neighbors in the lower layer. So these three would be part of that. And then I go on to this, and again I look at its neighbors at these three. I might have done this in a different color. And then I look at this, and it can look at itself and its neighbors. So a convolution is pretty smart. And then of course in the next layer, that repeats. Now if you think, what's the difference between doing this and a fully connected layer? So if I have a fully connected layer, a classic neural network, a fully connected layer, then this pixel here would incorporate information from all of the pixels here. And this pixel here would incorporate information from all the pixels. Now why might this be better? Because the information that I want here for this pixel might depend on this pixel over here. So I might benefit from a connection over there, or it might benefit from this pixel here, which it can't reach. And with a convolutional network, I can't do that. Why are then convolutional networks preferable? Because the convolutional network can do the same thing across multiple layers. So let's assume again that this pixel here needs information from this pixel right here. And as you can see in just one layer, it can only get information from those, right? But now take the next layer, so the same pixel here, it can attend to these three, right? Now these three can each in turn attend to their neighbors, right? And I'm not going to draw everything, but the resolution field for this pixel here will end up being all of this, right? Now we still don't have our desired pixel in here, but if we just go one layer more, then this pixel right here, a different color, this pixel right here, right? The resolution field across the layers increases, because it's always incorporating information from downstream, and the downstream again incorporates information from the downstream, so eventually you can aggregate the same information. So instead of having a single layer with all of these connections, we have convolutional layers, which seem like a worse idea, because they can only do less things, attend to less things, but across the layers they actually can do the same thing, right? And that turns out to be a huge advantage of these convolutional layers, and that's why convolutional layers are used for image processing, and not the multi-layer perceptrons. So the same exact thing happens with axial attention, just in a different form. It is a bit poorly drawn here, I believe, but this is how you have to imagine it. As before, this pixel, the red pixel here, if I just have a normal transformer layer, the red pixel can attend to all of the other pixels in the image, right? That's the, that's basically, and each of the pixels can do that, so that's your d squared computation right here. Now, what we want to do is, in a convolutional layer, what we would do is, okay, you can only attend to your neighbors, and then in the next layer the neighbors can attend to their neighbors, and thereby you go out and out. In axial attention, you say, okay, this thing can only attend to its row and its column, right? That's it. You can only do attention to your row and your column, and I believe they don't even do it at the same time. So in one layer you can attend to the row you're in, and in the other you can attend to the column you're in. Now, let's see how the same thing happens as for a convolutional layer. So in the, basically, how then, if the red pixel needs access to information in this green pixel, how does it do that? So in the first layer it can attend to its row and its column, right? And so can every other pixel, including, sorry, including, of course, the pixel where that, so let's say this square here can also attend to its row and its column, and its row happens to be including the green one, right? So in layer one, this red square here gets information from the green square via row attention, right? And then in layer two now, this, our red square of interest now can row attend to this other red square here, so they get connected in layer two. I'm sorry, I don't want that. So you see that within just two layers we've transferred information from the green square via this red square to that red square. So we can, in the same way as a convolution, you can replace the long-range arbitrary dependencies between pixels by simply having multiple layers of restricted dependence. The same goes for this axial attention. So you can replace the arbitrary attention in layers, right? You can replace that by a two-step process where you first transfer information via the column and then transfer it via the row. It's a bit like, you know, in chess you can have a queen that can move any direction, especially diagonally, and then if you just have a rook you kind of need to do two moves. So in the queen is like the full attention and the rook is the multi-layer axial attention. They can achieve the same thing, you just need more layers. But as a trade-off you get a super, super saving in requirement of memory and computation, right? So they stress that, you know, kind of you can represent the same distributions with the axial attention. And you know, the trade-off is you just have to do multiple layers of it. Right, so this is axial attention and they are now able to incorporate this into their model right here. So they have, I believe, eight blocks, so four row attention, you see this right here, and four column attention blocks in their model. And finally they output this distribution here across their region of interest. Now this again is your, I believe, this 64 by 64 resolution. So you can see how they kind of aggregated information across the 64 using this axial attention. And then that makes their prediction in this one hour. So this is this. Alright, so this was a long way. So recap, they have 15-minute snapshots of this input data across along with some features. They use a spatial down sampler, which is a CNN, on each of them individually. Then they use a convolutional LSTM to encode this across time to end up with a single representation here at the end. Then they use axial attention in order to aggregate information across the spatial dimensions. They do this in multiple stages and at the end they make a participation prediction, which is a distribution, as you can see here. So as an output you directly get a distribution of results, which is also cool because the physical simulation, you have to let it run many, many times in order to get a distribution of results. And this neural network can simply give you a distribution right away. That's what they say right here. So they go a bit into the architecture compared to baseline. I want to get back to what I showed you at the beginning. This here is just the picture, kind of the picture book example. So left is the ground truth, in the middle is MatNet, and on the right is a baseline method. This here is in, as you can see, in two hours, in four, six and eight. So you can see the MatNet gives you as an output this distribution. What I find interesting, for example, is this sample two right here. So in this sample one you can see there is a consistent difference and this is the forecast time, so how much in advance you want to get it? No, this would be a one hour, but it can go up to eight hours. Here is a consistent gap in F1, which means the MatNet does it better across this span of time, which is for the top sample right here. For the bottom sample though, you can see here, there is a big gap at the beginning, again, there is a big gap at the beginning, and then this gap gets smaller and smaller and smaller. And this, I think, might give you an indication of, let's say, the weakness of this approach, doing it with neural networks. So with neural networks you kind of rely on regularities, you kind of rely on broad scale, correct things that you can learn from the data, and this might work well as long as things are regular, which of course across shorter time spans things tend to be more regular, right? But if you go for longer time spans, I believe there is more of a chaos element to it, like weather can be very dependent on very subtle things, and the physics simulation that is really taking into account the actual physics might be able to much, much better account for that. And that's why I believe across time here you'll see that the two models get closer together. That being said, MetNet of course is still on top here. But it will be interesting to forecast for longer even, though I haven't actually dig through their results, through their numerical results, but you can do that if you want. Alright, so this was it for MetNet and axial attention. I hope you liked this, and bye bye.
 Hi there. So what you're looking at here is a weather forecast model. Specifically the
 very top row is a new weather forecast model called NetNet by Google Research. So the goal
 of weather prediction is pretty simple. You want to know what the weather is going to
 be in the future. Specifically here you want to know precipitation rates. And so this is
 a new work that uses neural network instead of physical models in order to predict precipitation.
 So in the middle here you see this is the ground truth of what really happened at that
 particular time. You see precipitation rates in red here moving across the country. Now
 the bottom there is a physical model and as far as I understand it, physical models have
 been used so far to make weather predictions. Which basically means that you simulate these
 rain clouds and the movement of them across the country. And you do a physical simulation
 like a particle simulation type of thing and then that allows you to predict and then you
 run that maybe multiple times and you get an idea of the kind of distribution you're
 going to get. Now what MetNet does is it simply uses a neural network to predict the outcome
 directly. So there's no physical simulation involved. There is just a neural network that
 takes as input what's the situation now and maybe over a stretch of time. And then you
 ask it please make a prediction in eight hours or something. And then the MetNet will make
 that prediction and it will just output it like snap. No physical simulation needed.
 And you also see here that MetNet outputs things in kind of a cloud way, in a probabilistic
 way. In one forward pass you don't need to run it multiple times. But we'll get to that.
 On the bottom here you see the measurement. So the axis is F1. F1 is kind of the overlap
 of how well you're able to predict the precipitation. And you see here the MetNet is above the HRR
 baseline for most of this time. Up to 480 minutes into the future. Which is eight hours
 I believe. All right. So the paper is the following. It's called MetNet, a Neural Weather
 Model for Precipitation Forecasting. And I'm not going to read all the names here. The
 main corresponding authors are Caspar K. Sonderby and Nalkalke Brenner. And it's a
 team of Google research. So specifically they use the input of these two things here. So
 one is this GOS16, which is what you see here on the left. And the precipitation rates are
 here depicted on the right. So you want to take these things as input into your model.
 Now how do you do that? Of course we want to build a neural network. And this is the
 architecture they come up with. So on the bottom here they feed in the data. And they
 feed in the data in 15 minute interval from 90 minutes into the pass. So you have to imagine
 it like this. So there's a timeline. I'm going to use a little bit of a finer thing. So there's
 a timeline. And you let's say are here. This is now. And then here in the future, this
 is maybe one hour into the future. This is your target, right? This is you are here and
 you're looking out. You would like to know what's the precipitation going to be in one
 hour from now. What Metnet does is it takes an input. And specifically it takes the last
 90 minutes before now as an input. And it samples it in frequencies of 15 minute intervals.
 So each one of these is going to be 15 minutes. And each 15 minutes you get like a snapshot
 of this entire of the input region. Now the input region, if I can jump back here to the
 website for a second, they show it what the input region is. The input region, if you
 want to predict in the middle of this small square, the input region is actually the entire
 1024 square kilometers around it. So it's very big input. Though the actual region you
 consider is the inside 64 square kilometers. But you take in information from the big region.
 And the main point of the paper, I believe, is how to do that. All right, so each 15 minutes
 you take in a snapshot. And these are these snapshots here on the bottom. So these are,
 and you have to imagine in here, every 15 minutes there's a stack of these inputs. So
 what are these inputs? These inputs are some kind of features that you have. So there is
 the target time, which in this case would be this one hour here. There is the month,
 day and hour, which is important for weather prediction, right? So the time of year, time
 of day and so on. Longitude latitude is probably pretty important. Elevation map is probably
 pretty important. So these you can see, these are all maps. Now sometimes, and this is how
 you encode things in these. Since it's a neural network, you know, all of these things must
 be of the same dimensions here. So if you have 256 dimensions here and probably 256
 dimensions here, then all of these things must be of the same dimension. And if you
 want to give a feature such as the target time, which in this case, let's say it's one
 hour, you just put here one hour, what's one hour? Let's say 60 minutes. So you just put
 the number 60 here, 60, 60, 60, 60, 60, 256 times and 256 times 250, 65, sorry, 265 times,
 no, 56. I'm confusing with German. So this is how you encode features. It's pretty primitive,
 but it turns out it works the best if you do it this way. All right. So you have these
 planes and some, as I said, are just features such as the target time, month, day and hour
 and so on. Elevation, I guess is a map, is like an elevation map of the region you consider.
 And this corresponds now to this, these 64 kilometers times 64 kilometers here. And that's
 exactly what these center crops are here. So this center crop thing, that now, this thing
 here, this plane, sorry, is these 64 by 64 region. That's this plane here. And also the,
 that's the precipitation and the GOES, that's this thing here. Now we also have these down
 sampled things, which these are the 1024 kilometers. So this here and this here, these are the
 1024 square kilometer patches, but they are down sampled. So everything is down sampled,
 I guess, to 256 by 256 pixels. So you don't really take into account every nuance of that
 very big, of that very big input, but you do down sample it. So you kind of get the
 big picture of the outer frame and in the inner frame, you take it in a much higher
 resolution in order to get the details. All right. So you stack all of this up into a
 big tensor and then you feed it into here into a spatial down sampler, which I guess,
 no, I have read is a, some just a convolutional neural network, right? So this is your typical
 image processing pipeline. So you do this for each of these stacks, right? And then
 what you get out of it is a lower size representation right here. So you get these representations
 and then you let a temporal encoder run over it. What does a temporal encoder do? This
 in particular is a convolutional LSTM. And if you already know what an LSTM is, a convolutional
 LSTM is nothing more than an LSTM that has as intermediate layers, convolutional layers.
 So it's pretty suited to do, for example, videos or any sort of image processing that
 goes over time like this one. So the temporal encoder simply starts out here with an initial
 state. My pens are screwing me today. So it starts out here with an initial state and
 then it simply inputs each of these representations, takes them one by one, runs across time, right?
 And each time producing a new intermediate representation of the input until it finally
 reaches this here, final representation. So this thing here is a single final representation
 of all of this input, right? Of this entire time span of all of these stacks here. Yeah,
 so you can press this into a single input with first a convolutional network to downsample
 each time point individually and then with a recurrent neural network, an LSTM, to integrate
 the information over time. You end up with this single piece here. And then what you
 do, so you still, here you still retain kind of an image sort of thing. So this representation
 here, you can see it in the background. Maybe I'll get down my scribbles here. This here
 is still sort of an image tensor, though I guess it's a hidden representation, so you
 couldn't really look at it. But it still has dimensions of images. So this here is still,
 I think, the same or corresponding to these dimensions here. So this still has some spatial
 information where this might be north-south here in this axis, might be east and west,
 right? And then these are just the hidden channels, the channels of the hidden representations,
 right? So what you would like to do now is to basically encode information from the space
 around you. If you look at, let's look at one of these, one of the big pictures. What
 you would like to do in weather prediction, let's say you are right here. What's a good
 example, you are right here, right? Now if you want to know if this particular cloud
 over here is going to move to your direction, what you want to know is, for example, is
 there a mountain range here, right? Because then it's more probable that this cloud is
 going maybe to move up there. You would also want to know how this cloud here moves, right?
 If this cloud here moves somewhere here around, then it's probably this cloud down here might
 be pulled with it or something like this. So you're very much, sorry, you're there. You're
 very much kind of want to look out into each of the directions here and you want to incorporate
 kind of what's happening across the space. We're already used to kind of convolutional
 networks being able to do this, but in here the authors use attention to do that. So if
 you don't know what attention is, my most popular video is about attention and you can
 do attention for images. So the way that works is that you have a series of images of stacked
 blocks of a neural network. Let me draw this here. So you have an image here and let's
 say it has just four pixels, right? So you have the next layer of these four pixels,
 right? So you have layers of this. So the next layers of the four pixels, they all emit
 what are called queries and queries are just vectors. So each pixel emits a single vector.
 Let's say this, that, that, this, right? And each of the lower layers emits what is called
 a key. This, this, this, this. And now the keys and the queries are routed together based
 on their inner product. So these two would be routed together. This would probably be
 routed here. This as well. This would probably routed here. So what in effect each of the
 pixels of the higher layer can look at specific pixels of the lower layer. Now you can imagine
 this is exactly what we want here in that if there is a mountain range here and we might
 be interested in that. So we'd be able from our, from our point here to specifically attend
 to that location using, using attention, right? So the authors here build basically a stacked
 model of attention layers. And that's what's happening in the third part here. And this
 is the attention is in order to incorporate long range dependencies. As I made the example
 with the mountain range, this might be far away, but it might actually influence your
 weather very much. So the attention is to incorporate these long range dependencies.
 But the problem with attention is, is as you saw in the example, each of these pixels can
 attend to each of the pixels in the lower layer. So what you'd end up with, so each
 can attend to that. This can attend to each. This can attend to each. You'll see you'll
 end up with 16 connections. Can't even draw them. So you end up with 16 connections. In
 general, if you have D here, you will end up with a D squared number of things you need
 to calculate, right? So if this here, and now of course we have images. So generally
 we'll think of D by D pixels. Now we have D by D pixels and that thing squared number
 of things we need to calculate. This quickly gets too much. So in, for example, MNIST,
 you have 28 by 28 pixel images. This is 780 or 2 or something. I don't quite remember.
 But you'll have to calculate this squared many connections between things. This is simply
 impossible pretty quickly, especially if you scale up the images and then have some channels
 in here as well. So attention for image processing has been a bit lagging compared to natural
 language processing. In natural language processing, you usually have maybe 500 tokens or something.
 Images you have much more. So attention is much more expensive. So you can't really do
 it on current hardware. Now this paper uses something called axial attention. Axial attention
 is kind of the trick of how to make this tension happen for images. And for that I want to
 switch over to this paper. It's called Axial Attention in Multidimensional Transformers
 by some of the same authors. So Jonathan Ho and Nell Coutt Brenner, also of Google Brain
 and UC Berkeley, proposed this axial transformer. Now they originally proposed axial attention
 for autoregressive models. If you know transformers, they also started by making autoregressive
 models, so language modeling and so on. But we can decouple the axial attention from the
 autoregressivity of these models. So I'm not going to talk about autoregressive models,
 it's just axial attention. So what is axial attention? It's pretty simple actually. And
 I want to start by talking about convolutions. So what does a convolution do? Let's just
 take a one-dimensional image, which is pretty boring, but let's say it has these eight pixels
 here. So this is an image, it just has one row of eight pixels. What do I do when I run
 a convolutional filter across that? This is the lower layer, and now this is the next
 layer that is produced by a convolution. So for each of the pixels in the next layer,
 what I can do with the convolutional layer, I can look at its neighbors in the lower layer.
 So these three would be part of that. And then I go on to this, and again I look at
 its neighbors at these three. I might have done this in a different color. And then I
 look at this, and it can look at itself and its neighbors. So a convolution is pretty
 smart. And then of course in the next layer, that repeats. Now if you think, what's the
 difference between doing this and a fully connected layer? So if I have a fully connected
 layer, a classic neural network, a fully connected layer, then this pixel here would incorporate
 information from all of the pixels here. And this pixel here would incorporate information
 from all the pixels. Now why might this be better? Because the information that I want
 here for this pixel might depend on this pixel over here. So I might benefit from a connection
 over there, or it might benefit from this pixel here, which it can't reach. And with
 a convolutional network, I can't do that. Why are then convolutional networks preferable?
 Because the convolutional network can do the same thing across multiple layers. So let's
 assume again that this pixel here needs information from this pixel right here. And as you can
 see in just one layer, it can only get information from those, right? But now take the next layer,
 so the same pixel here, it can attend to these three, right? Now these three can each in
 turn attend to their neighbors, right? And I'm not going to draw everything, but the
 resolution field for this pixel here will end up being all of this, right? Now we still
 don't have our desired pixel in here, but if we just go one layer more, then this pixel
 right here, a different color, this pixel right here, right? The resolution field across
 the layers increases, because it's always incorporating information from downstream,
 and the downstream again incorporates information from the downstream, so eventually you can
 aggregate the same information. So instead of having a single layer with all of these
 connections, we have convolutional layers, which seem like a worse idea, because they
 can only do less things, attend to less things, but across the layers they actually can do
 the same thing, right? And that turns out to be a huge advantage of these convolutional
 layers, and that's why convolutional layers are used for image processing, and not the
 multi-layer perceptrons. So the same exact thing happens with axial attention, just in
 a different form. It is a bit poorly drawn here, I believe, but this is how you have
 to imagine it. As before, this pixel, the red pixel here, if I just have a normal transformer
 layer, the red pixel can attend to all of the other pixels in the image, right? That's
 the, that's basically, and each of the pixels can do that, so that's your d squared computation
 right here. Now, what we want to do is, in a convolutional layer, what we would do is,
 okay, you can only attend to your neighbors, and then in the next layer the neighbors can
 attend to their neighbors, and thereby you go out and out. In axial attention, you say,
 okay, this thing can only attend to its row and its column, right? That's it. You can
 only do attention to your row and your column, and I believe they don't even do it at the
 same time. So in one layer you can attend to the row you're in, and in the other you
 can attend to the column you're in. Now, let's see how the same thing happens as for a convolutional
 layer. So in the, basically, how then, if the red pixel needs access to information
 in this green pixel, how does it do that? So in the first layer it can attend to its
 row and its column, right? And so can every other pixel, including, sorry, including,
 of course, the pixel where that, so let's say this square here can also attend to its
 row and its column, and its row happens to be including the green one, right? So in layer
 one, this red square here gets information from the green square via row attention, right?
 And then in layer two now, this, our red square of interest now can row attend to this other
 red square here, so they get connected in layer two. I'm sorry, I don't want that. So
 you see that within just two layers we've transferred information from the green square
 via this red square to that red square. So we can, in the same way as a convolution,
 you can replace the long-range arbitrary dependencies between pixels by simply having multiple layers
 of restricted dependence. The same goes for this axial attention. So you can replace the
 arbitrary attention in layers, right? You can replace that by a two-step process where
 you first transfer information via the column and then transfer it via the row. It's a bit
 like, you know, in chess you can have a queen that can move any direction, especially diagonally,
 and then if you just have a rook you kind of need to do two moves. So in the queen is
 like the full attention and the rook is the multi-layer axial attention. They can achieve
 the same thing, you just need more layers. But as a trade-off you get a super, super
 saving in requirement of memory and computation, right? So they stress that, you know, kind
 of you can represent the same distributions with the axial attention. And you know, the
 trade-off is you just have to do multiple layers of it. Right, so this is axial attention
 and they are now able to incorporate this into their model right here. So they have,
 I believe, eight blocks, so four row attention, you see this right here, and four column attention
 blocks in their model. And finally they output this distribution here across their region
 of interest. Now this again is your, I believe, this 64 by 64 resolution. So you can see how
 they kind of aggregated information across the 64 using this axial attention. And then
 that makes their prediction in this one hour. So this is this. Alright, so this was a long
 way. So recap, they have 15-minute snapshots of this input data across along with some
 features. They use a spatial down sampler, which is a CNN, on each of them individually.
 Then they use a convolutional LSTM to encode this across time to end up with a single representation
 here at the end. Then they use axial attention in order to aggregate information across the
 spatial dimensions. They do this in multiple stages and at the end they make a participation
 prediction, which is a distribution, as you can see here. So as an output you directly
 get a distribution of results, which is also cool because the physical simulation, you
 have to let it run many, many times in order to get a distribution of results. And this
 neural network can simply give you a distribution right away. That's what they say right here.
 So they go a bit into the architecture compared to baseline. I want to get back to what I
 showed you at the beginning. This here is just the picture, kind of the picture book
 example. So left is the ground truth, in the middle is MatNet, and on the right is a baseline
 method. This here is in, as you can see, in two hours, in four, six and eight. So you
 can see the MatNet gives you as an output this distribution. What I find interesting,
 for example, is this sample two right here. So in this sample one you can see there is
 a consistent difference and this is the forecast time, so how much in advance you want to get
 it? No, this would be a one hour, but it can go up to eight hours. Here is a consistent
 gap in F1, which means the MatNet does it better across this span of time, which is
 for the top sample right here. For the bottom sample though, you can see here, there is
 a big gap at the beginning, again, there is a big gap at the beginning, and then this
 gap gets smaller and smaller and smaller. And this, I think, might give you an indication
 of, let's say, the weakness of this approach, doing it with neural networks. So with neural
 networks you kind of rely on regularities, you kind of rely on broad scale, correct things
 that you can learn from the data, and this might work well as long as things are regular,
 which of course across shorter time spans things tend to be more regular, right? But
 if you go for longer time spans, I believe there is more of a chaos element to it, like
 weather can be very dependent on very subtle things, and the physics simulation that is
 really taking into account the actual physics might be able to much, much better account
 for that. And that's why I believe across time here you'll see that the two models get
 closer together. That being said, MetNet of course is still on top here. But it will be
 interesting to forecast for longer even, though I haven't actually dig through their results,
 through their numerical results, but you can do that if you want. Alright, so this was
 it for MetNet and axial attention. I hope you liked this, and bye bye.