Combining Token, Position, and Noise-Level Embeddings in PyTorch

Welcome back. Last time, you established what the denoiser may use: visible tokens on both sides of a gap, but not the hidden answers or padding. Now you’ll build the representation it receives before attention begins. In about 15–20 minutes, you’ll complete a small PyTorch module that combines corrupted token IDs, sequence positions, and one noise level per sequence.

Three pieces of information at each position

Suppose the input is c[MASK]t. A token embedding tells the model what is present at each position, including the fact that the middle character is masked. A position embedding distinguishes the first t-like token from one at another location. A noise-level embedding tells the model which masking probability was sampled for this sequence.

For position $j$ in batch item $b$, we’ll add three vectors of the same width:

Here $x_t$ is the corrupted sequence, not the clean target. The function $g$ turns a scalar noise level into a vector. The same $g(t_b)$ is added at every position of sequence $b$, while token and position vectors can differ by position.

Read two short parts of Positional Encodings in Transformer Models from Machine Learning Mastery. We’ll use its simplest implementation choice—learned position embeddings—rather than its sinusoidal alternative.

We’ll use nn.Embedding for both token and position lookup. For noise, we’ll use nn.Linear(1, d_model): it accepts a floating-point scalar and learns to turn it into a vector. This is a deliberately small design for the Colab model, not a required architecture for every diffusion language model.

Complete the input module

The essential part of the module is its forward method. With embedding layers already defined, its unfinished outline looks like this:

def forward(self, x_t, t):
    batch_size, length = x_t.shape
    positions = ...                 # integer positions 0 through length - 1
    token_vectors = ...             # [batch_size, length, d_model]
    position_vectors = ...          # [length, d_model]
    noise_vectors = ...             # [batch_size, 1, d_model]
    return ...                      # [batch_size, length, d_model]

Before looking at the completion, notice the two different kinds of sharing. Every sequence in a batch uses the same position indices 0, 1, ...; each sequence has its own noise level. PyTorch broadcasting will let us add [length, d_model] position vectors and [batch_size, 1, d_model] noise vectors to [batch_size, length, d_model] token vectors.

Here is the complete, runnable module:

import torch
from torch import nn


class DenoiserInputs(nn.Module):
    def __init__(self, vocab_size, max_length, d_model, pad_id):
        super().__init__()
        self.max_length = max_length

        self.token_embed = nn.Embedding(
            vocab_size, d_model, padding_idx=pad_id
        )
        self.position_embed = nn.Embedding(max_length, d_model)
        self.noise_embed = nn.Linear(1, d_model)

    def forward(self, x_t, t):
        # x_t: integer token IDs [batch_size, length]
        # t: one floating-point noise level per sequence [batch_size]
        batch_size, length = x_t.shape

        if length > self.max_length:
            raise ValueError("Sequence exceeds max_length")

        positions = torch.arange(length, device=x_t.device)

        token_vectors = self.token_embed(x_t)
        position_vectors = self.position_embed(positions)

        # [batch_size] -> [batch_size, 1] -> [batch_size, d_model]
        t = t.reshape(batch_size, 1).to(
            device=x_t.device, dtype=token_vectors.dtype
        )
        noise_vectors = self.noise_embed(t).unsqueeze(1)

        return token_vectors + position_vectors + noise_vectors

nn.Embedding requires integer IDs, so x_t should be a torch.long tensor containing your vocabulary IDs—including [MASK] where corruption hid a token. Positions must also be integers; torch.arange(length, device=x_t.device) creates them on the input’s device. By contrast, t is a floating-point value such as 0.25, not a vocabulary ID. Reshaping it to [batch_size, 1] gives nn.Linear the one-feature input it expects.

The noise vector’s extra dimension, added by unsqueeze(1), is important. It makes one vector apply across all positions within its own sequence, without mixing noise levels between batch items. The position lookup has no explicit batch dimension; broadcasting reuses those position vectors across the batch.

There is no special embedding lookup for a clean token hidden behind [MASK]: supplying that clean ID here would give away the answer. Also, padding_idx=pad_id makes the token embedding contribution at [PAD] zero, but position and noise vectors are still added there. That does not make padding meaningful text; the encoder will need a separate padding attention mask, which is the next lesson’s concern.

Check the shapes in Colab

This small check uses the same corrupted sequence twice, but gives the two copies different noise levels. The vocabulary IDs are illustrative: 0 is [PAD] and 1 is [MASK].

torch.manual_seed(0)

model = DenoiserInputs(
    vocab_size=8, max_length=6, d_model=16, pad_id=0
)

x_t = torch.tensor(
    [[2, 1, 3, 0],
     [2, 1, 3, 0]],
    dtype=torch.long,
)
t = torch.tensor([0.25, 0.75])

h = model(x_t, t)

print(h.shape)  # torch.Size([2, 4, 16])
assert h.shape == (2, 4, 16)
assert not torch.allclose(h[0], h[1])

The output is one 16-number vector for each of four positions in each of two sequences. Since the corrupted IDs and positions match across the two rows, their differing outputs demonstrate that the supplied noise levels affect the representations. It does not mean the model has learned to denoise yet: the embedding and linear-layer weights have only been initialized.

What to carry forward

You now have a concrete input for a denoiser: a vector at each position formed from the corrupted token, its position, and the sequence’s noise level. Keeping track of [batch_size, length, d_model] is the key implementation habit; broadcasting handles the shared position and per-sequence noise contributions.

This module prepares representations but does not yet let positions exchange information. Next, you’ll feed them to a Transformer encoder that allows bidirectional attention while excluding [PAD] from attention.