Repository: pytorch
Issue: [feature request] Simple and Efficient way to get gradients of each element of a sum
Labels: feature,module: autograd,triaged

For some application, I need to get gradients for each elements of a sum.
I am aware that this issue has already been raised previously, in various forms ([here](https://discuss.pytorch.org/t/gradient-w-r-t-each-sample/1433/2), [here](https://discuss.pytorch.org/t/efficient-per-example-gradient-computations/17204), [here](https://discuss.pytorch.org/t/quickly-get-individual-gradients-not-sum-of-gradients-of-all-network-outputs/8405) and possibly related to [here](https://github.com/pytorch/pytorch/issues/1407)) 
and has also been raised for other autodifferentiation libraries (some examples for TensorFlow: [here](https://github.com/tensorflow/tensorflow/issues/675), long discussion [here](https://github.com/tensorflow/tensorflow/issues/4897))

While the feature does exists in that there is a way to get the desired output, 
I have not found an efficient way of computing it after investing a significant amount of time.

I am bringing it up again, with some numbers showing the inefficiency of the existing solutions I have been able to find.

It is possible that the running time I observed is not an issue with the existing solutions but only with my implementation of them, 
or that my attempts have been misguided and that a simpler solution exists.
If there is a way to make that operation in an efficient way, I would love to know about it.

I can spend more time working on this issue, but as of now I do not know of a way forward.
The following contains a description of the desired feature, existing "workaround" and some evaluation of their performance.

---

## Feature description:

Given an objective that is a sum of functions, `f(x) = f_1(x) + ... + f_N(x)`, 
a _simple to use_ and _efficient_ way to compute the gradient with respect to `x` for each of the individual function `f_n`, i.e., getting
`[∇f_1(x), ..., ∇f_N(x)]`

If there already is a way to do so, an example would be very nice.

(I understand that this would involve additional memory overhead if `N` is large, as pointed out [here](https://discuss.pytorch.org/t/quickly-get-individual-gradients-not-sum-of-gradients-of-all-network-outputs/8405/2) - my setting is not bound by memory but by time)

## Use case:
- Computing approximate second-order information (Generalized Gauss-Newton type of algorithm).
- Computing gradient statistics such as the variance at a given point.
Some papers based on those ideas: [here](https://pdfs.semanticscholar.org/42e2/1cd78f578fa6ce61b06b99848697da85ed76.pdf), [here](http://www.cs.toronto.edu/~jmartens/docs/Deep_HessianFree.pdf)

## Existing workarounds:
- The `naive` implementation: do the forward pass in batch up to the last part of the function to minimize overhead, and then call backward on each individual `f_n(x)`.
- Goodfellow showed how to recover the individual gradients from gradients with respect to the activation functions in the case of feed-forward neural networks [here](https://arxiv.org/abs/1510.01799). This requires an additional derivation and performs the linear parts of the transformation in Python instead of C, making it scale badly. 
- What I call "Multiple Models": Define copies of the parameters `x_1, ..., x_N` and compute `grad(f_1(x_1) + ... + f_N(x_N))`. This ensures that the gradients are not accumulated, however it scales poorly. Adapted from the original formulation [here](https://github.com/tensorflow/tensorflow/issues/4897#issuecomment-290997283)


## Evaluation of existing workarounds
I have tried to make the simplest example of a sufficiently complex problem where this becomes an issue.
The code to reproduce these is available [here](https://github.com/fKunstner/fast-individual-gradients-with-autodiff).

Running on a simple Multi-Layer Perceptron, I am comparing the running time of
* `full`: computing the gradient of the full objective function
* `naive`, `goodf`, `multi`: computing the gradient of the sum by first computing individual gradients with the methods described above, and then taking the sum[1].

I am taking a sum over `N = 1000` elements, and tried different network configurations;
```
# D: Dimension of input and of each Layer
# L: Number of hidden layers
```
For a (relatively) wide and shallow network (`D, L = 500, 1`), I get the following running time
```
Full  :  0.04s
naive : 15.42s
goodf :  4.42s
multi :  1.89s
```
For a narrower and deeper network (`D, L = 100, 10`) I get
```
Full  :  0.03s
naive : 11.40s
goodf :  1.70s
multi :  1.60s
```

While `goodf` and `multi` definitely are improvements on the naive method, they are still ~50 time slower than simply computing the sum. I would not expect any method that stores individual gradients to perform as well as a method that can simply throw them away, but they are essentially doing the same operations so it should be possible to do better.
With those numbers, training a small-ish neural network that would normally take ~10 minutes would take 10 hours, making experimentation very difficult.

From the [information I have been able to find](https://discuss.pytorch.org/t/how-the-hook-works/2222/2), it seems difficult to get access to the un-aggregated gradients.
It is a bit frustrating to know that the backward pass on the function computes the sum in almost no time, and all that is needed for this feature is a way to intercept elements of this sum.

A hook that would allow to store individual gradients on the fly in some list or tensor - maybe with the use of those hooks or by defining a new sum function that intercepts gradients - would be amazing, but I do not understand the codebase and/or AD to do it. 
Some pointers in that direction would already be a big help

[1] _Note: I am not looking for an expensive way of computing the sum by computing individual gradients - I am taking the sum in the end so that all methods do the "same thing", but what is really needed is the intermediate matrix containing individual gradients_






cc @ezyang @SsnL @albanD @zou3519