How to do implicit differentiation for neural network parameters?

[JinraeKim]

Hi!

I'm trying to do implicit differentiation of the parameters for a given neural network.

For example, nn is a neural network constructed by Flux.jl. I can get the parameters by Flux.params(nn).

In this tutorial, I need to provide the parameters as the arguments of the forward solver function (here, it's lasso(data::ComponentArray)).

But I don't know how to do this for my case; namely, optimization_prob(parameters) = ...?

If I can overwrite network parameters with p = Flux.params(nn), then I would be able to do so like

function optimization_prob(parameters)  # will be provided by `Flux.params(nn)` outside this function
    load_parameters!(nn, parameters)
    # etc...
end

Is there any workaround for this?

[mohamed82008]

Once you create the implicit function, you can nest it with other functions. I don't understand what you are trying to do if it's not a simple nesting of 2 functions. Just make sure your second function you are nesting with the implicit function is Zygote compatible on its own first.

[mohamed82008]

@JinraeKim can you please post a complete example here?

[JinraeKim]

@mohamed82008
I was trying to do something like this before but I didn't complete it yet as I got busy and it's related to my side project :>

I may begin my side project this year, so if I figure it out or struggle with this, will post a comment!

[JinraeKim]

My workaround for this is to define the last layer of a neural network of interest (related to solving an optimization problem) by this package and put the intermediate values of the neural network to be the optimization parameters of the last layer, see #67