High-performance functional bayesian inf...

Scale By The Bay 2018: Avi Bryant, High-performance functional bayesian inference in Scala

Recording: Scale By The Bay 2018: Avi Bryant, High-performance functional bayesian inference in Scala

you I'm avi I work at stripe and we're gonna talk a little bit about an open source project I've been working on at stripe for the last year or so Rainier which is a system for doing Bayesian inference in Scala now ok I won't change that slide yet all right I am NOT a statistician I am NOT a Bayesian inference expert this is something I've been sort of coming up to speed on in the last year or so and this talk really neither assumes nor really teaches you very much about the statistics right this is a talk with an engineering hat on this is to talk about implementation so there are some other talks actually at this conference that I think go more into kind of the stats side the application side of Bayesian inference I know there's one tomorrow sort of along those lines that is not this this is going to be much more implementation focused and so I'm going to try to do my best to kind of motivate the work without sort of you know teaching a course on on what Bayesian inference is right so you know to set up the problem with an engineering hat on kind of this this is I think the the the most salient kind of compressed down into like a couple of scala interfaces way to think about it right we have some objective function that objective function takes a couple of kinds of input one of the pieces of input is some data that we have observed in the world right this is some values we know and let's just say that's like an array of doubles I mean maybe it's actually a matrix or a tensor or whatever right but but you know basically is an array of doubles it's a bunch of numbers and we have some parameters that we don't know the values of and traditionally that that parameter vector is often represented as a theta so that's what I've put it in as there so these are unknowns and we have some objective function that can take the data and can take a given value for the parameters right so we don't know this is the value we're just speculating what if this were the value and it gives you back some kind of score okay we're not gonna worry too much about what that score means just it gives you back sometimes it score and often people think about this from an optimization point of view right so you have an optimizer and the job of that optimizer is let's find the value for theta the particular parameter vector or if you like the particular point in the parameter space that minimizes this objective function right you know and often you think of this as I've got like a regression problem right so I've got a bunch of data I've got a bunch of weights I've got a loss function and I want to minimize the law I want to pick the weights that minimize the loss function okay so that's an optimizer there's a closely related thing you might want to do which is a sampler so a sampler is not trying to find a single best value of the parameters a sampler is admitting that we have uncertainty here and it's saying I want to give you a bunch of parameter values a bunch of parameter vectors but I want to give them to you kind of proportionally to the objective function right so this is kind of fuzzy when we're talking about continuous values but basically you know if you take any bucketing of this stream of parameter vectors right and you look at like how many parameter vectors end up in a particular that should be proportional to if you like summed up the objective function in that right if you can think of batter it as like integrating under the curve of these spaces or whatever right but that's that's basically what we're gonna do so that's the job of the sampler it's it's very similar to an optimizer but but it emits more uncertainty so this is a very hard problem to just kind of do in general right because like this could be a very big parameter space and you just search through it right how do you do that so a very common approach to help you with that is to place an additional requirement on the objective function you say I need to this objective function needs to be differentiable and you need to be able to tell me for any point in the parameter space what the derivative of each of the elements each of the dimensions in the parameter space is right and we call that a gradient and so if you do that then there are sort of well-known like fairly efficient ways of doing this problem right so in the optimization domain right very often we do what's called gradient descent right we're basically you just like you have your point in parameter space you have the gradient at that point you move just a tiny bit in the direction of the gradient and then you evaluate again and you get our new value and a new gradient right in gradient descent just kind of keeps doing that until it finds that you know whatever minimum it can find this is what it you know people know deep learning or sort of have heard of kind of all the stuff around deep learning deep learning basically makes use of like versions of gradient descent it uses tends to use stochastic gradient descent Ori actually not looking at all of your data once because you tend have a very large amount of data in the deep learning space right on the Bayesian inference side there's an algorithm that's very commonly used that is related to this called Hamiltonian Monte Carlo I'm not going to teach you all Hamiltonian Monte Carlo but it's a cool algorithm intuitively the way to think about Hamiltonian Monte Carlo is that if you think about gradient descent as being like I've got this like valley and I'm dropping a ball and the ball is like rolling down to the bottom of the valley and that's what I want to find is like where's the bottom of the valley Hamiltonian Monte Carlo you imagine that the valley is covered in ice and you have a hockey puck right and I'm dropping the hockey puck and the hockey pucks kind of sliding back and forth along the walls of the valley and I'm just like strobing right I'm just like taking snapshots of like where the puck ends up and yeah it's like more likely to be low than it is to be high right because like it's hard for it to climb off the walls but it's not like it's just gonna be sitting at the bottom it's actually gonna be sliding around everywhere right and that's kind of that covers the uncertainty right and if it's like a very deep steep canyon then yeah maybe it is just sitting at the bottom but usually that's not what's happening right so this talk we're not gonna talk about the optimization case which is the case where you have a huge amount of you want to find a single point estimate we're gonna be talking about the sampling case which is generally you have less data and more uncertainty and so you know you want to be using a sampler in that case and we're not going to talk about samplers or implementation of samplers we're just going to focus on the objective function okay on the implementation of the objective function and the question is how do you represent it how do you optimize it how do you evaluate it and how you make it easy to build so you know you could write an objective function and you could do this and use rainier samplers with it where you just literally like type out a bunch of Scala code they'll like takes an array of data and our array of parameters and like gives you back like a value right first of all that's gonna be like hard to do super efficiently you're gonna end up doing a bunch of like really ugly stuff just like the direct bit array access and making sure to like not allocate stuff and it's gonna be a pain second of all you're gonna have to write the gradient for that function which is at best gonna be really tedious right so you don't want to do that generally so what is commonly done in this space is that you represent the objective function in what's called a compute graph or an expression graph right so this is a dag where you have your leaf nodes that generally represent your parameters this figure is taken from a recent paper on automatic differentiation okay automatically drive ingredients so you have a compute graph where you have nodes that represent things like parameters or constant values and then you have interior nodes that are just operations right like some of these two things divide them whatever you can kind of see how that's going and then at the very top you have your output right and if you have that compute graph there are well-known ways to fairly efficiently you just kind of mechanically crank through it and produce the derivatives right you can in sort of a single pass through the graph from top to bottom you can produce the derivative functions for each of the parameters again I'm not gonna talk about how that works but you know understand that it's like a well-known thing and it's kind of showing the like derivation on the left there of how that would work right so rainier uses a computer F and I'm gonna start telling you a little bit about reiniers design of a compute graph for this so unsurprisingly we use like a Scala ADT for this so we've got a sealed trait which we call real because these are all functions the real valued functions from some parameter vector of kind of unknown width to a single real value okay and there's kind of a little ambiguity when you're working in these they kind of feel like you're actually working with values but they're not their functions right and the two types of leaf nodes are variable right we don't know the value of this or a constant we do know the value of this and I've put these beautiful visualizations before below of those these two graphs one which is just I created a single new variable and bound that to the variable X or the value X and I'm representing with a theta the variable node there and the other which is the constant two so you know that one is a function that always returns to one is a function that always returns the X parameter right it's basically the identity function and just you know to sort of there will be a bunch of these visualizations of graphs so the little houses represent like basically a variable binding right so that's not really part of the graph it's just like a label right and then the other things are like like nodes like objects Scala objects that are nodes in memory so let's let's look at a slightly more complicated so our graph does also have operations so for example has unary operations the two most commonly used ones are natural log and E to the X so we've got like a log method and an X method unreal and any real that you have any graph that you have you can take the log of it or X you know take the X pit and it just creates a unary node right that references the previous graph and then has this operation right I mean this is hopefully fairly straightforward for people and let's look at the visualizations of that right so we have our X and our two as we did before and we're going to take the X of each of them right so you can see down at the bottom we have our old X label which is a direct reference to the x-parameter node and then we have our e^x label which inserts this Expo right and the arrows I went back and forth on which direction the arrows should go here because people always show compute graphs this way because the arrows in the direction of the data flow right so when I want to evaluate one of these I get a value for that theta X node right for that fader node and then it flows and it gets you know we take e to the X and then it flows to the output but this is the opposite of the memory references right so when you look at this arrow you can see that as a data flow going forwards you can also remind yourself that what is happening in memory here is that we have a node representing X that has a reference back to the parameter right so anyway in case people get confused about that I wanted to point that out other thing that's interesting about this is that for the constant we don't have an expedite and we don't have that because we know what e to the 2 is and so if we try to take the X of 2 it's like we don't need to add another node we can just like compute it right away right and in compiler terms you might think of this as just kind of constant folding right it's like where I know what this is let's just do it now ok let's do something just slightly more complicated which is that we're gonna take our e to the X node and we're gonna take the log of that right so two questions one what do you expect that to look like in the graph right and you know very reasonably I think you might say you would expect to see you know a a parameter node and then that's connected to an expo and then that's connected to a log node right that would be sort of the normal thing second question what is log of e of the X log of U the X is X right and so that's what we actually get you can see that the e of X has the X node but log of e to the X we just go back to a direct reference to X right so it's doing what I'm calling here an eager partial evaluation so as you're constructing the graph if there's some computation it can do it does it okay and this is a design choice that when I talk to people whether they're like Bayesian stats people or compiler people they're like you're doing what so first of all why do I want to do a partial evaluation okay so two things to think about with this one is that we are building these graphs once we are then running these graphs maybe a million times right because we're iterating over this trying to search through the parameter space and we want to just crank through this over and over and over again so it's really valuable if we can do some optimization on the graph at the time we construct it that we're then going to be able to make use of a huge number of times when we run it okay the second is you're not writing this whole thing as one monolithic mathematical expression you're gonna be using abstractions you're going to be using building blocks in in methods that help you with this and so for example you might have a method here I've got a method up there for the log of the probability density of the normal distribution and you don't really have to care about why you would want that or what that is you might already know but the point is it takes in like a mean and a standard deviation for a distribution and then a value that we want to evaluate against that distribution and it like does a bunch of stuff right so I could call that with three constant nodes and it could just spit out the answer as a constant right I could call that with the mean and the standard deviation being unknowns but the X being known and that would be able to you know actually that wouldn't be able to pre-compute anything I could call it with the X being unknown but the mean and the standard deviation being known and then it would be able to like constant fold all of those SD terms right and but not you know not the X minus mean part right so to be able to kind of take these think of them as kind of like templates of a computation and specialize them to the particular usage can be very valuable compared to just you know sort of always including the full generality of the graph even when parts of that graph you know are already known right why do it eagerly though right so one observation is that theta the parameters they change every time but X does not right our known observations are known they're not going to change between iterations and so one design goal for this particular compute graph because we are working in a space unlike deep learning where the data is relatively small is often very small was to allow for the possibility of currying that X parameter right we want to end up with a function that is specialized to the data and only actually takes the theta as input is just to say we're in lining all of our data as constants in the compute graph right there's a risk when you do this that your compute graph is going to blow up and become huge right but if we have optimizations that protect against that we want to apply them eagerly so that rather than blowing up into a huge compute graph and then optimizing it down into a small one we want to just keep it small the whole time right we can't optimize realistically kind of doing some crazy whole problem and we could but we don't do some crazy whole program analysis every time instead we do what are sometimes called peephole optimizations where we just look at kind of the local area around whatever is going on right now but we do it eagerly right so let's talk about yes Oscar separate and then applies this kind of combination backwards I'm not doesn't it seem like most of what she talked about would be preserved so it makes me think that I'm missing something um I think I mean the thing I can say with certainty is that I found it easier to implement something that was efficient and possible to reason about by doing it this way right someone else might find a clever way of doing it a different way and that's great and like poor request welcomed right but but this was the way that like I could get to a working system and I totally acknowledge there's there's quite likely path dependence in that right I can try to ad hoc rationalize it right but you know right so so so because of this crazy path dependence that Oscar is referring to right this is the crazy place that I ended up and let's talk about it okay so you know you would probably expect that the next thing I'm going to talk about after unary operators is binary operators right and here's the crazy things we don't have binary operators instead what we have is linear operators and what I mean by that is that the node type that we use for addition is what we're calling a line it's a linear combination okay so it represents any number of additions of terms multiplied by constant coefficients okay so basically we have a constant vector and we have a vector of like graph nodes and we're taking the dot product of them and that's the value of that node and so that lets us represent additions of anything and multiplications of anything by a constant and then we have the kind of next level up in the rings of that of what we're calling a log line for lack of a better name which is the product of exponentiations with a constant right so it's any number of graphs raised to a constant power and then multiply them all and so that's how we represent multiplications of anything or exponentiations to a constant and we do have actually one binary node in the system which is if you want like to take a variable and raise it to a variable but you effectively never use that it's just there for completeness right so what does this look like so here's a few examples of like simple linear combinations so I've got my e to the X I'm going to multiply it by 2 then I'm gonna add 1 then I'm gonna add in an X and then I'm gonna add in another e to the X okay and let's see what that looks like so these are the nodes from top to bottom right so first of all I've got my e to the X it looks like it always did and then I've got a line here that just has a single coefficient of 2 because I'm multiplying that by 2 right that's my 2e of X when I add one I create a new node right that just has another element that's just the one it's just a one constant it's not you know coefficient for anything if I add in an X I need to add in another element right that has the X and if I add in another e to the X I don't need to add any new elements I can just bump the coefficient up right so I had to e to the X now I have 3 e to the X and the really nice thing about this is this order invariant as long as I'm just adding things you know whether I like have to e to the X is here and then an e to the X at the end there or whether I have like I'm summing them all together at the front it can it can figure out to combine those terms and just like have a like merged coefficient for them right so as long as I can kind of stay in line space I can keep kind of rolling like terms in to these coefficients and that's really nice what else can I do well so this is another example where I just like introduce another variable right so I've got Y here and then I end up with like I'm not gonna spend a lot of time in this slide actually but I end up with 3 y at the bottom and like our crazy x thing at the top and when I merged them again it's just like you're just merging the vectors right you're just merging the elements of that vector and actually we do use Scala maps for this and so even though you don't see it on these they are sharing structure even in the coefficients although I don't know how important that is we can also scale them right so if you take an existing one of these lines and you multiply it by a constant we don't need to create a new node all that we do is we scale up all of the individual coefficients right and so that's what's happening so what about nonlinear combinations what about when we want to multiply stuff right so here I've got my like function of X I've got my 3y and I want to multiply them together and for good measure let's multiply another Y in there right so you can see that we have this other type of node this logline node that just has references to those three things and the coefficients are all one because they're all just raised to the power of 1 right so you can see there there's an optimization that a kind of could be doing and isn't which is that I've got a 3y and I've got a y and it could be merging that into a 3y squared but because of the kind of people effect it can't kind of look through the three yet to see the Y and so it's not doing that so I mentioned that as long as you're in line space that's really nice because you can just kind of keep summing things into it and in general the structure of these functions sort of tends to have like a huge number of terms that you're summing together because basically you take log probabilities and you know normally you have a lot of probabilities you want to multiply if you have log probabilities you want to sum them so anyway you end up wanting to sum a lot of things and if we want to sum if we want to add something to this product there's kind of a ace that we want to drive this as much as possible towards the output being one of these lines one of these some vectors okay and so what it does is it tries to multiply out tries to distribute all the terms as much as possible to convert this product into a sum so that may be the thing we're summing into it is going to just like fold in to some existing element of the vector and so when we do that we get actually this like crazy rearrangement of the top of the graph and you can see that it's taken what was before a single product node and we want to add something to it and it's turned it into a single sum node and a bunch of product nodes right and you can kind of see the expression before that down there to see what's happened is that we've kind of taken this 3y squared and distribute it across the X function so that now we have salt like six terms that are like e to the x times y squared we've got three terms that are x times y squared and we've got three terms that are y squared a Madol just rolls up into a line right so that's nice because now if we have more things we want to add again you know probably will come across these terms again right there's a decent chance we'll see like another Y squared somewhere and we'll just be able to fold this in let me show you kind of a simple example of that kind of thing happening right so this is we're taking an l 2-norm right like a like a squared error similar to that normal PDF we saw before and so we have say e to the X is the parameter that we think is kind of you know the mean of something here right we have a couple of data points three and five right so we're going to take the difference between x the e to the x + 3 and square it we're gonna take the difference between a of x + 5 and square it so individually those are like fairly straightforward you know nodes there you can see there's like the linear like you know X minus 3 and then it gets squared right alright the linear X minus 5 and then it gets squared but if you add them together it's gonna do this thing of expanding it out it's gonna like foil out the X minus 3 right - like x squared minus 10x plus whatever 25 and so you end up with a single x squared term feeding into a linear combination of course if that happens for two of them that can happen for any number of them right so going back to this idea about folding the data currying the data in if this is my loss function right and I have an arbitrary number of data points that I'm folding in I can pre partially evaluate it all down into this one compute graph of constant size and so if you think about what that means that means that the complexity of my iterations just changed instead of my iterations each being order n for n data points my iterations are now order 1 I have order and complexity to build the compute graph but I only pay that once right and this kind of looks like a trick but it actually does like come up in in real models right that you can do this and so that's a massive speed up when you can get away with it and so that's that's you know the the compute graph now the question is how do you evaluate the compute graph right so you can build in interpreter obviously right for this thing that would be very easy it would be fairly slow you could try to like you know I don't know if I compile the GPU kernel or something like that right one of the design constraints that we had was that we wanted this to be like pure native JVM no J&I libs no external dependencies on on complicated like see math libs or anything like that we're on the JVM let's use the JVM so we use the ASM Lib to generate Java bytecode for these functions represented by the compute graph and any time we're compiling the first thing we need to do is build an intermediate representation so let's look at that so at the top you have the real compute graph for for this function we were just talking about and below you have something very similar but it's the IR so a few things to note about this right one is that we've now gotten rid of the like crazy line thing and we're back to binary operations right and that's because like the actual machine actually deals with like binary operations right another is that instead of having a dag we have a tree right so the first time that we compute e to the X we store it somewhere we store it into a temp variable you can see it's labeled t0 and then the next time that we come across it kind of thinking like depth or yeah depth first in the graph we reference it right we say ok now we have a local variable reference to t0 you can see that it happens to have like made this choice that instead of using like you know t0 squared it just as T 0 times T 0 right just kind of a small optimization it also does a useless small optimization which is that it pulls the thousand coefficient out we had like a thousand e to the x squared and it pulls that out factors that out of all the coefficients there's no reason and used to do that if there were multiple things with a coefficient of a thousand it would actually save an operation for it to like pull that out but it's just being over aggressive and that's sort of a bug it's fun to see your bugs turn up in your visualizations sometimes so this is what by the way if you generate that byte code and then you D compile it right what the Java code looks like so this Java code never actually existed we didn't like generate the byte code by like producing Java code and compiling it just to be clear but if you throw a T compiler at the byte code then this is the Java code that are producing you can see it's pretty straightforward right you have like this function it takes params array and it like pull something out of the prams array and you know math dot X and stores it an attempt and then it uses the temp a few times and very simple math operation right and returns it right this is just like very simple Java code which is nice because of the JIT likes very simple Java code okay another thing about the JIT interesting fact if your java method gets larger than eight thousand bytes larger than 8,000 byte codes it will refuse to get it and it will be interpreted and it will be very slow that is bad if your compute craft gets large you might produce a method that is larger than 8,000 byte codes and you do not want that so we have to be able to split these up so this obviously is not getting close to that range but we can pretend that we have a smaller constraint and see that we can split it up into multiple methods there they're just calling each other that's what those like boxes mean the larger boxes are the methods and we can see that the main change there is that where we had a T 0 we now have a G 0 so we're saying this is an intermediate value that is used by multiple methods so we have to store it in a global what do we mean by global what we mean is we do this very simple thing which is very easy to implement of just allocating a Global's array that we pass in as an extra argument to every one of these methods and they just use it as a scratch pad so you can see the highlighted line there when it computes e to the X it stashes it in global zero if any other method and this chain needs e to the X it retrieves it from global zero that's all we do we still only return a single value it's just kind of this like side channel that we can store stuff into right and that turns out to be very effective and lets us with the methods up right all of the methods have that identical signature of double array of params in double array of Global's and returning double ya we also have to split between multiple classes because class individual classes can get too big but it doesn't really matter it's just like static it's all static final stuff right the other place the Global's get used is in gradients so the gradients have separate methods compiled for them but almost always in fact probably basically always the gradient needs to refer to intermediate terms that were computed in the forward pass and so the forward pass will store those intermediate values in the Global's array and then we passed that same Global's array into the gradient methods and they are often like this look much much simpler they're just pulling a few intermediate values and doing some combination of those and then that ends up being here being your derivative right so is it fast after doing all of this the comparison point for me is Stan stan is a Bayesian inference system very well regarded very commonly used by scientists that has its own DSL that generates C++ from and then compiles with GCC okay and because of that it runs pretty fast right GCC is very good at generating efficient machine code turns out that the Java JIT if you give it like simple like double operations is also really good at generating efficient machine code the advantage that rainier has is that because stan does this very slow thing of generating c++ and compiling with GCC they do not look at the data because they assume you might run this on like multiple different data sets right and so they have to compile it ignorant of the data that they're gonna feed through this whereas rainier compiles extremely fast and so it's really no skin off our back to recompile it every time you give us a new data set and so we can specialize the compilation to the data and so what we've seen is that for many cases this makes us much faster than standing right so we see like a 10x speed up over stan on the other hand the JVM never uses in the instructions approximately and there's also some like fancy your algebra libraries that Stan uses that that can you know really optimize like heavily vectorized code and so we've also seen to go the other way where Stan is ten times faster than Rainier this is all benchmarking hard to take with a grain of salt right but but that's the kind of thing we've seen I was very happy with this result I did not go into this expecting to be anywhere near Stan's performance and so it's cool that this actually you know worked out I'm gonna really shift gears here and I'm gonna do a quick time check before I do that so I think we've got about 15 minutes left for me to talk which is probably pretty good for this so we're gonna start to talk about the functional in in the title right I mean I want to talk about modularity okay this is my like you know kind of hand wavy diagram of what deep learning models tend to look like okay is actually I was very gratified to see that almost this exact same diagram appeared in Tonk phase talk earlier if anyone was in that one so basically what deep learning models look like is you've got like some layer that creates a bunch of parameters and like combines them with the data in some way and then those feed into some other layer that like creates a bunch of parameters and combines them with the data in some other way and then you get out like some prediction which we're gonna call Y prime right and then you compare the prediction to the labels that you have from your data and you feed those into some loss function right that's kind of what it looks like in the salient thing from my point of view about this is that the loss function is almost an afterthought right you do all this work without thinking about the loss function at all and then you can just kind of like layer the loss function on on top and actually when you're using a deep learning model that you've already trained you drop that part of the function entirely out right it's like optional and stacked on top and that's great right that gives them great modularity properties in all kinds of ways Bayesian inference doesn't do it exactly that way because in Bayesian inference every time you create a parameter you have some prior right you're saying yes this is a parameter but these are some constraints I want to put on that prior parameter I have some beliefs about what this parameter should be that have nothing to do with my data right and so your loss function or you know you would say a joint density function in in Bayesian speak right that gets kind of smeared across your whole model right basically every module every layer or whatever you want to think about it has not just contribute some you know parameter values or transform parameter values or whatever it is feeding forward you know into the rest of your model but it's also producing off on the side some part of your objective function that you're gonna need to kind of propagate and computer and ultimately add up and that's your that's your objective function at the end right so this is maybe not intuitive but this has some modularity problems right because you don't really care about that addition to that you have to track it but you don't really care about it it's like it's just a term that's gonna get added in at the end it can't impact the rest of your model right and so the standard way for people to deal with it this is what Stan does is what Pines III doesn't so on is to have some global state and have some side effect where at your module that's creating some layer or parameter or whatever it is just kind of like plus equals you know it's like oh here's my objective function like plus equals this right so we can kind of like imagine world where rainier does this right we have some class that represents a normal distribution has a mean and a standard deviation and i want to just ask it for a parameter right like give me a normally distributed parameter thanks so it can create a variable node right great that's your parameter i'm gonna return that that's what you actually need going forward and i'm just gonna kind of like off on the side build a little compute graph that computes the like normal PDF of that and like stick it on some globally maintained like this is my objective function right so I don't like that I think that's a perfectly valid thing to do but you know I'm used to a functional programming style of Scala where that is not what I want to be doing right so of course everybody knows the solution of that is to introduce a monad right so rainier has the random variable monad and this is sort of a skeleton version of the random very old Monette so random variable and has two parts it has the value which can be any type and that is like this current state of our evolving model right is the way to think about that and then it's got a density which is a compute graph and that is like the thing we want to hide from you right which is like the accumulation of the objective function kind of off on the side and so map method on random variable just changes the state right just changes that value like if we're mapping we aren't changing the objective function at all we're just transforming in some way the the state of the model that we're trying to build here a flat map at take you know has a function that that goes from our state to a new random variable and there we do change the objective function but we change it in this way right we take the old objective function from what we started the flat map and we sum it with the new objective function the new piece of the objective function that we've gotten from this new random variable right so it looks like this so we create a new random variable with the value the new state from the new random variable but the density is the sum of the two right and that's basically it right everything else kind of follows from this formulation of the random variable monad so now we have no side effects so here's kind of like an equivalent of the previous slide and it's kind of ugly because I'm just doing like you know flat maps without a for notation or anything but you can kind of get the idea right we return from the normal distribution instead of a real we return a random variable real which includes within it the density the the the normal PDF and we can transform that thing with a map we can use a flat map to compose in some new distribution that depends on it right and and so on right whatever we want to do we also want to be able to combine these things again this kind of follows from flat map but from a convenience point of view it's worth pointing out you can implement zip right so you can take two random variables you know one parameter over here one parameter over here zip them and get a random variable where your models day is like the tuple of these two parameters or you can do use a traverse where you you know take a sequence of random variables and you end up with a random variable that has a sequence right so you know combine and compose these and every time you do that you're gonna be summing up their log densities right you're gonna be multiplying together their probabilities in your model one point to make about this is that it one unusual choice relative to like deep learning frameworks about reiniers compute graph is that it's scaler right usually people do these as tensor graphs and the problem in a scala context is that we are not used to working with tensors as data structures right we're used to using Scala data structures like sequences and maps and case classes and that kind of thing and like Tonk they did an amazing job of a typed tensor but what I was trying to do because I can have the luxury of doing this with somewhat smaller amounts of data is to say we're just going to completely use Scala data structures everywhere and the other interesting thing is that no matter what data structures you use they all just fall out at the end because all of this monads stuff right all of these flat maps and maps and zips and whatever else is actually just running once to construct this real compute graph that has none of them in it right none of the data structures appear in the final iteration code they're just a way sort of idiomatic and familiar way for you to build this compute raft that then gets compiled to like you know totally just like doubles and arrays Java code right or Java bytecode so you know as another piece of this we have a trait and this is quite simplified from the actual implementation but just to give you an idea called distribution and here we're getting a little bit more into like the Bayesian modeling stuff that I probably haven't quite like justified in a Forex set enough context for but we're getting short on time so the idea is that you have a distribution object representing things like uniform distributions and normal distributions or beta distributions or Poisson distributions or what have you and it can do two things it can hand you a parameter that you know your prior is it has that tribution or you can fit it against some data and so below is kind of you know maybe the simplest possible like full model which is in a in a for notation right which is that I have some mean that I'm saying is drawn from uniform distribution I have a standard deviation that I'm saying is log normally distributed which is to say it's drawn from a normal distribution and then we we take e to the X and then I'm gonna fit that I'm gonna construct a normal distribution with that mean and standard deviation and fit it against some data and so what I'm trying to do is find what the what the mean and standard deviation of that data is now of course you know this is probably not the best way to do that there are like other like you know closed form ways of doing that but if you you know happen to have a sampler to throw it it like this is a way that you can do that right so that gives you an idea what what actually using this stuff in practice looks like you might notice that it doesn't look that different I mean if you've worked with probabilistic models before it actually doesn't look that different from the way people like write down probabilistic models it's a fairly natural notation just using like Scala's like for flatmap notation right one subtlety that is interesting to point out so if you've got a diamond situation here where I've got say one parameter that then feeds into another parameter you know a couple of different parameters and then those ones come together to make my prediction or something right each of those at the time I created it you know has a prior so the way that I specified this earlier we would end up double counting that top parameter right because every time we combine something or transform it every time we flatmap we sum the previous thing right and so if you kind of walk the graph right we have two paths to get to that top one and so we would double count its prior and so we don't quite just sum what we do is we have a target object that is kind of a like identity wrapper around a compute graph and what the random variable actually holds is a set of those and we union those sets right so the only thing that achieves is that if we come again to a random variable or something that you know kind of was a random variable we've seen before we don't double count the other thing we have the opportunity to do with that target is in the case where we determine kind of just empirically that folding all of the data in is a bad idea right it's like we've got like a million data points and they're not reducing nicely to a compact compute graph we can bail out of doing that and just include it basically is what like tensorflow would call a placeholder right so we have a place in the target to just embed a bunch of like arrays of data associated with variable nodes in the graph and sort of say you know to the sampler like when you're evaluating this thing you're gonna need to evaluate this part of the compute graph over and over and over again with like this data right and so it's kind of like a very constrained loop that lets us scale to larger data sizes than we would be able to if we were actually trying to fold always all the data into the computer graph one last topic because I seem to have time for it is when you're sampling you don't necessarily just want to get out raw parameter values raw weights you probably actually want to use those to do something maybe you want to use them to make prediction and so the better way to think about the signature of the sample method is this bottom one right which is that I have some data I have some objective function and I have some sort of prediction function or output function which goes from parameter values to some value T that is what I actually care about right and when I sample I don't actually want to stream parameter values what I actually want sometimes is a stream of predictions right which can be of any type and so how do we kind of achieve that in this framework so we have one more type in here which is a generator and a generator is a way to sort of go from the abstract kind of real compute graphs back to a concrete value and it does that because it has this get method that gets two things one is a random number generator because sometimes these things are stochastic right and the other is a numeric instance on real which is to say a way to take a compute graph and get to double right in this one context what is the double value of that real and we're gonna call this generator with a bunch of different you know instances of numeric real and those will give you different doubles for it every time right and so we have a sample method on random variable which says if you have a way to convert my model state here into a generator then I will let you sample from this and what you will get out is a stream of the values this generator produces and also distribution any distribution like normal you can get a generator for right so let me show you sort of you know what that means in practice so this little model that I had before right the what the model gets out is a random variable with this tuple real real for mean and standard deviation board I want to get out when I'm sampling is a tuple of double double for for you know the actual values of that mean and standard deviation and so we have implicit that will take kind of any combination of reals like a tuple of them or a map of them or sequence of them or whatever and just you know map that numeric over them and give us doubles right so now we can take this model which is expressed in in terms of compute graph kind of abstract real objects and get a stream of actually concrete values out of it or the fit method returns the distribution we're fitting against and we can take that distribution and get a generator from it so that what we're doing here now is doing like forward prediction right what what people in deep learning sometimes call like hallucinating right so we have our distribution of the data we fit that distribution against our past data and now we're going to generate more data from that same distribution and so that gives us predictions right because those are the ways in which in which generator gets used and so I was giving a talk on on all of this earlier in the year and I said something about you know this is actually like three monads composing right so we have what I believe is the writer monad which is random variable wrapping around basically a free monad because that's basically what the real graph is right and also wrapping around generator which also has like map and flatmap although we didn't show them and is the probability monad right so if you like monads like Oscar does you got lots of own ads man so we're at time we have apparently a few minutes for questions we just cut a new release today with a bunch of the stuff I talked about on other stuff if you go to github stripe Rainier let me know if like the jars aren't actually unmuted and central or something or something is breaking because we just cut a new release and so it's probably you know who knows what the state is but but it's up there and does anyone have any questions yeah yeah so one of the reasons we wanted to run this on the JVM was that we wanted to be able to really easily run this on like a Hadoop cluster or a spark cluster without like any deployment hassles to be clear any individual model that you're sampling from runs on a single machine in this case what we wanted to be do do was parallel inference of like a hundred thousand different models independent models right and so that's why we want her to distribute across a cluster but this does know distributed inference and in fact it's a research problem right now to do this kind of thing in a distributed way right and it's you know as far as I know at least in the scientific community is a relatively rare to do that and people tend to just have like single big machines and run it for a long time right yeah there's things that real is old enough double so you just not be the objective function the compute graph is always always doing computations in double space yes the API for things like fit so we do support the idea of having say a distribution that has different types so you can have discrete distributions you can compute their likelihoods right your data can be in soar bowls or whatever you want but the parameters are all continuous right yeah someone at stripe has been hacking on a variational inference so for people who don't know variational inference is a way of turning the sampling problem into an optimization problem by saying well I can approximate the distribution of the parameters with some parameterised analytical distribution and then optimize its parameters I don't know if that made sense to anyone but so far we're pretty focused on Hamiltonian Monte Carlo is the sampler here and were you know probably if we want her to do variational stuff we would just use like tensorflow or something right but you could and and there is someone who's hacking on it in their spare time and we'll see if that that comes to anything yeah yeah no that's a very good question so there's a lot of like values of F that I could error or this that I could answer that question for the current application production application within stripe of this kind of Bayesian inference is basically cash flow forecasting for our customers or of our customers rather which has a bunch of implications for us for managing our own risk right and where us having kind of distributions of outcomes that we can then so the the basic framework that I really like of of Bayesian inference is you don't have a you don't have a classifier that tells you what to do right or a model that tells you what to do you have a model that tells you here's the distribution of the possible worlds that could exist right and then you have a bunch of actions you might want to take and you say for each of those actions if I pair this action with this possible state of the world is that good or bad and how good or bad isn't right what's my utility function and then if you integrate over the distribution of possible worlds for each of those actions then you get you know what the expected utility is for that action right and that for me is a really nice framework for thinking about how to make decisions in the face of uncertainty and has been really valuable for us for these kinds of cash flow problems so that's kind of why do we want to do this kind of modeling right why did we write our own system to do this the other production system that is very so there's two other production systems that are that are kind of commonly used for this one is Stan which is largely used in science and like has a dependency on generating C++ code and GCC and it's just like I've talked to people at Facebook who use it in production and they just like hate the deployment properties back right or even the like getting a developer up and running with is just it's bad the other is PI mc3 which i think is a great system but we don't tend to use Python for production systems at stripe we use it for like ad hoc model training that kind of thing but this had more of the profile of a production system and so python didn't feel like it was a possible choice also our early benchmarking was that we could do this much much faster in a Scala system than PI mc3 and we were doing enough different model inferences that that that actually seemed to matter for us the third reason is that I totally over fit on my experience at Twitter with scalding which is that like we built this like Scala data flow system that ended up like everyone at the company ended up using and so of course the correct thing to do in any case is to build your own Scala data flow system right I'd you know [Applause]