Package io.github.kirstenali.deepj.layers.transformer.attention

Class RoPEMultiHeadSelfAttention

java.lang.Object

io.github.kirstenali.deepj.layers.transformer.attention.MultiHeadSelfAttention

io.github.kirstenali.deepj.layers.transformer.attention.RoPEMultiHeadSelfAttention

All Implemented Interfaces:

Layer, Trainable

public final class RoPEMultiHeadSelfAttention
extends MultiHeadSelfAttention

Multi-head self-attention with Rotary Positional Embedding (RoPE).

RoPE attention is common in model families such as Llama, Mistral, Qwen, DeepSeek, and GPT-NeoX. This compact implementation extends MultiHeadSelfAttention and overrides only the two Q/K transform hooks — all attention mechanics, causal masking, and backpropagation are inherited unchanged.

The RotaryEmbedding is applied to Q and K after they are projected and split into per-head blocks, before the scaled dot-product is computed:

   Q_rot = RoPE(Q_heads)   K_rot = RoPE(K_heads)
   scores = softmax( Q_rot · K_rotᵀ / √d ) · V
 

Typical usage inside a Llama-style block:

 RotaryEmbedding rope = new RotaryEmbedding(dModel / nHeads, maxSeqLen);
 new GPTTransformerBlock(
     new RMSNorm1D(dModel),
     new RMSNorm1D(dModel),
     new RoPEMultiHeadSelfAttention(dModel, nHeads, true, rope, rnd),
     new SwiGLULayer(dModel, dFF, rnd)
 );
 

Field Summary

Fields inherited from class io.github.kirstenali.deepj.layers.transformer.attention.MultiHeadSelfAttention

headDim, nHeads

Constructor Summary

Constructors

Constructor	Description
RoPEMultiHeadSelfAttention(int dModel, int nHeads, boolean causalMask, RotaryEmbedding rope, Random rnd)

Method Summary

Modifier and Type	Method	Description
protected Tensor	transformQueryKey(Tensor heads, int seqLen)	Applies RoPE rotation to Q or K heads in the forward pass.
protected Tensor	transformQueryKeyBackward(Tensor gradHeads, int seqLen)	Applies the inverse (transpose) RoPE rotation to Q/K gradients in the backward pass.

Methods inherited from class io.github.kirstenali.deepj.layers.transformer.attention.MultiHeadSelfAttention

backward, forward, parameters

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Methods inherited from interface io.github.kirstenali.deepj.training.Trainable

zeroGrad

Constructor Details

RoPEMultiHeadSelfAttention

public RoPEMultiHeadSelfAttention(int dModel,
 int nHeads,
 boolean causalMask,
 RotaryEmbedding rope,
 Random rnd)

Parameters:

dModel - model (embedding) dimension

nHeads - number of attention heads; must divide dModel evenly

causalMask - true for autoregressive (decoder) attention

rope - pre-built rotary embedding sized for dModel / nHeads

rnd - random source for weight initialisation

Method Details

transformQueryKey

protected Tensor transformQueryKey(Tensor heads,
 int seqLen)

Applies RoPE rotation to Q or K heads in the forward pass. Input/output shape: [nHeads·seqLen × headDim].

Overrides:

transformQueryKey in class MultiHeadSelfAttention

transformQueryKeyBackward

protected Tensor transformQueryKeyBackward(Tensor gradHeads,
 int seqLen)

Applies the inverse (transpose) RoPE rotation to Q/K gradients in the backward pass. Because rotation matrices are orthogonal, the inverse is the transpose (negate sin).

Overrides:

transformQueryKeyBackward in class MultiHeadSelfAttention