so if you use a transformer, and you transform this, as we said, into a sequence of equal length,


Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation (Paper Explained)

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 encoding


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 zero zero, here is like zero one, here, it's zero, negative one,


negative one, zero, 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 1234,


right here, you simply say, well, here's a vector, here's a vector, here's a vector,


and here's 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,


358. 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 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 method, this thing right here, which is again an inner product between


the key rear query and the positional encodings. And this is the input position. And this is the


input position. 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.


And the query can decide I want to look at the input position, which is 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


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. All right, 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 one D sequences. If you do


this in multiple layers, and here you annotate the position, let's just go 1234. 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 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, it 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 the input x gets


the input x, and transforms it by means of these linear transformations right here into the


queries. These are the queries, let's call them 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


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


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 learn positional 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. Alright, 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 the convolution have this intrinsically.


Here, if you just do attention, that 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 of feed forward network has, so we need


to annotate it, we have to give it positional information. And learn 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 the so called axial attention. Now, axial attention is sort of a,


let's say a 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 will 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 an it would not


be called axial attention, but you could define it makes the same sense to say, well, that green


pixel 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 node to this node 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 is three long,


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. Now this 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 two, 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, we 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.


That's pretty good. So this axial attention, if you stack them on top of each other, you sacrifice


a little bit 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 one by m. So we've seen them going from if this is the the full input image,


and 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, right? 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


one by m. Okay, this is with axial attention. So in general, it's one by m. And then in the


next layer, we can go one 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 one 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 maps, 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 the


m value is now O of 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 still have the entire query space. So you multiply


query by keys. Now, even if you make the keys to be one by m, yes, you reduce, definitely,


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 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 atris 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's


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 will 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, blah. Ultimately, take a resonant 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 the resonant 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 resonant 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 at 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 failure cases here are pretty, you know,


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,


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, exactly 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,


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,


though 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 want to research what


their heads learn. And usually you can do this, right, you can kind of visualize what the attention


has learned. 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 noticed


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 may be 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 column heads are really nice,


because we couldn't like, 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're like, ah, 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


they contrast it with another row head. But then there's no row head that does this whole image. So


there's like, ah, column at five. Yeah, I'm not sure if there's, there is 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 write 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 realizing how much easier that is. All right,


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


A neurally plausible model learns successor representations in partially observable environments

All right. Hi there. Today we're looking at a Neurally Plausible Model, Learn Successor


Representations in Partially Observable Environments by Esther Vertes and Manish Sani. This paper


is a paper on a topic that has been interesting me for a while and that's successor representations.


So we'll dive into all of this. The title is fairly lengthy and complicated, but ultimately,


we're dealing with a setting of reinforcement learning. So if you know something about reinforcement


learning in reinforcement learning, usually you have a agent, which, yeah, let's just


say this is you. And there is an environment where it is a big black box that you don't


know anything about. This is environment. And what the environment gives you is what's


called an observation. So an observation could be anything. But in this case, let's just


assume it's a, you get a little picture, right, of what's in front of you. So in front of


you might be a tree and in front of you might be a house. And then you can perform an action,


a, and this action in this case might be to enter the house, right? And then the environment


in the next step, it gives you back a new picture and says, oh, you're now inside the


house. So here is a door that leads you to this room and the door that leads to that


room. And there's a little table in front of you. Right? So this, this, it's this, just


this cycle of action observation. And with that, you're trying to collect some reward


over time. Now, there are different ways of achieving this, this, this reward over time.


So basically, the reward is going to be, for example, you could get a reward for finding


the kitchen, or for going as into as many rooms as possible, or, you know, anything


like this. So the, the other objective is to learn a what's called a policy. So which


actions to take, so action one, action two, action three, given the observations that


maximizes your rewards. So there's mainly two ways to go about this. There's the model


free and the model based reinforcement learning approach. Let's split them. So in the model


free, model free approach, what you're trying to do is you're trying to simply learn a policy


and we call this here pi of s and s is your state and state, you can think of it as the


observation. So in this policy, we'll simply output an action. And this, this is the kind


of the simple setup of model free reinforcement learning. The important thing here is you're


trying to learn this, usually there's parameters theta of this policy pi, this could be a neural


network, and the theta are then the weights of the neural network. So you're trying to


learn the neural network such that if you give it a state, it just outputs the action.


So you have this neural network with your state, you input the state into layer, layer,


layer, layer, layer, and then it outputs one of maybe three actions. Go north, go south,


go west, maybe go east, right, this could be four, four actions, you're just trying


to train the neural network using backprop, and the reward signal through what's called


the reinforced trick or variants thereof. This is model free reinforcement learning.


It's very, it's very easy to implement, let's say, and it's very applicable. And it will


simply give you a mapping, you know, you have to know nothing about how the world works,


it will simply tell you at the end, if you're in this state, do that action and the reward


will be high. In contrast, there is the other world, this is the model based reinforcement


learning. So in model based reinforcement learning, what you have is a model of the


world. And the model of the world, let's is best described, for example, if you play chess,


right? So if you play chess, and this is a let's do a simplified chess board here, four


by four, and you have a pawn right here, right? You have a pawn. And, you know, if I do the


action of moving the pawn forward, I know the pawn will then be in this square right


here, right in the next time step, I know that because I have a model of the world,


I know how the world works. And I can predict basically the results of my actions. So if


you have a model based reinforcement learning setup, if you know how the world works, you


can do something like a search. So given you're here in the state, right, you know, if I do


action one, I go to this state, if I do action two, I go to that state. And if I do action


three, I go to this other state. And from each of the states, you can then say, Ah,


but again, I have three actions. And I can, you know, go into these three states go into


these, maybe here, too. And maybe here, I can go into these, actually, let's do three


as well. Right? And then the question more becomes, can you find a path through this


thing, such that at the end, you are in the state that you want to end up, right? So for


example, here is outside. And then here, you can go to the tree to the house, or to the


field. And in the house, you can go to the bedroom, the bathroom, or the kitchen. And


you know all of this, you have a model. So you can actually kind of compute what would


happen if I do something and then search for the best path. Whereas in the model free reinforcement


learning approach, what you'd simply do is you'd say here is a state and the state is


for example, I am in the house. And now give me the action that would maximize my future


reward. And you're trying to learn this directly. So it's a very different style of reinforcement


learning. Basically, one is a one is a pure machine learning approach. And the other one


is a search problem. Now you can of course, mix and match the two like, for example, people


