How to apply backpropagation when one layer of the network is a call-only function (no gradient)?

Question

I worked with Feed Forward Neural Network and VAE and understood backpropagation algorithm. Now I build a VAE network, one layer of it is a very complex vector-to-vector function $f(x)$ (a general 'method' in the programming sense instead of a 'math' expression).

Thus, there is no gradient info for this layer. I guess one cannot train the entire network with such a gap, though other layers are differentiable. Is there any nice way to train such a network?

One thing in my mind is to approximate the gradient by slightly changing $x$ and computing $$ \frac{f(x+\delta)-f(x)}{ \delta} $$

Answer 1

Well, you can specify a custom gradient by either being just the identity (i.e. returning the inputs in the gradient scope) or computing the gradient by hand if you know that expression.

Otherwise, you can approximate the gradient of non-differentiable functions by setting-up a score function estimator also called the REINFORCE estimator: also have a read here.

The finite difference method that you mentioned may work too, but it's often unstable for very small $\delta$ which are required in order to get good estimates.

Answer 2

You can maybe use a similar re-parameterisation trick. Where you may approximate the gradient rather than calculating it accurately.

Introduce auxiliary variables which can mimic the function an still be differentiable. For example, as you suggested in a VAE when we have to pick a random vector from a distribution, we don't actually pick a random vector. But we pick a $\mu$ and $\sigma$ (i.e mean and variance) to train on. And then supple the random vector as an output of that like random vector $z$ = $\mu$+ $\epsilon*\sigma$ where $\epsilon$ ~ 𝑁(0, 1).

Kindly check this reference: Reparameterisation trick. Also kindly provide an example of your vector function so that we can construct a more definite answer.