Machine Learning with JAX - From Hero to HeroPro+ | Tutorial #2
❤️ Become The AI Epiphany Patreon ❤️
https://www.patreon.com/theaiepiphany

👨‍👩‍👧‍👦 Join our Discord community 👨‍👩‍👧‍👦
https://discord.gg/peBrCpheKE

This is the second video in the JAX series of tutorials.

JAX is a powerful and increasingly more popular ML library built by the Google Research team. The 2 most popular deep learning frameworks built on top of JAX are Haiku (DeepMInd) and Flax (Google Research).

In this video, we continue on and learn additional components needed to train complex ML models (such as NNs) on multiple machines!

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
✅ My GitHub repo: https://github.com/gordicaleksa/get-started-with-JAX

✅ JAX GitHub: https://github.com/google/jax
✅ JAX docs: https://jax.readthedocs.io/
▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

⌚️ Timetable: 
00:00:00 My get started with JAX repo
00:01:25 Stateful to stateless conversion
00:11:00 PyTrees in depth
00:17:45 Training an MLP in pure JAX
00:27:30 Custom PyTrees
00:32:55 Parallelism in JAX (TPUs example)
00:40:05 Communication between devices
00:46:05 value_and_grad and has_aux
00:48:45 Training an ML model on multiple machines
00:58:50 stop grad, per example grads
01:06:45 Implementing MAML in 3 lines
01:08:35 Outro

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
💰 BECOME A PATREON OF THE AI EPIPHANY ❤️ 

If these videos, GitHub projects, and blogs help you,
consider helping me out by supporting me on Patreon! 

The AI Epiphany - https://www.patreon.com/theaiepiphany
One-time donation - https://www.paypal.com/paypalme/theaiepiphany 

Huge thank you to these AI Epiphany patreons:
Eli Mahler
Petar Veličković
Bartłomiej Danek
Zvonimir Sabljic

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

💼 LinkedIn - https://www.linkedin.com/in/aleksagordic/ 
🐦 Twitter - https://twitter.com/gordic_aleksa 
👨‍👩‍👧‍👦 Discord - https://discord.gg/peBrCpheKE

📺 YouTube - https://www.youtube.com/c/TheAIEpiphany/
📚 Medium - https://gordicaleksa.medium.com/ 
💻 GitHub - https://github.com/gordicaleksa 
📢 AI Newsletter - https://aiepiphany.substack.com/

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬

#jax #machinelearning #framework
 What's cracking guys? This is the second video in the machine learning with Jack's series of tutorials and in this one We're going to learn how to build a fully fledged neural networks So in the last video we went from zero to hero and in this one we're going to go from hero to hero Pro Plus Yeah, anyways before we even start let me show you this thing So I open sourced a repo that's gonna called get started with Jack's that basically contains aside from these tutorials I'll be making other useful like content that I personally used and found useful So things like various YouTube videos blogs, etc. So and the second thing I want to I want to tell you is that you can Basically open up these notebooks I'm using in these videos in Google collab and you can easily play with them So you can just open it up and the Python environment is already set up So Jackson everything else is pre-installed so you can just kind of run this and everything will work Just run anyway, and this is going to to work out of the box So basically this like the notebook for this second video is already checked in as the time I'm recording this so you'll be able to open it up and maybe Tweak the code as you are watching the video So I think the best approach to get the most out of these videos is to Tweak the code yourself play a bit and and see where you can learn something something interesting awesome with that out of the way Let's continue with the video So let's first import the packages. We saw this in the last one So we were just importing Jack's the NumPy API of Jack's Transform functions that we saw in the last video. So grad jet V map in this one We're going to learn about P map which helps us parallelize the computation across multiple devices Which is very useful as you can imagine and yeah other necessary libraries here So the goal as I said will be to learn how to build complex more complex machine learning models Parallelize the computation across multiple devices and yeah, basically after this video You should be able to to get started and build your own models In order to get to the point where we can build neural networks There is a couple of things we have to you to first solve the first one is the problem of state So when I say that I mean the following so Jack's is basically uses the functional program programming paradigm, which means We need pure functions, which means Jack's does not love state in the sense that you cannot use your your usual Object-oriented programming type of stateful classes or functions. So we're gonna see a way around it right now First short reminder from the first video. We saw that impure functions are problematic So we saw this global variable G and we can treat it as a state for the purpose of this of this of this section Basically what this impure uses global function does it accesses G which is external state Which is forbidden because this makes the function impure and what happens now is when the JIT when we call JIT on this function It's going to cache the current value of G because that's a side effect. So it's gonna just cache the value 0 So that means the first time we call it is gonna is gonna run correctly But the second time we call it it's gonna call we're going to call the cached version of the function Which means 5 will be added to 0 and not 10 So in the meantime in the meanwhile, we've updated G to 10 So the state is updated to 10, but that will not be reflected because of the because we have impure function So this is the the trick you have to be aware of and as you can see here We get the correct result for the first call, but we get the incorrect result for the second call So that's 5 and we should have gotten I guess 5 plus 10. That's 15, right? Okay, that's the first step Secondly, we saw this pattern already. So basically we saw that the pseudo random number generator in Jax is not stateful in contrast to two numPy's and we saw this pattern where we create the state the initial state By just passing in the seed and then in order to get the the next state We just call this split function and we pass we externally pass the state. So that's that's a big difference compared to numPy By pseudo random number generator. So I basically want you to remember this pattern So we are inputting the state inside of the function. So split is not saving state internally is just taking it as an input It's manipulating it somehow in a pure fashion and returning back some novel states Whereas just recall that we were using key and sub key Terminology previously. I'm just want to stress out that we are Basically, I want to stress out this pattern. So that's why I'm using this new terminology. But yeah, you get the point hopefully Okay. So now let's see a way around this and if I haven't made this already clear enough Let me just kind of motivate where we're doing this. Basically if you're trying to create a neural network, for example Like there is a lot of stateful code when you're trying to build neural networks Especially in other frameworks such as PyTorch or TensorFlow Basically, our model parameters are stateful your optimizer is stateful atom is for example storing internally a certain Statistics so that means it's stateful and also we have a layer such as batch norm, which are stateful Yeah in summary and any useful program will will have to have some some kind of state and we just need to find Alternative way how we can handle it. So in this first example, we have a stateful class. It's called counter We have this state called like n which is basically the internal count and we have the function count We just increments the this this state by one and returns the current value of that state Finally, once we call reset, we just set the state to zero and that I guess all makes sense here We just instantiate the an object of this class and then we we call we call three times We call this this count function and as expected Hopefully we'll get one two three and here they are Okay, so that all works like a charm and this is something you're used to now the problem arises here When you try to get the function, so let's reset the counter and let's so that means the state is now internally set to zero And now let's try and get the counter count function. So let's try and get this function here Now, let's see what happens with it Once we did it and we try to again run the same type of like a for loop here The results are kind of surprising maybe or maybe not even familiar with jet Basically we get all ones and the reason is as I said as this count function is not pure because we are accessing We are basically writing to an external state as you can see here. So self dot n is a state Basically, it's not passed in as an argument. We are we are just modifying it externally So that's bad. And what will happen is so once we call reset here And we'll be set to zero and the first time so once we get the first time we call a jitted function We're gonna do the tracing if you recall from the last video and that means that basically on the first call The function will return one and because we have a side effect is just gonna learn how to cache number one and that's it Now, let me just kind of try and clarify that if we were to use make Jack spur to make this more clear Let me run this and see what the like underlying Jack's expression Jack's expression is and we can see it learn to just return Like one so basically the function the cache function has like the whole body of the function is something like this return one That's it. And that's why we're getting ones even though we are we are expecting to get increments Okay, just to make this a bit more clear if we were to call counter count function here We'd expect the internal state to be modified to like one and then once we call jet It's gonna be returning twos now, hopefully So if we correctly understand this if I run this we get twos and the reason is again It's basically tracing this function the first time it's racing yet. The value was already set to one We add additionally one so we have two and we turn two and basically because we don't have any inputs It's just gonna learn how to cache and return those number two and again if if I were to do the same thing here So count and call the make Jack Jack spur. We're hopefully gonna get like just two Instead of one here. Let's see and that's correct. Okay, we got two here In any case, we have to find a way to avoid like creating these in pure function So let's see what the solution is we're going to define this counter state Like instead of defining a complex class, we're just gonna have like a simple primitive data type here. So integer So this version two of the counter is gonna work the following way So we're going to pass in the counter state and we're going to return back the output as well as the state Now this may be a bit confusing m plus one where we're turning a couple of m plus one two times I mean doesn't make any sense right? So the thing is in general this thing here will be state and this thing here will be some function of the state So here in the example of this simple counter the output as well as the state are the same thing So yeah, we're just returning m plus one, but yeah doesn't doesn't matter The whole point is we are passing in the state and we are processing it and then returning it back In the case of reset function, we're just returning back the state zero and that's it So here we instantiate the counter here. We call the reset and as you can see here compared to the version one of the counter We are not internally modifying the state. We are just returning back the state here externally and here we can see the pattern So we've called the counter count we pass in the state and we get the state and the value we print the value and we keep doing this And so yeah, I guess this should already be familiar Basically Jack's like run random number generators use this very same pattern So this should work like as expected now and as we can see we get one two three But but now let's see whether after we did the function this thing will work or not Let's call reset again. We get the initial state Let's get the function and let's try and run this And as expected we're getting correct results Even with this jaded version of the function because we are externally manipulating the state and not creating impure functions And yep, this is working as expected and in summary, let's see the pattern we we use to convert a stateful function into a stateless one So basically we had the state which was just number n in the case of a simple counter We have this stateful method. So the the especially like in particular we had this method called count We are passing some arguments in general case and we get back the output Now we converted this and instead of saving the state inside of the class We're passing the state inside here as an input and we return back the top also the output and the state And this pattern can be applied for to any a stateful class and to get a stateless one and that's cool Okay, that's that was step one towards building like a fully fledged neural networks Now the second step is we still need to find a way how to deal with gradients And you may wonder why we have to like learn anything new to deal with gradients Well, we already know about the grad function. Well, the thing is with our current knowledge We would go something like this. So imagine we have some loss function F and we have like the parameters of our neural network. So x y z and w So we basically only have four parameters for now and we do something like x squared plus y squared plus z squared Blah blah blah. So basically we get this is a L2 you can treat this as L2 loss So minimizing this loss would cause the weights to go closer to get closer to zero And now the problem here is if we had something like GPT-3 model Which has 175 billion parameters and some newer language models are even bigger Basically, you'd have to write 175 billion like terms here minus four that we already have here And that sucks, right? So the problem does not stop there Basically, if we now try to find the gradients because we want to do an update of our neural network parameters We'd have to get dfdx. So we have to find derivatives with respect to every single one of those arguments And grad as you recall usually by default returns the like a derivative with respect to only the first argument So that will be dfdx. So we have to explicitly pass this list of integers And in the case of 175 billion would have quite a lot of these, right? So this will not scale and yeah, let me just run this. This is going to work as expected But as I said, it's bothersome and it does not scale So we had twos because if we take argument dfdx is going to be 2x, 2y, 2z in the case of dfdz, etc. So basically because we input all once as the input we get twos for every single gradient So that makes sense. Now again, the problem does not stop here If you want to do an update would have to have 175 billion lines doing these types of SGD updates So basically learning rate gradients and we update our weights here So there is obviously a better way. This is just a motivation for why we need the following solution And the solution is called a Pytree So in general we are used to wrapping our model parameters in some more complex data structures dictionaries or whatever And Pytree is the Jack's solution to how to basically find gradients of arbitrarily nested data structures And let's see a simple example here. So this is just a contrived example obviously You can treat it as some type of model parameters, very esoteric but yes, nonetheless Just stick with me here. We're going to first learn what Pytree is and then we're going to see a real example We're going to train like a multi-layer perceptron really soon. Let's see how Pytree is defined in the official docs So if I open up the docs here, if I go to tutorialjacks101 and find the working with Pytrees here Okay, we have the definition here. So a Pytree is a container of leaf elements and or more Pytrees So it has this recursive definition. Containers include lists, tuples and dictionaries So let's see how that maps to our example here. So basically we have a container here So this whole thing is a Pytree, but this Pytree contains another Pytree here And finally this Pytree has a leaf 1 and the leaf string this character a And finally this object is also considered a leaf. The second, this is also a Pytree, this whole tuple here And tuple I guess. This is a leaf, one will be a leaf, whereas this thing here will again be a Pytree And two will be a leaf and three will be a leaf. There is a handy way to analyze this and it's called this From Jax, just import this tree leaves function. We're going to iterate through these Pytrees and we're going to get the leaves And we're going to finally print them in a nice format here. So let's see what the results here are Basically this first Pytree has three leaves as we can see here. So one, character a and a simple empty object Let me see what else is interesting. So as I said this is 1, 2, 3, all of those will be leaves In the case of a dictionary, keys are not considered leaves. Keys are just like metadata of the Pytree node So we're just going to have 1, 2, 3, 4 and 5 as leaves as you can see here And interesting to notice here is finally this array is going to be considered a single leaf So we won't have 1, 2, 3 as separate leaves. This whole thing will be treated as a leaf So that's something to keep in mind. Okay, so how do we manipulate these trees? There's a couple of neat and handy functions. So again we have an example here of a simple Pytree And what we do is we call this tree map function and we just pass in the Pytree And we define a function that is going to manipulate and work on top of these It's going to process the leaves of the Pytree and basically what we expect is that every leaf will be multiplied by 2 In the case of this simple lambda function. So let's run this and as expected we get everything is multiplied by 2 So that was in the case of these single argument functions. In the case when you have multiple Like when your function has multiple arguments, I'm just going to create a copy here So a reference basically and I'm going to do a simple addition and this time we're going to call treeMultiMap So we have 2 Pytrees as input. We manipulate them using treeMultiMap and as the output we get the Pytree again So if I run this we would expect same results and yep we have the same results here because Just adding two same structures is the same as doing multiply by 2, right? Okay, so finally worth noticing here is that the structures have to be the same So here the condition was trivially I guess true because I just made this a reference But in the case when you have two separate data structures you want to make sure they have the same structure Because if I artificially here introduce a change so I just do a deep copy of the original object And I just append this list that contains number 23 and if we try to run this We can see that list arityMismatch is reported as an error So yeah, keep that in mind and that was the last missing piece before we can train like a multilayer perceptron Hopefully you already know what MLP is if you don't let me just quickly show you some diagrams And whoops, okay, we are in a really bad part of the internet right now. Let me go to images here If this thing will load my internet is really really slow today Okay, so I'm gonna just type in neural network here and hopefully everything will work correctly And here we are. Here is a simple MLP basically you have your input data coming in from the left You're processing the inputs using these neurons. These are called hidden layers This is the output layer and outcomes to results So this thing may be like a binary classifier like a hot dog or not hot dog I don't know in any case that's MLP at least on a high level and Now, okay, we are we're ready to start training a toy multilayer perceptron model So we know how to handle state we know how to handle parameters I need to find gradients of our parameters which may be stored in certain types of nested data structures So that's everything we need to to train our first model in Jax, I guess So let's start here. We have this init MLP parameters function What it's going to do is return a pie tree object and here you can see we specify the type Basically the configuration of our MLP. So that's gonna it's gonna have three layers So the first one will map input. That's a scalar into 128 dimensional vector Then we're gonna have a mapping from 128 to 128 and finally we're gonna have a mapping from 128 to output scalar again So that's three layer deep MLP. So this function here just takes that list as an input as you can see here using this smart indexing We can just fetch tuples such as this So we first have the tuple 128 and we're gonna use that to form weights and biases of our neural of our like a linear layer And we're gonna wrap it up into the dictionary and just append to this list So this is gonna be a list of dictionaries which contain two keys weights and biases which then contain like a random meet randomly initialized matrices And you don't have to worry about this. This is just some specific type. I think gaming or something initialization doesn't matter And think to notice we're using numpy random here instead of using Jax PRNG doesn't matter for this simple example But just keep that in mind Anyways, so we get back the the the pie tree and because we have a pie tree we can nicely use the tree map to print certain like information about the pie tree And we can see we've done everything correctly. We have biases and weights. Everything is of correct shape In general, this is a good good practice to have to analyze the shapes of your if your data and model parameters and whatnot throughout the whole pipeline basically Okay, let's continue here. We're gonna have some update function here in order to train it This is the main part and I'm just going to before we even start digging into the training code. I'm just gonna comment on the training code I'm just gonna comment on this forward because we haven't defined it yet and I'm just gonna show you the input data So this is what we're trying to learn to regress So the MLP will have as an input a simple scalar and we're trying to regress the corresponding y which is simply x squared in this in this example here So now let's see how how this thing works. Let me uncomment this basically We'll have a couple of epochs. So 5000 epochs in this example. We're gonna call the update function and you see the the pattern again We are passing in the state which is the model parameters in this case and we return back the parameters again and we just iterate that and that's it That's going to update the parameters and finally we can regress using the forward method Let me show how those are defined. So update basically takes as an input parameters x and y. So that's the input and output data We calculate the gradients of the loss function with respect to basically this input data and the current model Just stick with me. I'm gonna analyze this in a bit more detail in a couple minutes. For now, let's just see what the loss function is It's a simple mean squared error loss. Basically we do here prediction. We find an error term that's basically predicted y minus the ground truth y We square it to get basically positive values error terms and then we do the mean across all of the data points and that's your MSC loss Forward is fairly simple. You have a this is basically your MLP feed forward pass. So remember params is a pie tree When we iterate layer is just a dictionary that contains weights and biases. So we take out the weights We do dot product with the input data and we just add the biases here and then apply the activation function rally in this particular case So there is a dichotomy we made here between hidden layers. Whoops hidden layers in the output layer The reason is the last layer will not apply an activation function. We're just going to return back the raw results That's the reason for this dichotomy. I guess we can rewrite this a bit nicer, but yeah, it works Okay that out of the way. Let's now focus on this part. We do the grad of our loss function and as you can recall Basically, it's gonna do a derivative with respect to the first parameter and since our first parameter is params and which is a pie tree Will basically get as the output the same structure as params just But instead of values of instead of the weights will have derivatives of those weights with respect to the loss function So that's the magic behind grad. Basically all of these built-in functions transform functions such as grad know how to handle pie trees And this is finally a concise way to solve this problem of figuring out the gradients Now that we have gradients we can use our loved tree multimap function and we just pass in parameters and gradients We have this so on every single leaf of these pie trees we're going to apply the following rule and that's from the weight Just subtract the learning rate times the gradient which is simple stochastic gradient descent update rule and that's it As the result again tree multimap as you recall it returns back the pie tree like structure again Which means we have the updated weights getting back as the output of update and we can see that here So params in params come out as well here. Now, let's try to run this and then we're going to play a bit with this code So let me run this cell. Let's see where I ran this one. Yeah, I did. So let me run this one and let's see what the results are Okay, this is the output we can see a pretty nice fit in this particular case We can now play with this function. We can try and see some other polynomials maybe X X cubed We can see it also manages to fit this curve as well. We can play and add some other types of polynomials something like this We just I just added plus X here and it's still be like you can see at the borders here because we have less data points The fittest has a bigger errors compared to where we have a dense sampling which is expected I guess right? Finally, let's try one more interesting example that sign us excess. Let's see what we get here Whoops, when you use jet numpy sign here, so this should work now And you can see the fit is way worse here. We can obviously improve that by increasing the capacity of the model or increasing the number of epochs But yeah, let's try some other frequency like 3 X and yeah, we have a much better fit here In any case you get the point. We trained our first neural network in Jax and it seems to be working So let me return back the original statement here. So that's a simple like a polynomial and now let's analyze a couple of things here So first, let's see the structure of grads. I'm gonna reduce this to one just a single epoch I'm gonna comment out the JIT part and I'm gonna print grads here Okay, so I'm gonna just print the gradients here so that we are confident that we get the same structure as params So I'm gonna run this cell. I'm gonna run this again. I guess the output is gonna be a bit unwieldy Okay, let me see what we got here Bunch of numbers, but you can basically see that we have a list of dictionaries where the keys are biases and weights And basically these here these numbers are derivatives and not the parameter values. That's the important part So that's this part. Let me just uncomment this, return back the JIT And I mean that's pretty much it. That's everything that goes into training an MLP As a quick mention here, you can notice that I called JIT on top of this high level update function Because it gives the XLA compiler much more freedom to optimize various functionalities inside of this update function You can obviously play and toggle off and on the JIT wrapper here and see how faster it is to train it with JIT compared to without JIT, I guess Okay, let's continue here. This was an example of a simple neural network But if we wanted to create our own neural network library, we'd need to create layers such as NNLinear That's just PyTorch syntax or conv layers, etc. And in order to get those to work, we have to create our custom PyTrees And let's see why that is. So first imagine this is this my container object here is just a simple NNLinear layer So let's imagine that's a linear layer and now imagine we have a list of two linear layers here So just create an example PyTree here and now let's fetch the leaves and see what we get as the output And we would expect, I guess, like eight leaves here, right? But the thing is it's going to return only two leaves So we have here that we this whole PyTree only contains two leaves and those are this my container objects That's obviously going to be undesirable because TreeMap and those other functions, they go through the leaves and do certain updates And if we had a linear layer like this, we wouldn't be able to do a gradient, like stochastic gradient descent update on top of these structures So we have to find a way around this. For example, if we were to run this simple manipulation, so just x plus one We'd expect to just increment these numerical values, but this is not going to work And I guess it makes sense because if you know how TreeMap works, it's going to iterate over the leaves And the leaves are my container objects and we're trying to add one to my container, which doesn't make any sense So let's run this and we can see that we get this error, we're trying to unsupport operand blah blah plus for my container and integer So that makes sense, right? So what's the solution to this? It's simple, you just have to define two functions One is flatten, one is unflatten and then you have to register PyTree nodes So we registered the my container with these two associated functions, flatten and unflatten So what flatten does basically just take the ABC, so those were the values we had here, right? We'll just wrap that information into an iterable here and we are not going to want to have a name as the content We want to have it as a metadata and that's why we include it in this auxiliary data structure And we just return this tuple of like content and auxiliary data Finally, the unflatten just receives those same inputs in reverse order though because of I think some legacy problems they had And you basically just construct the my container and you return it back, that's it So after I run this, now the my container class should work as expected So if we were to run the tree leaves function again, let's run it and we have the leaves 1, 2, 3, 4, 5, 6 as expected Now this thing will work, so if we were to add plus 1, we'll just get I guess 2, 3, 4 should be the output here And yes, that is the case Basically tree leaves, i.e. the tree map iterated over the leaves which are just the numerical values here, these ones And it just added plus 1, returned a PyTree again and we just print the leaves here, that's it There is one more gotcha I want to explain, tell you about and that's that we can oftentimes mistake nodes for leaves And when I say that, here's what I mean, so let's create a zero PyTree So basically we just have two nodes here, two leaves that contain zeros with these shapes And let's try to create the same type of PyTree that will have ones instead of zeros So what we do here is we create this intermediate PyTree, we're going to iterate through the zeros tree, through all of the leaves We're gonna take the shape and that's gonna be our new PyTree And then we can iterate through the shapes and apply the ones function using tree map to get like the tree of ones And let's see whether this works or not So if I print it like this, here is what we get We get correctly 2, 3, so we have two rows here and three columns, everything is fine We have three rows here and four columns, everything is fine here It seems that we have a correct result here as well, so we have 2, 3 shape and we have 3, 4 So that's the intermediate array, so this shapes tree here And finally we have some weird result here, so what happened? So the thing that happened is, if you can recall, basically this here, this whole thing is a PyTree But this tuple is a PyTree as well, it's not a leaf, that means that leaf is 2 and 3 and 3 and 4 Which means we're going to apply ones on top of all of these leaves and that's why we get two ones here and then three ones here and then three ones here and four ones here So simple work around here is to create, so we wanna make sure that this is a leaf and not a PyTree We can just basically wrap this x.shape with jmp array and this should now work Let me see, and yep, it works, so basically the reason is this is now a leaf, 2, 3 are a leaf And we pass that to ones and that's why we get correct results here, awesome Now we can design custom neural network layers and we can train end to end our neural networks Now what if we wanna train really, really big neural networks? Well, in that case we have to parallelize our computation and that's where this next section comes into play So I'm gonna introduce this pmap transform function and as we'll soon see it's very, very similar to vmap But before I start showing you how to use pmap we have to do a couple of things So the first thing is to set up, we have to use TPUs here for this section So I'm gonna set up hardware accelerator to a TPU, I'm gonna click save and hopefully after clicking reconnect it's gonna allocate a couple of TPUs for me Awesome, this worked, so now I just have to rerun the import statements because interrupting the runtime erases all of the whole state of the colab So we have to rerun this cell and then let's go back, I can open it up like this, just jump to parallelism injects in the table of contents And we are here, okay, so we're first going to just call this setup function and that's gonna set up TPUs in the background and hopefully we'll get 8 cores at disposal Okay, running this cell may take a while but finally we got 8 cores here as you can see, so 0 through 7 Now I'm gonna use this convolve function as a running example and we'll be parallelizing this particular piece of computation across multiple cores Although we'll later see training annual network, actually a simple ML model across parallel cores So let's start with convolve, so I'm gonna define a simple signal here, so we have a signal X and we have a kernel W And the convolve function works as expected, basically we are sliding W across X, we are doing dot products and we are storing the computations in this list And then we're just converting the list into like a Jax array, that's it, just a simple 1D convolution So if we run here, let's see the results and we get the correct results, so basically imagine, let me print X and let me print W as well here So that may make a bit more sense, so what we do is in the first computation we'll have 2 times 0, that's 0, 1 times 3, that's 3 And 2 times 4, that's 8, 8 plus 3 is 11 and that's why we get 11 as the first output, that's all nice and clear So now let's do the same thing, but this time let's parallelize the computation, so first this is just going to grab the number of devices on disposal We'll have 8 cores and then we're just gonna create a batch to simulate heavy loads, so some bigger computation compared to just having a single signal X and kernel W So let's run this and basically we have 8 devices and that's why we end up having a batch of 8 times 5, where 5 is basically the signal length and we're just replicating So here we have unique data for all of these 8 elements in the batch, on the other hand we're just duplicating the kernels, we're not creating novel kernels for each of the examples Ok, that's it, let's now first parallelize the computation across the memory of the host device, that means we'll be using a single core But VMAP is gonna parallelize the computation in the background across multiple threads, I'm not sure how it's exactly implemented Basically we wrap the convol function into VMAP, by default it's gonna assume that the in-axis argument is set to 0,0, so this is the equivalent to explicitly writing in-axis equal to this tuple of 0,0 And after we run this we'll get the result, and as you can see this part is the same as the example above, and then we have increasing numbers because the signal X is being, let me print X here Let me just print this batch, you can see we're going from 0 through 39 and that's why you can see increasing numbers here as well The cool thing about VMAP now is that we just need to swap VMAP for PMAP, nothing else changes and all of a sudden we're sharding our batch of data So the signal X across multiple cores and we're running the computation in parallel independently on each one of these cores So let me run this and we get the same results as you can see here above for VMAP, there is one difference you can notice here This thing is called device array, whereas here we have sharded device array, so this sharded alludes to the fact that the actual data and computation is distributed across multiple cores So that's super cool, so syntax same as VMAP and that's nice and sweet Let's continue, and I want to stress that all of this happens literally independently, so basically we have 8 signals, so recall that X was 8 times 5 What will happen is that array of 5 elements is going to be on each of the cores alongside the kernels, so we basically have a single computation So we have a single kernel and a single example on each of the cores and we are doing computations without communicating between the cores So that's super nice if the nature of our computation is such that we do not need a communication between the cores And here I'm just illustrating that we can do one pass of this computation, so we do a convolution one times And then we use the outputs as the new kernels and we again convolve them with the original signal X And basically the thing is during both of these two calls we did not have a single communication between the cores themselves So let me run this, and we get some even bigger numbers and that's not that important anymore Finally, just notice that instead of repeating, replicating W's, which is kind of wasteful, we can just use the in-axis argument here And as you recall, VMAP uses the same syntax, basically in-axis tells PMAP that the first argument, so that means W, basically should be broadcasted Because we set none here, and here zero just means that the first zero dimension of XS corresponds to the batch dimension So this way, PMAP will in the background replicate W and do the job for us, so that's way better And let's run this and we should see the same results as before, so 11, 20, 29, 326, that was the same number as here And elsewhere, so that's the same output again So again, this is super nice if your computation is such that you do not need communication between the cores This is usually not the case, and an example we'll soon see is the one of training ML model in a distributed fashion Where as we take our batch of data and we basically split the batch of data into mini-batches, we send each mini-batch to each core And then we do computations there, we calculate the gradients and then we communicate between the cores in order to do a mean of the gradients to update our ML model We're gonna see that example in a second, but just be aware that certain computations do require communication, certain do not, and Jack's got us covered for both of those use cases So now let's see a simple example, like just a modification on the above function that we had So this is going to be normalized convolution, and the only difference is that we'll have a communication between cores So let's see how this works This part of the function remains totally the same as before, so we have output here, we are doing like sliding across the signal, and we accumulate the results into output Now the only difference is that we are calling this, so remember, we're calling this Jack's LexPSum function, and just recall that Lex is this middle layer API of Jack's So what we're doing here is we're going to, along the batch dimension, we're going to sum up all of the outputs, and we're gonna divide output on this particular core with the sum of outputs across all of the cores And that's gonna normalize our output So a thing to notice is this axis name, it can be a bit confusing, so what you do is you again wrap your function of interest into PMAP You specify in axis, which means, hey, I want to broadcast W, I want to use the zero dimension of XS as the batch dimension And we just give the name to those mapped axis by putting an arbitrary name here, so this can be whatever Like we just need to have consistency between whatever we use here, we have to use here inside of the definition as well So I chose batch dimension because it's, I guess, indicative of the thing we're doing, right? Let's first run this and then we can analyze it additionally The first thing you can notice is that this output is the same as VMAP So I'm just going to, I just did that for you to understand that VMAP and PMAP are very, conceptually very similar The only difference is that PMAP is actually using physical cores, whereas VMAP is doing the same thing just on the single device So let me just comment out VMAP so that we can have a bit less output in the console here So basically what the normalization has done is that all of these columns are now normalized Meaning if we sum them up, if we take a single column like this one, so this one, like this element and this element and this element And we sum them up, we're going to get one, so that's what I've done here I took the result, I took all of the rows and zeroth column and we can see here that result sums up to one If I put one here, we're going to get the same result, so it's again one And if I put two here, we're going to get one as the output and that's it That's what this computation did Let me try and see whether we can print the output here I'm not sure whether this is going to work, let me just try and do this Nope, it's going to trace it because PMAP, that's an additional thing to keep in mind PMAP is automatically calling JIT in the background, so yeah To make this a bit more clear, I'm going to take the last example that didn't have any communication So let me take something like this I'm just going to copy paste it here So we're just going to copy paste that thing here Okay, now we have a result here that is not... I'm going to print it here So PMAP result and this result is... there is no normalization here So let me just print it out here So here are the results So what we have done is this PSUM basically alongside batch dimension, which is this dimension here So the vertical dimension is going to sum up all of these values And use them to divide the current output So for each core, for example, for core 0, this will be the output And we're going to divide this array here with the sum across this dimension, the vertical dimension So let me just kind of show you that that is the case So if I was to take the element number like this element here, 11 So that's going to be temp element, something like this, I don't know PMAP result Basically I'm going to take 0th So that's this first row, I'm going to take the 0th element, so that's 11 Let me confirm that's the case, so I'm going to print it And then I'm going to create a temporary sum here So basically just a simple sum of PMAP result across basically across the vertical dimension So that's the batch dimension, 0th column And I'm going to print the sum as well And then I'm going to just do division So I'm going to print temp element divide by temp sum And we should get the same result as here, so 0.0081 whatever, okay? So in order to make this a bit less crowded, let me see where I can do something No, let me just run this This should work as expected now So I took 11, I created a sum here And after dividing you can see we got the same result So hopefully this now is perfectly clear what happened And the whole magic was because of this middle level API's PSUM function And just giving a name to these mapped axes here Let's continue, let me just erase this output Okay, now before we get to the final training pipeline There's a couple more functions you should be familiar with And then you can train whatever in PureJax Pretty much whatever ML model you wish to train The first function is this value and grad function And it's going to return the value of the loss as well as the gradients of the function of interest So I have a simple loss function here It's basically just, we just create these error terms We square them and we sum them up We have a simple input here So let me just print it here So basically it's going to be 0 through 3 And finally the value will hopefully return the loss value as well as the gradients It's fairly trivial to show that the numbers make sense Let me just print y here as well So it's going to be the same as x just plus 0.1 So basically the value of the loss does make sense because we basically subtract The difference will be 0.1 for each element here Then we square it which is going to give us 0.01 And we have 4 times that because we have a sum So that's why we have 0.04 here Same goes for gradients, all of the gradients are minus 0.2 Because the derivative with respect to x will be 2 times x minus y The difference is 0.1 I mean the minus 0.1 times 2 That's why we get these numbers here So, and in any case if you want to log both the value of loss and you want to get your gradients You have this efficient way to do that is to just call this inbuilt function called value and grad The second function you should be familiar with is Not a function actually, just a functionality Is if you want to return sometimes from your loss function You want to return not just the loss itself But maybe also some intermediate values like maybe you want to return these raw error terms You can just instead of doing this So if I were to call gradient function on this loss function here It has two outputs which means a vector output We're going to get an error here because gradient only works on scalar output functions So you can see here gradient only defined for scalar output functions So there is this neat auxiliary flag If we set it to true and we run this It's going to return the gradient as well as the It's just going to pass this intermediate value as the output directly Without doing any type of manipulation on top of it Awesome, we are now ready to train a machine learning model completely in parallel Okay, let me see how I can go about explaining this part Because there is a lot of functionality here But I think it's going to be digestible Let me first start by analyzing our data So we have a simple model like the ground truth model is this one Basically a single input linear function here With some addition of noise, it's a noisy, noisy We'll have some noise observations here So the true value of W, so the coefficient that goes with X And of this bias term are 2 and minus 1 We generate the domain of our data Basically just sampling from the normal distribution We generate 128 elements We generate some noise, again, sampling from the normal distribution And finally we generate our noisy observations by applying the above function here So XS times the true value of the weight parameter Plus the true value of the bias parameter plus the noise And let me visualize that for you So we understand what we are trying to regress So we are going to try to fit a line to this data And we are going to do that the following way So first we'll need to initialize our model So let's jump to the init model function here It's going to initially create a random value for the weight and for the bias So this is our approximation of the true weight and the true bias values Initially the best we can do is just to create a random value And return them packed into this params class Which is just inherited from this name tuple So it has weight key and bias key and the values are our Jax arrays As you can see here I'm using Jax's built-in PRNG functionality I'm using the random number generator I'm splitting the keys into weight and bias And I'm consuming those here to generate two random numbers So that's what this init model function does Once we have that I'm going to fetch the number of devices That's going to be 8 if you recall from the previous exercises we had 8 cores And we are going to replicate the parameters because we want to shard this We want to replicate this particular model across all the cores and do the computation in parallel Since params, so the thing we just got from our init model function is a pie tree We can iterate over the leaves, so over the weight and over the bias And we can just multiply them by 8 which is going to replicate the values So we'll have, yeah, I'll just going to print it here so you're going to see it in a second Next thing we're going to do is just reshape the data for the PMAP function Because we're going to trade this in parallel again So we have the XS and YS, so that's the data we just saw here And we're going to basically just reshape them so that the leading dimension has the number 8 So let me run this and see what we have So, okay, I forgot to run this cell I'm going to run this cell and then I'm going to rerun this cell again And let's see what we got here So you can see a couple of things here First thing, weight and bias have the same number replicated 8 times as expected And finally the data is now reshaped into this 8 times 16 times 1 shape which is suitable for PMAP The Bayesian training pipeline will look very similar to the MLP we trained like 20 minutes ago Basically we're going to have a number of epochs, we're going to iterate for a certain number of epochs We're going to have an update function and again the familiar pattern We pass in the replicated parameters, we pass in the data And we get back the parameters, updated ones, as well as the loss value After that I'm going to just log some values here and we're going to see that in a second And finally we're going to take our final replicated parameters And we're just going to fetch parameters from a single core Because we'll basically have 8 duplicates across the cores So we just need to extract a single one and we get parameters here So TREEMAP will create a PyTree object and then DeviceCat just takes the data from the TPU and downloads it to the host memory That's it, now that's the high level picture of what is going to happen here Now let me start digesting these couple of functions here So what we need to understand is that the model is very simple We just have weight times XS plus the bias because that corresponds perfectly Because we know what generated our data so we can create the same model here Usually that's not the case, you usually don't know how your model structure looks like That's why you have a huge neural network and you let the data morph the structure of that neural network Loss function is very simple, we just basically do mean square error again So predictions, we subtract the ground truth, etc The update function is by far the hardest to understand so we're going to focus a lot on this function So let's see what update function does So we get our parameters, so that's the current version of the parameters, we get the data And what we do is first we decorate the update with this PMAP and we name the mapped axis to batch So this is just a convenient way to, if we do it like this then we don't have to wrap it inside of the PMAP here So let me show you, so here as you can see we don't have PMAP explicitly here We just decorated the update directly with PMAP and that's why we can avoid putting PMAP here So what will happen is that each TPU core is going to get a portion of the parameters So because those are replicated we're just going to have the same parameters spread out across every single core We're going to get a single mini batch out of the whole batch of data on each of the cores And then we're going to pass those into the loss function and find the value of the loss as well as the gradients of the loss function So that means because again grad does the derivative with respect to the first parameter Which means we're going to have derivatives with respect to W, the weight and the bias Next thing we're going to do is we're going to combine the gradients across all of the TPU cores And this is arguably one of the most important, the most important operation in this whole training pipeline Basically you take the gradients which are now different across each of the cores And they're different because each core got a different portion of the data set Whoops, I'm not a robot, I assure you So basically you're going to do a mean averaging across all of the gradients And finally each single core will have the same gradients because we do this p mean operation As a consequence because we also have the same parameters, that means that each core will have the parameters in sync So all of the parameters are initially the same, the gradients are the same which means that the updated parameters across each of the cores are also the same So that's cool We did the same thing for loss, we just averaged loss across the 8 cores And finally we applied the stochastic gradient descent like update rule Again we have pi tree here, pi tree here, we just do the usual from the current value of the parameter We subtract the learning rate times the gradient and that's how we get our new parameters Finally we're going to return those parameters as well as the loss value and that's it I made a short note here which may be important if you want to train with some more complex optimizers and not just with SGD So if we had an optimizer this is how the pattern would look like, we pass in the grads and we pass in the optimizer state And the optimizer because of remember our pattern is going to return the new optimizer state and the new basically updated version of gradients And then we can just apply the updates instead of grads here and that's it So it's very simple just another additional line following this stateless pattern that we saw in the beginning of this video Nice, let me now go and run the training So this is our training, let me run it and let's see what we get Okay it's working, we're printing some values, we're going to debug and understand, analyze what's going on in the console here But before that let's just plot the results Okay so let me run this and we see that we have a line that's fit to this particular data set we artificially built Now let's analyze the output in the console, what happened here First thing I want you to notice here is that replicated parameters because as I said those are going to stay on the cores They should be sharded, sharded device type so let's see whether that's the case So we can see here after printing after the 0th epoch we're going to have this is a class of blah blah blah sharded device array Which means that indeed our replicated parameters are on the cores themselves Loss is as well a sharded device array But the data is as you can see here I'm just printing the type of the data We can see it's a non-py array which means it's locally stored on the CPU Because we do not want to keep the data on the core that wouldn't make any sense especially for bigger data sets Finally what I'm printing here is the loss shape You can see it's an 8 dimensional vector because again it's sharded across different cores And we can see that the loss goes down as we are advancing here It gets into situation after already a couple hundred steps here but yeah And that's pretty much it we trained our first machine learning model in parallel on 8th epoch cores And as you saw it was very simple to do that Now you're almost ready to be able to train whatever type of machine learning model There's a couple more questions you have to get an answer to So how would you go about doing transfer learning injects? Relatedly how do we free certain layers and fine tune other layers? So we need to figure out how to do that, how to do stop gradients And finally how do we get per sample gradients instead of the usual per batch gradients I'm going to tell you all about that in this last section of this tutorial In order to demonstrate how stop gradients work in Jax I'm going to use a couple of examples So first things first Jax uses this JaxLax stop gradient primitive for this purpose, for stopping gradients Let me quickly go through this TD0 update rule here If you're not familiar with RL this may be super confusing but you basically you do not need to know what this is You just need to know some of the details which I'm going to explain to you right now So the whole point is we are trying to learn certain value functions So that value function in our example is going to be a simple linear value function So it's parameterized by theta and we are basically doing dot product with the state vector in order to get the value function So in order to find the gradients of theta we have to do this expression So basically we have a reward here at the present timestamp plus the value at the current state minus the value at the previous state And we multiply that by the gradient of the value function in the previous state So that's a mouthful, the whole point is the update is not a gradient of any loss function However if we reformulate this above expression into this thing And if we have a way to stop the gradients of this target here of this value at the present state We'll be able by doing a gradient of this pseudo loss function to get exactly this expression Again as you can see here if we ignore the dependency of the target on the parameter theta for our value function We're going to get exactly this So let's see why that is, if we do a gradient with respect to theta, 2 goes in front here We'll have this whole expression and then if we do a derivative with respect to theta And this does not depend on theta, that means we're just going to end up with the derivative of this previous value function at the previous state So that's exactly this expression here So just some basic rules of differentiation at the end, that's all So let's finally see how we can implement that We have the value function as I already explained, simple dot product We have some initial values for the theta We have some like, so this t1 just means timestamp minus 1, i.e. the previous timestamp We have the reward, we have the current state here And finally we have the td loss We calculate the value function at the previous state We calculate the target and this is the whole point here We just apply the stop gradient to target and this is going to be our loss And now when we do a gradient of that loss with respect to theta Because theta is the first argument here And remember again, gradient by default will do a derivative with respect to the first parameter We get the correct value of delta theta here, i.e. the gradients of our model I'm going to run this, although the number will not mean a lot And I don't think it's that important to go and analyze this further You get the point here The second example is this straight through estimator And I've seen this thing used, for example, in the VQVAE paper And I also covered that paper in one of my previous videos, you can check it out I've basically pasted the link here But the whole point is, if you have a certain non-differentiable function in the middle of your neural network, for example For example, this is a simple non-differentiable function You take the x, which is a float, and if you apply a rounding, that's simply not differentiable So now we can reformulate this non-differentiable function into this straight through f So basically, what we have done here is, if you take a look at this So x plus f of x minus x Basically that means in the forward pass we're going to have f of x Which means these two are equivalent in the forward pass On the other hand, if we do a derivative, if we find a gradient of this function and a gradient of this function They are going to be different So let me just take some simple example here So x is 5.6, just an arbitrary number We're going to print the value in the forward pass So we're going to see that those are going to be the same So 6 and 6, because when we round 5.6 we get 6 So that's clear Whereas, because this f is non-differentiable If we try to do a gradient at this point x, or whatever x is We get 0 all the time On the other hand, this straight through gradient is going to give us 1 all the time You can see why that is, because we have a stop gradient here So a derivative of this thing with respect to x is just going to be 1 Because, yeah, I mean, if a function is equal to x, then basically a gradient is 1 So why may this be useful? If you have an automatic differentiation in the reverse mode You're basically, if you have 1 here If a certain part of your whole function has a gradient 1 By multiplying the gradients that came from the deeper levels by 1 You're just going to pass them through, that's why it's called straight through You're just going to copy them to the second part of the network or whatever ML model you're using And that's why this thing is kind of cool But the whole point is, again, just be aware that we are using this stop gradient to achieve this functionality Let's see a couple more advanced things So, calculating per sample gradients And let me just read this through In many frameworks, PyTorch, TensorFlow, Theano, it is often not trivial to compute per example gradients Because the library directly accumulates the gradient over the batch Naive workarounds, such as computing a separate loss per example and then aggregating the resulting gradients Are typically very inefficient So you can go around this problem by taking a single sample, so for example, maybe a single image Passing it through, calculating the gradients for that single image And then doing that for the whole batch in a kind of a for loop and then aggregate the results Aggregate those gradients, but it's super expensive as you can imagine Jax basically has direct support for this type of a thing and it's very efficient So if we create a contrived example here by just stacking the state in the t-minus-1 timestamp And we batch the rewards, we batch the current state So basically using VMAP is going to allow us to get these per sample gradients Because remember, TD loss does not accept a batch of data It accepts a single example, so a single state, a single reward, and a single state here And we have arguments theta here So if we map it, if we wrap it inside of a VMAP using in-axis none, which means theta The arguments of our value, so the parameters of our value function are going to be broadcasted And then we say 000, which means the first zero dimension of these batches here is the batch dimension If we do that, we're going to trivially get the gradients for every single example And they are the same because we batch the same states here, so that's kind of trivial Wrapping up this tutorial, which is already quite lengthy In a summary, Jack's auto-differentiation engine is super powerful, you can do various stuff You can even define custom derivative rules if you have some numerical problems That's now already super advanced, you can check out the docs here I particularly linked this Autodeaf cookbook article, which you should read if you want to know all of the nitty-gritty details of the Autodeaf engine I'm going to end up by showing a simple example of this MAMBL This is this model agnostic metal learning algorithm And you can see it's fairly complicated, you can see the math here And the whole point is, this is how simply you can implement this MAMBL algorithm in Jack's It looks very much like mathematics, and that's why I think Jack's is super, super powerful and useful for researchers Not that beginner-friendly as I already mentioned in my previous video But even for beginners, if you go across a couple of hurdles, I think it's going to pay off Let me briefly go over MAMBL So the whole point here is, imagine you have a couple of tasks And while training, you do not want to minimize the loss on any particular task You want to minimize the loss, the potential losses on all of those tasks For example, maybe going in this direction will minimize the loss for the task 1 Going in this one will minimize the loss for the task 2 And similarly, going in this one will minimize the loss for task 3 That would lead us to have either theta 1, that's this These thetas are going to minimize the loss on the task 1 Similarly, this theta here, theta 2, is going to minimize the loss on this task 2 And similarly for theta 3 But we actually move here, which is somewhere in between all of those thetas And that's how MAMBL works And now, briefly going over this metal loss function We calculate the gradients here, and then we use the gradients to update Using SGD, we update the parameters here And we pass that into the loss function And then we take the gradients of that loss function It's kind of convoluted, I'm not going to try This is not the video about MAMBL, but I want you to just kind of appreciate how expressive Jax actually is Having said that, hopefully you liked this video If you did, share it out with your friends Consider subscribing to this channel And also join the Discord community, that's where the party is See you next time, bye bye!