in AlphaGo have done they have a model based reinforcement learning that also has kind


of learning machine learning elements. But in between now we have the successor features.


So the successor representations, they are, if you will, they are somewhere in between


the two. So they kind of trade off the advantages of model free, where you only have to learn


a function right from state to something with the advantages of model based the fact that


you actually have a bit of an idea of how the world works and can adjust quickly to


let's say different reward structures or things like this. So what do successor representations


do successor representations basically learn how states are connected. And this is a classic


successor representation. So the successor representation m here, of policy pi, the policy,


remember is what tells you which action you should take in a given state. You you define


it as a connection between state i and state j. And m of si, sj means given that I am in


si. So this could be the kitchen and your goal is to find the bedroom. And if this is


the kitchen, given that I am in state si, what's the probability that in the future


at some point I will transition to si? Right? Given that I'm in the kitchen, what's the


probability that I'll end up in the bedroom at some point in the future. And this is formally


expressed this is the expectation over your policy. And it's the it's the indicator function


that the future state, sorry, this is the future state t plus k, you see k goes from


zero to infinity. So for all of the future, and st is the one you're in now. So for any


future state, this is equal to sj. Now, of course, this makes no sense unless you kind


of discount, have a discount factor here. So if you're in state, if you're in the bedroom


further in the future, then this value would be lower. So that this value is high, if you


will transition from si to sj with high probability in the near future. And this is a successor


representation, right? It basically tells you, if you want to go from state si to state


sj, how likely is that in the near future, right? So if the if this number is high, you


know that that these two states are closely connected that you can expect to end up in


state sj somewhere down the line if you're in si now. And one more representation, if


you consider the vector m pi of si, given all of the sj's, some dot here, so this is


a vector, you can actually compare two states si. So if one is, if you plug in here, you


plug in the kitchen, and then also you plug in the, I don't know, the garage. If they


and you'll get out two vectors, right, you get two vectors, if those vectors are very


similar, then you know that if, if you're in the kitchen or in the garage, it doesn't


matter, you're going to end up you have a similar future trajectory, basically. However,


if those two vectors are far apart, you know that these two states are far apart with respect


to your policy. So this is pretty cool things you can do with successor representations.


And I hope this gives you kind of some insight. So another neat trick is that if you have


a value function, so and the value function, in this case, there's a simplified assumption,


but you don't actually need it. The simplified assumption is that the reward only depends


on the state you're in. Basically, it doesn't matter how you get to the state, like the


actions you perform. If you're in a given state, if you're in a given room in the house,


you'll get some reward, like, for example, if you find the bedroom, then you win. That's


a reward that would only be characterized by the state. If that's the case, you can


compute the value function of of the reinforcement learning problem simply by integrating over


the over the success representations. So for each state, you simply go over all of the


possible other states, and you ask how likely am I to go to that state? And what reward


will I have in that state? And that's your value function. So pretty, pretty simple.


You can actually learn the successor representations by TD learning by temporal difference learning,


which is a method that's applied for like throughout reinforcement learning, especially


in places like Q learning, and yeah, and also for learning value functions. So pretty neat


successor representations. This paper then goes from successor representations of individual


state to successor representations over continuous space. So right now we all have these states,


state kitchen, you go to the bedroom, you go to somewhere, right? And these states were


kind of discrete places. So there was a house and you have different, you know, different


rooms in the house and you can go between them. Now, we're dealing more with continuous


states. So you can generalize the successor representations to continue state by considering


not the states themselves, but features of the of the state and a feature. And this here


you have to kind of imagine as binary features. And the features, let me give like some really


dumb examples, but maybe it helps you like that one feature could be the smell. Does


it smell in the room? Like just binary? Does it smell or doesn't it smell? Right? And then


one feature could there be is there sunlight? And then one feature could be is it warm?


Right? Right? With, and these are all binary features. And so you could kind of you have


to build the features such that if the features are the same, then the states should be fairly,


you know, close in whatever sense. So, for example, if it smells, but there is no sunlight,


you're probably somewhere in the bathroom, like where exactly in x y coordinates you


are in the bathroom. It doesn't really matter to this as long as like the features are high.


And so if it smells and there is no sunlight, you're probably somewhere in the bathroom


and that makes all the states in the bathroom, all the coordinates close together. So this


is how you have to imagine these features. You can define your successor representations


exactly the same over these features, except that the representation is now not from state


i to state j, but from a state to a given feature. So that means if I am in state st


at the current time, what is the probability that in the near future, this feature will


be high, right? So if I am right now in the or close to the bathroom, let's say the the


probability that smell, oh, sorry, this should be a highlight, the probability that smell


is high in the future is very high, right? So this this number would be high. So exactly


the same except for these continuous features now. And you can do the same thing, including


defining the value function as a simple linear multiplication with these features. That is


an assumption under the assumption that the reward is a linear function of the features


of the states, which is the analogous assumption to saying that the reward only depends on


the state in the linear case, or somewhat of an analogous function, not entirely. All


right, so you can also learn this by temporal difference learning exactly the same. So this


is pretty cool. These these are the successor representations. And you can actually you


know, if you learn them, you have kind of a model of how the world works, not as much


a model as the model based reinforcement learning, where you know exactly how it works, right?


Here, you know exactly how the world works, you have this model. In model free, you don't


know how the world works at all, you simply know, oh, if I'm in this state and do this


action that that will turn out really well. But in the successor representation framework,


you have you have an idea of what states there are, we'll do the discrete case right now.


So this could be kitchen, this could be outdoor, this could be bedroom. And so you you have


an idea what states there are and so on, and how they connect to each other. Like you say


are from the kitchen, I can easily go to the bedroom, but I can not as well go to maybe


the bathroom. From outdoor, I can easily go to the kitchen, but I can't go to the bedroom


and so on. So you have kind of an idea of how all of these states connect to each other.


And that is the successor representation, you can already see how that helps learning


agent a lot, if you introduce the successor, if you have the successor representation.


And what this this paper deals with, in essence is it says, okay, the successor representations


are cool. But it has only so far been done in a case where you have full observability.


And the full observability is the case where you kind of know what state you're in, right,


you kind of know that, sorry, you are in the kitchen, you are outdoors, you are in the


bedroom, that is not known. But what if you don't? And in most problems, you don't? What


if you just have a picture like here, right? You just see a tree in the house, right? You


don't, you kind of have to infer that you are outdoor, right? And if you're here, you


just get this picture of a couple of doors and a table, and you have to infer that you


