Implementing an Optimizer Step for Denoising Training

Welcome back. Last lesson, you made the denoising loss grade only corrupted, non-padding positions. Now you can use that loss to change the denoiser’s parameters. In this 15–20 minute lesson, you’ll put one complete optimizer step into your Colab notebook and check that it really updates the model.

From a loss to an update

For one batch, the denoiser reads corrupted token IDs and produces vocabulary logits. Your loss compares those logits with the clean IDs at selected positions. PyTorch then computes how each model parameter contributed to that loss; the optimizer uses those gradients to adjust the parameters.

Watch the optimizer portion of PyTorch Tutorial 06 - Training Pipeline: Model, Loss, and Optimizer by Patrick Loeber. Its example uses a regression loss rather than our denoising loss, so focus on the calls that make training happen.

One terminology trap: loss.backward() is backpropagation through the model, not the diffusion model’s reverse-denoising operation. Backpropagation calculates gradients during training. Reverse denoising will later use the trained model to generate text.

Put the model and batch on the same device

Colab may provide a GPU, but the same code should work on CPU. The model and every tensor passed into its forward method must be on the same device. The clean targets and corruption mask must move there too, because the loss uses them.

For a short reference, read the accelerator section of What is torch.nn really? in the PyTorch documentation:

In your notebook, choose a device and move the denoiser before creating its optimizer:

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = model.to(device)  # your denoiser from the earlier lessons
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

print("Training on:", device)

Here, model is your existing denoiser. The optimizer receives model.parameters(), so it knows which tensors it is allowed to update.

Implement one denoising step

The function below expects a batch you have already corrupted. x0 holds clean IDs, x_t holds the corresponding corrupted IDs, corruption_mask identifies replaced positions, and noise_level is the level used to produce x_t. Keeping corruption outside this function makes the update itself easier to inspect.

It also assumes your model’s forward call has the form model(x_t, noise_level, valid_mask), where valid_mask is True at non-padding positions. If your earlier model takes these inputs in a different order—or expects a padding mask that is True at padding—adapt that one forward-call line.

def denoising_train_step(
    model, optimizer, x0, x_t, noise_level,
    corruption_mask, pad_id, device
):
    model.train()

    # Inputs, targets, and masks used together must share a device.
    x0 = x0.to(device)
    x_t = x_t.to(device)
    noise_level = noise_level.to(device)
    corruption_mask = corruption_mask.to(device).bool()

    valid_mask = x0.ne(pad_id)
    selected_count = (corruption_mask & valid_mask).sum().item()

    # Clear gradients left by any earlier batch.
    optimizer.zero_grad(set_to_none=True)

    # No supervised predictions means no useful update. Skipping step()
    # also avoids an optimizer changing parameters through weight decay.
    if selected_count == 0:
        return None

    # Keep gradient tracking on: do not use torch.no_grad() here.
    logits = model(x_t, noise_level, valid_mask)
    loss = denoising_loss(logits, x0, corruption_mask, pad_id)

    loss.backward()   # calculate parameter gradients
    optimizer.step()   # update parameters
    return loss.item()

denoising_loss is the safely normalized function from last lesson. It still decides which predictions count; valid_mask in the model call instead controls how padding is treated as input context. Neither replaces the other.

The order is worth remembering: clear old gradients, compute fresh logits and loss, backpropagate, then update. optimizer.step() cannot calculate gradients for you, and loss.backward() does not update parameters by itself.

Check the step in your notebook

Use one padded batch and its matching corrupted version from your earlier cells. In particular, noise_level must be the level used for that corruption, not a newly sampled value. Then compare parameter snapshots:

before = [p.detach().clone() for p in model.parameters()]

loss_value = denoising_train_step(
    model=model,
    optimizer=optimizer,
    x0=x0,
    x_t=x_t,
    noise_level=noise_level,
    corruption_mask=corruption_mask,
    pad_id=pad_id,
    device=device,
)

if loss_value is None:
    print("No valid tokens were masked; this batch was skipped.")
else:
    changed = any(
        not torch.equal(old, new.detach())
        for old, new in zip(before, model.parameters())
    )
    print("loss:", loss_value, "parameters changed:", changed)
    assert changed

If the batch is skipped, check how many eligible positions your corruption function selected. If you get a device error, check all four batch tensors as well as the model. A parameter change confirms that an update occurred; it does not require the loss on a single random batch to fall immediately.

What to carry forward

You now have the smallest complete training unit: a corrupted batch, a forward pass, the masked denoising loss, backpropagation, and an optimizer update. Device placement keeps its tensors compatible, while the empty-mask check avoids a step with no training signal.

Next, you’ll repeat this step on one tiny batch to see whether the denoiser can deliberately overfit it—a useful test before training on a larger corpus.