Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation (Paper Explained)
deep learning machine learning arxiv explained neural networks ai artificial intelligence paper google cnn resnet big bird bigbird attention attention mechanism attention for images transformer for images transformer bert convolutions window neighbors axial attention position embeddings positional encodings quadratic memory panoptic segmentation coco imagenet cityscapes softmax routing
#ai #machinelearning #attention

Convolutional Neural Networks have dominated image processing for the last decade, but transformers are quickly replacing traditional models. This paper proposes a fully attentional model for images by combining learned Positional Embeddings with Axial Attention. This new model can compete with CNNs on image classification and achieve state-of-the-art in various image segmentation tasks.

OUTLINE:
0:00 - Intro & Overview
4:10 - This Paper's Contributions
6:20 - From Convolution to Self-Attention for Images
16:30 - Learned Positional Embeddings
24:20 - Propagating Positional Embeddings through Layers
27:00 - Traditional vs Position-Augmented Attention
31:10 - Axial Attention
44:25 - Replacing Convolutions in ResNet
46:10 - Experimental Results & Examples

Paper: https://arxiv.org/abs/2003.07853
Code: https://github.com/csrhddlam/axial-deeplab

My Video on BigBird: https://youtu.be/WVPE62Gk3EM
My Video on ResNet: https://youtu.be/GWt6Fu05voI
My Video on Attention: https://youtu.be/iDulhoQ2pro

Abstract:
Convolution exploits locality for efficiency at a cost of missing long range context. Self-attention has been adopted to augment CNNs with non-local interactions. Recent works prove it possible to stack self-attention layers to obtain a fully attentional network by restricting the attention to a local region. In this paper, we attempt to remove this constraint by factorizing 2D self-attention into two 1D self-attentions. This reduces computation complexity and allows performing attention within a larger or even global region. In companion, we also propose a position-sensitive self-attention design. Combining both yields our position-sensitive axial-attention layer, a novel building block that one could stack to form axial-attention models for image classification and dense prediction. We demonstrate the effectiveness of our model on four large-scale datasets. In particular, our model outperforms all existing stand-alone self-attention models on ImageNet. Our Axial-DeepLab improves 2.8% PQ over bottom-up state-of-the-art on COCO test-dev. This previous state-of-the-art is attained by our small variant that is 3.8x parameter-efficient and 27x computation-efficient. Axial-DeepLab also achieves state-of-the-art results on Mapillary Vistas and Cityscapes.

Authors: Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille, Liang-Chieh Chen

Links:
YouTube: https://www.youtube.com/c/yannickilcher
Twitter: https://twitter.com/ykilcher
Discord: https://discord.gg/4H8xxDF
BitChute: https://www.bitchute.com/channel/yannic-kilcher
Minds: https://www.minds.com/ykilcher
Parler: https://parler.com/profile/YannicKilcher
LinkedIn: https://www.linkedin.com/in/yannic-kilcher-488534136/

If you want to support me, the best thing to do is to share out the content :)