are now in the living room. Right? So, in essence, there is an additional layer of complexity.


Not only do you go from state from state to state to state, but you don't actually observe


the states, what you observe is from each state, you observe what are called observations,


right? So you only observe these, and you have to infer what the kind of have to guess


what the underlying states are, in order to know what you should do to get to the next


state, right? You only ever observe the observations. So this here is the actual thing, this is


the kitchen. And this here could be a picture of the kitchen, right? There's a counter,


there's a stove. Yeah. And so you get kind of what I what I mean. In their example, they


have, they simplify this to kind of a toy data setup, where you have this environment.


And this is one beautiful picture. I don't why, oh, well, just do you have one this setup,


and it is this box, basically, this box, and it has this wall, right? And then you have


an agent that is able to walk around in here, like with whatever policy, the policy determines


how it works. But then what you observe is not the actual position. But what you observe


is, for example, for this position, you observe a random point here. So they basically add


noise to each observe to each state. And if you're in this state, you will observe one


of these points in this circle, right? So your, your trajectory might look to you as


you observe it much more, much Michael, for example, from here to here, to here to here,


right, and you kind of have to guess what the underlying state is. And you see this


here. This, this blue thing is what the agent actually does. But the gray thing is what


it observes. And the observations are sometimes even outside of this, of this boundary. And


this, this orange thing is now the inferred thing. And that's what we actually want is


to go from the observed to these inferred and we want to that the inferred is as close


as possible to this true latent state. Right? So the way they do it is they introduce these


distributional distributed coding for the for the features of basic for the expectation


of the features. And basically, what they say is they say, we will build a framework


where we can, where we represent the features as, as as expectations over some distribution,


and the expectation will call mu. And mu is simply the kind of mean of this, of this feature


under this distribution. This is very general. So let's look at what, at how to plug this


in. So what they now have to do is they have to learn these two things, right? They have


to first of all, if I draw this picture again, these are the underlying states, and they


kind of transition into each other. So this is state one, state two, state three, and


with action one, action two, we transition from state to state. But also, there are these


observations, observation one, observation two, observation three. So the agent needs


to learn two different things. First of all, it needs to learn, given an observation, what


state am I probably in? Right? This is the first thing it needs to learn. And then the


second thing it needs to learn is given this state, and this action, what's the next state


that I will go to? Right? And this is, and this, of course, these things down here, they're


not observed. So these things down here, you can only do in distribution. So I'm going


to represent this with a P here. You can only kind of do this in distribution. And the way


they handle it is they always maintain the expected value of these things. And that's,


they do this in this wake sleep algorithm. All right, so this is me rerecording this


part, because I have done a terrible job at the first time. So we want to understand this


wake sleep algorithm to compute the things that we don't know. Let me draw this actually


again. So the way this algorithm does it is actually pretty cool. It has two phases, a


sleep phase and a wake phase. And it alternates between the two constantly. It's kind of like


expectation maximization. And ultimately, what you want to learn are two different sets


of parameters, w and t. Now, you whenever you learn t, you use w, the one that you've


already learned. And whenever you learn w, you use the t that you've already learned.


So it's kind of a bootstrapping each other up. The two functions you learn here are this


f w, and the t here. So t is just a matrix, and f of w is a function. The function has


weights, the weights w. So you see in the sleep phase, you update w. And in the wake


phase, you update t. Now, why is this called wake and sleep? It's because in the wake phase,


you're actually so called awake and you use real observations. So in the wake phase, and


I find it easier to start actually at the wake phase. In the wake phase, you collect


observations. So you let your agent go around its environment and collect a bunch of observations.


You don't know what the states are. But what you do is simply you collect these observations.


Now, it's not that important what the policy is here. So you basically follow some policy


and you collect these observations, right? And then what you what you say is, okay, I


have the function f of w. And remember, since we're in the wake phase, we're learning t.


So we assume we already have the w. In essence, in practice, we start out with a random one


and right and then kind of alternate between the two phases until both get really good.


But so we already have a w and we use it to update t. How do we do this? This is we need


to understand what this function f of w does. f of w takes this mu and the current observation


and produces a new mu. So what is a mu? This mu here, this mu here, as we saw above here,


the mu is the expectation over the features. And in essence, the mu is a guess. The mu


is your best guess of what the features of the state are. Or in the in the discrete case,


you could also say a guess of what the state is, right? So you don't know the state, right?


But you what you want to maintain is a distribution over state. So you want to kind of maintain


this distribution. But you can't, you know, calculate, you can't properly efficiently


calculate with an entire distribution unless you assume it's some sort of Gaussian or so.


But what you can do is you can simply take its mean mu, right? And and that's your best


guess for what the state is. The state could be anywhere, anywhere here, right, according


to this distribution. But you simply come up with mu, which is your best guess. So the


function, the function f of w, it takes in the best guess of where you were up until


the last step. And it also takes as an argument your current observation. And it gives you


the output of f is mu t, right is the best guess of where you are now. Yeah, it's pretty,


it's pretty straightforward if you think about it. So for every, for every observation, you


want to have kind of a guess of where your what your state is, and that's mu, right?


So what f does is it takes whatever observations you had, these observations gave rise to a


mu that guess where you are, you take this mu, and you take this observation. And from


that you derive the next guess of where you are. Yeah, you just say, I guess I was in


the kitchen before. Now I moved, I observed that I moved through some sort of door and


there is some sort of table. So given that I get I thought I was in the kitchen and that


I observed this thing now I'm probably in the living room. That's what f w does. So


you input the observations that you had, and you input your current observation, right,


to get the guess of where you're next. And these are real observations, right? And then


you simply update t, what does t do? t relates your current and your next guess. And that's


important, we already said that f takes your your kind of last guess, and gives you the


next guess. T does kind of the same thing. But T does it without having without relying


on an additional observation, T simply says, Well, if I am here, or if my guess is that


I am in the kitchen, then what's the probability that in the next step, I'll be in the living


room without observing anything, right? It's T is simply relating states to each other,


or relating guesses of states to each other. Right? So it's simply, it's simply saying,


well, under the current policy that I am, what is the kind of distribution of going


from one room to the next room? Right, so you learn in the wake phase, you learn the


T, the T simply represents how you move from state to state. So it's exactly, basically,


this function here, except that it's not from state to state. But it relates your guess


about your guess your mu from the state one to the mu of the state two. Right? And then


in the sleep phase, so in the sleep phase, you you now assume that you have a good estimate


