Pipeline-level differentiable programming for the real world

Speaker 1 [00:07]

Thanks a lot everyone for coming here. Yes, I'm Alessandro, I work at Pasteur Labs, and at our company we're trying to bridge the gap between recent advances in academic research and use those into real-world engineering problems. And you see a few keywords here, like in my title, like differentiable programming and pipelines in the real world. And I hope in the next few minutes to show you, first of all, why these things matter. And before we do that, however, I want to show you what differentiable programming is. So I think in this community especially, a lot of us train the neural networks, for instance. And so a lot of us are already familiar with what differential programming is right like you know you define a neural network you just say what each layer does but you didn't need to code and cook up like your derivatives right but i want to show that in general uh differential programming is a concept that can be applied to a lot broader scope so let's take for instance this example here where we have a rendering of a picture where you see this uh coral-like structure that casts some some shadows and so on this is interesting and And the classical renderer, as an output, will just give you the output of the scene, given your parameters. And when you use, instead, a differentiable renderer, not only you get access to the rendering of this picture, but also how this picture changes when you vary some parameters in this scene. So, for instance, here on this right panel, you can see how will this picture change if you slightly rotate this color-like picture. And another instance that to me personally is like extremely interesting is when people try to apply machine learning to simulations of physics and in this case we see a recent example from Niels Tuery group where they develop basically a hybrid solver. So you want to simulate some physical system and you know you can have these classical methods like with numeric solver and as you all know this tend to make a lot of they have some sources of errors and one could have the idea that basically then you can just apply a neural network to try to learn these errors and then correct them before you you know apply the solver again to the states that you have to find the next state and here you can see like you know the result of a very coarse simulation And then when it's, you know, corrected in the neural network, this gives you a much better result. And something that is interesting here, however, is that when you want to train this network and you want to train this inside this loop that I just described, you will have to backpropagate your gradients not only through the network itself, but also through the solver. So your solver itself needs to be differentiable. and these examples here on the left can seem relatively simple examples but these ideas have been scaled to really complex systems like the whole atmosphere of the earth so this was a also a recent state-of-the-art model in the climate sense it's called neural GCM from Google and I mean here you see a picture at a high level of the various components of this model you can see that there are things like, you know, all these solvers, encoders, decoders, and the important part here is that all of that is fully differentiable, and this was kind of like a crucial component in making this work. And a lot of differentiable programming is powered basically by automatic differentiation, so let's recap a few salient features that, you know, are going to be useful for the rest of the talk. So in particular, automatic differentiation is a method to compute derivatives with your computer, typically. You can do it by hand if you want. And you can compute them accurately and efficiently. And it's very important that you don't confuse AD with neither symbolic differentiation, which is probably what you would do on a piece of paper if you ever do calculus, where you take a formula and manipulate them, apply some rule, and find an expression for the derivative that then, after you do these symbolic manipulations, you plug the values into. And it's also not to be confused with numerical differentiation, because that's like you know what what you do when you perturb slightly maybe the inputs to a function and then like look at how the outputs change and then by doing the ratio between these two things you uh you understand how the derivative you estimate the derivative instead it works in a funny way which we will see in the next slide but basically you would break up the computation into simple steps and then sort of apply the chain rule automatically as this thing runs and there are let's say two main flavors of automatic differentiation The one that probably a lot of people here are familiar with is reverse mode, because if you think about how a neural network works, then, you know, you first go from the ways of the network and the inputs of the network to, you know, the output of the network, and then to calculate the gradients, you back propagate. So you start from the end, and then you go to the start, right? There is also another flavor which is called forward mode, and this is instead where you start with your inputs and weights and something, and then, like, as you go down this computation, you calculate sort of simultaneously the derivatives. They have pros and cons, but let's say you, like, they are both very used. And as a very simple example, I promise this is going to be the part of the talk with, like, most formulas, then we're not going to see any anymore. So I took a very simple function just to really show how AD works. So we have this function, like, sine x times 2 plus e of sine x. And then, like, really, we start with x. Then we define a variable called a, which is just the sine of x. Then, you know, in b and c, one time we multiply it by 2. Another time we exponentiate it. And then we finally arrive at y, where we just sum these two things. It's, like, really simple. And you can see each of these steps are sort of elementary operations. And to calculate automatic derivatives of this thing in the forward mode, you will basically need to build a sort of copy of this graph. let's not focus maybe right now on the details but fundamentally you see that you can start from x you calculate a and then like you know a has this green arrow here which is like cos x and if you remember that cosine is the derivative of sine you see that that's how the derivatives basically propagate from this blue graph to the red one and then in the red is like how the derivatives basically gets calculated and you know forward mode as i said you start from the top of this graph you go down and you calculate derivatives at every step. Now, enough stuff recapped, let's try to then look at some real-world scenario where like all of all these things that I talked about can be interesting and let's imagine that you have some real engineering problem like you want to optimize the shape of a wing. So how could one go about that and like you want to optimize it means you want to optimize it in some scenario so here maybe when it's impinged on by some airflow and then also you know you typically would like to optimize maybe some sort of a loss function on this like maybe a function of drag lift weight of the wing and so on so very ideally the pipeline could look like something like this where you start with the shape parameters that define really the you know the shape of your ring and of course though this is very specific to your problem you cannot just give this value to a solver and and hope that this works. So you need some intermediate step of mass generation where you take these parameters and actually generate a 3D shape that corresponds to your wing. You pass this shape to a simulator that then takes this shape, knows the boundary condition with airflow, and so on and so forth. And then from this initial state, it generates the whole trajectory and dynamics of this thing. This is a lot of data, so probably you will need to calculate some function of this, maybe average stuff, and so on. And you can calculate from these states, this drag, lift, and so on. And finally, you arrive at the loss value. So this looks really simple, really linear, and also not linear, but like it's a line of these linear functions, right? And it would be nice to just calculate then the derivative of that loss value with respect to the shape parameters. And then you know how should I tune my initial parameters in order to get a slightly better wing and do this iteratively, and you get your optimized shape. And as I said, this is ideally true, because in practice, there are a lot of challenges. And I want to go through some of them. I'm sure, though, you can get a dozen more by just thinking a few minutes. But these are the ones that are closer to my heart. So first of all, some of these components might not be implemented in a differentiable way. The most typical thing that we see is, for instance, that the simulator is not component that you might want to write on your own with JAX, like maybe you're doing computational fluid dynamics, and you want to use some off-the-shelf solver, and you get what you get, like a black box solver that doesn't expose derivatives. Another issue that can arise here is that all of these components could live in distributed environments, so maybe they are mediated by HTTP requests, and then your AD framework doesn't really know how to process that because the derivative of an HTTP request is not a well-defined concept. And these are technical things. I think there are also sort of like sociological sort of challenges here as well, because if you think about the people who write some of these components, right, maybe the part that is mesh generation, maybe some really intense computational geometry niche that like write some meshing algorithms maybe, and then you have the simulator written by people who are really deep into CFD, and then, you know, the person at the outset is maybe an optimization expert and needs to understand all of the conventions and all of the sort of like assumptions that the people writing this code made into doing this. And it's really hard and there is no, I think it's probably the hardest problem to solve because technical problems, they can be solved like these other ones take a bit more effort. And also, you could have different frameworks between all of these components, like maybe some of that you want to really write in Julia and others in JAX. So you have all of these issues. And then also, I mean, this is a big classic in computer science, but all of these components could be having conflicting dependencies, both Python and system-wide. So yeah, you will need to containerize these things. But again, people writing a lot of this scientific code are not necessarily Docker experts or in general know a lot about deployment. And again, a lot of this stuff that I said is not a showstopper. For instance, the fact that the simulator is non-differentiable. If the physics is differentiable, again, you could sort of write your own logic to propagate derivatives. But basically, it makes it really hard to just use the AD framework that calculates these graphs for you. So here, I just show them for a simple operation that there are already a lot of arrows. And yeah, and it will seem from these points that there is no choice, basically, but build this graph on your own. And we face these challenges at Pasteur. And in order to solve some of them, or at least to ameliorate them, we developed a framework that I'm really happy to say that it has been open sourced. And it's called Tesseract Core. And it's based on the concept of Tesseracts, which are basically, the name can be a bit confusing, but it's like a Docker container, you can think of it as a box, right? And then like you add the extra dimension of automatic differentiation, and it's kind of like a 4D box to a Tesseract, if you're curious about the name. And the nice thing that they do, in my opinion, is that they automatically wrap just Python modules that typically in our work are CML components, but it can be other stuff as well. And basically they automatically wrap them So on their own, this Tesseract framework basically builds containers for you. And then these things expose the REST API, but you can also call them via the CLI. And we have a Python client, of course. Otherwise, I wouldn't be here at PyCon. And yeah, and designed to really be easy to create, to use, and to share. And as you can imagine, a project like this stands on the shoulders of a lot of open source packages and contributions, probably the ones that are directly most related to this package and that do a lot of the heavy lifting are Pydantic and FastAPI. And again, I'm showing this framework. I'm really proud of it. And feel free to start it. But I don't want to give a marketing pitch in this talk. I just want to show you the lessons we learned while building this and a bit of the design decisions that we tried to use to solve them. And to be very concrete, I'm going to show a minimal example of how you define a test seduct. So here, we're wrapping a very simple function It just accepts a 3D vector, sums all of its components, and returns a number, so extremely easy. And what do you need to do to wrap it in a tesseract? First of all, you need to define an input schema and an output schema. And you do these tools that are familiar to Python developers, like PyDantic. There are caveats here, because if you see how this x and the total are type annotated, you see that the hints here are a bit weird. Like, you know, we had to introduce our custom annotations here, and that's because Pydantic out of the box doesn't really support numeric types or types that you would use in computational science that well. So, okay, we had to do this, but then this gave us also a lot of flexibility in terms of how we serialize this stuff and do a lot of operations. And you see, for instance, here, I mean, it's sort of similar here also to NumPy type. So array 3 float32 is just a 3D array of float32 numbers. So nothing bad there. And also differentiable here, just you're marking some variables as the leaves that are differentiable or that you can differentiate with respect to, so that also not only these sort of primal schemas that you define here, but also all of the schemas for Jacobians and JVP, all of the differentiable stuff is automatically created for you here. And then once you have the input schema and the output schema, then you can define the meet of your function. So the apply function, which just connects the two and brings you from the inputs to the outputs. And something you can notice here is that here in apply, we have x.sum, and this is an indication that within this scope, x is a NumPy array. and literally like just these three like these two classes and the apply function will be enough for for defining a tesseract but then of course there are also other endpoints that you can define like Jacobian or JVP VJP and a lot of the other stuff that you will need for automatic differentiation and since you have all these rich annotations here like also in the pedantic scheme as you can put descriptions and you know doc strings when you want we did also a lot of effort into also piping this information when it's relevant and make sure that Redoc builds valid docs for all of these test tracks. So this is related to this sort of sociological problem that I mentioned before that maybe you really write a package that accepts maybe some vectors but these vectors have to be normalized according to some rule and I mean of course somebody is gonna have to write these things or maybe at some point we We'll just put LLMs there. I don't know, but at least this should make it slightly more frictionless to just write these descriptions and maybe package your code into a better way, even if you're not an expert on packaging and containerizing stuff. And again, these are the docs of the REST API, but also the command line interface, which is built, follows the same sort of design philosophy. Now, then, OK, there can be objections to this, then, because, OK, what I've shown so far is only how one could wrap these components into this Tesseract. But at the end of the day, I didn't say anything about assembling them into this graph. And it will seem that still these things that I described are containers or whatever. And then there are HTTP requests to move data around them. So it will seem that then you will still need to build this graph on your own. And this is where we started actually integrating Tesseracts with automatic differentiation frameworks. So here you can see an end-to-end example, really, of using these Tesseracts with JAX. You need to import this extension. And what this does is basically it registers the act of calling a tesseract as a JAX primitive. So when you do that, basically, JAX now knows how to trace through these things, even if there are software engineering things, like HTTP requests happening under the hood. And I would say, also, writing this package was extremely interesting, because we had to really delve deep into JAX internal. So it's also a relatively small package. And in my opinion, it's like an interesting read as well if you like to read source code. And it works out of the box. I mean, it works with Tesseract in this way that you can define, again, any arbitrary function. You call a Tesseract, you get a result, and maybe this result you pipe through another Tesseract or you do some local operation. And the important thing here is that not only you can calculate the output of this operation, but you can also calculate the gradient just via jax.grad, or you could use JAX.JVP and so on and so forth. Also, another minor detail, but I think it's cute, and it's that we have these context managers defined with Tesseract Core, where basically, again, since the communities we're writing stuff for, we don't really assume that they're like Docker Pro users or so, then you can just handle the lifecycle of these containers directly within Python. So in this context manager, when you enter it, this container is launched, and then, like, you know, once it's healthy like you start doing operations and then when this scope is left then your container is destroyed and also there are like a bit of technical caveats here, of course every AD engine tends to want to do stuff in a very specific way, for instance JAX basically does first shape inference before running the computation so it really requires this from its primitive to be able to say, okay, if your function accepts vectors of a given shape, your output should have shape, whatever. And so, of course, we cannot circumvent that, and then when you're using Tesseract with Tesseract JAX, you need to implement this endpoint as well. And, yeah, this is my final slide. So I wanted to show you just like this idealized scenario that I've shown before, how it looks like in a more realistic case. And, I mean, it is a bit intimidating, but, like, also I just want to show a few highlights of this. And, again, the problem here is that we wanted to optimize some shape immersed in some airflow, and we wanted to use OpenFOAM, which is an open-source computational fluid dynamics solver, which is not differentiable, but the physics underneath is. So what we do here, then, ideally, and I'm showing this a bit in this lower left panel, is that we would like to, you know, for some stuff, use the exact solver, but then when we have to maybe back-propagate things and so on, we would really like to switch the solver with a surrogate model that is like a machine learning model that learned how the solver works. And so here you see that this is all orchestrated with Tesseract, so one starts with these design parameters, some geometry is generated, and I mean, I really don't want to go into the details, But then, you know, from these design parameters, some machine learning-ready data is generated to then train the surrogate that I mentioned before. And then once this is done, then for the optimization loop, sort of one follows these blue arrows. So, you know, you go from design parameters to the objective. And then to backpropagate here, some of these components are easy to calculate derivatives of because they are differentiable. differentiable like you know machine learning models are you know inherently sort of differentiable and also this interpolator that we have here is as well and for instance though we had the issue that like our procedures generating the geometry were not differentiable but again since a lot of these rules you can insert at sort of like a component wise level then that one component we decided to go for a finite difference scheme which actually turns out is not necessarily inferior to autodiff in a lot of cases where you have maybe noisy data and so on and so i mean you can see some iteration loops here i mean it's not a particularly interesting example still but you know the orchestration and so on i think it's the interesting part here and yes this brings me to my conclusion so again i hope i've shown you that at least a bit that like differentiable programming has really a lot of applications beyond beyond neural networks Combining differentiable components into pipelines has great potential, but it is challenging. There are a lot of challenges related to technical stuff, but also I think there is the inherent challenge, let's say, of taking some scientific code coming from some research lab, packaging it, and sort of using it. There is a lot that I skimmed over, like which data types you will need to exchange between these components and so on and so forth. And, yeah, we built this open-source framework, Tesseract, so that at least people have an easier time to cope with this problem. And, yeah, Tesseract JAX also now allows you to orchestrate these things with JAX without really the need of doing much more. And, yeah, I think that's just my talk. Yeah, thanks a lot. Thank you. so thank you for this wonderful talk um are there any questions on Slido I forgot to mention that you can for sure ask question over Slido if not then raise your hands and I will come to you are any questions okay yeah then ah yeah so could you explain again maybe for an example if you have two components, let's say one written in JAX and the other one written in PyTorch that both have derivatives and you want the derivative of the whole thing, how do you glue them together in the end? Yes. So, very nice question. Thanks. So, the first thing is that, you know, let's say we have, like, just JAX and PyTorch here, right? So, in this case, these two components will be two separate modules that you will have to fill in here. We already do have templates for JAX and for PyTorch as well. So also a lot of the, let's say, annoying stuff of even writing down the Jacobian with the correct schema that you have here is sort of templated, and you don't have to do it. You literally just have to fill in the apply and the schemas. And this, then, is how it works at a component level. Now, to orchestrate these two different components, we have a bit of a bias with JAXA. That's the first integration that we built. So, yeah, you could then write a script like this one where, like, you know, you have one line where you apply Tesseract and this is, like, you know, the operation you wrote in JAX, let's say, and then, you know, you pipe the result here to, again, via this apply Tesseract to the PyTorch one. We would like to do work also on doing something similar for PyTorch. So, like, as the orchestrating language sort of and registering this as a custom primitive in PyTorch. But I have to say, right now, we didn't really start doing that, because we're pretty happy with JAX, but, yeah. Okay, anything else that you want to ask? Perfect. Then, I think we are through. Thank you for this presentation. Give another applause.