Jacobian Descent for Multi-Objective Optimization

arXiv:2406.16232 · cs.LG, cs.AI, math.OC · Submitted 2024-06-23 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Jacobian Descent for Multi-Objective Optimization".

Jane: Here's the gist of the paper. The gist The Jacobian descent (JD) algorithm is introduced as a direct generalization of gradient descent for multi-objective optimization, proposing AUPGrad to resolve conflicts while preserving gradient influence, and it enables instance-wise risk minimization (IWRM) in neural networks Jacobian Descent Formulation Jacobian descent is defined as an iterative update using the Jacobian matrix of a vector-valued objective function, where each row is the gradient of an individual objective <ref:2406.16232#pg2>.

Tom: Given a continuously differentiable function f: R n → R m, the first-order Taylor approximation suggests selecting an update y = −ηAJf(x) to reduce f(x + y), where A is an aggregator <ref:2406.16232#pg3>.

Jane: Algorithm 1 illustrates this process iteratively <ref:2406.16232#pg3>.

Tom: Stochastic Jacobian descent (SJD) can be implemented by aggregating estimates of the Jacobian, such as stochastic sub-Jacobian descent (SSJD) or stochastically estimated Jacobian descent (SEJD) <ref:2406.16232#pg3>.

Title and authors: Tom: If we boil down what this paper is actually doing, they are formalizing an iterative process called Jacobian descent and then proposing AUPGrad as the specific rule for combining gradients when you have multiple objectives <ref:2406.1632#pg2>.

Jane: Can you explain that in plain language for someone who isn't deep into optimization theory? What is this iterative update actually doing during training?

Lu: They define the update rule as x sub t plus one equals x sub t minus eta AUPGrad J f of x sub t, where the Jacobian matrix J has a row for each objective, and AUPGrad is that special aggregator <ref:2406.1632#pg3>.

Meng: So it’s like instead of just taking one gradient direction, you’re calculating a weighted combination of all those gradients simultaneously to take the next step <ref:2406.1632#pg3>.

Tom: Pretty much so. And they also introduce two variants of this algorithm for when things get tricky, which is why they have stochastic Jacobian descent, or SJD <ref:2406.1632#pg3>.

Jane: Stochastic? What does that mean in the context of training a model? Are we talking about using some kind of random sampling during the update step?

Lu: Yes, SJD relies on estimates of the aggregation of the Jacobian, and they have two approaches: stochastically estimated Jacobian descent or SEJD <ref:2406.1632#pg3>. Or another way, when you have a huge number of objectives, they propose to aggregate a matrix whose rows are a random subset of the rows of the true Jacobian <ref:2406.1632#pg3>.

Meng: That sounds computationally intense to run on real data, especially if the true Jacobian is huge. What’s the cost there?

Tom: The main computational headache in standard Jacobian descent is that you have to store a full Jacobian matrix <ref:2406.1632#pg11>, which can get really big fast, but they suggest using a Gramian-based method instead, which stores the Gramian Gf of x equals Jf(x) dot Jf(x) transpose <ref:2406.1632#pg12>.

Jane: So it’s a trade-off: you save memory by not storing the full matrix, but you still have to deal with these complex matrix operations during the update step <ref:2406.1632#pg13>.

The paper's summary: Tom: Let’s get into what they actually achieved beyond just defining the algorithm. They focus on showing that AUPGrad is superior to other methods we see in the literature <ref:2406.1632#pg13>.

Jane: They analyze several existing aggregators, and they show that standard ones like AMean are not non-conflicting because they conflict with other gradients, so AUPGrad satisfies all three properties <ref:2406.1632#pg13>. That’s a big win for managing conflicting objectives.

Lu: And the analysis shows that other methods, like AMGDA and ADualProj, are non-conflicting but they fail the linear under scaling property <ref:2406.1632#pg13>. AUPGrad manages to hit all three targets simultaneously <ref:2406.1632#pg5>.

