Repository: rwth-i6/returnn

Issue:
Extending Self Attention
I want to implement some changes to the self-attention used in the Transformer for MT, namely implement locality-sensitive hashing (https://arxiv.org/pdf/2001.04451.pdf).

Right now, self-attention is a single layer within RETURNN. While this is very convenient when using the default configuration, it is not very extensible: All options for it have been implemented as additional arguments for the layer and the code for it has become pretty messy over the time.
I could implement my changes within the layer by adding an additional parameter, but I think it might be better to not clutter the self-attention layer with even more (relatively specific) parameters.

Instead it might be nicer to implement them using existing other RETURNN Layers, similar to how encoder-decoder attention is implemented in our Trafo config.
For unmasked self-attention (where one can attend to the entire sequence, e.g. used in the encoder), I don't see an issue in implementing it completely analogous to the encoder-decoder attention:
Use three `linear` layers to obtain queries, keys and values. Compare all queries and keys against each other using a `dot` layer, and then use a `softmax_over_spatial` layer to turn these attention energies into attention weights. Finally use a `generic_attention` layer to compute a weighted sum of the attention values.

For masked self-attention (where one cannot attend to future positions, e.g. used in the decoder), there are two things to consider:
- We have to mask all future target positions by setting the attention energies to -\infty. This could for example be done in the softmax layer (which already considers the total sequence length anyway)
- When we are in a recurrent layer (e.g. during Trafo inference), then we would like to cache all previously computed attention keys, values and queries. The current self-attention layer does this, but it is also one of the reasons why it is messy to extend it currently: Both the recurrent and parallel case are handled somewhat differently.
I have no idea how that should look. What would the linear layers generating attention keys and values return in a recurrent loop (wouldn't that introduce a time axis even then)? How to handle that we do not need to recompute old keys/values?

What would be the best approach to extend self-attention? Stick to changing the Returnn code of the layer? Or implement it in multiple layers, but then how do I solve the problems I mentioned?
Thanks :)


Discussion:
It would definitely be better to implement this using the RETURNN layers instead of extending the `SelfAttentionLayer`. 

From my perspective, the problems you are mentioning might not be solvable without extending some of the other layers. If you want to access computed layer outputs from previous recurrent steps, you can use the [`WindowLayer`](https://returnn.readthedocs.io/en/latest/layer_reference/basic.html#window-layer) (unfortunately the docstring is not yet covering this special case). With the window layer, you could get e.g. the last N keys and values in the format `[B,N,D]` inside the recurrent subnetwork, and they would not be recomputed. The problem is, that you need all the states from the beginning, and the `WindowLayer` is currently not capable of doing so.

I am not sure if there is another approach currently existing. Extending the `WindowLayer` to have a "dynamic" window might not be the best approach, but it is the only one I can think of so far.
> Right now, self-attention is a single layer within RETURNN. While this is very convenient when using the default configuration, it is not very extensible: All options for it have been implemented as additional arguments for the layer and the code for it has become pretty messy over the time.

Thanks for mentioning this. This is what I keep saying, and also explain in the tutorial, and also in the introduction video, and in many other places.

This is why we definitely should not make it even more complex. I would even go as far as saying we should deprecate the `SelfAttentionLayer`, as it obviously tends to become messy, or already is, and this is one of the core principles of RETURNN, that layers should be simple.

This is what you should follow in principle. (If you have not seen the tutorial, or introduction video, please check it, to know the principles of RETURNN.) Layers should be as simple as possible, reduced mostly to a single operation (e.g. wrapping a single TF low-level function), such that you can basically define anything in your config by using the layers as building blocks. Also, such basic building blocks then should be useful for most other people as well, as they should be very generic.

This should mostly be straight forward. Except for automatic optimization, i.e. that you should think about the two cases, where the layer is inside a recurrent loop (at decoding time, or maybe if you have other recurrence in the model), or when it can be calculated in parallel (independent of the recurrent loop). I haven't really thought about how to split up `SelfAttentionLayer` into smaller building blocks such that it is still efficient for these two cases. This is the non-trivial part.

You could also implement another independent layer, `SelfLSHAttentionLayer` or so. But this might be very specific, so maybe not useful for other people (except if they want to use exactly the same LSH variant). But you could simply leave such layer inside your config and not push it back to RETURNN.

One principle rule of code in RETURNN is that this should be useful for many people, and would probably immediately use by many people, now and in the future. For generic basic building blocks, this is usually always the case. For sth like `SelfLSHAttentionLayer`, this is not clear currently.

So, to summarize:

* Try to think about basic building blocks you need, and whether that works correct with automatic optimization. Then add whatever basic building block is missing to RETURNN.
* Or: Put your custom `SelfLSHAttentionLayer` into your config. But this is maybe not so nice.

It seems that it is a better idea to write a new layer instead of hacking something into the `Window`-layer. @albertz suggested adding a `ConcatenateInTimeLayer` which returns the succeeding outputs in a recurrent layer.
Thanks for all of your comments!
I will start by implementing the 'normal' self attention as currently provided by the `SelfAttentionLayer` using multiple RETURNN layers, and then start from there - I do not seem to be the only person here interested in making self attention more configurable.

> It seems that it is a better idea to write a new layer instead of hacking something into the `Window`-layer. @albertz suggested adding a `ConcatenateInTimeLayer` which returns the succeeding outputs in a recurrent layer.

That's a good idea - that's the main part that's missing.
However, I would call the layer differently: `ConcatenateInTimeLayer` sounds very similar to `Prefix/PostfixInTimeLayer` to me (suggesting it concatenates multiple inputs along the time axis).
Perhaps something like `AccumulatePreviousOutputLayer`? That would be similar to the (obsolete) `GetRecAccumulatedOutputLayer`. But I'm not sure ..
Any thoughts? But that's just a name, I can start working on it nonetheless.

The idea of that layer would be then:
- If the layer is not in a recurrent loop, it just copies its input (and maybe also asserts that an recurrent time axis already exists).
- If it is in a recurrent loop, it adds a recurrent time axis which concatenates it's input at this time step and all previous time steps.
> However, I would call the layer differently: `ConcatenateInTimeLayer` sounds very similar to `Prefix/PostfixInTimeLayer` to me (suggesting it concatenates multiple inputs along the time axis).
> Perhaps something like `AccumulatePreviousOutputLayer`? That would be similar to the (obsolete) `GetRecAccumulatedOutputLayer`. But I'm not sure ..
> Any thoughts? But that's just a name, I can start working on it nonetheless.

Right, `ConcatenateInTimeLayer` is not so clear.

There is a very similar/related `CumsumLayer`, which does `tf.cumsum`.
Analog to that, maybe `CumConcatLayer`?
 
> * If it is in a recurrent loop, it adds a recurrent time axis which concatenates it's input at this time step and all previous time steps.

Yes. So let's be more specific. In the loop, it assumes an input of shape [B,D], and the output in first frame would be [B,1,D], in second frame would be [B,2,D], and so on, i.e. in general [B,i,D]. This assumes an initial state of [B,0,D], but that could be configured.

It also means when you know accumulate this frame-by-frame (for i=1...T), i.e. get the accumulated values from outside of the loop, you would get an output of shape [T,B,T,D], which has the usual mask on the second T axis.

> * If the layer is not in a recurrent loop, it just copies its input (and maybe also asserts that an recurrent time axis already exists).

This is not so clear. To be analog to the within-loop case, it should expect an input [T,B,D] (or whatever axis order), and then just produce the same output [T,B,T,D], with the usual mask on the second T axis.

To add on the outside-loop case:

However, this is unnecessarily inefficient. You are right, if outside the loop, it could also just return [T,B,D]. But this has two problems:

* This would be inconsistent to the within-loop operation, where the accumulated shape would be [T,B,T,D]. This is kind of the requirement: Whether automatic optimization is enabled or disabled should result in the same values. (Or at least the behavior in all possible cases should be the same.)

* Follow-up operations would need to be adjusted, esp. the energy mask need to be applied later on. But how should the follow-up layer know about this, if it just gets some data [T,B,D], that some mask needs to be applied?

`CumConcatLayer` or `AccumulatePreviousOutputLayer` are also then bad names, if it has this special behavior. Or maybe the behavior should be a flag?

And there need to be a custom layer which applies the mask only if outside the loop. But I wonder if this is really intuitive. How would this be called? `HistoryMaskLayer`? Maybe the newly introduced dimension, the second T in the case within-loop, its dimension tag can have a special flag `accumulated_over_time` or `historic` or so, and the `HistoryMaskLayer` only applies the mask if there is no dimension with such a flag. But I fear this becomes kind of non-intuitive.

We should draft maybe a net dict, how that would look like.

Analog to the `SelfAttentionLayer`, let's assume there is a layer `x` with the input:
```
{class: rec, unit: {
...,
x: ...,  # [B,D] inside, [T,B,D] outside
qkv: {class: linear, from: x, activation: None},  # [B,2*K+V] inside, [T,B,2*K+V] outside
qkv_split: {class: split, from: qkv, size_splits: [K,K,V]},
q: {class: copy, from: qkv_split/1},  # [B,K] inside, [T,B,K] outside
k: {class: copy, from: qkv_split/2},  # [B,K] inside, [T,B,K] outside
v: {class: copy, from: qkv_split/3},  # [B,V] inside, [T,B,V] outside
k_accum: {class: cum_concat, from: k, as_is_outside: True},  # [B,T,K] in both cases
v_accum: {class: cum_concat, from: v, as_is_outside: True},  # [B,T,V] in both cases
energy: {class: dot, from: [q, k_accum], red1:"static:-1", red2:"static:-1", var1:"T?", var2:"T"},  # [B,T] inside, [T,B,T] outside
energy_masked: {class: history_mask, from: energy, value:-inf},  # [B,T] inside, [T,B,T] outside (with mask applied)
att_weights: {class: softmax_over_spatial, from: energy_masked},  # [B,T] inside, [T,B,T] outside
att: {class: dot, from: [v_accum, att_weights], red1:"T", red2: "stag:history?", var1:"static:-1", var2:[]},  # [B,V] inside, [T,B,V] outside
}}
```

This is only a draft. There are a couple of details which probably need extra care. It misses:

* Scale on query (simple).
* Multiple heads (relatively simple).

Incomplete also: In the last dot-layer for `att`, I use `stag:history?`. This would be analog to `T?`, i.e. it means it optionally selects this axis if it exists, and otherwise not.

We cannot use `T?` there because this would probably select the wrong spatial/time axis! Optional `stag` support is missing, but should be simple to add. `history` is just a made-up name here. It depends how the `CumConcatLayer` calls the new dimension tag.

Or maybe instead of using `stag`, we could also have a special axis identifier `history`, which would also be used by `HistoryMaskLayer`?

I think there is sth wrong in the `att` at the end. It should not be `red2: "stag:history?"`. `red2` should always match to `red1`, so it cannot be optional. So it would just be `red2: "stag:history"`. To make that work, the `HistoryMaskLayer` actually could also introduce such history-dim-tag on the masked time axis.

Another idea: `CumConcatLayer`, in the outside-loop-case, could return shape [T,B,T',D], where T' is a special dimension of length 1, but which has the special meaning that it would expand to the other time-dim with masking. Such special meaning could be preserved in its dimension tag. Then it becomes somewhat more consistent to the inside-loop-case, and also `HistoryMaskLayer` becomes more straight-forward.

The shapes would become:
```
v_accum: [T',B,D] inside, [T,B,T',D] outside
```

However, the `DotLayer`s become less straight-forward. I think for `energy`, it just stays:
```
energy: {class: dot, from: [q, k_accum], red1:"static:-1", red2:"static:-1", var1:"T?", var2:"T"},  # [B,T'] inside, [T,B,T] outside
```
This would squeeze away the T' axis in case of outside-the-loop, as it would be part of the remaining axes and it assumes them to match. Not sure if it does it all correctly, or complains. Also not sure if it should do this implicitly like that, or we should make it more explicit.

Then for the `HistoryMaskLayer`, we would get:
```
energy_masked: [B,T'] inside, [T,B,T''] outside
```
T'' would be like T, but a separate dim-tag.

Also the final `DotLayer` for `att` becomes less straight-forward. Specifically, what should you put for `red2`? You want to reduce the axis if it can be matched to `T`, i.e. inside the loop, but not reduce it otherwise. You just want to squeeze it outside the loop, but not sure if that should be done in dot layer or separately. Maybe just like:
```
att: {class: dot, from: [v_accum, att_weights], red1:"stag:history", red2: "stag:history", var1:"static:-1", var2:[]},  # [B,V] inside, [T,B,V] outside
```
That does not work. The axes to be reduced must match, but they don't.
Inside the loop, red1 should be T', outside the loop, it should be T.

Maybe there is some way to overcome this. Or maybe this idea was just not good, and the first approach is better.

Right now, I think the first proposal (just return `[T,B,D]` in the outside loop case) is more intuitive and more straight forward to implement and use.
We can also just implement this behavior as a flag, and then have the benefits of both solutions.

As you said, it makes everything simpler if the `CumConcatLayer` marked the outputted time-axis in some special way, e.g. using `stag:history` as you mentioned.
Then we can use `stag` everywhere instead of just `T`. This is a bit verbose, but I think necessary to keep apart the query-axis (i.e. `stag:extern_data_classes` and the key/value-axis (i.e. `stag:history` or sth like that).

I am unsure whether a `HistoryMaskLayer` is really necessary.
We could also add an extra parameter `mask_future_of_axis` to the `SoftmaxOverSpatialLayer`.
Then we could simply say
```
att_weights: {class: softmax_over_spatial, from: energy,
              axis: stag:history, mask_future_of_axis: stag:extern_data:classes?}
```
and let the `softmax_over_spatial` layer apply masking exactly if it receives a `mask_future_of_axis`, which is exactly the case when we are outside a loop.
But then again this is less configurable ... so maybe this is a bad idea falling back into the loop-hole of having too many arguments to RETURNN layers.

Or instead, `HistoryMaskLayer` could receive two time-axes as argument: e.g. `axis` and `mask_future_of_axis`, where again `mask_future_of_axis` can be optional.
> Right now, I think the first proposal (just return `[T,B,D]` in the outside loop case) is more intuitive and more straight forward to implement and use.

Yes, I think I agree.
If outside the loop, `CumConcatLayer` would return [T_hist,B,D], for input [T,B,D], where T_hist=T, but with a different special dimension tag.

The new tag would be sth with `history` or so... maybe come up with a better name?

> We can also just implement this behavior as a flag, and then have the benefits of both solutions.

Let's not make it too complicated. Let's first think/draft one possibility, and only implement that.
 
> As you said, it makes everything simpler if the `CumConcatLayer` marked the outputted time-axis in some special way, e.g. using `stag:history` as you mentioned.
> Then we can use `stag` everywhere instead of just `T`. This is a bit verbose, but I think necessary to keep apart the query-axis (i.e. `stag:extern_data_classes` and the key/value-axis (i.e. `stag:history` or sth like that).

Yes, this is really necessary. Also that a layer like `HistoryMaskLayer` (or `SoftmaxOverSpatialLayer`) can operate correctly.

> I am unsure whether a `HistoryMaskLayer` is really necessary.
> We could also add an extra parameter `mask_future_of_axis` to the `SoftmaxOverSpatialLayer`.

Yes. But not sure what would be better. Remember, I don't like that layers become too complex, and more and more options are added. That is our problem with `SelfAttentionLayer` in the first place. It's usually better to separate things into building blocks, which makes it easy for the user to try out variations.

> Then we could simply say
> 
> ```
> att_weights: {class: softmax_over_spatial, from: energy,
>               axis: stag:history, mask_future_of_axis: stag:extern_data:classes?}
> ```

This is the wrong way around, right? The axis for the softmax is `T` or `stag:classes`.
 
> Or instead, `HistoryMaskLayer` could receive two time-axes as argument: e.g. `axis` and `mask_future_of_axis`, where again `mask_future_of_axis` can be optional.

Why two? It only needs to check for such a "history" dim tag, which you could maybe explicitly specify by `axis: "stag:history?"` or so. And if it doesn't find it, it does nothing.

I'm a bit afraid that this could potentially hide bugs, e.g. if you have a typo there for the history axis, you would not get any error, but it would silently just ignore this. This is bad. Any idea about this?

Maybe if outside a loop, the history axis must exist. And if inside the loop, it will ignore it. (This is not perfectly generic, though. E.g. what if you have a loop in a loop? Etc...)

> The new tag would be sth with `history` or so... maybe come up with a better name?
> If outside the loop, CumConcatLayer would return [T_hist,B,D], for input [T,B,D], where T_hist=T, but with a different special dimension tag.

Yes sounds good.

> Let's not make it too complicated. Let's first think/draft one possibility, and only implement that.

Okay sure. Would we still want the layer to have an attribute `as_is_outside` (or another name), and then for now `assert` that it is set to `True` for clarification (implying that `as_is_outside=False` might become the default argument in the future)? Or not?
Maybe we should also think of a better name if we want to include it.

> Yes. But not sure what would be better. Remember, I don't like that layers become too complex, and more and more options are added. That is our problem with `SelfAttentionLayer` in the first place. It's usually better to separate things into building blocks, which makes it easy for the user to try out variations.

Yes okay, you do have a point. Then I would go for the additional masking layer I think.

> > Then we could simply say
> > ```
> > att_weights: {class: softmax_over_spatial, from: energy,
> >               axis: stag:history, mask_future_of_axis: stag:extern_data:classes?}
> > ```
> 
> This is the wrong way around, right? The axis for the softmax is `T` or `stag:classes`.

Actually I think it is correct: `stag:history` would mark the history of all previous keys and values (including the current one), which we want to apply the softmax along. `stag:classes` (the axis of the queries) on the contrary does not exist if inside a loop.


> > Or instead, `HistoryMaskLayer` could receive two time-axes as argument: e.g. `axis` and `mask_future_of_axis`, where again `mask_future_of_axis` can be optional.
>
> Why two? It only needs to check for such a "history" dim tag, which you could maybe explicitly specify by `axis: "stag:history?"` or so. And if it doesn't find it, it does nothing.

I wanted it to be less specific to the `stag:history`-tag, in general my idea was that the layer took two time axes, `T1` and `T2`, and masks all the input to e.g. -infinity if `t1 > t2`.

> I'm a bit afraid that this could potentially hide bugs, e.g. if you have a typo there for the history axis, you would not get any error, but it would silently just ignore this. This is bad. Any idea about this?
> Maybe if outside a loop, the history axis must exist. And if inside the loop, it will ignore it. (This is not perfectly generic, though. E.g. what if you have a loop in a loop? Etc...)

Good point.
Maybe we could make `stag:history_mask?` the default argument? That's not a real fix of course, but should make typos less likely. Or add additional logic somewhere that ensures that one must not reference an stag that was never defined anywhere (I don't know though if there might be a valid reason to do that though)..
> Okay sure. Would we still want the layer [`CumConcatLayer`] to have an attribute `as_is_outside` (or another name), and then for now `assert` that it is set to `True` for clarification (implying that `as_is_outside=False` might become the default argument in the future)? Or not?

Hm, maybe. I'm not sure. The name is bad. Also, what would be the default? Or would this be a required argument? (If it has a default, we can also just leave it away.)

Also, we speak about these two cases inside loop / outside loop, but these are not really the only cases. It could be inside a loop, inside another loop. It would be good to abstract away a bit from that. Also, in an optimal world, the user should not really need to think about this, and just implement it in a way as it would be without automatic optimization, i.e. always inside the loop.

Actually, what are the other arguments of `CumConcatLayer`? It would be good to be able to explicitly specify the axis to be accumulated over. If inside a rec layer, we might need a special placeholder for the rec layer frames. If outside the loop, this axis would exist, but if inside the loop, this axis would not exist. Maybe sth like ":i" (that is also the name for the special layer which reflects the current frame index). This could be the default (`axis=":i"`). Or maybe the default `axis="T"`, and `T` can be treated somewhat special, i.e. namely if it does not exists, it is treated like ":i". And ":i" always implies that this is used inside a rec layer (no matter if inside the loop or outside).

I cannot really come up with a good name or other conceptual way to configure something like `as_is_outside`.
Maybe more like `historic_existing_axis: bool`? I.e. stressing the difference whether the specified axis exists or not -- this is already more generic than the concept whether within loop or not.
Or maybe we can leave that away for now.

> Yes okay, you do have a point. Then I would go for the additional masking layer [`HistoryMaskLayer` or so] I think.

Ok.

Btw, the naming for the current rec layer frame axis (":i" suggested above), the naming of `HistoryMaskLayer`, and the naming of the dimension tag (and/or its special flag) of the newly introduced axis by `CumConcatLayer`, these should be maybe more related, to make it more clear. Maybe:

* Current rec layer frame axis special name: "rec-frame" (would be same as "extern_data/classes" in training, or newly created axis in decoding)
* New dim-tag by `CumConcatLayer` when axis does not exist (e.g. when in rec layer, and inside loop): "rec-history" (or maybe just "rec-frame" as well? would that work?)
* New dim-tag by `CumConcatLayer` when axis exists (e.g. when in rec layer, and outside loop, or just outside rec layer): "left-masked" (or so?) (or also "rec-history"? but wouldn't that be confusing? maybe actually not...)
* "HistoryMaskLayer": `LeftMaskLayer` (this would operate on "left-masked", but not on "rec-history")

Btw, in the `SoftmaxOverSpatialLayer`, the axis you do softmax over, 

> > > ```
> > > att_weights: {class: softmax_over_spatial, from: energy,
> > >               axis: stag:history, mask_future_of_axis: stag:extern_data:classes?}
> > > ```
> > 
> > This is the wrong way around, right? The axis for the softmax is `T` or `stag:classes`.
> 
> Actually I think it is correct: `stag:history` would mark the history of all previous keys and values (including the current one), which we want to apply the softmax along. `stag:classes` (the axis of the queries) on the contrary does not exist if inside a loop.

I think I/we get confused now about the naming of "history" here. But you are right, depending on what the "history" axis refers to. If this refers to the dim tag by `CumConcatLayer` for both cases (cum-concat input axis exists or not -- e.g. inside loop vs outside loop), then this is correct.

> > > Or instead, `HistoryMaskLayer` could receive two time-axes as argument: e.g. `axis` and `mask_future_of_axis`, where again `mask_future_of_axis` can be optional.
> > 
> > Why two? It only needs to check for such a "history" dim tag, which you could maybe explicitly specify by `axis: "stag:history?"` or so. And if it doesn't find it, it does nothing.
> 
> I wanted it to be less specific to the `stag:history`-tag, in general my idea was that the layer took two time axes, `T1` and `T2`, and masks all the input to e.g. -infinity if `t1 > t2`.

The `HistoryMaskLayer` need to have the mode to be a no-op. Specifically when it is inside the loop. But in more generic terms, when you specify the axis to the "history" dim tag. Then it needs to distinguish the "history" dim tag, whether this is from within-loop, or without (the two cases of `CumConcatLayer` I described above).

But we can solve that, by adding a special flag/marker on the `DimensionTag` object, which explicitly tells that, sth like `is_history_of_axis: Optional[DimensionTag]`, and maybe additionally also `is_history_of_rec_loop: Optional[RecLayer]`. It can have the same name "rec-history", but depending on the case, either the first or the second attrib would be set.

However, then the `HistoryMaskLayer` doesn't need to have another axis argument, right? You specify `axis: "stag:rec-history"` or so, and it finds that axis, and its dim tag. If `is_history_of_axis` is set on the dim tag, it expects that this axis is also present in the input, and it does the masking. If `is_history_of_rec_loop` is set, it does nothing. Otherwise it will error. So it really expects such a special axis / dim tag. Then there is no easy way (e.g. by some typo) to get it wrong and silently use no-op.

If you instead have two axes, and do the no-op mode when one of the axis does not exists, this can potentially be problematic. E.g. just a typo could lead to the case that the axis never exists, and it is always a no-op.


> > Okay sure. Would we still want the layer [`CumConcatLayer`] to have an attribute `as_is_outside` (or another name), and then for now `assert` that it is set to `True` for clarification (implying that `as_is_outside=False` might become the default argument in the future)? Or not?
> 
> Hm, maybe. I'm not sure. The name is bad. Also, what would be the default? Or would this be a required argument? (If it has a default, we can also just leave it away.)

I think that either `as_is_outside=False` (or whatever name) should be the default or it should be required, as this is the behavior that works most consistently no matter if it is in a loop or not.
I also can't think of a better name..

> Actually, what are the other arguments of `CumConcatLayer`? It would be good to be able to explicitly specify the axis to be accumulated over. If inside a rec layer, we might need a special placeholder for the rec layer frames. If outside the loop, this axis would exist, but if inside the loop, this axis would not exist. Maybe sth like ":i" (that is also the name for the special layer which reflects the current frame index). This could be the default (`axis=":i"`). Or maybe the default `axis="T"`, and `T` can be treated somewhat special, i.e. namely if it does not exists, it is treated like ":i". And ":i" always implies that this is used inside a rec layer (no matter if inside the loop or outside).

I agree that the axis should be configurable.. But puh, I really can't tell what would be more intuitive :D
At least we agree on the default value, and then maybe we can just skip this for now and always use the most inner rec-time dim.


> Btw, the naming for the current rec layer frame axis (":i" suggested above), the naming of `HistoryMaskLayer`, and the naming of the dimension tag (and/or its special flag) of the newly introduced axis by `CumConcatLayer`, these should be maybe more related, to make it more clear. Maybe:
> 
> * Current rec layer frame axis special name: "rec-frame" (would be same as "extern_data/classes" in training, or newly created axis in decoding)
> * New dim-tag by `CumConcatLayer` when axis does not exist (e.g. when in rec layer, and inside loop): "rec-history" (or maybe just "rec-frame" as well? would that work?)
> * New dim-tag by `CumConcatLayer` when axis exists (e.g. when in rec layer, and outside loop, or just outside rec layer): "left-masked" (or so?) (or also "rec-history"? but wouldn't that be confusing? maybe actually not...)
> * "HistoryMaskLayer": `LeftMaskLayer` (this would operate on "left-masked", but not on "rec-history")

Don't we run into issues if the `CumConcatLayer` creates an axis with a different name depending whether it is in a loop or not?
I mean we would e.g. want to pass this identifier to the `SoftmaxOverSpatial` layer for the attention weights (but also other layers), but if it is called differently depending on whether we are in a loop or not we cannot do that .. right?

So if understood you correctly, you want to call the axis generated by the `CumConcatLayer` in either case (parallel or in-loop) `rec-history`, but add an `is_history_of_axis` and `is_history_of_rec_loop` attribute to the tag as you explained below (which is then used by e.g. `HistoryMaskLayer` to figure out if masking is necessary)?
That sounds like a good idea to me.

Maybe we could also postpone these `DimensionTag`-attributes and initially implement it using the simpler (but much more error prone) check if the query-time axis exists or not. For me it looks somewhat complicated to add these attributes to `DimensionTag`, as they are so widely used across the entire code base.
> I think that either `as_is_outside=False` (or whatever name) should be the default or it should be required, as this is the behavior that works most consistently no matter if it is in a loop or not.
> I also can't think of a better name..

So you think we should have this option on `CumConcatLayer`?
I tend to think not having it is better for now. Esp if it anyway would otherwise be the default, and the only implemented case. Esp also if we cannot find a better name. (The name should be intuitive, i.e. you should understand what it does, without needing to look at the code/documentation.)

> I agree that the axis should be configurable.. But puh, I really can't tell what would be more intuitive :D
> At least we agree on the default value, and then maybe we can just skip this for now and always use the most inner rec-time dim.

The question about the `CumContatLayer` `axis` option is quite clear I think. If it is ":i" / "rec-frame" (or "T", which would be treated special here, just like "rec-frame"), it would use the most inner rec-time dim. And this would also be the default.

> > * Current rec layer frame axis special name: "rec-frame" (would be same as "extern_data/classes" in training, or newly created axis in decoding)
> > * New dim-tag by `CumConcatLayer` when axis does not exist (e.g. when in rec layer, and inside loop): "rec-history" (or maybe just "rec-frame" as well? would that work?)
> > * New dim-tag by `CumConcatLayer` when axis exists (e.g. when in rec layer, and outside loop, or just outside rec layer): "left-masked" (or so?) (or also "rec-history"? but wouldn't that be confusing? maybe actually not...)
> > * "HistoryMaskLayer": `LeftMaskLayer` (this would operate on "left-masked", but not on "rec-history")
> 
> Don't we run into issues if the `CumConcatLayer` creates an axis with a different name depending whether it is in a loop or not?
> I mean we would e.g. want to pass this identifier to the `SoftmaxOverSpatial` layer for the attention weights (but also other layers), but if it is called differently depending on whether we are in a loop or not we cannot do that .. right?

Yes. So concluding from that, it should have the same name in both cases. So "rec-history" (or so) then in both cases.

This is about the name of the new dim tag. It still needs to be distinguished for the other logic (e.g. in `HistoryMaskLayer`). But this could easily be done via `is_history_of_axis`/`is_history_of_rec_loop` flags in the dim tag, as I outlined.

> So if understood you correctly, you want to call the axis generated by the `CumConcatLayer` in either case (parallel or in-loop) `rec-history`, but add an `is_history_of_axis` and `is_history_of_rec_loop` attribute to the tag as you explained below (which is then used by e.g. `HistoryMaskLayer` to figure out if masking is necessary)?
> That sounds like a good idea to me.

Yes exactly.

> Maybe we could also postpone these `DimensionTag`-attributes and initially implement it using the simpler (but much more error prone) check if the query-time axis exists or not.

I don't like that. We always should try to avoid error-prone code. And why not do it directly in a good way? We already have a possible solution, as outlined here.

> For me it looks somewhat complicated to add these attributes to `DimensionTag`, as they are so widely used across the entire code base.

They are widely used, but adding new attributes will not have any effect on existing code, so that shouldn't be a problem. And they would be optional, i.e. `None` by default.

The only problem I see with this is whether these two attribs/flags are maybe very specific for this specific case here, and will not be used in any other case. This is maybe a bit ugly. But maybe this is not too much a problem.

> So you think we should have this option on `CumConcatLayer`?
> I tend to think not having it is better for now. Esp if it anyway would otherwise be the default, and the only implemented case. Esp also if we cannot find a better name. (The name should be intuitive, i.e. you should understand what it does, without needing to look at the code/documentation.)

Okay, then we simply only support the "`as_is_outside=True`" case and omit the argument entirely, right?

> The question about the `CumContatLayer` `axis` option is quite clear I think. If it is ":i" / "rec-frame" (or "T", which would be treated special here, just like "rec-frame"), it would use the most inner rec-time dim. And this would also be the default.

Yeah okay, then let's do it this way.

> > Maybe we could also postpone these `DimensionTag`-attributes and initially implement it using the simpler (but much more error prone) check if the query-time axis exists or not.
> 
> I don't like that. We always should try to avoid error-prone code. And why not do it directly in a good way? We already have a possible solution, as outlined here.

Okay okay, I guess so, fine.

Then everything is more or less clear now, right?
I will then start work on an initial PR, and ask here for now if any further questions come up.

I would start by using the names we chose here (i.e. `CumConcatLayer`, `HistoryMaskLayer` and `rec-history` with `is_history_of_axis`/`is_history_of_rec_loop`). If we find better names we can of course easily change that before we merge this into the main branch.
> Okay, then we simply only support the "`as_is_outside=True`" case [for `CumConcatLayer`] and omit the argument entirely, right?

Yes.

> Then everything is more or less clear now, right?

I think so. At least the draft looks fine and like this could work.

> I will then start work on an initial PR, and ask here for now if any further questions come up.

Yea, I assume there will be subtle problems. But maybe not. We will see.
Create a new branch here such that I can maybe collaborate easily on some problems.

> I would start by using the names we chose here (i.e. `CumConcatLayer`, `HistoryMaskLayer` and `rec-history` with `is_history_of_axis`/`is_history_of_rec_loop`).

Use `RecHistoryMaskLayer`. But otherwise, I think it's good.


Resolving patch:
diff --git a/returnn/tf/layers/rec.py b/returnn/tf/layers/rec.py
--- a/returnn/tf/layers/rec.py
+++ b/returnn/tf/layers/rec.py
@@ -8531,3 +8531,173 @@ def get_out_data_from_opts(cls, name, sources, n_out, **kwargs):
         kind=DimensionTag.Types.Spatial, description="%s_rel_pos_enc_time" % name, dimension=None)
       data = data.copy_template_new_dim_tags((dummy_dim_tag, time_dim_tag, feature_dim_tag))
     return data
+
+
+class CumConcatLayer(_ConcatInputLayer):
+  """
+  Concatenates all previous frames of a time-axis.
+  Like :class:`CumsumLayer` uses `sum`, this layer uses `concat`.
+
+  This layer can be used as a base for auto-regressive self-attention.
+
+  This layer expects to be inside a :class:`RecLayer`.
+
+  Inside a rec loop (not optimized out),
+  this will concatenate the current input
+  to the previous accumulated inputs.
+  For an input of shape `input_shape`,
+  it will output a tensor of shape `[new_dim] + input_shape`.
+  `new_dim` is a special dimension, usually of length `i`,
+  where `i` is the current loop frame,
+  i.e. the length increases in every loop frame.
+  `new_dim` is specified by a separate own dim tag.
+  For example, in the first frame,
+  this will be of shape `[1] + input_shape`,
+  in the second frame shape `[2] + input_shape`,
+  and so on,
+  and in the last frame shape `[T] + input_shape`.
+
+  Outside the rec loop (optimized out),
+  this layer expects an input with the time dim of the rec layer,
+  and returns the input as-is,
+  but replacing the time dim tag with the dim tag `new_dim`
+  converted as outside the loop.
+
+  Normally the optimization should not matter for the user,
+  i.e. for the user, the logical behavior is always as being inside the rec loop.
+  Outside the loop,
+  the output represents a tensor of shape `[T, new_dim] + input_shape`,
+  although we actually have another `new_dim` outside the loop,
+  and `T` is not actually there,
+  but we still have all the information,
+  because the last frame has all information.
+  This `new_dim` outside the loop stores all the dynamic seq lengths
+  per frame of the loop, i.e. the dyn seq len are extended of shape [B,T] or [T]
+  (unlike usually just [B]).
+  This way following layers use different seq lengths of `new_dim` for different loop frames,
+  just like if the `T` dim would actually exist.
+  """
+  layer_class = "cum_concat"
+  recurrent = True  # order matters
+
+  def __init__(self, new_dim, **kwargs):
+    """
+    :param DimensionTag new_dim:
+    """
+    super(CumConcatLayer, self).__init__(**kwargs)
+    rec_layer = self.network.get_rec_parent_layer(inside_loop=False)
+    assert rec_layer, "%r must be used inside a RecLayer" % self
+    out_axis = self.output.get_axis_from_description(new_dim)
+    new_dim_ = self.output.dim_tags[out_axis]
+    assert new_dim_.control_flow_ctx == self.output.control_flow_ctx == self.network.get_control_flow_ctx()
+
+    if not self.input_data.has_axis(rec_layer.time_dim_tag):  # inside loop
+      current_data = self.input_data.copy_compatible_to(self.output, unbroadcast=False)
+      current_frame = current_data.placeholder  # [B, 1, ..., D]
+      last_frames = self._rec_previous_layer.rec_vars_outputs["state"]  # [B, t, ..., D]
+      concat_frames = tf.concat([last_frames, current_frame], axis=out_axis)  # [B, t+1, ..., D]
+      self.rec_vars_outputs["state"] = concat_frames
+      self.output.placeholder = concat_frames
+
+      if not new_dim_.dyn_size_ext:
+        # Unbroadcasting to [B] is not needed because any layers operating on this
+        # should be able to handle extended dyn sizes.
+        # Clipping it to the max length for sequences in the loop which are already ended
+        # (i.e. considering the end flag)
+        # is also not needed because any calculations after the end are irrelevant.
+        # Note: In case we have some initial state/output, this can be extended.
+        dyn_size = self.network.get_rec_step_index() + 1  # scalar
+        new_dim_.dyn_size_ext = Data(
+          name="%s:cum-concat:size-inside" % self.name,
+          dim_tags=[],  # scalar
+          placeholder=dyn_size, dtype="int32",
+          batch=self.output.batch, control_flow_ctx=self.network.get_control_flow_ctx())
+
+    else:  # outside loop
+      # If not inside a rec loop, this layer is a no-op on the tensor.
+      self.output.placeholder = self.input_data.placeholder
+
+      # However, we used new dim tags, which were already prepared.
+      # We now must fill in the extended dynamic size information.
+      if not new_dim_.dyn_size_ext:
+        # This must match the logic above for inside the loop.
+        # Note: In case we have some initial state/output, this can be extended.
+        dyn_size = tf.range(tf.math.reduce_max(rec_layer.time_dim_tag.dyn_size)) + 1  # [T]
+        new_dim_.dyn_size_ext = Data(
+          name="%s:cum-concat:size-outside" % self.name,
+          dim_tags=[rec_layer.time_dim_tag],
+          placeholder=dyn_size, dtype="int32",
+          batch=self.output.batch, control_flow_ctx=self.network.get_control_flow_ctx())
+
+  @classmethod
+  def get_out_data_from_opts(cls, name, network, sources, new_dim, **kwargs):
+    """
+    :param str name:
+    :param returnn.tf.network.TFNetwork network:
+    :param list[LayerBase] sources:
+    :param DimensionTag new_dim:
+    :rtype: Data
+    """
+    input_data = get_concat_sources_data_template(sources, name="%s_output" % name)
+    assert network.is_inside_rec_layer(inside_loop=False), "CumConcatLayer %r must be used inside a RecLayer" % name
+    rec_time_dim = network.get_inside_rec_time_dim(inside_loop=False)
+    assert rec_time_dim
+    ctx = network.get_control_flow_ctx()
+    assert ctx == input_data.control_flow_ctx
+    new_dim_in_ctx = new_dim.get_for_batch_ctx(batch=input_data.batch, ctx=ctx)
+
+    if not input_data.has_axis(rec_time_dim):  # inside loop
+      assert ctx and ctx.is_loop() and ctx.loop_spatial_dim == rec_time_dim
+      # Currently SelectSearchSourcesLayer assumes that all rec_vars_outputs are batch-major.
+      # Therefore we here copy the input as batch-major, and then add the time axis at axis 1.
+      # In the future, when SelectSearchSourcesLayer has support for this, we can change this to operate on axis 0,
+      # which should be more efficient
+      out = input_data.copy_as_batch_major()
+      out = out.copy_add_dim_by_tag(new_dim_in_ctx, unbroadcast=True, axis=1)
+      return out
+
+    else:  # outside loop
+      # Assume that the input has the time dim from the rec layer.
+      axis = input_data.get_axis_from_description(rec_time_dim)
+      return input_data.copy_template_replace_dim_tag(axis=axis, new_dim_tag=new_dim_in_ctx)
+
+  # noinspection PyMethodOverriding
+  @classmethod
+  def get_rec_initial_extra_outputs(cls, network, batch_dim, rec_layer, sources, output, new_dim, **kwargs):
+    """
+    :param returnn.tf.network.TFNetwork network:
+    :param tf.Tensor batch_dim:
+    :param returnn.tf.layers.rec.RecLayer|LayerBase rec_layer:
+    :param list[LayerBase] sources:
+    :param Data output:
+    :param DimensionTag new_dim:
+    :rtype: dict[str,tf.Tensor]
+    """
+    if network.is_inside_rec_layer():
+      shape = []
+      for tag in output.dim_tags:
+        if tag.is_batch_dim():
+          shape.append(batch_dim)
+        elif tag == new_dim:
+          shape.append(0)
+        elif tag.dimension is not None:
+          shape.append(tag.dimension)
+        else:
+          assert tag.dyn_size is not None
+          shape.append(tf.math.reduce_max(tag.dyn_size))
+      return {"state": tf.zeros(shape, dtype=output.dtype)}
+    else:
+      return {}
+
+  @classmethod
+  def get_rec_initial_extra_outputs_shape_invariants(cls, network, sources, output, **kwargs):
+    """
+    :param returnn.tf.network.TFNetwork network:
+    :param list[LayerBase] sources:
+    :param Data output:
+    :rtype: dict[str, tf.TensorShape]
+    """
+    if network.is_inside_rec_layer():
+      return {"state": tf.TensorShape(output.batch_shape)}
+    else:
+      return {}


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -3385,8 +3385,7 @@ def check_reclayer_optimize_out(subnet_layer_dict, other_subnet_layers=None, sha
     rec_layer_dict["unit"].update(other_subnet_layers)
   config = Config({
     "debug_print_layer_output_template": True,
-    "num_inputs": n_in,
-    "num_outputs": n_out
+    "extern_data": {"data": {"dim": n_in}},
   })
   from returnn.tf.layers.rec import _SubnetworkRecCell
   with make_scope() as session:
@@ -3463,6 +3462,40 @@ def test_reclayer_optimize_out_selfatt_left():
     "class": "self_attention", "attention_left_only": True, "num_heads": 2, "total_key_dim": 6, "n_out": 18})
 
 
+def test_reclayer_optimize_out_cum_concat_gen_self_att():
+  new_dim = DimensionTag(kind=DimensionTag.Types.Spatial, description="cum_concat_new_dim")
+  n_key = 5
+  n_value = 7
+  check_reclayer_optimize_out(
+    {"class": "linear", "from": "att", "activation": None},
+    {
+      # This is very much the vanilla self attention,
+      # implemented via the new generic way.
+      # See https://github.com/rwth-i6/returnn/issues/391 for a long discussion.
+      # Commented shapes are always for the layers inside the loop (not optimized).
+      "qkv": {"class": "linear", "from": "data:source", "activation": None, "n_out": n_key * 2 + n_value},  # [B,2*K+V]
+      "qkv_split": {"class": "split", "from": "qkv", "size_splits": [n_key, n_key, n_value]},
+      "q": {"class": "copy", "from": "qkv_split/0"},  # inside [B,K]. optimized out [T,B,K]
+      "k": {"class": "copy", "from": "qkv_split/1"},  # inside [B,K]. optimized out [T,B,K]
+      "v": {"class": "copy", "from": "qkv_split/2"},  # inside [B,V]. optimized out [T,B,V]
+      # cum_concat here. Note that the optimized-out shape is not as you might expect [T,max(t),B,K],
+      # but instead using the optimized format, with extended dyn size on the special dim tag,
+      # i.e. [t*,B,K], representing [T,t*,B,K].
+      "k_accum": {"class": "cum_concat", "new_dim": new_dim, "from": "k"},  # inside [t,B,K]. opt out [t*,B,K]
+      "v_accum": {"class": "cum_concat", "new_dim": new_dim, "from": "v"},  # inside [t,B,V]. opt out [t*,B,K]
+      "energy": {
+        "class": "dot", "from": ["q", "k_accum"],
+        "red1": "static:-1", "red2": "static:-1",
+        "var1": None, "var2": new_dim},  # inside [B,t]. optimized out [T,B,t*]
+      "att_weights": {
+        "class": "softmax_over_spatial", "from": "energy", "axis": new_dim},  # inside [B,t]. opt out [T,B,t*]
+      "att": {
+        "class": "dot", "from": ["att_weights", "v_accum"],
+        "red1": new_dim, "red2": new_dim,
+        "var1": None, "var2": "static:-1"},  # inside [B,V]. opt out [T,B,V]
+    })
+
+
 def test_reclayer_optimize_out_dot():
   # Used for multi-head dot-attention.
   AttNumHeads = 4