Implementing Token Corruption with Noise Sampling and Eligibility Masking

Hello! Last time, you built a padded batch and a valid_mask: True means a position contains text, while False means padding. Now you’ll use that mask to make a corrupted copy of the batch, hiding some characters behind [MASK].

By the end of this 15–20 minute lesson, you’ll have a PyTorch function that samples a noise level for each text and returns two tensors: the corrupted token IDs and a Boolean mask recording exactly which positions it hid.

From a clean sequence to a corrupted one

The top row of this diagram shows the process we’re implementing. In the forward direction, increasingly many tokens become [MASK]. The bottom row shows a different kind of discrete corruption—replacing tokens with other tokens—which we won’t use in our small model.

For this lesson, use the simple masking schedule $m(t)=t$, where the noise level $t$ lies between zero and one. Conditional on $t$, each eligible token has probability $t$ of being masked. We’ll sample one $t$ per sequence, then make a separate random masking decision at each position. A larger $t$ therefore tends to hide more characters; it does not guarantee an exact count.

This short video segment demonstrates the useful PyTorch pattern of comparing random numbers with a probability, excluding ineligible positions, and replacing selected IDs:

There are two masks to keep distinct. valid_mask says which positions belong to text. The new corruption_mask says which of those positions this call actually replaced. A hidden character remains a valid text position.

Implement batch corruption

Run this cell after the previous lesson’s make_batch cell. It uses the existing PAD_ID-based valid_mask and MASK_ID.

import torch

def corrupt_batch(input_ids, valid_mask):
    """Mask clean text at a sampled noise level for each batch row."""
    batch_size = input_ids.shape[0]

    # One level per row; shape (B, 1) broadcasts across its positions.
    noise_levels = torch.rand(
        (batch_size, 1), device=input_ids.device
    )

    # Make an independent decision at each position.
    random_values = torch.rand(
        input_ids.shape, device=input_ids.device
    )
    eligible = valid_mask & (input_ids != MASK_ID)
    corruption_mask = (random_values < noise_levels) & eligible

    corrupted_ids = input_ids.clone()
    corrupted_ids[corruption_mask] = MASK_ID

    return corrupted_ids, corruption_mask

The shape (B, 1) matters: every position in a row is compared with that row’s noise level. The comparison produces a (B, L) Boolean tensor. Combining it with eligible prevents padding—and any already-present [MASK] tokens—from being selected.

clone() also matters. We need the original input_ids intact as the clean answer the denoiser will later learn to recover. The returned corruption_mask records new replacements made by this call, not every position that happens to contain MASK_ID.

Check the result in Colab

Use the batch and valid_mask from last lesson. A seed makes this particular demonstration reproducible; the function still samples fresh randomness on each ordinary call.

torch.manual_seed(7)

original = batch.clone()
corrupted_ids, corruption_mask = corrupt_batch(batch, valid_mask)

print("Clean row:    ", batch[1].tolist())
print("Corrupted row:", corrupted_ids[1].tolist())
print("Hidden per row:", corruption_mask.sum(dim=1).tolist())

assert corrupted_ids.shape == batch.shape
assert corruption_mask.shape == batch.shape
assert corruption_mask.dtype == torch.bool
assert torch.equal(batch, original)  # The clean input was not changed.
assert not corruption_mask[~valid_mask].any()  # No padding was selected.
assert torch.equal(
    corrupted_ids[~corruption_mask],
    batch[~corruption_mask],
)
assert (corrupted_ids[corruption_mask] == MASK_ID).all()

The number hidden in a short row can be surprising: sampling with probability $t$ is not a promise to hide exactly a fraction $t$. A row may even have zero selected positions. Don’t force an extra mask here, because that would change the sampling rule; we’ll handle empty selections safely when we build the training loss.

Takeaway

You now have a direct forward-corruption function for masked text diffusion. It samples a noise level per sequence, hides only eligible positions, preserves the clean batch, and reports exactly what it hid. Next, you’ll start building the denoiser that reads these corrupted sequences and predicts the missing tokens.