Stop using implicit style differentiating

[ablaom]

It seems the style used here is being deprecated and won't work with Flux 0.14:

MLJFlux.jl/src/core.jl

Line 37 in 452c09d

gs = Flux.gradient(parameters) do

---

edit After discussion below, I suggest we wait on

• WeightDecay for L1 norm Optimisers.jl#159

and refactor to use a optimiser-based solution to weight regularisation, which will avoid current limitations of explicit differentiation outlined in the discussion. Note, this will likely mean the reported training_loss must change, as it will no longer include the weight penalty. So this will be breaking.

[mcabbott]

Relatedly, it would be nice if the MLJFlux models listed here https://github.com/FluxML/model-zoo#examples-elsewhere could be updated to use latest Flux, and avoid implicit gradients.

Examples of similar upgrades: https://github.com/FluxML/model-zoo/issues?q=is%3Aclosed+label%3Aupdate+explicit

In the end, Flux 0.14 did not drop support for implicit gradients, but 0.15 should.

[ablaom]

@pat-alt Would you have any time and interest in addressing this issue?

[pat-alt]

That actually syncs well with some of my other outstanding issues and I think I'll have to address this very same thing in CounterfactualExplanations.jl soon. So yes, please feel free to assign to this one to me and I'll look at it in the coming weeks 👍

[pat-alt]

I have added a draft for this with very minor changes here #230:

function train!(model::MLJFlux.MLJFluxModel, penalty, chain, optimiser, X, y)
    opt_state = Flux.setup(optimiser, chain)
    loss = model.loss
    n_batches = length(y)
    training_loss = zero(Float32)
    parameters = Flux.params(chain)
    for i in 1:n_batches
        batch_loss, gs = Flux.withgradient(chain) do m
            yhat = m(X[i])
            pen = penalty(parameters) / n_batches
            loss(yhat, y[i]) + pen
        end
        training_loss += batch_loss
        Flux.update!(opt_state, chain, gs[1])
    end
    return training_loss / n_batches
end

Currently, the following test fails:

[ Info: regularization has an effect:
[ Info: acceleration = CPU1{Nothing}(nothing)
regularization has an effect (typename(CPU1)): Test Failed at /Users/patrickaltmeyer/code/MLJFlux.jl/test/integration.jl:25
  Expression: !(loss2 ≈ loss3)
   Evaluated: !(0.8354643267207931 ≈ 0.8354643267207931)

I'm not quite sure what's happening. @mcabbott can you spot anything obviously wrong this?

[ToucheSir]

That's because the regularization term is still using implicit params. Something like FluxML/Flux.jl#2040 (comment) will be needed for explicit params.

[mcabbott]

parameters = Flux.params(chain) outside the gradient context will only work in the implicit style -- changing the explicit local m will not change pen. (Edit -- as ToucheSir says, while I was typing!)

What is penalty? For L2 it will be better to use WeightDecay like this: http://fluxml.ai/Flux.jl/stable/training/training/#Regularisation

[pat-alt]

Thanks both!

What is penalty? For L2 it will be better to use WeightDecay like this: http://fluxml.ai/Flux.jl/stable/training/training/#Regularisation

Currently, penalty functions are explicitly defined callable objects in MLJFlux (see here). I saw the note on WeightDecay in the Flux docs and was wondering if it's worth changing that.

In any case, I can't really get either of the approaches you suggest to work in this particular case, so we may indeed want to rethink the implementation of the penalty functions, for example by using WeightDecay instead. Will require a little extra work, but should be doable. @ablaom what do you think?

[ToucheSir]

I can't really get either of the approaches you suggest to work in this particular case

Can you elaborate? I'm not sure I understand why/how they wouldn't work.

[pat-alt]

Sure!

Moving the params call inside as follows

function train!(model::MLJFlux.MLJFluxModel, penalty, chain, optimiser, X, y)
    opt_state = Flux.setup(optimiser, chain)
    loss = model.loss
    n_batches = length(y)
    training_loss = zero(Float32)
    for i in 1:n_batches
        batch_loss, gs = Flux.withgradient(chain) do m
            yhat = m(X[i])
            pen = penalty(Flux.params(m)) / n_batches
            loss(yhat, y[i]) + pen
        end
        training_loss += batch_loss
        Flux.update!(opt_state, chain, gs[1])
    end
    return training_loss / n_batches
end

the tests just seem to get stuck at some point. I may try and commit this now, but at least locally on my machine things get stuck.

Alternatively, using the approach in FluxML/Flux.jl#2040 (comment) as follows

function train!(model::MLJFlux.MLJFluxModel, penalty, chain, optimiser, X, y)
    opt_state = Flux.setup(optimiser, chain)
    loss = model.loss
    n_batches = length(y)
    training_loss = zero(Float32)
    for i in 1:n_batches
        batch_loss, gs = Flux.withgradient(chain) do m
            yhat = m(X[i])
            l = loss(yhat, y[i])
            reg = Functors.fmap(penalty, m; exclude=Flux.trainable)
            l + reg / n_batches
        end
        training_loss += batch_loss
        Flux.update!(opt_state, chain, gs[1])
    end
    return training_loss / n_batches
end

I get the following error:

[ Info: acceleration = CPU1{Nothing}(nothing)
┌ Warning: Layer with Float32 parameters got Float64 input.
│   The input will be converted, but any earlier layers may be very slow.
│   layer = Dense(5 => 15)      # 90 parameters
│   summary(x) = "5×20 Matrix{Float64}"
└ @ Flux ~/.julia/packages/Flux/n3cOc/src/layers/stateless.jl:60
fit! and dropout (typename(CPU1)): Error During Test at /Users/patrickaltmeyer/code/MLJFlux.jl/test/test_utils.jl:38
  Got exception outside of a @test
  TypeError: non-boolean (NamedTuple{(:layers,), Tuple{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dropout{Float64, Colon, Random.TaskLocalRNG}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}) used in boolean context

Perhaps it has to do with the fact that the penalizers aren't Functors?

[ToucheSir]

Yeah I wouldn't try the first version you have there, was referring to the second one or @mcabbott's suggestion about moving things to the optimization step.

I get the following error: ...

Pretty sure that's due to a typo in the original example code snippet. See FluxML/Flux.jl#2040 (comment)

[pat-alt]

hmm in that case I get the following error: MethodError: no method matching Dense(::Float32, ::Float32, ::typeof(identity)). Any ideas?

[ablaom]

Thanks @pat-alt for this work!

In any case, I can't really get either of the approaches you suggest to work in this particular case, so we may indeed want to rethink the implementation of the penalty functions, for example by using WeightDecay instead. Will require a little extra work, but should be doable. @ablaom what do you think?

WeightDecay only provides a mechanism for L2 regularisation, but the current implementation provides for a combination of both L1 regularisation (good for feature selection) and L2 regularisation. It seems a pity to drop support of a feature to accomodate the new explicit syntax.

I don't know what the source of your current issue is.