Multi-Layer Perceptron (MLP) with GPU Backpropagation

Complete implementation of 2-layer MLP for MNIST with GPU-accelerated forward and backward passes.

Architecture

Input(784) → Dense(128) + ReLU → Dense(10) + Softmax → Output(10)

GPU Implementation

Forward Pass:

1. Layer 1: MatVec(784→128) + Bias + ReLU (fused kernel)
2. Layer 2: MatVec(128→10) + Bias (fused kernel)
3. Softmax(10)

Backward Pass:

1. Softmax gradient: dL/dlogits = probs - one_hot(label)
2. Layer 2 gradient: dW2 = h1^T @ dlogits, dB2 = dlogits, dH1 = W2^T @ dlogits
3. ReLU gradient: dH1pre = dH1 * (h1 > 0)
4. Layer 1 gradient: dW1 = input^T @ dH1pre, dB1 = dH1pre

SGD Update (GPU):

• W = W - lr * dW (element-wise on GPU)
• B = B - lr * dB (element-wise on GPU)

Type Definitions

structure Hesper.NN.MLP.MLPParams : Type

MLP parameters (weights and biases for 2 layers)

• w1 : Core.TensorData
• b1 : Core.TensorData
• w2 : Core.TensorData
• b2 : Core.TensorData

def Hesper.NN.MLP.instInhabitedMLPParams.default : MLPParams

Equations

• Hesper.NN.MLP.instInhabitedMLPParams.default = { w1 := default, b1 := default, w2 := default, b2 := default }

instance Hesper.NN.MLP.instInhabitedMLPParams : Inhabited MLPParams

Equations

• Hesper.NN.MLP.instInhabitedMLPParams = { default := Hesper.NN.MLP.instInhabitedMLPParams.default }

structure Hesper.NN.MLP.MLPInput : Type

MLP input (single sample + label for training)

• x : Core.TensorData
• label : Nat
• params : MLPParams

def Hesper.NN.MLP.instInhabitedMLPInput.default : MLPInput

Equations

• Hesper.NN.MLP.instInhabitedMLPInput.default = { x := default, label := default, params := default }

instance Hesper.NN.MLP.instInhabitedMLPInput : Inhabited MLPInput

Equations

• Hesper.NN.MLP.instInhabitedMLPInput = { default := Hesper.NN.MLP.instInhabitedMLPInput.default }

structure Hesper.NN.MLP.MLPOutput : Type

MLP output (predictions + loss)

• probs : Core.TensorData
• loss : Float

instance Hesper.NN.MLP.instInhabitedMLPOutput : Inhabited MLPOutput

Equations

• Hesper.NN.MLP.instInhabitedMLPOutput = { default := Hesper.NN.MLP.instInhabitedMLPOutput.default }

def Hesper.NN.MLP.instInhabitedMLPOutput.default : MLPOutput

Equations

• Hesper.NN.MLP.instInhabitedMLPOutput.default = { probs := default, loss := default }

structure Hesper.NN.MLP.MLPGradients : Type

MLP gradients (for backprop)

• dW1 : Core.TensorData
• dB1 : Core.TensorData
• dW2 : Core.TensorData
• dB2 : Core.TensorData

def Hesper.NN.MLP.instInhabitedMLPGradients.default : MLPGradients

Equations

• Hesper.NN.MLP.instInhabitedMLPGradients.default = { dW1 := default, dB1 := default, dW2 := default, dB2 := default }

instance Hesper.NN.MLP.instInhabitedMLPGradients : Inhabited MLPGradients

Equations

• Hesper.NN.MLP.instInhabitedMLPGradients = { default := Hesper.NN.MLP.instInhabitedMLPGradients.default }

CPU Reference Implementation

def Hesper.NN.MLP.cpuMatVec (mat vec : Array Float) (rows cols : Nat) : Array Float

Matrix-vector multiplication on CPU

Equations

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

def Hesper.NN.MLP.cpuAddBias (vec bias : Array Float) : Array Float

Add bias vector to result

Equations

• Hesper.NN.MLP.cpuAddBias vec bias = Array.zipWith (fun (x1 x2 : Float) => x1 + x2) vec bias

def Hesper.NN.MLP.cpuReLU (vec : Array Float) : Array Float

ReLU activation

Equations

• Hesper.NN.MLP.cpuReLU vec = Array.map (fun (x : Float) => max 0.0 x) vec

def Hesper.NN.MLP.cpuSoftmax (vec : Array Float) : Array Float

Softmax activation

Equations

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

def Hesper.NN.MLP.cpuCrossEntropyLoss (probs : Array Float) (label : Nat) : Float

Cross-entropy loss (approximated as 1 - prob for simplicity)

Equations

• Hesper.NN.MLP.cpuCrossEntropyLoss probs label = 1.0 - if probs[label]! < 1e-7 then 1e-7 else probs[label]!

