PonderNet: Learning to Ponder (Machine Learning Research Paper Explained)
deep learning machine learning arxiv explained neural networks ai artificial intelligence paper pondernet deepmind pondernet learning to ponder deepmind pondernet pondernet explained dynamic computation deep learning classic algorithms halting probability deep learning recurrent computation dynamic recurrent network broader impact deep network learning to stop
#pondernet #deepmind #machinelearning

Humans don't spend the same amount of mental effort on all problems equally. Instead, we respond quickly to easy tasks, and we take our time to deliberate hard tasks. DeepMind's PonderNet attempts to achieve the same by dynamically deciding how many computation steps to allocate to any single input sample. This is done via a recurrent architecture and a trainable function that computes a halting probability. The resulting model performs well in dynamic computation tasks and is surprisingly robust to different hyperparameter settings.

OUTLINE:
0:00 - Intro & Overview
2:30 - Problem Statement
8:00 - Probabilistic formulation of dynamic halting
14:40 - Training via unrolling
22:30 - Loss function and regularization of the halting distribution
27:35 - Experimental Results
37:10 - Sensitivity to hyperparameter choice
41:15 - Discussion, Conclusion, Broader Impact

Paper: https://arxiv.org/abs/2107.05407

Abstract:
In standard neural networks the amount of computation used grows with the size of the inputs, but not with the complexity of the problem being learnt. To overcome this limitation we introduce PonderNet, a new algorithm that learns to adapt the amount of computation based on the complexity of the problem at hand. PonderNet learns end-to-end the number of computational steps to achieve an effective compromise between training prediction accuracy, computational cost and generalization. On a complex synthetic problem, PonderNet dramatically improves performance over previous adaptive computation methods and additionally succeeds at extrapolation tests where traditional neural networks fail. Also, our method matched the current state of the art results on a real world question and answering dataset, but using less compute. Finally, PonderNet reached state of the art results on a complex task designed to test the reasoning capabilities of neural networks.1

Authors: Andrea Banino, Jan Balaguer, Charles Blundell

Links:
TabNine Code Completion (Referral): http://bit.ly/tabnine-yannick
YouTube: https://www.youtube.com/c/yannickilcher
Twitter: https://twitter.com/ykilcher
Discord: https://discord.gg/4H8xxDF
BitChute: https://www.bitchute.com/channel/yannic-kilcher
Minds: https://www.minds.com/ykilcher
Parler: https://parler.com/profile/YannicKilcher
LinkedIn: https://www.linkedin.com/in/yannic-kilcher-488534136/
BiliBili: https://space.bilibili.com/1824646584

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