[{"start": 0.0, "end": 6.98, "text": " What's cracking guys? This is the second video in the machine learning with Jack's series of tutorials and in this one"}, {"start": 6.98, "end": 10.620000000000001, "text": " We're going to learn how to build a fully fledged neural networks"}, {"start": 10.620000000000001, "end": 17.240000000000002, "text": " So in the last video we went from zero to hero and in this one we're going to go from hero to hero Pro Plus"}, {"start": 18.2, "end": 22.36, "text": " Yeah, anyways before we even start let me show you this thing"}, {"start": 22.36, "end": 29.76, "text": " So I open sourced a repo that's gonna called get started with Jack's that basically contains aside from these tutorials"}, {"start": 29.76, "end": 34.480000000000004, "text": " I'll be making other useful like content that I personally used and found useful"}, {"start": 34.480000000000004, "end": 42.86, "text": " So things like various YouTube videos blogs, etc. So and the second thing I want to I want to tell you is that you can"}, {"start": 43.800000000000004, "end": 45.86, "text": " Basically open up these notebooks"}, {"start": 45.86, "end": 50.7, "text": " I'm using in these videos in Google collab and you can easily play with them"}, {"start": 50.7, "end": 54.1, "text": " So you can just open it up and the Python environment is already set up"}, {"start": 54.1, "end": 58.28, "text": " So Jackson everything else is pre-installed so you can just kind of run this and everything will work"}, {"start": 58.28, "end": 61.300000000000004, "text": " Just run anyway, and this is going to to work out of the box"}, {"start": 61.4, "end": 67.56, "text": " So basically this like the notebook for this second video is already checked in as the time"}, {"start": 67.56, "end": 71.72, "text": " I'm recording this so you'll be able to open it up and maybe"}, {"start": 72.12, "end": 74.32, "text": " Tweak the code as you are watching the video"}, {"start": 74.32, "end": 77.48, "text": " So I think the best approach to get the most out of these videos is to"}, {"start": 77.72, "end": 84.24000000000001, "text": " Tweak the code yourself play a bit and and see where you can learn something something interesting awesome with that out of the way"}, {"start": 84.24000000000001, "end": 86.24000000000001, "text": " Let's continue with the video"}, {"start": 86.24, "end": 90.32, "text": " So let's first import the packages. We saw this in the last one"}, {"start": 90.32, "end": 93.64, "text": " So we were just importing Jack's the NumPy API of Jack's"}, {"start": 94.36, "end": 99.64, "text": " Transform functions that we saw in the last video. So grad jet V map in this one"}, {"start": 99.64, "end": 104.91999999999999, "text": " We're going to learn about P map which helps us parallelize the computation across multiple devices"}, {"start": 104.91999999999999, "end": 110.28, "text": " Which is very useful as you can imagine and yeah other necessary libraries here"}, {"start": 110.28, "end": 116.16, "text": " So the goal as I said will be to learn how to build complex more complex machine learning models"}, {"start": 116.16, "end": 121.16, "text": " Parallelize the computation across multiple devices and yeah, basically after this video"}, {"start": 121.16, "end": 124.1, "text": " You should be able to to get started and build your own models"}, {"start": 125.08, "end": 128.51999999999998, "text": " In order to get to the point where we can build neural networks"}, {"start": 128.51999999999998, "end": 132.85999999999999, "text": " There is a couple of things we have to you to first solve the first one is the problem of state"}, {"start": 132.85999999999999, "end": 135.8, "text": " So when I say that I mean the following so Jack's is basically"}, {"start": 136.72, "end": 140.28, "text": " uses the functional program programming paradigm, which means"}, {"start": 140.28, "end": 147.28, "text": " We need pure functions, which means Jack's does not love state in the sense that you cannot use your your usual"}, {"start": 147.8, "end": 154.6, "text": " Object-oriented programming type of stateful classes or functions. So we're gonna see a way around it right now"}, {"start": 155.72, "end": 161.68, "text": " First short reminder from the first video. We saw that impure functions are problematic"}, {"start": 161.76, "end": 168.8, "text": " So we saw this global variable G and we can treat it as a state for the purpose of this of this of this section"}, {"start": 168.8, "end": 175.36, "text": " Basically what this impure uses global function does it accesses G which is external state"}, {"start": 175.36, "end": 182.64000000000001, "text": " Which is forbidden because this makes the function impure and what happens now is when the JIT when we call JIT on this function"}, {"start": 182.92000000000002, "end": 188.04000000000002, "text": " It's going to cache the current value of G because that's a side effect. So it's gonna just cache the value 0"}, {"start": 188.24, "end": 191.52, "text": " So that means the first time we call it is gonna is gonna run correctly"}, {"start": 191.52, "end": 195.8, "text": " But the second time we call it it's gonna call we're going to call the cached version of the function"}, {"start": 195.8, "end": 199.98000000000002, "text": " Which means 5 will be added to 0 and not 10"}, {"start": 199.98000000000002, "end": 203.48000000000002, "text": " So in the meantime in the meanwhile, we've updated G to 10"}, {"start": 203.48000000000002, "end": 209.38000000000002, "text": " So the state is updated to 10, but that will not be reflected because of the because we have impure function"}, {"start": 209.38000000000002, "end": 213.42000000000002, "text": " So this is the the trick you have to be aware of and as you can see here"}, {"start": 213.42000000000002, "end": 217.84, "text": " We get the correct result for the first call, but we get the incorrect result for the second call"}, {"start": 217.84, "end": 224.08, "text": " So that's 5 and we should have gotten I guess 5 plus 10. That's 15, right? Okay, that's the first step"}, {"start": 224.08, "end": 230.70000000000002, "text": " Secondly, we saw this pattern already. So basically we saw that the pseudo random number generator in"}, {"start": 231.04000000000002, "end": 238.28, "text": " Jax is not stateful in contrast to two numPy's and we saw this pattern where we create the state the initial state"}, {"start": 238.32000000000002, "end": 242.88000000000002, "text": " By just passing in the seed and then in order to get the the next state"}, {"start": 242.88000000000002, "end": 248.76000000000002, "text": " We just call this split function and we pass we externally pass the state. So that's that's a big difference compared to numPy"}, {"start": 248.76, "end": 253.44, "text": " By pseudo random number generator. So I basically want you to remember this pattern"}, {"start": 253.44, "end": 260.14, "text": " So we are inputting the state inside of the function. So split is not saving state internally is just taking it as an input"}, {"start": 260.14, "end": 264.84, "text": " It's manipulating it somehow in a pure fashion and returning back some novel states"}, {"start": 265.32, "end": 268.26, "text": " Whereas just recall that we were using key and sub key"}, {"start": 269.0, "end": 271.86, "text": " Terminology previously. I'm just want to stress out that we are"}, {"start": 272.52, "end": 277.84, "text": " Basically, I want to stress out this pattern. So that's why I'm using this new terminology. But yeah, you get the point hopefully"}, {"start": 277.84, "end": 283.44, "text": " Okay. So now let's see a way around this and if I haven't made this already clear enough"}, {"start": 284.03999999999996, "end": 289.67999999999995, "text": " Let me just kind of motivate where we're doing this. Basically if you're trying to create a neural network, for example"}, {"start": 290.4, "end": 294.34, "text": " Like there is a lot of stateful code when you're trying to build neural networks"}, {"start": 294.34, "end": 297.12, "text": " Especially in other frameworks such as PyTorch or TensorFlow"}, {"start": 297.28, "end": 304.62, "text": " Basically, our model parameters are stateful your optimizer is stateful atom is for example storing internally a certain"}, {"start": 304.62, "end": 310.58, "text": " Statistics so that means it's stateful and also we have a layer such as batch norm, which are stateful"}, {"start": 310.58, "end": 315.9, "text": " Yeah in summary and any useful program will will have to have some some kind of state and we just need to find"}, {"start": 316.3, "end": 322.34000000000003, "text": " Alternative way how we can handle it. So in this first example, we have a stateful class. It's called counter"}, {"start": 322.58, "end": 329.02, "text": " We have this state called like n which is basically the internal count and we have the function count"}, {"start": 329.02, "end": 335.64, "text": " We just increments the this this state by one and returns the current value of that state"}, {"start": 335.94, "end": 341.9, "text": " Finally, once we call reset, we just set the state to zero and that I guess all makes sense here"}, {"start": 341.9, "end": 348.09999999999997, "text": " We just instantiate the an object of this class and then we we call we call three times"}, {"start": 348.09999999999997, "end": 351.02, "text": " We call this this count function and as expected"}, {"start": 352.65999999999997, "end": 355.02, "text": " Hopefully we'll get one two three and here they are"}, {"start": 355.02, "end": 360.74, "text": " Okay, so that all works like a charm and this is something you're used to now the problem arises here"}, {"start": 361.02, "end": 368.76, "text": " When you try to get the function, so let's reset the counter and let's so that means the state is now internally set to zero"}, {"start": 368.76, "end": 373.15999999999997, "text": " And now let's try and get the counter count function. So let's try and get this function here"}, {"start": 373.15999999999997, "end": 375.09999999999997, "text": " Now, let's see what happens with it"}, {"start": 375.09999999999997, "end": 380.5, "text": " Once we did it and we try to again run the same type of like a for loop here"}, {"start": 380.5, "end": 384.22, "text": " The results are kind of surprising maybe or maybe not even familiar with jet"}, {"start": 384.22, "end": 390.66, "text": " Basically we get all ones and the reason is as I said as this count function is not pure because we are accessing"}, {"start": 390.66, "end": 396.42, "text": " We are basically writing to an external state as you can see here. So self dot n is a state"}, {"start": 396.42, "end": 400.5, "text": " Basically, it's not passed in as an argument. We are we are just modifying it externally"}, {"start": 400.5, "end": 405.5, "text": " So that's bad. And what will happen is so once we call reset here"}, {"start": 405.5, "end": 411.5, "text": " And we'll be set to zero and the first time so once we get the first time we call a jitted function"}, {"start": 411.5, "end": 416.5, "text": " We're gonna do the tracing if you recall from the last video and that means that basically on the first call"}, {"start": 417.5, "end": 424.5, "text": " The function will return one and because we have a side effect is just gonna learn how to cache number one and that's it"}, {"start": 424.5, "end": 431.5, "text": " Now, let me just kind of try and clarify that if we were to use make Jack spur to make this more clear"}, {"start": 431.5, "end": 439.5, "text": " Let me run this and see what the like underlying Jack's expression Jack's expression is and we can see it learn to just return"}, {"start": 439.5, "end": 446.5, "text": " Like one so basically the function the cache function has like the whole body of the function is something like this return one"}, {"start": 446.5, "end": 452.5, "text": " That's it. And that's why we're getting ones even though we are we are expecting to get increments"}, {"start": 452.5, "end": 457.5, "text": " Okay, just to make this a bit more clear if we were to call counter count function here"}, {"start": 457.5, "end": 464.5, "text": " We'd expect the internal state to be modified to like one and then once we call jet"}, {"start": 464.5, "end": 466.5, "text": " It's gonna be returning twos now, hopefully"}, {"start": 466.5, "end": 472.5, "text": " So if we correctly understand this if I run this we get twos and the reason is again"}, {"start": 472.5, "end": 478.5, "text": " It's basically tracing this function the first time it's racing yet. The value was already set to one"}, {"start": 478.5, "end": 483.5, "text": " We add additionally one so we have two and we turn two and basically because we don't have any inputs"}, {"start": 483.5, "end": 489.5, "text": " It's just gonna learn how to cache and return those number two and again if if I were to do the same thing here"}, {"start": 489.5, "end": 497.5, "text": " So count and call the make Jack Jack spur. We're hopefully gonna get like just two"}, {"start": 497.5, "end": 503.5, "text": " Instead of one here. Let's see and that's correct. Okay, we got two here"}, {"start": 503.5, "end": 508.5, "text": " In any case, we have to find a way to avoid like creating these in pure function"}, {"start": 508.5, "end": 513.5, "text": " So let's see what the solution is we're going to define this counter state"}, {"start": 513.5, "end": 519.5, "text": " Like instead of defining a complex class, we're just gonna have like a simple primitive data type here. So integer"}, {"start": 519.5, "end": 522.5, "text": " So this version two of the counter is gonna work the following way"}, {"start": 522.5, "end": 529.5, "text": " So we're going to pass in the counter state and we're going to return back the output as well as the state"}, {"start": 529.5, "end": 534.5, "text": " Now this may be a bit confusing m plus one where we're turning a couple of m plus one two times"}, {"start": 534.5, "end": 542.5, "text": " I mean doesn't make any sense right? So the thing is in general this thing here will be state and this thing here will be some function of the state"}, {"start": 542.5, "end": 547.5, "text": " So here in the example of this simple counter the output as well as the state are the same thing"}, {"start": 547.5, "end": 551.5, "text": " So yeah, we're just returning m plus one, but yeah doesn't doesn't matter"}, {"start": 551.5, "end": 556.5, "text": " The whole point is we are passing in the state and we are processing it and then returning it back"}, {"start": 556.5, "end": 561.5, "text": " In the case of reset function, we're just returning back the state zero and that's it"}, {"start": 561.5, "end": 568.5, "text": " So here we instantiate the counter here. We call the reset and as you can see here compared to the version one of the counter"}, {"start": 568.5, "end": 571.5, "text": " We are not internally modifying the state. We are just"}, {"start": 571.5, "end": 576.5, "text": " returning back the state here externally and here we can see the pattern"}, {"start": 576.5, "end": 583.5, "text": " So we've called the counter count we pass in the state and we get the state and the value we print the value and we keep doing this"}, {"start": 583.5, "end": 586.5, "text": " And so yeah, I guess this should already be familiar"}, {"start": 586.5, "end": 592.5, "text": " Basically Jack's like run random number generators use this very same pattern"}, {"start": 592.5, "end": 597.5, "text": " So this should work like as expected now and as we can see we get one two three"}, {"start": 597.5, "end": 601.5, "text": " But but now let's see whether after we did the function this thing will work or not"}, {"start": 601.5, "end": 605.5, "text": " Let's call reset again. We get the initial state"}, {"start": 605.5, "end": 609.5, "text": " Let's get the function and let's try and run this"}, {"start": 609.5, "end": 612.5, "text": " And as expected we're getting correct results"}, {"start": 612.5, "end": 620.5, "text": " Even with this jaded version of the function because we are externally manipulating the state and not creating impure functions"}, {"start": 620.5, "end": 628.5, "text": " And yep, this is working as expected and in summary, let's see the pattern we we use to convert a stateful function into a stateless one"}, {"start": 628.5, "end": 633.5, "text": " So basically we had the state which was just number n in the case of a simple counter"}, {"start": 633.5, "end": 638.5, "text": " We have this stateful method. So the the especially like in particular we had this method called count"}, {"start": 638.5, "end": 643.5, "text": " We are passing some arguments in general case and we get back the output"}, {"start": 643.5, "end": 648.5, "text": " Now we converted this and instead of saving the state inside of the class"}, {"start": 648.5, "end": 655.5, "text": " We're passing the state inside here as an input and we return back the top also the output and the state"}, {"start": 655.5, "end": 662.5, "text": " And this pattern can be applied for to any a stateful class and to get a stateless one and that's cool"}, {"start": 662.5, "end": 668.5, "text": " Okay, that's that was step one towards building like a fully fledged neural networks"}, {"start": 668.5, "end": 672.5, "text": " Now the second step is we still need to find a way how to deal with gradients"}, {"start": 672.5, "end": 677.5, "text": " And you may wonder why we have to like learn anything new to deal with gradients"}, {"start": 677.5, "end": 682.5, "text": " Well, we already know about the grad function. Well, the thing is with our current knowledge"}, {"start": 682.5, "end": 686.5, "text": " We would go something like this. So imagine we have some loss function"}, {"start": 686.5, "end": 692.5, "text": " F and we have like the parameters of our neural network. So x y z and w"}, {"start": 692.5, "end": 698.5, "text": " So we basically only have four parameters for now and we do something like x squared plus y squared plus z squared"}, {"start": 698.5, "end": 703.5, "text": " Blah blah blah. So basically we get this is a L2 you can treat this as L2 loss"}, {"start": 703.5, "end": 708.5, "text": " So minimizing this loss would cause the weights to go closer to get closer to zero"}, {"start": 708.5, "end": 712.5, "text": " And now the problem here is if we had something like GPT-3 model"}, {"start": 712.5, "end": 718.5, "text": " Which has 175 billion parameters and some newer language models are even bigger"}, {"start": 718.5, "end": 725.5, "text": " Basically, you'd have to write 175 billion like terms here minus four that we already have here"}, {"start": 725.5, "end": 730.5, "text": " And that sucks, right? So the problem does not stop there"}, {"start": 730.5, "end": 736.5, "text": " Basically, if we now try to find the gradients because we want to do an update of our neural network parameters"}, {"start": 736.5, "end": 743.5, "text": " We'd have to get dfdx. So we have to find derivatives with respect to every single one of those arguments"}, {"start": 743.5, "end": 752.5, "text": " And grad as you recall usually by default returns the like a derivative with respect to only the first argument"}, {"start": 752.5, "end": 757.5, "text": " So that will be dfdx. So we have to explicitly pass this list of integers"}, {"start": 757.5, "end": 763.5, "text": " And in the case of 175 billion would have quite a lot of these, right?"}, {"start": 763.5, "end": 770.5, "text": " So this will not scale and yeah, let me just run this. This is going to work as expected"}, {"start": 770.5, "end": 773.5, "text": " But as I said, it's bothersome and it does not scale"}, {"start": 773.5, "end": 782.5, "text": " So we had twos because if we take argument dfdx is going to be 2x, 2y, 2z in the case of dfdz, etc."}, {"start": 782.5, "end": 787.5, "text": " So basically because we input all once as the input we get twos for every single gradient"}, {"start": 787.5, "end": 790.5, "text": " So that makes sense. Now again, the problem does not stop here"}, {"start": 790.5, "end": 798.5, "text": " If you want to do an update would have to have 175 billion lines doing these types of SGD updates"}, {"start": 798.5, "end": 802.5, "text": " So basically learning rate gradients and we update our weights here"}, {"start": 802.5, "end": 809.5, "text": " So there is obviously a better way. This is just a motivation for why we need the following solution"}, {"start": 809.5, "end": 811.5, "text": " And the solution is called a Pytree"}, {"start": 811.5, "end": 818.5, "text": " So in general we are used to wrapping our model parameters in some more complex data structures dictionaries or whatever"}, {"start": 818.5, "end": 828.5, "text": " And Pytree is the Jack's solution to how to basically find gradients of arbitrarily nested data structures"}, {"start": 828.5, "end": 834.5, "text": " And let's see a simple example here. So this is just a contrived example obviously"}, {"start": 834.5, "end": 840.5, "text": " You can treat it as some type of model parameters, very esoteric but yes, nonetheless"}, {"start": 840.5, "end": 847.5, "text": " Just stick with me here. We're going to first learn what Pytree is and then we're going to see a real example"}, {"start": 847.5, "end": 853.5, "text": " We're going to train like a multi-layer perceptron really soon. Let's see how Pytree is defined in the official docs"}, {"start": 853.5, "end": 860.5, "text": " So if I open up the docs here, if I go to tutorialjacks101 and find the working with Pytrees here"}, {"start": 860.5, "end": 867.5, "text": " Okay, we have the definition here. So a Pytree is a container of leaf elements and or more Pytrees"}, {"start": 867.5, "end": 872.5, "text": " So it has this recursive definition. Containers include lists, tuples and dictionaries"}, {"start": 872.5, "end": 878.5, "text": " So let's see how that maps to our example here. So basically we have a container here"}, {"start": 878.5, "end": 884.5, "text": " So this whole thing is a Pytree, but this Pytree contains another Pytree here"}, {"start": 884.5, "end": 891.5, "text": " And finally this Pytree has a leaf 1 and the leaf string this character a"}, {"start": 891.5, "end": 897.5, "text": " And finally this object is also considered a leaf. The second, this is also a Pytree, this whole tuple here"}, {"start": 897.5, "end": 904.5, "text": " And tuple I guess. This is a leaf, one will be a leaf, whereas this thing here will again be a Pytree"}, {"start": 904.5, "end": 910.5, "text": " And two will be a leaf and three will be a leaf. There is a handy way to analyze this and it's called this"}, {"start": 910.5, "end": 916.5, "text": " From Jax, just import this tree leaves function. We're going to iterate through these Pytrees and we're going to get the leaves"}, {"start": 916.5, "end": 921.5, "text": " And we're going to finally print them in a nice format here. So let's see what the results here are"}, {"start": 921.5, "end": 930.5, "text": " Basically this first Pytree has three leaves as we can see here. So one, character a and a simple empty object"}, {"start": 930.5, "end": 935.5, "text": " Let me see what else is interesting. So as I said this is 1, 2, 3, all of those will be leaves"}, {"start": 935.5, "end": 942.5, "text": " In the case of a dictionary, keys are not considered leaves. Keys are just like metadata of the Pytree node"}, {"start": 942.5, "end": 947.5, "text": " So we're just going to have 1, 2, 3, 4 and 5 as leaves as you can see here"}, {"start": 947.5, "end": 953.5, "text": " And interesting to notice here is finally this array is going to be considered a single leaf"}, {"start": 953.5, "end": 958.5, "text": " So we won't have 1, 2, 3 as separate leaves. This whole thing will be treated as a leaf"}, {"start": 958.5, "end": 962.5, "text": " So that's something to keep in mind. Okay, so how do we manipulate these trees?"}, {"start": 962.5, "end": 968.5, "text": " There's a couple of neat and handy functions. So again we have an example here of a simple Pytree"}, {"start": 968.5, "end": 973.5, "text": " And what we do is we call this tree map function and we just pass in the Pytree"}, {"start": 973.5, "end": 978.5, "text": " And we define a function that is going to manipulate and work on top of these"}, {"start": 978.5, "end": 985.5, "text": " It's going to process the leaves of the Pytree and basically what we expect is that every leaf will be multiplied by 2"}, {"start": 985.5, "end": 992.5, "text": " In the case of this simple lambda function. So let's run this and as expected we get everything is multiplied by 2"}, {"start": 992.5, "end": 1000.5, "text": " So that was in the case of these single argument functions. In the case when you have multiple"}, {"start": 1000.5, "end": 1005.5, "text": " Like when your function has multiple arguments, I'm just going to create a copy here"}, {"start": 1005.5, "end": 1012.5, "text": " So a reference basically and I'm going to do a simple addition and this time we're going to call treeMultiMap"}, {"start": 1012.5, "end": 1018.5, "text": " So we have 2 Pytrees as input. We manipulate them using treeMultiMap and as the output we get the Pytree again"}, {"start": 1018.5, "end": 1023.5, "text": " So if I run this we would expect same results and yep we have the same results here because"}, {"start": 1023.5, "end": 1028.5, "text": " Just adding two same structures is the same as doing multiply by 2, right?"}, {"start": 1028.5, "end": 1035.5, "text": " Okay, so finally worth noticing here is that the structures have to be the same"}, {"start": 1035.5, "end": 1042.5, "text": " So here the condition was trivially I guess true because I just made this a reference"}, {"start": 1042.5, "end": 1047.5, "text": " But in the case when you have two separate data structures you want to make sure they have the same structure"}, {"start": 1047.5, "end": 1053.5, "text": " Because if I artificially here introduce a change so I just do a deep copy of the original object"}, {"start": 1053.5, "end": 1058.5, "text": " And I just append this list that contains number 23 and if we try to run this"}, {"start": 1058.5, "end": 1063.5, "text": " We can see that list arityMismatch is reported as an error"}, {"start": 1063.5, "end": 1069.5, "text": " So yeah, keep that in mind and that was the last missing piece before we can train like a multilayer perceptron"}, {"start": 1069.5, "end": 1074.5, "text": " Hopefully you already know what MLP is if you don't let me just quickly show you some diagrams"}, {"start": 1074.5, "end": 1080.5, "text": " And whoops, okay, we are in a really bad part of the internet right now. Let me go to images here"}, {"start": 1080.5, "end": 1085.5, "text": " If this thing will load my internet is really really slow today"}, {"start": 1085.5, "end": 1090.5, "text": " Okay, so I'm gonna just type in neural network here and hopefully everything will work correctly"}, {"start": 1090.5, "end": 1097.5, "text": " And here we are. Here is a simple MLP basically you have your input data coming in from the left"}, {"start": 1097.5, "end": 1103.5, "text": " You're processing the inputs using these neurons. These are called hidden layers"}, {"start": 1103.5, "end": 1106.5, "text": " This is the output layer and outcomes to results"}, {"start": 1106.5, "end": 1111.5, "text": " So this thing may be like a binary classifier like a hot dog or not hot dog"}, {"start": 1111.5, "end": 1115.5, "text": " I don't know in any case that's MLP at least on a high level and"}, {"start": 1115.5, "end": 1121.5, "text": " Now, okay, we are we're ready to start training a toy multilayer perceptron model"}, {"start": 1121.5, "end": 1126.5, "text": " So we know how to handle state we know how to handle parameters"}, {"start": 1126.5, "end": 1133.5, "text": " I need to find gradients of our parameters which may be stored in certain types of nested data structures"}, {"start": 1133.5, "end": 1139.5, "text": " So that's everything we need to to train our first model in Jax, I guess"}, {"start": 1139.5, "end": 1144.5, "text": " So let's start here. We have this init MLP parameters function"}, {"start": 1144.5, "end": 1151.5, "text": " What it's going to do is return a pie tree object and here you can see we specify the type"}, {"start": 1151.5, "end": 1156.5, "text": " Basically the configuration of our MLP. So that's gonna it's gonna have three layers"}, {"start": 1156.5, "end": 1162.5, "text": " So the first one will map input. That's a scalar into 128 dimensional vector"}, {"start": 1162.5, "end": 1170.5, "text": " Then we're gonna have a mapping from 128 to 128 and finally we're gonna have a mapping from 128 to output scalar again"}, {"start": 1170.5, "end": 1178.5, "text": " So that's three layer deep MLP. So this function here just takes that list as an input as you can see here using this smart indexing"}, {"start": 1178.5, "end": 1181.5, "text": " We can just fetch tuples such as this"}, {"start": 1181.5, "end": 1190.5, "text": " So we first have the tuple 128 and we're gonna use that to form weights and biases of our neural of our like a linear layer"}, {"start": 1190.5, "end": 1195.5, "text": " And we're gonna wrap it up into the dictionary and just append to this list"}, {"start": 1195.5, "end": 1204.5, "text": " So this is gonna be a list of dictionaries which contain two keys weights and biases which then contain like a random meet randomly initialized matrices"}, {"start": 1204.5, "end": 1212.5, "text": " And you don't have to worry about this. This is just some specific type. I think gaming or something initialization doesn't matter"}, {"start": 1212.5, "end": 1218.5, "text": " And think to notice we're using numpy random here instead of using Jax PRNG doesn't matter for this simple example"}, {"start": 1218.5, "end": 1220.5, "text": " But just keep that in mind"}, {"start": 1220.5, "end": 1230.5, "text": " Anyways, so we get back the the the pie tree and because we have a pie tree we can nicely use the tree map to print certain like information about the pie tree"}, {"start": 1230.5, "end": 1236.5, "text": " And we can see we've done everything correctly. We have biases and weights. Everything is of correct shape"}, {"start": 1236.5, "end": 1247.5, "text": " In general, this is a good good practice to have to analyze the shapes of your if your data and model parameters and whatnot throughout the whole pipeline basically"}, {"start": 1247.5, "end": 1253.5, "text": " Okay, let's continue here. We're gonna have some update function here in order to train it"}, {"start": 1253.5, "end": 1261.5, "text": " This is the main part and I'm just going to before we even start digging into the training code. I'm just gonna comment on the training code"}, {"start": 1261.5, "end": 1267.5, "text": " I'm just gonna comment on this forward because we haven't defined it yet and I'm just gonna show you the input data"}, {"start": 1267.5, "end": 1270.5, "text": " So this is what we're trying to learn to regress"}, {"start": 1270.5, "end": 1281.5, "text": " So the MLP will have as an input a simple scalar and we're trying to regress the corresponding y which is simply x squared in this in this example here"}, {"start": 1281.5, "end": 1287.5, "text": " So now let's see how how this thing works. Let me uncomment this basically"}, {"start": 1287.5, "end": 1295.5, "text": " We'll have a couple of epochs. So 5000 epochs in this example. We're gonna call the update function and you see the the pattern again"}, {"start": 1295.5, "end": 1303.5, "text": " We are passing in the state which is the model parameters in this case and we return back the parameters again and we just iterate that and that's it"}, {"start": 1303.5, "end": 1308.5, "text": " That's going to update the parameters and finally we can regress using the forward method"}, {"start": 1308.5, "end": 1316.5, "text": " Let me show how those are defined. So update basically takes as an input parameters x and y. So that's the input and output data"}, {"start": 1316.5, "end": 1325.5, "text": " We calculate the gradients of the loss function with respect to basically this input data and the current model"}, {"start": 1325.5, "end": 1331.5, "text": " Just stick with me. I'm gonna analyze this in a bit more detail in a couple minutes. For now, let's just see what the loss function is"}, {"start": 1331.5, "end": 1342.5, "text": " It's a simple mean squared error loss. Basically we do here prediction. We find an error term that's basically predicted y minus the ground truth y"}, {"start": 1342.5, "end": 1351.5, "text": " We square it to get basically positive values error terms and then we do the mean across all of the data points and that's your MSC loss"}, {"start": 1351.5, "end": 1360.5, "text": " Forward is fairly simple. You have a this is basically your MLP feed forward pass. So remember params is a pie tree"}, {"start": 1360.5, "end": 1366.5, "text": " When we iterate layer is just a dictionary that contains weights and biases. So we take out the weights"}, {"start": 1366.5, "end": 1373.5, "text": " We do dot product with the input data and we just add the biases here and then apply the activation function rally in this particular case"}, {"start": 1373.5, "end": 1379.5, "text": " So there is a dichotomy we made here between hidden layers. Whoops hidden layers in the output layer"}, {"start": 1379.5, "end": 1386.5, "text": " The reason is the last layer will not apply an activation function. We're just going to return back the raw results"}, {"start": 1386.5, "end": 1391.5, "text": " That's the reason for this dichotomy. I guess we can rewrite this a bit nicer, but yeah, it works"}, {"start": 1391.5, "end": 1398.5, "text": " Okay that out of the way. Let's now focus on this part. We do the grad of our loss function and as you can recall"}, {"start": 1398.5, "end": 1407.5, "text": " Basically, it's gonna do a derivative with respect to the first parameter and since our first parameter is params and which is a pie tree"}, {"start": 1407.5, "end": 1412.5, "text": " Will basically get as the output the same structure as params just"}, {"start": 1412.5, "end": 1420.5, "text": " But instead of values of instead of the weights will have derivatives of those weights with respect to the loss function"}, {"start": 1420.5, "end": 1428.5, "text": " So that's the magic behind grad. Basically all of these built-in functions transform functions such as grad know how to handle pie trees"}, {"start": 1428.5, "end": 1437.5, "text": " And this is finally a concise way to solve this problem of figuring out the gradients"}, {"start": 1437.5, "end": 1445.5, "text": " Now that we have gradients we can use our loved tree multimap function and we just pass in parameters and gradients"}, {"start": 1445.5, "end": 1453.5, "text": " We have this so on every single leaf of these pie trees we're going to apply the following rule and that's from the weight"}, {"start": 1453.5, "end": 1460.5, "text": " Just subtract the learning rate times the gradient which is simple stochastic gradient descent update rule and that's it"}, {"start": 1460.5, "end": 1466.5, "text": " As the result again tree multimap as you recall it returns back the pie tree like structure again"}, {"start": 1466.5, "end": 1472.5, "text": " Which means we have the updated weights getting back as the output of update and we can see that here"}, {"start": 1472.5, "end": 1478.5, "text": " So params in params come out as well here. Now, let's try to run this and then we're going to play a bit with this code"}, {"start": 1478.5, "end": 1486.5, "text": " So let me run this cell. Let's see where I ran this one. Yeah, I did. So let me run this one and let's see what the results are"}, {"start": 1486.5, "end": 1493.5, "text": " Okay, this is the output we can see a pretty nice fit in this particular case"}, {"start": 1493.5, "end": 1502.5, "text": " We can now play with this function. We can try and see some other polynomials maybe X X cubed"}, {"start": 1502.5, "end": 1510.5, "text": " We can see it also manages to fit this curve as well. We can play and add some other types of polynomials something like this"}, {"start": 1510.5, "end": 1517.5, "text": " We just I just added plus X here and it's still be like you can see at the borders here because we have less data points"}, {"start": 1517.5, "end": 1523.5, "text": " The fittest has a bigger errors compared to where we have a dense sampling which is expected I guess right?"}, {"start": 1523.5, "end": 1530.5, "text": " Finally, let's try one more interesting example that sign us excess. Let's see what we get here"}, {"start": 1530.5, "end": 1537.5, "text": " Whoops, when you use jet numpy sign here, so this should work now"}, {"start": 1537.5, "end": 1549.5, "text": " And you can see the fit is way worse here. We can obviously improve that by increasing the capacity of the model or increasing the number of epochs"}, {"start": 1549.5, "end": 1558.5, "text": " But yeah, let's try some other frequency like 3 X and yeah, we have a much better fit here"}, {"start": 1558.5, "end": 1565.5, "text": " In any case you get the point. We trained our first neural network in Jax and it seems to be working"}, {"start": 1565.5, "end": 1573.5, "text": " So let me return back the original statement here. So that's a simple like a polynomial and now let's analyze a couple of things here"}, {"start": 1573.5, "end": 1578.5, "text": " So first, let's see the structure of grads. I'm gonna reduce this to one just a single epoch"}, {"start": 1578.5, "end": 1584.5, "text": " I'm gonna comment out the JIT part and I'm gonna print grads here"}, {"start": 1584.5, "end": 1591.5, "text": " Okay, so I'm gonna just print the gradients here so that we are confident that we get the same structure as params"}, {"start": 1591.5, "end": 1597.5, "text": " So I'm gonna run this cell. I'm gonna run this again. I guess the output is gonna be a bit unwieldy"}, {"start": 1597.5, "end": 1601.5, "text": " Okay, let me see what we got here"}, {"start": 1601.5, "end": 1610.5, "text": " Bunch of numbers, but you can basically see that we have a list of dictionaries where the keys are biases and weights"}, {"start": 1610.5, "end": 1617.5, "text": " And basically these here these numbers are derivatives and not the parameter values. That's the important part"}, {"start": 1617.5, "end": 1621.5, "text": " So that's this part. Let me just uncomment this, return back the JIT"}, {"start": 1621.5, "end": 1625.5, "text": " And I mean that's pretty much it. That's everything that goes into training an MLP"}, {"start": 1625.5, "end": 1632.5, "text": " As a quick mention here, you can notice that I called JIT on top of this high level update function"}, {"start": 1632.5, "end": 1640.5, "text": " Because it gives the XLA compiler much more freedom to optimize various functionalities inside of this update function"}, {"start": 1640.5, "end": 1651.5, "text": " You can obviously play and toggle off and on the JIT wrapper here and see how faster it is to train it with JIT compared to without JIT, I guess"}, {"start": 1651.5, "end": 1656.5, "text": " Okay, let's continue here. This was an example of a simple neural network"}, {"start": 1656.5, "end": 1662.5, "text": " But if we wanted to create our own neural network library, we'd need to create layers such as NNLinear"}, {"start": 1662.5, "end": 1665.5, "text": " That's just PyTorch syntax or conv layers, etc."}, {"start": 1665.5, "end": 1671.5, "text": " And in order to get those to work, we have to create our custom PyTrees"}, {"start": 1671.5, "end": 1679.5, "text": " And let's see why that is. So first imagine this is this my container object here is just a simple NNLinear layer"}, {"start": 1679.5, "end": 1684.5, "text": " So let's imagine that's a linear layer and now imagine we have a list of two linear layers here"}, {"start": 1684.5, "end": 1691.5, "text": " So just create an example PyTree here and now let's fetch the leaves and see what we get as the output"}, {"start": 1691.5, "end": 1698.5, "text": " And we would expect, I guess, like eight leaves here, right? But the thing is it's going to return only two leaves"}, {"start": 1698.5, "end": 1705.5, "text": " So we have here that we this whole PyTree only contains two leaves and those are this my container objects"}, {"start": 1705.5, "end": 1713.5, "text": " That's obviously going to be undesirable because TreeMap and those other functions, they go through the leaves and do certain updates"}, {"start": 1713.5, "end": 1722.5, "text": " And if we had a linear layer like this, we wouldn't be able to do a gradient, like stochastic gradient descent update on top of these structures"}, {"start": 1722.5, "end": 1730.5, "text": " So we have to find a way around this. For example, if we were to run this simple manipulation, so just x plus one"}, {"start": 1730.5, "end": 1735.5, "text": " We'd expect to just increment these numerical values, but this is not going to work"}, {"start": 1735.5, "end": 1740.5, "text": " And I guess it makes sense because if you know how TreeMap works, it's going to iterate over the leaves"}, {"start": 1740.5, "end": 1746.5, "text": " And the leaves are my container objects and we're trying to add one to my container, which doesn't make any sense"}, {"start": 1746.5, "end": 1755.5, "text": " So let's run this and we can see that we get this error, we're trying to unsupport operand blah blah plus for my container and integer"}, {"start": 1755.5, "end": 1761.5, "text": " So that makes sense, right? So what's the solution to this? It's simple, you just have to define two functions"}, {"start": 1761.5, "end": 1766.5, "text": " One is flatten, one is unflatten and then you have to register PyTree nodes"}, {"start": 1766.5, "end": 1772.5, "text": " So we registered the my container with these two associated functions, flatten and unflatten"}, {"start": 1772.5, "end": 1777.5, "text": " So what flatten does basically just take the ABC, so those were the values we had here, right?"}, {"start": 1777.5, "end": 1784.5, "text": " We'll just wrap that information into an iterable here and we are not going to want to have a name as the content"}, {"start": 1784.5, "end": 1789.5, "text": " We want to have it as a metadata and that's why we include it in this auxiliary data structure"}, {"start": 1789.5, "end": 1793.5, "text": " And we just return this tuple of like content and auxiliary data"}, {"start": 1793.5, "end": 1801.5, "text": " Finally, the unflatten just receives those same inputs in reverse order though because of I think some legacy problems they had"}, {"start": 1801.5, "end": 1806.5, "text": " And you basically just construct the my container and you return it back, that's it"}, {"start": 1806.5, "end": 1810.5, "text": " So after I run this, now the my container class should work as expected"}, {"start": 1810.5, "end": 1818.5, "text": " So if we were to run the tree leaves function again, let's run it and we have the leaves 1, 2, 3, 4, 5, 6 as expected"}, {"start": 1818.5, "end": 1827.5, "text": " Now this thing will work, so if we were to add plus 1, we'll just get I guess 2, 3, 4 should be the output here"}, {"start": 1827.5, "end": 1829.5, "text": " And yes, that is the case"}, {"start": 1829.5, "end": 1838.5, "text": " Basically tree leaves, i.e. the tree map iterated over the leaves which are just the numerical values here, these ones"}, {"start": 1838.5, "end": 1844.5, "text": " And it just added plus 1, returned a PyTree again and we just print the leaves here, that's it"}, {"start": 1844.5, "end": 1854.5, "text": " There is one more gotcha I want to explain, tell you about and that's that we can oftentimes mistake nodes for leaves"}, {"start": 1854.5, "end": 1859.5, "text": " And when I say that, here's what I mean, so let's create a zero PyTree"}, {"start": 1859.5, "end": 1865.5, "text": " So basically we just have two nodes here, two leaves that contain zeros with these shapes"}, {"start": 1865.5, "end": 1873.5, "text": " And let's try to create the same type of PyTree that will have ones instead of zeros"}, {"start": 1873.5, "end": 1881.5, "text": " So what we do here is we create this intermediate PyTree, we're going to iterate through the zeros tree, through all of the leaves"}, {"start": 1881.5, "end": 1885.5, "text": " We're gonna take the shape and that's gonna be our new PyTree"}, {"start": 1885.5, "end": 1893.5, "text": " And then we can iterate through the shapes and apply the ones function using tree map to get like the tree of ones"}, {"start": 1893.5, "end": 1895.5, "text": " And let's see whether this works or not"}, {"start": 1895.5, "end": 1899.5, "text": " So if I print it like this, here is what we get"}, {"start": 1899.5, "end": 1904.5, "text": " We get correctly 2, 3, so we have two rows here and three columns, everything is fine"}, {"start": 1904.5, "end": 1907.5, "text": " We have three rows here and four columns, everything is fine here"}, {"start": 1907.5, "end": 1912.5, "text": " It seems that we have a correct result here as well, so we have 2, 3 shape and we have 3, 4"}, {"start": 1912.5, "end": 1915.5, "text": " So that's the intermediate array, so this shapes tree here"}, {"start": 1915.5, "end": 1919.5, "text": " And finally we have some weird result here, so what happened?"}, {"start": 1919.5, "end": 1926.5, "text": " So the thing that happened is, if you can recall, basically this here, this whole thing is a PyTree"}, {"start": 1926.5, "end": 1932.5, "text": " But this tuple is a PyTree as well, it's not a leaf, that means that leaf is 2 and 3 and 3 and 4"}, {"start": 1932.5, "end": 1942.5, "text": " Which means we're going to apply ones on top of all of these leaves and that's why we get two ones here and then three ones here and then three ones here and four ones here"}, {"start": 1942.5, "end": 1948.5, "text": " So simple work around here is to create, so we wanna make sure that this is a leaf and not a PyTree"}, {"start": 1948.5, "end": 1955.5, "text": " We can just basically wrap this x.shape with jmp array and this should now work"}, {"start": 1955.5, "end": 1961.5, "text": " Let me see, and yep, it works, so basically the reason is this is now a leaf, 2, 3 are a leaf"}, {"start": 1961.5, "end": 1965.5, "text": " And we pass that to ones and that's why we get correct results here, awesome"}, {"start": 1965.5, "end": 1972.5, "text": " Now we can design custom neural network layers and we can train end to end our neural networks"}, {"start": 1972.5, "end": 1976.5, "text": " Now what if we wanna train really, really big neural networks?"}, {"start": 1976.5, "end": 1983.5, "text": " Well, in that case we have to parallelize our computation and that's where this next section comes into play"}, {"start": 1983.5, "end": 1990.5, "text": " So I'm gonna introduce this pmap transform function and as we'll soon see it's very, very similar to vmap"}, {"start": 1990.5, "end": 1994.5, "text": " But before I start showing you how to use pmap we have to do a couple of things"}, {"start": 1994.5, "end": 2000.5, "text": " So the first thing is to set up, we have to use TPUs here for this section"}, {"start": 2000.5, "end": 2012.5, "text": " So I'm gonna set up hardware accelerator to a TPU, I'm gonna click save and hopefully after clicking reconnect it's gonna allocate a couple of TPUs for me"}, {"start": 2012.5, "end": 2024.5, "text": " Awesome, this worked, so now I just have to rerun the import statements because interrupting the runtime erases all of the whole state of the colab"}, {"start": 2024.5, "end": 2033.5, "text": " So we have to rerun this cell and then let's go back, I can open it up like this, just jump to parallelism injects in the table of contents"}, {"start": 2033.5, "end": 2047.5, "text": " And we are here, okay, so we're first going to just call this setup function and that's gonna set up TPUs in the background and hopefully we'll get 8 cores at disposal"}, {"start": 2047.5, "end": 2054.5, "text": " Okay, running this cell may take a while but finally we got 8 cores here as you can see, so 0 through 7"}, {"start": 2054.5, "end": 2065.5, "text": " Now I'm gonna use this convolve function as a running example and we'll be parallelizing this particular piece of computation across multiple cores"}, {"start": 2065.5, "end": 2071.5, "text": " Although we'll later see training annual network, actually a simple ML model across parallel cores"}, {"start": 2071.5, "end": 2078.5, "text": " So let's start with convolve, so I'm gonna define a simple signal here, so we have a signal X and we have a kernel W"}, {"start": 2078.5, "end": 2090.5, "text": " And the convolve function works as expected, basically we are sliding W across X, we are doing dot products and we are storing the computations in this list"}, {"start": 2090.5, "end": 2096.5, "text": " And then we're just converting the list into like a Jax array, that's it, just a simple 1D convolution"}, {"start": 2096.5, "end": 2108.5, "text": " So if we run here, let's see the results and we get the correct results, so basically imagine, let me print X and let me print W as well here"}, {"start": 2108.5, "end": 2119.5, "text": " So that may make a bit more sense, so what we do is in the first computation we'll have 2 times 0, that's 0, 1 times 3, that's 3"}, {"start": 2119.5, "end": 2127.5, "text": " And 2 times 4, that's 8, 8 plus 3 is 11 and that's why we get 11 as the first output, that's all nice and clear"}, {"start": 2127.5, "end": 2136.5, "text": " So now let's do the same thing, but this time let's parallelize the computation, so first this is just going to grab the number of devices on disposal"}, {"start": 2136.5, "end": 2148.5, "text": " We'll have 8 cores and then we're just gonna create a batch to simulate heavy loads, so some bigger computation compared to just having a single signal X and kernel W"}, {"start": 2148.5, "end": 2162.5, "text": " So let's run this and basically we have 8 devices and that's why we end up having a batch of 8 times 5, where 5 is basically the signal length and we're just replicating"}, {"start": 2162.5, "end": 2175.5, "text": " So here we have unique data for all of these 8 elements in the batch, on the other hand we're just duplicating the kernels, we're not creating novel kernels for each of the examples"}, {"start": 2175.5, "end": 2184.5, "text": " Ok, that's it, let's now first parallelize the computation across the memory of the host device, that means we'll be using a single core"}, {"start": 2184.5, "end": 2194.5, "text": " But VMAP is gonna parallelize the computation in the background across multiple threads, I'm not sure how it's exactly implemented"}, {"start": 2194.5, "end": 2209.5, "text": " Basically we wrap the convol function into VMAP, by default it's gonna assume that the in-axis argument is set to 0,0, so this is the equivalent to explicitly writing in-axis equal to this tuple of 0,0"}, {"start": 2209.5, "end": 2225.5, "text": " And after we run this we'll get the result, and as you can see this part is the same as the example above, and then we have increasing numbers because the signal X is being, let me print X here"}, {"start": 2225.5, "end": 2234.5, "text": " Let me just print this batch, you can see we're going from 0 through 39 and that's why you can see increasing numbers here as well"}, {"start": 2234.5, "end": 2245.5, "text": " The cool thing about VMAP now is that we just need to swap VMAP for PMAP, nothing else changes and all of a sudden we're sharding our batch of data"}, {"start": 2245.5, "end": 2252.5, "text": " So the signal X across multiple cores and we're running the computation in parallel independently on each one of these cores"}, {"start": 2252.5, "end": 2260.5, "text": " So let me run this and we get the same results as you can see here above for VMAP, there is one difference you can notice here"}, {"start": 2260.5, "end": 2273.5, "text": " This thing is called device array, whereas here we have sharded device array, so this sharded alludes to the fact that the actual data and computation is distributed across multiple cores"}, {"start": 2273.5, "end": 2279.5, "text": " So that's super cool, so syntax same as VMAP and that's nice and sweet"}, {"start": 2279.5, "end": 2294.5, "text": " Let's continue, and I want to stress that all of this happens literally independently, so basically we have 8 signals, so recall that X was 8 times 5"}, {"start": 2294.5, "end": 2303.5, "text": " What will happen is that array of 5 elements is going to be on each of the cores alongside the kernels, so we basically have a single computation"}, {"start": 2303.5, "end": 2311.5, "text": " So we have a single kernel and a single example on each of the cores and we are doing computations without communicating between the cores"}, {"start": 2311.5, "end": 2317.5, "text": " So that's super nice if the nature of our computation is such that we do not need a communication between the cores"}, {"start": 2317.5, "end": 2326.5, "text": " And here I'm just illustrating that we can do one pass of this computation, so we do a convolution one times"}, {"start": 2326.5, "end": 2335.5, "text": " And then we use the outputs as the new kernels and we again convolve them with the original signal X"}, {"start": 2335.5, "end": 2343.5, "text": " And basically the thing is during both of these two calls we did not have a single communication between the cores themselves"}, {"start": 2343.5, "end": 2349.5, "text": " So let me run this, and we get some even bigger numbers and that's not that important anymore"}, {"start": 2349.5, "end": 2360.5, "text": " Finally, just notice that instead of repeating, replicating W's, which is kind of wasteful, we can just use the in-axis argument here"}, {"start": 2360.5, "end": 2371.5, "text": " And as you recall, VMAP uses the same syntax, basically in-axis tells PMAP that the first argument, so that means W, basically should be broadcasted"}, {"start": 2371.5, "end": 2382.5, "text": " Because we set none here, and here zero just means that the first zero dimension of XS corresponds to the batch dimension"}, {"start": 2382.5, "end": 2390.5, "text": " So this way, PMAP will in the background replicate W and do the job for us, so that's way better"}, {"start": 2390.5, "end": 2400.5, "text": " And let's run this and we should see the same results as before, so 11, 20, 29, 326, that was the same number as here"}, {"start": 2400.5, "end": 2404.5, "text": " And elsewhere, so that's the same output again"}, {"start": 2404.5, "end": 2411.5, "text": " So again, this is super nice if your computation is such that you do not need communication between the cores"}, {"start": 2411.5, "end": 2418.5, "text": " This is usually not the case, and an example we'll soon see is the one of training ML model in a distributed fashion"}, {"start": 2418.5, "end": 2425.5, "text": " Where as we take our batch of data and we basically split the batch of data into mini-batches, we send each mini-batch to each core"}, {"start": 2425.5, "end": 2434.5, "text": " And then we do computations there, we calculate the gradients and then we communicate between the cores in order to do a mean of the gradients to update our ML model"}, {"start": 2434.5, "end": 2443.5, "text": " We're gonna see that example in a second, but just be aware that certain computations do require communication, certain do not, and Jack's got us covered for both of those use cases"}, {"start": 2443.5, "end": 2450.5, "text": " So now let's see a simple example, like just a modification on the above function that we had"}, {"start": 2450.5, "end": 2457.5, "text": " So this is going to be normalized convolution, and the only difference is that we'll have a communication between cores"}, {"start": 2457.5, "end": 2458.5, "text": " So let's see how this works"}, {"start": 2458.5, "end": 2469.5, "text": " This part of the function remains totally the same as before, so we have output here, we are doing like sliding across the signal, and we accumulate the results into output"}, {"start": 2469.5, "end": 2479.5, "text": " Now the only difference is that we are calling this, so remember, we're calling this Jack's LexPSum function, and just recall that Lex is this middle layer API of Jack's"}, {"start": 2479.5, "end": 2494.5, "text": " So what we're doing here is we're going to, along the batch dimension, we're going to sum up all of the outputs, and we're gonna divide output on this particular core with the sum of outputs across all of the cores"}, {"start": 2494.5, "end": 2496.5, "text": " And that's gonna normalize our output"}, {"start": 2496.5, "end": 2506.5, "text": " So a thing to notice is this axis name, it can be a bit confusing, so what you do is you again wrap your function of interest into PMAP"}, {"start": 2506.5, "end": 2515.5, "text": " You specify in axis, which means, hey, I want to broadcast W, I want to use the zero dimension of XS as the batch dimension"}, {"start": 2515.5, "end": 2522.5, "text": " And we just give the name to those mapped axis by putting an arbitrary name here, so this can be whatever"}, {"start": 2522.5, "end": 2528.5, "text": " Like we just need to have consistency between whatever we use here, we have to use here inside of the definition as well"}, {"start": 2528.5, "end": 2534.5, "text": " So I chose batch dimension because it's, I guess, indicative of the thing we're doing, right?"}, {"start": 2534.5, "end": 2538.5, "text": " Let's first run this and then we can analyze it additionally"}, {"start": 2538.5, "end": 2543.5, "text": " The first thing you can notice is that this output is the same as VMAP"}, {"start": 2543.5, "end": 2549.5, "text": " So I'm just going to, I just did that for you to understand that VMAP and PMAP are very, conceptually very similar"}, {"start": 2549.5, "end": 2557.5, "text": " The only difference is that PMAP is actually using physical cores, whereas VMAP is doing the same thing just on the single device"}, {"start": 2557.5, "end": 2565.5, "text": " So let me just comment out VMAP so that we can have a bit less output in the console here"}, {"start": 2565.5, "end": 2573.5, "text": " So basically what the normalization has done is that all of these columns are now normalized"}, {"start": 2573.5, "end": 2580.5, "text": " Meaning if we sum them up, if we take a single column like this one, so this one, like this element and this element and this element"}, {"start": 2580.5, "end": 2583.5, "text": " And we sum them up, we're going to get one, so that's what I've done here"}, {"start": 2583.5, "end": 2589.5, "text": " I took the result, I took all of the rows and zeroth column and we can see here that result sums up to one"}, {"start": 2589.5, "end": 2594.5, "text": " If I put one here, we're going to get the same result, so it's again one"}, {"start": 2594.5, "end": 2598.5, "text": " And if I put two here, we're going to get one as the output and that's it"}, {"start": 2598.5, "end": 2601.5, "text": " That's what this computation did"}, {"start": 2601.5, "end": 2604.5, "text": " Let me try and see whether we can print the output here"}, {"start": 2604.5, "end": 2608.5, "text": " I'm not sure whether this is going to work, let me just try and do this"}, {"start": 2608.5, "end": 2613.5, "text": " Nope, it's going to trace it because PMAP, that's an additional thing to keep in mind"}, {"start": 2613.5, "end": 2618.5, "text": " PMAP is automatically calling JIT in the background, so yeah"}, {"start": 2618.5, "end": 2623.5, "text": " To make this a bit more clear, I'm going to take the last example that didn't have any communication"}, {"start": 2623.5, "end": 2628.5, "text": " So let me take something like this"}, {"start": 2628.5, "end": 2631.5, "text": " I'm just going to copy paste it here"}, {"start": 2631.5, "end": 2634.5, "text": " So we're just going to copy paste that thing here"}, {"start": 2634.5, "end": 2639.5, "text": " Okay, now we have a result here that is not... I'm going to print it here"}, {"start": 2639.5, "end": 2645.5, "text": " So PMAP result and this result is... there is no normalization here"}, {"start": 2645.5, "end": 2649.5, "text": " So let me just print it out here"}, {"start": 2649.5, "end": 2650.5, "text": " So here are the results"}, {"start": 2650.5, "end": 2657.5, "text": " So what we have done is this PSUM basically alongside batch dimension, which is this dimension here"}, {"start": 2657.5, "end": 2661.5, "text": " So the vertical dimension is going to sum up all of these values"}, {"start": 2661.5, "end": 2665.5, "text": " And use them to divide the current output"}, {"start": 2665.5, "end": 2669.5, "text": " So for each core, for example, for core 0, this will be the output"}, {"start": 2669.5, "end": 2675.5, "text": " And we're going to divide this array here with the sum across this dimension, the vertical dimension"}, {"start": 2675.5, "end": 2678.5, "text": " So let me just kind of show you that that is the case"}, {"start": 2678.5, "end": 2687.5, "text": " So if I was to take the element number like this element here, 11"}, {"start": 2687.5, "end": 2692.5, "text": " So that's going to be temp element, something like this, I don't know"}, {"start": 2692.5, "end": 2695.5, "text": " PMAP result"}, {"start": 2695.5, "end": 2698.5, "text": " Basically I'm going to take 0th"}, {"start": 2698.5, "end": 2703.5, "text": " So that's this first row, I'm going to take the 0th element, so that's 11"}, {"start": 2703.5, "end": 2706.5, "text": " Let me confirm that's the case, so I'm going to print it"}, {"start": 2706.5, "end": 2710.5, "text": " And then I'm going to create a temporary sum here"}, {"start": 2710.5, "end": 2719.5, "text": " So basically just a simple sum of PMAP result across basically across the vertical dimension"}, {"start": 2719.5, "end": 2722.5, "text": " So that's the batch dimension, 0th column"}, {"start": 2722.5, "end": 2725.5, "text": " And I'm going to print the sum as well"}, {"start": 2725.5, "end": 2727.5, "text": " And then I'm going to just do division"}, {"start": 2727.5, "end": 2732.5, "text": " So I'm going to print temp element divide by temp sum"}, {"start": 2732.5, "end": 2737.5, "text": " And we should get the same result as here, so 0.0081 whatever, okay?"}, {"start": 2737.5, "end": 2742.5, "text": " So in order to make this a bit less crowded, let me see where I can do something"}, {"start": 2742.5, "end": 2744.5, "text": " No, let me just run this"}, {"start": 2744.5, "end": 2745.5, "text": " This should work as expected now"}, {"start": 2745.5, "end": 2749.5, "text": " So I took 11, I created a sum here"}, {"start": 2749.5, "end": 2751.5, "text": " And after dividing you can see we got the same result"}, {"start": 2751.5, "end": 2755.5, "text": " So hopefully this now is perfectly clear what happened"}, {"start": 2755.5, "end": 2761.5, "text": " And the whole magic was because of this middle level API's PSUM function"}, {"start": 2761.5, "end": 2765.5, "text": " And just giving a name to these mapped axes here"}, {"start": 2765.5, "end": 2768.5, "text": " Let's continue, let me just erase this output"}, {"start": 2768.5, "end": 2771.5, "text": " Okay, now before we get to the final training pipeline"}, {"start": 2771.5, "end": 2774.5, "text": " There's a couple more functions you should be familiar with"}, {"start": 2774.5, "end": 2776.5, "text": " And then you can train whatever in PureJax"}, {"start": 2776.5, "end": 2779.5, "text": " Pretty much whatever ML model you wish to train"}, {"start": 2779.5, "end": 2782.5, "text": " The first function is this value and grad function"}, {"start": 2782.5, "end": 2789.5, "text": " And it's going to return the value of the loss as well as the gradients of the function of interest"}, {"start": 2789.5, "end": 2791.5, "text": " So I have a simple loss function here"}, {"start": 2791.5, "end": 2794.5, "text": " It's basically just, we just create these error terms"}, {"start": 2794.5, "end": 2796.5, "text": " We square them and we sum them up"}, {"start": 2796.5, "end": 2798.5, "text": " We have a simple input here"}, {"start": 2798.5, "end": 2800.5, "text": " So let me just print it here"}, {"start": 2800.5, "end": 2803.5, "text": " So basically it's going to be 0 through 3"}, {"start": 2803.5, "end": 2809.5, "text": " And finally the value will hopefully return the loss value as well as the gradients"}, {"start": 2809.5, "end": 2812.5, "text": " It's fairly trivial to show that the numbers make sense"}, {"start": 2812.5, "end": 2815.5, "text": " Let me just print y here as well"}, {"start": 2815.5, "end": 2818.5, "text": " So it's going to be the same as x just plus 0.1"}, {"start": 2818.5, "end": 2826.5, "text": " So basically the value of the loss does make sense because we basically subtract"}, {"start": 2826.5, "end": 2830.5, "text": " The difference will be 0.1 for each element here"}, {"start": 2830.5, "end": 2833.5, "text": " Then we square it which is going to give us 0.01"}, {"start": 2833.5, "end": 2836.5, "text": " And we have 4 times that because we have a sum"}, {"start": 2836.5, "end": 2839.5, "text": " So that's why we have 0.04 here"}, {"start": 2839.5, "end": 2843.5, "text": " Same goes for gradients, all of the gradients are minus 0.2"}, {"start": 2843.5, "end": 2850.5, "text": " Because the derivative with respect to x will be 2 times x minus y"}, {"start": 2850.5, "end": 2852.5, "text": " The difference is 0.1"}, {"start": 2852.5, "end": 2855.5, "text": " I mean the minus 0.1 times 2"}, {"start": 2855.5, "end": 2857.5, "text": " That's why we get these numbers here"}, {"start": 2857.5, "end": 2862.5, "text": " So, and in any case if you want to log both the value of loss and you want to get your gradients"}, {"start": 2862.5, "end": 2868.5, "text": " You have this efficient way to do that is to just call this inbuilt function called value and grad"}, {"start": 2868.5, "end": 2870.5, "text": " The second function you should be familiar with is"}, {"start": 2870.5, "end": 2873.5, "text": " Not a function actually, just a functionality"}, {"start": 2873.5, "end": 2876.5, "text": " Is if you want to return sometimes from your loss function"}, {"start": 2876.5, "end": 2878.5, "text": " You want to return not just the loss itself"}, {"start": 2878.5, "end": 2884.5, "text": " But maybe also some intermediate values like maybe you want to return these raw error terms"}, {"start": 2884.5, "end": 2886.5, "text": " You can just instead of doing this"}, {"start": 2886.5, "end": 2890.5, "text": " So if I were to call gradient function on this loss function here"}, {"start": 2890.5, "end": 2893.5, "text": " It has two outputs which means a vector output"}, {"start": 2893.5, "end": 2898.5, "text": " We're going to get an error here because gradient only works on scalar output functions"}, {"start": 2898.5, "end": 2902.5, "text": " So you can see here gradient only defined for scalar output functions"}, {"start": 2902.5, "end": 2908.5, "text": " So there is this neat auxiliary flag"}, {"start": 2908.5, "end": 2911.5, "text": " If we set it to true and we run this"}, {"start": 2911.5, "end": 2915.5, "text": " It's going to return the gradient as well as the"}, {"start": 2915.5, "end": 2920.5, "text": " It's just going to pass this intermediate value as the output directly"}, {"start": 2920.5, "end": 2924.5, "text": " Without doing any type of manipulation on top of it"}, {"start": 2924.5, "end": 2929.5, "text": " Awesome, we are now ready to train a machine learning model completely in parallel"}, {"start": 2929.5, "end": 2932.5, "text": " Okay, let me see how I can go about explaining this part"}, {"start": 2932.5, "end": 2934.5, "text": " Because there is a lot of functionality here"}, {"start": 2934.5, "end": 2936.5, "text": " But I think it's going to be digestible"}, {"start": 2936.5, "end": 2939.5, "text": " Let me first start by analyzing our data"}, {"start": 2939.5, "end": 2943.5, "text": " So we have a simple model like the ground truth model is this one"}, {"start": 2943.5, "end": 2947.5, "text": " Basically a single input linear function here"}, {"start": 2947.5, "end": 2950.5, "text": " With some addition of noise, it's a noisy, noisy"}, {"start": 2950.5, "end": 2952.5, "text": " We'll have some noise observations here"}, {"start": 2952.5, "end": 2957.5, "text": " So the true value of W, so the coefficient that goes with X"}, {"start": 2957.5, "end": 2960.5, "text": " And of this bias term are 2 and minus 1"}, {"start": 2960.5, "end": 2963.5, "text": " We generate the domain of our data"}, {"start": 2963.5, "end": 2966.5, "text": " Basically just sampling from the normal distribution"}, {"start": 2966.5, "end": 2968.5, "text": " We generate 128 elements"}, {"start": 2968.5, "end": 2974.5, "text": " We generate some noise, again, sampling from the normal distribution"}, {"start": 2974.5, "end": 2979.5, "text": " And finally we generate our noisy observations by applying the above function here"}, {"start": 2979.5, "end": 2984.5, "text": " So XS times the true value of the weight parameter"}, {"start": 2984.5, "end": 2987.5, "text": " Plus the true value of the bias parameter plus the noise"}, {"start": 2987.5, "end": 2990.5, "text": " And let me visualize that for you"}, {"start": 2990.5, "end": 2993.5, "text": " So we understand what we are trying to regress"}, {"start": 2993.5, "end": 2996.5, "text": " So we are going to try to fit a line to this data"}, {"start": 2996.5, "end": 3000.5, "text": " And we are going to do that the following way"}, {"start": 3000.5, "end": 3003.5, "text": " So first we'll need to initialize our model"}, {"start": 3003.5, "end": 3006.5, "text": " So let's jump to the init model function here"}, {"start": 3006.5, "end": 3010.5, "text": " It's going to initially create a random value for the weight and for the bias"}, {"start": 3010.5, "end": 3014.5, "text": " So this is our approximation of the true weight and the true bias values"}, {"start": 3014.5, "end": 3018.5, "text": " Initially the best we can do is just to create a random value"}, {"start": 3018.5, "end": 3021.5, "text": " And return them packed into this params class"}, {"start": 3021.5, "end": 3024.5, "text": " Which is just inherited from this name tuple"}, {"start": 3024.5, "end": 3029.5, "text": " So it has weight key and bias key and the values are our Jax arrays"}, {"start": 3029.5, "end": 3035.5, "text": " As you can see here I'm using Jax's built-in PRNG functionality"}, {"start": 3035.5, "end": 3038.5, "text": " I'm using the random number generator"}, {"start": 3038.5, "end": 3041.5, "text": " I'm splitting the keys into weight and bias"}, {"start": 3041.5, "end": 3044.5, "text": " And I'm consuming those here to generate two random numbers"}, {"start": 3044.5, "end": 3047.5, "text": " So that's what this init model function does"}, {"start": 3047.5, "end": 3050.5, "text": " Once we have that I'm going to fetch the number of devices"}, {"start": 3050.5, "end": 3055.5, "text": " That's going to be 8 if you recall from the previous exercises we had 8 cores"}, {"start": 3055.5, "end": 3060.5, "text": " And we are going to replicate the parameters because we want to shard this"}, {"start": 3060.5, "end": 3067.5, "text": " We want to replicate this particular model across all the cores and do the computation in parallel"}, {"start": 3067.5, "end": 3072.5, "text": " Since params, so the thing we just got from our init model function is a pie tree"}, {"start": 3072.5, "end": 3076.5, "text": " We can iterate over the leaves, so over the weight and over the bias"}, {"start": 3076.5, "end": 3081.5, "text": " And we can just multiply them by 8 which is going to replicate the values"}, {"start": 3081.5, "end": 3086.5, "text": " So we'll have, yeah, I'll just going to print it here so you're going to see it in a second"}, {"start": 3086.5, "end": 3091.5, "text": " Next thing we're going to do is just reshape the data for the PMAP function"}, {"start": 3091.5, "end": 3093.5, "text": " Because we're going to trade this in parallel again"}, {"start": 3093.5, "end": 3097.5, "text": " So we have the XS and YS, so that's the data we just saw here"}, {"start": 3097.5, "end": 3104.5, "text": " And we're going to basically just reshape them so that the leading dimension has the number 8"}, {"start": 3104.5, "end": 3107.5, "text": " So let me run this and see what we have"}, {"start": 3107.5, "end": 3111.5, "text": " So, okay, I forgot to run this cell"}, {"start": 3111.5, "end": 3115.5, "text": " I'm going to run this cell and then I'm going to rerun this cell again"}, {"start": 3115.5, "end": 3118.5, "text": " And let's see what we got here"}, {"start": 3118.5, "end": 3120.5, "text": " So you can see a couple of things here"}, {"start": 3120.5, "end": 3125.5, "text": " First thing, weight and bias have the same number replicated 8 times as expected"}, {"start": 3125.5, "end": 3135.5, "text": " And finally the data is now reshaped into this 8 times 16 times 1 shape which is suitable for PMAP"}, {"start": 3135.5, "end": 3142.5, "text": " The Bayesian training pipeline will look very similar to the MLP we trained like 20 minutes ago"}, {"start": 3142.5, "end": 3147.5, "text": " Basically we're going to have a number of epochs, we're going to iterate for a certain number of epochs"}, {"start": 3147.5, "end": 3151.5, "text": " We're going to have an update function and again the familiar pattern"}, {"start": 3151.5, "end": 3154.5, "text": " We pass in the replicated parameters, we pass in the data"}, {"start": 3154.5, "end": 3159.5, "text": " And we get back the parameters, updated ones, as well as the loss value"}, {"start": 3159.5, "end": 3163.5, "text": " After that I'm going to just log some values here and we're going to see that in a second"}, {"start": 3163.5, "end": 3167.5, "text": " And finally we're going to take our final replicated parameters"}, {"start": 3167.5, "end": 3170.5, "text": " And we're just going to fetch parameters from a single core"}, {"start": 3170.5, "end": 3173.5, "text": " Because we'll basically have 8 duplicates across the cores"}, {"start": 3173.5, "end": 3177.5, "text": " So we just need to extract a single one and we get parameters here"}, {"start": 3177.5, "end": 3187.5, "text": " So TREEMAP will create a PyTree object and then DeviceCat just takes the data from the TPU and downloads it to the host memory"}, {"start": 3187.5, "end": 3191.5, "text": " That's it, now that's the high level picture of what is going to happen here"}, {"start": 3191.5, "end": 3197.5, "text": " Now let me start digesting these couple of functions here"}, {"start": 3197.5, "end": 3200.5, "text": " So what we need to understand is that the model is very simple"}, {"start": 3200.5, "end": 3206.5, "text": " We just have weight times XS plus the bias because that corresponds perfectly"}, {"start": 3206.5, "end": 3212.5, "text": " Because we know what generated our data so we can create the same model here"}, {"start": 3212.5, "end": 3217.5, "text": " Usually that's not the case, you usually don't know how your model structure looks like"}, {"start": 3217.5, "end": 3224.5, "text": " That's why you have a huge neural network and you let the data morph the structure of that neural network"}, {"start": 3224.5, "end": 3229.5, "text": " Loss function is very simple, we just basically do mean square error again"}, {"start": 3229.5, "end": 3234.5, "text": " So predictions, we subtract the ground truth, etc"}, {"start": 3234.5, "end": 3239.5, "text": " The update function is by far the hardest to understand so we're going to focus a lot on this function"}, {"start": 3239.5, "end": 3241.5, "text": " So let's see what update function does"}, {"start": 3241.5, "end": 3247.5, "text": " So we get our parameters, so that's the current version of the parameters, we get the data"}, {"start": 3247.5, "end": 3255.5, "text": " And what we do is first we decorate the update with this PMAP and we name the mapped axis to batch"}, {"start": 3255.5, "end": 3262.5, "text": " So this is just a convenient way to, if we do it like this then we don't have to wrap it inside of the PMAP here"}, {"start": 3262.5, "end": 3266.5, "text": " So let me show you, so here as you can see we don't have PMAP explicitly here"}, {"start": 3266.5, "end": 3272.5, "text": " We just decorated the update directly with PMAP and that's why we can avoid putting PMAP here"}, {"start": 3272.5, "end": 3278.5, "text": " So what will happen is that each TPU core is going to get a portion of the parameters"}, {"start": 3278.5, "end": 3283.5, "text": " So because those are replicated we're just going to have the same parameters spread out across every single core"}, {"start": 3283.5, "end": 3289.5, "text": " We're going to get a single mini batch out of the whole batch of data on each of the cores"}, {"start": 3289.5, "end": 3299.5, "text": " And then we're going to pass those into the loss function and find the value of the loss as well as the gradients of the loss function"}, {"start": 3299.5, "end": 3306.5, "text": " So that means because again grad does the derivative with respect to the first parameter"}, {"start": 3306.5, "end": 3311.5, "text": " Which means we're going to have derivatives with respect to W, the weight and the bias"}, {"start": 3311.5, "end": 3316.5, "text": " Next thing we're going to do is we're going to combine the gradients across all of the TPU cores"}, {"start": 3316.5, "end": 3322.5, "text": " And this is arguably one of the most important, the most important operation in this whole training pipeline"}, {"start": 3322.5, "end": 3329.5, "text": " Basically you take the gradients which are now different across each of the cores"}, {"start": 3329.5, "end": 3334.5, "text": " And they're different because each core got a different portion of the data set"}, {"start": 3334.5, "end": 3337.5, "text": " Whoops, I'm not a robot, I assure you"}, {"start": 3337.5, "end": 3342.5, "text": " So basically you're going to do a mean averaging across all of the gradients"}, {"start": 3342.5, "end": 3348.5, "text": " And finally each single core will have the same gradients because we do this p mean operation"}, {"start": 3348.5, "end": 3356.5, "text": " As a consequence because we also have the same parameters, that means that each core will have the parameters in sync"}, {"start": 3356.5, "end": 3363.5, "text": " So all of the parameters are initially the same, the gradients are the same which means that the updated parameters across each of the cores are also the same"}, {"start": 3363.5, "end": 3364.5, "text": " So that's cool"}, {"start": 3364.5, "end": 3368.5, "text": " We did the same thing for loss, we just averaged loss across the 8 cores"}, {"start": 3368.5, "end": 3374.5, "text": " And finally we applied the stochastic gradient descent like update rule"}, {"start": 3374.5, "end": 3383.5, "text": " Again we have pi tree here, pi tree here, we just do the usual from the current value of the parameter"}, {"start": 3383.5, "end": 3389.5, "text": " We subtract the learning rate times the gradient and that's how we get our new parameters"}, {"start": 3389.5, "end": 3393.5, "text": " Finally we're going to return those parameters as well as the loss value and that's it"}, {"start": 3393.5, "end": 3400.5, "text": " I made a short note here which may be important if you want to train with some more complex optimizers and not just with SGD"}, {"start": 3400.5, "end": 3407.5, "text": " So if we had an optimizer this is how the pattern would look like, we pass in the grads and we pass in the optimizer state"}, {"start": 3407.5, "end": 3415.5, "text": " And the optimizer because of remember our pattern is going to return the new optimizer state and the new basically updated version of gradients"}, {"start": 3415.5, "end": 3419.5, "text": " And then we can just apply the updates instead of grads here and that's it"}, {"start": 3419.5, "end": 3427.5, "text": " So it's very simple just another additional line following this stateless pattern that we saw in the beginning of this video"}, {"start": 3427.5, "end": 3431.5, "text": " Nice, let me now go and run the training"}, {"start": 3431.5, "end": 3435.5, "text": " So this is our training, let me run it and let's see what we get"}, {"start": 3435.5, "end": 3444.5, "text": " Okay it's working, we're printing some values, we're going to debug and understand, analyze what's going on in the console here"}, {"start": 3444.5, "end": 3447.5, "text": " But before that let's just plot the results"}, {"start": 3447.5, "end": 3459.5, "text": " Okay so let me run this and we see that we have a line that's fit to this particular data set we artificially built"}, {"start": 3459.5, "end": 3462.5, "text": " Now let's analyze the output in the console, what happened here"}, {"start": 3462.5, "end": 3470.5, "text": " First thing I want you to notice here is that replicated parameters because as I said those are going to stay on the cores"}, {"start": 3470.5, "end": 3475.5, "text": " They should be sharded, sharded device type so let's see whether that's the case"}, {"start": 3475.5, "end": 3482.5, "text": " So we can see here after printing after the 0th epoch we're going to have this is a class of blah blah blah sharded device array"}, {"start": 3482.5, "end": 3489.5, "text": " Which means that indeed our replicated parameters are on the cores themselves"}, {"start": 3489.5, "end": 3494.5, "text": " Loss is as well a sharded device array"}, {"start": 3494.5, "end": 3498.5, "text": " But the data is as you can see here I'm just printing the type of the data"}, {"start": 3498.5, "end": 3502.5, "text": " We can see it's a non-py array which means it's locally stored on the CPU"}, {"start": 3502.5, "end": 3510.5, "text": " Because we do not want to keep the data on the core that wouldn't make any sense especially for bigger data sets"}, {"start": 3510.5, "end": 3513.5, "text": " Finally what I'm printing here is the loss shape"}, {"start": 3513.5, "end": 3518.5, "text": " You can see it's an 8 dimensional vector because again it's sharded across different cores"}, {"start": 3518.5, "end": 3522.5, "text": " And we can see that the loss goes down as we are advancing here"}, {"start": 3522.5, "end": 3528.5, "text": " It gets into situation after already a couple hundred steps here but yeah"}, {"start": 3528.5, "end": 3533.5, "text": " And that's pretty much it we trained our first machine learning model in parallel on 8th epoch cores"}, {"start": 3533.5, "end": 3536.5, "text": " And as you saw it was very simple to do that"}, {"start": 3536.5, "end": 3540.5, "text": " Now you're almost ready to be able to train whatever type of machine learning model"}, {"start": 3540.5, "end": 3543.5, "text": " There's a couple more questions you have to get an answer to"}, {"start": 3543.5, "end": 3550.5, "text": " So how would you go about doing transfer learning injects? Relatedly how do we free certain layers and fine tune other layers?"}, {"start": 3550.5, "end": 3554.5, "text": " So we need to figure out how to do that, how to do stop gradients"}, {"start": 3554.5, "end": 3559.5, "text": " And finally how do we get per sample gradients instead of the usual per batch gradients"}, {"start": 3559.5, "end": 3563.5, "text": " I'm going to tell you all about that in this last section of this tutorial"}, {"start": 3563.5, "end": 3568.5, "text": " In order to demonstrate how stop gradients work in Jax I'm going to use a couple of examples"}, {"start": 3568.5, "end": 3575.5, "text": " So first things first Jax uses this JaxLax stop gradient primitive for this purpose, for stopping gradients"}, {"start": 3575.5, "end": 3580.5, "text": " Let me quickly go through this TD0 update rule here"}, {"start": 3580.5, "end": 3586.5, "text": " If you're not familiar with RL this may be super confusing but you basically you do not need to know what this is"}, {"start": 3586.5, "end": 3590.5, "text": " You just need to know some of the details which I'm going to explain to you right now"}, {"start": 3590.5, "end": 3593.5, "text": " So the whole point is we are trying to learn certain value functions"}, {"start": 3593.5, "end": 3597.5, "text": " So that value function in our example is going to be a simple linear value function"}, {"start": 3597.5, "end": 3606.5, "text": " So it's parameterized by theta and we are basically doing dot product with the state vector in order to get the value function"}, {"start": 3606.5, "end": 3612.5, "text": " So in order to find the gradients of theta we have to do this expression"}, {"start": 3612.5, "end": 3621.5, "text": " So basically we have a reward here at the present timestamp plus the value at the current state minus the value at the previous state"}, {"start": 3621.5, "end": 3626.5, "text": " And we multiply that by the gradient of the value function in the previous state"}, {"start": 3626.5, "end": 3632.5, "text": " So that's a mouthful, the whole point is the update is not a gradient of any loss function"}, {"start": 3632.5, "end": 3637.5, "text": " However if we reformulate this above expression into this thing"}, {"start": 3637.5, "end": 3644.5, "text": " And if we have a way to stop the gradients of this target here of this value at the present state"}, {"start": 3644.5, "end": 3649.5, "text": " We'll be able by doing a gradient of this pseudo loss function to get exactly this expression"}, {"start": 3649.5, "end": 3656.5, "text": " Again as you can see here if we ignore the dependency of the target on the parameter theta for our value function"}, {"start": 3656.5, "end": 3658.5, "text": " We're going to get exactly this"}, {"start": 3658.5, "end": 3663.5, "text": " So let's see why that is, if we do a gradient with respect to theta, 2 goes in front here"}, {"start": 3663.5, "end": 3668.5, "text": " We'll have this whole expression and then if we do a derivative with respect to theta"}, {"start": 3668.5, "end": 3676.5, "text": " And this does not depend on theta, that means we're just going to end up with the derivative of this previous value function at the previous state"}, {"start": 3676.5, "end": 3678.5, "text": " So that's exactly this expression here"}, {"start": 3678.5, "end": 3681.5, "text": " So just some basic rules of differentiation at the end, that's all"}, {"start": 3681.5, "end": 3684.5, "text": " So let's finally see how we can implement that"}, {"start": 3684.5, "end": 3687.5, "text": " We have the value function as I already explained, simple dot product"}, {"start": 3687.5, "end": 3690.5, "text": " We have some initial values for the theta"}, {"start": 3690.5, "end": 3700.5, "text": " We have some like, so this t1 just means timestamp minus 1, i.e. the previous timestamp"}, {"start": 3700.5, "end": 3703.5, "text": " We have the reward, we have the current state here"}, {"start": 3703.5, "end": 3705.5, "text": " And finally we have the td loss"}, {"start": 3705.5, "end": 3708.5, "text": " We calculate the value function at the previous state"}, {"start": 3708.5, "end": 3711.5, "text": " We calculate the target and this is the whole point here"}, {"start": 3711.5, "end": 3718.5, "text": " We just apply the stop gradient to target and this is going to be our loss"}, {"start": 3718.5, "end": 3722.5, "text": " And now when we do a gradient of that loss with respect to theta"}, {"start": 3722.5, "end": 3724.5, "text": " Because theta is the first argument here"}, {"start": 3724.5, "end": 3729.5, "text": " And remember again, gradient by default will do a derivative with respect to the first parameter"}, {"start": 3729.5, "end": 3736.5, "text": " We get the correct value of delta theta here, i.e. the gradients of our model"}, {"start": 3736.5, "end": 3739.5, "text": " I'm going to run this, although the number will not mean a lot"}, {"start": 3739.5, "end": 3742.5, "text": " And I don't think it's that important to go and analyze this further"}, {"start": 3742.5, "end": 3743.5, "text": " You get the point here"}, {"start": 3743.5, "end": 3746.5, "text": " The second example is this straight through estimator"}, {"start": 3746.5, "end": 3751.5, "text": " And I've seen this thing used, for example, in the VQVAE paper"}, {"start": 3751.5, "end": 3755.5, "text": " And I also covered that paper in one of my previous videos, you can check it out"}, {"start": 3755.5, "end": 3758.5, "text": " I've basically pasted the link here"}, {"start": 3758.5, "end": 3764.5, "text": " But the whole point is, if you have a certain non-differentiable function in the middle of your neural network, for example"}, {"start": 3764.5, "end": 3767.5, "text": " For example, this is a simple non-differentiable function"}, {"start": 3767.5, "end": 3772.5, "text": " You take the x, which is a float, and if you apply a rounding, that's simply not differentiable"}, {"start": 3772.5, "end": 3777.5, "text": " So now we can reformulate this non-differentiable function into this straight through f"}, {"start": 3777.5, "end": 3780.5, "text": " So basically, what we have done here is, if you take a look at this"}, {"start": 3780.5, "end": 3783.5, "text": " So x plus f of x minus x"}, {"start": 3783.5, "end": 3787.5, "text": " Basically that means in the forward pass we're going to have f of x"}, {"start": 3787.5, "end": 3790.5, "text": " Which means these two are equivalent in the forward pass"}, {"start": 3790.5, "end": 3796.5, "text": " On the other hand, if we do a derivative, if we find a gradient of this function and a gradient of this function"}, {"start": 3796.5, "end": 3798.5, "text": " They are going to be different"}, {"start": 3798.5, "end": 3800.5, "text": " So let me just take some simple example here"}, {"start": 3800.5, "end": 3802.5, "text": " So x is 5.6, just an arbitrary number"}, {"start": 3802.5, "end": 3805.5, "text": " We're going to print the value in the forward pass"}, {"start": 3805.5, "end": 3807.5, "text": " So we're going to see that those are going to be the same"}, {"start": 3807.5, "end": 3811.5, "text": " So 6 and 6, because when we round 5.6 we get 6"}, {"start": 3811.5, "end": 3813.5, "text": " So that's clear"}, {"start": 3813.5, "end": 3816.5, "text": " Whereas, because this f is non-differentiable"}, {"start": 3816.5, "end": 3820.5, "text": " If we try to do a gradient at this point x, or whatever x is"}, {"start": 3820.5, "end": 3822.5, "text": " We get 0 all the time"}, {"start": 3822.5, "end": 3826.5, "text": " On the other hand, this straight through gradient is going to give us 1 all the time"}, {"start": 3826.5, "end": 3829.5, "text": " You can see why that is, because we have a stop gradient here"}, {"start": 3829.5, "end": 3833.5, "text": " So a derivative of this thing with respect to x is just going to be 1"}, {"start": 3833.5, "end": 3839.5, "text": " Because, yeah, I mean, if a function is equal to x, then basically a gradient is 1"}, {"start": 3839.5, "end": 3841.5, "text": " So why may this be useful?"}, {"start": 3841.5, "end": 3844.5, "text": " If you have an automatic differentiation in the reverse mode"}, {"start": 3844.5, "end": 3846.5, "text": " You're basically, if you have 1 here"}, {"start": 3846.5, "end": 3851.5, "text": " If a certain part of your whole function has a gradient 1"}, {"start": 3851.5, "end": 3855.5, "text": " By multiplying the gradients that came from the deeper levels by 1"}, {"start": 3855.5, "end": 3858.5, "text": " You're just going to pass them through, that's why it's called straight through"}, {"start": 3858.5, "end": 3864.5, "text": " You're just going to copy them to the second part of the network or whatever ML model you're using"}, {"start": 3864.5, "end": 3866.5, "text": " And that's why this thing is kind of cool"}, {"start": 3866.5, "end": 3871.5, "text": " But the whole point is, again, just be aware that we are using this stop gradient to achieve this functionality"}, {"start": 3871.5, "end": 3873.5, "text": " Let's see a couple more advanced things"}, {"start": 3873.5, "end": 3875.5, "text": " So, calculating per sample gradients"}, {"start": 3875.5, "end": 3877.5, "text": " And let me just read this through"}, {"start": 3877.5, "end": 3883.5, "text": " In many frameworks, PyTorch, TensorFlow, Theano, it is often not trivial to compute per example gradients"}, {"start": 3883.5, "end": 3887.5, "text": " Because the library directly accumulates the gradient over the batch"}, {"start": 3887.5, "end": 3894.5, "text": " Naive workarounds, such as computing a separate loss per example and then aggregating the resulting gradients"}, {"start": 3894.5, "end": 3896.5, "text": " Are typically very inefficient"}, {"start": 3896.5, "end": 3901.5, "text": " So you can go around this problem by taking a single sample, so for example, maybe a single image"}, {"start": 3901.5, "end": 3904.5, "text": " Passing it through, calculating the gradients for that single image"}, {"start": 3904.5, "end": 3909.5, "text": " And then doing that for the whole batch in a kind of a for loop and then aggregate the results"}, {"start": 3909.5, "end": 3912.5, "text": " Aggregate those gradients, but it's super expensive as you can imagine"}, {"start": 3912.5, "end": 3917.5, "text": " Jax basically has direct support for this type of a thing and it's very efficient"}, {"start": 3917.5, "end": 3927.5, "text": " So if we create a contrived example here by just stacking the state in the t-minus-1 timestamp"}, {"start": 3927.5, "end": 3930.5, "text": " And we batch the rewards, we batch the current state"}, {"start": 3930.5, "end": 3935.5, "text": " So basically using VMAP is going to allow us to get these per sample gradients"}, {"start": 3935.5, "end": 3943.5, "text": " Because remember, TD loss does not accept a batch of data"}, {"start": 3943.5, "end": 3948.5, "text": " It accepts a single example, so a single state, a single reward, and a single state here"}, {"start": 3948.5, "end": 3950.5, "text": " And we have arguments theta here"}, {"start": 3950.5, "end": 3958.5, "text": " So if we map it, if we wrap it inside of a VMAP using in-axis none, which means theta"}, {"start": 3958.5, "end": 3963.5, "text": " The arguments of our value, so the parameters of our value function are going to be broadcasted"}, {"start": 3963.5, "end": 3970.5, "text": " And then we say 000, which means the first zero dimension of these batches here is the batch dimension"}, {"start": 3970.5, "end": 3974.5, "text": " If we do that, we're going to trivially get the gradients for every single example"}, {"start": 3974.5, "end": 3978.5, "text": " And they are the same because we batch the same states here, so that's kind of trivial"}, {"start": 3978.5, "end": 3983.5, "text": " Wrapping up this tutorial, which is already quite lengthy"}, {"start": 3983.5, "end": 3989.5, "text": " In a summary, Jack's auto-differentiation engine is super powerful, you can do various stuff"}, {"start": 3989.5, "end": 3993.5, "text": " You can even define custom derivative rules if you have some numerical problems"}, {"start": 3993.5, "end": 3997.5, "text": " That's now already super advanced, you can check out the docs here"}, {"start": 3997.5, "end": 4006.5, "text": " I particularly linked this Autodeaf cookbook article, which you should read if you want to know all of the nitty-gritty details of the Autodeaf engine"}, {"start": 4006.5, "end": 4010.5, "text": " I'm going to end up by showing a simple example of this MAMBL"}, {"start": 4010.5, "end": 4013.5, "text": " This is this model agnostic metal learning algorithm"}, {"start": 4013.5, "end": 4016.5, "text": " And you can see it's fairly complicated, you can see the math here"}, {"start": 4016.5, "end": 4022.5, "text": " And the whole point is, this is how simply you can implement this MAMBL algorithm in Jack's"}, {"start": 4022.5, "end": 4031.5, "text": " It looks very much like mathematics, and that's why I think Jack's is super, super powerful and useful for researchers"}, {"start": 4031.5, "end": 4034.5, "text": " Not that beginner-friendly as I already mentioned in my previous video"}, {"start": 4034.5, "end": 4038.5, "text": " But even for beginners, if you go across a couple of hurdles, I think it's going to pay off"}, {"start": 4038.5, "end": 4041.5, "text": " Let me briefly go over MAMBL"}, {"start": 4041.5, "end": 4045.5, "text": " So the whole point here is, imagine you have a couple of tasks"}, {"start": 4045.5, "end": 4050.5, "text": " And while training, you do not want to minimize the loss on any particular task"}, {"start": 4050.5, "end": 4054.5, "text": " You want to minimize the loss, the potential losses on all of those tasks"}, {"start": 4054.5, "end": 4059.5, "text": " For example, maybe going in this direction will minimize the loss for the task 1"}, {"start": 4059.5, "end": 4062.5, "text": " Going in this one will minimize the loss for the task 2"}, {"start": 4062.5, "end": 4065.5, "text": " And similarly, going in this one will minimize the loss for task 3"}, {"start": 4065.5, "end": 4070.5, "text": " That would lead us to have either theta 1, that's this"}, {"start": 4070.5, "end": 4074.5, "text": " These thetas are going to minimize the loss on the task 1"}, {"start": 4074.5, "end": 4079.5, "text": " Similarly, this theta here, theta 2, is going to minimize the loss on this task 2"}, {"start": 4079.5, "end": 4081.5, "text": " And similarly for theta 3"}, {"start": 4081.5, "end": 4085.5, "text": " But we actually move here, which is somewhere in between all of those thetas"}, {"start": 4085.5, "end": 4087.5, "text": " And that's how MAMBL works"}, {"start": 4087.5, "end": 4091.5, "text": " And now, briefly going over this metal loss function"}, {"start": 4091.5, "end": 4097.5, "text": " We calculate the gradients here, and then we use the gradients to update"}, {"start": 4097.5, "end": 4100.5, "text": " Using SGD, we update the parameters here"}, {"start": 4100.5, "end": 4103.5, "text": " And we pass that into the loss function"}, {"start": 4103.5, "end": 4106.5, "text": " And then we take the gradients of that loss function"}, {"start": 4106.5, "end": 4109.5, "text": " It's kind of convoluted, I'm not going to try"}, {"start": 4109.5, "end": 4115.5, "text": " This is not the video about MAMBL, but I want you to just kind of appreciate how expressive Jax actually is"}, {"start": 4115.5, "end": 4118.5, "text": " Having said that, hopefully you liked this video"}, {"start": 4118.5, "end": 4120.5, "text": " If you did, share it out with your friends"}, {"start": 4120.5, "end": 4122.5, "text": " Consider subscribing to this channel"}, {"start": 4122.5, "end": 4126.5, "text": " And also join the Discord community, that's where the party is"}, {"start": 4126.5, "end": 4150.5, "text": " See you next time, bye bye!"}]