TorchLean API

NN.Spec.Generative.Diffusion.Loss

Diffusion training losses (spec layer) #

This module defines the standard diffusion losses as named wrappers around existing spec losses (primarily MSE).

Why keep this separate from NN.Spec.Layers.Loss:

def Generative.Diffusion.epsPredLoss {α : Type} [Context α] {T : } {s : Spec.Shape} (sched : VPSchedule α T) (model : EpsModel α s) (x0 : Spec.Tensor α s) (t : Fin (T + 1)) (eps : Spec.Tensor α s) :
α

ε-prediction (DDPM-style) loss: MSE( ε̂_θ(x_t, t), ε ).

Inputs:

  • x0: clean sample,
  • t: discrete time index,
  • eps: explicit noise used to build x_t (intended as standard normal),
  • model: ε-prediction model.

This is the classic DDPM training objective (up to constant weighting choices).

Instances For