Training Generative Models

Entrenamiento de modelos generativos

Desarrollo y Evaluación de Modelos Propios

Especialización en IA Generativa y Desarrollo de Negocios

Today

Agenda · 3 hours

  1. Warm-up: three questions on Session 01 10'
  2. Theory: the training loop, generalization and data, from autoencoders to VAEs 60'
  3. Break 10'
  4. Guided lab on Kaggle: train a VAE on Fashion-MNIST 60'
  5. Session challenge (graded, individual) 40'

Same five blocks as every session. Today's lab runs best on a Kaggle GPU (CPU also works, just slower).

Warm-up

Three questions to reactivate Session 01: sampling, the adaptation spectrum and model families.

Warm-up · 10'
Warm-up · Q1

You lower the sampling temperature from 1.0 to 0.2. What happens to the next-token distribution?

Softmax divides the logits by \(T\) before normalising: with \(T < 1\) the differences between logits grow, so the top tokens take almost all the probability. Temperature changes decoding, never the weights.
Warm-up · Q2

Andes Bank's agents need answers from 2,000 pages of product manuals that change every month. Best position on the adaptation spectrum?

The knowledge lives in documents that change: retrieving them at question time keeps answers current without retraining. Fine-tuning is for behaviour, format and jargon, and would go stale every month.
Warm-up · Q3

Which model family compresses each example into a latent code and decodes new samples from that code space?

A VAE has an encoder (data → latent code) and a decoder (latent code → data). Sampling a new code and decoding it produces new data. You will train one today on images of clothes.

Learning objectives

By the end of this session you can…

  1. Write the training loop (forward → loss → backward → step) in PyTorch RAE 1
  2. Choose a loss and an optimizer, and diagnose learning-rate problems RAE 1
  3. Read training and validation curves to detect overfitting and decide when to stop RAE 2
  4. Spot preprocessing pitfalls and data leakage in business data RAE 2
  5. Explain how a VAE learns (ELBO, reparameterization, β) and train one RAE 1

RAE numbers refer to the learning outcomes in the course syllabus (guía de asignatura).

Part 1

The training loop

Losses, gradients, learning rates and optimizers: the engine under every model you will train.

Theory · 60'

What does “training” mean?

Search for weights that make the loss small

A model is a function \(f_\theta\) with parameters \(\theta\) (the weights). A loss measures how wrong it is on the training data:

$$ \mathcal{L}(\theta) = \frac{1}{N} \sum_{i=1}^{N} \ell\big(f_\theta(x_i),\, y_i\big) \qquad \text{training:}\ \ \theta^\star \approx \arg\min_\theta \mathcal{L}(\theta) $$

mini-batch → 1 · forward → 2 · loss → 3 · backward (gradients) → 4 · step (update θ) ↺

Generative twist

The “label” is often the data itself: the next token for a language model, the input image for an autoencoder. That is why unlabelled data is enough (self-supervision).

Vocabulary

Step (iteration): one mini-batch update. Epoch: one full pass over the training set. 20,000 images in batches of 128 ≈ 157 steps per epoch.

The loop in PyTorch

Five lines you will write in every project


model = VAE(latent_dim=8).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(EPOCHS):
    model.train()                        # training mode: dropout on
    for x, _ in train_loader:            # one mini-batch at a time
        x = x.to(DEVICE)
        optimizer.zero_grad()            # 1 · clear the previous gradients
        logits, mu, logvar = model(x)    # 2 · forward pass
        loss = vae_loss(logits, x, mu, logvar)   # 3 · loss on this batch
        loss.backward()                  # 4 · backward: gradients for every weight
        optimizer.step()                 # 5 · update the weights
    val_loss = evaluate(model, val_loader)       # 6 · check unseen data
    

evaluate() switches to model.eval() and torch.no_grad(): no dropout, no gradients. Forget zero_grad() and gradients from old batches pile up.

Loss functions

Three losses you will meet all course long

LossFormulaUse it when…In generative models
Cross-entropy\(-\sum_k y_k \log p_k\)The target is a category (a class, a token)Next-token prediction in every LLM; binary cross-entropy for pixels in [0, 1]
MSE\(\frac{1}{n}\sum_i (y_i - \hat{y}_i)^2\)The target is a continuous valueReconstructing real-valued data; diffusion models predict the added noise with MSE
KL divergence\(D_{\mathrm{KL}}(q \,\|\, p) = \sum_z q(z) \log \frac{q(z)}{p(z)}\)You compare two probability distributionsThe VAE regulariser; keeping a fine-tuned model close to the original one
KL is always ≥ 0 and equals 0 only when \(q = p\). It is not symmetric: \(D_{\mathrm{KL}}(q\|p) \neq D_{\mathrm{KL}}(p\|q)\) in general.
Key idea: minimising cross-entropy is the same as maximising the likelihood of the training data (maximum likelihood).
Check · losses