of how the states relate to each other. And what you can then do is you can actually sample


trajectories. And this is why it's called sleeping. It's kind of like dreaming. So given


that you have a model T of how states transition to each other, or your your guesses about


states more precisely, you can now sample state trajectories. So you can dream up how


you would move in an environment, right. And the assumption here is that you know the process


that that if you have a state that gives you an observation, for example, in their experiments


is always the state is xy coordinates, and that's corrupted by Gaussian noise. There


is also ways to learn this transition, this is what's called the this is what's called


the observation process, right. But you assume you know it. So you can sample trajectories


of states and corresponding observations, right. Now, this is not the real world, but


this is using this T down here. You kind of know how or you kind of have some sort of


model can you learn a model of how you move about the world. So you sample these trajectories.


And from these trajectories, you can now learn the f of w function. So you see, since you


know what the state is, right, you can compute these features exactly. And then you can learn


this f of w function that gives you a takes in the guess of the last state and the current


observation and gives you the next the guess of the next state. And that you can then use


temporal difference learning. This is always here. Also with the T here, we have temporal


difference kind of temporal difference learning to learn the parameters w. So it's very kind


of convoluted. But ultimately, it's a simple process in the wake phase, you go into the


world and actually collect real observations, right. And you have a method of deriving from


these observations, deriving the guesses about the states. And so what you can do is you


can learn a transition between the states, right, if you have a good guess of what the


state are, given each observation, you can learn how to transition from one state to


the next state, except you don't do it in actual states, you do it in guesses about


states. Then once you have a model of how you move from one state to the next state,


you can go and dream up such state trajectories, right, you can dream state trajectories, and


therefore, also, you can dream how you would observe them. And given that you can learn


then a better function that relates your guess, your guess about a state given the observation


to the actual features of the state, since for this particular thing, you know what the


state is, right. So this, this is this two step process. Notice the cool thing, we've


never actually had to learn this mu explicitly, we never had to learn how to go from one state


observations to your guesses about states, because we can compute this recursively, right.


So you simply start out with mu zero, which is a guess about the initial state. And then


you go to mu one and mu two, and you never actually have to learn that function. All


right. All right. So that's how they, that's how they kind of learn these, these success


representations and experiments of this are fairly cool. Here is another diagram of how


that looks like you have a state, this gives you an observation. And from that you derive


a guess of what this state is. Right. So you can now look at what the agent learned, the


agent actually learns dynamics of this room means if you're here, you probably go somewhere,


right, there is no clear direction. But if you're close to the wall, your next states


are probably going to be inwards of this wall. Right. And yeah, I've already shown you this


picture. So they have a last cool experiment here, where what they do is they specify a


reward, and the reward is down here. And from each state, you want to know, which way do


I have to go to write get the reward. Now, if they give the agent the value of the latent


state, and the latent state here are just your x y coordinates. If they give this to


the agent, and they let it run, they let it learn the structure of the world, it will


correctly conclude a hollow, these are the kind of here is high value states, lower,


lower, lower, lower, lower value states, right, up until over here are the most low value


states, because you travel the longest to go to the reward. If you just give it the


observation, the noisy observation, it will actually assign high value to states here,


because of course, it, it doesn't infer the latent state, it simply takes the observation


as face value says, Well, I was here, and I reached here pretty quickly. So it must


be a good state. But in fact, it wasn't here, it was here, and the added noise would just


corrupt the observation. So you see, it learns kind of a wrong model of the world. Whereas


if you use this DDC, you see, sorry about that. If you use this DDC, you see you're


much closer to the true state of the world, like to the to the one on the left here, where


you so on the left here, you actually kind of cheat, right, you give it the actual state.


But on here, you give it the observation, but tell it, it's actually a noisy observation.


And you use what this paper proposes. And again, it will learn to assign a low value


to these states, because it needs to go all the way around, even though it has seen, supposedly


seen the agent go from here to here directly. But it, it kind of understands that it's just


a noisy observation. Alright, so this was this. From this paper, it's very, very cool


approach, I think, to reinforcement learning. And there's some more experiments where you


can see that this DDC actually helps. And I'm excited about successor representations


and how to incorporate them in reinforcement learning, because it seems a perfect kind


of middle ground between model based and model free RL. Alright, with that, thanks for listening


CURL: Contrastive Unsupervised Representations for Reinforcement Learning

Hi there. Today we're going to look at CURL, Contrastive Unsupervised Representations for


Reinforcement Learning by Aravind Srinivas, Michael Laskin and Pieter Abbeel. So this is


a general framework for unsupervised representation learning for RL. So let's untangle the title a


little bit. It is for reinforcement learning, which if you don't know what reinforcement learning is,


I've done a bunch of videos on RL framework. So it's for general reinforcement learning. That


means it can be paired with almost any RL algorithm out there. So we're not going to, you


know, dive into specific RL algorithms today. It is unsupervised, which means it doesn't need any


sort of labels. And it also doesn't need a reward signal for RL, which is pretty cool, because


usually the entire RL pipelines rely on some sort of a reward or auxiliary reward signal. Now there


is a training objective here, but it doesn't have to do with the RL reward. And then in the, it is


learning representations, which means it learns, it learns intermediate representations of the


input data that is useful. And in the end, it is contrastive. And that is the kind of secret sauce


in here. The training objective, it's what's called contrastive learning. And that's what we're going


to spend most of our time on today exploring what that means. Alright, so here's the general


framework. You can see it down here. Sorry about that. So you can see that reinforcement learning


is just a box, which is we don't care about the RL algorithm you use. That's just, you know,


what what comes at the end, what comes at the beginning, oh, here is the observation. So the


observation in an RL algorithm is kind of fundamental. Now, if someone explains RL to you,


reinforcement learning, usually what they'll say is there is some kind of actor, and there is some


kind of environment, right? And the environment will give you an observation, right, observation,


observation, oh, which is some sort of, let's say, here is an image, right. So in this in this RL


framework, specifically, the examples they give are of image based reinforcement learning. So


let's say the Atari game, where you have this little spaceship here, and there are meteorites


up here. And you need to shoot them. So there is a little shot here, right, you need to shoot those


meteorites. Right. So this is the observation Oh, and then as an age, as an actor, you have to come


up with some sort of action. And the actions here can be something like move to the left, move to


the right, press the button that you know, does the shooting. So you have to come up with an action


somehow given this observation, and then the environment will give you back a reward along with


