What Happens During Training?

Chapter 3 examined what a trained music model does when it predicts tokens, denoises latents, follows conditions, and generates audio.

This chapter asks a different question: how did the model learn those behaviours?

Training is a repeated correction process. The model receives examples, produces predictions, measures its errors, calculates how its parameters contributed to those errors, and changes the parameters slightly.

One update usually changes the model only a little. Useful behaviour emerges after this loop is repeated across many batches.

Loading a batch

Training datasets can contain thousands or millions of examples. Processing the complete dataset before every update would be expensive, while updating from one example at a time may produce noisy and inefficient hardware use.

A batch groups several examples into one model computation. The loss is commonly averaged across valid examples or target positions, and the resulting gradients estimate how the parameters should change.

Shuffling examples without scrambling the music

Training examples are often shuffled between dataset passes so successive batches contain different combinations of examples.

Shuffling the dataset does not mean randomising the token order inside each song. The internal order of a sequence must remain intact for next-token training.

The distinction is important: randomise the order in which examples are presented, not the musical order within each example.

The forward pass

The batch is passed through the model using its current parameters.

For an autoregressive music model, the output may be logits for every possible codec token at each valid position. For a diffusion model, the output may be an estimate of the noise added at a chosen timestep.

This computation is called the forward pass because information moves from the model inputs toward its predictions.

Turning many errors into one objective

A model may make millions of predictions inside one batch. The loss function converts those prediction errors into a value that optimisation can minimise.

For next-token music generation, cross-entropy can compare the logits at each valid token position with the real token. For diffusion training, a regression loss may compare predicted noise with the noise that was actually added.

The loss is a mathematical training objective. It is not a complete measurement of whether generated music sounds good.

The backward pass

The forward pass created a chain of mathematical operations connecting the parameters to the loss.

Backpropagation follows that chain in reverse. Using the chain rule, it calculates how a small change in each trainable parameter would affect the loss.

Automatic differentiation systems such as PyTorch autograd record the operations used during the forward pass and use that record to calculate gradients during the backward pass.

Why gradients must be cleared

In PyTorch, gradients accumulate by default. Calling the backward pass again adds new gradient values to those already stored.

This behaviour supports intentional gradient accumulation, where several small batches contribute to one optimiser update. Without that intention, failing to clear gradients mixes unrelated training steps.

A typical loop therefore clears gradients before or after each optimiser update.

The optimiser decides the update

Gradients do not update the parameters by themselves. The optimiser converts them into parameter changes.

The simplest form of gradient descent moves parameters against the gradient. Practical optimisers may also use momentum, adaptive scaling, weight decay, or a learning-rate schedule.

The optimiser determines how the model responds to the same gradient information.

The learning rate controls update scale

The learning rate is one of the most important optimisation settings.

If it is too small, training may progress slowly or remain near the initial solution. If it is too large, updates may overshoot useful regions, cause unstable loss, or produce non-finite values.

Many training systems change the learning rate over time using warm-up, decay, or another schedule.

Limiting unusually large gradients

Some batches can produce very large gradient norms. Large updates may destabilise training, particularly in deep sequence models.

Gradient clipping limits the total norm or individual values before the optimiser step. It is a safety mechanism, not a substitute for choosing a suitable model, learning rate, and data pipeline.

Meta's AudioCraft MusicGen solver supports gradient-norm clipping before its optimiser update.

Mixed-precision training

Large audio models require substantial accelerator memory and computation. Mixed-precision training performs selected operations using lower-precision numerical formats while keeping sensitive operations at higher precision.

This can reduce memory use and improve throughput on suitable hardware. Gradient scaling may be used with float16 training to reduce the risk that very small gradients become zero through numerical underflow.

Steps, epochs, and tokens

A step usually refers to one batch iteration and optimiser update. An epoch usually means one pass through the training dataset.

If a dataset contains 10,000 examples and the batch size is 100, one epoch contains roughly 100 batch steps before accounting for distributed sampling or incomplete final batches.

For very large, repeated, or streaming datasets, researchers may report training progress in optimiser steps, examples, audio hours, or processed tokens rather than relying only on epochs.

Training loss is not enough

The optimiser directly reduces loss on training batches. A separate validation set helps estimate whether the learned behaviour transfers to examples not used for parameter updates.

During validation, the model makes predictions and calculates metrics without running the backward pass or changing parameters.

If training loss continues falling while validation performance worsens, the model may be fitting the training data more closely without improving general behaviour.

Saving progress

Training large models can take many updates. A checkpoint records enough state to evaluate the model or continue training later.

A resumable training checkpoint may contain model parameters, optimiser state, learning-rate scheduler state, gradient-scaler state, current step, and configuration information.

Saving only the model parameters may be enough for generation but not enough to resume the exact optimisation process.

Real-world example: one MusicGen training step

Meta's official AudioCraft implementation makes the abstract loop concrete.

For each batch, the training solver:

1. Receives audio and associated metadata.
2. Converts metadata into conditioning attributes.
3. Encodes audio into several streams of codec tokens.
4. Creates a mask for valid, non-padded token positions.
5. Runs the language model to produce logits.
6. Calculates cross-entropy separately for each codec codebook.
7. Averages the codebook losses.
8. Runs backpropagation.
9. Optionally unscales and clips gradients.
10. Runs the optimiser and learning-rate scheduler.
11. Clears the gradients.
12. Logs cross-entropy and token perplexity.

The model is not trained by listening to complete generated songs after every step. It learns from token-level prediction errors across the batch.

Where do the initial parameters come from?

The optimisation loop can begin with randomly initialised parameters or with parameters already trained on another dataset and task.

Starting from scratch gives full control but requires substantial data, compute, and experimentation. Fine-tuning begins with an existing model and adapts some or all of its parameters.

The next section, Training From Scratch Versus Fine-Tuning, compares these two starting points.