A language model is pretrained to predict the next token among 50,000 possible tokens. Which loss fits?

Token ids are categories, not quantities: token 812 is not “closer” to 813 than to 40,000. Cross-entropy on the softmax output is the standard choice. Accuracy is a useful metric, but it is not differentiable, so it cannot drive gradient descent.

Gradient descent

Walk downhill, one step at a time

$$ \theta_{t+1} = \theta_t - \eta \, \nabla_\theta \mathcal{L}(\theta_t) $$

The gradient points uphill; we step the other way. The learning rate \(\eta\) sets the step size.

🐢 Too small

Safe but slow: many epochs, more compute, and it may stop in the first valley it finds.

✅ About right

The loss falls quickly and then settles near a minimum.

💥 Too large

Steps overshoot the valley: the loss oscillates, or grows until it becomes inf or nan.

Stochastic gradient descent estimates the gradient on a mini-batch: noisier, but far cheaper than using all the data at every step.

Simulator · learning-rate playground

Find the learning rate that reaches the global minimum

After – steps: x = – · loss = –

Optimizers

From plain SGD to AdamW

OptimizerUpdate ideaIn practice
SGD\(\theta \leftarrow \theta - \eta\, g\)Simple and well understood; needs a well-tuned learning rate and schedule
SGD + momentum\(v \leftarrow \mu v + g,\ \ \theta \leftarrow \theta - \eta\, v\)A “heavy ball” (Polyak, 1964): smooths zig-zags and rolls past small bumps
AdamStep per weight from running averages of \(g\) and \(g^2\)Robust default; the paper's default rate is \(10^{-3}\) (Kingma & Ba, 2015)
AdamWAdam + weight decay applied directly to the weightsStandard for Transformers and fine-tuning (Loshchilov & Hutter, 2019)
\(g\) = gradient of the loss on the current mini-batch; \(\mu\) = momentum coefficient (often 0.9).

Adam adapts the step size of every weight separately, which is why it usually makes progress faster than SGD with the same learning rate. In the lab you will compare both for one epoch.

Learning-rate schedules and batch size

The learning rate does not have to stay fixed

  • Warm-up: start small for the first steps, while the optimizer's statistics are still unreliable.
  • Decay: big steps early to make progress, small steps late to settle. Cosine decay: Loshchilov & Hutter (2017).
  • Batch size: larger batches give smoother gradients and use the GPU better, but need more memory and give fewer steps per epoch. If you change it a lot, re-tune the learning rate.
Start from known defaults: Adam at \(10^{-3}\) for small models trained from scratch; smaller rates for fine-tuning (Session 04).

Cost of experimentation

Every training run has a price

96 runs

for a small grid: 4 learning rates × 4 batch sizes × 2 optimizers × 3 random seeds (illustrative arithmetic). At 2 GPU-hours per run, that is 192 GPU-hours before the first result reaches a customer.

Your stopping rule

What would convince you to stop paying for more tuning?

NovaTel (a fictional telecom) trains a model to draft replies to customers. The team asks for “one more week of tuning”, again and again.

  • Start from known defaults; run a pilot on a data subset.
  • Use early stopping and log every run (settings, curves, cost).
  • Agree on a stopping rule: the metric the business needs, and a budget.

Part 2

Generalization and data

Overfitting, regularization, early stopping, and the data mistakes that no optimizer can fix.

Generalization

Underfitting, good fit, overfitting

Underfitting

Training and validation loss both stay high. The model is too simple or trained too little.

Good fit

Both losses are low and close to each other; validation loss is flat near its minimum.

Overfitting

Training loss keeps falling while validation loss rises: the model memorises the training set.

SplitUsed to…Typical share
TrainFit the weights70–80%
ValidationChoose hyperparameters and when to stop10–15%
TestEstimate final performance, once, at the end10–15%
Shares are common conventions, not rules. If you tune on the test set, it stops being a test.

Simulator · overfitting explorer

How big should the model be?

Early stopping picks: –

Gap at epoch 30 (val − train): –

Synthetic curves precomputed in the browser to illustrate typical shapes; not real training runs.

Regularization

Tools to close the gap

Weight decay

Penalise large weights (L2). In AdamW, the weight_decay argument.