the next observation, like the next frame of the game. And you're going to have to come up with


another action in response to that, and the environment is going to give you back another


reward, and the next observation, and so on. So what you want to do is you want to find a mapping


from observation to action, such that your reward is going to be as high as possible, right? This is


the fundamental problem of RL. And usually what people do is they take this act this mapping here


from observation to action, to be some sort of function, some sort of function that is parameterized


maybe. And nowadays, of course, it's often a neural network. But you're trying to learn, given the


input observation, what output action you need to do. And you can think of the same here. So you have


this input observation up here. And down here, after the reinforcement learning, the output is going


to be an action, right. And so this this function we talked about up here is usually implemented,


sorry, is usually implemented like this, you put the observation into the RL framework, and then the


RL framework learns this f of theta function to give you an action. Now, here, you can see the


pipeline is a bit different, we don't want to shove the observation in directly, right, we don't want


the observation directly. But what we put into the RL framework is this queue thing. Now, the queue is


supposed to be a representation of the observation and a useful representation. So if we think of


this, of this game here of this Atari game up here, what could be the what could be a useful


representation? If if I had to craft one by hand, how would I construct a useful representation? Keep


in mind, the representation, the goal is to have a representation of the observation that is more


useful to the RL algorithm than just the pure pixels of the image, right. So if I had to craft a


representation, let's say it's a vector, right, let's say our our representations need to be


vectors, what I would do is I would probably take the x and y coordinates of the little spaceship,


right, x and y and put it in the vector, that's pretty useful. And then I would probably take the


x and y coordinates of the meteorites that are around, right, let's say there's maximum two, so


x, y, x, y here, I would probably take the angle, right, the angle where my spaceship is pointing


to, that should be pretty useful, because if I shoot, I want to know where I shoot, right. So


theta here, and then probably, maybe the x and y coordinate of the of the shot here of the red


shot that I fired, if there is one, right, also going to put that into my representation, so x,


and y, and maybe delta x, delta y, something like this, right. So you can see if I had to handcraft


something, if I, I can pretty much guarantee that if I put in this representation right here, into


the RL algorithm, right, I put this in here, it would turn out guaranteed, it would turn out to be


a better RL agent that learns faster than if I put in the original observation, which is the the


pixel image of the game, right. Because, of course, in order to play the game correctly, in order to


play the game to win, you need to extract this information, right, you need to get others


something like a spaceship, there's something like meteorites, this is all things that the RL


algorithm doesn't know, per se, and would have to learn from the pixels, right. But if I already


give it the information that is useful, it can learn much faster. Right. So you can see, if I


handcraft a good representation, it's pretty easy for the RL algorithm to improve. Now, we want to


come up with a framework that automatically comes up with a good representation, right. So it


alleviates the RL algorithm here, the reinforcement learning, it alleviates that from learn from


having to learn a good representation, right, it already is burdened with learning the what a good


action is in any given situation, right, we want to alleviate it of the burden to also extract


useful information from the from the observation space. Right. So how do we do this, this this cue


here is supposed to be exactly that is supposed to be a good representation, but not one that we


handcrafted, but a used with a technique that can be employed pretty much everywhere. And the goal


at sorry, that the secret sauce here is this contrastive loss thing. Okay, this bombed,


contrastive learning is this this kind of magic thing that will make us good representations. So


what is contrastive learning? In this case, I'm going to explain it in this case, for this kind of


image based for image based reinforcement learning, but just for image based neural networks, how can


we come up with a contrastive loss. So you see, there's kind of a two pipeline thing going on


here, there is like this and this, and then one of them is going to be the the good encoding. All


right. So let's check it out. Let's say we have this image that we had before, right? Draw it again.


This little spaceship, this and this, and shot, right. And we want to, we want to do this, what we


need to do is we need to produce three different things from it. We need to produce an anchor,


what's called an anchor. So we need to produce a positive sample, positive sample, and we need to


produce negative samples. Let's just go with one negative sample for now. Right. So the goal is to


come up with a task that where we produce our own labels, right. So we want since we're training a


encoder, and the encoder is a neural network that's parameterized, we need some sort of loss


function. So the goal is to come up with a method where we can create our own labels to a task, but


that we construct the task in a way such that the neural network has no choice, but learn something


meaningful, even though we made the task up ourselves. All right, I hope this was kind of


clear. So how are we going to do this, our method of choice here is going to be random cropping. Now


random cropping means that I just I take an image, right, and I crop a piece from it. So a smaller


piece from the image, I just take a view inside the image. So in case of the anchor, right, I'm


going to draw the same picture here. Bear with me, I'm going to draw the same picture here, a couple


of times, this all supposed to be the same picture. And with the negative sample, I'm just going to


leave it empty for now. To meteorites, to meteorites shot shot, right. So for the anchor, we're going


to actually not random crop, but center crop. Right. So we're going to take here the center


image. Right. So the assumption is kind of that if I center, if I center crop, I won't lose, you


know, too much of the image, I can actually make the crop bigger, such that almost everything of


the image is somewhat contained in this. And that, yeah. Right. So this is going to be my anchor,


and then the positive sample is going to be a random crop of the same image. So I'm just


randomly going to select a same size, same size section from that image. Let's say this is up


right here. All right. And the negative sample is going to be around the crop from a different


image, right. So a different image might be from the same game, right, but might be there is a


meteorite here, right. And there is no shot. I don't I don't shoot. And I'm going to take a


random crop from this. Let's say I'm going to take a random crop here. Let's put a meteorite here as


well, just for fun. All right. So these are going to be our three samples. And now the question is


going to be, if I give the anchor to the neural network, I'm going to say I give you the anchor,


right. But I'm also going to give you this and this thing. And I'm not going to give any of this,


I'm just going to give whatever I cropped, right. So just just these things. So I asked the neural


network neural network, I give you the anchor. Now, which one of these two, which one of these


two crops comes from the same image? Right? So as a human, you look at this, and if you just see


the center crop, you see, okay, down here, there's this this tip of this thing. And then there's the


shot right. And in relation to the shot, there is a meteor here, right. And then you look at the


second one. And you say, Okay, I don't see the spaceship. But there's the same relation here from


the shot to the meteor. And I can kind of see the meteor up here. And this also fits with that,


right. And the spaceship must be, you know, down here somewhere. And then I go over here, and I


try to do the same thing as Okay, here's the meteor. And, you know, it might be, it might be


in the original image, it might be over here somewhere. So that's possible, I don't see it,