Meng: So what does that mean when we think about deploying this in a production environment? Is it just theoretical nice, or does it actually help us train models better?

Tom: The empirical results on simple image classification tasks show that AUPGrad provides a significant improvement over the baseline AMean, which is basically empirical risk minimization with SGD <ref:2406.1632#pg9>.

Jane: What’s the practical implication of those results? It seems to suggest that AUPGrad prevents gradients from hard examples from being completely overwhelmed by gradients from easier examples early on in the training process <ref:2406.1632#pg9>.

Lu: They also found that increasing the batch size refines the dual cone, which means you get even more benefit from that non-conflicting property as you increase your batch size <ref:2406.1632#pg9>.

Meng: That’s interesting. So, if we can improve how the aggregator works with larger batches, does that mean it’s viable for practical use?

Tom: It shows that it’s not just a theoretical exercise. They also mentioned applications like instance-wise risk minimization or IWRM, which is a new way to look at loss vectors <ref:2406.1632#pg7>.

Jane: Instance-wise risk minimization—that sounds like thinking about every single training example as its own separate objective, where the loss of each example is treated as a distinct part of the overall objective function f(x) =

f1(x) f2(x) · · · fm(x): T <ref:2406.1632#pg7>.

Lu: Yes, they suggest that you can train these instance-wise objectives using stochastic sub-Jacobian descent, SSJD <ref:2406.1632#pg7>.

The paper's improvements: Tom: So to wrap up on "Jacobian Descent for Multi-Objective Optimization," the main thing is that they gave us a robust way to handle conflicting objectives using AUPGrad, which resolves conflicts while keeping the gradient influence proportional to its norm <ref:2406.1632#pg5>.

Jane: They also showed that this method converges to the Pareto front in the smooth convex case, which is a key theoretical result for understanding how optimization behaves under these conditions <ref:2406.1632#pg6>.

Lu: The paper suggests that they've introduced IWRM as a novel learning paradigm where you consider the loss of each training example as a separate objective, which is something that could open up new ways to structure model training <ref:2406.1632#pg7>.

Meng: From an engineering standpoint, the Gramian-based approach for storing the Jacobian helps reduce memory overhead by only storing the Gramian Gf of x equals Jf(x) dot Jf(x) transpose <ref:2406.1632#pg12>.

Lalam: I’m confident that this work provides a solid framework for multi-objective optimization by providing AUPGrad, and the convergence proof to the Pareto front is strong <ref:2406.1632#pg6>.

Tom: Exactly. It gives us a concrete algorithm, and it also shows how practical improvements on things like batch size can actually make a difference in real training performance <ref:2406.1632#pg9>.

Jane: So, this paper introduces Jacobian descent with AUPGrad as a framework for training AI systems where multiple objectives are competing, and it also hints at instance-wise risk minimization as a new way to think about those loss vectors <ref:2406.1632#pg7>.

Conclusion: Tom: So we've looked at "Jacobian Descent for Multi-Objective Optimization," Pierre Quinton and Valérian Rey’s work, and the main thing is that they gave us AUPGrad, which resolves conflicts between multiple objectives while keeping the gradient influence proportional to its norm.

Jane: It also showed that this method converges to the Pareto front in the smooth convex case, which is a key theoretical result for understanding how optimization behaves when you have conflicting goals.

Lu: They introduced IWRM as a way to think about loss vectors where you consider every single training example as its own separate objective, which could open up new ways to structure model training.

Meng: The Gramian-based approach for storing the Jacobian helps reduce memory overhead by only storing the Gramian Gf of x equals Jf(x) dot Jf(x) transpose.

Lalam: I’m confident that this work provides a solid framework for multi-objective optimization by providing AUPGrad, and the convergence proof to the Pareto front is strong.

Tom: It gives us a concrete algorithm, and it also shows how practical improvements on things like batch size can actually make a difference in real training performance.

Jane: So this paper introduces Jacobian descent with AUPGrad as a framework for training AI systems where multiple objectives are competing, and it also hints at instance-wise risk minimization as a new way to think about those loss vectors.

