Question about Rotary Embedding Sequence in Model Code vs. Diagrams

[littlepsilon]

Hello,

I have been studying the Gemma-Pytorch implementation and noticed a potential discrepancy in the code related to the application of rotary embeddings in the attention mechanism.

Specifically, the rotary embeddings are applied before calculating the attention scores in the code between lines 213 and 331, which seems to contradict several diagrams that depict these embeddings being applied after the attention calculation.
code: https://github.com/google/gemma_pytorch/blob/main/gemma/model.py#L213
diagram: https://developers.googleblog.com/en/gemma-explained-new-in-gemma-2/

Could you please clarify whether the embeddings should indeed be applied before the attention scores are calculated, or if the diagrams might have represented this differently?

Thank you!

[Balakrishna-Chennamsetti]

Hi @littlepsilon ,

Apologies for the late reply, yes your observation on this Gemma model is correct. Rotary embeddings (RoPE) are applied to the Query (Q) and Key (K) vectors before the attention scores are calculated via the dot product. This is fundamental to how RoPE enables relative positional encoding within the attention mechanism. The diagrams are likely simplifying the visual representation or focusing on the overall data flow rather than the precise low-level computational sequence within the attention head.

Thanks.