right? That's possible. But then there should be, there should be a shot right somewhere here, or,


sorry, further up, whoopsie. There should be a shot somewhere here, right? I'm pretty sure because


there's there's one over here. And I don't see it, right. So I am fairly sure, Mr. Task Asker, that


this image here is the positive sample, while this image here is the negative sample. Right? So this


is the task that you ask of the neural network, give it the anchor, and you ask, which one of the


of these two comes from the same image? Right? This is called contrastive learning. Now, there is a


bit more complicated in that, of course, what you do is you encode these things using neural


networks. And then, so each of the things you encode, so the anchor, you're going to encode all


of these things using a neural network, right. And then this is what's going to become the query. And


these are becoming the keys, so key one or key two. And then you're going to feed it always two of


them into a bilinear product, right. And a bilinear product is simply you can think of it as an inner


product in a perturbed space that you can learn. So you're going to have this, have these two here,


these go into q w k one, and then these two here, sorry, this and this go into q w k two. Now w here


is a learnable parameter, right? So you have some freedom. And then you basically take whichever one


of those two is highest, right? So this might be this high, and this might only be this high. And


then you say, aha, cool, this one's higher. So this one must be the positive, right. And you train the


w specifically to make this higher, to make the positive ones higher and the negative ones lower.


So this is a supervised learning task, right, where these things here are going to be the logits or


or the the so their inner products, but you basically then pick the one that is highest as a


in a softmax way. And they put this in the paper. So if we go down here, the objective that they use


to do the contrastive learning is this one. So as you can see, it's a softmax, like in multi class


classification of the inner product, the bilinear product with the positive samples over the


bilinear product with the positive samples plus the bilinear product with all of the negative


samples, so you're going to come up with more than one negative sample. Alright, now the only thing


left that we don't have here is that the encoding, how you're going to come from the image space to


this space here is going to be slightly different. And depending on whether you're talking on the


anchor, or on the what are called the keys, the things you compare to. And this is out of a kind


of a stability criterion, you already maybe you don't you know, like something like double Q


learning or things like this, it sometimes when you train with your own thing, so in Q learning,


you're kind of trying to come up with an actor and a critic or it's not the same thing, but you're


you're kind of using the same neural network twice in in your in your setup, and then you compare


the outputs toward to each other, which isn't, you know, it leads to instability. So in our case, we


took it three times here or multiple times, especially for the same objective here, we have


twice something that was encoded by the same neural network and is on the two sides of this


bilinear product. So if we were to use the same neural network that tends to be somewhat unstable.


So we have different neural networks, one that will encode the query, which is this FQ, and one


which will encode the keys, sorry, F, K. Now, we don't want to learn to neural networks. And that's


why there's a bit of a compromise where we say, it is the same neural network. But But basically,


this one is the one we learn. And then we always every now and then we transfer over the parameters


to that one. And in fact, each step, we transfer over the parameters and do an exponentially moving


average with the parameters of this momentum encoder from the step before. So the momentum


encoder parameters are a moving average of the parameters of the query encoder. And that is, so


you get kind of get the best of both worlds, you don't have to learn a second neural network. But


your second neural network is not the same as your first neural network. But it it kind of lags


behind but it is also it is also performing almost as well. So that is, I don't know if that makes


sense, but it is the best I can explain it. So to recap, you take your observation, you encode it


as a query, sorry, you crop crop here for your anchor that gets your query, and then you random


crop for your keys right into positive and negative samples, right? So you random crop from


the same observation, or from different observations, right, these become your positive


and negative samples, then you take you take you push these through your encoders for the query and


for the keys respectively, you end up with the Q which is the encoded anchor and case which are


the encoded positive and negative samples. And then you learn you update this encoder here using


the contrastive loss. Right. And at the same time, you feed the Q, you feed the Q here into the


reinforcement learning algorithm, and you learn your reinforcement learning algorithm. Instead of


giving having the observation directly as an input here, you now have the Q here as an input,


right, that is it, the reinforcement learning works exactly the same, but except having the pixel


input, oh, you now have the representation input Q. And you don't have to worry about anything else


in terms of the reinforcement learning algorithm stays exactly the same, right? The the this whole


thing here can actually run either in parallel, or you can think of it before you can think of it


off policy on policy, it is sort of modular how you how you fit this in simply comes up with good


representation. So that is that is basically the deal here, right? And you hope that the whole


procedure of this contrastive learning, then gives you good representation of this anchor thing


here. If you encode that to the queue, you hope that this representation now is a good representation


as a basis for the RL algorithm. And it turns out, at least in their experiments, it is. So here you


see the same thing, they actually, they do something more where in RL usually deal with a stack of


observations, not just a single observation, because, so for example, in Atari, people always


concatenate something like for the four last frames, right? And their their point is, okay, if


we have to stack here, if we do this data augmentation, you know, these crops, we kind of


need to do them consistently, right? We need to crop every single image at the same point for the


query. And also, if we do a random crop, let's say a random crop down here, we need to do this same


random crop for all of the of the stack of images here, right? So that is kind of the additional


thing they introduce it with respect to RL that deals with with stacked timeframes. But it's kind


of the same, the same diagram as above here. Right, so they explain the the RL algorithms they use


and exactly their, their, their thing. And here you can see that the anchor is a crop. And the


positive sample is a random crop from the same image, this would be up here somewhere, the anchor


is cropped from the middle. And then the negative would be a random crop from a different image.


Or a different stack of images. They have a pseudocode here, where, so it's pretty simple,


we'll just go through it quickly. Right, you start off with FQ and FK, these are the encoders for


the query and keys, you start them off the same. Then you go through your data loader, you do this


random augmentation of query and keys. And I don't not even sure if the random augmentation needs


actually to be a central crop for the anchor, or just two different two different crops from the


same image that might be as well. So you know, just I guess I guess it's a thing you could choose, I


don't know what exactly is the best thing. Alright, then I forward the query through the FQ, and I


forward the keys through the FK. Then important, I detach this so I don't train, I don't want to


train the FK, I only want to train the FQ, right. Then I do the bilinear product here with the W.


These are the bilinear product. And then I put this all of this into a cross entropy loss, right.


In the end, I update my FQ and my W and I do this exponentially moving average for my key encoder.


And they test on two different things. They test on the DeepMind control tasks, and they always test


100k time steps. So their big point is data efficiency, right? They claim they can use learn


useful representations with not much data. So the task is here, how good are you at 100k step