Lu: It’s interesting that they tackled the three properties of the aggregator—non-conflicting, linear under scaling, and weighted—that other methods just couldn't handle together.

Meng: From an engineering standpoint, I’m still thinking about how much overhead is involved in setting up that Gramian matrix calculation before we even get to training.

Lalam: The AUPGrad approach seems like it solves that setup problem by making the update step structured, which is a big plus for scaling up these kinds of multi-objective tasks.

Tom: Yeah, and those empirical results on image classification show AUPGrad gives a significant improvement over the standard methods we use right now.

Jane: It seems to suggest that this method helps prevent gradients from hard examples being completely overwhelmed by gradients from easier examples early on in the training process.

Lu: They also found that increasing the batch size refines that dual cone, meaning you get even more benefit from that non-conflicting property as your batch size gets bigger.

Meng: If we can improve how the aggregator works with larger batches, it definitely makes this viable for production deployment instead of just being a theoretical exercise.

Lalam: This Jacobian Descent for Multi-Objective Optimization paper gives us a solid framework, and I think that’s exactly what we needed to move forward with these complex AI setups.

Tom: It gives us a concrete algorithm, and it also shows how practical improvements on things like batch size can actually make a difference in real training performance.

Jane: So this paper introduces Jacobian descent with AUPGrad as a framework for training AI systems where multiple objectives are competing, and it also hints at instance-wise risk minimization as a new way to think about those loss vectors.

Lu: It’s interesting that they tackled the three properties of the aggregator—non-conflicting, linear under scaling, and weighted—that other methods just couldn't handle together.

Meng: From an engineering standpoint, I’m still thinking about how much overhead is involved in setting up that Gramian matrix calculation before we even get to training.

Lalam: The AUPGrad approach seems like it solves that setup problem by making the update step structured, which is a big plus for scaling up these kinds of multi-objective tasks.

Tom: Yeah, and those empirical results on image classification show AUPGrad gives a significant improvement over the standard methods we use right now.

Jane: It seems to suggest that this method helps prevent gradients from hard examples being completely overwhelmed by gradients from easier examples early on in the training process.

Lu: They also found that increasing the batch size refines that dual cone, meaning you get even more benefit from that non-conflicting property as your batch size gets bigger.

Meng: If we can improve how the aggregator works with larger batches, it definitely makes this viable for production deployment instead of just being a theoretical exercise.

Lalam: This Jacobian Descent for Multi-Objective Optimization paper gives us a solid framework, and I think that’s exactly what we needed to move forward with these complex AI setups.

Tom: Exactly. It gives us a concrete algorithm, and it also shows how practical improvements on things like batch size can actually make a difference in real training performance.

Jane: So this paper introduces Jacobian descent with AUPGrad as a framework for training AI systems where multiple objectives are competing, and it also hints at instance-wise risk minimization as a new way to think about those loss vectors.

Pierre Quinton, Valérian Rey

EPFL

cs.LG, cs.AI, math.OC

Submitted: 2024-06-23

Updated: 2026-10-19

Comments: 53 pages, 14 figures, Accepted at NeurIPS 2026 (poster)

Code: https://github.com/TorchJD/torchjd

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 82/100

The gist: The gist The Jacobian descent (JD) algorithm is introduced as a direct generalization of gradient descent for multi-objective optimization, proposing AUPGrad to resolve conflicts while preserving

Key concepts

