How to do L2 regularization with new Flux and Zygote

I think this MWE used to work. Now I don’t know how to compute a norm penalty because the example in Flux documentation here is still using the old Tracked stuff.

using Flux

m = Dense(5,1)

gs = Flux.gradient(Flux.params(m)) do
    sum(norm, Flux.params(m))
end

throws

ERROR: Mutating arrays is not supported
Stacktrace:
 [1] (::Zygote.var"#986#987")(::Nothing) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\lib\array.jl:42
 [2] (::Zygote.var"#2607#back#988"{Zygote.var"#986#987"})(::Nothing) at C:\Users\Henri\.julia\packages\ZygoteRules\6nssF\src\adjoint.jl:49
 [3] push! at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface.jl:65 [inlined]
 [4] (::typeof(∂(push!)))(::Nothing) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface2.jl:0
 [5] params! at C:\Users\Henri\.julia\packages\Flux\oX9Pi\src\functor.jl:42 [inlined]
 [6] (::typeof(∂(params!)))(::Nothing) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface2.jl:0
 [7] params! at C:\Users\Henri\.julia\packages\Flux\oX9Pi\src\functor.jl:48 [inlined]
 [8] (::typeof(∂(params!)))(::Nothing) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface2.jl:0
 [9] params! at C:\Users\Henri\.julia\packages\Flux\oX9Pi\src\functor.jl:48 [inlined]
 [10] (::typeof(∂(params!)))(::Nothing) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface2.jl:0
 [11] params! at C:\Users\Henri\.julia\packages\Flux\oX9Pi\src\functor.jl:48 [inlined]
 [12] (::typeof(∂(params!)))(::Nothing) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface2.jl:0
 [13] params! at C:\Users\Henri\.julia\packages\Flux\oX9Pi\src\functor.jl:45 [inlined]
 [14] (::typeof(∂(params!)))(::Nothing) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface2.jl:0
 [15] params at C:\Users\Henri\.julia\packages\Flux\oX9Pi\src\functor.jl:54 [inlined]
 [16] (::typeof(∂(params)))(::NamedTuple{(:order, :params),Tuple{Array{Array{Float64,N} where N,1},Nothing}}) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface2.jl:0
 [17] #48 at .\none:2 [inlined]
 [18] (::typeof(∂(#48)))(::Float64) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface2.jl:0
 [19] (::Zygote.var"#38#39"{Zygote.Params,Zygote.Context,typeof(∂(#48))})(::Float64) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface.jl:101
 [20] gradient(::Function, ::Zygote.Params) at C:\Users\Henri\.julia\packages\Zygote\8dVxG\src\compiler\interface.jl:47
 [21] top-level scope at none:0
3 Likes

Sum with function argument is also known to cause bugs with zygote
https://github.com/FluxML/Zygote.jl/issues/231

1 Like

Better use Optimiser(WeightDecay(lambda), opt) for l2 regularization.

2 Likes