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

Class MultiHeadSelfAttention

java.lang.Object

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

All Implemented Interfaces:

Layer, Trainable

Direct Known Subclasses:

RoPEMultiHeadSelfAttention

public class MultiHeadSelfAttention
extends Object
implements Layer

Multi-head causal self-attention for a single sequence (no batch dimension). Input/Output shape: [seqLen x dModel]

Field Summary

Fields

Modifier and Type	Field	Description
protected final int	headDim	Exposed to subclasses that need per-head dimension for custom Q/K transforms.
protected final int	nHeads	Exposed to subclasses that need head count for custom Q/K transforms (e.g.

Constructor Summary

Constructors

Constructor	Description
MultiHeadSelfAttention(int dModel, int nHeads, boolean causalMask, Random rnd)

Method Summary

Modifier and Type	Method	Description
Tensor	backward(Tensor dOut)
Tensor	forward(Tensor x)
List<Parameter>	parameters()
protected Tensor	transformQueryKey(Tensor heads, int seqLen)	Hook called on split-head Q and K tensors ([nHeads·seqLen × headDim]) before the scaled dot-product.
protected Tensor	transformQueryKeyBackward(Tensor gradHeads, int seqLen)	Inverse of transformQueryKey(io.github.kirstenali.deepj.tensor.Tensor, int), called on Q/K gradients in backward.

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

Field Details

nHeads

protected final int nHeads

Exposed to subclasses that need head count for custom Q/K transforms (e.g. RoPE).

headDim

protected final int headDim

Exposed to subclasses that need per-head dimension for custom Q/K transforms.

Constructor Details

MultiHeadSelfAttention

public MultiHeadSelfAttention(int dModel,
 int nHeads,
 boolean causalMask,
 Random rnd)

Method Details

forward

public Tensor forward(Tensor x)

Specified by:

forward in interface Layer

backward

public Tensor backward(Tensor dOut)

Specified by:

backward in interface Layer

transformQueryKey

protected Tensor transformQueryKey(Tensor heads,
 int seqLen)

Hook called on split-head Q and K tensors ([nHeads·seqLen × headDim]) before the scaled dot-product. No-op by default; override to apply RoPE etc.

transformQueryKeyBackward

protected Tensor transformQueryKeyBackward(Tensor gradHeads,
 int seqLen)

Inverse of transformQueryKey(io.github.kirstenali.deepj.tensor.Tensor, int), called on Q/K gradients in backward. No-op by default; override to invert whatever forward transform was applied.

parameters

public List<Parameter> parameters()

Specified by:

parameters in interface Trainable