Jacobian Descent (JD)
JD is an iterative optimization algorithm that uses the Jacobian matrix of a vector-valued objective function to determine the next step. It's a generalization of standard gradient descent, allowing optimization across multiple objectives simultaneously by using the relationship between inputs and outputs.
AUPGrad
AUPGrad is a specific aggregator used in JD that manages conflicting objectives effectively. It achieves this by projecting gradients onto the dual cone of the Jacobian rows, ensuring it remains non-conflicting, linear under scaling, and weighted. This helps preserve gradient influence while balancing multiple goals.
Instance-wise Risk Minimization (IWRM)
IWRM is a paradigm where each training example's loss is treated as a separate objective function. The total loss is represented as a vector of individual losses. JD with AUPGrad can be used to train models under this paradigm, optimizing the entire instance-wise loss vector.
Gramian-based Jacobian Descent
This implementation optimizes JD by only storing the Gramian matrix Gf(x) instead of the full Jacobian. The Gramian is a smaller matrix derived from $J_f(x) imes J_f(x)^ op$, which significantly reduces memory overhead, making it computationally efficient for high-dimensional problems.

Terminology

Summary

The gist The Jacobian descent (JD) algorithm is introduced as a direct generalization of gradient descent for multi-objective optimization, proposing AUPGrad to resolve conflicts while preserving gradient influence, and it enables instance-wise risk minimization (IWRM) in neural networks

Jacobian Descent Formulation

Jacobian descent is defined as an iterative update using the Jacobian matrix of a vector-valued objective function, where each row is the gradient of an individual objective <ref:2406.16232#pg2>. Given a continuously differentiable function f: R n → R m, the first-order Taylor approximation suggests selecting an update y = −ηAJf(x) to reduce f(x + y), where A is an aggregator <ref:2406.16232#pg3>. Algorithm 1 illustrates this process iteratively <ref:2406.16232#pg3>. Stochastic Jacobian descent (SJD) can be implemented by aggregating estimates of the Jacobian, such as stochastic sub-Jacobian descent (SSJD) or stochastically estimated Jacobian descent (SEJD) <ref:2406.16232#pg3>.

Desirable Aggregator Properties and AUPGrad

To manage conflicting objectives, desirable properties for an aggregator A were defined:

  1. Non-conflicting: If for all J ∈ R m×n, 0 ĺ J · A(J), then A is said to be non-conflicting <ref:2406.16232#pg3>.

  2. Linear under scaling: If the mapping from any 0 ă c ∈ R m to A diag(c) · J is linear in c, then A is said to be linear under scaling <ref:2406.16232#pg3>.

  3. Weighted: If for all J ∈ R m×n, there exists w ∈ R m satisfying A(J) = J⊤ · w, then A is said to be weighted <ref:2406.16232#pg3>.

AUPGrad is proposed as a specific aggregator designed to satisfy these three properties by projecting gradients onto the dual cone of the rows of the Jacobian and averaging the results <ref:2406.16232#pg5>. This approach ensures that AUPGrad is always non-conflicting, linear under scaling, and weighted <ref:2406.16232#pg5>.

Convergence Guarantees

Theorem 1 provides convergence guarantees for JD with AUPGrad when minimizing a β-smooth and ĺ-convex function f: R n → R m <ref:2406.16232#pg6>. The update rule is xt+1 = xt − ηAUPGradJf(xt), where η = 1/β√m <ref:2406.16232#pg6>. If the weights wt defining AUPGradJf(xt) are bounded, then f(xt) converges to the Pareto front f(x∗) for some x∗ ∈ X∗ <ref:2406.16232#pg6>. This result is significant as it is presented as the first non-conflicting aggregator that provably converges to the Pareto front in the smooth convex case <ref:2406.16232#pg6>.

Applications and Empirical Results

The paper highlights two main applications:

  1. Instance-wise risk minimization (IWRM): This paradigm considers the loss of each training example as a separate objective, defined by f(x) = [f1(x) f2(x) · · · fm(x)]⊤ <ref:2406.16232#pg7>. IWRM can be trained using stochastic sub-Jacobian descent (SSJD), which uses the Jacobian of a random batch of training example losses <ref:2406.16232#pg7>.

  2. Multi-task learning: JD with a non-conflicting aggregator is suggested as a natural setting for multi-task learning where task conflicts are inherent <ref:2406.16232#pg7>.

