llm

Day 7/50: Building a Small Language Model from Scratch: Code Positional Embeddings

Yesterday, we explored what positional embeddings are and why they matter. Today, it’s time to dive into the code and see them in action!

We’ll use code from this excellent open-source model I’ve been experimenting with: Tiny Children Stories 30M. It’s a lightweight GPT-style model designed for generating children’s stories, and it beautifully demonstrates how positional embeddings are used in practice.

Quick Recap: Why Do Transformers Need Positional Embeddings?

Transformer models process input tokens in parallel, unlike RNNs, which process them sequentially. That’s great for speed, but it also means they don’t inherently understand the order of words.

To a plain transformer, these two sentences look the same:

• “The cat sat on the mat”
• “The mat sat on the cat”

Same tokens, different meaning, and no way to tell the difference unless we explicitly tell the model the position of each word. That’s the job of positional embeddings.

What Are Positional Embeddings?

Think of them as little “position tags” that say:

“Hey, this token is at position 0.”
“This one’s at position 1.”
“Here’s position 2…”

These tags are added to the token embeddings, enabling the model to understand what the token represents and its context.

Let’s Look at the Code

We’ll step through how positional embeddings are implemented in the Tiny GPT model.

1. Model Configuration

@dataclass
class GPTConfig:
    vocab_size: int = 50257
    block_size: int = 1024
    n_layer: int = 6
    n_head: int = 8
    n_embd: int = 512
    dropout: float = 0.1
    bias: bool = True

The key parameter here is block_size, which defines the maximum sequence length the model can handle and determines the number of positional embeddings required.

2. Defining Positional Embedding Layer

self.transformer = nn.ModuleDict(dict(
    wte=nn.Embedding(config.vocab_size, config.n_embd),  # Word/token embeddings
    wpe=nn.Embedding(config.block_size, config.n_embd),  # Positional embeddings
    ...
))

• wte = word token embeddings
• wpe = word position embeddings

These embeddings are of the same dimension (512), so we can add them together later.

3. Using Positional Embeddings in Forward Pass

pos = torch.arange(0, t, dtype=torch.long, device=device)
tok_emb = self.transformer.wte(idx)  # (batch_size, seq_len, embed_dim)
pos_emb = self.transformer.wpe(pos)  # (seq_len, embed_dim)
x = self.transformer.drop(tok_emb + pos_emb)

Step-by-step:

• Generate position indices: [0, 1, 2, ..., t-1]
• Look up the token embeddings (tok_emb)
• Look up the positional embeddings (pos_emb)
• Add them together and apply dropout

Real Example

Imagine our input is:

“The cat sat”
Token IDs: [464, 2368, 3290]
Sequence length: 3

Here’s what happens:

Token Token Embedding Positional Embedding Combined Embedding The [0.1, -0.3, 0.7, ...] [0.0, 0.1, -0.2, ...] [0.1, -0.2, 0.5, ...] cat [0.5, 0.2, -0.1, ...] [0.1, 0.0, 0.3, ...] [0.6, 0.2, 0.2, ...] sat [-0.2, 0.8, 0.4, ...] [0.2, -0.1, 0.1, ...] [0.0, 0.7, 0.5, ...]

Now, each token carries both its meaning and its position in the sequence.

Why This Works So Well

Adding positional embeddings to token embeddings gives the model:

• Semantic meaning (what the word is)
• Positional context (where the word is)

The model learns how positions matter for different patterns, particularly in text generation tasks such as storytelling.

Limitations of Learned Positional Embeddings

1. Fixed Length: You can’t go beyond the block_size (e.g., 1024 tokens)
2. No Relative Awareness: Doesn’t capture how far one token is from another
3. Sparse Coverage: If you never use long sequences in training, the model won’t learn to handle them well

Alternatives: Sinusoidal and Relative Embeddings

Sinusoidal Positional Embeddings

Instead of learning embeddings, you can generate them using sine and cosine functions:

def get_sinusoidal_embeddings(seq_len, embed_dim):
    pos = torch.arange(seq_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, embed_dim, 2) * -(math.log(10000.0) / embed_dim))
    pe = torch.zeros(seq_len, embed_dim)
    pe[:, 0::2] = torch.sin(pos * div_term)
    pe[:, 1::2] = torch.cos(pos * div_term)
    return pe

Pros:

• Infinite extendability
• No need to learn parameters

Relative Positional Embeddings

These indicate to the model how far two tokens are from each other, rather than simply stating “this is token 5.”

Useful in long documents and tasks where distance matters, like reasoning or QA.

Tips from Practice

• Don’t go overboard with block_size — larger = more memory
• Make sure your training data covers different sequence lengths
• If you’re dealing with long-form generation, consider RoPE or relative embeddings

Final Thoughts

Positional embeddings may seem like a small detail, but they’re the secret sauce that gives transformer models their power to handle sequence data.

Just by adding two vectors — token + position — we make it possible for a model to understand not only what it’s reading, but where everything fits.

In a storytelling model, this means knowing that “Once upon a time” belongs at the beginning, and “The End” belongs… well, at the end.

Pretty amazing.

Coming Up Next…

Tomorrow, we’ll look at Rotary Positional Embeddings (RoPE) — a more efficient and expressive way to encode positions using rotations in embedding space. Stay tuned!

Extras

• 📘 Reference Code: Tiny Children Stories 30M
• 🎥 YouTube Series: DevOps Boot Camp

If this post helped you, feel free to share it!
And if you’re following along the journey, let’s connect on LinkedIn.