Cross-Entropy Loss for Language Model Training

Implements numerically stable cross-entropy loss for teacher-forcing:

loss = -log(softmax(logits)[target])
     = -logits[target] + log(sum(exp(logits - max(logits))))

Backward:

dLogits[i] = softmax(logits)[i] - (i == target ? 1 : 0)

This elegant form means the backward pass is just softmax minus one-hot.

Forward: Cross-Entropy Loss

def Hesper.Training.Loss.crossEntropyForwardKernel (vocabSize : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

GPU kernel: Compute cross-entropy loss for a single token.

Uses two-pass approach for numerical stability:

1. Find max(logits) via parallel reduction
2. Compute log-sum-exp and loss

Input: logits [vocabSize], target [1] (u32 token ID) Output: loss [1] (scalar float)

Uses workgroup shared memory for reductions.

Equations

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

@[inline]

def Hesper.Training.Loss.executeCrossEntropyForward {β : Type} [GPUBackend β] (ctx : β) (logitsBuf targetBuf lossBuf : GPUBackend.Buf β) (vocabSize : Nat) : IO Unit

Execute cross-entropy loss forward. Returns loss value by reading back from GPU.

Equations

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

def Hesper.Training.Loss.crossEntropyForwardAccumKernel (vocabSize : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Cross-entropy forward with GPU-side loss accumulation. Adds the per-token loss to an accumulator buffer instead of overwriting. This allows batching all tokens' loss computation without CPU readback.

Equations

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

@[inline]

def Hesper.Training.Loss.executeCrossEntropyForwardAccum {β : Type} [GPUBackend β] (ctx : β) (logitsBuf targetBuf lossAccumBuf : GPUBackend.Buf β) (vocabSize : Nat) : IO Unit

Execute cross-entropy forward with GPU-side loss accumulation. Call this per-token; loss accumulates on GPU. Read once at end of example.

Equations

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

Backward: dLogits = softmax(logits) - one_hot(target)

def Hesper.Training.Loss.crossEntropyBackwardKernel (vocabSize : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

GPU kernel: Compute gradient of cross-entropy loss w.r.t. logits.

dLogits[i] = softmax(logits)[i] - (i == target ? 1 : 0)

Two phases:

1. Compute max and sum-exp (same as forward) via shared memory
2. Each thread computes its softmax value and subtracts one-hot

Input: logits [vocabSize], target [1] (u32) Output: dLogits [vocabSize]

Equations

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

@[inline]

def Hesper.Training.Loss.executeCrossEntropyBackward {β : Type} [GPUBackend β] (ctx : β) (logitsBuf targetBuf dLogitsBuf : GPUBackend.Buf β) (vocabSize : Nat) : IO Unit

Execute cross-entropy backward: dLogits = softmax(logits) - one_hot(target)

Equations

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