Empirical experiments on simple image classification tasks (SVHN, CIFAR-10, EuroSAT, MNIST, Fashion-MNIST and Kuzushiji-MNIST) show that AUPGrad provides a significant improvement over the baseline AMean (which is equivalent to empirical risk minimization with SGD) <ref:2406.16232#pg9>. The results demonstrate that AUPGrad prevents gradients of hard examples from being dominated by those of easier examples early into the training process <ref:2406.16232#pg9>. Furthermore, increasing the batch size refines the dual cone, pushing even further the benefits of the non-conflicting property <ref:2406.16232#pg9>.

Computational Efficiency

The main computational overhead in JD comes from storing a full Jacobian matrix <ref:2406.16232#pg11>. The Gramian-based method offers an alternative by only storing the Gramian Gf(x) = Jf(x)· Jf(x)⊤, which is typically much smaller and similar in memory requirement to SGD <ref:2406.16232#pg12>. The time complexity for extracting the weights of AUPGrad is O(m5), though solving the m quadratic programs in parallel can reduce this to O(m4) <ref:2406.16232#pg13>.

Existing Aggregators Analysis

The paper analyzes several existing aggregators, summarizing their properties in Table 1 <ref:2406.16232#pg13>. It demonstrates that AMean is not non-conflicting because it conflicts with other gradients, while AUPGrad satisfies all three properties <ref:2406.16232#pg13>. Other methods like AMGDA and ADualProj are shown to be non-conflicting but fail the linear under scaling property <ref:2406.16232#pg13>. The analysis also shows that AGradDrop is not non-conflicting or weighted, and ANash-MTL is not linear under scaling <ref:2406.16232#pg13>.

Conclusion

The paper introduces Jacobian descent (JD) with AUPGrad to provide a robust framework for multi-objective optimization by resolving conflicts between conflicting gradients while maintaining desirable properties. The work also proposes IWRM as a novel learning paradigm and outlines an efficient implementation using the Gramian of the Jacobian <ref:2406.16232#pg11>.

The paper is 2406.16232.

Improvements for AI systems

  1. Development of Instance-wise Risk Minimization (IWRM) for training models by considering each training example as a separate objective function, which can be represented by a vector of loss values f(x) = [f1(x) f2(x) · · · fm(x)] T.

  2. Implementation of Jacobian descent with AUPGrad to resolve conflicts between conflicting objectives while ensuring that they preserve an influence proportional to their norm, leading to convergence guarantees to the Pareto front in the smooth convex case.

  3. Application of Gramian-based Jacobian descent, which reduces computational overhead by only storing the Gramian instead of the full Jacobian, enabling efficient optimization when the number of objectives is dominated by model parameters.

  4. Integration of AUPGrad into momentum-based optimizers like Adam to combine gradients and momentum in a way that preserves non-conflicting properties, potentially leading to faster convergence than standard SGD on multi-objective tasks.

  5. Utilization of stochastic sub-Jacobian descent (SSJD) with AUPGrad across multiple training examples, which enables the training of neural networks with a large number of objectives by considering the minimization of instance-wise loss vectors.

Abstract

Many optimization problems require balancing multiple conflicting objectives. To generalize gradient descent, which is limited to single-objective optimization, we introduce Jacobian descent (JD) and its stochastic variants. This algorithm iteratively updates parameters using the Jacobian matrix of a vector-valued objective function, in which each row is the gradient of an individual objective. While several methods to combine gradients already exist in the literature, they generally struggle when the objectives conflict. In contrast, we propose projecting gradients to fully resolve conflict. We prove convergence to the Pareto front in the smooth convex case with this approach, a stronger guarantee than the usual convergence to a weakly Pareto stationary point. Our method also enables instance-wise risk minimization (IWRM), a novel learning paradigm in which the loss of each training example is considered a separate objective. Early experimentation on image classification shows promising results for IWRM. Lastly, we provide an efficient implementation of JD using the Gramian of the Jacobian matrix to drastically reduce memory requirements.

Sources

Related papers