BackPACK: per-sample gradients and second-order estimates from one backward pass

One backward pass, several quantities that are not gradients

BackPACK is a PyTorch extension whose purpose is to compute quantities other than the gradient. The name is a pun on the original idea, packing more into backprop, and the listed outputs are specific. Individual gradients from a mini-batch, so the per-sample gradient of the loss rather than the sum. Estimates of the gradient variance or second moment. And approximate second-order information, in both diagonal and Kronecker approximations.

For a practitioner, the first of those is the common case. When you want to know which training example caused a large loss, or you want to weight samples, the gradient of the batch is not enough, because the batch gradient averages the samples together and the individual contributions are lost. The usual workaround is to run the backward pass once per sample, which multiplies the cost of the training step by the batch size.

BackPACK's claim is that this is not the only way. The motivation section states that computing most such quantities is not necessarily expensive, often just a small modification of the existing backward pass where backpropagated information can be reused, but that it is difficult to do in the current software environment. That is the framing to hold on to: the gap is not arithmetic, it is that the standard PyTorch autograd interface gives you one gradient and no hook for the intermediates that a different quantity would need.

The project is the work of Felix Dangel and Frederik Kunstner, with Philipp Hennig as co-author of the paper, and it is hosted at backpack.pt with documentation at docs.backpack.pt.

The quantities are cheap, the interface is the hard part

The interesting engineering claim in the project is not that per-sample gradients are cheap to compute in principle, but that they can be obtained without giving up the vectorised backward pass.

Autograd computes derivatives by propagating information backward through the graph. At each layer, the backward operation receives the incoming gradient and produces contributions to the earlier parameters. A per-sample gradient requires keeping those contributions separated by sample rather than summed over the batch, which means the backward computation has to be shaped differently. The variance and second-order estimates need more than the gradient, since they need the squared or outer-product terms of the same intermediate quantities.

If all of that can be produced during one pass, the extra cost is the memory to hold the per-sample terms and the arithmetic of the extra moments. If it cannot, the fallback is a loop over the batch, and the cost scales with the batch size. This is why the phrase in the description, efficiently computes quantities other than the gradient, is the load-bearing claim in the project, and why the estimates are described as approximate rather than exact.

The project's own honesty about that is worth noting. Individual gradients, variance estimates and the diagonal and Kronecker second-order approximations are all named as provided quantities, but nothing in the documentation claims they are exact or free. A user who needs exact per-sample second derivatives is not the audience, and the approximate label on the second-order entries is doing real work.

The examples are organised as a basic usage page and a use cases section in the documentation, which is where the practical shape of each quantity is worked out.

Installing backpack-for-pytorch and the version floor it sets

Installation is a single pip command:

pip install backpack-for-pytorch

The package name on PyPI is `backpack-for-pytorch`, which is worth noting since the repository and the website both use the shorter BackPACK name. The declared dependencies are compact:

dependencies = [
    "torch>=2.2.0",
    "torchvision>=0.7.0",
    "einops>=0.3.0,<1.0.0",
    "unfoldNd>=0.2.0,<1.0.0",
]

Those four requirements say a lot about the design. `torch` and `torchvision` are the base, `einops` is used for tensor rearrangements, which is what per-sample bookkeeping tends to need, and `unfoldNd` is a small helper for unfolding tensors into windows of a given shape, which is the kind of operation that shows up when the backward pass has to be restructured around batch elements.

The Python requirement is `>=3.9` in the packaging metadata, and the classifiers list Python 3.9 through 3.12. The badge in the project header still says Python 3.8 or later, which is an older claim than the packaging metadata supports, so the metadata is the more reliable of the two. The project is also classified as `Development Status :: 4 - Beta`, and it is not endorsed by or affiliated with Facebook, Inc, whose trademarks cover PyTorch and its logo.

For research code this is a reasonable setup, but the `torch>=2.2.0` floor is a real constraint. An environment pinned to an older PyTorch cannot install the package without upgrading, and upgrading PyTorch in a locked research environment is a larger change than installing a helper library.

What you can compute once BackPACK is installed

The documentation is organised around two entry pages, a basic usage example that shows everything in one place and a use cases section that shows the quantities applied to specific problems. The categories of output are the ones named in the project description.

Individual gradients from a mini-batch give you, per training example, the gradient of the loss with respect to each parameter. That is the quantity behind per-example attribution, outlier detection in a batch, and any reweighting scheme that needs to know which samples contributed what.