Dropout

Randomly switch off units during training only (Srivastava et al., 2014).

More or augmented data

Often the strongest regulariser, when you can get it.

Smaller model

Less capacity to memorise. Fewer parameters also run cheaper.

Early stopping

Watch validation loss; keep the best checkpoint; stop after patience epochs without improvement.


if val_loss < best - 1e-4:
    best, bad = val_loss, 0
    save_checkpoint(model)
else:
    bad += 1
if bad >= patience:
    break
        
Check · curves

Training loss keeps falling. Validation loss reached its minimum at epoch 6 and has risen ever since. You are at epoch 15. What do you do?

Rising validation loss with falling training loss is the signature of overfitting. Keep the best validation checkpoint. Option d turns the test set into a second validation set: its number is no longer an honest estimate.

Data preprocessing

Most training problems start in the data

1 · Clean

Duplicates, missing values, impossible values, inconsistent formats and units.

2 · Split

Before any fitting. By customer or by time when rows are related.

3 · Normalise

Fit scalers on train only, then apply them to validation and test. Pixels: divide by 255 → [0, 1].

4 · Encode

Text → tokens (S01); categories → ids or one-hot; images → tensors.

Why normalise? Features on very different scales (age vs. balance in pesos) make gradients unbalanced and training slow or unstable.
Order matters: computing a mean, a vocabulary or a scaler on all the data before splitting leaks test information into training.

Data leakage

When the model sees the answer in advance

Leakage: information is available during training that will not be available when the model is used (Kaufman et al., 2012).

TypeExample in a CRMFix
Target leakageA “cancellation call date” feature, recorded only after the customer asked to leaveKeep only fields that exist at prediction time
Train–test contaminationThe same customer appears in both sets; a scaler fitted on all rowsSplit by customer; fit preprocessing on train only
Temporal leakageA random split of monthly snapshots: the model trains on 2024 to “predict” 2023Split by time: train on the past, validate on the future
Symptom: validation results that look too good to be true, then collapse in production.

Apply it

Find the pitfalls in Andes Bank's churn dataset

Find 5 pitfalls · 6 min

A teammate prepared 250,000 rows (one per customer per month, Jan 2023–Dec 2024) to predict churn and to generate retention emails. Columns: customer_id, full_name, national_id, age, income, balance, complaint_text, retention_call_date, churned. Steps taken:

  1. Filled missing income with 0.
  2. Standardised every numeric column with the mean and std of all 250,000 rows.
  3. Split rows at random: 80% train, 20% test.
  4. Kept retention_call_date: “the most predictive feature!”
  5. Passed full_name, national_id and complaint_text to the email generator.
  6. Reported 97% test accuracy; 3% of customers churn.

For each pitfall: what goes wrong, and one fix. One room reports in 30 seconds.

Decisions built on leaks

A 97% model that catches nobody

Andes Bank (fictional) funds a retention campaign because the churn model “scored 97%”. In production it flags almost no one who actually leaves: the score came from leaked fields and a trivial baseline.

The cost is not the model: it is the campaign budget, the lost customers and the lost trust in data projects.

Questions a manager should ask

  • Which fields exist at the moment of prediction?
  • How was the split made: by customer? by time?
  • What does a trivial baseline score on the same data?
  • Was the test set used only once, at the end?

Part 3

From autoencoders to VAEs

A latent space you can sample from: the model you will train in the lab.

Autoencoders

Compress, then rebuild

x · 784 pixels → encoder → z · 8 numbers → decoder → x̂ · 784 pixels

What it learns

Trained to make \(\hat{x} \approx x\) through a narrow bottleneck, the network must keep only the essential features.

Uses: compression, denoising, anomaly detection (unusual inputs reconstruct badly).

Why it is not (yet) a good generator

Nothing organises the latent space: codes can sit anywhere, with gaps between them. Decode a random \(z\) and you often get noise.

Goodfellow et al. (2016), ch. 14.

Variational autoencoder

Encode to a distribution, not to a point

x → encoder → μ(x), σ(x) → z = μ + σ ⊙ ε → decoder → x̂

Encoder \(q_\phi(z \mid x)\)

Outputs a Gaussian \(\mathcal{N}(\mu(x), \operatorname{diag}\sigma^2(x))\) for each input: a small cloud in latent space.

Prior \(p(z)\)

A standard normal \(\mathcal{N}(0, I)\). Training pulls every cloud towards it, so the space has no gaps.

Decoder \(p_\theta(x \mid z)\)

Turns a latent code into data. To generate: sample \(z \sim \mathcal{N}(0, I)\) and decode it.

