Thanks for your time.
The code is very long and our paper isn’t out yet but hopefully I can bring some clarifications instead. When doing backpropagation through time for gradient based hyperparameter optimization, forward mode differentiation is used because its memory cost scales as O(KD) instead of O(TD) for reverse mode, where K=n_hyperparameters, D=n_weights, T=n_iterations. Here my goal is to get the (hyper)gradient of some validation loss after T steps with respect to K hyperparameters used in those steps.
In practice this means we keep updating a matrix Z of shape (D,K) for each step, such that Z_t=dparams_t/dhyperparams. This update rule requires the Hessian product shown above (see Eqt (11) and (14) of this paper if interested http://proceedings.mlr.press/v70/franceschi17a/franceschi17a.pdf 4) which accounts for most of the compute time. Once we have Z_T at the final iteration we can easily get the hypergradients desired with the chain rule.
My feeling is that since I am doing exactly the same operations as the reverse mode but in a different order, I should be able to get similar times.