A convolutional variational autoencoder trained on PathMNIST (colorectal cancer histology, 9 classes, 28x28 RGB), with a short study of how the KL weight trades reconstruction quality against the structure of the latent space.
This is a small demonstrator built to work through generative modelling on biomedical images end to end. It is not a research contribution, and the limitations section says where it stops.
Requires uv.
uv run check_setup.py # verify environment, download data
uv run train_vae.py --epochs 30 --beta 0.1check_setup.py downloads PathMNIST into data/ (206 MB, gitignored) and writes
one real sample per class to results/real_samples.png.
infer.py runs the model on your own images, or samples from the prior. It does not
depend on the training dataset.
python infer.py --input tiles/ --out results/ # encode and reconstruct
python infer.py --sample 16 --out results/ # generate from N(0, I)Encoding writes one reconstruction per input image plus latents.csv, holding the
encoder's mean vector for each image: 32 numbers, no pixels. That distinction is the
point. A site can run the model locally and share latents for comparison while the
images themselves never leave the machine.
Reconstructions decode mu rather than a sampled z, so the same input always produces
the same output.
PNG, JPEG, TIFF and BMP, plus DICOM (.dcm) and NIfTI (.nii, .nii.gz).
Volumes and multi-frame series are reduced to their middle slice.
Medical formats do not arrive in the range the model expects. DICOM is commonly 12- or 16-bit and NIfTI is often float on an arbitrary scale, so each image is normalised by its own minimum and maximum. That is lossy on purpose: it discards absolute intensity, which for CT in Hounsfield units is clinically meaningful. A tool that needed those values would apply the rescale slope and intercept and a fixed window instead of a per-image stretch.
Worth stating plainly: reading a format is not the same as the model being meaningful on it. This model was trained on 28x28 RGB histology tiles. It will happily encode a CT slice and return a reconstruction, and that output means nothing diagnostically. The format support exists so the tool fits into a medical imaging pipeline, not because the weights generalise.
docker build -t pathmnist-vae .
docker run --rm -v "$PWD/tiles:/data:ro" -v "$PWD/out:/out" \
pathmnist-vae --input /data --out /outThe weights are baked into the image and nothing reaches the network at runtime, and it reads the same formats as the local tool including DICOM and NIfTI. The image is 1.37 GB, essentially all PyTorch; a smaller runtime would need ONNX or TorchScript rather than the full framework.
Encoding the same images on CPU inside the container and on Apple MPS on the host gives latents agreeing to five or six decimal places, which is the level of reproducibility to expect across hardware without pinning further.
PathMNIST, from MedMNIST v2. Hematoxylin & eosin stained colorectal tissue tiles, 28x28 RGB, nine classes: adipose, background, debris, lymphocytes, mucus, smooth muscle, normal colon mucosa, cancer-associated stroma, and colorectal adenocarcinoma epithelium.
| Split | Images |
|---|---|
| train | 89,996 |
| val | 10,004 |
| test | 7,180 |
Encoder: two stride-2 convolutions (3→32→64), giving a 64×7×7 feature map, then
two linear heads producing mu and logvar over a 32-dimensional latent.
Decoder: linear back to 64×7×7, two transposed convolutions, sigmoid output.
373k parameters.
Loss is per-image binary cross-entropy summed over pixels, plus beta times the
analytic Gaussian KL against N(0, I). Summing rather than averaging over pixels
keeps the reconstruction term on a comparable scale to the KL term; averaging
shrinks it by a factor of 2352 and the model collapses to predicting the dataset
mean.
Adam, lr 1e-3, batch size 128, 30 epochs, seed 0. Trains in about 2.5 minutes on an M-series GPU via MPS.
Raw BCE is close to uninterpretable on continuous data, and this turned out to matter for reading the results at all.
Binary cross-entropy has an irreducible floor when targets are not binary: even an exact reconstruction scores the entropy of the data itself, not zero. Measured on a fixed validation batch:
| BCE per image | |
|---|---|
| Perfect reconstruction (floor) | 1359.2 |
| Predicting the mean image for every input | 1480.5 |
So the entire learnable range is about 121 nats. A first training run reaching 1389 looked like a flat failure against an absolute scale, when in fact it had captured roughly three quarters of the achievable range. All reconstruction numbers below are therefore also reported as explained, the fraction of that range captured:
explained = (baseline - achieved) / (baseline - floor)
| beta | val BCE | KL (nats) | explained |
|---|---|---|---|
| 1.0 | 1388.80 | 8.79 | 75.6% |
| 0.3 | 1382.16 | 21.25 | 81.1% |
| 0.1 | 1379.03 | 39.76 | 83.7% |
| 0.03 | 1378.11 | 65.32 | 84.4% |
At beta = 1.0 the latent carries 8.8 nats across 32 dimensions, about a quarter
of a nat per dimension. The regularisation is strong enough that the model encodes
little beyond average colour and brightness. Lowering beta lets the latent carry
more information and reconstructions improve monotonically.
The returns flatten sharply, though. Going from beta = 0.1 to 0.03 more than
doubles the information in the latent, from 39.8 to 65.3 nats, and buys 0.7
percentage points of explained reconstruction. The extra latent capacity is not
being spent on anything the reconstruction metric can see, which is consistent
with the texture loss being a property of the pixel-wise loss rather than a
capacity limit.
Whether sample fidelity degrades at low beta, as the aggregate posterior drifts
away from the N(0, I) prior being sampled from, is not something this study can
answer. The samples at 0.03 and 0.1 are not visibly different to me, and
distinguishing them would need a perceptual metric that is not implemented here.
beta = 0.1 is used for the figures below as a reasonable point on the curve, not
as a tuned optimum.
Reconstructions at beta = 0.1 (top row real, bottom row reconstructed):
Coarse structure survives, including the branching adipose boundaries and the large pale region in the last tile. Fine cellular texture does not.
Samples decoded from z ~ N(0, I) at beta = 0.1:
These reproduce the H&E colour distribution and plausible coarse morphology, but no cell-level detail. A pathologist would not mistake them for tissue.
This is the model, not a training failure. The decoder has to produce a
low-error reconstruction across a neighbourhood of latents, because z is
resampled every forward pass, and the loss is computed per pixel. Where the
decoder is uncertain, the value minimising expected per-pixel error is the average
over plausible options, and averaging sharp texture produces smooth blur.
This is intrinsic to pixel-wise reconstruction with a stochastic latent, and it is the reason GANs and diffusion models produce sharper images at the same resolution.
--report adds a recon_mse column to latents.csv: the mean squared error between
each input and its reconstruction. The intent was a transfer check, so a site could flag
images unlike the data the model was trained on.
It does not do that, and the failure is worth recording.
Tested against four controlled probes, generated by make_probes.py and compared with
real PathMNIST tiles scoring 0.004 to 0.010:
| Probe | Texture | Tissue colour | recon_mse |
|---|---|---|---|
| checkerboard | high | no | 0.013595 |
| noise_pink | high | yes | 0.000326 |
| flat_pink | none | yes | 0.000182 |
| flat_white | none | no | 0.000162 |
flat_pink is the decisive case. It is exactly tissue-coloured and scores roughly thirty
times better than real tissue. Colour carries no weight. The checkerboard, the least
tissue-like image in the set, scores highest, purely because it has hard edges the model
cannot reproduce.
So the metric tracks how much high-frequency detail an image contains, not whether it belongs to the training distribution. A blank slide would score better than any real tissue tile, which is backwards for a quality check. That deep generative models assign favourable likelihoods or errors to out-of-distribution inputs is a known result; this is a small instance of it.
One caveat on the experiment itself. noise_pink did not test what it was meant to.
It is generated at 224x224 and resized to 28x28, which averages neighbouring pixels and
removes the noise almost entirely, so the model saw a nearly flat image. The checkerboard
survived only because its blocks alias into a visible pattern at 28x28. The conclusion
rests on flat_pink and checkerboard, which were not affected. Preprocessing quietly
destroying the property under test is a hazard worth naming.
The column is still worth having, described accurately: it identifies images the model reconstructs poorly. It is not a validity check, and using it as one would be worse than having no check at all.
- 28x28 is far below diagnostic resolution. Nothing here transfers directly to full-resolution histology without re-testing.
- The floor and baseline are computed on a single fixed validation batch, so the explained percentages carry some sampling noise.
- No perceptual metric. FID or a comparable feature-space distance would measure sample quality better than reconstruction BCE, which cannot see texture loss.
- Only the KL weight was varied. Latent dimension, depth, and training length were held fixed and not tuned.
- Single seed per configuration, so small differences between adjacent
betavalues should not be over-read.
The motivating question was whether VAE samples are useful as training-data augmentation: train a classifier on real data, then on real plus synthetic, and compare. Given the sample quality above I would expect no improvement, and possibly a small degradation, since the samples lack the texture a tissue classifier depends on. That experiment is not in this repository yet.
A genuine transfer check still needs building, since reconstruction error turned out not to be one. A feature-space distance against the training distribution would be the next thing to try.
- Yang et al., MedMNIST v2: A Large-Scale Lightweight Benchmark for 2D and 3D Biomedical Image Classification, Scientific Data, 2023.
- Kingma & Welling, Auto-Encoding Variational Bayes, ICLR 2014.
- Higgins et al., beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework, ICLR 2017.