def Hesper.NN.MLP.cpuMLPForward (input : MLPInput) : MLPOutput

CPU forward pass

Equations

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

def Hesper.NN.MLP.cpuMLPBackward (input : MLPInput) (output : MLPOutput) : MLPGradients

CPU backward pass (returns gradients)

Equations

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

GPU Kernels (using existing WGSL helpers)

def Hesper.NN.MLP.generateLayer1Kernel : String

Generate Layer 1 kernel (MatVec + Bias + ReLU)

Equations

• Hesper.NN.MLP.generateLayer1Kernel = Hesper.WGSL.Helpers.generateMatVecBiasActivationShader 784 128 Hesper.WGSL.Helpers.reluActivation

def Hesper.NN.MLP.generateLayer2Kernel : String

Generate Layer 2 kernel (MatVec + Bias)

Equations

• Hesper.NN.MLP.generateLayer2Kernel = Hesper.WGSL.Helpers.generateMatVecBiasActivationShader 128 10 Hesper.WGSL.Helpers.identityActivation

def Hesper.NN.MLP.generateSoftmaxKernel : String

Generate Softmax kernel

Equations

• Hesper.NN.MLP.generateSoftmaxKernel = Hesper.WGSL.Helpers.generateSoftmaxShader 10

GPU Backward Pass Kernels (using WGSL DSL)

def Hesper.NN.MLP.genSoftmaxGradKernel (outputSize : Nat) : String

Softmax gradient kernel: dLogits[i] = probs[i] - (i == label ? 1 : 0)

Equations

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

def Hesper.NN.MLP.genLayer2BackwardKernel (hiddenSize outputSize : Nat) : String

Layer 2 backward kernel: computes dW2, dB2, dH1

Equations

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

def Hesper.NN.MLP.genReLUBackwardKernel (size : Nat) : String

ReLU backward kernel: dH1pre[i] = dH1[i] * (h1[i] > 0)

Equations

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

def Hesper.NN.MLP.genLayer1BackwardKernel (inputSize hiddenSize : Nat) : String

Layer 1 backward kernel: computes dW1, dB1

Equations

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

SGD Update on GPU

def Hesper.NN.MLP.genSGDKernel (size : Nat) (lr : Float) : String

SGD kernel: param[i] -= lr * grad[i]

Equations

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

def Hesper.NN.MLP.genAdamKernel (size : Nat) (lr beta1 beta2 epsilon : Float) (step : Nat) : String

Generate GPU kernel for Adam optimizer update.

Adam (Adaptive Moment Estimation) combines momentum and RMSprop.

Algorithm:

m[i] = beta1 * m[i] + (1 - beta1) * grad[i]           # Update first moment
v[i] = beta2 * v[i] + (1 - beta2) * grad[i]^2         # Update second moment
m_hat = m[i] / (1 - beta1^step)                       # Bias correction
v_hat = v[i] / (1 - beta2^step)                       # Bias correction
params[i] -= lr * m_hat / (sqrt(v_hat) + epsilon)    # Update parameters

Parameters:

• size: Number of parameters
• lr: Learning rate (typical: 0.001)
• beta1: Momentum decay rate (typical: 0.9)
• beta2: Variance decay rate (typical: 0.999)
• epsilon: Numerical stability constant (typical: 1e-8)
• step: Current step number (for bias correction)

Buffers:

• @binding(0): params (read-write) - parameters to update
• @binding(1): grads (read-only) - gradients
• @binding(2): m (read-write) - first moment estimates
• @binding(3): v (read-write) - second moment estimates

Example:

let kernel := genAdamKernel 1000 0.001 0.9 0.999 1e-8 10

Equations

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

SGD Update on CPU

def Hesper.NN.MLP.cpuSGDUpdate (params : MLPParams) (grads : MLPGradients) (lr : Float) : MLPParams

CPU SGD parameter update

Equations

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

VerifiedOpFusion Instance

instance Hesper.NN.MLP.instVerifiedOpFusionOfNatNatMLPInputMLPOutputExpScalarF32 : Core.VerifiedOpFusion 256 1 1 MLPInput MLPOutput (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

Note: Full VerifiedOpFusion instance would require:

1. Kernel type signatures matching the pattern
2. WGSL expression builders
3. Proper workgroup dimensions

For now, we provide the CPU spec which can be called directly from the training loop. GPU kernels can be invoked via Compute API.

Equations

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

Training Helpers

def Hesper.NN.MLP.trainStep (input : MLPInput) (lr : Float) : MLPOutput × MLPParams

Single training step: forward + backward + update

Equations

• Hesper.NN.MLP.trainStep input lr = (Hesper.NN.MLP.cpuMLPForward input, Hesper.NN.MLP.cpuSGDUpdate input.params (Hesper.NN.MLP.cpuMLPBackward input (Hesper.NN.MLP.cpuMLPForward input)) lr)