Skip to content

Diffusion (DDPM) — Visual Guide

A diffusion model learns to undo noise. In training, a clean ECG is mixed with Gaussian noise at a random level t and the network learns to predict what was added (here: the velocity v). At sampling time it starts from pure noise and removes it in about 50 DDIM steps. Code: src/ecgen/families/continuous/module.py (DiffusionModule), process src/ecgen/processes/ddpm.py, guidance src/ecgen/processes/base.py, backbone src/ecgen/nn/backbones/unet1d.py, conditioning src/ecgen/conditioning/composite.py, EMA src/ecgen/training/ema.py. Written reference: Diffusion and flow matching.

One training step

clean x₀ (B, 12, 5000) noise ε N(0, I) + √ᾱ_t · x₀ √(1−ᾱ_t) · ε t ~ U{0…999} x_t UNet1d(x_t, t, cond) conditioner cond v̂ w(t) · MSEtarget v schematic traces · shapes for the default 12-lead, 5000-sample data
Hover over a block for details.

The forward (noising) process

DDPMProcess precomputes alphas_cumprod ᾱt for num_steps = 1000. Drag t to see how much of the ECG survives. The plot shows ᾱt for the default cosine schedule and the linear one, plus the min-SNR loss weight for v-prediction (γ = 5).

The ᾱt values are computed exactly as in _schedule(); the ECG is schematic. The linear schedule destroys the signal much earlier, so many of its 1000 steps are spent on near-pure noise.

Architecture: UNet1d

The default backbone is an ADM-style 1-D U-Net with base_channels = 64, channel_mults = [1, 2, 4, 4], num_res_blocks = 2 and attention at level 3. The three stride-2 convolutions take 5000 → 2500 → 1250 → 625, so the length must be divisible by 8.

Hover over a block for the exact layers. Every skip tensor is concatenated on the channel axis before the next up-path ResBlock.

Each ResBlock: where time and condition enter

h GN SiLU Conv k=3 AdaGNγ, β from emb SiLU Dropout Conv k=3zero init + out→ attn (lvl 3) shortcut: identity or 1×1 conv emb (B, 256)
Hover over a layer. The time step and the global condition vector act only through AdaGN.

Conditioning: three slots

CompositeConditioner turns each condition field into one of three slots. During training each field is replaced by a learned null embedding with its own probability (conditioning.dropout), so the same network also learns the unconditional model needed for guidance.

fieldencoderdropped? → nullslotused by UNet1d measurementsage, labels, … encoder + Linear→ 256 null_vectorlearned, zero init vector (B,256) AddToTimeEmb→ AdaGN report textas_tokens encoder + Linear→ (N, 256) null_token+ token mask tokens (B,N,256) cross-attentionlevel 3 + mid obs, obs_maskpartial ECG masked_signal[x·m, m] zeros“nothing observed” dense (B,24,5000) concat to x_tbefore conv_in training: field dropped with probability dropout[field] · sampling with guidance: drop="all" for the unconditional pass
Hover over a block. The ssm backbone accepts vector and dense slots only; unet1d and dit1d accept all three.

Other backbones

The process and the conditioning do not care which network predicts v. Two other backbones take the same forward(x, t, cond):

  • dit1d (src/ecgen/nn/backbones/dit1d.py): cuts the signal into patches of 25 samples (5000 → 200 tokens of width 256), runs 8 transformer blocks with AdaLN-Zero modulation from the time + condition embedding and cross-attention to condition tokens in every block, then un-patches with a zero-initialised Linear.
  • ssm (src/ecgen/nn/backbones/ssm.py): SSSD-style stack of 8 residual blocks with bidirectional S4D layers at full resolution (128 channels, no downsampling; the receptive field is the whole signal). mamba swaps S4D for bidirectional Mamba.

Classifier-free guidance

With a conditional model and guidance = w ≠ 1, every sampling step runs the backbone twice: once with the condition (c) and once with all fields replaced by their nulls (u, encode_condition(..., drop="all")). cfg_model_fn combines them as u + w·(c − u).

Schematic traces shown as predicted clean ECGs. Guidance is applied to the raw network output v, but x̂₀ = √ᾱ·xt − √(1−ᾱ)·v is affine in v, so the picture is the same. Default guidance = 1.0 (off).

