Rec design for recurrent definitions / loops

[albertz]

This issue is to collect some thoughts on the recurrent loops design, which wraps the RecLayer with an explicit subnetwork in RETURNN.

The main goal is to have this very straight-forward and simple for the user. We can abstract away from the underlying RecLayer if that makes things easier. We can also extend RETURNN itself if needed.

Related is also #6 (rec prev mechanism), and this issue here might fix/resolve #6, although not necessarily.

This also needs some mechanism for unrolling/unstacking, i.e. when we iterate over input x with some time-axis, i.e. to get x[t]. This is rwth-i6/returnn#552.

---

To define a loop like this pseudo Python code:

x  # given, shape {batch, time, dim}
h = Zeros({batch,dim})()  # initial state, shape {batch,dim}
out = []
for t in range(x.max_seq_len):
  x_lin = Linear(dim)(x[t])
  h_prev = h
  h = Linear(dim)(x_lin + h_prev)
  out.append(h)

h  # final state
out  # shape {time, batch, dim}

Current design:

There is Loop() which can be used in a with context, which corresponds to the for-loop in the example, or in general to a while-loop. Like:

with Loop() as loop:
  ...

There is State() which can define hidden state (for any module or any code).

The example above can be written as:

h = State({batch, dim}, initial=0)
with Loop() as loop:  # this introduces a new loop
  x_t = loop.unstack(x)  # shape {batch, dim}

  x_lin = Linear(dim)(x_t)
  h_prev = h.get()
  h_ = Linear(dim)(x_lin + h_prev)  # shape {batch, dim}
  h.assign(h_)

  out = loop.stack(h_)  # shape {time,batch,dim}
  h_last = loop.last(h_)

# h.get() would now return the last state
# h_last is an alternative

Or with a module as:

class MyRec(Module):
  def __init__(self):
    super().__init__()
    self.x_linear = Linear(dim)
    self.h_linear = Linear(dim)
    self.h = State({batch, dim}, initial=0)

  def forward(self, x):
    # x shape is {batch, dim}
    x_lin = self.x_linear(x)
    h_prev = self.h.get()
    h = self.h_linear(x_lin + h_prev)  # shape {batch, dim}
    self.h.assign(h)
    return h

rec = MyRec()
with Loop() as loop:  # this introduces a new loop
  x_t = loop.unstack(x)  # shape {batch, dim}
  h_ = rec(x_t)  # shape {batch,dim}. this represents the inner value
  h = loop.last(h_)  # shape {batch,dim}
  out = loop.stack(h_)  # shape {time,batch,dim}

For the TF name scopes (and variable scopes), we should follow #25, i.e. make it exactly as the module hierarchy.

