Visual Guide

This section explains every implemented model family with diagrams: what goes in, what the network does to the tensor, what comes out, and how the training loss is computed. Most diagrams are animated or interactive. Hover over (or tap) a block to see the exact layer, and use the controls under a diagram to change its inputs.

Page	Model	Generates from
VAE	VAE1D / VAELightning	a Gaussian latent sequence z ~ N(0, I)
VQ-VAE + Prior	VQVAE1D + PixelCNNPrior	code indices sampled by an autoregressive prior
Pulse2Pulse GAN	WaveGANGenerator + WaveGANDiscriminator	uniform noise with the shape of an ECG
Diffusion (DDPM)	DiffusionModule + DDPMProcess + UNet1d	Gaussian noise (n, 12, 5000), denoised in ~50 DDIM steps
Flow Matching	FlowMatchingModule + RectifiedFlowProcess + UNet1d	Gaussian noise, moved along a learned velocity field with Euler / Heun steps

All shapes and hyper-parameters in the diagrams are the defaults in src/ecgen/families/ and src/ecgen/nn/. The written reference for each model is under Models (VAE, VQ-VAE, Pulse2Pulse, Diffusion and flow matching).

Legend

tensor / data trainable module frozen module loss term forward pass gradient skip connection sg[·] stop-gradient (.detach())

The input: a 12-lead ECG

Every model works on a tensor of shape (batch, leads, samples). MIMICIVECG yields 12 leads × 5000 samples (10 s at 500 Hz) by default. Pulse2Pulse uses 8 channels (num_channels=8, selected with leads: independent8: I, II, V1–V6).

12 leads — VAE, VQ-VAE8 leads — Pulse2Pulse

Schematic traces (not real patient data), shown on the standard 3 × 4 layout. Channel order in the tensor is I, II, III, aVR, aVL, aVF, V1–V6.

How each family generates an ECG

Training differs a lot between the three families, but at sampling time each one is a fixed pipeline from random numbers to a (n, leads, 5000) tensor:

Hover over a block for details.

Shared backbone: Encoder1D / Decoder1D

The VAE and the VQ-VAE use the same convolutional encoder and decoder from src/ecgen/nn/. They differ only in the latent bottleneck: the VAE samples a Gaussian latent with 8 channels, and the VQ-VAE snaps a 64-channel latent to a codebook. Switch between them below:

VAE bottleneck (latent_channels = 8)VQ-VAE bottleneck (latent_channels = 64)

Hover over a block for the exact layer. Thin blocks are the stride-2 down/up-sampling convolutions.

Every coloured block in the stages is a ResidualBlock1D:

Hover over a layer.

How long is the latent?

Only the first len(channel_multipliers) − 1 stages downsample, each by 2 (Conv1d(k=4, stride=2, padding=1)), and the decoder upsamples by the same factor. If the input length is not divisible by the factor the reconstruction comes out shorter than the input:

stages

seq_length 2500500040001000010005003

Use VQVAELightning.downsample_factor / latent_length() in code instead of hard-coding 8.

At a glance

	VAE	VQ-VAE + prior	Pulse2Pulse	Diffusion	Flow matching
Leads × samples	12 × 5000	12 × 5000	8 × 5000	12 × 5000	12 × 5000
Latent	continuous, (8, 625)	discrete, 625 indices into a 512 × 64 codebook	none: noise (8, 5000)	none: noise (12, 5000)	none: noise (12, 5000)
Training stages	1	2 (VQ-VAE, then prior on frozen codes)	1 (adversarial, 2 optimisers)	1	1
Loss	MSE + kl_weight · KL	MSE + VQ loss; then cross-entropy	WGAN-GP critic loss; −D(G(z))	min-SNR-weighted MSE on v	MSE on velocity x₁ − x₀
Sampling cost	1 decoder pass	625 sequential prior steps + 1 decoder pass	1 generator pass	50 UNet passes (DDIM; × 2 with guidance)	10–50 UNet passes (Euler / Heun)
Output range	unbounded	unbounded	tanh, [−1, 1]	unbounded	unbounded
Conditioning	—	—	—	vector, tokens, dense; CFG	vector, tokens, dense; CFG

Not covered yet

Latent diffusion and latent flow matching (src/ecgen/families/continuous/latent.py) run the same processes as the Diffusion and Flow Matching pages in the latent space of a frozen VAE or VQ-VAE; see Diffusion and flow matching. They do not have their own visual page yet.