Kingma & Welling (2014); a longer tutorial is Kingma & Welling (2019).

The training objective

ELBO = reconstruction − KL

$$ \log p_\theta(x) \;\ge\; \underbrace{\mathbb{E}_{q_\phi(z\mid x)}\big[\log p_\theta(x \mid z)\big]}_{\text{reconstruction}} \;-\; \underbrace{D_{\mathrm{KL}}\big(q_\phi(z \mid x)\,\|\,p(z)\big)}_{\text{regulariser}} \;=\; \text{ELBO} $$

Loss = −ELBO = recon + KL

Reconstruction: binary cross-entropy summed over the 784 pixels (pixel values in [0, 1]).

KL for Gaussians, in closed form:

\( D_{\mathrm{KL}} = -\tfrac{1}{2}\sum_{j=1}^{d}\big(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2\big) \)

A tug of war

Reconstruction wants precise, well-separated codes. KL wants every code close to \(\mathcal{N}(0, I)\). Training finds a balance.

Sanity check: KL ≥ 0 always. A negative KL in your logs means a sign error.

The reparameterization trick

Move the randomness out of the way

We cannot back-propagate through “draw \(z\) at random from \(\mathcal{N}(\mu, \sigma^2)\)”. Rewrite the sample as a deterministic function plus independent noise:

$$ z = \mu + \sigma \odot \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I) $$

def reparameterize(mu, logvar):
    std = torch.exp(0.5 * logvar)  # σ
    eps = torch.randn_like(std)    # ε ~ N(0, I)
    return mu + std * eps          # z = μ + σ ⊙ ε
      

The randomness now lives in \(\varepsilon\); gradients flow through \(\mu\) and \(\sigma\) back to the encoder.

The encoder predicts \(\log \sigma^2\) (any real number) instead of \(\sigma\) (must be positive), and \(\sigma = \exp(\tfrac{1}{2}\log \sigma^2)\): simpler and numerically stable.

β-VAE

One knob for the trade-off

$$ \mathcal{L}_\beta = \text{recon} + \beta \cdot D_{\mathrm{KL}} \qquad \text{(Higgins et al., 2017)} $$
SettingReconstructionLatent spaceSamples from the prior
\(\beta < 1\)SharperLess regular, can have gapsMore often implausible
\(\beta = 1\)The standard VAE: the loss is exactly −ELBO
\(\beta > 1\)BlurrierSmoother; fewer active dimensionsSmoother, less detailed
Compare models with the same yardstick: report reconstruction and KL separately, plus recon + KL (−ELBO with β = 1). The β-weighted loss is not comparable across different β.

Self-check · define it before you flip it

VAE vocabulary

Reconstruction term
How well the decoder rebuilds \(x\) from \(z\); here binary cross-entropy summed over pixels.
KL term
How far each encoded Gaussian is from the prior \(\mathcal{N}(0, I)\). Never negative.
Reparameterization
\(z = \mu + \sigma \odot \varepsilon\): sampling rewritten so gradients reach the encoder.
Sampling from the prior
Draw \(z \sim \mathcal{N}(0, I)\) and decode it: a brand-new example.
Latent interpolation
Decode points on the line between two codes: one garment morphs into another.
Active dimension
A latent dimension whose mean changes with the input; inactive ones just copy the prior.
Check · VAEs

Your VAE reconstructs test images well, but images decoded from random \(z \sim \mathcal{N}(0, I)\) look like noise. Which change is most likely to help?

Good reconstructions with poor samples mean the codes the decoder learned sit far from the prior: the KL term is too weak. A larger β pulls the codes towards \(\mathcal{N}(0, I)\), where we sample, at some cost in sharpness. β = 0 is a plain autoencoder: the worst option for sampling.

Three families of image generators

VAE vs GAN vs diffusion

VAEGANDiffusion
TrainingMaximise the ELBO (recon + KL)Generator vs discriminator gameLearn to remove noise step by step (MSE on the noise)
StabilityStable, one lossCan be unstable; mode collapseStable
SamplesOften blurrySharpSharp and diverse
Sampling costOne decoder passOne generator passMany denoising passes (1,000 in Ho et al., 2020)
Latent spaceExplicit and smooth, with an encoderNo encoder by defaultNoise of the same size as the data
Kingma & Welling (2014) · Goodfellow et al. (2014) · Ho et al. (2020). Latent diffusion (Rombach et al., 2022) combines them: diffusion runs inside an autoencoder's latent space.

Choosing a generator