time steps, right? So you don't you don't optimize until the end, you just you get 100k time steps.


And then the question is, how how good are you? And the curl here outperforms all of the baselines


handily in the DeepMind control tasks. And it also outperforms a lot of the baselines in the Atari


tasks. And it actually if you look at the results, it doesn't outperform everything. But for example,


here the red is curl, and the dashed gray is state as AC. Now state as AC, the important thing to note


here is it has access to the state, whereas curl only works from pixels, right? So that what I said


before, like if I had to craft the use for representation, basically state as AC has access


to that. And you see that in many of the tasks that the curl comes close or performs equally well


to the state as AC, right? So that's pretty impressive. Especially if you look at pixel as AC,


sorry, which is the same algorithm, but does not have access to the state just the pixels,


it often fails terribly, right? So that is pretty interesting to see. And even to me, it's pretty


interesting to see that this kind of this kind of algorithm, this kind of self labeled algorithm,


this kind of self labeled algorithm comes up with such useful representations. Alright, so I hope I


have explained this satisfactorily. And check out the paper for more experiments, ablation studies,


Competition-Level Code Generation with AlphaCode (Paper Review)

AlphaCode is a system by DeepMind that does automated competitive programming.


You're able to give this system a lead code style problem in natural language and it will


come up with code by itself that solves the problem.


It does this by using a combination of language modeling, sampling, filtering and clustering


before it finally decides on the solutions that it's going to try out to submit to the


server.


What is mind blowing is that this system was able to perform in human competitions and


be about as good as the average programmer in these competitions, which is crazy because


previous systems were nowhere near human level.


So here's how it goes.


This video right here is a comprehensive paper review where I will go through the paper with


you and explain to you the most important parts of the paper, what's in there and what


I think is good and what I think is bad.


After this video, you'll have a good understanding of the paper and of how the system works and


what its potential weaknesses are.


However, in the next video released tomorrow, I will interview the authors of AlphaCode,


which is a huge privilege and I'll be able to ask them anything I want and they will


have seen my paper review and they'll be directly able to respond to any criticism that I've


raised there to any questions that I had and to whatever I did wrong in my paper review.


On top of that, you're able to get a behind the scenes look into their work.


Even at places like DeepMind, things go wrong, things don't work out.


They've had results that they thought were too good to be true and they turned out not


to be true and many more things.


On top of that, we talked about how the project came to be and also how they've dealt with


media reception because this paper has made big waves.


So I absolutely invite you to watch both this video and the interview part because they're


very much complimentary.


Let me know how I can improve these videos for you.


If you like, leave a like, tell someone to subscribe and I'll see you around.


Bye.


Hello there.


Today we're going to look at competition level code generation with AlphaCode.


This is by researchers of DeepMind and presents a novel system that can take part in competitive


programming challenges.


These are challenges where you as a user you'd register and then you'd be given lead code


style problems to solve and these aren't easy problems.


These aren't just solving some or writing down some SQL statement.


These are legitimate, difficult programming challenges where you need to think of algorithms


and solutions to problems and so on.


So having a system that can actually take part and compete against humans is very remarkable.


They've submitted this system to 10 of these challenges.


And as you can see, the orange lines here is AlphaCode's relation to other humans.


They perform about as well as a median human would like an average middle of the road competitive


programmer, if you will.


So this is pretty remarkable, especially since the baseline system so far had been sort of


in the third or fourth percentile, not very good.


So this represents a significant boost.


And today we're going to find out how they did it.


But first, here is what such a problem might look like.


So this is one problem.


This is one data point in this in this data set or one such challenge that you have to


solve.


You can see it starts with a description.


So the title is Backspace.


It starts with a description.


You're given two strings, S and T, both consisting of lowercase English letters, yada, yada,


yada.


What you should note right here is that the description is in natural language.


It's made for humans and therefore it's just natural that it is a natural language.


There is no other form.


There's no machine readable form right here.


This is it.


This is what the algorithm alpha code sees and gets as an input.


There's also a description of the input again in natural language.


There's description of the output.


And there is also this part right here.


This is an important part.


It consists of a bunch of example inputs and outputs.


So here is an example input.


For example, there are four problems in this problem set.


All of this will be described in the input section.


So the input section here says the first line is a single integer, the number of test cases


and so on.


So that's the four.


Then we have this is a problem.


So this is S and this is T of the first problem.


The goal is to type S and strategically type the backspace button instead of the letter


at S to go from S to T.


So in this case, we start with S. So the first letter is A, but we choose to type the backspace


button, which would not type A and would delete what we have, but we have nothing.


So yeah, then we would type B. Sorry about that.


And we would type B. Then we would type A, then we would type B. And instead of the last


A, we get again type the backspace button, which would delete the letter before it.


And we'd end up with B, A.


Therefore, we got from S to T. And therefore we output the letter the word Yes.


So we are tasked with writing an algorithm that automatically determines whether it's


possible to go from S to T in any of these test cases, and output the corresponding answer.


This is challenging by itself, but you only get the problem right, if you can do it for


all the test cases.


And the way these problems are evaluated is that on the on the test server, they have


a whole bunch more of these test cases, and including checking all the corner cases like


very long inputs, no input at all, only inputs containing the letter A. If for some reason


you expected a B to be there.


And so they test all the edge cases, and you need to be correct in all of them in order


to get the points.


This is extremely challenging even for a human, the output that you're supposed to give is


an algorithm like this, you can see it's not an easy thing.


It's not just a snippet, it's a full blown algorithm, it contains input.


So you read, you read the inputs, even that to program an algorithm to come up with that


piece of code is already challenging by itself to firstly read that first, let the first


line and then read as many inputs, then you need to build lists and reverse lists, and


you go into a while loop where you pop off things of list depending on comparisons.


And in the end, you output the correct thing depending on whether that list is zero, or


empty or not empty.


So as you can see, this is this is a challenging, challenging task.


And this is just one data point.


The next data point isn't going to be another variant on two strings and and typing the


backspace button.


The next data point is going to be a completely different problem, right?


Like searching for shortest paths and some graph or something with denominators of numbers


or numerators or, or something like this, right?


It is, it is very diverse set of problems and very challenging even for humans.


And the fact that an algorithm can tackle it is very remarkable.


So how do they do it?


That's that's our question today.


If you guessed that it has something to do with large language models, then and transformers


and so on, then yes, you know, kudos, you got it.


But there is a lot more to it.


And this is really an engineering effort.


