Adding additional embeddings after the Q and K are projected, but before the RoPE

[JavaStudentAlex]

Hi,
I would like to use your library for implementing some transformer layers. It is nothing hard, just a cross-attention with ffn and also with RoPE. And here is the example of my code:

attention = AttentionLayers(
    dim=64,
    depth=1,
    heads=8,
    cross_attn_dim_head=8,
    cross_attend=True,
    only_cross=True,
    ff_dropout=0.1,
    rotary_pos_emb=True,
    ff_mult=8

But also what I need is to add some additional embeddings(not positional) to the projected Q and K values but before the RoPE is applied. I have carefully checked the forward parameters from the code and till now could not find any option to do it. Could you please guide me.

[lucidrains]

@JavaStudentAlex hey Alex, thanks for sharing your unique use case

do you mean something like q = self.to_q(x) + q_embed and k = self.to_k(x) + k_embed per layer? could you explain more on what you are trying to do?

[JavaStudentAlex]

Hi, @lucidrains thank you for your answer, basically it is, I am just experimenting with my data and considered also such an option. Could you please advise me or at least give the direction from your POV as a developer how I can implement such a functionality?

[lucidrains]

@JavaStudentAlex yea, that's really interesting, the only time i've seen this pattern was a couple years ago, from shortformer (but the embeddings were positional, and for self-attention)

let me think about it

[lucidrains]

@JavaStudentAlex thought about it some more and think it is worth it

will add it tomorrow morning

[lucidrains]

@JavaStudentAlex oh my mistake, shortformer actually adds it before the projection

the embedding you are adding to projected queries and keys, are they shared across all the attention heads or no?

[JavaStudentAlex]

@lucidrains sorry for a late reply, the embeddings I am adding are distinct, not shared across heads.