TorchJD Gets a Major Update: Multi-Loss Training No Longer Relies on Trial-and-Error Weight Tuning

Anyone who has done multi-task learning knows how nasty multi-loss training can be: you take several losses, weight them, and sum them together—but how do you choose the weights? You start with arbitrary numbers like 0.3, 0.5, and 0.2, then halfway through training realize one task never learned anything, so you go back and tune again. This TorchJD update is a major step toward turning that process from black magic into engineering.

This open-source library, focused on Jacobian descent in PyTorch, has just merged a batch of new aggregator implementations, covering nearly all major multi-objective optimization methods proposed in recent years. For developers constantly struggling with multi-task objectives, constraints, and regularization terms, it’s worth a serious look.

First: What’s the Difference Between Scalarization and Jacobian Descent?

If a model has only one loss, training is simple—compute gradients, backpropagate, update. But in real-world scenarios, you usually have a pile of losses: main-task losses, auxiliary-task losses, constraint terms, regularizers, and so on.

There are two mainstream ways to handle them:

• Scalarization: Combine N losses into a single scalar. The simplest approach is averaging; more sophisticated versions use dynamic weighting methods like GradNorm, DWA, or Uncertainty Weighting. The advantage is efficiency—one backward pass, memory-friendly.
• Jacobian Descent: Compute gradients for each loss separately, forming a Jacobian matrix (one gradient per row), then aggregate those gradients using some strategy to produce an update vector. The goal is not to reduce the “average loss,” but to reduce “every loss.”

What’s the practical difference? Take multi-objective optimization in recommendation systems as an example: click-through rate and conversion rate. If the gradient directions are similar, scalarization works fine—just average them. But if the gradients strongly conflict—one pointing east, the other west—the weighted average may optimize neither task properly, or even make them interfere with each other.

The core idea of Jacobian descent is: if gradients conflict, the aggregation step should explicitly account for that conflict, finding a direction that decreases all losses (or at least avoids worsening any single one). Classic methods like MGDA, PCGrad, and CAGrad all belong to this family.

The tradeoff is memory and computation. A Jacobian means you need a separate backward pass for each loss (or some approximation trick), with memory overhead on the order of O(N × parameter count).

What’s Included in This Update

According to the project team, after v0.2 they focused on filling out the major aggregators from the literature. The library now directly supports a fairly broad set of methods:

• UPGrad: TorchJD’s flagship aggregator. The idea is to project each gradient into directions that “do not harm other tasks” before aggregation, with theoretical Pareto optimality guarantees.
• MGDA / MGDA-UB: Classic multi-objective optimization methods that search for the minimum-norm direction among linear combinations of gradients.
• PCGrad: From Google’s well-known paper—when gradients conflict, project one onto the orthogonal complement of the other.
• CAGrad: Searches near the average gradient for an update direction that guarantees no task gets worse.
• GradDrop, Nash-MTL, Aligned-MTL, DualProj, and a series of later methods.

Combined with the traditional Scalarization family (Mean, Sum, weighted methods, Random Weighting, GradNorm, etc.), almost anything you see in papers and want to reproduce can now be swapped in with a one-line API change.

That’s genuinely valuable in practical R&D. Previously, comparing PCGrad and CAGrad meant digging through separate author repositories with inconsistent interfaces, mismatched dependencies, and potentially incompatible PyTorch versions. Now everything is unified under one library, one interface, and one consistent abstraction, dramatically lowering the barrier for A/B testing.

What Using It Feels Like

TorchJD’s API follows standard PyTorch conventions. You replace loss.backward() with the library’s backward, passing in a loss vector and an aggregator. The pseudocode looks roughly like this:

from torchjd import backward
from torchjd.aggregation import UPGrad

aggregator = UPGrad()

for batch in loader:
    losses = model.compute_losses(batch)  # returns a loss tensor of shape [num_losses]
    optimizer.zero_grad()
    backward(losses, aggregator)
    optimizer.step()

Compared to manually implementing PCGrad—registering hooks, computing per-task gradients, projecting them, writing them back into parameters—it’s much cleaner. And if you want to switch algorithms, you just change one line to aggregator = CAGrad().

For the classic “shared backbone + multi-task head” MTL architecture, the library also provides mtl_backward, which applies Jacobian descent only to shared parameters while task-specific parameters use standard backpropagation, saving a significant amount of memory and compute. That detail shows the team has real MTL engineering experience—the academic version of PCGrad computes the Jacobian over the entire model, which can easily blow up GPU memory in practice.

Who Should Use It—and Who Shouldn’t

A few practical judgments.

Scenarios where it’s worth trying:

• Multi-task learning: Especially when tasks conflict to some degree, such as CTR + CVR + dwell time in recommendation systems, or detection + segmentation + depth estimation in autonomous driving.
• Constrained optimization: Main loss plus multiple hard constraints (fairness constraints, safety constraints, etc.). Traditional Lagrangian tuning is painful, while the “none of the objectives should worsen” semantics of Jacobian descent fit more naturally.
• Physics-Informed Neural Networks (PINNs): Balancing data loss and physics residual loss is notoriously difficult, and the community already has a number of works using Jacobian-descent-style methods.
• Contrastive learning with auxiliary tasks: Main-task loss plus various auxiliary losses. Letting the algorithm decide aggregation directions is often more reliable than manually tuning weights.

Scenarios where you probably shouldn’t force it:

• Single dominant loss with weak regularization: If regularization weights are two orders of magnitude smaller than the main loss, scalarization is enough. Jacobian descent is overkill and costs extra memory.
• Very large models: For LLM fine-tuning scales, full-parameter Jacobians are basically impractical. You’d have to restrict them to key layers, and the actual benefit needs empirical validation.
• Huge numbers of losses: With dozens or hundreds of losses, the Jacobian matrix itself becomes the bottleneck.

The Bigger Context

Multi-loss training has long occupied an awkward spot: there are tons of papers and methods, but most industry teams still rely on basic weighted sums. The problem isn’t that the methods are ineffective—it’s that the engineering barrier is too high. Every paper comes with its own codebase, reproduction is difficult, hyperparameters are sensitive, and unified benchmarks are lacking.

Libraries like TorchJD—which collect an entire subfield’s methods, unify the interfaces, and integrate deeply with PyTorch—have real value in helping these methods reach practical adoption. It’s somewhat analogous to what timm did for vision models or transformers did for NLP—not inventing the methods themselves, but lowering the cost of using them so more people are willing to try.

Judging from project activity, the contributor count has clearly increased since v0.2, and issue responses are fast. The documentation isn’t exceptional but it’s sufficient, and the examples cover typical scenarios like MTL, PINNs, and adversarial training. The project uses the MIT license, so commercial usage is unrestricted.

Summary

If you’re currently running a multi-loss project and are exhausted from tuning scalarization weights, TorchJD is worth spending half a day integrating and testing. In the worst case, performance is similar and you confirm scalarization is sufficient. In a better case, it could save you a month of hyperparameter tuning later on.

Methods are one thing; engineering is another. The biggest value of this library is not that any particular aggregator is overwhelmingly superior—academia still hasn’t reached consensus on which method is “best”—but that it brings this entire line of techniques from paper territory into easy practical use. That’s a good thing for both research and deployment across multi-task learning and constrained optimization.

References

• TorchJD GitHub Repository — source code, documentation, and example implementations
• Reddit r/MachineLearning Discussion Thread — update announcement by the author and community discussion
• Earlier Reddit Discussion — community feedback when v0.2 was released