How the loss is computed

# ContinuousGenerativeModule._loss (simplified: no masks)
x0 = self.to_model_space(signal)                       # identity for DiffusionModule
cond = self.adapt_condition(self.condition(batch), x0.shape[-1])  # with per-field dropout
pair = self.process.training_pair(x0, torch.randn_like(x0))      # x_t, t, target, weight
out = self.backbone(pair.xt, pair.t, cond)
per_sample = ((out - pair.target) ** 2).flatten(1).mean(1)
loss = (pair.weight * per_sample).mean()

With the default prediction = "v" and min_snr_gamma = 5:

\[ x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\varepsilon,\qquad v = \sqrt{\bar\alpha_t}\,\varepsilon - \sqrt{1-\bar\alpha_t}\,x_0,\qquad t \sim \mathcal{U}\{0,\dots,999\} \]
\[ \mathcal{L} = \frac{1}{B}\sum_{b} w(t_b)\,\operatorname{mean}\big[(\hat v_b - v_b)^2\big],\qquad w(t) = \frac{\min(\mathrm{SNR}_t,\ \gamma)}{\mathrm{SNR}_t + 1},\quad \mathrm{SNR}_t = \frac{\bar\alpha_t}{1-\bar\alpha_t} \]
∂ℒ → UNet1d and conditioner (the process has no parameters) x₀ ε ~ N(0, I) training_pair t ~ U{0, …, 999} x_t = √ᾱ x₀ + √(1−ᾱ) ε v = √ᾱ ε − √(1−ᾱ) x₀ SNR = ᾱ / (1 − ᾱ) w = min(SNR, 5)/(SNR+1) UNet1d(x_t, t, cond) conditioner v̂ target v MSE per sample(B,) weight w(t) (B,) × ℒ (mean) EMA purple = gradient path · orange = loss terms · the target and weight are constants (no gradient)
Hover over a block.

Partial-ECG options

With data.masks, loss_on: unobserved averages the error only where obs_mask = 0, and clean_observed: true (SSSD-style) overwrites the observed entries of xt with the clean signal before the backbone, in training and at every sampling step.

Sampling

sample() swaps in the EMA weights, encodes the condition (and its null version for guidance), then lets DDPMProcess.sample run DDIM from pure noise over steps evenly spaced time steps (default 50 of the 1000).

EMA weightsema.swap() encode c, and u (drop="all")only for conditional models x ~ N(0, I)(n, 12, 5000) for t, t′ in zip(timesteps, timesteps[1:] + [−1]): out = u + w·(c − u) # cfg_model_fn; just c when w = 1 x̂₀ = √ᾱ_t·x − √(1−ᾱ_t)·out ; ε̂ = √(1−ᾱ_t)·x + √ᾱ_t·out x = √ᾱ_t′·x̂₀ + √(1−ᾱ_t′−σ²)·ε̂ + σ·z, σ = η·√((1−ᾱ_t′)/(1−ᾱ_t)·(1−ᾱ_t/ᾱ_t′)) if inpaint: x = m·q_sample(obs, t′) + (1 − m)·x paste observedif obs given η = 0: deterministic DDIM (default) · η = 1: ancestral (DDPM-like) · x_init / start_step: SDEdit-style editing
Hover over a block.

The widget runs that loop on one lead. Change the number of steps, η, and the inpainting mode: replace overwrites the observed part (here the first 35 %) with the observation noised to the next level after every step; RePaint (inpaint="repaint", jumps=5) additionally diffuses back from t′ to t and redoes the step, 5 times per step, so the generated part can adjust to the observed part.

Toy: the “denoiser” is a hand-made function that knows the target ECG, so this shows the update rule, not model quality. Early steps predict a blurry average; detail appears as ᾱt → 1.
Option Where Effect
steps sample(steps=), default sample_steps = 50 network calls (× 2 with guidance, × jumps with RePaint)
eta sample(eta=), default 0 0 = deterministic DDIM, 1 = ancestral
guidance sample(guidance=), default 1.0 classifier-free guidance scale w
inpaint "replace", "repaint"; "clean" is automatic for clean_observed models training-free partial-ECG completion
use_ema default True sample with the EMA weights

Flow matching uses the same module, backbone and conditioning with a different process; see Flow matching.