Lightning Talk: TorchJD: Jacobian Descent in PyTorch - Pierre Quinton & Valérian Rey

This talk introduces Torch JD, a library designed for multi-objective optimization using PyTorch, focusing on the Jacobian descent algorithm. The speakers discuss the importance of making decisions that balance conflicting objectives, such as price and quality in everyday scenarios or economic and ecological impacts in broader societal contexts. They explain how traditional gradient descent can be adapted for multiple objectives through techniques like unconflicting projections of gradients. The presentation also includes a practical demonstration of Torch JD's features, such as computing Jacobians and aggregating them into gradients within PyTorch's framework. The speakers highlight various applications of the library in deep learning contexts, including multitask and federated learning, while encouraging attendees to experiment with the tool and provide feedback.

Welcome to our presentation. I'm Pierre, and this is Valerian, and today we'll be presenting a bit our library uh Torch JD that enables doing multi-objective optimization in PyTorch. And specifically, we can do the Jacobian descent algorithm with it. Um All right. Let me start with a bit of motivation first. So, it's 11:45-ish, and I claim that you've already done multi-objective optimization today unless you just woke up.

But, uh for instance, if you went to the grocery store and bought some food, then probably you optimized for price, maybe for uh flavor, or nutritive content, many things, and this is multi-objective optimization. At a different scale, if we take a decision as a society and we want it to be good maybe for economics, or ecological impact, social impact, and many other factors, then this is multi-objective

optimization, obviously. Now, we might see those objectives as being conflicting, but surely there must be ways of taking decisions that are good for all of them. Um actually, if you think about it, we have a tendency when we formulate a problem to frame the objectives as being conflicting, and the reason for that, I think, is because we want to pinpoint exactly where the tension is and where

we need to be smart in the decision process. So, in that sense, conflict is good because it's informative, and if you use this information right, then you can make better decisions. All right. Uh like taking a decision, resolving the conflict, I think it's pretty hard, but this is a very important problem, probably one of the most important uh of the century, and this is why we started

working on this in the first place. Of course, this is more applied to optimization, but I think it's an important topic. Now, if you have several objectives, what you want to be careful about is you want to know if the objectives are and then how to resolve the conflict between objectives and then when is it possible to do so. All right. Maybe let's step back into optimization

and let's start with gradient descent. So, if you have one objective then what you typically do is you compute the gradient and then you would update the parameters of the model into the opposite direction of the gradient and this will minimize at least locally the the objective. Okay, if we have several objectives then we have several gradients. Let's say this one. Now, if we take the same

decision as before so we update in the opposite direction of the first gradient because they have a negative inner product what you would end up doing is optimizing the first one so that's good but de-optimizing the second one. So, this is a bad decision. Another thing you could try which is quite natural is maybe you do the average of the objectives and you compute the gradient and

you would get this. If you update the parameters of the then this is pretty good for the second objective because it's aligned because it's like a positive inner product but it's pretty bad for the first one this time because they have a negative inner product. So, you need to be a bit careful and if you think about it you reason about the set of vectors that have

a positive inner product with all of your gradients basically and this is this green cone. And if you take any decision outside of this cone then you would actually locally de-optimize at least one objective and this is probably not something you want to do. Now, there are many decisions you could take and the one that we recommend because it works very well in practice but you could

find your own is to project the gradients onto the and then average them and this will fall inside the cone. It's quite natural. Now, this we call the unconflicting projection of gradients and before I will let Valerian present a bit the library, let me do a tiny bit of terminology. So, if you have several objectives, you have several gradients. If you stack those gradients into a matrix,

we call this the Jacobian matrix and this is why we call the algorithm Jacobian descent. The Jacobian is a generalization of the gradient. So, Jacobian descent is a generalization of gradient descent. But, the Jacobian itself is not enough to update the parameters. You need to take a decision in what direction you want to move your parameters. And so, you need to aggregate that Jacobian into a direction.

Thank you, Pierre. Um so, now I'm going to show you how you can do Jacobian descent in PyTorch using our library torch-jd. So, the good thing about torch-jd is that its interface is quite similar to that of torch. Um so, in torch you have autograd.grad to compute some gradients and you have autograd.backward to compute some gradients and accumulate them in the grad field of the inputs. Um

in torch-jd, you have autoja.jac to compute directly some Jacobians and also gradients at the same time. And you have autoja.backward to compute some Jacobians and accumulate them in the jac field of the inputs. And lastly, we have a function autoja.jac_to_grad to do the aggregation, so to take this Jacobian from the jac fields, accumulate it aggregate it so into a vector and accumulate this vector in the of

the tensors. Now, we're going to see that with a very simple example. So we start with some imports from torch JD. And in particular, you can see that we Oh, okay. Yeah. Yeah. Yeah, it's working. Yeah. So in particular, um you have to you see that we import up something called up grad. So it stands for unconflicting projection of gradients, and it's exactly what Pierre showed just

before. So the algorithm where you project the rows of the Jacobian onto this green cone, and then you take the average. The example is going to be very simple. It's just supervised learning, but instead of uh averaging the losses over the batch, what if we don't? And we consider each loss from the batch separately. So here we have a batch size of 16, and we have a

loss function where we set the reduction to none, so that we don't average over the batch. And then we iterate over some data loader, we compute some outputs, and we compute a batch of losses. So we have 16 losses, and we would like to compute a gradient for all of them. So we would like to compute the Jacobian of the losses with respect to the model parameters,

and we can do that with auto jack.backward. And here, as you can notice, we don't even have to specify that we differentiate with respect to the model parameters, because the model parameters are the leaves of the autograd graph, so it's as in autograd, we don't even have to specify that, it's just the default. And now you have model parameters that have a jack field. But what we

need to train a model, as you know, is to have a grad field. So we can simply call auto jack.jack_to_grad to do this aggregation and turn this Jacobian into something that will be used as a gradients and that will be stored inside of the grad field. And as you can see here, we specify the unconflicting projection of method, but you could actually use any method here. Uh

we support more than 15 from the literature, and you could actually implement your own if you want to try something different. And at this point, you have a grad field in your model parameters, so you can just use PyTorch as usual, call your optimizer, and it's very It's really simple. You plug a few lines from Torch-JA into your PyTorch training loop, and you end up we've seen

how you can use Torch-JA to do multi-objective optimization in Now, why should you care? Thing is that in deep learning, there are many ways your problems can be framed as multi-objective optimization. For example, if you do multitask learning, you train a single model uh with several different tasks. Or if you do federated learning, you have different entities that collaborate to train a common model based on their

own data. Uh even if you do AI alignment, for example, you try to train a model to be as good as possible at its job, but while retaining some human principles. And even if your work doesn't fall in any of those categories, as we've seen, we can actually transform a simple Sorry, a simple um supervised learning by considering the loss from each element of the batch separately.

Or even you could think of the loss of each element of the training set And actually, we've been experimenting on that quite a lot, so if you're interesting if you're interested about results for that, uh you can just refer to our paper. also, if you find all of this interesting, the best thing you could do to help us would be to just try it out and give

us some feedback. And also we'd be really happy to discuss with you offline, to answer any questions, so really feel free to come to us. Uh thanks everyone for listening.