TorchJD routes a multi-loss PyTorch loop through a scalarizer or an aggregator

Scalarization or Jacobian descent, decided before the loop starts

A multi-task loop already exists in most codebases: compute one loss per head, sum them, call backward. TorchJD keeps that shape and changes what happens between the loss values and the parameter update.

The first route is scalarization. A scalarizer such as GeometricMean, imported from torchjd.scalarization, turns the vector of losses into a single scalar before backprop, so nothing else in the loop has to move. The methods come out of the multi-objective optimization literature, and the project itself calls this route a good baseline, which is a sound thing to measure against before taking on anything heavier.

The second route is Jacobian descent. Instead of one backward pass over a combined scalar, each loss is differentiated on its own, the Jacobian of the losses with respect to the parameters is formed, and an aggregator collapses that matrix into a single update direction. UPGrad, MGDA and CAGrad are the three names given in the summary, with a longer aggregator list and its weightings held behind a link. UPGrad is described as built to find directions that benefit all losses simultaneously, and that conflict-free property is the reason to accept the extra matrix arithmetic, since the project states it can resolve problems that may be impossible for standard scalarizers.

Both routes sit on the same autograd graph, so the difference is which object you construct and which call you make.

autojac backward fills the jac field, and the optimizer reads grad

The autojac engine deliberately mirrors the autograd interface, which makes it easy to read and easy to misjudge. autojac.jac stands in for autograd.grad but returns Jacobians instead of gradients. autojac.backward stands in for autograd.backward, and here is the catch: it accumulates into the .jac fields of the parameters, not into the .grad fields the optimizer is built to read.

A loop that calls autojac.backward and then steps the optimizer has nothing for it to consume until the Jacobians have been aggregated and moved across. That step is autojac.jac_to_grad, which takes the accumulated Jacobians, aggregates them into a gradient, and puts it in .grad. Leave it out and the parameter update runs on an empty gradient, and because the values involved are ordinary PyTorch fields, nothing in the stack complains about it.

The multi-task case splits state a second time. mtl_backward is built for a model with a shared trunk and separate task heads: it computes and accumulates gradients with respect to the task-specific parameters and Jacobians with respect to the shared parameters. One call, two kinds of accumulated state, and the distinction matters as soon as a loss touches both the trunk and its own head.

The pip extra covers UPGrad and DualProj, and the other aggregators are not covered

Installation is one pip command, and the extra inside it is doing real work:

pip install "torchjd[quadprog_projector]"

That extra pulls the dependencies required by UPGrad and DualProj, and nothing more is promised. The project is direct that some other aggregators carry additional dependencies and points readers at the installation page under torchjd.org/stable/installation for those cases. In other words, the aggregator you intend to use decides whether the single line above is enough, and the name of the extra gives a reader no way to guess which ones those are.

This is the practical trap of a library whose entry points are all optional. An install that succeeds can still leave an aggregator failing later because a solver dependency was never pulled in, and the repair is another install rather than a change to the training loop. Matching the pinned version to the PyTorch build on the machine is the other half of the job.

Two destinations are linked from the project page for the same package: the PyPI project page for torchjd, and the documentation site at torchjd.org, where every scalarizer and aggregator has its own page.

The torch floor reads 2.3.0 while the comment beside it reads 2.4.0

pyproject.toml depends on torch>=2.3.0 and carries a comment on that very line: problems before 2.4.0, especially with autogram. The floor written in the requirement string is therefore not the version the authors would pick for themselves, and autogram, the engine whose entire purpose is to make Jacobian descent workable on large models, is named as the part that suffers on the older releases. A reader who takes 2.3.0 at face value adopts the one configuration the comment warns about.

Python is tidier. requires-python is >=3.10, and the classifiers run from 3.10 through 3.14, so five interpreter versions are declared. The operating system is listed as independent, and the intended audience is developers and science or research, which fits a library that lives inside a training script rather than behind a service boundary.

The project status is declared as Development Status 4, Beta, at version 0.17.1. On a library whose whole surface is a set of numerical aggregators, a beta label is a fair warning about how much the API can still move. A CHANGELOG.md does sit at the top level, so the record needed to pin a version is in the tree.

autogram keeps the Gramian and drops the Jacobian

Differentiating one loss is cheap. Holding the Jacobian of many losses over millions of parameters is not, and the aggregators need the Gramian of that Jacobian rather than the matrix in full. That single substitution is what the autogram engine exists for: it computes the Gramian incrementally and never stores the full Jacobian in memory, which is what makes Jacobian descent feasible on large models where storing the matrix would be too expensive.

