Diffusion Models

in progressThis lesson is still being built - content, demos, and checkpoints may change before release.

Denoising diffusion

Loading interactive demo…

A cloud of points walks from pure noise back to a shape. At every timestep a small network guesses the noise hiding in each point, the point moves to the mean that guess implies, and a smaller fresh noise is added back, none on the last step. The network learned by guessing the noise in targets corrupted to random timesteps, in untraced updates up to the ladder's rung. Points that fly off the plot pin to its edge.Take the tour, or press and explore. Legend:

The panel beside the stage is the real DenoisingDiffusion class from the published library executing in your browser, with the current line highlighted.

The cloud as it stands; the loss over the untraced updates.
The pull of the guessed noise: a stroke from each point to the mean it would move to.
The cloud at its means, before the fresh noise.
The cloud just after the fresh noise; in the forward direction, the targets corrupted to the same timestep.
The shape's outline; the finished cloud, faint, where the walk will end.

timestep counts down the reverse walk, alpha bar is how much clean signal survives at it, noise scale the sigma of the fresh noise, distance to shape the mean distance from the cloud's points to their nearest targets. Drag the loss strip to change how much training came first; flip the direction to see the forward corruption at the same timestep.

The code panel shows bare executable code by default. The bar above it names the file (hover the name for the full published path); options holds the toggles: comments, the tracer instrumentation that produces this visualization, and verbose names for the terse identifiers (display only: the default is the byte-exact published file). The view switch adds plain-English narration: both interleaves it with the code, pseudo shows only the pseudocode. hide code collapses the panel; the rail it leaves behind brings it back.

Keys: Space play/pause · ←/→ step · Home/End jump · R reset · F fullscreen · G guided tour

The forward process. Each timestep adds a little Gaussian noise and fades what was there: $x_t = \sqrt{\alpha_t}\, x_{t-1} + \sqrt{\beta_t}\, \epsilon$ with $\alpha_t = 1 - \beta_t$ . Because Gaussians compose, any timestep is reachable in one draw from the clean point,

x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\,\epsilon, \qquad \bar\alpha_t = \prod_{s \le t} \alpha_s ,

which is what corrupt computes and the forward direction shows. The schedule is chosen so $\bar\alpha_T$ is near zero: the last timestep is indistinguishable from pure noise, so sampling can start there.

The objective. Rather than model the data directly, train a network $\epsilon_\theta(x_t, t)$ to guess the noise that was added:

\mathcal L = \mathbb E_{x_0,\, t,\, \epsilon} \big\lVert \epsilon - \epsilon_\theta(x_t, t) \big\rVert^2 .

A guessed noise is a guessed clean point, and a guessed clean point is a guessed direction back. This is the simple loss of the 2020 paper, a reweighting of the variational bound.

The reverse step. From $x_t$ the mean of $x_{t-1}$ given the guess is

\mu_\theta = \frac{1}{\sqrt{\alpha_t}} \Big( x_t - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}}\, \epsilon_\theta(x_t, t) \Big),

and the sample is $x_{t-1} = \mu_\theta + \sigma_t z$ with $\sigma_t^2 = \beta_t$ and $z$ fresh noise, except on the last step, where $z$ is dropped. Strip the noise's share, undo the fade, add a smaller noise back: the three traced events of every timestep.

Linear against cosine. Linear betas rise evenly, so most of the signal is gone before the chain is half done and the late timesteps teach the network little. The cosine schedule sets $\bar\alpha_t = \cos^2\!\big(\tfrac{t/T + s}{1 + s}\cdot\tfrac{\pi}{2}\big)$ directly, keeping signal through the middle, and derives the betas, clipping the last ones so no single step destroys everything. The readout of alpha-bar shows the difference at any timestep.

Cost. Training is one network pass per example per update, cheap; sampling is $T$ network passes per point, which is why a thousand timesteps took minutes per image and why the implicit sampler, learned variances, and latent spaces mattered. The stage keeps $T$ between ten and forty, its network at two layers of thirty-two, and its ladder at three thousand updates; a spiral is beyond a network that small, and the distance readout says so honestly.