And I think we should appreciate just, you know, how far you can push a system to get


continuous improvements.


What they do first, though, is they collect a data set.


They do train on a open source code from GitHub.


That is the pre training data set.


This is very similar to open AI codex model.


So open as codex model is trained on code from GitHub.


And you can simply do next token prediction on code.


And I have to say I've tried codex and I'm pretty happy with its suggestions is very


good.


But it can give me like longer snippets than an auto complete.


But it cannot like solve any kind of problems like this, it can just continue code.


In any case, they collect this pre training data set, they have whatever 700 gigabytes


of code that they train on.


And they do run their regular language modeling objective on that piece of code.


Then they fine tune on an appropriate data set of code contests.


So this is a mixture data set that they scrape from multiple websites, for example, code


forces description to code, code net, these are these are papers, previous papers, or


competition settings that they have collected these data sets from and the data sets again,


this here is one data point, right?


This is a problem description.


And usually these, these data sets, they contain one or multiple solutions, not all of them


might be correct.


But they contain about an order of magnitude more solutions than they contain texts or


problem descriptions.


So first, they collect a data set, and then they train on that data set.


So that could be the story right here, but it is not.


The entire pipeline is a bit more complicated.


You can see first, there's GitHub, we collect pre training data, we do pre training, then


fine tuning on pairs of problems and solutions of these code contests data set.


This is as I said, a collection of various data sets that contain these that contain


these code challenge type of problems, lead code style problems, and they do fine tuning.


By the way, their model is a transformer model, you could guess it, they do have a special


they have an encoder decoder model.


So you have some sort of an encoder, and they choose to make the encoder shallow and the


decoder, the decoder deep.


And there are specific reasons for that, which we'll get to in a second.


But the encoder mainly handles the this description, which is so the description is natural language


mostly contains, you know, some some code snippets and so on.


However, it contains mostly the description.


That's the encoder.


The benefit of using an encoder decoder architecture over a decoder only is that you do get bi


directionality in the encoder.


And as they do here, you can make them different sizes, which means that you can shrink the


encoder, which makes you sample able to sample faster and sampling is going to be very important


for the system right here in just a second.


And then the decoder will be a autoregressive decoder, where they just well int j equals


five, yada yada yada.


So this is this is actually going to produce the code token by token in sort of a language


modeling way.


Their objective is is they have a masked language model objective at the encoder.


And then the decoder obviously there's cross attention right here.


There's there's self attention in the encoder, there's self attention, causal self attention


in the decoder.


And then there is cross attention from the decoder to the encoder.


And they have a language modeling objective in the decoder.


And they do say, it's quite important to have the masked language modeling loss additionally


in the encoder, because it apparently makes the encoder understand this, the stuff in


inside of it, the stuff that it's fed a lot better, I'm just going to believe them right


here.


So now that we have this model, we can, we can fine tune it on these data sets, right,


we can feed a description right here.


And we can feed one of the solutions, and that could already be it.


However, that's not it.


It turns out that most of the time, this doesn't actually solve the problem.


So you feed in a description, and you sample the solution, it is not it does not go well.


So what do they do?


Well, there are two ways.


The first way is you try to make your model a lot better at like thinking and coming up


with solutions and reasoning abstractly and so on.


But that doesn't sound very deep learning and transformer like so what do we do is we


just do large scale sampling.


That essentially means you have a problem, you get a new problem, you feed this into


your decoder right here.


And then you just sample like a bunch of solutions from your decoder.


Sorry, I just said decoder over here.


You store this into the encoder, you let the decoder run, and you generate a ginormous,


a ginormous amount of outputs.


So you can do this with language models, you can sample according to some temperature,


you can do some other stuff, you do nuclear sampling and whatnot.


But you can generate diverse outputs from the decoder.


And they do, they sample thousands up to a million different outputs from the decoder.


So now they have this large set of potential solutions.


And what do they do with it?


This is very important, they do filter, and they cluster.


So first the filtering happens.


And it might not surprise you, but the filtering happens on these example inputs that we saw


right here.


So for every problem, you get a tiny amount of example inputs and corresponding example


outputs.


They simply let all of the programs they generate run on these example inputs.


And the ones that don't crash, they evaluate whether they do get the example outputs.


And if they do get the example outputs correctly, they keep them around, otherwise they discard


them.


This is obviously vastly different from how humans solve these things.


Humans don't just generate giant amounts of solutions and then let them run on this tiny


amount of example problems.


But this eliminates, as they say, it eliminates over 99% of these sampled things.


So you end up with a slice right here of this data that you've generated by simply evaluating


on these example cases that you had.


So it's quite important that these are there for the system to work.


I wonder if like, we could replace this, because we have this approach as well in for example


Dali, where a lot of stuff is generated and then clip is used to rerank.


I wonder if something like this could be done here.


But they have several helper models in here in order to help the system during training.


So I don't know if another helper model might be even appropriate.


So this leaves them with a tiny amount of solutions, which could still be a lot, right?


99% out of a million is still a lot of solutions.


And they keep themselves to just submitting 10 of them.


As a human, sometimes these code platforms, they have actually a limit on how many things


you can try to submit and 10 is like a reasonable limit.


It gives you a little bit of as a human, a little bit of you're not anxious to submit


a solution.


If you think it's the correct one, sorry.


But you also you can submit a few times, but not too often.


Like you can't brute force the test set that's on the server.


So they need to get down from these still large amount of solutions to 10 solutions.


And that's where this clustering comes in.


So the goal is to end up with this small select set of candidates to execute and evaluate.


And what do they do with the clustering?


This is where one of these helper models gets in.


So all of these things right here, they are programs, they're programs that could take


inputs and outputs, and there are many, many of them.


What we want to do is we want to cluster them, a lot of these programs are going to be different


in the tokens that they use, like in the exact code, but they're going to be essentially


the equivalent program to each other, like they're going to be the same program isomorphic


to each other.


However, graph isomorphism, like let's say we parse them in a syntax tree and check graph


isomorphism, that's, I do believe that's like a really hard problem.


I might be mistaken, but I think that's used in cryptography to show like a really hard


problem.


So it's not really graph isomorphism on the syntax tree, it might not even get all the


isomorphic programs.


So what do we do?


Our plan is going to be we want to group these programs into the same one.


So maybe these three here are actually the same.


And this one here is actually the same.


So we'd like to figure that out.


How do we do it?


We just feed like a whole bunch of inputs, like we just generate a whole bunch of inputs


to the programs.