Permutation Equivariance of Self-Attention (No Positional Encoding)

Self-attention (without positional information) should be equivariant to permutations of the token axis:

If we reorder the input tokens, the output is reordered in the same way.

This file formalizes that statement for TorchLean’s spec-layer Spec.selfAttention over ℝ.

The reindexing adapters below are exported in the model-proof namespace, not the general TorchLean.Tensor API. They describe token permutations used by the attention theorems.

Token reindexing

The adapters rebuild tensor slices in the order selected by the permutation.

def NN.Proofs.Models.Attention.reindexOuter {α : Type} [TorchLean.Storage α] {n : ℕ} {s : Spec.Shape} (σ : Equiv.Perm (Fin n)) : TorchLean.Tensor α (Spec.Shape.dim n s) → TorchLean.Tensor α (Spec.Shape.dim n s)

Reindex the outermost axis of a tensor by a permutation.

@[simp]

theorem NN.Proofs.Models.Attention.get_reindexOuter {α : Type} [TorchLean.Storage α] {n : ℕ} {s : Spec.Shape} (σ : Equiv.Perm (Fin n)) (t : TorchLean.Tensor α (Spec.Shape.dim n s)) (i : Fin n) : Spec.get (reindexOuter σ t) i = Spec.get t (σ i)

Slice i of the permuted tensor is slice σ i of the original.

def NN.Proofs.Models.Attention.reindexCols {α : Type} [TorchLean.Storage α] {m n : ℕ} (σ : Equiv.Perm (Fin n)) : TorchLean.Tensor α [m, n] → TorchLean.Tensor α [m, n]

Reindex the column axis of a matrix by a permutation.

@[simp]

theorem NN.Proofs.Models.Attention.get2_reindexOuter {α : Type} [TorchLean.Storage α] {m n : ℕ} (σ : Equiv.Perm (Fin m)) (A : TorchLean.Tensor α [m, n]) (i : Fin m) (j : Fin n) : Spec.get2 (reindexOuter σ A) i j = Spec.get2 A (σ i) j

Matrix form: permuting the outer axis permutes rows.

@[simp]

theorem NN.Proofs.Models.Attention.get2_reindexCols {α : Type} [TorchLean.Storage α] {m n : ℕ} (σ : Equiv.Perm (Fin n)) (A : TorchLean.Tensor α [m, n]) (i : Fin m) (j : Fin n) : Spec.get2 (reindexCols σ A) i j = Spec.get2 A i (σ j)

And permuting the inner axis permutes columns.

def NN.Proofs.Models.Attention.permMatrix {α : Type} [TorchLean.Storage α] {n : ℕ} (σ : Equiv.Perm (Fin n)) (A : TorchLean.Tensor α [n, n]) : TorchLean.Tensor α [n, n]

Simultaneously permute rows and columns of an n×n matrix by the same permutation.

@[simp]

theorem NN.Proofs.Models.Attention.get2_permMatrix {α : Type} [TorchLean.Storage α] {n : ℕ} (σ : Equiv.Perm (Fin n)) (A : TorchLean.Tensor α [n, n]) (i j : Fin n) : Spec.get2 (permMatrix σ A) i j = Spec.get2 A (σ i) (σ j)

Entries of a simultaneously permuted matrix: (σ i, σ j) of the original.

This is the form attention scores take under a token permutation, since both the query index and the key index move with the same σ.

Softmax equivariance

TorchLean's spec softmax on vectors is implemented in a stabilized way (x ↦ exp(x - m) / Σ exp(x - m)), but over ℝ it agrees with the plain exp(x) / Σ exp(x) formula (Proofs.getScalar_softmaxVecSpec_eq_exp_div). This lets us prove permutation equivariance without reasoning about how the stabilizing shift m is chosen.

theorem NN.Proofs.Models.Attention.SoftmaxEquivariance.softmaxVecSpec_reindexOuter {n : ℕ} (σ : Equiv.Perm (Fin n)) (t : TorchLean.Tensor ℝ [n]) : Activation.softmaxVecSpec (reindexOuter σ t) = reindexOuter σ (Activation.softmaxVecSpec t)

Spec vector softmax commutes with reindexing, including for an empty vector.