Estimates of the gradient variance or second moment give a sense of how noisy each sample's contribution is, which is a different question from how large it is. A sample with a small gradient and high variance is not the same as a sample with a small gradient and low variance, and the second-order information is what lets you tell them apart in a Hessian-like sense.

The approximate second-order information, in diagonal and Kronecker approximations, is the most specialised of the three. A diagonal approximation keeps only the diagonal entries of the second-order matrix, and a Kronecker approximation factorises it into smaller pieces under an assumption about the structure of the curvature. Both are approximations, which is why the project labels them that way, and both exist because the exact object does not fit in memory for a real network.

The decision to try the package is therefore cheap to make and easy to reverse: the outputs are additive to what PyTorch already gives you, and the documentation links to a paper describing the method rather than asking you to trust a benchmark table.

Against looping over the batch, and against a second-order library

The obvious alternative to BackPACK is the loop. If you need per-sample gradients, you can compute the loss for one sample at a time and call backward for each. The result is exact, the implementation takes a few lines, and the cost is a training step that runs as many backward passes as there are samples.

That approach is the baseline BackPACK is measured against, and the difference in approach is precise. The loop recomputes the forward pass for every sample, while BackPACK runs the batch once and keeps the intermediate quantities that the backward pass already produces. The second-order quantities have no loop equivalent that is cheap at all, since an exact Hessian by looping is not a thing anyone does, which is where the diagonal and Kronecker approximations earn their place.

The other alternative is to bring in a dedicated second-order library rather than extend PyTorch itself. The trade-off there is dependency weight and generality: a general second-order tool computes a curvature approximation for an arbitrary model, while BackPACK is scoped to the PyTorch model and computes a family of quantities that include per-sample gradients. The price of that focus is the tight version coupling already noted, since the package has to know the internals of a specific PyTorch version to restructure the backward pass.

The practical question is which quantities you actually need. If it is only per-sample gradients for a few specific layers in a research script, the loop is often the faster route to an answer, because it has no version floor and no assumptions. If it is variance estimates or curvature for a large model, the loop is not an option and a PyTorch-coupled package is the realistic path.

Release history, development state and the upgrade question

The release history is short and dated. Version 1.7.1 was published on 2024-11-15, 1.7.0 on 2024-11-12 and 1.6.0 on 2023-07-12. The repository is not archived, and the last push was on 2026-07-17, which is well after the newest tag. The project description also states that BackPACK is actively being developed and that help is appreciated, with a developer README describing the development procedure.

So there is a gap between the released versions and the working tree. Whether that matters depends on why you want the package. If you need a version that works with your pinned PyTorch, the published releases are the reference point, and the `torch>=2.2.0` floor is the constraint to check first. If you need a change that has not been released, you are reading a moving target.

The licence is MIT, recorded in the packaging metadata and in a `LICENSE.txt` file at the repository root. That is permissive, with no copyleft obligation attached to code that uses the package, and it is the usual situation for academic tooling.

The repository also shows the ordinary infrastructure of a research package: a `pyproject.toml` and `setup.cfg` for packaging, a `pytest.ini` and a `test/` directory, a `makefile`, a `changelog.md`, a `docs/` tree alongside `docs_src/`, a `black.toml` for formatting, and a `.conda_env.yml` for the conda environment used in development. The `fully_documented.txt` file and the `docs_src/` directory suggest a documentation coverage check in the workflow, which matters for a library whose value is in explaining what each quantity means.

For citation, the project asks that you cite the paper, given as the ICLR 2020 proceedings entry for BackPACK: Packing more into Backprop by Felix Dangel, Frederik Kunstner and Philipp Hennig.

Editorial conclusion

BackPACK suits PyTorch researchers who need per-sample gradients, gradient variance estimates or diagonal and Kronecker second-order approximations and would rather reuse the backward pass than loop over the batch. It does not suit a production service that needs a stable API, since the package is classified as beta, requires PyTorch 2.2 or newer, and its newest release on record is 1.7.1 from 2024-11-15 even though the last push was on 2026-07-17. Before adopting, check that your PyTorch version satisfies the declared torch>=2.2.0 floor, and read the linked paper before assuming the variance and second-order estimates are exact.

Frequently asked questions

Official sources

1. f-dangel/backpack on GitHub
2. License: MIT
3. Project website
4. README
5. Releases