Issue with backprop in a custom recurrent `Lux` layer

pietro-butti
Tue, 13 May 2025 14:52:09 +0000
Machine Learning

I am trying to implement an autoencoder in Lux, inspired by the excellent one that I found in githun.com/lkulowski/LSTM_encoder_decoder, coded with PyTorch.

(diagram taken from the link)

Model overview

The encoder is a simple recursive LSTM (I used RNN for simplicity), that can be easily implemented with

encoder = Recurrence(RNN(1=>15)) # hidden_dim = 15

The decoder is slightly more complicated:

input: the last hidden state and the last time step (\vec h_\tau,x^{(\tau)})

recurrent structure: Generates outputs \hat y by passing each hidden state through a dense layer (not depicted in the original diagram). Hidden states are built from previous predictions:

\hat y^{(1)} = \text{Dense}(\text{RNN}((\vec h_\tau,x^{(\tau)}))) \\ \hat y^{(2)} = \text{Dense}(\text{RNN}((\vec h_{\tau+1},\hat y^{(1)}))) \\ \cdots

(Note: the input at each step is the previous prediction, after the dense pass).

Custom decoder layer implementation

I am trying to implement this in Lux, and I could not find any custom way to do it, therefore I am trying to implement a custom layer.

struct Decoder{R,T} <: LuxCore.AbstractLuxContainerLayer{(:cell, :dense)}
    cell::R
    dense::T
end

# Initializer
function Decoder(in_dim, hidden_dim, out_dim)
    return Decoder(RNNCell(in_dim=>hidden_dim), Dense(hidden_dim=>out_dim))
end

To implement the forward pass, a key challenge is efficiently accumulating the subsequent outputs \hat y^{(1)},\hat y^{(2)},\hat y^{(3)},… during the decoding phase.

Since Zygote.jl does not support array mutation, it is not possible to initialize an empty vector ([]) and append outputs within the loop, nor can we preallocate an array and fill it progressively. This limitation complicates the design of a stateful decoder, as shown in my current approach:

# Forward pass
function (net::Decoder)(ydata::AbstractArray{T,3}, (h_last,x_last), ps::NamedTuple, st::NamedTuple) where T
    function f_recur((output, (hₜ,ŷₜ), s), x_t)
        (_,(hₜ₊₁,)), st_cell = net.cell((ŷₜ,(hₜ,)), ps.cell, s.cell)
        ŷₜ₊₁, st_dense       = net.dense(hₜ₊₁    , ps.dense, s.dense)

        # Stack output
        out = (output...,ŷₜ₊₁)

        # Update carry
        carry = (hₜ₊₁,ŷₜ₊₁)

        return out, carry, (cell=st_cell, dense=st_dense)
    end

    # We do recursion by hand because foldl gives me an error caused by subarray types
    buffer = ((), (h_last,x_last), st)
    for x_t in Lux.eachslice(ydata,dims=2)
        buffer = f_recur(buffer, x_t)
    end

    # unpacking for clarity
    (out, memory, state) = buffer

    # stacking results to compute loss
    out = hcat([reshape(y,size(y,1),1,size(y,2)) for y in out]...)

    return out, state
end

Problem: Gradient Issues with Zygote

However, this approach fails with Zygote when calling the gradient function.

For a minimal example of usage

rnx = Xoshiro(1994)

model = Decoder(1,15,1)
pars,state = Lux.setup(rng,model)

h_last = randn(rng, Float32, 15, 2)
x_last = randn(rng, Float32, 1, 2)
encoded_state = (h_last,x_last)

y_test = randn(rng, Float32, 1, 5, 2)

The error can be reproduced by running

train_state = Training.TrainState(model, pars, state, Adam(3.0f-4))
const lossMSE = MSELoss()
generic_loss(_model, _ps, _st, (y,enc_state)) = begin
    ŷ,st = _model(y,enc_state, _ps, _st)
    return lossMSE(ŷ,y), st, 0
end
Training.compute_gradients(
    AutoZygote(), generic_loss, (y_test,encoded_state), train_state
)

or, directly with the pullback

loss(p,s) = begin
    ŷ,st = model(y_test,encoded_state, p, s)
    lossMSE(ŷ,y_test), st
end

(loss_eval, state, ), loss_back = pullback(loss,pars,state)
grad, _ = loss_back((1.,nothing))

The same issues occur when using another AD backend, like AutoEnzyme().

Question

How can I efficiently accumulate outputs in the forward pass without breaking the differentiability with Zygote? Are there better design patterns for this type of recurrent decoding in Lux?

7 posts - 4 participants