Charles Frye on using machine learning to multiply by one

Key Takeaways

what i wanted to talk about today is one of my uh sort of favorite test problems one of my favorite sort of basic um simple uh simple examples that i have in my head when i'm thinking about neural network optimization and neural networks in general so the motivating idea here is that neural network optimization is really complicated it's got a lot of pieces to it and even though it might feel easy now it's not actually that easy of a thing to do so you know nowadays we've got so many good architectural ideas we've got so many good just sort of you know built-ins in our in the neural network libraries that we use like pytorch and tensorflow that it doesn't strike us as that difficult to spin up a neural network and train it on something like image classification and this is something that the optimization for deep learning review paper by ro you sun at uh u of i handles really uh mentions and i think makes a a good point about that if you if it was 1980 and you sat down to try and solve an image classification problem with neural networks you probably would fail because you don't have all these nice things in your toolkit that we have uh and so because of that i think it's easy for us to forget that there's a lot of complexity underneath something like dot fit uh and so when it comes to these kind when it comes to trying to understand something difficult uh in in mathematics which is in the end i think uh one of the best ways to understand what's going on in deep learning in this you know numerical computing kind of field i think the most important thing that is underappreciated by people who come from a software background is the importance of simple test problems uh and so the i think there's a really great explanation of this that i i got from a friend who is who finished their phd in in math at berkeley around the same time i was doing mine and what they said was you know everyone i know who does math does basically the exact same thing to make their lives easier whenever they hear a conjecture somebody's thought about uh how something might work in some branch of math they compare it to two examples specific to that branch one that's well behaved and simple and another that's pathological and counter-intuitive and if the conjecture is true on both then it's probably true and if it's not then you've got yourself a nice counter example so when it comes to trying to understand and think about machine learning like what's going to work well what's going to not work well why uh to use a particular piece or why not to use a particular piece or if somebody makes a claim about uh about the nature of deep neural network training i i use this same i use the same tactic and so my two examples one of them is memorization tasks that's one of my pathological test examples that if somebody makes a claim about how neural networks work and how they generalize i think okay does this explain memorization in neural networks well so specifically a memorization task is one where you destroy all structure relating the inputs to the labels so if you were to take say the endness data set and then instead of having the actual targets you just shuffle them uh so this little bit of python code here takes the endness targets and replaces them with a shuffled version of them turns out neural networks can still minimize the training loss in this example even though there is no you know there's no information there anymore uh it's just uh the the labels are effectively random uh so this is something that people have demonstrated and it shows that neural networks are capable of memorizing their data sets even if they don't necessarily do that in real life uh so the the two sides of this coin or you have one example that's sort of the canonical gnarly bad example and then you have one example that's the well-behaved and simple example and for me i use one that is effectively just multiplying by one with a deep neural network so let's go through this it's got a mathematical uh inspiration and so i took a mathematical approach here so a fully connected network let's define that it's a function of its parameters and it and its inputs that alternates between matrix multiplication and function application so it's got all these non-linearities sigma 1 2 up to k and then it's also got these parameter matrices theta up to k and we alternate between applying the matrix to the input and then applying a function to it a possibly non-linear function to it an auto-encoder network is a specific type of network one that compares its outputs directly to its inputs as opposed to comparing them to something else so d is some function that takes two inputs if they're the same it returns zero otherwise it returns something else and what d is doing here is it's comparing the input x to the output of the network f applied to its parameters and its input x and that function there uh that i'm defining here through this is the loss function it says okay the loss function compares the inputs to the outputs and the loss is what we make small and an autoencoder therefore returns something that's close to its inputs so the case of autoencoders includes just multiplying numbers by one as an important special case so if we just take all of those functions to be the identity right just take all the non-linearities out of our neural network we'll end up with something called a linear network and linear networks i'm not the only one who likes to use these as a test case for doing theory on neural networks or for thinking about neural networks so there's a nice paper by sax mcclelland and ganguly that's basically what can we do to understand not only how neural networks went but actually a little bit of the psychology of how animals and humans learn just by studying what happens in neural in linear networks it's also the case that networks with rectifying linear units are actually piecewise linear on their inputs so this gives you a sense of like okay if i only change the input a little bit a a rectified linear network will behave very close to a linear network so it's another reason why this case is useful for helping you build your intuition about what's going on uh the difference being that because it's simple it's and because it's dropped that non-linear piece it's a lot easier to understand so if we also take another step and make this even simpler and choose the dimension of everything to be one so every hidden layer has dimension one the input has dimension one and so the output has dimension one we replace all that matrix multiplication with just a series of things multiplying together uh and we get a function looks a little bit like this i happened to uh i was using this case actually to help me understand jax better so that's kind of where this this talk came from and so in jack style numpy that looks like okay f of theta and x takes the product of all the thetas and multiplies them by x it's the same thing as what's written above and it's got this nice little bit of extra jack stuff here to make it so that i can now apply this function on a single input or multiple inputs so batch or non-batch and i can apply it to one parameter set or many parameters that's all at once so i can effectively apply 10 networks to 10 data points and get to get 100 outputs with just the same function that i can use for normal forward prop so jax is pretty great is one thing i learned in the process of building this example so the reason why i wanted to do this simple setting and one of the things that makes this a really powerful example is because it's so simple that we can actually see everything that we normally can't see when we're doing deep learning you know for real when we're doing it on uh images with thousands of pixels in them and trying to predict one of thousands of classes at the end the dimension of everything is so big the dimension of that input x the dimension of the hidden uh of the of the parameters theta these things are in the thousands in the millions in the billions nowadays uh and so you're you're hosed for trying to think about that spatially and for trying to visualize it but if we in this case choose the number of hidden layers to be one so there's one hidden layer and then one output layer and so two layers total then we've reduced this down so much that now there's only essentially two dimensions two parameter dimensions to think about and so anything that has to do with the parameters we can we can visualize uh we can visualize because we have two dimensions for the parameters and then one dimension for whatever it is we're trying to measure so for example we can look at the network's output as a function of the parameters here so this is a surface here so the values there it gets red as the value gets higher and it gets blue as the value goes lower and then in between it is gray uh and that z axis which is also the color axis is the output of the network on the number one so if i put in the number one to this network uh what do the outputs look like and it looks like the outputs get big as the parameters get as the parameters have the same sign so as we go to 2 2 in the upper right quadrant and it gets uh and similarly the output gets big if both inputs and if both parameters are negative and large that's the bottom left corner which is also red so the uh in order to finally get a loss function and get ourselves to the point where we're looking at a neural network optimization problem we just need to choose that difference there that d function that compares the inputs the outputs uh and so if we choose the squared error then the loss function is just the square difference between the inputs and the outputs averaged over a batch that's easy to code up we just apply the the mean of the squared error function uh like that uh and that loss is actually simple enough that it's something we can write out explicitly so if we just say all right let's say that our only input is the number one there's only one data point then if we look at that as a function of theta let's see what the loss looks like so first theta 1 times theta 0 times 1 is what f of theta colon one is up there if i think a semicolon one up there is that's the definition of what our network does to its inputs so that simplifies things and then we can take that and and write it out it sort of expand it out if we'd like it makes it a little bit easier to do other manipulations uh and it's very clear now that this is just a polynomial right and that's helpful especially if you have a bit more of a math background to recognize something that you may have seen in other places in mathematics and other problems that you've run into you've probably run into polynomials before whereas you know neural networks are maybe a new and uh and curious thing on the other hand maybe if you know a little bit less about math and are a little bit less comfortable with polynomials it's interesting to see that the simplest version of a network of a neural network ends up being a polynomial uh in in terms of its loss function uh so not only can we write it out explicitly we can also sort of look at it we can directly look at it we can visualize this loss function this is the thing that we are moving on when we do gradient descent this is the thing that we are this you know because it's in 2d with this one extra dimension it's sort of a lot more physical right it looks a little bit if you look at it uh like a mountain range right there's two of those those two big red mountains coming up off the side a nice big valley in the middle you know it's a little bit more physical and we can think about it in these sort of more uh physical terms and um this you know in a way that we couldn't if we were looking at a neural network problem where this is tens of thousands of dimensions and uh even more so you know we can also combine that with by hand mathematical calculations to get the values of the gradient so if you sort of if you crank through that that polynomial that we had previously you can calculate the gradient by hand uh uses only the rules of of the grading calculus for polynomials which are on the easier side of calculus when it comes to when it comes to doing calculus and you can actually get an explicit expression for the gradient uh of course in python we rely on automatic differentiation but it's helpful especially in cases where we use automation to take some of the difficulty out of the actual problems we want to work on to be able to go back to a simple problem like this one and and actually crank it out work on it step by step and see what uh what exactly is being automated away and the utility here is that we can actually then look at not just the loss but also the gradient as an overall vector field we can look at what the gradient is at every single point uh and get a better sense maybe for how gradient descent is going to behave so you can see if you take a look at this uh at this gradient field a little bit more closely the highest values are very far away from the the center and then the values drop precipitously so the the these little tiny arrows represent the gradient at each point they point in the direction gradient descent would go so it's the negative gradient and these you can see that they're much much tinier as we get down to the bottom here so the gradient the scale of the gradient it gets really really big when you're far away from a solution and then once you get even a little bit close to it it shrinks back down this is something that's going to generally be the case with neural networks and this is one of the reasons why your loss goes down really really quickly when you first start because you tend to start in some area where you're not doing so uh do so well and there's obvious you know big gradients that help you improve really really quickly so gradient descent is the method that gets used to optimize neural networks right we take that gradient and we apply it we calculate it uh take it scale it flip it around and add it to the value of the parameters so we follow those arrows there that i was showing on that on the surface and so uh that looks like you know in python code it's theta minus eta times the gradient uh and i grabbed it took jax here and i actually vectorized that function which gives you the ability to run gradient descent on multiple networks across batches at the same time all with the same function that you would use to train one again jax has some pretty cool stuff in it uh though i feel obligated to point out that some things are a little bit harder in jacks namely uh that if you want to try and and update the values of the matrix say pre-allocate memory and then touch it later makes jax makes that a little bit harder uh so you know that's i i felt like i couldn't be completely positive without pointing out the one place i ran into some trouble but uh you know that tiny bit of trouble aside uh that that that little um little detail avoided uh it was you know straightforward to set up this nice bashed gradient descent so to run gradient descent for a couple of hundred steps from a couple of hundred different positions was pretty fast uh and that gave me uh this uh this visualization here on this one uh and you can see the a large number of the trajectories are getting sucked towards the center there so they're getting sucked towards the value zero zero right in the middle at the origin uh so multiplying the inputs by zero and that's not exact that's not what we wanted to do right we wanted to get a solution where we multiplied the inputs by one we want the inputs the outputs to match one another and so just looking at this one it would look like okay well there's you know sometimes when i do grady descent it's gonna fail on this problem and sometimes it's uh it's gonna um it's gonna succeed uh and so we can see that if we take a look at the losses over time so this is the kind of plot that you're maybe actually able to do when you're doing you know machine learning at scale you can just plot the loss that you calculate through training and so this looks you know actually quite different from what you would see during training in that there are you know all these uh there are all these traces that are not converging down to to the lowest values and this is just across different initializations but when we work with typical neural networks we use random normalized starting positions right we don't we don't start on these in these special grid positions uh so we need to randomly initialize things and just for fun here randomly doing things in jacks is actually quite a bit more challenging than it is to uh than it is in other frameworks you have to be very explicit about where you're being random and non-random uh but once you uh once i did that it was pretty straightforward to do some once i figured out how to do that it was pretty straightforward to do random initializations and things look actually quite a bit different so this is i think the same number of initial trajectories or at least the same number as the number that i plotted the loss of before but now we can see that actually you know none of them are going towards that center value and all of them are actually starting off at lower levels and going down to at about the same height this is also something that we see in other neural networks in bigger neural networks so long as we're careful to randomly initialize and we're careful to set that scale in a thoughtful way so to make sure that things are not just random but also normalized if you look at the starting positions of those you'll see they form a nice circle in the x y plane uh and training looks much better it looks you know they all seem to be going down to the same value of the loss perfect score of zero seems like things are going in that direction um and with the except and one interesting thing here so you can see one gets a little bit stuck so it heads towards that center there actually maybe if we head back a slide you can see that a couple of these trajectories head towards the center of the um uh they head towards the center of the of the parameter area and they uh before heading down and joining the rest of the trajectories that are going to the correct solution so that's what this that's what this trace here represents that's one of those examples um and the this is something that we do also sometimes see when we're doing uh when we're applying neural networks at uh at scale we see sometimes there are there are trajectories which get stuck for a while before jumping off of a cliff so it seems like they're probably doing something kind of similar to what's happening in this problem where there's a solution that is close to the um that is that looks kind of like this one that has this kind of pringle shape where in one direction the values the loss goes up and in one direction the value goes down okay so so that's another little phenomenon that we can see in this simple problem so due to the simplicity of this loss we can actually also calculate all the possible solutions so at all of the solutions it will be the case that the gradient is zero right so we want the gradient to go down over time as we're optimizing uh and that when the gradient is zero that could be because we're at a solution at least all solutions will have gradient zero uh so let's check what are some values that'll give us gradient zero well um we'll we'll apply it to that uh to the value of the gradient which we can calculate by hand uh it's relatively straightforward i skipped the math here just to keep things going forward and unsurprisingly it's going to happen when theta 1 and theta 2 are multiplied together equal one so that means that if we multiply the input by those values we'll get one um uh we'll get the same value out as in and so those are our solutions and we can actually calculate those and the code there at the bottom now shows the um the uh how you would calculate the positions of all of those of all those solutions x and x and 1 over x for positive values of x and for negative values of x and now we can actually visualize that solution set so we can see that what's happening is that gradient is taking us towards that solution uh and not towards that uh the value in the center where the um where both of the parameters are zero uh so we can see that the they're landing at different points on that solution set and interestingly they're landing on both the values that are positive and the values that are negative depending on where they were initialized to so the interesting thing about that solution set is basically one corresponds to flipping the value and then flipping it back to multiply by one and one corresponds to just you know multiplying it by two positive numbers not flipping the sign and so what's interesting about that there's a couple of interesting sort of i don't know um takeaways i see one is that the only one of those solutions is one that a human would have thought of i wouldn't think i've given i've been given two numbers to multiply by to get one i wouldn't think of using uh i wouldn't think of using two negative numbers to do that uh that wouldn't have immediately jumped out at me but neural networks can find solutions that are very unintuitive to humans uh and the uh the other side of uh of that is also that neural networks don't tend to find the simplest solution that maybe humans would immediately jump for which is to make both of those values equal to one to make you know theta one and uh theta 2 both equal to 1 so you're just multiplying by 1 the entire time that's not something there's no obvious reason for neural networks to do things the simple way and so if they can avoid doing things the simple way they'll do it so i think that's another nice lesson about the way neural networks and machine learning work uh that you can that you can take from this simple example uh the other kind of fun thing about this is if you look at those they've got this nice shape they've got this nice hyperboloid shape those solution values in green in that chart uh and that is you know it's not a coincidence i can take any solution uh i can i can get another solution by multiplying one weight by a value a and the other weight by a value one over a so that allows me to move along that curve there or to switch sides if i use a negative sign uh and so i can there's a there's a very clear algebraic and geometric structure to this set of solutions and that's something that generalizes to more realistic settings as well so for rectified linear networks if you make them big enough they actually have that same algebraic structure it's called a lead group and it's in fact a lead group based off of the linear group and so you know so long as your network is really really big and uses rectified linear units then it will have something like that you can do the same thing where you switch between the possible solutions in that fashion the other thing is that those were two essentially two separate basins so what i mean by basin is if you look at those two sets of solutions they're separated by a hump in the middle the hump isn't that high especially relative to the really high values of the loss but there is this nice hump there that would be hard to get over with gradient descent and so solutions in different basins can indeed have very different representations and sort of performance uh even when they're separated by only a small barrier so this is something that people have found by looking very closely at ensembles or are we stop smartly ensemble neural networks and then also i used a sort of high dimensional version of this same problem to improve techniques for finding zero gradient points of neural networks so finding all the places where the gradient is zero they're called critical points of the loss and it's something that's really useful for understanding optimization i used this multiplying by one example and it's higher dimensional analogs to understand the problem of finding critical points better and to develop better solutions so in general the takeaway here to remind is that you can use simple problems like this one to build your understanding of how a complicated thing works useful ideas and technologies like neural networks are usually compositions of simpler pieces or scaling up those simpler pieces in some way and those simpler pieces tend to admit base cases that are particularly intuitive or easy to understand um and so the uh the last um yeah i guess the last thing is of course you want to combine that with pathological counter examples uh so you know we've got an example like this where things are all looking nice working nicely and things are very uh straightforward and um you know easy to work with then there you've got your pathological counter examples like the memorization task and then you've got this nice one-two punch for tackling difficult concepts and this general approach i think applies to things not just machine learning but more broadly in both software engineering and in math so i'll take a second to answer some of the questions that i saw in the chat there so one question what kinds of problems where jacks is better suited than other frameworks so that question that's an interesting one i think right now jax is more suited for if you want to do fast math that isn't well handled by the other like neural network gpu acceleration frameworks so it's if you want to do things that it's generally it's going to be something weird that people other people in machine learning are not doing so if you want the jacobian of your neural network which is like the gradient uh but for the outputs of the neural network as opposed to the loss that's something that you can do really well and um with you know it's well supported in jacks in a way that it for a long time was not supported in other frameworks it's sort of an inheritor to auto grad in a lot of ways if you're familiar with auto grad which was the automatic differentiation software unaccelerated automatic differentiation software for a while there um so i would say what's most useful about jax actually is that it's very principled and very thoughtful on a computational level uh it's all about composition it uses a lot of ideas from uh uses a lot of ideas from you know the theory of functional programming and from essentially you know category theory and to these branches of abstract math and while it might not necessarily help you directly being able to program things in jacks the kind of mental exercise you get from it is certainly helpful i felt like i learned quite a bit uh just by changing away around the way that i was thinking about my the various functions and pieces of my neural network pipeline uh then the last question could you train something like a language model that has a more dynamic structure that's not neatly divisible into layers absolutely yeah you might run it if you're if you're sort of overriding a value over and over again you might find that you need to write it slightly differently in jax than you would write it in a different programming paradigm in a different more you know more pythonic paradigm like pytorch but you can certainly you can certainly do it and you'd be able to make things i think very fast [Music] you