Positional encodings (spec layer)

This file provides learnable and sinusoidal positional embeddings, added to token embeddings, and RoPE rotations of query/key head vectors.

PyTorch analogy:

• This is the same idea as having an nn.Embedding(max_len, d_model) (or a parameter tensor [max_len, d_model]) and doing x + pos[:seqLen].

Why learnable positional encodings show up a lot in practice:

• they are easy to train and tend to work well for fixed-length settings (e.g. ViT with a chosen patch grid, or language models trained with a fixed max sequence length),
• they keep the spec algebraic: there is no trigonometry, complex numbers, or special casing for even/odd dimensions.

The sinusoidal helpers below construct fixed additive embeddings. The RoPE helpers instead combine each head vector with its pairwise rotation using position-dependent cosine and sine factors.

Reference (sinusoidal): "Attention Is All You Need" (Vaswani et al., 2017): https://arxiv.org/abs/1706.03762

structure Spec.PositionalEncodingSpec (seqLen embedDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Learnable positional encoding parameters for a fixed (seqLen, embedDim).

This record stores the trainable positional table. Higher-level models decide how to initialize it and whether to share or resize it across different sequence lengths.

• pos : TorchLean.Tensor α [seqLen, embedDim]

  The learnable positional table, one embedDim vector per position.

def Spec.addPositionalEncodingSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen embedDim : ℕ} (pe : PositionalEncodingSpec seqLen embedDim α) (x : TorchLean.Tensor α [seqLen, embedDim]) : TorchLean.Tensor α [seqLen, embedDim]

Add positional encodings: y = x + pos.

Both x and pos have the same shape, so this is just elementwise addition (no broadcasting).

Gradients

Learnable positional encoding is just an elementwise addition:

y = x + pos

So the adjoint is just:

• δx = δy
• δpos = δy

The same upstream tensor supplies both gradients.

Sinusoidal positional encodings (pure functions)

These are the classic Transformer sinusoidal encodings from:

Vaswani et al. (2017), "Attention Is All You Need".

We implement them as pure tensor generators so they can be reused in multiple model specs without adding new trainable parameters.

def Spec.posencFreqDenomSpec {α : Type} [Context α] (i d : ℕ) : α

Frequency denominator used by both sinusoidal PE and RoPE:

denom(i; d) = 10000^(2*i / d)

implemented as:

exp(log(10000) * (2*i / d))

Each operation uses the selected scalar backend. On a grid that cannot represent 10000, the logarithm receives the rounded base; this is not a single rounding of the exact real logarithm.

def Spec.posencAngleSpec {α : Type} [Context α] (pos i d : ℕ) : α

Common angle used by sinusoidal PE and RoPE:

θ(pos, i; d) = pos / 10000^(2*i/d).

def Spec.sinusoidalPositionalEncodingSpec {α : Type} [TorchLean.Storage α] [Context α] (seqLen embedDim : ℕ) (startPos : ℕ := 0) : TorchLean.Tensor α [seqLen, embedDim]

Pure sinusoidal positional encoding tensor with shape (seqLen, embedDim).

Definition (Transformer):

• PE[pos, 2i] = sin(θ(pos, i; embedDim))
• PE[pos, 2i+1] = cos(θ(pos, i; embedDim))

startPos is an offset for the absolute positions; use it when generating a chunk of positions for cached decoding (e.g. tokens startPos .. startPos+seqLen-1).

This definition is total for all seqLen/embedDim:

• if embedDim = 0, the inner dimension is empty, so no scalar computations are observed.
• if embedDim is odd, the last column uses the same i = floor(j/2) convention as usual.

def Spec.addSinusoidalPositionalEncodingSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen embedDim : ℕ} (x : TorchLean.Tensor α [seqLen, embedDim]) (startPos : ℕ := 0) : TorchLean.Tensor α [seqLen, embedDim]

Add sinusoidal positional encodings: y = x + sinusoidal(startPos, seqLen, embedDim).

Rotary positional embeddings (RoPE) utilities (pure functions)

RoPE (Su et al., 2021) encodes position by applying a 2D rotation to each pair of features in the last dimension.

In most transformer implementations, RoPE is applied to query/key head vectors:

• per head: (seqLen, headDim)
• all heads: (numHeads, seqLen, headDim)

This file provides pure RoPE helpers. Attention modules can apply these helpers before the query/key dot product.

References:

• Su et al. (2021), "RoFormer: Enhanced Transformer with Rotary Position Embedding".
• Many implementations use the same θ(pos,i;d) frequency schedule as sinusoidal PE.

def Spec.ropeRotatePairsSpec {α : Type} [TorchLean.Storage α] [Context α] {headDim : ℕ} (x : TorchLean.Tensor α [headDim]) : TorchLean.Tensor α [headDim]

Rotate pairs on the last dimension:

(x0, x1, x2, x3, ...) ↦ (-x1, x0, -x3, x2, ...).

This corresponds to multiplying each 2-vector (x_even, x_odd) by the matrix: [[0, -1], [1, 0]].

Design note:

• Standard RoPE assumes headDim is even.
• This spec function is total: if headDim is odd, the last (unpaired) entry is left unchanged.

def Spec.ropeCosVectorSpec {α : Type} [TorchLean.Storage α] [Context α] (pos headDim : ℕ) : TorchLean.Tensor α [headDim]

Broadcast RoPE cos(θ) factors to a full (headDim) vector for one position.

For an odd width, the final unpaired coordinate receives cosine factor 1.

def Spec.ropeSinVectorSpec {α : Type} [TorchLean.Storage α] [Context α] (pos headDim : ℕ) : TorchLean.Tensor α [headDim]

Broadcast RoPE sin(θ) factors to a full (headDim) vector for one position.

For an odd width, the final unpaired coordinate receives sine factor 0. Together with its cosine factor 1, ropeApplySpec evaluates that coordinate as x * 1 + x * 0.

def Spec.ropeApplySpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen headDim : ℕ} (x : TorchLean.Tensor α [seqLen, headDim]) (startPos : ℕ := 0) : TorchLean.Tensor α [seqLen, headDim]

Apply RoPE to a single head matrix x : (seqLen, headDim).

Implementation matches the standard identity:

rope(x) = x * cos + rotatePairs(x) * sin

where cos and sin are position-dependent vectors broadcast across the last dimension.

startPos is an absolute-position offset (useful for KV-cache decoding). When headDim is odd, the final unpaired coordinate evaluates to x * 1 + x * 0. This is an identity over the reals. IEEE arithmetic still performs both products and the addition: an infinite entry produces NaN through ∞ * 0, and a NaN entry remains NaN.

def Spec.ropeApplyHeadsSpec {α : Type} [TorchLean.Storage α] [Context α] {numHeads seqLen headDim : ℕ} (x : TorchLean.Tensor α [numHeads, seqLen, headDim]) (startPos : ℕ := 0) : TorchLean.Tensor α [numHeads, seqLen, headDim]

Apply RoPE to (numHeads, seqLen, headDim) by applying ropeApplySpec independently per head.