-
|
Hey, I am new to Jax - up to now I've been doing some fairly hobby-level coding using mygrad, which I think I'm now reaching the limits of. I have one main Q about jax before I dive in - can I implement metagradient descent? That is can I treat some function that implements gradient descent using jax as an opaque solver for some outer gradient descent algorithm? While JAX's grad tracers support this or will they get confused as the outer ones are included in the inner's computation graph? Am I barking mad? |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 4 replies
-
|
In general JAX transformations are composable, so you can use autodiff to compute the gradient of another function that uses autodiff. That said, there may be some limitations: for example, JAX does not implement reverse-mode autodiff of unbounded loops, so if your inner solver uses a |
Beta Was this translation helpful? Give feedback.
-
|
Gradient-based meta-learning was one of the first projects built in the original Autograd, which is JAX's direct predecessor (JAX is basically "Autograd 2"). So being able to do this kind of composable higher-order AD is close to our hearts! |
Beta Was this translation helpful? Give feedback.
-
|
I just remembered that for a lecture I made a colab with a toy model of gradient-based metalearning. Look for the section "What's meta-learning look like?" |
Beta Was this translation helpful? Give feedback.
-
|
Thanks both, now I just need to work out how on earth to convert my highly OOP to be jax compatible - I may be asking some pytree questions as the whole flatten/unflatten thing is quite confusing! |
Beta Was this translation helpful? Give feedback.
-
|
Returning to my party here. I've got a pretty sweet 'inner' solver but am struggling with 'noise' in my outer solver. Are there limitations around how many iterations the inner loop could have? I did use a while loop in my inner loop (a 'Sweeping' my loss for the 'outer' optimiser shows a super smooth loss function, but the gradients of this are extremely noisy. At this point I'm in quite deep and am running out of ideas of how to debug this - any pointer/common gotchas to look out for? (i've already switched on 64 bit for more precision) I note that reducing the iterations in the inner loop seems to decrease the noise - am I suffering from some kind of vanishing gradient issue? |
Beta Was this translation helpful? Give feedback.
In general JAX transformations are composable, so you can use autodiff to compute the gradient of another function that uses autodiff.
That said, there may be some limitations: for example, JAX does not implement reverse-mode autodiff of unbounded loops, so if your inner solver uses a
while_loopfor convergence, you'll be limited to forward-mode autodiff.