TorchJD: Jacobian Descent for Multi-Loss Training in PyTorch

The multi-loss problem TorchJD is built around

Most training loops assume one number to minimize. Multi-task learning breaks that assumption: a shared trunk feeds several heads, each with its own loss, and someone has to decide how those losses become a single gradient. The usual answer is a weighted sum, and the usual consequence is that the weights become hyperparameters you tune by hand. When two tasks want the shared parameters to move in opposite directions, the sum can cancel itself out and the shared trunk stops learning anything useful. TorchJD targets exactly that situation. Its README frames the library as supporting two complementary approaches, scalarization and Jacobian descent, and states that Jacobian descent 'allows taking conflict-free optimization directions, which can resolve problems that may be impossible to solve with standard scalarizers.' The audience is therefore narrow and specific: people already doing multi-task or multi-objective training in PyTorch who have hit the limits of manual weighting, not people who want a general optimizer replacement.

Scalarization as the baseline path

The first approach TorchJD offers does not change the shape of training at all. A scalarizer takes the vector of per-task losses and returns one scalar, which you then call backward on as usual. The README's diff shows this concretely: you import GeometricMean from torchjd.scalarization, construct it once, and replace the line loss = loss1 + loss2 with loss = scalarizer(torch.stack([loss1, loss2])). Everything downstream, including loss.backward(), optimizer.step() and optimizer.zero_grad(), stays as it was. The README lists geometric mean and softmax weighting among the supported scalarizers and notes that this approach 'is often a good baseline.' That phrasing is worth taking at face value. Scalarization here is not the pitch, it is the control group. If a scalarizer solves your problem, you have avoided the cost of computing a Jacobian entirely.

How Jacobian descent changes the update direction

The second approach is the one the library is named for. Instead of collapsing the losses before backprop, TorchJD computes the Jacobian matrix of the losses with respect to the parameters, then aggregates that matrix into an update direction using an aggregator. The README names UPGrad, MGDA and CAGrad among the supported aggregators, and links to an arXiv paper on Jacobian descent for the underlying method. The distinction from scalarization is architectural, not cosmetic. A scalarizer commits to a combination rule before any gradient information exists, so it cannot see whether the tasks conflict. A Jacobian-based aggregator sees the full matrix of per-loss parameter sensitivities and chooses a direction from it, which is what makes a conflict-free direction expressible at all. The trade is that you now materialize and process a Jacobian each step rather than a single gradient vector, and the aggregator's own algorithm sits between you and the optimizer.

Installing it and the dependency question

The README gives one installation command: pip install "torchjd[quadprog_projector]". That extra is not decorative. The README states it includes the dependencies required by UPGrad and DualProj, and adds that some other aggregators may have additional dependencies, pointing to the installation page at torchjd.org/stable/installation for those. This is the first thing to check before committing, because your aggregator choice determines your dependency footprint. The README also carries a PyTorch badge reading >=2.3, so the floor on the framework side is stated rather than implied. Python version support is exposed through the PyPI classifiers linked from the README, not restated in the text. The base install is a single pip command; the aggregator-specific extras are where the setup stops being trivial, and the README does not enumerate them.

Where the documentation stops and you have to decide

The supplied README is truncated mid-section, at a heading that reads Jacobian and nothing after it. That matters for anyone evaluating the library from this page alone: the scalarization diff is shown in full, but the Jacobian descent usage example is not present in the material, and neither is the list of supported aggregators and weightings that the README references by anchor. The same applies to the installation extras beyond quadprog_projector. None of this is evidence that the documentation is incomplete on torchjd.org, which the README presents as the full documentation site. It does mean that the GitHub README is a landing page, not a reference. Plan to read the site before you write code, and treat the README's diff as the only worked example you have in hand.

The licence is unresolved, and that is a real blocker

The repository metadata reports the licence as NOASSERTION, which means the licence could not be identified automatically from the repository contents. The README does not state a licence either. For a library you intend to vendor into a training stack, that is a concrete unknown rather than a formality. There is no SPDX identifier to check, no licence badge in the README, and nothing in the supplied material that names terms. If your organization has any policy gate on third-party dependencies, this is the item to resolve before anything else, and it is a question for the maintainers or for the repository's licence file, not something to infer. I am not giving legal advice here, only noting that the metadata as reported does not answer the question.

A real alternative: plain weighted-sum training

The honest alternative is what most multi-task PyTorch code already does: assign weights to each loss, sum them, and call backward once. The difference in approach is not about implementation quality, it is about when the combination decision is made. A weighted sum fixes the trade-off before the backward pass, using numbers you chose in advance, and produces one gradient that the optimizer consumes directly. TorchJD's Jacobian path defers that decision: it computes the per-loss Jacobian, then lets an aggregator such as UPGrad, MGDA or CAGrad pick a direction from the matrix. The weighted sum has no extra dependencies, no Jacobian to materialize, and no aggregator to tune, which is why the README itself calls scalarization 'often a good baseline.' Reach for the Jacobian path when the weighted sum has demonstrably failed, not as a first move. There is also a middle option inside TorchJD: swap the manual weights for a scalarizer like GeometricMean and keep the rest of the loop unchanged.

Maintenance, releases and what to verify before adopting

The release history shows v0.17.0, v0.16.0 and v0.15.0 dated within roughly ten days of each other in June 2026, and the repository's last push is dated 2026-09-03. Frequent minor releases before a 1.0 tag mean the API surface is still moving, so pin the version you develop against and read the release notes when you bump it. The README points to a Discord community channel and a documentation site, and the repository carries CI and coverage badges, so there is an active maintenance surface. What the material does not give you is a support commitment or a deprecation policy. Before adopting, verify three things: the licence terms, the extra dependencies your chosen aggregator needs beyond quadprog_projector, and whether the Jacobian computation fits your step budget on your actual model size. The last one you can only answer by running it, which is the point at which the README stops helping.

Editorial conclusion

Adopt TorchJD when your multi-task model has losses that pull the shared parameters in opposing directions and a weighted sum has already failed you. Do not adopt it as a default replacement for loss weighting on a model with a single objective, or when the extra dependency set and per-step Jacobian computation are more than your training budget allows. Before wiring it into a real run, verify two things: which aggregator you intend to use and what that aggregator requires beyond the base install, and whether your PyTorch version satisfies the >=2.3 floor stated in the README badge. Then reproduce the scalarization diff from the README on your own model and compare it against the Jacobian descent path on the same architecture.

Official sources

1. Issues
2. Project website
3. README
4. Releases
5. SimplexLab/TorchJD on GitHub