theorem NN.Proofs.Models.Attention.SoftmaxEquivariance.softmaxSpec_permMatrix {n : ℕ} (σ : Equiv.Perm (Fin n)) (A : TorchLean.Tensor ℝ [n, n]) : Activation.softmaxSpec 1 (permMatrix σ A) = permMatrix σ (Activation.softmaxSpec 1 A)

Matrix axis-1 softmax commutes with simultaneous row/column permutations.

Linear algebra: matmul/transpose commute with permutations

theorem NN.Proofs.Models.Attention.matMulSpec_reindexOuter_reindexCols {m n p : ℕ} (σ : Equiv.Perm (Fin m)) (τ : Equiv.Perm (Fin p)) (A : TorchLean.Tensor ℝ [m, n]) (B : TorchLean.Tensor ℝ [n, p]) : Spec.matMulSpec (reindexOuter σ A) (reindexCols τ B) = reindexOuter σ (reindexCols τ (Spec.matMulSpec A B))

Matrix multiplication commutes with independent output-row and output-column reindexing.

This is the most general bookkeeping lemma used below: reindexing rows of the left factor controls the rows of the product, while reindexing columns of the right factor controls the columns of the product.

theorem NN.Proofs.Models.Attention.matMulSpec_reindexOuter_left {m n p : ℕ} (σ : Equiv.Perm (Fin m)) (A : TorchLean.Tensor ℝ [m, n]) (B : TorchLean.Tensor ℝ [n, p]) : Spec.matMulSpec (reindexOuter σ A) B = reindexOuter σ (Spec.matMulSpec A B)

Left-multiplying by a row-permuted matrix only row-permutes the matrix-product output.

This is the projection-layer version of token equivariance: the same learned weight matrix is applied independently to every token, so changing token order before the projection merely changes the output token order.

theorem NN.Proofs.Models.Attention.swapAdjacentAxes_reindexOuter {m n : ℕ} (σ : Equiv.Perm (Fin m)) (A : TorchLean.Tensor ℝ [m, n]) : (reindexOuter σ A).swapAdjacentAxes 0 = reindexCols σ (A.swapAdjacentAxes 0)

Transposition converts a row permutation into a column permutation.

The attention-score proof uses this to turn (P Q) (P K)ᵀ into a simultaneous row/column permutation of Q Kᵀ.

Main theorem: self-attention equivariance

theorem NN.Proofs.Models.Attention.scaleSpec_permMatrix {n : ℕ} (σ : Equiv.Perm (Fin n)) (A : TorchLean.Tensor ℝ [n, n]) (c : ℝ) : (permMatrix σ A).scaleSpec c = permMatrix σ (A.scaleSpec c)

Elementwise scaling commutes with simultaneous row/column reindexing.

Scaled dot-product attention divides all score entries by the same scalar, so this lemma lets the token-permutation proof move the scale step past the score-matrix conjugation.

theorem NN.Proofs.Models.Attention.matMulSpec_permMatrix_reindexOuter {n d : ℕ} (σ : Equiv.Perm (Fin n)) (A : TorchLean.Tensor ℝ [n, n]) (B : TorchLean.Tensor ℝ [n, d]) : Spec.matMulSpec (permMatrix σ A) (reindexOuter σ B) = reindexOuter σ (Spec.matMulSpec A B)

Multiplying a simultaneously row/column-permuted attention matrix by a row-permuted value matrix produces a row-permuted output.

This is the final linear-algebra step in self-attention equivariance: the permutation of the attention weights and the permutation of the value rows cancel on the internal summation index.

theorem NN.Proofs.Models.Attention.selfAttention_reindexOuter {n dModel projDim : ℕ} (σ : Equiv.Perm (Fin n)) (x : TorchLean.Tensor ℝ [n, dModel]) (Wq Wk Wv : TorchLean.Tensor ℝ [dModel, projDim]) (Wo : TorchLean.Tensor ℝ [projDim, dModel]) (h1 : n ≠ 0) : Spec.selfAttention (reindexOuter σ x) Wq Wk Wv Wo h1 = reindexOuter σ (Spec.selfAttention x Wq Wk Wv Wo h1)

Self-attention without positional encodings is equivariant to any token permutation.

If the input token axis is reordered by σ, then Q, K, and V are reordered in the same way, the score matrix is conjugated by the corresponding permutation matrix, row-wise softmax commutes with that conjugation, and the final output is exactly the same token reordering of the original output. This theorem is intentionally stated for the spec-layer block, not for CUDA kernels.