If you want to support me financially (completely optional and voluntary, but a lot of people have asked for this):
SubscribeStar: https://www.subscribestar.com/yannickilcher
Patreon: https://www.patreon.com/yannickilcher
Bitcoin (BTC): bc1q49lsw3q325tr58ygf8sudx2dqfguclvngvy2cq
Ethereum (ETH): 0x7ad3513E3B8f66799f507Aa7874b1B0eBC7F85e2
Litecoin (LTC): LQW2TRyKYetVC8WjFkhpPhtpbDM4Vw7r9m
Monero (XMR): 4ACL8AGrEo5hAir8A9CeVrW8pEauWvnp1WnSDZxW7tziCDLhZAGsgzhRQABDnFy8yuM9fWJDviJPHKRjV4FWt19CJZN9D4n
 Hello there! Today we'll look at PonderNet Learning to Ponder by Andrea Bonino, Jan Ballager and Charles Blundell. This paper on a high level introduces a recurrent architecture or a principle of recurrent computation for deep networks that essentially says the network recurrently computes its output at each step and at each step it can decide to stop now because it is satisfied with the answer that it has. The idea is that at a complex task you can compute for many steps because it requires many steps of thinking and then give the output and for an easy task the network can decide to output right away because it already has computed the solution. This decision can be done on a per sample basis so for each sample the network can decide when it's time to give the final output. This is not necessarily a paper that just makes something bigger and then pushes state-of-the-art on some benchmark and that's why it piqued my interest is that it tries to rephrase a little bit how we think about the connection of deep learning and algorithms like classic algorithms by themselves. Essentially this is a dynamic if condition in this algorithm that decides when it's when it's time to stop and I appreciate that you know it not everything has to be state-of-the-art pushing here this is simply a cool method to do something that's relatively new. Of course things like this have been done before and they are discussed at length in this paper how this paper is different from other papers that do similar things and it does push state-of-the-art just not on benchmarks that you might be super duper familiar with. But yeah it's it's a cool paper it's a short paper the idea is pretty simple and it appears to work and yeah that's exciting stuff. So we're gonna dive into this paper have a look have a look at what's new in this particular model how it works and as always if you have feedback leave a comment subscribe I'd be happy for that and yeah thanks for being here. Okay so in the abstract here they say that in a standard neural network the amount of computation used grows with the size of the inputs but not with the complexity of the problem being learned. So which is true right in a standard neural network you have a forward pass be that in a fully connected neural network where you have you know you have your input and then you go layer layer layer layer layer and then you have your output. This computation here is always the same no matter the input even in a recurrent neural network right you have kind of an input right here at the beginning you have a layer then you have an input again and then you have this that goes into the same layer and then you have the next input that goes into the same layer even a recurrent neural network usually usually just does the same forward pass. This is a little bit different if you have something like a language model that can emit at some point a you know a stop token or an end of sentence token at which point the computation essentially stops but it's a little bit of a different thing than we consider right here. Right here we consider a neural network that has to find the answer to a particular problem and we're gonna see the problems down but one problem that they present is the parity problem. So the parity problem is you get a string of zeros and ones I think there is also negative ones in there but I think they're a bit for a distraction and the answer you're looking for is as a whole is the parity so the amount of ones in this string odd or even right so this requires a let's say an integrated view of computation this is essentially a classic algorithm that you have to perform over this string and neural networks as good as they are in computer vision and speech recognition they are having trouble with simple algorithmic tasks like this so the idea of this paper here is that well it doesn't make sense to apply a neural network that always does the same amount of compute right I shove this sequence just like in here it doesn't make sense because you know if there is just a single one in the string and I see that right away I can give the answer right away however if it's a long string and it has a bunch of ones I might need to think about this problem for a while and thus adapt the number of computation steps I do in my head I might you know first if I look at this string I might first connect these two you know and then that's two and then I might connect these two that's two again and then I might connect these two that's four there's nothing here there's nothing here right okay four so that's kind of like one two three steps of computation so that's the the rough idea whereas this if the string was shorter and and more regular I might need less computation so they say to overcome this limitation we introduce ponder net a new algorithm that learns to adapt the amount of computation based on the complexity of the problem at hand ponder net learns end to end the number of computational steps to achieve an effective compromise between training prediction accuracy computational cost and generalization so we are going to see how they do this yeah exactly so they then they go into the the tasks their experimental tasks in this paper are sort of these constructed tasks where people know you need this dynamic computation they're not gonna they're not gonna compete on like image net or something like this so the majority of the paper is in in contra posing their model against this a CT model the adaptive computation time I believe so there have been previous attempts at doing dynamic computation time yet either they have so it turns out they're kind of finicky and this model here this pondernet model has a bunch of advantages they say they present pondernet that builds on the previous ideas it's fully differentiable which allows for low variance gradient estimates unlike reinforce so a couple of previous attempts have been with reinforcement learning so let's just learn the number of steps or when to stop using reinforcement learning and that as you might know is very very noisy it has unbiased gradient estimates which is also unlike other models in the past and yeah so they say this has consequences in all three in all aspects of the model in pondernet the halting node predicts the probability of halting conditional on not having halted before this kind of seems obvious but apparently that no one has done this so far so what do we need for an architecture for pondernet they say this down here essentially that's the architecture it's an inline formula which you know but that's the architecture so what you need is you need an input okay you need an input which is X your input and X is transformed into a hidden state this is let's say the hidden state at step one those two or you can also reformulate this as just a hidden state the hidden state is going into s the so-called step function and that's the recurrent function right here so into this step function you can put anything you want you can put like a CNN inside you can treat this as an LSTM since we're going to apply it recursively sorry recurrently and anything you want can be the step function as long as it can be applied recurrently so this step function is going to give you the next hidden state right so you can see it's a recurrent neural network however it is also going to give you the output at that particular point in time so y1 I guess that be here and it's also going to give you this number lambda n now what are these so from here you could apply the step function again you'd get h3 you get the output 2 and you'd get lambda sorry that's that's a 1 that's a 2 so it seems like it's just a recurrent neural network and if I were to put push this to the end right I go give my H H H and then at the end I get my Y N and I treat that as the output of the computation then it's just a recurrent neural network however as we said the network can in this case decide to stop anywhere in between for example if it decides to stop at this particular step then that would be the output of the computation so every computation step the network computes and a potential output a suggestion for an output and then it also thinks about whether or not it really wants to answer with that output or whether it wants to continue and to do another step essentially take another shot at answering the question because it doesn't yet have the correct answer and that's where this lambda thing comes in so the lambda is a probability of stopping essentially so here you can see the output lambda is a number between zero and one and that is the probability of halting this is the output considered that the network halts so whenever this is one the network will halt conditioned on the fact that it hasn't previously halted yeah it seemed as I said it seems obvious to formulate it like this because you can you know you can only halt if you haven't previously halted but apparently previous models have simply output a number that is sort of the probability of halting in general which doesn't give you a bias sorry an unbiased gradient if you try to back propagate through it so if you consider the lambdas to be like this if you unroll for an entire training run then you get we get the probability of halting at any particular step this one so this is what this is what the previous networks would have estimated directly however this network estimates these lambdas these ones here you can see how you can compute the probability that for example the network halts after three steps by multiplying up the probability that network has not halted which is this one at step one has not halted at step two and then the probability that network halts at step three that it given that it hasn't halted at the previous steps so that is a valid probability distribution it's a generalization of the geometric distribution and essentially it encapsulates a decision tree right so at you're at the beginning you can halt sorry let's go a halt or not or continue if you continue then again you can halt or you can continue if again you can halt or continue and so on and all of this so if you want the probability that the network halts after you know this the third step then you would consider this node which means that you'd multiply that you multiply up these paths right here and that's the probability that it holds after three steps okay so the network can output this lambda at every step if the lambda is high then the network halts of course at inference this is done probabilistically now at training time this is done a little bit differently so you I hope you can see at inference time you simply go forward and you get a lambda maybe the lambda in the first step is point one and then you flip the coin a biased coin right if if it comes up heads you stop with the probability of point one it comes up tails which is a point nine probability you continue then maybe at the second step it's it's point zero five so maybe maybe you stop but probably you won't stop and then at the third step it like comes up point nine the network thinks yeah I should probably stop here and you sample from that and yes you you might indeed in nine out of ten cases you actually stop there so that's inference how about training how about we train this thing during training what we do is again we input X our input into an encoder for a hidden state and as I said you can also input X all the time into your step function as you see right here but what you do is you unroll the network for a number of steps right independent of these output nodes independent of the sorry if the halting probability let's say we we unroll it for for five steps right here and at every point we get a output and a value y3 y4 this is lambda 2 lambda 3 lambda 4 so at training we simply unroll until a given step now there are some technical difficulties with doing with unrolling for a finite amount of step like how do you normalize the probability distribution because essentially this tree can go on until infinity they find okay we we can simply unroll until kind of the rest probability the probability we haven't used yet is is really small and then just load that all onto the last step but these are technical difficulties that you really only care when you then go and implement however so we unroll for a number of steps and then our we consider all the outputs at the same time now this is one big difference I believe to one of the previous networks to this a CT so what a CT does is it always unrolls and then the the output of the network so for a CT the output of the network would simply be a weighted output of the lambda I y I so the output of the network is always a waiting between the different steps okay and the network can decide okay how do I want to wait the individual outputs whereas here it's different here the output is really either y1 or y2 or y3 or y4 and to in order to pack this into a single loss function what we can do sorry I should probably leave this in order to pack this into a single loss function we simply take okay what's the loss what would be the loss if we answered y1 right what would be the loss and we weigh that by the probability and we say okay what would be the loss of y2 we weighed by the probability that the network output so now if we and so on so plus essentially we compute the expected loss given the probabilities that the network has output so now if we back prop this we back prop through these losses we have of course two paths of back propping so we back prop through the wise which means it's at some so there is a loss right and both these things and these things go into the loss right so the loss is well how bad is this times how probably it was so on so the back propagation path would actually attack at two different paths you can see so the back prop goes into why because you want the network to compute a a better output but the propagation also goes into the lambda because you want the network to get better at estimating when its output is good and when not this I see a little bit as a tricky situation because usually this this seems a little bit unstable just from experience from other papers and so on if you have a back prop through two different things especially that are appear to be multiplied together and that you know the network can now trade off one versus the other which might you might think is desirable right it can either choose to make its output better if it wants to keep the probability high of outputting this thing or it can just reduce the probability that it's going to output whatever it wants to output and you know then it doesn't have to necessarily make the output itself correct because the loss the loss won't be as high for that particular thing because the probability of outputting it is low so network essentially has a choice as I said this might be desirable but usually that's kind of unstable and I think this is just my personal opinion I think a lot of them why this might work might rest on whether or not or let's say the complexity itself of assessing of making why better versus adjusting these probabilities of course yeah so you see if the output y is very complex right then this you know the same gradient signal for that might mean much less than simply reducing the probability okay so if the output is very very complex right not the problem but just the output itself right how to arrive at an output if the output is an entire pixel map or something like this and that has dependencies and so on the network might just choose to always reduce the probability because it's like well how am I gonna how am I gonna make this better at all I don't know I can just reduce the probability I'm going to output this crap right and it will probably do this then for every you know single step which you know if it's complex problem makes sense but still that's it that would be a bit my my fear here and that this is not really discussed in the paper itself so I think the fact that this works might rely on sort of a balance of the of the complexity or information content that you get from the loss at the output node versus the loss at the probability node so okay enough about that so in yeah during training you simply compute the expected loss weighted by the probabilities and then you can back prop through that and I hope you can see the difference between these two one is a they both seem to sum up somehow the outputs weighted by these these factors however one considers the actual output of the network to be a weighted combination of outputs of the individual steps where the other one says no no no the network output is actually one of them we don't know which one ergo for the loss we need to compute the expectation of the loss that seems to be a bit of a let's just say yeah it seems to be a more reasonable formulation though in hindsight you can say many things are reasonable if they work better right yeah so they discuss things like maximum number of pondering steps and so on again which I think is a technical detail and this is interesting so there you have the training loss as we just discussed now we've discussed this part right here which they call the reconstruction loss because you have some kind of desired y and you have a y that comes from this and I was a little bit wrong here in my formulation of course the expectation you don't have you don't want to take the lambdas you actually want to take the probabilities that each thing happens which means that you need to compute this P number you know going along this tree as we did because the P is the actual probability that you reach that node whereas the lambda is only the conditional probability that you reach a node given you were at the previous node so yeah consider there that if you if you are crazy enough to implement things straight as I speak in the videos lucid rains shout out the second part of the loss here and you can see this is a hyper parameter so you you're gonna trade off two of two losses right here because right now we saw okay you can either continue or not continue and for the network you know it might actually be easier as I said if the loss of the output comes reasonably complex right here it might be easier to simply say well in this case I'm just always going to reduce my probabilities you might counteract this with having this number of steps not like maximum number of steps but essentially this term here is what counteracts that really there is a regularization term on these probabilities as you can see right here so we regularize with the KL divergence which is sort of a distance measure don't tell this to a mathematician it's a it's a divergence it's a sort of a distance measure between the distribution that the network outputs for the steps and this thing right here which is a geometric distribution with this parameter and this parameter lambda p is another hyper parameter so what does that mean essentially if you consider here the number of steps that the network thinks right think things for what you regularize for this distribution right here is a geometric distribution I'll go something like maybe no something like this so essentially a geometric distribution is set exactly computes this tree that we computed right so at each step you can essentially stop and the question is after you know this distribution gives you a indication after what's the probability that you stop after one step two steps three steps four steps considering the fact that in order to stop after four steps you already have to have made three non-stopping steps except in the geometric distribution the probability of continuing is always the same whereas in our network our network for each node and the tree it can output a different probability otherwise you know there'd be no point we can simply put in the fixed distribution now what that probability is of stopping at each point that's exactly this lambda p hyper parameter right here so you regularize for a KL for this which means that you tell the network look here is a a reasonable reasonable distribution of when you should stop so you should stop so it should be you know somewhat probable that you stop after one step and somewhat probable if you've already done one step that you stop after two steps and so on so you give it sort of a default probability of stopping after each step so if this is 0.1 for example you tell the network essentially look at any given step there's like a default 10% chance that you should stop I as a designer of the algorithm think that's a reasonable prior to have now the network can decide differently the network can decide no no no no no I actually want to stop way earlier right like like this it puts much more emphasis on the first steps which of course in turn because you need to normalize put less emphasis on the latter steps so the network can still decide to violate this prior if the if it may reduce the loss for enough so this is as I said a trade-off there are two hyper parameters the geometric distribution shape and the amount that you regularize by this KL divergence and yeah so now we come into the experimental results and these are pretty pretty neat because yeah they I think these are straightforward experimental results they're not super big large-scale results or anything like this but they show that look on tasks where we sort of know that this dynamic computation has an advantage our model will outperform both previous attempts at dynamic computation and especially networks that have no dynamic computation built in whatsoever so this is the parity task which we're going to look at as you can see here the orange is this a CT which is the previous work that they compare most with that is most similar to them you can see in terms of accuracy pondir net beats this network by quite a bit also appreciate the error bars in this one they almost overlap but they don't so you can say that you're definitely better and interestingly the number of compute steps even though yeah the error bars overlap as well here but pondir net itself needs less compute steps which might be you know I don't I don't know why why exactly that happens but you can speculate that it is because pondir net sort of fixes on a single like it outputs a single answer whereas the a CT it outputs this weighing of things and therefore when it when it outputs that say the first step answer it always needs to consider that this needs to be compatible with potential future steps so just formulating so just formulating how a CT output stuff it seems like it becomes a lot less dynamic because the output is always a waiting of different outputs and therefore the first steps they have to they can't just output what they think is the correct solution but they sort of already have to incorporate the future and estimate well if I'm going to continue computing then you know there's going to be stuff added to my output right here and they have to take this into account so it can be ironically less dynamic of a network and that's why I think pondir net might need less steps here I might be totally wrong though so this is the parity task and specifically they train with string lengths between you know so this is a string length of one and then string length of we've before we had like eight right something like this so they train up from one until 49 lengths one until 49 and this is a little bit important I think because their training set contains all of them which you know this is a little bit of an experimental trick right so in order for your network what you wanted to learn is kind of the general principle of parity independent of string length so you construct the training data set to be sort of a distribution of lengths of string rather than just strings of a fixed length and then you assess their parity so yeah that that's maybe a bit of a lesson for if you do experiments construct your tasks themselves already such that they help find the correct solution right so they train with strings of length one up up until 49 and then they try to extrapolate which is this B right here so this is extrapolation where then they test so first here they test they train on small strings they test on small strings here in B they train on the same small strings up till length 49 but then as I understand it they give it length 50 to what 99 or so in 2 or 96 it says it somewhere just longer strings that it has been trained with right and now that the setup is you know clear it's clear why they did the different length strings in the training set and not just fixed length strings because there's a reasonable chance the network does not learn to extrapolate just from one particular or two particular lengths of string nevertheless they test how does the network extrapolate to longer strings and you can see right here that a CT even though it also has been trained on the dynamic length strings it is that's 50% right that's pure chance so it's a parity test right it's the output is either odd or even so a CT just gets a pure random chance as a result whereas the pondernet as you can see has like an accuracy of 0.9 which I guess is pretty good especially on strings that are so long you've never seen them so what can we read from this I'm not exactly sure there's always the possibility that you know they've just trained a CT wrong or something like this but it's also it's also reasonable to say that just how the previous models were constructed either they didn't learn the concept or their their output is just weird in the way a CT is or since a CT has biased gradients estimates and pondernet doesn't yada yada we don't know what we do know is that in their experiments this pondernet was actually able to solve the extrapolation task right here the interesting thing is that if you look at the number of compute steps done you can see that pondernet in contrast to what it was trained with during inference sorry that's an alarm in in contrast to what it was trained with during inference during inference it has like two point between 2.5 and three steps let's say three steps computes for about three steps during inference time that's what it decides on for the smaller strings yet the same model right train on the same strings this is the same model during inference time on the longer strings all of a sudden it raises its compute to five steps whereas a CT okay a CT doesn't work in the in this one it just decides to stick around two or three steps as it does in training right so the authors sort of claim that this is good evidence that pondernet learns to solve the actual task right here and as the task gets more complex pondernet needs more steps to think about the task and this might be exactly you know what we saw that you have some sort of a string of zeros and ones and you learn during training you learn a how to take one of these maybe in multiple steps and get an output but now you all of a sudden you have a longer string right well so now what you can do is you can also learn an output for this one and now you have two outputs right and now you can learn a series of steps to transform the two outputs here into a single output and that might just need one or two more computation steps which is exactly what we see right here happening so it's a good it's a good indication that something like this is happening I would be wondering pondering one might say haha if you know how this actually happens like like what do the individual computation steps represent is it in fact a for example in this parity task is the network going about this task in a hierarchical fashion you know like like I've shown here is it something different is it going about it in sort of a purely recurrent fashion where even though we as I understand it we input the entire string at the beginning does it only look at the string position by position or you know how does this work how does the scaling behave in general if you know they only show small strings large strings but how does it behave in general as you go up the length and so on it would be really interesting to introspect this model a little bit more than simply showing kind of end results here of the individual tasks okay what they also find is that the hyper parameter how you regularize the shape we've seen this up here how you regularize this shape is you know that is a hyper parameter but it doesn't seem to be terribly important again they compare to a CT which has another hyper parameter that does the similar thing that regularizes the shape of the of the desired halting distribution which they call tau tau doesn't mean a particular thing in so they say it does not have any straightforward interpretation though I guess the authors of a CT might disagree but as you can see here so if I draw the the means there is a region where the tau where a selection of tau performs high though you have to say see that is all around sort of the same value of like 5e minus 4 or something like this and then for the other values that you might set it for it simply doesn't work at all so you the authors claim you have to hit this tau pretty correctly in order to even get the network to do anything whereas they claim in pondernet this variable right here first of all it's between 0 and 1 and not just an arbitrary value right because it's a probability and they claim that you know it kind of works for for most things except this one right here where essentially you bias the network to just output everything after one step so the trick is for the geometric distribution you have to take the inverse so one over this lambda p and that will give you the expected number of steps that the network would compute according to this prior so when you put in 0.9 that would essentially be a single step that you ask the network to do but for all the other things well you you judge for yourself whether this here is really good but what you can say is that look it goes from 0 to 1 so you have a clear range and for most of that range the the thing seems to work okay ish and what they highlight is even down here so even if they do this even if they said lambda p to 1 or sorry to point 1 which would essentially bias the network towards 10 steps that the prior is please do 10 steps of computation in this parity task as I understand it even for that point one you can see the network it doesn't do 10 steps it actually also goes towards 3 4 or 5 steps most of the time so the network learns to be sort of somewhat robust to this prior distribution I mean I guess that's also a function largely of the hyper parameter here where you trade it off we don't know the effect of that just from the paper but even you know even if they set that to really low it's it it of course then the network is kind of robust to the choice of the lambda p yet it's still good news because that means you would mean you wouldn't have to regularize the the model super heavily in order to get it to work okay they go into two other tasks right here again these aren't tasks that you might necessarily know they are tasks where this type of computation shines particularly and yeah as I said I see the paper more as sort of an interesting an interesting task an interesting niche tasks subtask you might say of of connecting deep learning and classic algorithms there are a number of things that I think you can do right here to extend this so it's completely thinkable that you know the loss might be a bit different that you don't ask the network to output the direct answer at each point but you know you might you might want to attach memories and so on at at these output nodes you might want it want them to output intermediate results or something like this another thing you could do is you could work with sort of adversarial losses instead of of you know kind of reconstruction losses or whatnot so you could you could have some sort of a GAN going on inside of this in order to decide on the on the stopping probability that there's lots of stuff one can fiddle around with this type of network and you can even think of crazier architectures I don't know hopfield like structures where you decide you know how far you iterate because you don't you may not always want to iterate until fixed points I don't know I'm just I'm just talking crap right now okay one last shout out to the broader impact statement of this paper what a beautiful beautiful piece of of writing so essentially they say well this enables neural networks to adapt their computational complexity to the tasks they are trying to solve you know neural networks are good but currently they require much time expensive hardware they often fail pondernet expands the capabilities they say look it you know it can do this it can do that makes it particularly well suited for platforms with limited resources such as mobile phones which is a good thing right it can also generalize better that means it's better for real-world problems and they say it we encourage other researchers to pursue the questions we have considered on this work we believe that biasing neural network architectures to behave more like algorithms and less like flat mappings will help developing deep learning methods to their full potential and that is indeed the broader impact of this work like that is that's the impact it had on me and that's the impact that it it should have yeah I'm not like at today's conferences that must might be kicked out because of course it doesn't say technology good technology bad technology biased but you know respect for that and that was it for me let me know what you think and bye bye
 Hello there! Today we'll look at PonderNet Learning to Ponder by Andrea Bonino,
 Jan Ballager and Charles Blundell. This paper on a high level introduces a
 recurrent architecture or a principle of recurrent computation for deep networks
 that essentially says the network recurrently computes its output at each
 step and at each step it can decide to stop now because it is satisfied with
 the answer that it has. The idea is that at a complex task you can compute for
 many steps because it requires many steps of thinking and then give the
 output and for an easy task the network can decide to output right away because
 it already has computed the solution. This decision can be done on a per
 sample basis so for each sample the network can decide when it's time to
 give the final output. This is not necessarily a paper that just
 makes something bigger and then pushes state-of-the-art on some benchmark and
 that's why it piqued my interest is that it tries to rephrase a little bit how we
 think about the connection of deep learning and algorithms like classic
 algorithms by themselves. Essentially this is a dynamic if condition in this
 algorithm that decides when it's when it's time to stop and I appreciate that
 you know it not everything has to be state-of-the-art pushing here this is
 simply a cool method to do something that's relatively new. Of course things
 like this have been done before and they are discussed at length in this paper
 how this paper is different from other papers that do similar things and it
 does push state-of-the-art just not on benchmarks that you might be super duper
 familiar with. But yeah it's it's a cool paper it's a short paper the idea is
 pretty simple and it appears to work and yeah that's exciting stuff. So we're gonna
 dive into this paper have a look have a look at what's new in this particular
 model how it works and as always if you have feedback leave a comment subscribe
 I'd be happy for that and yeah thanks for being here. Okay so in the abstract
 here they say that in a standard neural network the amount of computation used
 grows with the size of the inputs but not with the complexity of the problem
 being learned. So which is true right in a standard neural network you have a
 forward pass be that in a fully connected neural network where you have
 you know you have your input and then you go layer layer layer layer layer
 and then you have your output. This computation here is always the same no
 matter the input even in a recurrent neural network right you have kind of an
 input right here at the beginning you have a layer then you have an input
 again and then you have this that goes into the same layer and then you have
 the next input that goes into the same layer even a recurrent neural network
 usually usually just does the same forward pass. This is a little bit
 different if you have something like a language model that can emit at some
 point a you know a stop token or an end of sentence token at which point the
 computation essentially stops but it's a little bit of a different thing than we
 consider right here. Right here we consider a neural network that has to
 find the answer to a particular problem and we're gonna see the problems down
 but one problem that they present is the parity problem. So the parity problem is
 you get a string of zeros and ones I think there is also negative ones in
 there but I think they're a bit for a distraction and the answer you're
 looking for is as a whole is the parity so the amount of ones in this string odd
 or even right so this requires a let's say an integrated view of computation
 this is essentially a classic algorithm that you have to perform over this
 string and neural networks as good as they are in computer vision and speech
 recognition they are having trouble with simple algorithmic tasks like this so
 the idea of this paper here is that well it doesn't make sense to apply a neural
 network that always does the same amount of compute right I shove this sequence
 just like in here it doesn't make sense because you know if there is just a
 single one in the string and I see that right away I can give the answer right
 away however if it's a long string and it has a bunch of ones I might
 need to think about this problem for a while and thus adapt the number of
 computation steps I do in my head I might you know first if I look at this
 string I might first connect these two you know and then that's two and then I
 might connect these two that's two again and then I might connect these two
 that's four there's nothing here there's nothing here right okay four so that's
 kind of like one two three steps of computation so that's the the rough idea
 whereas this if the string was shorter and and more regular I might need less
 computation so they say to overcome this limitation we introduce ponder net a new
 algorithm that learns to adapt the amount of computation based on the
 complexity of the problem at hand ponder net learns end to end the number of
 computational steps to achieve an effective compromise between training
 prediction accuracy computational cost and generalization so we are going to
 see how they do this yeah exactly so they then they go into the the tasks
 their experimental tasks in this paper are sort of these constructed tasks
 where people know you need this dynamic computation they're not gonna they're
 not gonna compete on like image net or something like this so the majority of
 the paper is in in contra posing their model against this a CT model the
 adaptive computation time I believe so there have been previous attempts at
 doing dynamic computation time yet either they have so it turns out they're
 kind of finicky and this model here this pondernet model has a bunch of
 advantages they say they present pondernet that builds on the previous
 ideas it's fully differentiable which allows for low variance gradient
 estimates unlike reinforce so a couple of previous attempts have been with
 reinforcement learning so let's just learn the number of steps or when to
 stop using reinforcement learning and that as you might know is very very
 noisy it has unbiased gradient estimates which is also unlike other models in the
 past and yeah so they say this has consequences in all three in all aspects
 of the model in pondernet the halting node predicts the probability of halting
 conditional on not having halted before this kind of seems obvious but
 apparently that no one has done this so far so what do we need for an
 architecture for pondernet they say this down here essentially that's the
 architecture it's an inline formula which you know but that's the
 architecture so what you need is you need an input okay you need an input
 which is X your input and X is transformed into a hidden state this is
 let's say the hidden state at step one those two or you can also reformulate
 this as just a hidden state the hidden state is going into s the so-called step
 function and that's the recurrent function right here so into this step
 function you can put anything you want you can put like a CNN inside you can
 treat this as an LSTM since we're going to apply it recursively sorry recurrently
 and anything you want can be the step function as long as it can be applied
 recurrently so this step function is going to give you the next hidden state
 right so you can see it's a recurrent neural network however it is also going
 to give you the output at that particular point in time so y1 I guess
 that be here and it's also going to give you this number lambda n now what are
 these so from here you could apply the step function again you'd get h3 you get
 the output 2 and you'd get lambda sorry that's that's a 1 that's a 2 so it seems
 like it's just a recurrent neural network and if I were to put push this
 to the end right I go give my H H H and then at the end I get my Y N and I treat
 that as the output of the computation then it's just a recurrent neural
 network however as we said the network can in this case decide to stop anywhere
 in between for example if it decides to stop at this particular step then that
 would be the output of the computation so every computation step the network
 computes and a potential output a suggestion for an output and then it
 also thinks about whether or not it really wants to answer with that output
 or whether it wants to continue and to do another step essentially take another
 shot at answering the question because it doesn't yet have the correct answer
 and that's where this lambda thing comes in so the lambda is a probability of
 stopping essentially so here you can see the output lambda is a number between
 zero and one and that is the probability of halting this is the output considered
 that the network halts so whenever this is one the network will halt conditioned
 on the fact that it hasn't previously halted yeah it seemed as I said it seems
 obvious to formulate it like this because you can you know you can only
 halt if you haven't previously halted but apparently previous models have simply
 output a number that is sort of the probability of halting in general which
 doesn't give you a bias sorry an unbiased gradient if you try to back
 propagate through it so if you consider the lambdas to be like this if you
 unroll for an entire training run then you get we get the probability of
 halting at any particular step this one so this is what this is what the
 previous networks would have estimated directly however this network estimates
 these lambdas these ones here you can see how you can compute the probability
 that for example the network halts after three steps by multiplying up the
 probability that network has not halted which is this one at step one has not
 halted at step two and then the probability that network halts at step
 three that it given that it hasn't halted at the previous steps so that is
 a valid probability distribution it's a generalization of the geometric
 distribution and essentially it encapsulates a decision tree right so
 at you're at the beginning you can halt sorry let's go a halt or not or continue
 if you continue then again you can halt or you can continue if again you can
 halt or continue and so on and all of this so if you want the probability
 that the network halts after you know this the third step then you would
 consider this node which means that you'd multiply that you multiply up
 these paths right here and that's the probability that it holds after three
 steps okay so the network can output this lambda at every step if the lambda
 is high then the network halts of course at inference this is done
 probabilistically now at training time this is done a little bit differently so
 you I hope you can see at inference time you simply go forward and you get a
 lambda maybe the lambda in the first step is point one and then you flip the
 coin a biased coin right if if it comes up heads you stop with the probability
 of point one it comes up tails which is a point nine probability you continue
 then maybe at the second step it's it's point zero five so maybe maybe you stop
 but probably you won't stop and then at the third step it like comes up point
 nine the network thinks yeah I should probably stop here and you sample from
 that and yes you you might indeed in nine out of ten cases you actually stop
 there so that's inference how about training how about we train this thing
 during training what we do is again we input X our input into an encoder for a
 hidden state and as I said you can also input X all the time into your step
 function as you see right here but what you do is you unroll the network for a
 number of steps right independent of these output nodes independent of the
 sorry if the halting probability let's say we we unroll it for for five steps
 right here and at every point we get a output and a value y3 y4 this is lambda
 2 lambda 3 lambda 4 so at training we simply unroll until a given step now
 there are some technical difficulties with doing with unrolling for a finite
 amount of step like how do you normalize the probability distribution because
 essentially this tree can go on until infinity they find okay we we can simply
 unroll until kind of the rest probability the probability we haven't
 used yet is is really small and then just load that all onto the last step but
 these are technical difficulties that you really only care when you then go
 and implement however so we unroll for a number of steps and then our we consider
 all the outputs at the same time now this is one big difference I believe to
 one of the previous networks to this a CT so what a CT does is it always unrolls
 and then the the output of the network so for a CT the output of the network
 would simply be a weighted output of the lambda I y I so the output of the
 network is always a waiting between the different steps okay and the network can
 decide okay how do I want to wait the individual outputs whereas here it's
 different here the output is really either y1 or y2 or y3 or y4 and to in
 order to pack this into a single loss function what we can do sorry I should
 probably leave this in order to pack this into a single loss function we
 simply take okay what's the loss what would be the loss if we answered y1
 right what would be the loss and we weigh that by the probability and we say
 okay what would be the loss of y2 we weighed by the probability that the
 network output so now if we and so on so plus essentially we compute the expected
 loss given the probabilities that the network has output so now if we back
 prop this we back prop through these losses we have of course two paths of
 back propping so we back prop through the wise which means it's at some so
 there is a loss right and both these things and these things go into the loss
 right so the loss is well how bad is this times how probably it was so on so
 the back propagation path would actually attack at two different paths you can
 see so the back prop goes into why because you want the network to compute
 a a better output but the propagation also goes into the lambda because you
 want the network to get better at estimating when its output is good and
 when not this I see a little bit as a tricky situation because usually this
 this seems a little bit unstable just from experience from other papers and so
 on if you have a back prop through two different things especially that are
 appear to be multiplied together and that you know the network can now trade
 off one versus the other which might you might think is desirable right it can
 either choose to make its output better if it wants to keep the probability high
 of outputting this thing or it can just reduce the probability that it's going
 to output whatever it wants to output and you know then it doesn't have to
 necessarily make the output itself correct because the loss the loss won't
 be as high for that particular thing because the probability of outputting it
 is low so network essentially has a choice as I said this might be desirable
 but usually that's kind of unstable and I think this is just my personal opinion
 I think a lot of them why this might work might rest on whether or not or
 let's say the complexity itself of assessing of making why better versus
 adjusting these probabilities of course yeah so you see if the output y is very
 complex right then this you know the same gradient signal for that might mean
 much less than simply reducing the probability okay so if the output is
 very very complex right not the problem but just the output itself right how to
 arrive at an output if the output is an entire pixel map or something like this
 and that has dependencies and so on the network might just choose to
 always reduce the probability because it's like well how am I gonna how am I
 gonna make this better at all I don't know I can just reduce the
 probability I'm going to output this crap right and it will probably do this
 then for every you know single step which you know if it's complex
 problem makes sense but still that's it that would be a bit my my fear here and
 that this is not really discussed in the paper itself so I think the fact that
 this works might rely on sort of a balance of the of the complexity or
 information content that you get from the loss at the output node versus the
 loss at the probability node so okay enough about that so in yeah during
 training you simply compute the expected loss weighted by the probabilities and
 then you can back prop through that and I hope you can see the difference between
 these two one is a they both seem to sum up somehow the outputs weighted by these
 these factors however one considers the actual output of the network to be a
 weighted combination of outputs of the individual steps where the other one
 says no no no the network output is actually one of them we don't know which
 one ergo for the loss we need to compute the expectation of the loss that seems
 to be a bit of a let's just say yeah it seems to be a more reasonable
 formulation though in hindsight you can say many things are reasonable if they
 work better right yeah so they discuss things like maximum number of pondering
 steps and so on again which I think is a technical detail and this is interesting
 so there you have the training loss as we just discussed now we've discussed
 this part right here which they call the reconstruction loss because you have
 some kind of desired y and you have a y that comes from this and I was a little
 bit wrong here in my formulation of course the expectation you don't have
 you don't want to take the lambdas you actually want to take the probabilities
 that each thing happens which means that you need to compute this P number you
 know going along this tree as we did because the P is the actual probability
 that you reach that node whereas the lambda is only the conditional probability
 that you reach a node given you were at the previous node so yeah consider
 there that if you if you are crazy enough to implement things straight as I
 speak in the videos lucid rains shout out the second part of the loss here and
 you can see this is a hyper parameter so you you're gonna trade off two of two
 losses right here because right now we saw okay you can either continue or not
 continue and for the network you know it might actually be easier as I said if
 the loss of the output comes reasonably complex right here it might be easier to
 simply say well in this case I'm just always going to reduce my probabilities
 you might counteract this with having this number of steps not like maximum
 number of steps but essentially this term here is what counteracts that
 really there is a regularization term on these probabilities as you can see right
 here so we regularize with the KL divergence which is sort of a distance
 measure don't tell this to a mathematician it's a it's a divergence
 it's a sort of a distance measure between the distribution that the
 network outputs for the steps and this thing right here which is a geometric
 distribution with this parameter and this parameter lambda p is another hyper
 parameter so what does that mean essentially if you consider here the
 number of steps that the network thinks right think things for what you
 regularize for this distribution right here is a geometric distribution I'll
 go something like maybe no something like this so essentially a geometric
 distribution is set exactly computes this tree that we computed right so at
 each step you can essentially stop and the question is after you know this
 distribution gives you a indication after what's the probability that you
 stop after one step two steps three steps four steps considering the fact
 that in order to stop after four steps you already have to have made three
 non-stopping steps except in the geometric distribution the probability
 of continuing is always the same whereas in our network our network for each node
 and the tree it can output a different probability otherwise you know there'd
 be no point we can simply put in the fixed distribution now what that
 probability is of stopping at each point that's exactly this lambda p hyper
 parameter right here so you regularize for a KL for this which means that you
 tell the network look here is a a reasonable reasonable distribution of
 when you should stop so you should stop so it should be you know somewhat
 probable that you stop after one step and somewhat probable if you've already
 done one step that you stop after two steps and so on so you give it sort of a
 default probability of stopping after each step so if this is 0.1 for example
 you tell the network essentially look at any given step there's like a default
 10% chance that you should stop I as a designer of the algorithm think that's a
 reasonable prior to have now the network can decide differently the network can
 decide no no no no no I actually want to stop way earlier right like like this it
 puts much more emphasis on the first steps which of course in turn because
 you need to normalize put less emphasis on the latter steps so the network can
 still decide to violate this prior if the if it may reduce the loss for enough
 so this is as I said a trade-off there are two hyper parameters the geometric
 distribution shape and the amount that you regularize by this KL divergence
 and yeah so now we come into the experimental results and these are
 pretty pretty neat because yeah they I think these are straightforward
 experimental results they're not super big large-scale results or anything like
 this but they show that look on tasks where we sort of know that this dynamic
 computation has an advantage our model will outperform both previous attempts
 at dynamic computation and especially networks that have no dynamic
 computation built in whatsoever so this is the parity task which we're going to
 look at as you can see here the orange is this a CT which is the previous work
 that they compare most with that is most similar to them you can see in terms of
 accuracy pondir net beats this network by quite a bit also appreciate the error
 bars in this one they almost overlap but they don't so you can say that you're
 definitely better and interestingly the number of compute steps even though yeah
 the error bars overlap as well here but pondir net itself needs less compute
 steps which might be you know I don't I don't know why why exactly that happens
 but you can speculate that it is because pondir net sort of fixes on a single like
 it outputs a single answer whereas the a CT it outputs this weighing of things and
 therefore when it when it outputs that say the first step answer it always
 needs to consider that this needs to be compatible with potential future steps so
 just formulating so just formulating how a CT output stuff it seems like it
 becomes a lot less dynamic because the output is always a waiting of different
 outputs and therefore the first steps they have to they can't just output what
 they think is the correct solution but they sort of already have to incorporate
 the future and estimate well if I'm going to continue computing then you know
 there's going to be stuff added to my output right here and they have to take
 this into account so it can be ironically less dynamic of a network and that's why
 I think pondir net might need less steps here I might be totally wrong though so
 this is the parity task and specifically they train with string lengths between
 you know so this is a string length of one and then string length of we've
 before we had like eight right something like this so they train up from one
 until 49 lengths one until 49 and this is a little bit important I think
 because their training set contains all of them which you know this is a little
 bit of an experimental trick right so in order for your network what you wanted
 to learn is kind of the general principle of parity independent of
 string length so you construct the training data set to be sort of a
 distribution of lengths of string rather than just strings of a fixed length and
 then you assess their parity so yeah that that's maybe a bit of a lesson for
 if you do experiments construct your tasks themselves already such that they
 help find the correct solution right so they train with strings of length one up
 up until 49 and then they try to extrapolate which is this B right here
 so this is extrapolation where then they test so first here they test they train
 on small strings they test on small strings here in B they train on the same
 small strings up till length 49 but then as I understand it they give it length
 50 to what 99 or so in 2 or 96 it says it somewhere just longer strings that it
 has been trained with right and now that the setup is you know clear it's clear
 why they did the different length strings in the training set and not just
 fixed length strings because there's a reasonable chance the network does not
 learn to extrapolate just from one particular or two particular lengths of
 string nevertheless they test how does the network extrapolate to longer strings
 and you can see right here that a CT even though it also has been trained on
 the dynamic length strings it is that's 50% right that's pure chance so it's a
 parity test right it's the output is either odd or even so a CT just gets a
 pure random chance as a result whereas the pondernet as you can see has like an
 accuracy of 0.9 which I guess is pretty good especially on strings that are so
 long you've never seen them so what can we read from this I'm not exactly sure
 there's always the possibility that you know they've just trained a CT wrong or
 something like this but it's also it's also reasonable to say that just how the
 previous models were constructed either they didn't learn the concept or their
 their output is just weird in the way a CT is or since a CT has biased gradients
 estimates and pondernet doesn't yada yada we don't know what we do know is
 that in their experiments this pondernet was actually able to solve the
 extrapolation task right here the interesting thing is that if you look at
 the number of compute steps done you can see that pondernet in contrast to what
 it was trained with during inference sorry that's an alarm in in contrast to
 what it was trained with during inference during inference it has like
 two point between 2.5 and three steps let's say three steps computes for about
 three steps during inference time that's what it decides on for the smaller
 strings yet the same model right train on the same strings this is the same
 model during inference time on the longer strings all of a sudden it raises
 its compute to five steps whereas a CT okay a CT doesn't work in the in this
 one it just decides to stick around two or three steps as it does in training
 right so the authors sort of claim that this is good evidence that pondernet
 learns to solve the actual task right here and as the task gets more complex
 pondernet needs more steps to think about the task and this might be exactly
 you know what we saw that you have some sort of a string of zeros and ones and
 you learn during training you learn a how to take one of these maybe in
 multiple steps and get an output but now you all of a sudden you have a longer
 string right well so now what you can do is you can also learn an output for this
 one and now you have two outputs right and now you can learn a series of steps
 to transform the two outputs here into a single output and that might just need
 one or two more computation steps which is exactly what we see right here
 happening so it's a good it's a good indication that something like this is
 happening I would be wondering pondering one might say haha if you know how this
 actually happens like like what do the individual computation steps represent is
 it in fact a for example in this parity task is the network going about this
 task in a hierarchical fashion you know like like I've shown here is it
 something different is it going about it in sort of a purely recurrent fashion
 where even though we as I understand it we input the entire string at the
 beginning does it only look at the string position by position or you know
 how does this work how does the scaling behave in general if you know they only
 show small strings large strings but how does it behave in general as you go up
 the length and so on it would be really interesting to introspect this model a
 little bit more than simply showing kind of end results here of the individual
 tasks okay what they also find is that the hyper parameter how you regularize
 the shape we've seen this up here how you regularize this shape is you know
 that is a hyper parameter but it doesn't seem to be terribly important again they
 compare to a CT which has another hyper parameter that does the similar thing
 that regularizes the shape of the of the desired halting distribution which they
 call tau tau doesn't mean a particular thing in so they say it does not have
 any straightforward interpretation though I guess the authors of a CT might
 disagree but as you can see here so if I draw the the means there is a region
 where the tau where a selection of tau performs high though you have to say see
 that is all around sort of the same value of like 5e minus 4 or something
 like this and then for the other values that you might set it for it simply
 doesn't work at all so you the authors claim you have to hit this tau pretty
 correctly in order to even get the network to do anything whereas they
 claim in pondernet this variable right here first of all it's between 0 and 1
 and not just an arbitrary value right because it's a probability and they
 claim that you know it kind of works for for most things except this one right
 here where essentially you bias the network to just output everything after
 one step so the trick is for the geometric distribution you have to take
 the inverse so one over this lambda p and that will give you the expected
 number of steps that the network would compute according to this prior so when
 you put in 0.9 that would essentially be a single step that you ask the network
 to do but for all the other things well you you judge for yourself whether
 this here is really good but what you can say is that look it goes from 0 to 1
 so you have a clear range and for most of that range the the thing seems to
 work okay ish and what they highlight is even down here so even if they do this
 even if they said lambda p to 1 or sorry to point 1 which would essentially bias
 the network towards 10 steps that the prior is please do 10 steps of
 computation in this parity task as I understand it even for that point one
 you can see the network it doesn't do 10 steps it actually also goes towards 3
 4 or 5 steps most of the time so the network learns to be sort of somewhat
 robust to this prior distribution I mean I guess that's also a function largely
 of the hyper parameter here where you trade it off we don't know the effect of
 that just from the paper but even you know even if they set that to really low
 it's it it of course then the network is kind of robust to the choice of the
 lambda p yet it's still good news because that means you would mean you
 wouldn't have to regularize the the model super heavily in order to get it
 to work okay they go into two other tasks right here again these aren't
 tasks that you might necessarily know they are tasks where this type of
 computation shines particularly and yeah as I said I see the paper more as sort
 of an interesting an interesting task an interesting niche tasks subtask you
 might say of of connecting deep learning and classic algorithms there are a
 number of things that I think you can do right here to extend this so it's
 completely thinkable that you know the loss might be a bit different that you
 don't ask the network to output the direct answer at each point but you know
 you might you might want to attach memories and so on at at these output
 nodes you might want it want them to output intermediate results or
 something like this another thing you could do is you could work with sort of
 adversarial losses instead of of you know kind of reconstruction losses or
 whatnot so you could you could have some sort of a GAN going on inside of this in
 order to decide on the on the stopping probability that there's lots of stuff
 one can fiddle around with this type of network and you can even think of
 crazier architectures I don't know hopfield like structures where you
 decide you know how far you iterate because you don't you may not always want
 to iterate until fixed points I don't know I'm just I'm just talking crap
 right now okay one last shout out to the broader impact statement of this paper
 what a beautiful beautiful piece of of writing so essentially they say well
 this enables neural networks to adapt their computational
 complexity to the tasks they are trying to solve you know neural networks are
 good but currently they require much time expensive hardware they often fail
 pondernet expands the capabilities they say look it you know it can do this it
 can do that makes it particularly well suited for platforms with limited
 resources such as mobile phones which is a good thing right it can also
 generalize better that means it's better for real-world problems and they say it
 we encourage other researchers to pursue the questions we have considered on this
 work we believe that biasing neural network architectures to behave more
 like algorithms and less like flat mappings will help developing deep
 learning methods to their full potential and that is indeed the broader impact of
 this work like that is that's the impact it had on me and that's the impact that
 it it should have yeah I'm not like at today's conferences that must might be
 kicked out because of course it doesn't say technology good technology bad
 technology biased but you know respect for that and that was it for me let me
 know what you think and bye bye