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

Class MultiHeadLatentAttention

java.lang.Object

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

All Implemented Interfaces:

Layer, Trainable

public final class MultiHeadLatentAttention
extends Object
implements Layer

Compact Multi-Head Latent Attention (MLA) inspired by DeepSeek-V2/V3.

Unlike standard MHA which projects Q, K, V directly from x, MLA first compresses through a low-rank bottleneck then expands:

   cQ  = x   · Wdq      [seqLen × qRank]   — Q  compression
   Q   = cQ  · Wuq      [seqLen × dModel]  — Q  expansion

   cKV = x   · Wdkv     [seqLen × kvRank]  — shared KV compression
   K   = cKV · Wuk      [seqLen × dModel]  — K  expansion
   V   = cKV · Wuv      [seqLen × dModel]  — V  expansion
 

RoPE is applied to Q and K after expansion. Scaled dot-product attention and the output projection Wo are then identical to standard MHA.

The factorisation can support caching only cKV in an incremental decoder. This layer currently computes a complete sequence and does not own an inference cache.

Parameters: Wdq, Wuq, Wdkv, Wuk, Wuv, Wo (6 total, vs 4 in standard MHA).

Constructor Summary

Constructors

Constructor	Description
MultiHeadLatentAttention(int dModel, int nHeads, int qRank, int kvRank, RotaryEmbedding rope, Random rnd)

Method Summary

Modifier and Type	Method	Description
Tensor	backward(Tensor dOut)
Tensor	forward(Tensor x)
List<Parameter>	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

MultiHeadLatentAttention

public MultiHeadLatentAttention(int dModel,
 int nHeads,
 int qRank,
 int kvRank,
 RotaryEmbedding rope,
 Random rnd)

Parameters:

dModel - model dimension

nHeads - number of attention heads; must divide dModel

qRank - Q latent dimension (e.g. 1536 in DeepSeek-V2; use dModel/2 for small models)

kvRank - KV latent dimension (e.g. 512 in DeepSeek-V2; use dModel/4 for small models)

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

rnd - random source for weight initialisation

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

parameters

public List<Parameter> parameters()

Specified by:

parameters in interface Trainable