Skip to content

VAE — Visual Guide

VAE1D compresses a 12-lead ECG into a short sequence of Gaussian latent vectors and reconstructs it. After training, decoding random draws from N(0, I) generates new ECGs. Code: src/ecgen/families/vae/, loss: src/ecgen/families/vae/losses.py. Written reference: VAE.

Forward pass during training

input x (B, 12, 5000) Encoder μ log σ² (B, 8, 625) z ε ~ N(0, I) Decoder reconstruction x̂ (B, 12, 5000) Time is compressed 8× (5000 → 625) and the 12 leads are mixed into 8 latent channels.
Hover over a block for details.

Inside the encoder and decoder

The VAE uses the shared backbone with latent_channels = 8. Hover over a block to see the layer:

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

See Shared backbone for the residual block and a latent-length calculator.

The reparameterisation trick

Sampling is not differentiable, so VAE1D.reparameterize moves the randomness into an input ε that needs no gradient:

μ log σ² exp(0.5·) σ ε ~ N(0, I) ⊙ + z ∂L/∂μ, ∂L/∂logvar no gradient
z = μ + exp(0.5·log σ²) ⊙ ε. Gradients (purple) reach μ and log σ², and through them the encoder.

How the loss is computed

total, recon, kl = vae_loss(recon, x, mean, logvar, kl_weight=1e-4)
\[ \mathcal{L} \;=\; \underbrace{\frac{1}{BCL}\sum (\hat{x}-x)^2}_{\texttt{recon\_loss}} \;+\; \texttt{kl\_weight}\cdot \underbrace{\frac{-\tfrac12\sum\left(1+\log\sigma^2-\mu^2-\sigma^2\right)}{B\cdot C\cdot L}}_{\texttt{kl\_loss}} \]
x x̂ μ log σ² MSE(x̂, x)mean over B·C·L −½ Σ(1 + log σ² − μ² − σ²)sum over latent ÷ (B·C·L) × kl_weight kl_loss recon_loss + ℒ
Hover over a loss term. Logged metrics: train/recon_loss, train/kl_loss, train/total_loss and the val/… equivalents.

What kl_weight does

The KL term pulls each ECG's posterior q(z | x) towards the prior N(0, I). The sketch below is schematic (one latent dimension, four ECGs); it is not computed from a trained model.

Small weight: posteriors are narrow and far apart, so reconstructions are sharp but random z from N(0, I) may land between them. Large weight: posteriors collapse onto the prior, so samples are smooth but blurry.

Sampling

VAELightning.sample(n_samples, seq_length=5000) skips the encoder:

z ~ N(0, I)(n, 8, seq_length // 8) Decoder
No-grad, eval mode. The latent length is computed from the encoder's downsampling factor.