lucidrains/block-recurrent-transformer-pytorch

Implementation of Block Recurrent Transformer - Pytorch

block-recurrent-transformer-pytorch's Issues

It seems line 942 states = [] is override states from params

def generate(
    self,
    prime,
    length = None,
    xl_memories: List[torch.Tensor] = [],
    states: List[torch.Tensor] = [],
    temperature = 1.,
    filter_thres = 0.9,
    return_memories_and_states = False
):
    length = default(length, self.max_seq_len + 1)
    start_len = prime.shape[-1]

    assert start_len < self.max_seq_len
    assert length <= (self.max_seq_len + 1)
    assert start_len < length

    output = prime

    memories = []
    states = []

    for ind in range(length - start_len):

Question

Is it supposed to detach?

Inside block_recurrent_transformer_pytorch.py line 815

if exists(layer_next_states):
next_states.append(layer_next_states.detach())

How would the gradients flow through the states?

recurrent state

Hi, i noticed that recurrent state keys and values are derived from self.init_state rather than previous recurrent state. In this case, the keys and value embedding layers, along with the projection and MLP layers for horizontal attention will not be trained? I am not sure if this is the intended behavior as I could not find any details regarding this in the paper.

Incorrect pip install in README

Hello, I believe there is a typo in the "Install" section of the readme.

Instead of:

pip install block-recurrent-transformer

I used this:

pip install block-recurrent-transformer-pytorch

Suggestion

Can you try to input the recurrent state to the first transformer layer? Like a recurrent layer but without next state? Maybe that will improve the performance since high level representations don't have to be recomputed and remembered info can be passed through the whole model again without any additional compute cost

Cross attention wrong query vector

Hi, in the AttentionBlock forward function, the vertical attention is computing cross-attention with Qe^v instead of Qs^v? Same goes for the horizontal direction, it is computing using Qs^h instead of Qe^h.