The RETURNN layer name of the created RecLayer via Loop does not matter too much. It could be arbitrary, or some clever (but simple) logic to use the first module name or so. The RETURNN layer hierarchy can be independent from the actual TF name scopes (via #25).

Special options for the RecLayer like include_eos can be options for Loop, like Loop(include_eos=True). Or as a method, like loop.set_include_eos(True).

Loop (potential) methods:

• unstack.
  We need RecLayer multiple inputs via explicit unstacking returnn#552 for this.
  unstack also implicitly implies that the loop runs over the time-axis of x.
• last
• stack
• idx: to return some layer which wraps RETURNN ':i'

State has methods get and assign. (... See discussion below for more ...)

Current reasonings:

Why no special base class Rec which derives from Module? We want to easily allow to use any kind of module inside a loop. We think the current API makes this more straight-forward.

Why is h not an argument of forward, and why State instead? This allows to call other sub modules, which might define their own hidden state. So the root recurrent module does not need to know about all the hidden states of sub modules.

Why to have the hidden state explicit, and not use sth more close to self.prev? To make the behavior more straight-forward.

The current design allows for nested loops and sub modules with hidden state.
Only the Loop() call actually introduces a new loop.

class MySubRec(Module):
  def __init__(self):
    super().__init__()
    self.h = State({batch,dim})

  def forward(self, a):
    # assume a shape {batch,dim}
    h = self.h.get() + a
    self.h.assign(h)
    return h

class MyRec(Module):
  def __init__(self):
    super().__init__()
    self.sub = MySubRec()
    self.h = State({batch,dim})

  def forward(self, x):
    a = self.h.get() + x

    # example with sub as nested loop
    with Loop() as loop:
      y = self.sub(a)
      y = loop.last(y)

    # or: example with sub in same loop
    y = self.sub(a)
    
    self.h.assign(y)
    return y

There should not be any special handling needed for the Choice layer.
Note that the search flag and train flag logic is a separate thing (#18).

There should not be any special handling needed whether the input to a rec module call would be inside the current/same loop or not. unstack on some value which is already inside the loop would not make sense, though, and should result in an error. But this would all be covered by RETURNN logic already.

RETURNN rec automatic optimization should not cause any problems. RETURNN already should guarantee that it is equivalent. From the user view point, it never ever should matter whether it is optimized. Otherwise this is rwth-i6/returnn#573. On this returnn-common level, it should not matter.

---

Example for LSTM for a single step:

class Lstm(Module):
  def __init__(self):
    super().__init__()
    self.h = State({batch,dim})
    self.c = State({batch,dim})
    self.ff_linear = Linear(dim * 4)
    self.rec_linear = Linear(dim * 4)

  def forward(self, x):
    # x shape is {batch,dim} (single frame)
    x_ = self.ff_linear(x)
    h_ = self.rec_linear(self.h.get())
    x_in, g_in, g_forget, g_out = split(x_ + h_, 4)
    c = self.c.get() * sigmoid(g_forget) + tanh(x_in) * sigmoid(g_in)
    self.c.assign(c)
    h = tanh(c) * sigmoid(g_out)
    self.h.assign(h)
    return h

[albertz]

The example should be extended:

How are nested loops handled? This should be straight-forward as well. And nothing must be ambiguous about it, e.g. the (implicit or explicit) definition and handling of hidden state.

[albertz]

What are reasonable sub modules with hidden state? How would you define them? Just normal Modules? Or other Recs where you call step?

[albertz]

Custom ending condition, how to define?

[albertz]

How to get the last hidden state?

[albertz]

Also related: rwth-i6/returnn#391 rwth-i6/returnn#545

I.e. the question about how to define masked self-attention in a generic/flexible way in this Rec concept. E.g. HistoryMaskLayer or so. This could all be wrapped directly but is this intuitive then? How does it look like? Can we make it more intuitive?

[albertz]

Also: Is there anything special needed for Choice (ChoiceLayer)? And beam search? Maybe not and it is all straight-forward already.

How to define whether search is enabled? Just not do any special handling on this level at all and leave it to the current RETURNN behavior? Or make it more explicit?

[Atticus1806]

How to define whether search is enabled? Just not do any special handling on this level at all and leave it to the current RETURNN behavior? Or make it more explicit?

From what I can see RETURNN handling should be enough. Assuming using this with automation changes in config can be done elsewhere I think.

Custom ending condition, how to define?

Maybe give an optional parameter for Rec like Rec.custom_break which will be checked internally at the beginning of the loop?

self.h = self.state({batch, dim})

How do you concretely think of this init? This requires an explicit value for the batch dimension right? I am not sure if this is what we should aim for, RETURNN should be able to handle this already? Or is this right now just for example sake?

with rec.loop() as loop:  # this introduces a new loop
 h_ = rec(x)  # shape {batch,dim}. this represents the inner value
 h = loop.last(h_)  # shape {batch,dim}
 out = loop.stack(h_)  # shape {time,batch,dim}

Maybe this example (esp. for the doc.) could be extended with another (non-rec) module, showing how to use it in a Module, so that someone who did not work with Rec in this way yet directly has an example which can be modified.
Maybe something like (assuming this is how its intended):

def EncDec(Module):
  def __init__():
    self.enc = Linear(dim)
    self.dec = MyRec()
  
  def forward(x):
    x = self.enc(x)
    with self.dec.loop():
        h_ = rec(x)  # shape {batch,dim}. this represents the inner value
        h = loop.last(h_)  # shape {batch,dim}
        out = loop.stack(h_)  # shape {time,batch,dim}
    return out

Also: Is there anything special needed for Choice (ChoiceLayer)? And beam search? Maybe not and it is all straight-forward already.

We are mainly aiming for conversion to RETURNN here, so maybe if we choose a smart way to include them in the loop and then some internal handling this can be converted (since in returnn they are part of the rec unit anyways).

[albertz]

How to define whether search is enabled? Just not do any special handling on this level at all and leave it to the current RETURNN behavior? Or make it more explicit?

From what I can see RETURNN handling should be enough. Assuming using this with automation changes in config can be done elsewhere I think.

Note that this is not just about what is enough to be able to define all networks. Or not sure how you mean it.

It's about being straight forward and clear. I.e. at no time, it should be unclear to the user when search is used.

We do not have to follow exactly the behavior of RETURNN. There are also multiple ways in RETURNN. We can restrict it to one clean way. We can also change it. Or introduce a simpler variant here.

I'm tending to make it explicit. But not sure.

PyTorch also has a similar concept for the train flag (as we do have as well in RETURNN). I.e. some PyTorch modules behave differently depending if they are in train or eval mode (e.g. Dropout). We have exactly the same in RETURNN. And search is a flag like train.

The difference is how these flags are set:

• In RETURNN, this is all globally, and for search flag, there are some additional (maybe unintuitive) ways to overwrite it. And the flags are implied automatically in RETURNN, depending e.g. on the task, and the user has not much control over it. It is quite hidden.

• In PyTorch, there are no implicit automatic implied global flags. Every module has its own flag, and it is set explicitly (and easily recursively for all sub modules). Every module has always the train flag set initially, and you can disable it explicitly. So to the user, it's always clear how the flags are set, because the user sets them, and no automatic behavior. The user explicitly writes model.train() or model.eval().

Maybe again, here in returnn-common, we can follow the PyTorch style a bit for this, and also copy it for the search flag? Not sure...

Edit: I think this is actually a separate thing, which we can discuss/handle independently. So I opened #18 about this. And this can be marked as off-topic here.

[albertz]

Custom ending condition, how to define?

Maybe give an optional parameter for Rec like Rec.custom_break which will be checked internally at the beginning of the loop?

In RETURNN, the layer is called end. So maybe additionally to step, the user can define a method end (kind of similar as tf.while_loop gets body and cond).

This end method can be optional. E.g. when unstack is used at some point, then the sequence length is implied from this. The same is true when cross-entropy loss is used, as unstack would be used on the targets.

[albertz]

self.h = self.state({batch, dim})

How do you concretely think of this init? This requires an explicit value for the batch dimension right? I am not sure if this is what we should aim for, RETURNN should be able to handle this already? Or is this right now just for example sake?

batch here would be the dimension tag (#17). I.e. this is kind of a placeholder. This does not require an explicit value. And yes, RETURNN has all that logic already.

[albertz]

with rec.loop() as loop:  # this introduces a new loop
 h_ = rec(x)  # shape {batch,dim}. this represents the inner value
 h = loop.last(h_)  # shape {batch,dim}
 out = loop.stack(h_)  # shape {time,batch,dim}

Maybe this example (esp. for the doc.) could be extended with another (non-rec) module, showing how to use it in a Module, so that someone who did not work with Rec in this way yet directly has an example which can be modified.
Maybe something like (assuming this is how its intended):

def EncDec(Module):
  def __init__():
    self.enc = Linear(dim)
    self.dec = MyRec()
  
  def forward(x):
    x = self.enc(x)
    with self.dec.loop():
        h_ = rec(x)  # shape {batch,dim}. this represents the inner value
        h = loop.last(h_)  # shape {batch,dim}
        out = loop.stack(h_)  # shape {time,batch,dim}
    return out

Yes sure. Although I thought this part was clear.

[albertz]

Also: Is there anything special needed for Choice (ChoiceLayer)? And beam search? Maybe not and it is all straight-forward already.

We are mainly aiming for conversion to RETURNN here, so maybe if we choose a smart way to include them in the loop and then some internal handling this can be converted (since in returnn they are part of the rec unit anyways).

No, simplicity, clear and straight-forward behavior has priority over everything. In the end, everything can be converted to RETURNN, because RETURNN should be generic to allow for everything (and if that is not the case, we should simply extend RETURNN itself). Esp if the RETURNN behavior is unintuitive or not straight-forward, we should not just copy that here.

I'm not sure what you mean by "smart way", and "internal handling" for what?

Right now, you would just use Choice as you would use other layers like Linear. I was wondering whether we need any special behavior for Choice here in returnn-common, or not. Currently I don't really see a reason for special behavior. But maybe I'm overlooking sth.

This is of course related to the search aspect. Although, even when we can set the search-flag now explicitly in returnn-common, the logic how the ChoiceLayer behaves with search-flag or without, that is already all part of RETURNN, so no special handling needed for that here.

[albertz]

Note: I changed the original suggestion of automatically stacking the output, and automatically implying the construction of the loop, because it was not clear for sub modules with hidden state when they should create a new nested loop or reuse the same loop. So now this is explicit, via loop().

However, I'm not sure if this is now the most intuitiv, clean, straight-forward way. Although other ways I can think of right now are also not really better.

When we keep using Module everywhere and not use Rec as a base class, and Rec.State() and Rec.Loop() explicitly, this is would work already in this draft.
(We can still use Rec as a namespace to cover all rec related logic. Or maybe another Python module/package for this. Or just use RecLoop() instead of Rec.Loop(). But this is not relevant for now.)
It's just a technical question then about how to handle this code:

# Definition for a single step:
class MyRec(Module):
  ...

# Now the loop:
rec = MyRec()
with Rec.Loop() as loop:
  x_ = loop.unstack(x)
  y_ = rec(x_)
  y = loop.stack(y_)

Or, alternatively to Rec.Loop(), maybe Rec stays a Module where you define the inner step, which would just exactly contain that code:

# Definition for a single step:
class MyRec(Module):
  ...

# The loop:
class MyRecLoop(Rec):
  def __init__(self):
    self.rec = MyRec()

  def step(self, x):
    x_ = self.unstack(x)
    y_ = self.rec(x_)
    return self.stack(y_)

loop = MyRecLoop()
y = loop(x)

[Atticus1806]

Currently I don't really see a reason for special behavior. But maybe I'm overlooking sth.

I agree with that, right now it does not seem like we need it, but again I can report back once I finished my config wether there is some weird behavior we didnt think of.

I'm not sure what you mean by "smart way", and "internal handling" for what?

That was weird phrasing, what I meant pretty much is that I would have all of that Choice and beam search handling in ´returnn_common´ and not leave it to the user to find the "correct" way to use the layers, best case it should just be used as you said similar to a Linear layer imo.