RETURNN layers with hidden state should make it explicit

[albertz]

As it was discussed in #16, RETURNN layers with (hidden) state (e.g. RecLayer with unit="lstm") should make the state explicit in the API. E.g. the Rec module should get two arguments input and prev_state and return output and state. So the usage would look like this in a loop:

lstm = Lstm(...)
with Loop() as loop:
  ...
  out, loop.state.lstm = lstm(x, loop.state.lstm)

Or like this outside a loop (using default initial state, ignoring last state):

lstm = Lstm(...)
out, _ = lstm(x)

This applies for all RETURNN layers with rec hidden state, and further modules like Lstm.
See RETURNN layers with rec hidden state.

Relevant modules here:

• _Rec based, e.g. Lstm (only one so far)
• window
• cumsum
• ken_lm_state
• edit_distance_table
• unmask
• _TwoDLSTM
• cum_concat

[albertz]

Implemented partly now. All recurrent layers do return a tuple (output, state), and accept a state (initial or previous state) argument.

The returned state is via GetLastHiddenStateLayer. This is incomplete:

• It needs to have a specific n_out.
• For cases where the state is more structured (e.g. LSTM with (h,c)), this should return the same structure.

This might need to have special implementations for each case. Not sure.

[albertz]

Another issue: Some layers like WindowLayer often would not use the state and then become unnatural to use. I'm not really sure about the best way to solve this.

Maybe we can make two versions of the module, one which is not recurrent and one which is recurrent.

Also unclear whether this should be automatically generated, or whether we explicitly handle these cases.

[albertz]

We should be a bit more specific, and list the layers or modules which need custom hidden state. E.g. Lstm, or Rec in general. Or basically every rec layer, but we should specify the hidden state.

We also should list the layers where we might consider non-recurrent variants.

It should be collected by editing the main post.

[albertz]

It's not totally obvious.

CumsumLayer and CumConcatLayer have state = output, so one single return is enough, not a tuple.
Same for RecLayer with many units, such as GRU, RNN, etc.

RecLayer with LSTM will return (h, c), where h = output.

state = output or h = output, that is the case only inside the loop.
Outside the loop, state or h is just the last state, where output is the whole seq. But the whole seq includes h or state then.

The initial_state argument however needs to be consistent with the returned state.

So, you could argue, RecLayer inside the loop could just return (output, c) instead of (output, (h, c)), but the initial_state argument is (h, c) anyway.
Note that in PyTorch, LSTM (operating on the seq) returns (output, (h, c)).
In PyTorch, LSTMCell (operating on the frame) just returns (h, c). Although its prev state input argument is (h, c).
So it's a bit inconsistent, as you cannot generalize for any PyTorch module with state, what you need to keep from the returned output to feed back as prev state input in the next frame. There is no generic API as we intent here.

[albertz]

I introduced the explicit LayerState (which derives from dict currently, although that might maybe change), the code can explicitly test and see whether some layer returned a state.

By convention, I suggest that the return in such case should be a tuple, and the last item should always be this LayerState instance.

Note that LayerState also avoids that the user accidentally does automatic concat, such as in this example:

x = rnn(...)
y = linear(x)

When x is a tuple (out, hidden), where hidden is just a layer-ref, this might do an automatic concat (#40, #41).
When hidden is a LayerState, such automatic concat cannot happen here.

[albertz]

Passing state to a RETURNN layer is a bit tricky. It is not really the same as initial_state. initial_state defines the initial state only (both when it is optimized out or inside the loop), not the state per frame. See: rwth-i6/returnn#732

So we need a new way on RETURNN side to allow passing the state per frame on such layers, via a new state option (rwth-i6/returnn#732).

But we also want to be able to make use of RETURNN rec layer automatic optimization, which is tricky as well. We need to figure out if the layer uses its own prev state as input state. If that is the case, it would set initial_state on the layer, pointing to the initial argument from the State object. Otherwise it would use the new state layer option.

However, at the time we call the module (layer maker), it can not know this. Only when it returns and when State.assign happens, it can figure out that the state argument passed to the module was the same.

So this means we need to delay the layer dict creation, or post edit it to apply the optimization and remove the state argument in this case, and replace it with initial_state.

This is not only for the Lstm module but any hierarchical module where state might be nested, which makes this more complicated because it might not be just a single layer dict.

We can recursively go through a State.assign value and check each referred layers whether they are Layer with layer dict with class: get_last_hidden_state and then resolve the name and find the layer and its layer dict and then check the state argument whether it points to the same layer or layers.