LSTM cannot be trained successfully with the latest release version

[AStupidBear]

The following script works fine using Flux.jl#tracker, but get stuck on a "plateou" of loss=0.11 using Flux.jl#zygote.

using Flux, Random, Statistics

Random.seed!(1234)

x = randn(Float32, 10, 1, 100)
y = mean(x, dims = 1)

model = Chain(LSTM(10, 100), LSTM(100, 1))

function loss(x, y)
    xs = Flux.unstack(x, 3)
    ys = Flux.unstack(y, 3)
    ŷs = model.(xs)
    l = 0f0
    for t in 1:length(ŷs)
        l += Flux.mse(ys[t], ŷs[t])
    end
    return l / length(ŷs)
end

opt = ADAM(1e-3, (0.9, 0.999))

cb = () -> (println(Flux.data(loss(x, y))); Flux.reset!(model))

Flux.@epochs 10 Flux.train!(loss, params(model), repeat([(x, y)], 100), opt, cb = cb)

This is also reproducible on GPU, JuliaBox, and other computers. You can verify this by yourself by running this script in two environments with Flux.jl#v0.9.0 and Flux.jl#v0.10.4 installed seperately.

Since LSTM implementations are the same across Flux.jl#v0.9.0 and Flux.jl#v0.10.4, I suspect that this is a Zygote issue, but cannot find the root cause and fix it.

[AStupidBear]

I digged into a little bit and found that the gradient of recurrent layers calculated by Zygote.jl deviated too much.

using Flux, Random, Statistics

Random.seed!(1234)

model = Chain(LSTM(10, 100), LSTM(100, 1)) |> f64

θ0, re = Flux.destructure(model)

function loss(x, y, model = model)
    xs = Flux.unstack(x, 3)
    ys = Flux.unstack(y, 3)
    ŷs = model.(xs)
    l = 0.0
    for t in 1:length(ŷs)
        l += Flux.mse(ys[t], ŷs[t])
    end
    return l / length(ŷs)
end

loss_sciml(θ, x, y) = loss(x, y, re(θ))

function ngradient(f, x, n = length(x))
    Δ = zero(x)
    for i in 1:min(n, length(x))
      δ = sqrt(eps())
      tmp = x[i]
      x[i] = tmp - δ / 2
      y1 = f(x)
      x[i] = tmp + δ / 2
      y2 = f(x)
      x[i] = tmp
      Δ[i] = (y2 - y1) / δ
    end
    return Δ
end
  
for seqlen in (1, 2, 3, 4, 5)
    x = randn(Float64, 10, 1, seqlen)
    y = mean(x, dims = 1)
    g1 = Flux.gradient(θ0) do θ
        loss_sciml(θ, x, y)
    end |> first
    g2 = ngradient(θ0, 1000) do θ
        loss_sciml(θ, x, y)
    end
    @show seqlen, maximum(abs, g2), maximum(abs, g1[1:1000] .- g2[1:1000])
end

(seqlen, maximum(abs, g2), maximum(abs, g1[1:1000] .- g2[1:1000])) = (1, 0.0004209382750559598, 4.859688407167929e-11)
(seqlen, maximum(abs, g2), maximum(abs, g1[1:1000] .- g2[1:1000])) = (2, 0.01700042188167572, 0.005579658870112846)
(seqlen, maximum(abs, g2), maximum(abs, g1[1:1000] .- g2[1:1000])) = (3, 0.023597870022058487, 0.024526089218849207)
(seqlen, maximum(abs, g2), maximum(abs, g1[1:1000] .- g2[1:1000])) = (4, 0.02711454499512911, 0.019288377511780964)
(seqlen, maximum(abs, g2), maximum(abs, g1[1:1000] .- g2[1:1000])) = (5, 0.025720552541315556, 0.031243588046897124)

Two steps of reccurency should not have vanishing gradient bu ut the result shows once seqlen > 1, the gradient is no longer accurate.

[AStupidBear]

@MikeInnes @dhairyagandhi96

[AStupidBear]

Maybe this is related to #876

[AStupidBear]

I should avoid using broadcasting and write in a old loop fashion. The following script has no problem.

using Flux, Random, Statistics

Random.seed!(1234)

x = randn(Float32, 10, 1, 100)
y = mean(x, dims = 1)

model = Chain(LSTM(10, 100), LSTM(100, 1))

function loss(x, y)
    xs = Flux.unstack(x, 3)
    ys = Flux.unstack(y, 3)
    l = 0f0
    for t in 1:length(xs)
        ŷ = model(xs[t])
        l += Flux.mse(ys[t], ŷ)
    end
    return l / length(xs)
end

opt = ADAM(1e-3, (0.9, 0.999))

cb = () -> (println(Flux.data(loss(x, y))); Flux.reset!(model))

Flux.@epochs 10 Flux.train!(loss, params(model), repeat([(x, y)], 100), opt, cb = cb)

[AStupidBear]

Duplicate of #1209