VAE-based augmentation pipeline for a small image dataset

What's actually being asked

A step-by-step GenAI design: train a VAE on the small dataset, sample latents, generate synthetic images, filter for quality, and use a downstream classifier to prove the synthetic data is actually useful.

Answer

Problem framing

Small datasets (say, a few hundred images per class in a medical-imaging or defect-detection task) make classifiers overfit. A VAE trained on the available images can generate additional, varied-but-plausible samples to pad out training data — cheaper than collecting more real data, provided the synthetic samples are validated properly. This last part is where most of the marks are; examiners like augmentation designs that don't just trust the generator blindly.

Step 1 — Train the VAE

Encoder (Conv2D stack) → μ(x), σ(x) → reparameterize z = μ + σ⊙ε, ε~N(0,I) → Decoder (transposed-Conv stack) → x̂. Train by minimizing L(θ,φ;x) = −𝔼_{q_φ(z|x)}[log p_θ(x|z)] + D_KL(q_φ(z|x) ‖ p(z)) over the small dataset, using the closed-form D_KL = −½ Σ_j(1+log σ_j² − μ_j² − σ_j²). With little data, use heavier regularization (dropout, weight decay, random crop/flip applied to the real images first) to stop the VAE itself from overfitting/memorizing — an overfit VAE will just reproduce near-exact copies of training images as "synthetic" data, adding no real diversity.

Step 2 — Latent sampling strategy

Two complementary sampling modes:

• Prior sampling: z ~ N(0, I) directly, decoded fresh — maximum diversity, but risks drifting into less class-faithful regions for a small, imperfectly-learned latent space.
• Perturbation sampling: encode a real training image to get its μ, then sample z = μ + σ⊙ε with the learned σ — stays closer to a known real example's identity/class while still varying pose/texture, better for targeted class-balanced augmentation.

Step 3 — Generate synthetic images

Decode the sampled z's in batches: x_synth = Decoder(z). Generate noticeably more than needed (e.g., 3–5× the final required count), since the filtering step will discard a chunk of them.

Step 4 — Quality filtering

Not every decoded image is usable. Filters, cheapest to most expensive:

• Reconstruction-error check: for perturbation-sampled images, compare x_synth to the source real image via pixel distance/SSIM — drop near-blank, overly blurred, or artifacted outputs.
• Classifier-confidence check: run a provisional classifier (even one trained so far on real data only) on each x_synth; keep only images it labels with reasonably high confidence as the intended class — this discards ambiguous or malformed samples that wouldn't teach the model anything correct.
• Diversity check: cluster/feature-compare kept samples (e.g., with a pretrained feature extractor) and drop near-duplicates, so augmentation doesn't just pad with many copies of the same mode.
• Manual spot-check: cheap for a small dataset, and catches failure modes automated filters miss.

Step 5 — Downstream evaluation

This is the step that actually proves usefulness:

1. Split real data into train/val/test, keeping val and test 100% real and untouched.
2. Train classifier A on real-train only; train classifier B on real-train + filtered synthetic images.
3. Compare A vs B on the same real val/test set — accuracy, F1, and per-class recall (especially for the minority class the augmentation targeted).
4. If B outperforms A, especially on the minority/weak class, the synthetic data earned its keep. If not — or if it only improves train accuracy while hurting val/test, a sign the classifier is learning VAE-specific artifacts rather than real features — tighten the quality filters or reduce the synthetic:real ratio and retest.

Pipeline summary: real images → train VAE → sample z (prior or perturbation) → decode → quality filter → augmented dataset → train classifier → compare vs real-only baseline.

Exam tip

Always close the loop back to a real, held-out test set — a common mistake is evaluating the classifier on a validation set that itself contains synthetic images, which inflates apparent performance without proving real-world usefulness.