Transformer XL - understanding paper's illustration

Question

If I understand correctly, the Key hidden layer in the Transformer XL is of size 2L * d, where L is the segment length and d is the embedding dimension.

concatenation of two hidden sequences along the length dimension

Therefore, the size of the attention matrix would be L X 2L, where row i represents the attention Query i should apply to each of the 2L Keys.

That is, the self attention window length = 2 X segment length.

However, in the following image from the paper, the segment length is 4 and there are only 4 lines linked to each node. Shouldn't there be 4 * 2 = 8 lines from each node?

Link to transformer XL paper

Answer 1

If you look at the Github code, there is 2xNxD in the Multi head attention function indeed:

class MultiHeadAttn(nn.Module):
    def __init__(self, n_head, d_model, d_head, dropout, dropatt=0, 
                 pre_lnorm=False):
        super(MultiHeadAttn, self).__init__()

        self.n_head = n_head
        self.d_model = d_model
        self.d_head = d_head
        self.dropout = dropout

        self.q_net = nn.Linear(d_model, n_head * d_head, bias=False)
        self.kv_net = nn.Linear(d_model, 2 * n_head * d_head, bias=False)

        self.drop = nn.Dropout(dropout)
        self.dropatt = nn.Dropout(dropatt)
        self.o_net = nn.Linear(n_head * d_head, d_model, bias=False)

        self.layer_norm = nn.LayerNorm(d_model)

        self.scale = 1 / (d_head ** 0.5)

        self.pre_lnorm = pre_lnorm 

But kv refers to the key AND value vectors.

Source: https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/mem_transformer.py