Package io.github.kirstenali.deepj.transformer.embeddings

Class RotaryEmbedding

java.lang.Object
io.github.kirstenali.deepj.transformer.embeddings.RotaryEmbedding
public final class RotaryEmbedding extends Object
Rotary Positional Embedding (RoPE) — used in Llama, Mistral, Qwen, DeepSeek, and GPT-NeoX.

Unlike additive positional embeddings, RoPE has no learnable parameters. It encodes position by rotating pairs of Q and K head dimensions by position-dependent angles, which causes relative-position information to appear naturally in dot-product attention scores.

Applied inside MultiHeadSelfAttention after the Q/K projections are split into heads, before the scaled dot-product is computed.

Math (per position t, pair index i):

   θ_{t,i}  = t / 10000^(2i / headDim)

   x_rot[2i]   =  x[2i]  · cos θ  −  x[2i+1] · sin θ
   x_rot[2i+1] =  x[2i]  · sin θ  +  x[2i+1] · cos θ

Backward (transpose rotation — negate sin):

   dx[2i]   =  d[2i]  · cos θ  +  d[2i+1] · sin θ
   dx[2i+1] = −d[2i]  · sin θ  +  d[2i+1] · cos θ

Pairing convention: this implementation uses the interleaved pairing of adjacent dimensions (2i, 2i+1) — the original RoPE / GPT-NeoX formulation — rather than the "half-split" pairing (i, i + headDim/2) used by some Llama/HF checkpoints. Both are mathematically valid and self-consistent as long as apply(io.github.kirstenali.deepj.tensor.Tensor, int, int) and applyBackward(io.github.kirstenali.deepj.tensor.Tensor, int, int) agree; however the two conventions are not weight-compatible with each other, so externally-trained RoPE weights must match this pairing.

Input tensor shape convention (split-head layout): [nHeads × seqLen, headDim] — head h occupies rows [h·seqLen .. (h+1)·seqLen), position t is at row h·seqLen + t.

  • Constructor Details

    • RotaryEmbedding

      public RotaryEmbedding(int headDim, int maxSeqLen)
      Pre-computes the cos/sin rotation tables.
      Parameters:
      headDim - dimension of each attention head (dModel / nHeads); must be even
      maxSeqLen - maximum sequence length to support

  • Method Details

    • apply

      public Tensor apply(Tensor t, int seqLen, int nHeads)
      Apply rotary embeddings to a split-head tensor (forward direction).
      Parameters:
      t - split-head tensor, shape [nHeads·seqLen × headDim]
      seqLen - number of positions
      nHeads - number of attention heads
      Returns:
      rotated tensor with the same shape

    • applyBackward

      public Tensor applyBackward(Tensor t, int seqLen, int nHeads)
      Apply the transpose (inverse) rotation — used in the backward pass.
      Parameters:
      t - gradient tensor, same shape as the forward input
      seqLen - number of positions
      nHeads - number of attention heads
      Returns:
      un-rotated gradient with the same shape

    • headDim

      public int headDim()