Gradient calculation.

[magus96]

In Neural_ode.py, we're using the tf automatic differentiation API for calculating the differential, so do we use the adjoint method to calculate the gradient? Sorry but getting a little mixed up here. Thank you.

[kmkolasinski]

Yes, adjoint method is implemented and it requires a gradient computation of Neural Net output w.r.t to its input. You are probably asking about this part of code:

    def _backward_dynamics(self, state):
        t = state[0]
        ht = state[1]
        at = -state[2]

        with tf.GradientTape() as g:
            g.watch(ht)
            ht_new = self._model(inputs=[t, ht])
        # compute at^T df/dz, at^T df / dtheta c.f aug_dynamics from paprt
        gradients = g.gradient(
            target=ht_new, sources=[ht] + self._model.weights,
            output_gradients=at
        )

        return [1.0, ht_new, *gradients]

This is exactly the same what aug_dynamics function from Algorithm 1 from paper. Is it now clear ?

[magus96]

Just a little confused, doesn't the tf.gradient tape calculate the gradients? Is this the gradient during forward mode differentiation? I'll brush up on the paper once again. Just these couple of doubts.

[kmkolasinski]

Yes, gradient tape computes gradients, but in forward mode we do not compute gradients (at least in regular Neural ODEs), we just integrate NN with some selected solver: odeint(net, h0, t_steps). However in backward mode we have to compute jacobian of NN w.r.t its inputs (see the snippet from my previous answer) to accumulate gradients using adjoint method, but note we do it once in each backward step and we don't need to create and memorize whole computational graph. To summarize in forward mode, we solve ODE by integrating NNetwork with selected solver. Then one must define some loss function which generates gradients. These gradients are then plug in into adjoint ODE which additionally requires evaluation of the jacobian-vector-product (JVP): a(t)^T * dNet(h)/dh at every step of the solver.

[magus96]

Thank you. That cleared it out.