Choices worth reading closely in the implementation (the code panel's comments and tracer checkboxes reveal the full published file):

The schedule's arrays are plural, betas, alphas, alphaBars, so at(t) can hand back one timestep's alphaBar by its singular name, and the sigma is zero at timestep zero by construction.
The timestep enters as a positional encoding, the transformer's sines and cosines from Kernels.positionalEncoding, so a network of plain relu layers can tell the timesteps apart.
corrupt is shared by the untraced fit and the forward view, so the training examples and the picture of the forward process are the same closed-form draw.
A clone restarts the random draws from the seed, keeping only the weights, so every rung of the trained ladder samples from the same starting noise and the slider compares like with like.
Fresh noise is emitted with its sigma, and the last step's sigma is zero, so the trace shows where the walk stops being random.

The idea came from physics before it came from pictures. A 2015 paper borrowed from non-equilibrium thermodynamics: destroy the structure of a data distribution slowly, by an iterative forward diffusion, then learn to run the destruction backward. It stayed a curiosity until 2019 and 2020, when score matching and then the denoising diffusion probabilistic model made the reverse walk a matter of guessing the noise, and the samples matched the best adversarial networks. Within two years the walk ran in the latent space of an autoencoder and drew pictures from text on a home graphics card.

2015Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli - 'Deep Unsupervised Learning using Nonequilibrium Thermodynamics': systematically and slowly destroy structure in a data distribution through an iterative forward diffusion process, then learn the reverse.source
2019Yang Song and Stefano Ermon - Perturb the data with different levels of Gaussian noise, estimate the score at every level, and sample by annealed Langevin dynamics over gradually decreasing noise levels.source
2020Jonathan Ho, Ajay Jain, and Pieter Abbeel - 'Denoising Diffusion Probabilistic Models' connects diffusion models with denoising score matching and reaches an FID of 3.17 on CIFAR-10, the walk this stage runs.source
2020Jiaming Song, Chenlin Meng, and Stefano Ermon - Denoising Diffusion Implicit Models sample ten to fifty times faster in wall-clock time than DDPMs, with the same training procedure.source
2021Alex Nichol and Prafulla Dhariwal - 'Improved Denoising Diffusion Probabilistic Models': learned variances allow sampling with an order of magnitude fewer forward passes, and log-likelihoods become competitive.source
2021Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjorn Ommer - Latent diffusion applies the walk in the latent space of a pretrained autoencoder, a near-optimal point between complexity reduction and detail preservation.source
2022CompVis at LMU Munich, Runway, and Stability AI - Stable Diffusion, a latent diffusion model, is released on August 22 with its code and weights public; an optimized version runs on a consumer GPU with 2.4 GB of memory.source

Used today

Text-to-image generation - Stable Diffusion and DALL-E 2 combine the reverse walk with text encoders and cross-attention to draw what a prompt describes.
Audio, video, and molecules - The same noise-guessing objective trains models over waveforms, frames, and atom positions; only the network's shape changes.
Image editing and restoration - Start the walk from a corrupted version of a real picture instead of pure noise and it fills in, sharpens, or restyles what is there.

timesteps20

schedule

shape

trained1200 updates

direction

timestep: walk done
alpha bar: 0.9900
noise scale: 0.000
distance to shape: 0.315

diffusion.ts

import { Mulberry32 } from "../random";import type { Tracer } from "../trace";import { Relu } from "./activation";import { Kernels } from "./kernels";export type ScheduleKind = "linear" | "cosine";/** The schedule's numbers at one timestep. */export interface ScheduleStep {  beta: number;  alpha: number;  alphaBar: number;  /** The fresh noise added after the mean step; zero on the last step. */  sigma: number;}/** One example of the noise-prediction task: a corrupted point, its timestep, and the noise that corrupted it. */export interface NoiseExample {  input: readonly number[];  t: number;  target: readonly number[];}export interface DiffusionOptions {  /** Timesteps of the forward chain, T. */  steps: number;  schedule: ScheduleKind;  /** Neurons in each of the predictor's two hidden layers. */  hidden: number;  learningRate: number;  seed?: number;}/** How many sinusoidal features the timestep is fed as. */export const TIME_FEATURES = 8;/** The offset of the cosine schedule, so the first step is not exactly noiseless. */const COSINE_OFFSET = 0.008;/** No step may destroy everything: the cosine schedule's last betas are clipped here. */const MAX_BETA = 0.999;/** * How much noise each timestep adds. `betas` are the per-step variances, * `alphas` what each step keeps, and `alphaBars` the running product: * how much of the clean signal survives to step t. Linear betas rise * evenly and are scaled by the number of steps so the chain ends near * pure noise however short it is; the cosine schedule sets the products * directly, along a cosine squared, and derives the betas. */export class NoiseSchedule {  readonly betas: number[];  readonly alphas: number[];  readonly alphaBars: number[];  constructor(    readonly kind: ScheduleKind,    readonly steps: number,  ) {    if (!Number.isInteger(steps) || steps < 2) throw new RangeError("a schedule needs at least two steps");    if (kind === "linear") {      // Betas from 0.2/T to 6/T: the sum stays about three, so the last product is under 0.05 for T in 10..40.      this.betas = Array.from({ length: steps }, (_, t) => Math.min(MAX_BETA, (0.2 + (5.8 * t) / (steps - 1)) / steps));    } else {      const f = (u: number) => Math.cos((((u / steps + COSINE_OFFSET) / (1 + COSINE_OFFSET)) * Math.PI) / 2) ** 2;      this.betas = Array.from({ length: steps }, (_, t) => Math.min(MAX_BETA, 1 - f(t + 1) / f(t)));    }    this.alphas = this.betas.map((b) => 1 - b);    this.alphaBars = [];    let product = 1;    for (const a of this.alphas) {      product *= a;      this.alphaBars.push(product);    }  }  at(t: number): ScheduleStep {    if (!Number.isInteger(t) || t < 0 || t >= this.steps) throw new RangeError(`timestep ${t} is outside 0..${this.steps - 1}`);    const beta = this.betas[t]!;    return { beta, alpha: this.alphas[t]!, alphaBar: this.alphaBars[t]!, sigma: t === 0 ? 0 : Math.sqrt(beta) };  }}interface Forward {  features: number[];  hidden: number[][];  outputs: number[];}/** * The network that guesses the noise in a corrupted point: the point's * coordinates and a sinusoidal code of the timestep go through two relu * layers to one linear output per coordinate. Trained by plain gradient * descent on the squared error against the noise that was actually * added, with the backward pass written out. */export class NoisePredictor {  private w: number[][][];  private b: number[][];  readonly features: number;  private readonly relu = new Relu();  constructor(    readonly inputs: number,    readonly hidden: number,    seed = 1,  ) {    if (!Number.isInteger(inputs) || inputs < 1) throw new RangeError("a point needs at least one coordinate");    if (!Number.isInteger(hidden) || hidden < 1) throw new RangeError("at least one hidden neuron");    this.features = inputs + TIME_FEATURES;    const random = new Mulberry32(seed);    const sizes = [this.features, hidden, hidden, inputs];    this.w = [];    this.b = [];    for (let l = 0; l < sizes.length - 1; l++) {      const from = sizes[l]!;      const into = sizes[l + 1]!;      const spread = Math.sqrt(6 / (from + into));      this.w.push(Array.from({ length: into }, () => Array.from({ length: from }, () => random.between(-spread, spread))));      this.b.push(Array.from({ length: into }, () => 0));    }  }  /** Takes another predictor's weights, for snapshots along a training run. */  copyFrom(other: NoisePredictor): void {    if (other.inputs !== this.inputs || other.hidden !== this.hidden) throw new RangeError("the predictors differ in shape");    this.w = other.w.map((layer) => layer.map((row) => [...row]));    this.b = other.b.map((row) => [...row]);  }  /** The coordinates followed by the timestep's sinusoidal code. */  encode(point: readonly number[], t: number): number[] {    return [...point, ...Kernels.positionalEncoding(t, TIME_FEATURES)];  }  private forward(point: readonly number[], t: number): Forward {    const features = this.encode(point, t);    const hidden: number[][] = [];    let a = features;    let outputs: number[] = [];    for (let l = 0; l < this.w.length; l++) {      const rows = this.w[l]!;      const sums: number[] = new Array<number>(rows.length);      for (let j = 0; j < rows.length; j++) {        const row = rows[j]!;        let sum = this.b[l]![j]!;        for (let i = 0; i < row.length; i++) sum += row[i]! * a[i]!;        sums[j] = sum;      }      if (l === this.w.length - 1) {        outputs = sums;      } else {        a = sums.map((v) => this.relu.apply(v));        hidden.push(a);      }    }    return { features, hidden, outputs };  }  /** The guessed noise in a point at timestep t. */  predict(point: readonly number[], t: number): number[] {    if (point.length !== this.inputs) throw new RangeError(`expected ${this.inputs} coordinates, got ${point.length}`);    return this.forward(point, t).outputs;  }  /** One gradient step on a batch; returns the mean squared error before the update. */  step(batch: readonly NoiseExample[], learningRate: number): number {    if (batch.length < 1) throw new RangeError("a step needs at least one example");    if (!(learningRate > 0)) throw new RangeError("the learning rate must be positive");    const gW = this.w.map((layer) => layer.map((row) => row.map(() => 0)));    const gB = this.b.map((row) => row.map(() => 0));    let loss = 0;    for (const example of batch) {      const f = this.forward(example.input, example.t);      // Blame at the output: the slope of the squared error, averaged over the batch.      let delta = f.outputs.map((v, k) => {        const error = v - example.target[k]!;        loss += (error * error) / batch.length;        return (2 * error) / batch.length;      });      for (let l = this.w.length - 1; l >= 0; l--) {        const below = l === 0 ? f.features : f.hidden[l - 1]!;        const rows = this.w[l]!;        const next: number[] = new Array<number>(below.length).fill(0);        for (let j = 0; j < rows.length; j++) {          const g = delta[j]!;          if (g === 0) continue;          const row = rows[j]!;          const gRow = gW[l]![j]!;          gB[l]![j] = gB[l]![j]! + g;          for (let i = 0; i < row.length; i++) {            gRow[i] = gRow[i]! + g * below[i]!;            next[i] = next[i]! + g * row[i]!;          }        }        // Through the relu: blame passes only where the neuron fired.        if (l > 0) delta = next.map((g, i) => g * this.relu.derivative(f.hidden[l - 1]![i]!));      }    }    this.w.forEach((layer, l) => layer.forEach((row, j) => row.forEach((_, i) => (row[i] = row[i]! - learningRate * gW[l]![j]![i]!))));    this.b.forEach((row, l) => row.forEach((_, j) => (row[j] = row[j]! - learningRate * gB[l]![j]!)));    return loss;  }}/** * A denoising diffusion model on points of the plane. The forward * process corrupts a clean point toward pure noise along the schedule * in one closed-form draw; training asks the predictor to guess the * noise in a corrupted point at a random timestep; sampling starts * from noise and walks the timesteps backward, at each one subtracting * the guessed noise to reach the mean and adding a smaller fresh noise, * none on the last step. `sample` traces one cloud's walk, one event * per timestep per stage; `fit` runs the training untraced. A clone * restarts the random draws from the seed, so a cloud sampled from * every snapshot of a run starts from the same noise. * * Time: O(count · T · hidden²) per sample. Space: O(weights + count). */export class DenoisingDiffusion {  readonly schedule: NoiseSchedule;  readonly predictor: NoisePredictor;  private readonly random: Mulberry32;  private spare: number | null = null;  constructor(    readonly options: DiffusionOptions,    private readonly tracer?: Tracer,  ) {    if (!(options.learningRate > 0)) throw new RangeError("the learning rate must be positive");    this.schedule = new NoiseSchedule(options.schedule, options.steps);    this.predictor = new NoisePredictor(2, options.hidden, options.seed ?? 1);    this.random = new Mulberry32((options.seed ?? 1) * 7919 + 13);  }  /** A standard normal draw by the Box-Muller transform, two at a time. */  private gaussian(): number {    if (this.spare !== null) {      const value = this.spare;      this.spare = null;      return value;    }    const u = 1 - this.random.next();    const v = this.random.next();    const radius = Math.sqrt(-2 * Math.log(u));    this.spare = radius * Math.sin(2 * Math.PI * v);    return radius * Math.cos(2 * Math.PI * v);  }  /** A copy with the same weights and the random draws restarted from the seed; a tracer can be attached to the copy. */  clone(tracer: Tracer | undefined = this.tracer): DenoisingDiffusion {    const copy = new DenoisingDiffusion(this.options, tracer);    copy.predictor.copyFrom(this.predictor);    return copy;  }  /** The forward process in one draw: the clean point faded by root alpha-bar plus noise scaled by root one minus alpha-bar. */  corrupt(x0: readonly number[], t: number): { noisy: number[]; noise: number[] } {    const s = this.schedule.at(t);    const noise = x0.map(() => this.gaussian());    const noisy = x0.map((v, k) => Math.sqrt(s.alphaBar) * v + Math.sqrt(1 - s.alphaBar) * noise[k]!);    return { noisy, noise };  }  /** Untraced training: each step corrupts a random batch of targets at random timesteps and moves the predictor; returns the loss per step. */  fit(targets: readonly (readonly number[])[], steps: number, batch: number): number[] {    if (targets.length < 1) throw new RangeError("training needs at least one target");    if (!Number.isInteger(steps) || steps < 0) throw new RangeError("steps must be a whole number");    if (!Number.isInteger(batch) || batch < 1) throw new RangeError("a batch needs at least one example");    const history: number[] = [];    for (let s = 0; s < steps; s++) {      const examples: NoiseExample[] = [];      for (let b = 0; b < batch; b++) {        const t = this.random.index(this.schedule.steps);        const { noisy, noise } = this.corrupt(targets[this.random.index(targets.length)]!, t);        examples.push({ input: noisy, t, target: noise });      }      history.push(this.predictor.step(examples, this.options.learningRate));    }    return history;  }  /** The traced reverse process: a cloud of `count` points walked from pure noise back to the shape. */  sample(count: number): number[][] {    if (!Number.isInteger(count) || count < 1) throw new RangeError("a sample needs at least one point");    let x = Array.from({ length: count }, () => [this.gaussian(), this.gaussian()]);    let cloud = Array.from({ length: count }, () => [this.gaussian(), this.gaussian()]);    for (let t = this.schedule.steps - 1; t >= 0; t--) {      const s = this.schedule.at(t);      // ▸ this timestep's numbers: how much signal is kept, how much noise comes back      this.tracer?.emit("schedule", { t, beta: s.beta, alpha: s.alpha, alphaBar: s.alphaBar, sigma: s.sigma });      const eps = x.map((p) => this.predictor.predict(p, t));      const guessedNoise = cloud.map((p) => this.predictor.predict(p, t));      const meanNorm = eps.reduce((sum, e) => sum + Math.hypot(e[0]!, e[1]!), 0) / count;      const meanNorm = guessedNoise.reduce((sum, e) => sum + Math.hypot(e[0]!, e[1]!), 0) / count;      // ▸ the network guesses the noise hiding in every point of the cloud      this.tracer?.emit("predict", { t, points: x.map((p) => [...p]), noise: eps.map((e) => [...e]), meanNorm });      this.tracer?.emit("predict", { t, points: cloud.map((p) => [...p]), noise: guessedNoise.map((e) => [...e]), meanNorm });      // The mean of the step back: strip the guessed noise's share and undo this step's fade.      const pull = s.beta / Math.sqrt(1 - s.alphaBar);      x = x.map((p, i) => p.map((v, k) => (v - pull * eps[i]![k]!) / Math.sqrt(s.alpha)));      cloud = cloud.map((p, i) => p.map((v, k) => (v - pull * guessedNoise[i]![k]!) / Math.sqrt(s.alpha)));      // ▸ strip the guessed noise's share and undo this step's fade      this.tracer?.emit("mean", { t, points: x.map((p) => [...p]) });      this.tracer?.emit("mean", { t, points: cloud.map((p) => [...p]) });      const z = x.map((p) => p.map(() => this.gaussian()));      const freshNoise = cloud.map((p) => p.map(() => this.gaussian()));      x = x.map((p, i) => p.map((v, k) => v + s.sigma * z[i]![k]!));      cloud = cloud.map((p, i) => p.map((v, k) => v + s.sigma * freshNoise[i]![k]!));      // ▸ add a smaller fresh noise, none at all on the last step      this.tracer?.emit("noise", { t, points: x.map((p) => [...p]), sigma: s.sigma });      this.tracer?.emit("noise", { t, points: cloud.map((p) => [...p]), sigma: s.sigma });    }    return x;    return cloud;  }}

distance =: 0.315

64 points end a mean distance 0.315 from the shape's nearest target.

import { Mulberry32 } from "../random";import type { Tracer } from "../trace";import { Relu } from "./activation";import { Kernels } from "./kernels";export type ScheduleKind = "linear" | "cosine";/** The schedule's numbers at one timestep. */export interface ScheduleStep { beta: number; alpha: number; alphaBar: number; /** The fresh noise added after the mean step; zero on the last step. */ sigma: number;}/** One example of the noise-prediction task: a corrupted point, its timestep, and the noise that corrupted it. */export interface NoiseExample { input: readonly number[]; t: number; target: readonly number[];}export interface DiffusionOptions { /** Timesteps of the forward chain, T. */ steps: number; schedule: ScheduleKind; /** Neurons in each of the predictor's two hidden layers. */ hidden: number; learningRate: number; seed?: number;}/** How many sinusoidal features the timestep is fed as. */export const TIME_FEATURES = 8;/** The offset of the cosine schedule, so the first step is not exactly noiseless. */const COSINE_OFFSET = 0.008;/** No step may destroy everything: the cosine schedule's last betas are clipped here. */const MAX_BETA = 0.999;/** * How much noise each timestep adds. `betas` are the per-step variances, * `alphas` what each step keeps, and `alphaBars` the running product: * how much of the clean signal survives to step t. Linear betas rise * evenly and are scaled by the number of steps so the chain ends near * pure noise however short it is; the cosine schedule sets the products * directly, along a cosine squared, and derives the betas. */export class NoiseSchedule { readonly betas: number[]; readonly alphas: number[]; readonly alphaBars: number[]; constructor( readonly kind: ScheduleKind, readonly steps: number, ) { if (!Number.isInteger(steps) || steps < 2) throw new RangeError("a schedule needs at least two steps"); if (kind === "linear") { // Betas from 0.2/T to 6/T: the sum stays about three, so the last product is under 0.05 for T in 10..40. this.betas = Array.from({ length: steps }, (_, t) => Math.min(MAX_BETA, (0.2 + (5.8 * t) / (steps - 1)) / steps)); } else { const f = (u: number) => Math.cos((((u / steps + COSINE_OFFSET) / (1 + COSINE_OFFSET)) * Math.PI) / 2) ** 2; this.betas = Array.from({ length: steps }, (_, t) => Math.min(MAX_BETA, 1 - f(t + 1) / f(t))); } this.alphas = this.betas.map((b) => 1 - b); this.alphaBars = []; let product = 1; for (const a of this.alphas) { product *= a; this.alphaBars.push(product); } } at(t: number): ScheduleStep { if (!Number.isInteger(t) || t < 0 || t >= this.steps) throw new RangeError(`timestep ${t} is outside 0..${this.steps - 1}`); const beta = this.betas[t]!; return { beta, alpha: this.alphas[t]!, alphaBar: this.alphaBars[t]!, sigma: t === 0 ? 0 : Math.sqrt(beta) }; }}interface Forward { features: number[]; hidden: number[][]; outputs: number[];}/** * The network that guesses the noise in a corrupted point: the point's * coordinates and a sinusoidal code of the timestep go through two relu * layers to one linear output per coordinate. Trained by plain gradient * descent on the squared error against the noise that was actually * added, with the backward pass written out. */export class NoisePredictor { private w: number[][][]; private b: number[][]; readonly features: number; private readonly relu = new Relu(); constructor( readonly inputs: number, readonly hidden: number, seed = 1, ) { if (!Number.isInteger(inputs) || inputs < 1) throw new RangeError("a point needs at least one coordinate"); if (!Number.isInteger(hidden) || hidden < 1) throw new RangeError("at least one hidden neuron"); this.features = inputs + TIME_FEATURES; const random = new Mulberry32(seed); const sizes = [this.features, hidden, hidden, inputs]; this.w = []; this.b = []; for (let l = 0; l < sizes.length - 1; l++) { const from = sizes[l]!; const into = sizes[l + 1]!; const spread = Math.sqrt(6 / (from + into)); this.w.push(Array.from({ length: into }, () => Array.from({ length: from }, () => random.between(-spread, spread)))); this.b.push(Array.from({ length: into }, () => 0)); } } /** Takes another predictor's weights, for snapshots along a training run. */ copyFrom(other: NoisePredictor): void { if (other.inputs !== this.inputs || other.hidden !== this.hidden) throw new RangeError("the predictors differ in shape"); this.w = other.w.map((layer) => layer.map((row) => [...row])); this.b = other.b.map((row) => [...row]); } /** The coordinates followed by the timestep's sinusoidal code. */ encode(point: readonly number[], t: number): number[] { return [...point, ...Kernels.positionalEncoding(t, TIME_FEATURES)]; } private forward(point: readonly number[], t: number): Forward { const features = this.encode(point, t); const hidden: number[][] = []; let a = features; let outputs: number[] = []; for (let l = 0; l < this.w.length; l++) { const rows = this.w[l]!; const sums: number[] = new Array<number>(rows.length); for (let j = 0; j < rows.length; j++) { const row = rows[j]!; let sum = this.b[l]![j]!; for (let i = 0; i < row.length; i++) sum += row[i]! * a[i]!; sums[j] = sum; } if (l === this.w.length - 1) { outputs = sums; } else { a = sums.map((v) => this.relu.apply(v)); hidden.push(a); } } return { features, hidden, outputs }; } /** The guessed noise in a point at timestep t. */ predict(point: readonly number[], t: number): number[] { if (point.length !== this.inputs) throw new RangeError(`expected ${this.inputs} coordinates, got ${point.length}`); return this.forward(point, t).outputs; } /** One gradient step on a batch; returns the mean squared error before the update. */ step(batch: readonly NoiseExample[], learningRate: number): number { if (batch.length < 1) throw new RangeError("a step needs at least one example"); if (!(learningRate > 0)) throw new RangeError("the learning rate must be positive"); const gW = this.w.map((layer) => layer.map((row) => row.map(() => 0))); const gB = this.b.map((row) => row.map(() => 0)); let loss = 0; for (const example of batch) { const f = this.forward(example.input, example.t); // Blame at the output: the slope of the squared error, averaged over the batch. let delta = f.outputs.map((v, k) => { const error = v - example.target[k]!; loss += (error * error) / batch.length; return (2 * error) / batch.length; }); for (let l = this.w.length - 1; l >= 0; l--) { const below = l === 0 ? f.features : f.hidden[l - 1]!; const rows = this.w[l]!; const next: number[] = new Array<number>(below.length).fill(0); for (let j = 0; j < rows.length; j++) { const g = delta[j]!; if (g === 0) continue; const row = rows[j]!; const gRow = gW[l]![j]!; gB[l]![j] = gB[l]![j]! + g; for (let i = 0; i < row.length; i++) { gRow[i] = gRow[i]! + g * below[i]!; next[i] = next[i]! + g * row[i]!; } } // Through the relu: blame passes only where the neuron fired. if (l > 0) delta = next.map((g, i) => g * this.relu.derivative(f.hidden[l - 1]![i]!)); } } this.w.forEach((layer, l) => layer.forEach((row, j) => row.forEach((_, i) => (row[i] = row[i]! - learningRate * gW[l]![j]![i]!)))); this.b.forEach((row, l) => row.forEach((_, j) => (row[j] = row[j]! - learningRate * gB[l]![j]!))); return loss; }}/** * A denoising diffusion model on points of the plane. The forward * process corrupts a clean point toward pure noise along the schedule * in one closed-form draw; training asks the predictor to guess the * noise in a corrupted point at a random timestep; sampling starts * from noise and walks the timesteps backward, at each one subtracting * the guessed noise to reach the mean and adding a smaller fresh noise, * none on the last step. `sample` traces one cloud's walk, one event * per timestep per stage; `fit` runs the training untraced. A clone * restarts the random draws from the seed, so a cloud sampled from * every snapshot of a run starts from the same noise. * * Time: O(count · T · hidden²) per sample. Space: O(weights + count). */export class DenoisingDiffusion { readonly schedule: NoiseSchedule; readonly predictor: NoisePredictor; private readonly random: Mulberry32; private spare: number | null = null; constructor( readonly options: DiffusionOptions, private readonly tracer?: Tracer, ) { if (!(options.learningRate > 0)) throw new RangeError("the learning rate must be positive"); this.schedule = new NoiseSchedule(options.schedule, options.steps); this.predictor = new NoisePredictor(2, options.hidden, options.seed ?? 1); this.random = new Mulberry32((options.seed ?? 1) * 7919 + 13); } /** A standard normal draw by the Box-Muller transform, two at a time. */ private gaussian(): number { if (this.spare !== null) { const value = this.spare; this.spare = null; return value; } const u = 1 - this.random.next(); const v = this.random.next(); const radius = Math.sqrt(-2 * Math.log(u)); this.spare = radius * Math.sin(2 * Math.PI * v); return radius * Math.cos(2 * Math.PI * v); } /** A copy with the same weights and the random draws restarted from the seed; a tracer can be attached to the copy. */ clone(tracer: Tracer | undefined = this.tracer): DenoisingDiffusion { const copy = new DenoisingDiffusion(this.options, tracer); copy.predictor.copyFrom(this.predictor); return copy; } /** The forward process in one draw: the clean point faded by root alpha-bar plus noise scaled by root one minus alpha-bar. */ corrupt(x0: readonly number[], t: number): { noisy: number[]; noise: number[] } { const s = this.schedule.at(t); const noise = x0.map(() => this.gaussian()); const noisy = x0.map((v, k) => Math.sqrt(s.alphaBar) * v + Math.sqrt(1 - s.alphaBar) * noise[k]!); return { noisy, noise }; } /** Untraced training: each step corrupts a random batch of targets at random timesteps and moves the predictor; returns the loss per step. */ fit(targets: readonly (readonly number[])[], steps: number, batch: number): number[] { if (targets.length < 1) throw new RangeError("training needs at least one target"); if (!Number.isInteger(steps) || steps < 0) throw new RangeError("steps must be a whole number"); if (!Number.isInteger(batch) || batch < 1) throw new RangeError("a batch needs at least one example"); const history: number[] = []; for (let s = 0; s < steps; s++) { const examples: NoiseExample[] = []; for (let b = 0; b < batch; b++) { const t = this.random.index(this.schedule.steps); const { noisy, noise } = this.corrupt(targets[this.random.index(targets.length)]!, t); examples.push({ input: noisy, t, target: noise }); } history.push(this.predictor.step(examples, this.options.learningRate)); } return history; } /** The traced reverse process: a cloud of `count` points walked from pure noise back to the shape. */ sample(count: number): number[][] { if (!Number.isInteger(count) || count < 1) throw new RangeError("a sample needs at least one point"); let x = Array.from({ length: count }, () => [this.gaussian(), this.gaussian()]); let cloud = Array.from({ length: count }, () => [this.gaussian(), this.gaussian()]); for (let t = this.schedule.steps - 1; t >= 0; t--) { const s = this.schedule.at(t); // ▸ this timestep's numbers: how much signal is kept, how much noise comes back this.tracer?.emit("schedule", { t, beta: s.beta, alpha: s.alpha, alphaBar: s.alphaBar, sigma: s.sigma }); const eps = x.map((p) => this.predictor.predict(p, t)); const guessedNoise = cloud.map((p) => this.predictor.predict(p, t)); const meanNorm = eps.reduce((sum, e) => sum + Math.hypot(e[0]!, e[1]!), 0) / count; const meanNorm = guessedNoise.reduce((sum, e) => sum + Math.hypot(e[0]!, e[1]!), 0) / count; // ▸ the network guesses the noise hiding in every point of the cloud this.tracer?.emit("predict", { t, points: x.map((p) => [...p]), noise: eps.map((e) => [...e]), meanNorm }); this.tracer?.emit("predict", { t, points: cloud.map((p) => [...p]), noise: guessedNoise.map((e) => [...e]), meanNorm }); // The mean of the step back: strip the guessed noise's share and undo this step's fade. const pull = s.beta / Math.sqrt(1 - s.alphaBar); x = x.map((p, i) => p.map((v, k) => (v - pull * eps[i]![k]!) / Math.sqrt(s.alpha))); cloud = cloud.map((p, i) => p.map((v, k) => (v - pull * guessedNoise[i]![k]!) / Math.sqrt(s.alpha))); // ▸ strip the guessed noise's share and undo this step's fade this.tracer?.emit("mean", { t, points: x.map((p) => [...p]) }); this.tracer?.emit("mean", { t, points: cloud.map((p) => [...p]) }); const z = x.map((p) => p.map(() => this.gaussian())); const freshNoise = cloud.map((p) => p.map(() => this.gaussian())); x = x.map((p, i) => p.map((v, k) => v + s.sigma * z[i]![k]!)); cloud = cloud.map((p, i) => p.map((v, k) => v + s.sigma * freshNoise[i]![k]!)); // ▸ add a smaller fresh noise, none at all on the last step this.tracer?.emit("noise", { t, points: x.map((p) => [...p]), sigma: s.sigma }); this.tracer?.emit("noise", { t, points: cloud.map((p) => [...p]), sigma: s.sigma }); } return x; return cloud; }}