I think of an nn.Module to be a function f that transforms an input x into some output y as y = f(x). Assuming all my layers are invertible, is there a direct way to compute the backward pass for x = f^-1(y)?

Hi Zeeshan!

If your layers are indeed invertible and I interpret “compute the backward pass”

to mean compute the Jacobian of f^-1, then yes, this Jacobian is the matrix

inverse of the Jacobean of f, J[f^-1] = (J[f])^-1.

This statement, together with important relevant context, can be found in

Wikipedia’s inverse-function-theorem entry.

If instead you are asking how to compute f^-1 itself, that, in general, can be

difficult. You could either compute the inverses of each of your layers separately

and then chain those inverses together, or you would compute the inverse of f,

x = f^-1 (y), by numerically solving for x in f (x) = y.

Best.

K. Frank

Thank you for the reply. I was asking about the backward pass f^-1, but thanks for pointing the inverse-function-theorem. Is there a way to use Pytorch’s autodiff for such a computation?

Hi Zeeshan!

You may use jacobian() or jacrev() to compute the Jacobian of your `Module`

and

inv() to compute its inverse. Note the Jacobian will be returned as a tuple of tensors

and you will have to reorganize it into a two-dimensional matrix in order to pass it

to `inv()`

.

Bear in mind that the mapping from input to output for a typical neural network

will not be invertible. Unless your neural network / `Module`

is in fact invertible

you won’t be able to invert the Jacobian.

Best.

K. Frank