The two engines are not rivals. autojac covers the case a reader can hold in their head, where per-loss Jacobians are accumulated into .jac fields and handed to an aggregator. autogram covers the case where that memory is the binding constraint, and the version floor discussed above is where the two paths diverge most sharply, since the pyproject comment singles out autogram as the piece that misbehaves on older PyTorch releases.

The practical half of this is in the examples, where the project names instance-wise risk minimization and partial Jacobian descent, both collected under the examples path of the documentation site. Those two names are the ones to look up when a model is only partly multi-task, and the aggregations they exercise are the ones that decide whether a direction is conflict free.

Three scalarizers carry a dash where the others carry a publication

The README's scalarizer table pairs each name with a publication, and that pairing is where the design intent shows through. Constant, Mean and Sum carry a dash in the publication column, which is the table saying that the trivial combinations need nothing defended. Everything else in the visible table brings a citation: COSMOS, DWA, FAMO, GeometricMean, IMTL-L, PBI, Random and STCH.

Two of those links are not arXiv entries. PBI points at an IEEE document, and IMTL-L points at a university discovery page. Random cites a paper arguing that random weighting is a reasonable default and a litmus test for multi-task learning, which is a useful precedent for treating scalarization as a baseline to beat rather than a fallback to apologize for. STCH and COSMOS point at arXiv identifiers numbered 2402 and 2303, and GeometricMean traces back to a 2019 multi-stream feature aggregation paper.

The visible portion of the table ends at Sum, so the menu of scalarizers is longer than the eleven names shown here, and the aggregator and weighting tables linked from the opening summary do not appear on the project page at all. The documentation site is where the rest of the list can be read.

Both walkthroughs are diffs that stop before the library is called

Each usage walkthrough is written as a diff against a loop that already works, which is a good default: the shared trunk of two linear layers with ReLU between them, the two task heads, the parameter list, the MSELoss, and SGD at lr 0.1 all stay untouched. The scalarization diff adds one import and one object:

  import torch
  from torch.nn import Linear, MSELoss, ReLU, Sequential
  from torch.optim import SGD

+ from torchjd.scalarization import GeometricMean

and then places `scalarizer = GeometricMean()` next to the optimizer. The descent diff adds two imports and one aggregator:

+ from torchjd.autojac import jac_to_grad, mtl_backward
+ from torchjd.aggregation import UPGrad

What the visible text does not reach is the problem. The scalarization walkthrough ends mid statement at `task1_targets = torch.randn(8`, and the descent walkthrough stops inside the comment on the inputs line. Neither one shows the call that combines the losses or the call that aggregates the Jacobians, which are the two lines the library exists to supply. Those calls live on the autojac and scalarization pages of the documentation site, so the project page teaches the shape of the change and the docs carry the change itself.

src, docs, tests and a skills directory sit in one tree

The top level of the repository is a full project rather than a package dump. There is src/ for the code, docs/ for the site that torchjd.org is built from, and tests/ for the suite. Next to them is a skills/ directory that the project page never explains. The process files are written down too: AGENTS.md, CODEOWNERS, CODE_OF_CONDUCT.md, CONTRIBUTING.md and GOVERNANCE.md sit alongside a .pre-commit-config.yaml and a codecov.yml, and the CI badge points at a workflow called checks.yml.

Licensing is the one place where packaging data and the repository disagree. The license attached to the repository reads NOASSERTION, while pyproject.toml embeds the full MIT text with a copyright line naming Valerian Rey and Pierre Quinton. A LICENSE file and a NOTICES file both exist at the top level, so the permissive terms are not in doubt, even though the field a package index reads is empty.

Release numbering is plain and the dates sit close together: v0.16.0 on 22 June 2026, v0.17.0 two days later on 24 June, and v0.17.1 on 23 September 2026, which is also the date of the most recent push to main. A 0.x line with three releases inside one quarter, the last of them shipped on the day the tree last moved, is a repository being worked on, though the version number and the beta classifier are the only status signals the packaging metadata offers.

Editorial conclusion

TorchJD earns a place in a research training script that already has a multi-task trunk and task heads, and a poor one in a serving or production pipeline, since every route here ends in a numerical aggregator that has to be read before it is trusted. Before adopting it, check which aggregator you need against the dependencies the pip extra actually installs, hold your PyTorch at 2.4.0 or above given the note about problems before that release, and read the autojac pages closely enough to know that .grad is not written until jac_to_grad runs.

Frequently asked questions

Official sources

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