If you want to support me financially (completely optional and voluntary, but a lot of people have asked for this):
SubscribeStar: https://www.subscribestar.com/yannickilcher
Patreon: https://www.patreon.com/yannickilcher
Bitcoin (BTC): bc1q49lsw3q325tr58ygf8sudx2dqfguclvngvy2cq
Ethereum (ETH): 0x7ad3513E3B8f66799f507Aa7874b1B0eBC7F85e2
Litecoin (LTC): LQW2TRyKYetVC8WjFkhpPhtpbDM4Vw7r9m
Monero (XMR): 4ACL8AGrEo5hAir8A9CeVrW8pEauWvnp1WnSDZxW7tziCDLhZAGsgzhRQABDnFy8yuM9fWJDviJPHKRjV4FWt19CJZN9D4n
 Transformers are quickly coming for your favorite models. Yesterday they replaced LSTMs in NLP. They used to be good at NLP but we now have transformers. Think again. Today we're going to see that maybe in the near future transformers will replace convolutions in image processing. So this paper is a step in towards this direction. You just wonder what is it going to be tomorrow. Maybe linear regression is going to be replaced just by giant transformers trained on 5,000 TPUs. Who knows? We'll see. In any case we're looking at Axial Deep Lab standalone axial attention for panoptic segmentation by Hui Yu Wang, Yuh-Kun Chu, Bradley Green, Hartwick Adam, Alan Yuel and Liang-Qi Chen of John Hopkins University and Google Research. So this paper combines a bunch of techniques that have been introduced recently to deal with attention in problems where you would traditionally use a convolution. So in this particular case they deal with this problem of panoptic segmentation which basically you'll see you'll get an image and there's a bunch of stuff on the image like a cat here and a house right here and you're supposed to color the pixels of the same object the same so you see all these pixels here are house and then all these pixels these pixels right here are cat and so on and then there's also the background so all these pixels right here I know beautiful beautiful beautiful our background. So for this problem it's kind of important that there you you you're very precise first of all so you can look at you know pixels or clusters of pixels and also that you take long-range dependencies into account because if you for example recognize that this is a house and you recognize that here's a wall right here you might be able to much better classify what is wall over here and what isn't. So the kind of long-range dependencies play a role in these problems across images and usually attention mechanisms are pretty good for these long-range dependencies but they're also expensive and that's what this paper deals with. So they use this axial attention that has been introduced for exactly resolving this problem in types of data like images or higher order tensors and they also combine this together with learned positional encodings which we've also seen time and time again throughout the kind of transformer and attention literature. So the combination of axial attention these learn positional embeddings allows them to replace the ResNet backbone that usually is found in panoptic segmentation models with a standalone attention. So they build models that are partial replace the convolutions with attention modules or replace them entirely so the entire model is going to be just an attention model so no more convolutions in it and they perform pretty well in classic tasks like they test on ImageNet classification they perform pretty well and they achieve state-of-the-art on some of these segmentation tasks. So we'll go through the model right here this is a very very extensive paper in terms of experimental evaluation what I want to get into is mainly how the method works and show you what their model looks like. So we'll go through it and as always let me know what you think in the comments and tell me if you liked it or not share it out if you did. Alright so they go over a very long list of prior work which is you know pretty pretty cool and here they say their contributions so their contributions are fourfold. First of all the proposed method is the first attempt to build standalone attention models with larger large or a global receptive field and we'll see what that means. We propose position sensitive attention layer that makes better use of positional information without adding much computational cost. We show that axial attention works well not only as a standalone model on image classification but also as a backbone on panoptic segmentation, instance segmentation and semantic segmentation. Maybe what I did before described before was instance or semantic segmentation and not panoptic segmentation. Excuse me if that's the case. As you can see it can be used for various various image tasks. Lastly our axial deep lab improved significantly over bottom-up state-of-the-art on Cocoa achieving comparable performance of two stage methods. We also surpassed previous state-of-the-art methods on mapillary, vistas and city scapes. So these are various tasks as I said and also what they don't mention here is that they perform fairly well on ImageNet. In fact in the abstract they formulate this as in particular our model outperforms all existing standalone self-attention models on ImageNet. That's a way to phrase it. You just exclude all of the other models until you're the best. Outperforms all existing standalone self-attention models on ImageNet. That's good. There's something to be said of comparing apples to apples but you can also go overboard if you want to make your work look as good as possible. Of course everyone does that and there's no particular shame in it. We're going to build up our model right here and the basic element of this model is going to be this self-attention mechanism. Quickly because I know you all know what it is but very quickly you want to perform this action right here over a region right here. There is always a query and now the subscripts here are going to be important in this paper. The query is at a given position, position O and you can see that's the O right here. I'm going to call it the output. I guess that's what they said as well. So the output position. You want to go over all of the input positions and you want to aggregate data from all of the input positions. That's right here. How do you aggregate data? By this softmax operator right here. You can see the key also has a P right here and the softmax is over the axis of P. In particular case of the images what does that mean? If you have an image right here it's made into pixels. You have pixels. Now a transformer or generally these attention models, what you can imagine is they always transform a data point into a data point of the same dimensions. This doesn't have to be actually and I think one of the developments that is going to come in coming years or months or weeks, maybe someone's already doing it, is in fact to play more with this arbitrary constraint that we're imposing on ourselves because it's not really clear that this is the best thing. But for now an attention layer is always transforming a data point, here a 4x4 image, into a data point of the same size. Also a 4x4 image right here. Now this is, as I said, this is quite simplified but it is true in NLP where we always transform our whatever 512 sequence, a token sequence into a 512 token sequence and it is true here. Now the output is going to be here on the right and the question always is, okay so I'll go over these pixels right here and for every pixel, let's say for this pixel, I'm going to ask what data goes there? What's the output of the layer at that particular pixel? And the output of the layer is going to be somehow dependent on the input right here. Now if you know classic convolutional models, what the classic convolutional model says, the output of this is going to be dependent on this region right here, if it's like a 3x3 filter. So you have this convolutional filter and that means that blue dot on the right is going to pay attention to its own location in the input plus everything around it. And then every single data point here is going to do that. So for example this green data point is going to pay attention to this region right here. Now there's a border so there's maybe some padding but the question is always where does the information come from and how is it aggregated? In a convolution layer, what happens in a convolution layer? In a convolution layer you simply have your filter and the filter has numbers in it like three and five and eight and so on. And what you're going to do is you're going to take this region right here, this blue region of the lower layer and that's also filled with numbers like seven, what's a good number? Zero. Zero is a nice number and you're going to multiply those and then you're going to sum them up and then you're going to put that on where the blue dot is. So where does the information come from in the convolution? From around the location, from around the output location but in the input. So you go to the input at the same location as where you want the output to be, you take the neighborhood and there is a fixed scheme of aggregating the neighborhood. And then you multiply and you sum across it. In contrast to this, in a fully attentional model where does the information come from? Let's again look at the blue dot and let's consider it fully attentional. Where does the information come from? Everywhere, anywhere, anywhere at all. The information comes from everywhere. Now how do I know how to aggregate the information? So it's no longer in a neighborhood. How do I know how to aggregate the information? That's also different. So two things are different. Now in a convolution I would have another four by four grid here that's pre-specified but in the attention model this here is basically all filled with question marks. Question mark, question mark. What number goes here? In the end I also do this multiply and I sum it up and I put it right here. But how do these numbers come to be? Well these numbers also come, these are dynamically computed also from the input. It's a bit special but this is how attention works. So every pixel gets to decide where information comes from and how it is aggregated. Basically it comes from anywhere and how it is aggregated is dynamic depending on the pixel. If you still don't understand it maybe pay out to watch a video on attention itself. I happen to have made one but you can watch any one when you understand that you will understand the extension here to the image. It's the exact same thing as with the sequence except the pixels are basically one long sequence in the image. So this would be a fully attentional model down here. Now what's the problem here? The problem is that pictures are pretty large. So even something like MNIST which is like 28 by 28 is like 700 pixels plus. I don't remember exactly but it's like about 700 pixels. And our big transformers now, so BERT, a very famous transformer, takes inputs that are like 512 in length. You already need pretty decent hardware to run this. And the requirements on memory and compute scale quadratically with the input length. So already with MNIST you're in pretty pretty shady territory. If you go up to something like ImageNet which is like 225 by 225 that's bad. That's not good. So you have to come up with something else. So people have been playing around, the reason why I introduced it this way, is people have been playing around a bit with sort of coming up with an intermediate with a compromise between the two. So the compromise that this paper here focuses on is going to be a compromise where we, you remember when I said where does the information for a given pixel come from? And we said okay it can come from anywhere in the attention framework and that's good because that allows us to make super long-range connections. So any pixel can aggregate information from any other pixel and not even in a fixed way but in a dynamic way. So depending on the pixel value itself and the other values it can decide how it wants to aggregate information. That turns out to be expensive, right? Every pixel together with every pixel, well that's quadratic. Okay so what do we do? We make a third method that's going to be a compromise and the compromise is going to be the following. The compromise is going to be alright we still do the dynamic aggregation which means that we still do the attention thing. However we're going to restrict back to this neighborhood region of the convolution. So in this model where does information for the blue dot come from? It again comes from this neighborhood right here and this number, the size here is going to be called m. So it still comes from that m by m neighborhood so a pixel can only aggregate information from its neighbors but contrary to a convolution how it aggregates the information like this what in convolution would be a kernel. The kernel is made dynamically by the attention module and it's made dynamically on a case-by-case basis. So we restrict it to a neighborhood, multiply, sum it up and then put it into the output and we do that for every pixel. Now it resembles much more a convolution, simply a convolution with this dynamic matrix right here. And that's the starting point for this paper. So this paper does two things to this. It says okay we can augment this by so-called positional embeddings. A positional embedding you might know from the sequence transformers. So if I have a sequence my cat is tall, I don't even know what that means for a cat. But okay what in a positional encoding so if you use a transformer and you transform this as we said into a sequence of equal length and then transformers basically information routing the transformer simply sees the lower layer sequence as a set not as a sequence. It has no notion of what's neighboring to what, what comes from where. So it pays to tell the transformer by the way this is word one, this is word two, this is word three, this is word four. There are various ways to do it. Transformers usually have fairly complicated kind of sine wave based positional encodings that bring many advantages with them. In this case they say well it might pay pay off to learn where actually these things are in this neighborhood. So they experiment with relative positional encoding which means they annotate this neighborhood with something like look here in the middle it's a 0 0 here is like 0 1 here it's 0 negative 1 negative 1 0 and so on. So they annotate it with these positional encodings. Now this is this would be the easy way what they actually do is they simply they give the model a matrix like this and they learn that matrix by heart let's say. So the positional encodings are relative positional encodings and they are learned okay so you can do that you can learn positional encoding so if you don't want to do the one two three four right here you simply say well here is a vector here is a vector here is a vector and here is also a vector. Now model you're already learning like all the weights to make this thing here happen and you're already learning your output weights up here right using back propagation. Why don't you learn yourself what you would like for position one like what kind of information you would like to be to have there using back propagation right so the model you provide them up you always provide the same vector so this is the same vector for position one and you have a different vector for position two and you have a different vector for position three right so but across all of the data points these vectors are going to be the same so the vector one is always going to be that same vector for all of the data points so the model somehow must learn independent of the data point what it means to be in position one so the model must learn how it wants to fill that vector that's called a learned positional embeddings we've seen this in many models so far it usually works pretty well and I guess here is it works especially well if you have these relative positional encodings and so this thing here is not going to be an actual matrix filled with these numbers it's going to be a learned matrix a trainable matrix that is filled that the network is allowed to fill with numbers right like three five eight and you might be you might notice that we've seen this before right so ultimately the information in this blue thing right here is going to depend on this dynamically created aggregating of information through the neighborhood and this statically learned aggregation of information throughout the neighborhood which is a con which is sort of a convolution right because in the convolution you've already seen here this is a statically learned map of how to aggregate information from the neighborhood of a pixel so I think even though there are slight differences they for example say this these are the same across attention heads and so on however I suspect that you can think of these learned positional embeddings to be to be kind of like what you learn in a convolution not exactly so no I I think I made a mistake and we'll see it in the formula we'll see it in the formula yeah okay so here they introduce these positional embeddings okay so you see that we previously we had the softmax previously we had this and this okay so this is the lower layer this is the information that comes into the layer and now it's it's transformed into values by a linear matrix but essentially this is the lower layer and for each of the output locations you want to know how should I aggregate information from that lower layer and you do this by this thing here this thing here is this dynamically constructed attention matrix using also the softmax okay so how should you aggregate information this comes from this query at the output position and the keys at the input position and now you add to that this math this thing right here which is again an inner product between the career query and the positional encodings okay so the positional encodings are going to be learned and hard coded but they still are modified by the queries so the query can still pay attention the difference is the keys depend on the input while the positional encoding does not depend on the input so the queries can decide I want to gather information from this and this and this type of information so that would be the key or it can decide I would like very much to look at pixels that are somehow on the bottom right of the pixel that I am now that would be the positional encodings and that's that's the mistake I made when I said it's equivalent to a convolution it is not because the query can still it's still modulated by that query vector of how to aggregate information otherwise you would have this to be a standalone multiplied by the input right here but it sort of pays off to think of it like what you do in the convolution so in the convolution you learn how to aggregate information basically based on position relative position to the position that you want to output and here you do a similar thing you learn static position embeddings that you then can attend to with your queries alright so these are the position embeddings and they make use of those position embeddings in fact they attend them to the following in this work we enable the output to retrieve relative positions beside the content based on query key affinities formally so the problem up here is that okay you have these position embeddings and here are the outputs but if you do this in multiple layers right if you do let's let's go with 1d sequences if you do this in multiple layers and here you annotate the position let's just go one two three four and okay this layer can make use of that right we gather stuff from here but then when this layer when this layer gathers information from here the where the information comes from in the layer below is some is how somehow getting lost right so it cannot kind of pull through this information to here or at least it's very complicated this model extends this positional embeddings in order to pull through that information so as you can see there are two new things right here the biggest important new thing is that right here we don't so here is how we aggregate information okay and here is the information that we aggregate over now you can see previously this was just this value vector and now it is extended to the position to positional embeddings learned positional embeddings okay so the this with this you're able to route the positional embeddings to the output and also here you can see the attention gets fairly complex so you have query key attention which is classic attention the queries can attend to positional encodings but also the keys can attend to positional encodings so not only can not only can the the node on top say I would like to attend to position three position three can also say well together with me positions two and four are are fairly important I guess that's what that's what that is maybe I'm mistaken here but you can see right here there is an interaction between the keys and the positional encoding right here now these positional encodings they are different for the queries keys and values but ultimately we don't that doesn't make too much of a difference so here is a contrast between what a traditional attention layer would do and what they would do so a traditional attention layer gets the input X and transforms it by means of these linear transformations right here into the queries these are the queries it's called Q into the keys and into the values okay then it does a matrix multiplication with the keys and the queries and puts that through a softmax so this here is going to be our attention matrix this is the attention matrix and the attention matrix is multiplied here by the values and that determines our output okay again the attention matrix defines how we aggregate information and the values is what information do we aggregate you know for the output in contrast when we introduce these positional encodings you can see right here again we have query key and value now it gets a little bit more more more complex right here namely we do this query key multiplication right here but we also multiply the query by these positional embeddings for Q we also multiply the keys by the positional embeddings for K and all of this together so this is a big plus right here all of this together is routed through the softmax okay and now the diagram is a little bit complicated now you can see the softmax aggregates information from here and from this learned position embeddings I would rather have they would just use it like they did in the formula do V plus R and say that's going to be the information that we are aggregating and the softmax here the output of the softmax is going to be how we aggregate information this is the attention all right I hope that's sort of clear you introduce these positional embeddings for queries keys and values and that allows the model to have a sense of where the information is coming from basically what positions which if you drop the convolutions so the convolution had this intrinsically because in your convolutional kernel right can I I'm dumb if in your convolutional kernel the number right here if there was a seven right here that meant that wherever you are whatever is on the bottom right is seven important okay so that's that was the convolution have this intrinsically here if you just do attention the we as humans we see it in a in this kind of grid form but the machine doesn't the machine simply sees a set of pixels it simply sees you can this is to the attention mechanism this is exactly the same as a long list of pixels or a discontinued set it doesn't matter to the machine so it's like the problems a feed-forward network has so we need to annotate it we have to give it positional information and learned positional information seems to work very well right here though you could think of static positional information okay this is the first thing the positional embeddings that now help the attention mechanism see where the information is coming from that's really important in pictures so we add that the second thing they do is this so-called axial attention now axial attention is sort of a let's say trick in order to reduce the load on a the load on an attention mechanism so what does it mean we've already we've already seen in sequences right if I have a sequence a sequence layer that's going to be n squared connections between the two now there are various ways to restrict that so instead of having all of these connections let's say from one node we've already seen wait if we just restrict it to let's say only this thing right here only this stuff that can be that is lower right that is lower in complexity and this in this case it would be just a neighborhood so that's what we've done that's this this M thing right here however we can also do it in different ways since this is a set anyway we can simply say maybe we should just always skip one we could like do attention like this and that would be just fine too right that would also leave away some of the information but you gain in computational efficiency there are various trade-offs now in a picture you have the same options right so you can do the neighborhood thing as we did or you can say where should the green pixel pay attention to axial attention says the green pixel should pay attention to only the row where it is in okay that's it should ignore the rest of the input it should only pay attention to that row where it is in and then in the next layer we'll flip it then the green pixel the same green pixel will pay attention to only the column it is in okay so that's that's called axial attention but don't think like don't don't there is nothing special about this being an axis or whatnot you could also define and it would not be called axial attention but you could define it makes the same sense to say well that green pixel it just depends on this diagonal right here just in the in this layer it just does this diagonal and then in the next layer it does like the anti diagonal you can say I just choose five random pixels in this layer and five random pixels in the next layer and that would work as well we've already seen this in this paper called big bird right the big big big bird but big bird so big bird explicitly used random connections in the attention mechanism and their argument was well if we use different random connections in each layer then information can travel pretty fast through the network so what's the problem with these neighborhoods right here what's the problem with neighborhood attention like this the problem is that you break the long-range dependencies so let's see what happens if information needs to go from this pixel to this pixel or this node to this node but if information needs to travel from this note to this note in a classic attention mechanism everything's connected to everything so that node in the next layer can simply aggregate information from here well that's not possible if you do this kind of neighborhood attention as we've done here if I do neighborhood attention then at most right because the neighborhood is three long at most this node right here can aggregate information from this node and then again it's three long in the next step so now this node can aggregate information from this node okay because the in the neighborhood is three long and you can only attend to within your neighborhood this means that if I want to send information to something that's really far away I need to I need to go many many layers right I need to go layer layer layer layer and this has been well known this has already been a like a problem this has already been a property of convolutional neural networks so convolutions specifically traded off the fully connectedness of fully connected layers to local connections convolutions but that means that you have to go very deep in order to make long-range connections you can't just make them in one step the same problem right here that is paper Big Bird argued that if you have random connections instead of neighborhood connections just the property of random graphs mean that you you are pretty fast in sending information around so because in a random graph of size n you on average all two nodes are connected by path lengths of log n this is much faster because in this neighborhood thing two nodes are connected in a path length of order of n right you can you can pretty easily see that if I make the sequence longer I need that many more steps in order to send it around in fact it's like something like n divided by m this neighborhood size in a random graph it's log n and in this axial attention that's why I introduced it it's 2 okay every every two nodes are connected by two steps if if node if this node right here needs to send information to this node right here in a classic attention mechanism you could do some one step because every pixel attends to every other pixel however right now we have to we have to see so this node attends in this layer sorry I have to think so how do we send information between the two we select this node right here in the first layer this node pays attention to this row okay which includes the red dot so the red dot can send information to the X in this layer in the next layer we select this node right here which is our target node where the information should go to it pays attention to all of this column which includes that X that before right this this X right here where we send information to so it takes two layers two steps to send information from any node to any other node well that's pretty good so this axial attention if you stack them on top of each other you sacrifice a little bit of of being able to send information from anywhere to anywhere for the pleasure of not having this quadratic attention anymore as you can see your attention mechanism is now as long or as big as your column or is wide or your row is high again this isn't this isn't specific to rows or columns you could do this as I said with these kind of diagonals you could do it with any other sort of sub pattern where you can sort of guarantee that the overlap between the layers is enough so you can send information around pretty efficiently and they use this right here so this axial attention you can see the formula is exactly the same the only change from before is this part right here you can see that the neighborhood that they aggregate over is no longer M by M it is now 1 by M so we've seen them going from if this is the the full input image and you want to you want to see where to attend what this paper does is it says a classic sorry a convolutional neural network would be attending to some sub part right this is convolution an attention mechanism pure attention would attend to everything but this is attention then what we are doing sorry that was a mistake what other people were doing were reverting back this attention to a sub part this kind of neighborhood attention okay but that was still you know you still have M squared you still have O of M squared because of the attention mechanism now what we are doing is we are going even lower we're actually going 1 by M okay this this is with with axial attention so in general it's 1 by M and then in the next layer we can go 1 by M in this direction and have that property and because it's so cheap now right because it's now O of M to compute this we might as well make M as long as the row itself okay so their last step is going to be to say okay we have 1 by M right here and that's going to be the row itself now you can see right here that they say axial attention reduces the complexity to HWM this enables global receptive field which is achieved by setting the span M directly to the whole input features optionally one could also use a fixed M value in order to reduce memory footprint on huge feature apps which is something that they're going to do later on ImageNet I believe so when they have big inputs or big outputs they actually do use a smaller M what you can see right here is that I wasn't really that wasn't really correct of me to say that it's now O of M because you you still have the entire query space so you multiply query by by keys now even if you make the keys to be 1 by M yes you reduce definitely you reduce this from height times width to times height times width to this but then you can see this thing right here if you take it and let's say we have this kind of row pattern and we replace M by the width then we have width squared so again the square appears however it's smaller than the original attention the original attention was H squared W squared right because HW is the image and you need that squared in order to do the attention mechanism now we've basically reduced one of the factors it is still an attention mechanism so there's still attention going but we've basically transformed the the image we've reduced it to one column now the one column is still attention so this is still attention like here so this now reduces to the attention that you see in a in a single sequence okay if you see the image as a long stretch of pixels what this does is basically it's up it simply subdivides that into neighborhoods so we're back to neighborhoods basically but we shift the neighborhoods from layer to layer so in the next layer the neighborhoods are going to be just alternating right the neighborhoods is going to be this is one neighborhood connected to this neighborhood connected to this neighborhood I hope this makes sense so it's going to be it's basically a mix between if you if you if you were to do this in convolution you could do one layer where it's neighborhood convolution and then one layer where it's like convolution with holes in it I think they're called atras convolutions or something like this with like giant holes in it that are exact is exactly the anti pattern of the neighborhood convolution from before that's what this is so you see their axial attention block right here their axial attention block replaces the ResNet block so if you know ResNet I've done a paper on ResNet ResNet basically takes the input pipes it through straight and adds to it whatever comes out of this operation okay that's a residual block now usually this thing here would be convolutions and convolutions and they are now replaced by these multi head axial attention you can see there is a multi head attention in the height and there is a multi head attention in the width and that gives us the property that every node can send around information to every other node in two steps I don't like the fact that there is only two because what this I guess this gives a significant bias to one or the other direction depending on the order that you do them in if if I had done this I maybe would have used three of them because it depends on how you want to aggregate information right like here you train the network specifically to aggregate information first in this direction and then in this direction which might work and it'll give you that sending around information anywhere so maybe they've actually tried and it just performed the same so I just might have a dumb suggestion right here in any case they simply replace in we've come a long way right we've gone to like neighborhoods and blah blah blah blah ultimately take a ResNet place the convolutions with the height axis attention and the width axis attention and we're good and then we come to results so that's it you have these positional embeddings you have the axial attention and it turns out that on ImageNet they perform fairly fairly well so you can see that models like a ResNet 50 model will get a 76.9 on ImageNet which is not state-of-the-art but it's also not it's not bad right the ResNet 50 is pretty good model you can see the full axial attention right here achieves a 78.1 also not state-of-the-art but still pretty good and as they say it's the best fully attentional model on ImageNet or standalone attention model on ImageNet so where this model really shines is where you really have to make long-range connections between pixels and that's these kind of segmentation tasks and I want to skip the tables right here they're best and everything and go to the appendix where they have some examples of this so here you can see specifically this is the original image you have a ground truth and you have the differences between their model this axial deep lab and the panoptic deep lab that is a baseline for them and you can see that the the failure cases here are are pretty you know show how show how the axial deep lab is better I don't know if they are cherry-picked or not but at least you can see that at some point so it handles occlusions better it handles instances better so here you see that the ground truth separates the person from the tie and the axial attention is able to do this but the the baseline is not able to do this correctly because it labels part of that white shirt also as and you can see why there's kind of a delimiter line here here here here but if you have long-range dependencies right if you have long-range dependencies in the model the model will recognize wait wait that's that must be the same thing as this thing here and this thing here and this thing here so that must be the same object it's simply that the shirt was occluded by the tie and goes beneath it and now appears again it's not a different it's not part of the tie and it's not part of the of a different object it's actually part of the shirt so the long-range attention you can see at these examples sometimes here okay this might not be an instance of super duper long-range dependencies this is simply where the model performs better so you can see here the ground truth has that surfboard segmented and the baseline does not that this can also just be you know there are a lot of tricks to make this work of course and you throw a lot of compute at it and sometimes you just get better numbers or part of the better numbers because of the additional compute right here what do we have so you can see occlusions it appears to handle occlusions in a better way and this might be due to this axial attention it might be due to the positional embeddings but you can see that the ground truth here has the laptop between the person's hands segmented the baseline cannot do that but the axial attention does do that and I don't know what this is honestly this is you can you can see though the axial attention also misses the fact that it should segment this in the background and if this occlusion handling you can see best in this example where the person in the back reappears on both sides of that person so you can see that the axial attention manages to segment that where that is just a mutant person right here the ground truth is equally shaky I think there is might be some ambiguity of how you can segment these images obviously but you can see the fact that there are long-range dependencies probably helped with this saying that wait in this image there's this white stuff right here and there's this white stuff right here and connecting these two regions with attention probably helped in segmenting these to be the same object even though you can see there is a break in the object so there is a break no at no point is the object on the left touching or the segment on the left touching the segment on the right and still the model manages to put those into the same label category there is the last last thing where they they want to research what their heads learn and usually you can do this right you can kind of visualize what the attention heads learn so in this case right here in the column heads the way you have to read this is that this particular head right here aggregates information from its column so everywhere where it lights up it there's a lot of information being routed you can see specifically in this here the heads of the people or the heads of the persons in the picture light up fairly well so for example this head right here is probably aggregating information a lot from this position right here and this head here is aggregating information from this position so you can deduce that that particular attention head probably deals with people's faces whereas that particular attention head probably deals you can see the attention is mostly on the grass right here and you can see the same with the for the row heads now their description here is that we notice that column head one corresponds to human heads while column head four course correlates with the field only which you know you can interpret it as this this seemed pretty clear but then they say something like row head six focuses on relatively large relatively local regions where column head five pools all over the image so row head six which is this thing right here you can see that okay it maybe focuses on small regions though you can see okay what like here you can get it that's a person but in other places I don't know where column head five pools over the whole image and this I don't know maybe they just needed something more to say because they put these pictures here they were like okay the the column heads are really nice because we couldn't like these this one's really nice because it you know just pays attention to the people and this one looks really nice because it pays attention to the field and but we can't really put the column head attention without putting the row head attention but then none of the row heads really are like super distinctive on a particular thing in the image so we need to come up with something that we can say and then you look at this one this is there's not a lot of attention so we need to contrast this with something then you would think that they contrast it with another row head but then there's no row head that does this whole image so there's like a column at five yeah I'm not sure if there's there's a bit of there is a bit of tactical writing going on here I suspect I mean still you know it's doing something cool but yeah there's there's definitely an element of sales in when you do when you do where I research papers and just not to this data but just props to the lines in front of the histograms makes it so much easier to read how big the stupid bars are why does everyone put the lines behind the histogram I probably do that myself and now I'm just I'm realizing how much easier that is alright there is a big big big experimental section right here and there's a big appendix where you can read up all of the different numbers comparisons ablations whatnot ultimately I just wanted to go over the method basically putting this into context with other things like putting this into context with stuff like Big Bird axial attention other positional encodings how it how it relates to convolutions how it relates to feed-forward networks and what convolutions did to feed-forward networks and so on I hope you have at least a little bit gained an understanding of what's going on here and with that said I see you next time bye bye
 Transformers are quickly coming for your favorite models. Yesterday they replaced
 LSTMs in NLP. They used to be good at NLP but we now have transformers. Think again.
 Today we're going to see that maybe in the near future transformers will
 replace convolutions in image processing. So this paper is a step in towards
 this direction. You just wonder what is it going to be tomorrow. Maybe linear
 regression is going to be replaced just by giant transformers trained on 5,000
 TPUs. Who knows? We'll see. In any case we're looking at Axial Deep Lab
 standalone axial attention for panoptic segmentation by Hui Yu Wang, Yuh-Kun Chu,
 Bradley Green, Hartwick Adam, Alan Yuel and Liang-Qi Chen of John Hopkins
 University and Google Research. So this paper combines a bunch of techniques that
 have been introduced recently to deal with attention in problems where you
 would traditionally use a convolution. So in this particular case they deal with
 this problem of panoptic segmentation which basically you'll see you'll get an
 image and there's a bunch of stuff on the image like a cat here and a house
 right here and you're supposed to color the pixels of the same object the same
 so you see all these pixels here are house and then all these pixels these
 pixels right here are cat and so on and then there's also the background so all
 these pixels right here I know beautiful beautiful beautiful our background.
 So for this problem it's kind of important that there you you you're very
 precise first of all so you can look at you know pixels or clusters of pixels and
 also that you take long-range dependencies into account because if you
 for example recognize that this is a house and you recognize that here's a
 wall right here you might be able to much better classify what is wall over
 here and what isn't. So the kind of long-range dependencies play a role in
 these problems across images and usually attention mechanisms are pretty good for
 these long-range dependencies but they're also expensive and that's what
 this paper deals with. So they use this axial attention that has been introduced
 for exactly resolving this problem in types of data like images or higher
 order tensors and they also combine this together with learned positional
 encodings which we've also seen time and time again throughout the kind of
 transformer and attention literature. So the combination of axial attention these
 learn positional embeddings allows them to replace the ResNet backbone that
 usually is found in panoptic segmentation models with a standalone
 attention. So they build models that are partial replace the convolutions with
 attention modules or replace them entirely so the entire model is going to
 be just an attention model so no more convolutions in it and they perform
 pretty well in classic tasks like they test on ImageNet classification they
 perform pretty well and they achieve state-of-the-art on some of these
 segmentation tasks. So we'll go through the model right here this is a very very
 extensive paper in terms of experimental evaluation what I want to get into is
 mainly how the method works and show you what their model looks like. So we'll go
 through it and as always let me know what you think in the comments and tell
 me if you liked it or not share it out if you did. Alright so they go over a
 very long list of prior work which is you know pretty pretty cool and here they
 say their contributions so their contributions are fourfold. First of all
 the proposed method is the first attempt to build standalone attention models
 with larger large or a global receptive field and we'll see what that means. We
 propose position sensitive attention layer that makes better use of
 positional information without adding much computational cost. We show that
 axial attention works well not only as a standalone model on image
 classification but also as a backbone on panoptic segmentation, instance
 segmentation and semantic segmentation. Maybe what I did before described before
 was instance or semantic segmentation and not panoptic segmentation. Excuse me
 if that's the case. As you can see it can be used for various various image tasks.
 Lastly our axial deep lab improved significantly over bottom-up state-of-the-art
 on Cocoa achieving comparable performance of two stage methods. We also surpassed
 previous state-of-the-art methods on mapillary, vistas and city scapes. So these
 are various tasks as I said and also what they don't mention here is that they
 perform fairly well on ImageNet. In fact in the abstract they formulate this as
 in particular our model outperforms all existing standalone self-attention
 models on ImageNet. That's a way to phrase it. You just exclude
 all of the other models until you're the best. Outperforms all existing
 standalone self-attention models on ImageNet. That's good.
 There's something to be said of comparing apples to apples but you can
 also go overboard if you want to make your work look as good as
 possible. Of course everyone does that and there's no particular
 shame in it. We're going to build up our model right here and the
 basic element of this model is going to be this self-attention mechanism.
 Quickly because I know you all know what it is but very quickly you want to
 perform this action right here over a region right here. There is always a
 query and now the subscripts here are going to be important in this paper.
 The query is at a given position, position O and you can see that's the O
 right here. I'm going to call it the output. I guess that's what they
 said as well. So the output position. You want to go over all of the input
 positions and you want to aggregate data from all of the input positions.
 That's right here. How do you aggregate data? By this softmax operator right
 here. You can see the key also has a P right here and the softmax is over the
 axis of P. In particular case of the images what does that mean? If you have
 an image right here it's made into pixels. You have pixels. Now a
 transformer or generally these attention models, what you can imagine is
 they always transform a data point into a data point of the same dimensions.
 This doesn't have to be actually and I think one of the developments that is
 going to come in coming years or months or weeks, maybe someone's already doing
 it, is in fact to play more with this arbitrary constraint that we're
 imposing on ourselves because it's not really clear that this is the best thing.
 But for now an attention layer is always transforming a data point, here a 4x4
 image, into a data point of the same size. Also a 4x4 image right here.
 Now this is, as I said, this is quite simplified but it is true in NLP where we
 always transform our whatever 512 sequence, a token sequence into a 512
 token sequence and it is true here. Now the output is going to be here on the
 right and the question always is, okay so I'll go over these pixels
 right here and for every pixel, let's say for this pixel, I'm going to ask what
 data goes there? What's the output of the layer at that particular pixel? And the
 output of the layer is going to be somehow dependent on the input right
 here. Now if you know classic convolutional models, what the classic
 convolutional model says, the output of this is going to be dependent on this
 region right here, if it's like a 3x3 filter. So you have this
 convolutional filter and that means that blue dot on the right is going to pay
 attention to its own location in the input plus everything around it.
 And then every single data point here is going to do that. So for example this
 green data point is going to pay attention to this region right here. Now
 there's a border so there's maybe some padding but the question is always where
 does the information come from and how is it aggregated? In a convolution
 layer, what happens in a convolution layer? In a convolution layer you simply
 have your filter and the filter has numbers in it like
 three and five and eight and so on. And what you're going to do is you're going
 to take this region right here, this blue region of the lower layer and that's
 also filled with numbers like seven, what's a good number? Zero.
 Zero is a nice number and you're going to multiply those and then you're
 going to sum them up and then you're going to put that on where the blue dot
 is. So where does the information come from in the convolution? From around
 the location, from around the output location but in the input. So you go
 to the input at the same location as where you want the output to be, you take
 the neighborhood and there is a fixed scheme of aggregating the
 neighborhood. And then you multiply and you sum across it. In
 contrast to this, in a fully attentional model where does the information come
 from? Let's again look at the blue dot and let's consider it fully
 attentional. Where does the information come from? Everywhere,
 anywhere, anywhere at all. The information comes from everywhere. Now
 how do I know how to aggregate the information? So it's no longer in a
 neighborhood. How do I know how to aggregate the information? That's also
 different. So two things are different. Now in a convolution I would have
 another four by four grid here that's pre-specified but in the attention model
 this here is basically all filled with question marks. Question mark, question
 mark. What number goes here? In the end I also do this multiply and I
 sum it up and I put it right here. But how do these numbers come to be?
 Well these numbers also come, these are dynamically computed also from the
 input. It's a bit special but this is how attention works. So every pixel
 gets to decide where information comes from and how it is aggregated.
 Basically it comes from anywhere and how it is aggregated is dynamic depending
 on the pixel. If you still don't understand it maybe pay out to watch a
 video on attention itself. I happen to have made one but you can watch any one
 when you understand that you will understand the extension here to the
 image. It's the exact same thing as with the sequence except the pixels are
 basically one long sequence in the image. So this would be a fully
 attentional model down here. Now what's the problem here? The problem is that
 pictures are pretty large. So even something like MNIST which is like
 28 by 28 is like 700 pixels plus. I don't remember exactly but it's like about
 700 pixels. And our big transformers now, so BERT, a very famous transformer, takes
 inputs that are like 512 in length. You already need pretty decent hardware
 to run this. And the requirements on memory and compute scale quadratically
 with the input length. So already with MNIST you're in pretty pretty shady
 territory. If you go up to something like ImageNet which is like 225 by 225
 that's bad. That's not good. So you have to come up with something else.
 So people have been playing around, the reason why I introduced it this way, is
 people have been playing around a bit with sort of coming up with an
 intermediate with a compromise between the two. So the compromise that this
 paper here focuses on is going to be a compromise where we, you
 remember when I said where does the information for a given pixel come from?
 And we said okay it can come from anywhere in the attention framework and
 that's good because that allows us to make super long-range connections. So
 any pixel can aggregate information from any other pixel and not even in a fixed
 way but in a dynamic way. So depending on the pixel value itself and the other
 values it can decide how it wants to aggregate information. That turns out to
 be expensive, right? Every pixel together with every pixel, well that's
 quadratic. Okay so what do we do? We make a third method that's going to be a
 compromise and the compromise is going to be the following. The compromise is
 going to be alright we still do the dynamic aggregation which means that we
 still do the attention thing. However we're going to restrict back to
 this neighborhood region of the convolution. So in this model where does
 information for the blue dot come from? It again comes from this neighborhood
 right here and this number, the size here is going to be called m. So it still
 comes from that m by m neighborhood so a pixel can only aggregate information
 from its neighbors but contrary to a convolution how it aggregates the
 information like this what in convolution would be a kernel. The kernel
 is made dynamically by the attention module and it's made dynamically on a
 case-by-case basis. So we restrict it to a neighborhood, multiply, sum it up
 and then put it into the output and we do that for every pixel. Now it resembles
 much more a convolution, simply a convolution with this dynamic
 matrix right here. And that's the starting point for this paper. So this
 paper does two things to this. It says okay we can augment this by so-called
 positional embeddings. A positional embedding you might know from the
 sequence transformers. So if I have a sequence my cat is tall, I don't even
 know what that means for a cat. But okay what in a positional encoding so if you
 use a transformer and you transform this as we said into a sequence of equal
 length and then transformers basically information routing the transformer
 simply sees the lower layer sequence as a set not as a sequence. It has no
 notion of what's neighboring to what, what comes from where. So it pays to tell
 the transformer by the way this is word one, this is word two, this is word three,
 this is word four. There are various ways to do it. Transformers usually have
 fairly complicated kind of sine wave based positional encodings that bring
 many advantages with them. In this case they say well it might pay pay off to
 learn where actually these things are in this neighborhood. So they experiment
 with relative positional encoding which means they annotate this
 neighborhood with something like look here in the middle it's a 0 0 here is
 like 0 1 here it's 0 negative 1 negative 1 0 and so on. So they annotate it with
 these positional encodings. Now this is this would be the easy way what they
 actually do is they simply they give the model a matrix like this and they learn
 that matrix by heart let's say. So the positional encodings are relative
 positional encodings and they are learned okay so you can do that you can
 learn positional encoding so if you don't want to do the one two three four
 right here you simply say well here is a vector here is a vector here is a vector
 and here is also a vector. Now model you're already learning like all the
 weights to make this thing here happen and you're already learning your output
 weights up here right using back propagation. Why don't you learn yourself
 what you would like for position one like what kind of information you would
 like to be to have there using back propagation right so the model you
 provide them up you always provide the same vector so this is the same vector
 for position one and you have a different vector for position two and you
 have a different vector for position three right so but across all of the
 data points these vectors are going to be the same so the vector one is always
 going to be that same vector for all of the data points so the model somehow
 must learn independent of the data point what it means to be in position one so
 the model must learn how it wants to fill that vector that's called a learned
 positional embeddings we've seen this in many models so far it usually works
 pretty well and I guess here is it works especially well if you have these
 relative positional encodings and so this thing here is not going to be an
 actual matrix filled with these numbers it's going to be a learned matrix a
 trainable matrix that is filled that the network is allowed to fill with numbers
 right like three five eight and you might be you might notice that we've
 seen this before right so ultimately the information in this blue thing right
 here is going to depend on this dynamically created aggregating of
 information through the neighborhood and this statically learned aggregation of
 information throughout the neighborhood which is a con which is sort of a
 convolution right because in the convolution you've already seen here
 this is a statically learned map of how to aggregate information from the
 neighborhood of a pixel so I think even though there are slight differences
 they for example say this these are the same across attention heads and so on
 however I suspect that you can think of these learned positional embeddings to
 be to be kind of like what you learn in a convolution not exactly so no I I
 think I made a mistake and we'll see it in the formula we'll see it in the
 formula yeah okay so here they introduce these positional embeddings okay so you
 see that we previously we had the softmax previously we had this and this
 okay so this is the lower layer this is the information that comes into the
 layer and now it's it's transformed into values by a linear matrix but
 essentially this is the lower layer and for each of the output locations you
 want to know how should I aggregate information from that lower layer and
 you do this by this thing here this thing here is this dynamically
 constructed attention matrix using also the softmax okay so how should you
 aggregate information this comes from this query at the output position and
 the keys at the input position and now you add to that this math this thing
 right here which is again an inner product between the career query and the
 positional encodings okay so the positional encodings are going to be
 learned and hard coded but they still are modified by the queries so the
 query can still pay attention the difference is the keys depend on the
 input while the positional encoding does not depend on the input so the queries
 can decide I want to gather information from this and this and this type of
 information so that would be the key or it can decide I would like very much to
 look at pixels that are somehow on the bottom right of the pixel that I am now
 that would be the positional encodings and that's that's the mistake I made
 when I said it's equivalent to a convolution it is not because the query
 can still it's still modulated by that query vector of how to aggregate
 information otherwise you would have this to be a standalone multiplied by
 the input right here but it sort of pays off to think of it like what you do in
 the convolution so in the convolution you learn how to aggregate information
 basically based on position relative position to the position that you want
 to output and here you do a similar thing you learn static position embeddings
 that you then can attend to with your queries alright so these are the
 position embeddings and they make use of those position embeddings in fact they
 attend them to the following in this work we enable the output to retrieve
 relative positions beside the content based on query key affinities formally
 so the problem up here is that okay you have these position embeddings and here
 are the outputs but if you do this in multiple layers right if you do let's
 let's go with 1d sequences if you do this in multiple layers and here you
 annotate the position let's just go one two three four and okay this layer can
 make use of that right we gather stuff from here but then when this layer when
 this layer gathers information from here the where the information comes from in
 the layer below is some is how somehow getting lost right so it cannot kind of
 pull through this information to here or at least it's very complicated this
 model extends this positional embeddings in order to pull through that
 information so as you can see there are two new things right here the biggest
 important new thing is that right here we don't so here is how we aggregate
 information okay and here is the information that we aggregate over now
 you can see previously this was just this value vector and now it is extended
 to the position to positional embeddings learned positional embeddings okay so
 the this with this you're able to route the positional embeddings to the output
 and also here you can see the attention gets fairly complex so you have query
 key attention which is classic attention the queries can attend to positional
 encodings but also the keys can attend to positional encodings so not only can
 not only can the the node on top say I would like to attend to position three
 position three can also say well together with me positions two and four
 are are fairly important I guess that's what that's what that is maybe I'm
 mistaken here but you can see right here there is an interaction between the
 keys and the positional encoding right here now these positional encodings they
 are different for the queries keys and values but ultimately we don't that
 doesn't make too much of a difference so here is a contrast between what a
 traditional attention layer would do and what they would do so a traditional
 attention layer gets the input X and transforms it by means of these linear
 transformations right here into the queries these are the queries it's
 called Q into the keys and into the values okay then it does a matrix
 multiplication with the keys and the queries and puts that through a softmax
 so this here is going to be our attention matrix this is the attention
 matrix and the attention matrix is multiplied here by the values and that
 determines our output okay again the attention matrix defines how we
 aggregate information and the values is what information do we aggregate you
 know for the output in contrast when we introduce these positional encodings you
 can see right here again we have query key and value now it gets a little bit
 more more more complex right here namely we do this query key multiplication
 right here but we also multiply the query by these positional embeddings for
 Q we also multiply the keys by the positional embeddings for K and all of
 this together so this is a big plus right here all of this together is
 routed through the softmax okay and now the diagram is a little bit complicated
 now you can see the softmax aggregates information from here and from this
 learned position embeddings I would rather have they would just use it like
 they did in the formula do V plus R and say that's going to be the information
 that we are aggregating and the softmax here the output of the softmax is going
 to be how we aggregate information this is the attention all right I hope that's
 sort of clear you introduce these positional embeddings for queries keys
 and values and that allows the model to have a sense of where the information is
 coming from basically what positions which if you drop the convolutions so
 the convolution had this intrinsically because in your convolutional kernel
 right can I I'm dumb if in your convolutional kernel the number right
 here if there was a seven right here that meant that wherever you are
 whatever is on the bottom right is seven important okay so that's that was the
 convolution have this intrinsically here if you just do attention the we as
 humans we see it in a in this kind of grid form but the machine doesn't the
 machine simply sees a set of pixels it simply sees you can this is to the
 attention mechanism this is exactly the same as a long list of pixels or a
 discontinued set it doesn't matter to the machine so it's like the problems a
 feed-forward network has so we need to annotate it we have to give it
 positional information and learned positional information seems to work
 very well right here though you could think of static positional information
 okay this is the first thing the positional embeddings that now help the
 attention mechanism see where the information is coming from that's really
 important in pictures so we add that the second thing they do is this so-called
 axial attention now axial attention is sort of a let's say trick in order to
 reduce the load on a the load on an attention mechanism so what does it mean
 we've already we've already seen in sequences right if I have a sequence a
 sequence layer that's going to be n squared connections between the two now
 there are various ways to restrict that so instead of having all of these
 connections let's say from one node we've already seen wait if we just
 restrict it to let's say only this thing right here only this stuff that can be
 that is lower right that is lower in complexity and this in this case it
 would be just a neighborhood so that's what we've done that's this this M thing
 right here however we can also do it in different ways since this is a set
 anyway we can simply say maybe we should just always skip one we could like do
 attention like this and that would be just fine too right that would also
 leave away some of the information but you gain in computational efficiency
 there are various trade-offs now in a picture you have the same options right
 so you can do the neighborhood thing as we did or you can say where should the
 green pixel pay attention to axial attention says the green pixel should
 pay attention to only the row where it is in okay that's it should ignore the
 rest of the input it should only pay attention to that row where it is in and
 then in the next layer we'll flip it then the green pixel the same green
 pixel will pay attention to only the column it is in okay so that's that's
 called axial attention but don't think like don't don't there is nothing special
 about this being an axis or whatnot you could also define and it would not be
 called axial attention but you could define it makes the same sense to say
 well that green pixel it just depends on this diagonal right here just in the in
 this layer it just does this diagonal and then in the next layer it does like
 the anti diagonal you can say I just choose five random pixels in this layer
 and five random pixels in the next layer and that would work as well we've
 already seen this in this paper called big bird right the big big big bird but
 big bird so big bird explicitly used random connections in the attention
 mechanism and their argument was well if we use different random connections in
 each layer then information can travel pretty fast through the network so what's
 the problem with these neighborhoods right here what's the problem with
 neighborhood attention like this the problem is that you break the long-range
 dependencies so let's see what happens if information needs to go from this
 pixel to this pixel or this node to this node but if information needs to travel
 from this note to this note in a classic attention mechanism everything's
 connected to everything so that node in the next layer can simply aggregate
 information from here well that's not possible if you do this kind of
 neighborhood attention as we've done here if I do neighborhood attention then
 at most right because the neighborhood is three long at most this node right
 here can aggregate information from this node and then again it's three long in
 the next step so now this node can aggregate information from this node okay
 because the in the neighborhood is three long and you can only attend to within
 your neighborhood this means that if I want to send information to something
 that's really far away I need to I need to go many many layers right I need to
 go layer layer layer layer and this has been well known this has already been a
 like a problem this has already been a property of convolutional neural
 networks so convolutions specifically traded off the fully connectedness of
 fully connected layers to local connections convolutions but that means
 that you have to go very deep in order to make long-range connections you can't
 just make them in one step the same problem right here that is paper Big
 Bird argued that if you have random connections instead of neighborhood
 connections just the property of random graphs mean that you you are pretty fast
 in sending information around so because in a random graph of size n you on
 average all two nodes are connected by path lengths of log n this is much
 faster because in this neighborhood thing two nodes are connected in a path
 length of order of n right you can you can pretty easily see that if I make the
 sequence longer I need that many more steps in order to send it around in fact
 it's like something like n divided by m this neighborhood size in a random graph
 it's log n and in this axial attention that's why I introduced it it's 2 okay
 every every two nodes are connected by two steps if if node if this node right
 here needs to send information to this node right here in a classic attention
 mechanism you could do some one step because every pixel attends to every
 other pixel however right now we have to we have to see so this node attends in
 this layer sorry I have to think so how do we send information between the two
 we select this node right here in the first layer this node pays attention to
 this row okay which includes the red dot so the red dot can send information to
 the X in this layer in the next layer we select this node right here which is our
 target node where the information should go to it pays attention to all of this
 column which includes that X that before right this this X right here where we
 send information to so it takes two layers two steps to send information
 from any node to any other node well that's pretty good so this axial
 attention if you stack them on top of each other you sacrifice a little bit of
 of being able to send information from anywhere to anywhere for the pleasure of
 not having this quadratic attention anymore as you can see your attention
 mechanism is now as long or as big as your column or is wide or your row is
 high again this isn't this isn't specific to rows or columns you could do
 this as I said with these kind of diagonals you could do it with any other
 sort of sub pattern where you can sort of guarantee that the overlap between
 the layers is enough so you can send information around pretty efficiently and
 they use this right here so this axial attention you can see the formula is
 exactly the same the only change from before is this part right here you can
 see that the neighborhood that they aggregate over is no longer M by M it is
 now 1 by M so we've seen them going from if this is the the full input image and
 you want to you want to see where to attend what this paper does is it says a
 classic sorry a convolutional neural network would be attending to some sub
 part right this is convolution an attention mechanism pure attention
 would attend to everything but this is attention then what we are doing sorry
 that was a mistake what other people were doing were reverting back this
 attention to a sub part this kind of neighborhood attention okay but that was
 still you know you still have M squared you still have O of M squared because of
 the attention mechanism now what we are doing is we are going even lower we're
 actually going 1 by M okay this this is with with axial attention so in general
 it's 1 by M and then in the next layer we can go 1 by M in this direction and
 have that property and because it's so cheap now right because it's now O of M
 to compute this we might as well make M as long as the row itself okay so their
 last step is going to be to say okay we have 1 by M right here and that's going
 to be the row itself now you can see right here that they say axial attention
 reduces the complexity to HWM this enables global receptive field which is
 achieved by setting the span M directly to the whole input features optionally
 one could also use a fixed M value in order to reduce memory footprint on huge
 feature apps which is something that they're going to do later on ImageNet I
 believe so when they have big inputs or big outputs they actually do use a
 smaller M what you can see right here is that I wasn't really that wasn't really
 correct of me to say that it's now O of M because you you still have the entire
 query space so you multiply query by by keys now even if you make the keys to be
 1 by M yes you reduce definitely you reduce this from height times width to
 times height times width to this but then you can see this thing right here if
 you take it and let's say we have this kind of row pattern and we replace M by
 the width then we have width squared so again the square appears however it's
 smaller than the original attention the original attention was H squared W
 squared right because HW is the image and you need that squared in order to
 do the attention mechanism now we've basically reduced one of the factors it
 is still an attention mechanism so there's still attention going but we've
 basically transformed the the image we've reduced it to one column now the
 one column is still attention so this is still attention like here so this now
 reduces to the attention that you see in a in a single sequence okay if you see
 the image as a long stretch of pixels what this does is basically it's up it
 simply subdivides that into neighborhoods so we're back to
 neighborhoods basically but we shift the neighborhoods from layer to layer so in
 the next layer the neighborhoods are going to be just alternating right the
 neighborhoods is going to be this is one neighborhood connected to this
 neighborhood connected to this neighborhood I hope this makes sense so
 it's going to be it's basically a mix between if you if you if you were to do
 this in convolution you could do one layer where it's neighborhood convolution
 and then one layer where it's like convolution with holes in it I think
 they're called atras convolutions or something like this with like giant
 holes in it that are exact is exactly the anti pattern of the neighborhood
 convolution from before that's what this is so you see their axial attention
 block right here their axial attention block replaces the ResNet block so if
 you know ResNet I've done a paper on ResNet ResNet basically takes the input
 pipes it through straight and adds to it whatever comes out of this
 operation okay that's a residual block now usually this thing here would be
 convolutions and convolutions and they are now replaced by these multi head
 axial attention you can see there is a multi head attention in the height and
 there is a multi head attention in the width and that gives us the property
 that every node can send around information to every other node in two
 steps I don't like the fact that there is only two because what this I guess
 this gives a significant bias to one or the other direction depending on the
 order that you do them in if if I had done this I maybe would have used three
 of them because it depends on how you want to aggregate information right like
 here you train the network specifically to aggregate information first in this
 direction and then in this direction which might work and it'll give you that
 sending around information anywhere so maybe they've actually tried and it just
 performed the same so I just might have a dumb suggestion right here in any case
 they simply replace in we've come a long way right we've gone to like
 neighborhoods and blah blah blah blah ultimately take a ResNet place the
 convolutions with the height axis attention and the width axis attention
 and we're good and then we come to results so that's it you have these
 positional embeddings you have the axial attention and it turns out that on
 ImageNet they perform fairly fairly well so you can see that models like a
 ResNet 50 model will get a 76.9 on ImageNet which is not state-of-the-art
 but it's also not it's not bad right the ResNet 50 is pretty good model you can
 see the full axial attention right here achieves a 78.1 also not state-of-the-art
 but still pretty good and as they say it's the best fully attentional model on
 ImageNet or standalone attention model on ImageNet so where this model really
 shines is where you really have to make long-range connections between pixels
 and that's these kind of segmentation tasks and I want to skip the tables
 right here they're best and everything and go to the appendix where they have
 some examples of this so here you can see specifically this is the original
 image you have a ground truth and you have the differences between their model
 this axial deep lab and the panoptic deep lab that is a baseline for them and
 you can see that the the failure cases here are are pretty you know show how
 show how the axial deep lab is better I don't know if they are cherry-picked or
 not but at least you can see that at some point so it handles occlusions
 better it handles instances better so here you see that the ground truth
 separates the person from the tie and the axial attention is able to do this
 but the the baseline is not able to do this correctly because it labels part of
 that white shirt also as and you can see why there's kind of a delimiter line here
 here here here but if you have long-range dependencies right if you
 have long-range dependencies in the model the model will recognize wait wait
 that's that must be the same thing as this thing here and this thing here and
 this thing here so that must be the same object it's simply that the shirt was
 occluded by the tie and goes beneath it and now appears again it's not a
 different it's not part of the tie and it's not part of the of a different
 object it's actually part of the shirt so the long-range attention you can see
 at these examples sometimes here okay this might not be an instance of super
 duper long-range dependencies this is simply where the model performs better
 so you can see here the ground truth has that surfboard segmented and the
 baseline does not that this can also just be you know there are a lot of
 tricks to make this work of course and you throw a lot of compute at it and
 sometimes you just get better numbers or part of the better numbers because of
 the additional compute right here what do we have so you can see occlusions it
 appears to handle occlusions in a better way and this might be due to this axial
 attention it might be due to the positional embeddings but you can see
 that the ground truth here has the laptop between the person's hands
 segmented the baseline cannot do that but the axial attention does do that and
 I don't know what this is honestly this is you can you can see though the axial
 attention also misses the fact that it should segment this in the background
 and if this occlusion handling you can see best in this example where the
 person in the back reappears on both sides of that person so you can see that
 the axial attention manages to segment that where that is just a mutant person
 right here the ground truth is equally shaky I think there is might be some
 ambiguity of how you can segment these images obviously but you can see the
 fact that there are long-range dependencies probably helped with this
 saying that wait in this image there's this white stuff right here and there's
 this white stuff right here and connecting these two regions with
 attention probably helped in segmenting these to be the same object even though
 you can see there is a break in the object so there is a break no at no
 point is the object on the left touching or the segment on the left touching the
 segment on the right and still the model manages to put those into the same label
 category there is the last last thing where they they want to research what
 their heads learn and usually you can do this right you can kind of visualize
 what the attention heads learn so in this case right here in the column heads
 the way you have to read this is that this particular head right here aggregates
 information from its column so everywhere where it lights up it there's
 a lot of information being routed you can see specifically in this here the
 heads of the people or the heads of the persons in the picture light up fairly
 well so for example this head right here is probably aggregating information a
 lot from this position right here and this head here is aggregating information
 from this position so you can deduce that that particular attention head
 probably deals with people's faces whereas that particular attention head
 probably deals you can see the attention is mostly on the grass right here and
 you can see the same with the for the row heads now their description here is
 that we notice that column head one corresponds to human heads while column
 head four course correlates with the field only which you know you can
 interpret it as this this seemed pretty clear but then they say something like
 row head six focuses on relatively large relatively local regions where column
 head five pools all over the image so row head six which is this thing right
 here you can see that okay it maybe focuses on small regions though you can
 see okay what like here you can get it that's a person but in other places I
 don't know where column head five pools over the whole image and this I don't
 know maybe they just needed something more to say because they put these
 pictures here they were like okay the the column heads are really nice because
 we couldn't like these this one's really nice because it you know just pays
 attention to the people and this one looks really nice because it pays
 attention to the field and but we can't really put the column head attention
 without putting the row head attention but then none of the row heads really
 are like super distinctive on a particular thing in the image so we need
 to come up with something that we can say and then you look at this one this
 is there's not a lot of attention so we need to contrast this with something
 then you would think that they contrast it with another row head but then there's
 no row head that does this whole image so there's like a column at five yeah
 I'm not sure if there's there's a bit of there is a bit of tactical writing going
 on here I suspect I mean still you know it's doing something cool but yeah
 there's there's definitely an element of sales in when you do when you do
 where I research papers and just not to this data but just props to the lines in
 front of the histograms makes it so much easier to read how big the stupid bars
 are why does everyone put the lines behind the histogram I probably do that
 myself and now I'm just I'm realizing how much easier that is alright there is
 a big big big experimental section right here and there's a big appendix where
 you can read up all of the different numbers comparisons ablations whatnot
 ultimately I just wanted to go over the method basically putting this into
 context with other things like putting this into context with stuff like Big
 Bird axial attention other positional encodings how it how it relates to
 convolutions how it relates to feed-forward networks and what convolutions
 did to feed-forward networks and so on I hope you have at least a little bit
 gained an understanding of what's going on here and with that said I see you
 next time bye bye