Product mock-ups for a fashion retailer

Hilo Norte (a fictional apparel retailer) wants quick mock-ups of new garments so designers can explore ideas and blend two best-sellers.

A small VAE on its own catalogue is cheap, fast to retrain, private and controllable (interpolate between two products). Photo-realistic images call for diffusion, with higher compute and licensing review.

Questions before choosing

  • Sketch-level or photo-level fidelity?
  • Who owns the training images, and may we use them?
  • How often must the model be retrained, and at what cost?
  • Who reviews outputs before they reach a customer?

Break

Stretch, refill your coffee, and switch on the GPU accelerator in Kaggle before the lab.

Hands-on

Guided lab

Train a variational autoencoder on Fashion-MNIST and make it draw new clothes.

Lab · 60'

Lab bridge

What we will build

1Fashion-MNIST → tensors in [0, 1]; train and validation splits10'
2An MLP VAE (latent 8), the reparameterization trick and the ELBO loss12'
3The training loop and its loss curves13'
4Reconstructions, samples from the prior and latent interpolation15'
5SGD vs Adam: one epoch each10'

Open the notebook

Kaggle checklist: Settings → Accelerator → GPU T4 x2 · Internet on. CPU also works, only slower.

Open in Kaggle Colab

Session challenge

Individual, graded, 40 minutes in class. Submit by 23:59 today.

Challenge de la sesión · evaluado

Challenge S02: tu propio VAE

  1. Ejecución: entrena un VAE con tu latent_dim, tu β, tu tasa de aprendizaje y tus 3 prendas asignadas; reporta reconstrucción y KL.
  2. Análisis: compáralo con la línea base (latente 8, β = 1) e interpreta el equilibrio reconstrucción–KL con tus números.
  3. Explica y decide: configura el generador de bocetos de producto de Hilo Norte.
  4. Crítica a la IA: encuentra y corrige 3 errores en un ciclo de entrenamiento escrito por una IA, y vuelve a entrenar.
  • Rúbrica: ejecución 40% · análisis 40% · reflexión y registro de IA 20%.
  • IA permitida con registro; micro-sustentaciones: si no explicas tu entrega, máximo 3.0.
  • Entrega: Archivo → Descargar notebook → e-Aulas como S02_<codigo>.ipynb, 23:59.

Wrap-up

Key takeaways

  • Every model trains with the same loop: forward, loss, backward, step.
  • The learning rate is the first knob to check when a loss curve looks wrong.
  • Judge a model on validation data; stop early, regularise, and hunt for leakage.
  • A VAE balances reconstruction against KL; β moves the balance.

Exit question

Which data from your organisation would you trust to train a model today, and which leakage risk worries you most?

Next session: Pretrained models and transfer learning, today at 10:00 after a short break. Final-project groups of 3 are formed in S03 (milestone H1). Challenge due by 23:59 today; the S03 pre-work was done before Saturday.

References

Sources for this session

  • Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep learning. MIT Press. Ch. 8 (optimization), 14 (autoencoders), 20 (generative models).
  • Kingma, D. P., & Welling, M. (2014). Auto-encoding variational Bayes. ICLR. arXiv:1312.6114.
  • Kingma, D. P., & Welling, M. (2019). An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 12(4). arXiv:1906.02691.
  • Higgins, I., et al. (2017). β-VAE: Learning basic visual concepts with a constrained variational framework. ICLR.
  • Kingma, D. P., & Ba, J. (2015). Adam: A method for stochastic optimization. ICLR. arXiv:1412.6980.
  • Loshchilov, I., & Hutter, F. (2019). Decoupled weight decay regularization. ICLR. arXiv:1711.05101.
  • Loshchilov, I., & Hutter, F. (2017). SGDR: Stochastic gradient descent with warm restarts. ICLR.
  • Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5).
  • Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., & Salakhutdinov, R. (2014). Dropout. JMLR, 15.
  • Kaufman, S., Rosset, S., Perlich, C., & Stitelman, O. (2012). Leakage in data mining: Formulation, detection, and avoidance. ACM TKDD, 6(4).
  • Goodfellow, I., et al. (2014). Generative adversarial nets. NeurIPS.
  • Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. NeurIPS.
  • Rombach, R., et al. (2022). High-resolution image synthesis with latent diffusion models. CVPR.
  • Xiao, H., Rasul, K., & Vollgraf, R. (2017). Fashion-MNIST: A novel image dataset for benchmarking machine learning algorithms. arXiv:1708.07747.
  • PyTorch. Learn the Basics (tutorial): pytorch.org/tutorials