Matrix-Based Network Operations

This module connects the optimized IntervalMatrix infrastructure with the neural network layer definitions, enabling:

1. Batch Processing: Forward pass on multiple inputs simultaneously
2. Attention Mechanism: Q × K^T computation for Transformer verification
3. Efficient Layer Application: Using quantized integer arithmetic

Key Definitions

• toIntervalMatrix - Convert weight matrix to IntervalMatrix format
• batchForward - Apply layer to multiple inputs (batch dimension)
• attentionScores - Compute Q × K^T for self-attention

Design Notes

The IntervalMatrix uses Structure-of-Arrays layout for cache efficiency. For Transformer verification, the key operations are:

• Linear projections (matmul with weight matrices)
• Attention score computation (Q × K^T)
• Softmax bounds (separate module)

Conversion Functions

def LeanCert.ML.Optimized.rationalToIntMatrix (rows cols : ℕ) (weights : List (List ℚ)) (exp : ℤ) : IntMatrix rows cols

Convert a rational weight matrix (List of Lists) to IntMatrix with given exponent. Each rational q is converted to an integer q * 2^(-exp) (rounded).

Equations

• One or more equations did not get rendered due to their size.

def LeanCert.ML.Optimized.layerToIntervalMatrix (l : Layer) (exp : ℤ) : IntervalMatrix l.outputDim l.inputDim

Convert Layer weights to IntervalMatrix (point intervals). Useful for propagating through linear layers with matrix multiply.

Equations

• One or more equations did not get rendered due to their size.

Batch Forward Pass

@[reducible, inline]

abbrev LeanCert.ML.Optimized.InputBatch (inputDim batchSize : ℕ) : Type

A batch of inputs represented as an IntervalMatrix. Each column is one input vector, rows are features. Shape: (input_dim, batch_size)

Equations

• LeanCert.ML.Optimized.InputBatch inputDim batchSize = LeanCert.ML.Optimized.IntervalMatrix inputDim batchSize

@[reducible, inline]

abbrev LeanCert.ML.Optimized.OutputBatch (outputDim batchSize : ℕ) : Type

A batch of outputs represented as an IntervalMatrix. Shape: (output_dim, batch_size)

Equations

• LeanCert.ML.Optimized.OutputBatch outputDim batchSize = LeanCert.ML.Optimized.IntervalMatrix outputDim batchSize

def LeanCert.ML.Optimized.batchLinear {outDim inDim batchSize : ℕ} (W : IntervalMatrix outDim inDim) (X : InputBatch inDim batchSize) : OutputBatch outDim batchSize

Apply weight matrix to a batch of inputs: W × X where W : (output_dim, input_dim) and X : (input_dim, batch_size) Result: (output_dim, batch_size)

Equations

• LeanCert.ML.Optimized.batchLinear W X = W.matmul X

Attention Mechanism Support

def LeanCert.ML.Optimized.attentionScores {seqLen dK : ℕ} (Q K : IntervalMatrix seqLen dK) : IntervalMatrix seqLen seqLen

Compute attention scores: Q × K^T Q : (seq_len, d_k) - Query matrix K : (seq_len, d_k) - Key matrix Result: (seq_len, seq_len) - Attention score matrix

Equations

• LeanCert.ML.Optimized.attentionScores Q K = Q.matmul K.transpose

structure LeanCert.ML.Optimized.AttentionParams (dModel numHeads : ℕ) : Type

Structure for Multi-Head Attention parameters

• wQ : IntervalMatrix dModel dModel

  Query projection weights (dModel, dModel)

• wK : IntervalMatrix dModel dModel

  Key projection weights (dModel, dModel)

• wV : IntervalMatrix dModel dModel

  Value projection weights (dModel, dModel)

• wO : IntervalMatrix dModel dModel

  Output projection weights (dModel, dModel)

def LeanCert.ML.Optimized.projectQKV {numHeads seqLen dModel : ℕ} (X : IntervalMatrix seqLen dModel) (params : AttentionParams dModel numHeads) : IntervalMatrix seqLen dModel × IntervalMatrix seqLen dModel × IntervalMatrix seqLen dModel

Project input to Q, K, V matrices. Input X : (seq_len, d_model) Returns (Q, K, V) each of shape (seq_len, d_model)

Equations

• LeanCert.ML.Optimized.projectQKV X params = (X.matmul params.wQ.transpose, X.matmul params.wK.transpose, X.matmul params.wV.transpose)

ReLU for Matrices

def LeanCert.ML.Optimized.reluIntMatrix {r c : ℕ} (M : IntMatrix r c) : IntMatrix r c

Element-wise ReLU on an IntMatrix: max(0, x)

Equations

• LeanCert.ML.Optimized.reluIntMatrix M = LeanCert.ML.Optimized.IntMatrix.ofFn fun (i : Fin r) (j : Fin c) => max 0 (M.get ↑i ↑j ⋯ ⋯)

def LeanCert.ML.Optimized.reluIntervalMatrix {r c : ℕ} (M : IntervalMatrix r c) : IntervalMatrix r c

Element-wise ReLU on IntervalMatrix. For interval [lo, hi]: ReLU([lo, hi]) = [max(0, lo), max(0, hi)]

Equations

• LeanCert.ML.Optimized.reluIntervalMatrix M = { lo := LeanCert.ML.Optimized.reluIntMatrix M.lo, hi := LeanCert.ML.Optimized.reluIntMatrix M.hi }

Bias Addition for Matrices

def LeanCert.ML.Optimized.addBiasMatrix {outDim batchSize : ℕ} (M : IntervalMatrix outDim batchSize) (bias_lo bias_hi : Fin outDim → ℤ) : IntervalMatrix outDim batchSize

Add bias vector to each column of a matrix. M : (out_dim, batch_size), bias : out_dim vector Result: M[i,j] + bias[i] for all i,j

Equations

• One or more equations did not get rendered due to their size.

Full Layer Forward (Matrix version)

def LeanCert.ML.Optimized.layerForwardBatch {inDim outDim batchSize : ℕ} (W : IntervalMatrix outDim inDim) (bias_lo bias_hi : Fin outDim → ℤ) (X : InputBatch inDim batchSize) : OutputBatch outDim batchSize

Forward pass through a layer with batch input. Computes ReLU(W × X + b) where X is a batch.

Equations

• LeanCert.ML.Optimized.layerForwardBatch W bias_lo bias_hi X = LeanCert.ML.Optimized.reluIntervalMatrix (LeanCert.ML.Optimized.addBiasMatrix (W.matmul X) bias_lo bias_hi)

Soundness (Placeholder)

theorem LeanCert.ML.Optimized.mem_batchLinear {outDim inDim batchSize : ℕ} (W : IntervalMatrix outDim inDim) (X : InputBatch inDim batchSize) (W_real : Matrix (Fin outDim) (Fin inDim) ℝ) (X_real : Matrix (Fin inDim) (Fin batchSize) ℝ) (exp : ℤ) (hW : W.mem exp W_real) (hX : IntervalMatrix.mem X exp X_real) : IntervalMatrix.mem (batchLinear W X) (exp + exp) (W_real * X_real)

Soundness of batch linear transformation

theorem LeanCert.ML.Optimized.mem_attentionScores {seqLen dK : ℕ} (Q K : IntervalMatrix seqLen dK) (Q_real K_real : Matrix (Fin seqLen) (Fin dK) ℝ) (exp : ℤ) (hQ : Q.mem exp Q_real) (hK : K.mem exp K_real) : (attentionScores Q K).mem (exp + exp) (Q_real * K_real.transpose)

Soundness of attention score computation