TinyTorch/modules/source/11_embeddings/embeddings_dev.ipynb

Module 11: Embeddings - Converting Tokens to Learnable Representations

Welcome to Module 11! You're about to build embedding layers that convert discrete tokens into dense, learnable vectors - the foundation of all modern NLP models.

🔗 Prerequisites & Progress

You've Built: Tensors, layers, tokenization (discrete text processing) You'll Build: Embedding lookups and positional encodings for sequence modeling You'll Enable: Foundation for attention mechanisms and transformer architectures

Connection Map:

Tokenization → Embeddings → Positional Encoding → Attention (Module 12)
(discrete)     (dense)      (position-aware)     (context-aware)

Learning Objectives

By the end of this module, you will:

1. Implement embedding layers for token-to-vector conversion
2. Understand learnable vs fixed positional encodings
3. Build both sinusoidal and learned position encodings
4. Analyze embedding memory requirements and lookup performance

Let's transform tokens into intelligence!

📦 Where This Code Lives in the Final Package

Learning Side: You work in modules/11_embeddings/embeddings_dev.py
Building Side: Code exports to tinytorch.text.embeddings

# How to use this module:
from tinytorch.text.embeddings import Embedding, PositionalEncoding, create_sinusoidal_embeddings

Why this matters:

• Learning: Complete embedding system for converting discrete tokens to continuous representations
• Production: Essential component matching PyTorch's torch.nn.Embedding with positional encoding patterns
• Consistency: All embedding operations and positional encodings in text.embeddings
• Integration: Works seamlessly with tokenizers for complete text processing pipeline

In [ ]:

"""
## 1. Essential Imports and Setup

Setting up our embedding toolkit with tensor operations and mathematical functions.
"""

#| default_exp text.embeddings
#| export

import numpy as np
import math
from typing import List, Optional, Tuple

# Core tensor operations - our foundation
### BEGIN SOLUTION
# For this educational implementation, we'll create a simple Tensor class
# In practice, this would import from tinytorch.core.tensor

class Tensor:
    """Educational tensor for embeddings module."""

    def __init__(self, data, requires_grad=False):
        self.data = np.array(data)
        self.shape = self.data.shape
        self.requires_grad = requires_grad
        self.grad = None

    def __repr__(self):
        return f"Tensor({self.data})"

    def __getitem__(self, idx):
        return Tensor(self.data[idx])

    def __add__(self, other):
        if isinstance(other, Tensor):
            return Tensor(self.data + other.data)
        return Tensor(self.data + other)

    def size(self, dim=None):
        if dim is None:
            return self.shape
        return self.shape[dim]

    def reshape(self, *shape):
        return Tensor(self.data.reshape(shape))

    def expand(self, *shape):
        return Tensor(np.broadcast_to(self.data, shape))

    def parameters(self):
        return [self] if self.requires_grad else []

# Simple Linear layer for this module
class Linear:
    """Educational linear layer."""

    def __init__(self, in_features, out_features, bias=True):
        # Xavier initialization
        limit = math.sqrt(6.0 / (in_features + out_features))
        self.weight = Tensor(
            np.random.uniform(-limit, limit, (in_features, out_features)),
            requires_grad=True
        )
        self.bias = Tensor(np.zeros(out_features), requires_grad=True) if bias else None

    def forward(self, x):
        result = Tensor(np.dot(x.data, self.weight.data))
        if self.bias is not None:
            result = result + self.bias
        return result

    def parameters(self):
        params = [self.weight]
        if self.bias is not None:
            params.append(self.bias)
        return params
### END SOLUTION

2. Understanding Token Embeddings - From Discrete to Dense

Before we implement embeddings, let's understand what problem they solve and how the lookup process works.

The Fundamental Challenge

When dealing with text, we start with discrete symbols (words, characters, tokens) but neural networks need continuous numbers. Embeddings bridge this gap by creating a learned mapping from discrete tokens to dense vector representations.

Token-to-Vector Transformation Visualization

Traditional One-Hot Encoding (Sparse):
Token "cat" (index 42) → [0, 0, ..., 1, ..., 0]  (50,000 elements, mostly zeros)
                           position 42

Modern Embedding Lookup (Dense):
Token "cat" (index 42) → [0.1, -0.3, 0.7, 0.2, ...]  (512 dense, meaningful values)

How Embedding Lookup Works

Embedding Table (vocab_size × embed_dim):
    Token ID    Embedding Vector
    ┌─────┐    ┌─────────────────────────┐
 0  │  0  │ →  │ [0.2, -0.1,  0.3, ...] │  "the"
 1  │  1  │ →  │ [0.1,  0.4, -0.2, ...] │  "cat"
 2  │  2  │ →  │ [-0.3, 0.1,  0.5, ...] │  "sat"
... │ ... │    │        ...              │   ...
42  │ 42 │ →  │ [0.7, -0.2,  0.1, ...] │  "dog"
... │ ... │    │        ...              │   ...
    └─────┘    └─────────────────────────┘

Lookup Process:
Input tokens: [1, 2, 42] → Output: Matrix (3 × embed_dim)
Row 0: embedding[1]  → [0.1,  0.4, -0.2, ...]  "cat"
Row 1: embedding[2]  → [-0.3, 0.1,  0.5, ...]  "sat"
Row 2: embedding[42] → [0.7, -0.2,  0.1, ...]  "dog"

Why Embeddings Are Powerful

1. Dense Representation: Every dimension can contribute meaningful information
2. Learnable: Vectors adjust during training to capture semantic relationships
3. Efficient: O(1) lookup time regardless of vocabulary size
4. Semantic: Similar words learn similar vector representations

Memory Implications

For a vocabulary of 50,000 tokens with 512-dimensional embeddings:

• Storage: 50,000 × 512 × 4 bytes = ~100MB (in FP32)
• Scaling: Memory grows linearly with vocab_size × embed_dim
• Trade-off: Larger embeddings capture more nuance but require more memory

This is why embedding tables often dominate memory usage in large language models!

3. Implementing Token Embeddings

Now let's build the core embedding layer that performs efficient token-to-vector lookups.

In [ ]:

class Embedding:
    """
    Learnable embedding layer that maps token indices to dense vectors.

    This is the fundamental building block for converting discrete tokens
    into continuous representations that neural networks can process.

    TODO: Implement the Embedding class

    APPROACH:
    1. Initialize embedding matrix with random weights (vocab_size, embed_dim)
    2. Implement forward pass as matrix lookup using numpy indexing
    3. Handle batch dimensions correctly
    4. Return parameters for optimization

    EXAMPLE:
    >>> embed = Embedding(vocab_size=100, embed_dim=64)
    >>> tokens = Tensor([[1, 2, 3], [4, 5, 6]])  # batch_size=2, seq_len=3
    >>> output = embed.forward(tokens)
    >>> print(output.shape)
    (2, 3, 64)

    HINTS:
    - Use numpy advanced indexing for lookup: weight[indices]
    - Embedding matrix shape: (vocab_size, embed_dim)
    - Initialize with Xavier/Glorot uniform for stable gradients
    - Handle multi-dimensional indices correctly
    """

    ### BEGIN SOLUTION
    def __init__(self, vocab_size: int, embed_dim: int):
        """
        Initialize embedding layer.

        Args:
            vocab_size: Size of vocabulary (number of unique tokens)
            embed_dim: Dimension of embedding vectors
        """
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim

        # Xavier initialization for better gradient flow
        limit = math.sqrt(6.0 / (vocab_size + embed_dim))
        self.weight = Tensor(
            np.random.uniform(-limit, limit, (vocab_size, embed_dim)),
            requires_grad=True
        )

    def forward(self, indices: Tensor) -> Tensor:
        """
        Forward pass: lookup embeddings for given indices.

        Args:
            indices: Token indices of shape (batch_size, seq_len) or (seq_len,)

        Returns:
            Embedded vectors of shape (*indices.shape, embed_dim)
        """
        # Handle input validation
        if np.any(indices.data >= self.vocab_size) or np.any(indices.data < 0):
            raise ValueError(
                f"Index out of range. Expected 0 <= indices < {self.vocab_size}, "
                f"got min={np.min(indices.data)}, max={np.max(indices.data)}"
            )

        # Perform embedding lookup using advanced indexing
        # This is equivalent to one-hot multiplication but much more efficient
        embedded = self.weight.data[indices.data.astype(int)]

        return Tensor(embedded)

    def parameters(self) -> List[Tensor]:
        """Return trainable parameters."""
        return [self.weight]

    def __repr__(self):
        return f"Embedding(vocab_size={self.vocab_size}, embed_dim={self.embed_dim})"
    ### END SOLUTION

In [ ]:

def test_unit_embedding():
    """🔬 Unit Test: Embedding Layer Implementation"""
    print("🔬 Unit Test: Embedding Layer...")

    # Test 1: Basic embedding creation and forward pass
    embed = Embedding(vocab_size=100, embed_dim=64)

    # Single sequence
    tokens = Tensor([1, 2, 3])
    output = embed.forward(tokens)

    assert output.shape == (3, 64), f"Expected shape (3, 64), got {output.shape}"
    assert len(embed.parameters()) == 1, "Should have 1 parameter (weight matrix)"
    assert embed.parameters()[0].shape == (100, 64), "Weight matrix has wrong shape"

    # Test 2: Batch processing
    batch_tokens = Tensor([[1, 2, 3], [4, 5, 6]])
    batch_output = embed.forward(batch_tokens)

    assert batch_output.shape == (2, 3, 64), f"Expected batch shape (2, 3, 64), got {batch_output.shape}"

    # Test 3: Embedding lookup consistency
    single_lookup = embed.forward(Tensor([1]))
    batch_lookup = embed.forward(Tensor([[1]]))

    # Should get same embedding for same token
    assert np.allclose(single_lookup.data[0], batch_lookup.data[0, 0]), "Inconsistent embedding lookup"

    # Test 4: Parameter access
    params = embed.parameters()
    assert all(p.requires_grad for p in params), "All parameters should require gradients"

    print("✅ Embedding layer works correctly!")

test_unit_embedding()

4. Understanding Positional Encoding - Teaching Models About Order

Sequences have inherent order, but embeddings by themselves are orderless. We need to explicitly encode positional information so the model understands that "cat chased dog" is different from "dog chased cat".

Why Position Matters in Sequences

Unlike images where spatial relationships are built into the 2D structure, text sequences need explicit position encoding:

Word Order Changes Meaning:
"The cat chased the dog" ≠ "The dog chased the cat"
"Not good" ≠ "Good not"
"She told him" ≠ "Him told she"

Two Approaches to Position Encoding

1. Learned Positional Embeddings:
   ┌─────────────────────────────────────┐
   │ Position  │  Learned Vector        │
   ├─────────────────────────────────────┤
   │    0      │ [0.1, -0.2, 0.4, ...]  │  (trained)
   │    1      │ [0.3,  0.1, -0.1, ...] │  (trained)
   │    2      │ [-0.1, 0.5, 0.2, ...]  │  (trained)
   │   ...     │        ...              │
   │   511     │ [0.4, -0.3, 0.1, ...]  │  (trained)
   └─────────────────────────────────────┘
   ✓ Can learn task-specific patterns
   ✗ Fixed maximum sequence length
   ✗ Requires additional parameters

2. Sinusoidal Position Encodings:
   ┌─────────────────────────────────────┐
   │ Position  │  Mathematical Pattern   │
   ├─────────────────────────────────────┤
   │    0      │ [0.0,  1.0, 0.0, ...]   │  (computed)
   │    1      │ [sin1, cos1, sin2, ...] │  (computed)
   │    2      │ [sin2, cos2, sin4, ...] │  (computed)
   │   ...     │        ...              │
   │   N       │ [sinN, cosN, sin2N,...] │  (computed)
   └─────────────────────────────────────┘
   ✓ No additional parameters
   ✓ Can extrapolate to longer sequences
   ✗ Cannot adapt to specific patterns

How Positional Information Gets Added

Token Embeddings + Positional Encodings = Position-Aware Representations

Input Sequence: ["The", "cat", "sat"]
Token IDs:      [  1,    42,    7 ]

Step 1: Token Embeddings
[1] → [0.1, 0.4, -0.2, ...]
[42]→ [0.7, -0.2, 0.1, ...]
[7] → [-0.3, 0.1, 0.5, ...]

Step 2: Position Encodings
pos 0 → [0.0, 1.0, 0.0, ...]
pos 1 → [0.8, 0.6, 0.1, ...]
pos 2 → [0.9, -0.4, 0.2, ...]

Step 3: Addition (element-wise)
Result:
[0.1+0.0, 0.4+1.0, -0.2+0.0, ...] = [0.1, 1.4, -0.2, ...]  "The" at position 0
[0.7+0.8, -0.2+0.6, 0.1+0.1, ...] = [1.5, 0.4, 0.2, ...]   "cat" at position 1
[-0.3+0.9, 0.1-0.4, 0.5+0.2, ...] = [0.6, -0.3, 0.7, ...]  "sat" at position 2

This way, the same word gets different representations based on its position in the sentence!

5. Implementing Learned Positional Encoding

Let's build trainable positional embeddings that can learn position-specific patterns for our specific task.

In [ ]:

class PositionalEncoding:
    """
    Learnable positional encoding layer.

    Adds trainable position-specific vectors to token embeddings,
    allowing the model to learn positional patterns specific to the task.

    TODO: Implement learnable positional encoding

    APPROACH:
    1. Create embedding matrix for positions: (max_seq_len, embed_dim)
    2. Forward pass: lookup position embeddings and add to input
    3. Handle different sequence lengths gracefully
    4. Return parameters for training

    EXAMPLE:
    >>> pos_enc = PositionalEncoding(max_seq_len=512, embed_dim=64)
    >>> embeddings = Tensor(np.random.randn(2, 10, 64))  # (batch, seq, embed)
    >>> output = pos_enc.forward(embeddings)
    >>> print(output.shape)
    (2, 10, 64)  # Same shape, but now position-aware

    HINTS:
    - Position embeddings shape: (max_seq_len, embed_dim)
    - Use slice [:seq_len] to handle variable lengths
    - Add position encodings to input embeddings element-wise
    - Initialize with smaller values than token embeddings (they're additive)
    """

    ### BEGIN SOLUTION
    def __init__(self, max_seq_len: int, embed_dim: int):
        """
        Initialize learnable positional encoding.

        Args:
            max_seq_len: Maximum sequence length to support
            embed_dim: Embedding dimension (must match token embeddings)
        """
        self.max_seq_len = max_seq_len
        self.embed_dim = embed_dim

        # Initialize position embedding matrix
        # Smaller initialization than token embeddings since these are additive
        limit = math.sqrt(2.0 / embed_dim)
        self.position_embeddings = Tensor(
            np.random.uniform(-limit, limit, (max_seq_len, embed_dim)),
            requires_grad=True
        )

    def forward(self, x: Tensor) -> Tensor:
        """
        Add positional encodings to input embeddings.

        Args:
            x: Input embeddings of shape (batch_size, seq_len, embed_dim)

        Returns:
            Position-encoded embeddings of same shape
        """
        if len(x.shape) != 3:
            raise ValueError(f"Expected 3D input (batch, seq, embed), got shape {x.shape}")

        batch_size, seq_len, embed_dim = x.shape

        if seq_len > self.max_seq_len:
            raise ValueError(
                f"Sequence length {seq_len} exceeds maximum {self.max_seq_len}"
            )

        if embed_dim != self.embed_dim:
            raise ValueError(
                f"Embedding dimension mismatch: expected {self.embed_dim}, got {embed_dim}"
            )

        # Get position embeddings for this sequence length
        pos_embeddings = self.position_embeddings.data[:seq_len]  # (seq_len, embed_dim)

        # Broadcast to match batch dimension: (1, seq_len, embed_dim)
        pos_embeddings = pos_embeddings[np.newaxis, :, :]

        # Add positional information to input embeddings
        result = x.data + pos_embeddings

        return Tensor(result)

    def parameters(self) -> List[Tensor]:
        """Return trainable parameters."""
        return [self.position_embeddings]

    def __repr__(self):
        return f"PositionalEncoding(max_seq_len={self.max_seq_len}, embed_dim={self.embed_dim})"
    ### END SOLUTION

In [ ]:

def test_unit_positional_encoding():
    """🔬 Unit Test: Positional Encoding Implementation"""
    print("🔬 Unit Test: Positional Encoding...")

    # Test 1: Basic functionality
    pos_enc = PositionalEncoding(max_seq_len=512, embed_dim=64)

    # Create sample embeddings
    embeddings = Tensor(np.random.randn(2, 10, 64))
    output = pos_enc.forward(embeddings)

    assert output.shape == (2, 10, 64), f"Expected shape (2, 10, 64), got {output.shape}"

    # Test 2: Position consistency
    # Same position should always get same encoding
    emb1 = Tensor(np.zeros((1, 5, 64)))
    emb2 = Tensor(np.zeros((1, 5, 64)))

    out1 = pos_enc.forward(emb1)
    out2 = pos_enc.forward(emb2)

    assert np.allclose(out1.data, out2.data), "Position encodings should be consistent"

    # Test 3: Different positions get different encodings
    short_emb = Tensor(np.zeros((1, 3, 64)))
    long_emb = Tensor(np.zeros((1, 5, 64)))

    short_out = pos_enc.forward(short_emb)
    long_out = pos_enc.forward(long_emb)

    # First 3 positions should match
    assert np.allclose(short_out.data, long_out.data[:, :3, :]), "Position encoding prefix should match"

    # Test 4: Parameters
    params = pos_enc.parameters()
    assert len(params) == 1, "Should have 1 parameter (position embeddings)"
    assert params[0].shape == (512, 64), "Position embedding matrix has wrong shape"

    print("✅ Positional encoding works correctly!")

test_unit_positional_encoding()

6. Understanding Sinusoidal Position Encodings

Now let's explore the elegant mathematical approach to position encoding used in the original Transformer paper. Instead of learning position patterns, we'll use trigonometric functions to create unique, continuous position signatures.

The Mathematical Intuition

Sinusoidal encodings use sine and cosine functions at different frequencies to create unique position signatures:

PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))     # Even dimensions
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))     # Odd dimensions

Why This Works - Frequency Visualization

Position Encoding Pattern (embed_dim=8, showing 4 positions):

Dimension:  0     1     2     3     4     5     6     7
Frequency:  High  High  Med   Med   Low   Low   VLow  VLow
Function:   sin   cos   sin   cos   sin   cos   sin   cos

pos=0:    [0.00, 1.00, 0.00, 1.00, 0.00, 1.00, 0.00, 1.00]
pos=1:    [0.84, 0.54, 0.01, 1.00, 0.00, 1.00, 0.00, 1.00]
pos=2:    [0.91, -0.42, 0.02, 1.00, 0.00, 1.00, 0.00, 1.00]
pos=3:    [0.14, -0.99, 0.03, 1.00, 0.00, 1.00, 0.00, 1.00]

Notice how:
- High frequency dimensions (0,1) change quickly between positions
- Low frequency dimensions (6,7) change slowly
- Each position gets a unique "fingerprint"

Visual Pattern of Sinusoidal Encodings

Frequency Spectrum Across Dimensions:
High Freq ← - - - - - - - - - - - - - - - - - - - - - → Low Freq
Dim:  0   1   2   3   4   5   6   7   8   9  ...  510 511

Wave Pattern for Position Progression:
Dim 0: ∿∿∿∿∿∿∿∿∿∿∿∿∿∿∿∿∿∿∿∿  (rapid oscillation)
Dim 2: ∿---∿---∿---∿---∿---∿  (medium frequency)
Dim 4: ∿-----∿-----∿-----∿--  (low frequency)
Dim 6: ∿----------∿----------  (very slow changes)

This creates a unique "barcode" for each position!

Advantages of Sinusoidal Encodings

1. No Parameters: Zero additional memory overhead
2. Extrapolation: Can handle sequences longer than training data
3. Unique Signatures: Each position gets a distinct encoding
4. Smooth Transitions: Similar positions have similar encodings
5. Mathematical Elegance: Clean, interpretable patterns

7. Implementing Sinusoidal Positional Encodings

Let's implement the mathematical position encoding that creates unique signatures for each position using trigonometric functions.

In [ ]:

def create_sinusoidal_embeddings(max_seq_len: int, embed_dim: int) -> Tensor:
    """
    Create sinusoidal positional encodings as used in "Attention Is All You Need".

    These fixed encodings use sine and cosine functions to create unique
    positional patterns that don't require training and can extrapolate
    to longer sequences than seen during training.

    TODO: Implement sinusoidal positional encoding generation

    APPROACH:
    1. Create position indices: [0, 1, 2, ..., max_seq_len-1]
    2. Create dimension indices for frequency calculation
    3. Apply sine to even dimensions, cosine to odd dimensions
    4. Use the transformer paper formula with 10000 base

    MATHEMATICAL FORMULA:
    PE(pos, 2i) = sin(pos / 10000^(2i/embed_dim))
    PE(pos, 2i+1) = cos(pos / 10000^(2i/embed_dim))

    EXAMPLE:
    >>> pe = create_sinusoidal_embeddings(512, 64)
    >>> print(pe.shape)
    (512, 64)
    >>> # Position 0: [0, 1, 0, 1, 0, 1, ...] (sin(0)=0, cos(0)=1)
    >>> # Each position gets unique trigonometric signature

    HINTS:
    - Use np.arange to create position and dimension arrays
    - Calculate div_term using exponential for frequency scaling
    - Apply different formulas to even/odd dimensions
    - The 10000 base creates different frequencies for different dimensions
    """

    ### BEGIN SOLUTION
    # Create position indices [0, 1, 2, ..., max_seq_len-1]
    position = np.arange(max_seq_len, dtype=np.float32)[:, np.newaxis]  # (max_seq_len, 1)

    # Create dimension indices for calculating frequencies
    div_term = np.exp(
        np.arange(0, embed_dim, 2, dtype=np.float32) *
        -(math.log(10000.0) / embed_dim)
    )  # (embed_dim//2,)

    # Initialize the positional encoding matrix
    pe = np.zeros((max_seq_len, embed_dim), dtype=np.float32)

    # Apply sine to even indices (0, 2, 4, ...)
    pe[:, 0::2] = np.sin(position * div_term)

    # Apply cosine to odd indices (1, 3, 5, ...)
    if embed_dim % 2 == 1:
        # Handle odd embed_dim by only filling available positions
        pe[:, 1::2] = np.cos(position * div_term[:-1])
    else:
        pe[:, 1::2] = np.cos(position * div_term)

    return Tensor(pe)
    ### END SOLUTION

In [ ]:

def test_unit_sinusoidal_embeddings():
    """🔬 Unit Test: Sinusoidal Positional Embeddings"""
    print("🔬 Unit Test: Sinusoidal Embeddings...")

    # Test 1: Basic shape and properties
    pe = create_sinusoidal_embeddings(512, 64)

    assert pe.shape == (512, 64), f"Expected shape (512, 64), got {pe.shape}"

    # Test 2: Position 0 should be mostly zeros and ones
    pos_0 = pe.data[0]

    # Even indices should be sin(0) = 0
    assert np.allclose(pos_0[0::2], 0, atol=1e-6), "Even indices at position 0 should be ~0"

    # Odd indices should be cos(0) = 1
    assert np.allclose(pos_0[1::2], 1, atol=1e-6), "Odd indices at position 0 should be ~1"

    # Test 3: Different positions should have different encodings
    pe_small = create_sinusoidal_embeddings(10, 8)

    # Check that consecutive positions are different
    for i in range(9):
        assert not np.allclose(pe_small.data[i], pe_small.data[i+1]), f"Positions {i} and {i+1} are too similar"

    # Test 4: Frequency properties
    # Higher dimensions should have lower frequencies (change more slowly)
    pe_test = create_sinusoidal_embeddings(100, 16)

    # First dimension should change faster than last dimension
    first_dim_changes = np.sum(np.abs(np.diff(pe_test.data[:10, 0])))
    last_dim_changes = np.sum(np.abs(np.diff(pe_test.data[:10, -1])))

    assert first_dim_changes > last_dim_changes, "Lower dimensions should change faster than higher dimensions"

    # Test 5: Odd embed_dim handling
    pe_odd = create_sinusoidal_embeddings(10, 7)
    assert pe_odd.shape == (10, 7), "Should handle odd embedding dimensions"

    print("✅ Sinusoidal embeddings work correctly!")

test_unit_sinusoidal_embeddings()

8. Building the Complete Embedding System

Now let's integrate everything into a production-ready embedding system that handles both token and positional embeddings, supports multiple encoding types, and manages the full embedding pipeline used in modern NLP models.

Complete Embedding Pipeline Visualization

Complete Embedding System Architecture:

Input: Token IDs [1, 42, 7, 99]
         ↓
    ┌─────────────────────┐
    │   Token Embedding   │  vocab_size × embed_dim table
    │   Lookup Table      │
    └─────────────────────┘
         ↓
    Token Vectors (4 × embed_dim)
    [0.1, 0.4, -0.2, ...]  ← token 1
    [0.7, -0.2, 0.1, ...]  ← token 42
    [-0.3, 0.1, 0.5, ...]  ← token 7
    [0.9, -0.1, 0.3, ...]  ← token 99
         ↓
    ┌─────────────────────┐
    │ Positional Encoding │  Choose: Learned, Sinusoidal, or None
    │  (Add position info) │
    └─────────────────────┘
         ↓
    Position-Aware Embeddings (4 × embed_dim)
    [0.1+pos0, 0.4+pos0, ...]  ← token 1 at position 0
    [0.7+pos1, -0.2+pos1, ...] ← token 42 at position 1
    [-0.3+pos2, 0.1+pos2, ...] ← token 7 at position 2
    [0.9+pos3, -0.1+pos3, ...] ← token 99 at position 3
         ↓
    Optional: Scale by √embed_dim (Transformer convention)
         ↓
    Ready for Attention Mechanisms!

Integration Features

• Flexible Position Encoding: Support learned, sinusoidal, or no positional encoding
• Batch Processing: Handle variable-length sequences with padding
• Memory Efficiency: Reuse position encodings across batches
• Production Ready: Matches PyTorch patterns and conventions

In [ ]:

class EmbeddingLayer:
    """
    Complete embedding system combining token and positional embeddings.

    This is the production-ready component that handles the full embedding
    pipeline used in transformers and other sequence models.

    TODO: Implement complete embedding system

    APPROACH:
    1. Combine token embedding + positional encoding
    2. Support both learned and sinusoidal position encodings
    3. Handle variable sequence lengths gracefully
    4. Add optional embedding scaling (Transformer convention)

    EXAMPLE:
    >>> embed_layer = EmbeddingLayer(
    ...     vocab_size=50000,
    ...     embed_dim=512,
    ...     max_seq_len=2048,
    ...     pos_encoding='learned'
    ... )
    >>> tokens = Tensor([[1, 2, 3], [4, 5, 6]])
    >>> output = embed_layer.forward(tokens)
    >>> print(output.shape)
    (2, 3, 512)

    HINTS:
    - First apply token embedding, then add positional encoding
    - Support 'learned', 'sinusoidal', or None for pos_encoding
    - Handle both 2D (batch, seq) and 1D (seq) inputs gracefully
    - Scale embeddings by sqrt(embed_dim) if requested (transformer convention)
    """

    ### BEGIN SOLUTION
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int,
        max_seq_len: int = 512,
        pos_encoding: str = 'learned',
        scale_embeddings: bool = False
    ):
        """
        Initialize complete embedding system.

        Args:
            vocab_size: Size of vocabulary
            embed_dim: Embedding dimension
            max_seq_len: Maximum sequence length for positional encoding
            pos_encoding: Type of positional encoding ('learned', 'sinusoidal', or None)
            scale_embeddings: Whether to scale embeddings by sqrt(embed_dim)
        """
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.max_seq_len = max_seq_len
        self.pos_encoding_type = pos_encoding
        self.scale_embeddings = scale_embeddings

        # Token embedding layer
        self.token_embedding = Embedding(vocab_size, embed_dim)

        # Positional encoding
        if pos_encoding == 'learned':
            self.pos_encoding = PositionalEncoding(max_seq_len, embed_dim)
        elif pos_encoding == 'sinusoidal':
            # Create fixed sinusoidal encodings (no parameters)
            self.pos_encoding = create_sinusoidal_embeddings(max_seq_len, embed_dim)
        elif pos_encoding is None:
            self.pos_encoding = None
        else:
            raise ValueError(f"Unknown pos_encoding: {pos_encoding}. Use 'learned', 'sinusoidal', or None")

    def forward(self, tokens: Tensor) -> Tensor:
        """
        Forward pass through complete embedding system.

        Args:
            tokens: Token indices of shape (batch_size, seq_len) or (seq_len,)

        Returns:
            Embedded tokens with positional information
        """
        # Handle 1D input by adding batch dimension
        if len(tokens.shape) == 1:
            tokens = Tensor(tokens.data[np.newaxis, :])  # (1, seq_len)
            squeeze_batch = True
        else:
            squeeze_batch = False

        # Get token embeddings
        token_embeds = self.token_embedding.forward(tokens)  # (batch, seq, embed)

        # Scale embeddings if requested (transformer convention)
        if self.scale_embeddings:
            token_embeds = Tensor(token_embeds.data * math.sqrt(self.embed_dim))

        # Add positional encoding
        if self.pos_encoding_type == 'learned':
            # Use learnable positional encoding
            output = self.pos_encoding.forward(token_embeds)
        elif self.pos_encoding_type == 'sinusoidal':
            # Use fixed sinusoidal encoding
            batch_size, seq_len, embed_dim = token_embeds.shape
            pos_embeddings = self.pos_encoding.data[:seq_len]  # (seq_len, embed_dim)
            pos_embeddings = pos_embeddings[np.newaxis, :, :]  # (1, seq_len, embed_dim)
            output = Tensor(token_embeds.data + pos_embeddings)
        else:
            # No positional encoding
            output = token_embeds

        # Remove batch dimension if it was added
        if squeeze_batch:
            output = Tensor(output.data[0])  # (seq_len, embed_dim)

        return output

    def parameters(self) -> List[Tensor]:
        """Return all trainable parameters."""
        params = self.token_embedding.parameters()

        if self.pos_encoding_type == 'learned':
            params.extend(self.pos_encoding.parameters())

        return params

    def __repr__(self):
        return (f"EmbeddingLayer(vocab_size={self.vocab_size}, "
                f"embed_dim={self.embed_dim}, "
                f"pos_encoding='{self.pos_encoding_type}')")
    ### END SOLUTION

In [ ]:

def test_unit_complete_embedding_system():
    """🔬 Unit Test: Complete Embedding System"""
    print("🔬 Unit Test: Complete Embedding System...")

    # Test 1: Learned positional encoding
    embed_learned = EmbeddingLayer(
        vocab_size=100,
        embed_dim=64,
        max_seq_len=128,
        pos_encoding='learned'
    )

    tokens = Tensor([[1, 2, 3], [4, 5, 6]])
    output_learned = embed_learned.forward(tokens)

    assert output_learned.shape == (2, 3, 64), f"Expected shape (2, 3, 64), got {output_learned.shape}"

    # Test 2: Sinusoidal positional encoding
    embed_sin = EmbeddingLayer(
        vocab_size=100,
        embed_dim=64,
        pos_encoding='sinusoidal'
    )

    output_sin = embed_sin.forward(tokens)
    assert output_sin.shape == (2, 3, 64), "Sinusoidal embedding should have same shape"

    # Test 3: No positional encoding
    embed_none = EmbeddingLayer(
        vocab_size=100,
        embed_dim=64,
        pos_encoding=None
    )

    output_none = embed_none.forward(tokens)
    assert output_none.shape == (2, 3, 64), "No pos encoding should have same shape"

    # Test 4: 1D input handling
    tokens_1d = Tensor([1, 2, 3])
    output_1d = embed_learned.forward(tokens_1d)

    assert output_1d.shape == (3, 64), f"Expected shape (3, 64) for 1D input, got {output_1d.shape}"

    # Test 5: Embedding scaling
    embed_scaled = EmbeddingLayer(
        vocab_size=100,
        embed_dim=64,
        pos_encoding=None,
        scale_embeddings=True
    )

    # Use same weights to ensure fair comparison
    embed_scaled.token_embedding.weight = embed_none.token_embedding.weight

    output_scaled = embed_scaled.forward(tokens)
    output_unscaled = embed_none.forward(tokens)

    # Scaled version should be sqrt(64) times larger
    scale_factor = math.sqrt(64)
    expected_scaled = output_unscaled.data * scale_factor
    assert np.allclose(output_scaled.data, expected_scaled, rtol=1e-5), "Embedding scaling not working correctly"

    # Test 6: Parameter counting
    params_learned = embed_learned.parameters()
    params_sin = embed_sin.parameters()
    params_none = embed_none.parameters()

    assert len(params_learned) == 2, "Learned encoding should have 2 parameter tensors"
    assert len(params_sin) == 1, "Sinusoidal encoding should have 1 parameter tensor"
    assert len(params_none) == 1, "No pos encoding should have 1 parameter tensor"

    print("✅ Complete embedding system works correctly!")

test_unit_complete_embedding_system()

9. Systems Analysis - Embedding Memory and Performance

Understanding the systems implications of embedding layers is crucial for building scalable NLP models. Let's analyze memory usage, lookup performance, and trade-offs between different approaches.

Memory Usage Analysis

Embedding Memory Scaling:
Vocabulary Size vs Memory Usage (embed_dim=512, FP32):

 10K vocab: 10,000 × 512 × 4 bytes = 20 MB
 50K vocab: 50,000 × 512 × 4 bytes = 100 MB
100K vocab: 100,000 × 512 × 4 bytes = 200 MB
  1M vocab: 1,000,000 × 512 × 4 bytes = 2 GB

GPT-3 Scale: 50,257 × 12,288 × 4 bytes ≈ 2.4 GB just for embeddings!

Memory Formula: vocab_size × embed_dim × 4 bytes (FP32)

Performance Characteristics

Embedding Lookup Performance:
- Time Complexity: O(1) per token (hash table lookup)
- Memory Access: Random access pattern
- Bottleneck: Memory bandwidth, not computation
- Batching: Improves throughput via vectorization

Cache Efficiency:
Repeated tokens → Cache hits → Faster access
Diverse vocab → Cache misses → Slower access

In [ ]:

def analyze_embedding_memory():
    """📊 Analyze embedding memory requirements and scaling behavior."""
    print("📊 Analyzing Embedding Memory Requirements...")

    # Vocabulary and embedding dimension scenarios
    scenarios = [
        ("Small Model", 10_000, 256),
        ("Medium Model", 50_000, 512),
        ("Large Model", 100_000, 1024),
        ("GPT-3 Scale", 50_257, 12_288),
    ]

    print(f"{'Model':<15} {'Vocab Size':<12} {'Embed Dim':<12} {'Memory (MB)':<15} {'Parameters (M)':<15}")
    print("-" * 80)

    for name, vocab_size, embed_dim in scenarios:
        # Calculate memory for FP32 (4 bytes per parameter)
        params = vocab_size * embed_dim
        memory_mb = params * 4 / (1024 * 1024)
        params_m = params / 1_000_000

        print(f"{name:<15} {vocab_size:<12,} {embed_dim:<12} {memory_mb:<15.1f} {params_m:<15.2f}")

    print("\n💡 Key Insights:")
    print("• Embedding tables often dominate model memory (especially for large vocabularies)")
    print("• Memory scales linearly with vocab_size × embed_dim")
    print("• Consider vocabulary pruning for memory-constrained environments")

    # Positional encoding memory comparison
    print(f"\n📊 Positional Encoding Memory Comparison (embed_dim=512, max_seq_len=2048):")

    learned_params = 2048 * 512
    learned_memory = learned_params * 4 / (1024 * 1024)

    print(f"Learned PE:     {learned_memory:.1f} MB ({learned_params:,} parameters)")
    print(f"Sinusoidal PE:  0.0 MB (0 parameters - computed on-the-fly)")
    print(f"No PE:          0.0 MB (0 parameters)")

    print("\n🚀 Production Implications:")
    print("• GPT-3's embedding table: ~2.4GB (50K vocab × 12K dims)")
    print("• Learned PE adds memory but may improve task-specific performance")
    print("• Sinusoidal PE saves memory and allows longer sequences")

analyze_embedding_memory()

In [ ]:

def analyze_lookup_performance():
    """📊 Analyze embedding lookup performance characteristics."""
    print("\n📊 Analyzing Embedding Lookup Performance...")

    import time

    # Test different vocabulary sizes and batch configurations
    vocab_sizes = [1_000, 10_000, 100_000]
    embed_dim = 512
    seq_len = 128
    batch_sizes = [1, 16, 64, 256]

    print(f"{'Vocab Size':<12} {'Batch Size':<12} {'Lookup Time (ms)':<18} {'Throughput (tokens/s)':<20}")
    print("-" * 70)

    for vocab_size in vocab_sizes:
        # Create embedding layer
        embed = Embedding(vocab_size, embed_dim)

        for batch_size in batch_sizes:
            # Create random token batch
            tokens = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))

            # Warmup
            for _ in range(5):
                _ = embed.forward(tokens)

            # Time the lookup
            start_time = time.time()
            iterations = 100

            for _ in range(iterations):
                output = embed.forward(tokens)

            end_time = time.time()

            # Calculate metrics
            total_time = end_time - start_time
            avg_time_ms = (total_time / iterations) * 1000
            total_tokens = batch_size * seq_len * iterations
            throughput = total_tokens / total_time

            print(f"{vocab_size:<12,} {batch_size:<12} {avg_time_ms:<18.2f} {throughput:<20,.0f}")

    print("\n💡 Performance Insights:")
    print("• Lookup time is O(1) per token - vocabulary size doesn't affect individual lookups")
    print("• Larger batches improve throughput due to vectorization")
    print("• Memory bandwidth becomes bottleneck for large embedding dimensions")
    print("• Cache locality important for repeated token patterns")

analyze_lookup_performance()

In [ ]:

def analyze_positional_encoding_trade_offs():
    """📊 Compare learned vs sinusoidal positional encodings."""
    print("\n📊 Analyzing Positional Encoding Trade-offs...")

    max_seq_len = 512
    embed_dim = 256

    # Create both types of positional encodings
    learned_pe = PositionalEncoding(max_seq_len, embed_dim)
    sinusoidal_pe = create_sinusoidal_embeddings(max_seq_len, embed_dim)

    # Analyze memory footprint
    learned_params = max_seq_len * embed_dim
    learned_memory = learned_params * 4 / (1024 * 1024)  # MB

    print(f"📈 Memory Comparison:")
    print(f"Learned PE:     {learned_memory:.2f} MB ({learned_params:,} parameters)")
    print(f"Sinusoidal PE:  0.00 MB (0 parameters)")

    # Analyze encoding patterns
    print(f"\n📈 Encoding Pattern Analysis:")

    # Test sample sequences
    test_input = Tensor(np.random.randn(1, 10, embed_dim))

    learned_output = learned_pe.forward(test_input)

    # For sinusoidal, manually add to match learned interface
    sin_encodings = sinusoidal_pe.data[:10][np.newaxis, :, :]  # (1, 10, embed_dim)
    sinusoidal_output = Tensor(test_input.data + sin_encodings)

    # Analyze variance across positions
    learned_var = np.var(learned_output.data, axis=1).mean()  # Variance across positions
    sin_var = np.var(sinusoidal_output.data, axis=1).mean()

    print(f"Position variance (learned):    {learned_var:.4f}")
    print(f"Position variance (sinusoidal): {sin_var:.4f}")

    # Check extrapolation capability
    print(f"\n📈 Extrapolation Analysis:")
    extended_length = max_seq_len + 100

    try:
        # Learned PE cannot handle longer sequences
        extended_learned = PositionalEncoding(extended_length, embed_dim)
        print(f"Learned PE: Requires retraining for sequences > {max_seq_len}")
    except:
        print(f"Learned PE: Cannot handle sequences > {max_seq_len}")

    # Sinusoidal can extrapolate
    extended_sin = create_sinusoidal_embeddings(extended_length, embed_dim)
    print(f"Sinusoidal PE: Can extrapolate to length {extended_length} (smooth continuation)")

    print(f"\n🚀 Production Trade-offs:")
    print(f"Learned PE:")
    print(f"  + Can learn task-specific positional patterns")
    print(f"  + May perform better for tasks with specific position dependencies")
    print(f"  - Requires additional memory and parameters")
    print(f"  - Fixed maximum sequence length")
    print(f"  - Needs training data for longer sequences")

    print(f"\nSinusoidal PE:")
    print(f"  + Zero additional parameters")
    print(f"  + Can extrapolate to any sequence length")
    print(f"  + Provides rich, mathematically grounded position signals")
    print(f"  - Cannot adapt to task-specific position patterns")
    print(f"  - May be suboptimal for highly position-dependent tasks")

analyze_positional_encoding_trade_offs()

10. Module Integration Test

Final validation that our complete embedding system works correctly and integrates with the TinyTorch ecosystem.

In [ ]:

def test_module():
    """
    Comprehensive test of entire embeddings module functionality.

    This final test ensures all components work together and the module
    is ready for integration with attention mechanisms and transformers.
    """
    print("🧪 RUNNING MODULE INTEGRATION TEST")
    print("=" * 50)

    # Run all unit tests
    print("Running unit tests...")
    test_unit_embedding()
    test_unit_positional_encoding()
    test_unit_sinusoidal_embeddings()
    test_unit_complete_embedding_system()

    print("\nRunning integration scenarios...")

    # Integration Test 1: Realistic NLP pipeline
    print("🔬 Integration Test: NLP Pipeline Simulation...")

    # Simulate a small transformer setup
    vocab_size = 1000
    embed_dim = 128
    max_seq_len = 64

    # Create embedding layer
    embed_layer = EmbeddingLayer(
        vocab_size=vocab_size,
        embed_dim=embed_dim,
        max_seq_len=max_seq_len,
        pos_encoding='learned',
        scale_embeddings=True
    )

    # Simulate tokenized sentences
    sentences = [
        [1, 15, 42, 7, 99],        # "the cat sat on mat"
        [23, 7, 15, 88],           # "dog chased the ball"
        [1, 67, 15, 42, 7, 99, 34] # "the big cat sat on mat here"
    ]

    # Process each sentence
    outputs = []
    for sentence in sentences:
        tokens = Tensor(sentence)
        embedded = embed_layer.forward(tokens)
        outputs.append(embedded)

        # Verify output shape
        expected_shape = (len(sentence), embed_dim)
        assert embedded.shape == expected_shape, f"Wrong shape for sentence: {embedded.shape} != {expected_shape}"

    print("✅ Variable length sentence processing works!")

    # Integration Test 2: Batch processing with padding
    print("🔬 Integration Test: Batched Processing...")

    # Create padded batch (real-world scenario)
    max_len = max(len(s) for s in sentences)
    batch_tokens = []

    for sentence in sentences:
        # Pad with zeros (assuming 0 is padding token)
        padded = sentence + [0] * (max_len - len(sentence))
        batch_tokens.append(padded)

    batch_tensor = Tensor(batch_tokens)  # (3, 7)
    batch_output = embed_layer.forward(batch_tensor)

    assert batch_output.shape == (3, max_len, embed_dim), f"Batch output shape incorrect: {batch_output.shape}"

    print("✅ Batch processing with padding works!")

    # Integration Test 3: Different positional encoding types
    print("🔬 Integration Test: Position Encoding Variants...")

    test_tokens = Tensor([[1, 2, 3, 4, 5]])

    # Test all position encoding types
    for pe_type in ['learned', 'sinusoidal', None]:
        embed_test = EmbeddingLayer(
            vocab_size=100,
            embed_dim=64,
            pos_encoding=pe_type
        )

        output = embed_test.forward(test_tokens)
        assert output.shape == (1, 5, 64), f"PE type {pe_type} failed shape test"

        # Check parameter counts
        if pe_type == 'learned':
            assert len(embed_test.parameters()) == 2, f"Learned PE should have 2 param tensors"
        else:
            assert len(embed_test.parameters()) == 1, f"PE type {pe_type} should have 1 param tensor"

    print("✅ All positional encoding variants work!")

    # Integration Test 4: Memory efficiency check
    print("🔬 Integration Test: Memory Efficiency...")

    # Test that we're not creating unnecessary copies
    large_embed = EmbeddingLayer(vocab_size=10000, embed_dim=512)
    test_batch = Tensor(np.random.randint(0, 10000, (32, 128)))

    # Multiple forward passes should not accumulate memory (in production)
    for _ in range(5):
        output = large_embed.forward(test_batch)
        assert output.shape == (32, 128, 512), "Large batch processing failed"

    print("✅ Memory efficiency check passed!")

    print("\n" + "=" * 50)
    print("🎉 ALL TESTS PASSED! Module ready for export.")
    print("📚 Summary of capabilities built:")
    print("  • Token embedding with trainable lookup tables")
    print("  • Learned positional encodings for position awareness")
    print("  • Sinusoidal positional encodings for extrapolation")
    print("  • Complete embedding system for NLP pipelines")
    print("  • Efficient batch processing and memory management")
    print("\n🚀 Ready for: Attention mechanisms, transformers, and language models!")
    print("Export with: tito module complete 11")

In [ ]:

if __name__ == "__main__":
    """Main execution block for module validation."""
    print("🚀 Running Embeddings module...")
    test_module()
    print("✅ Module validation complete!")

🤔 ML Systems Thinking: Embedding Foundations

Question 1: Memory Scaling

You implemented an embedding layer with vocab_size=50,000 and embed_dim=512.

• How many parameters does this embedding table contain? _____ million
• If using FP32 (4 bytes per parameter), how much memory does this use? _____ MB
• If you double the embedding dimension to 1024, what happens to memory usage? _____ MB

Question 2: Lookup Complexity

Your embedding layer performs table lookups for token indices.

• What is the time complexity of looking up a single token? O(_____)
• For a batch of 32 sequences, each of length 128, how many lookup operations? _____
• Why doesn't vocabulary size affect individual lookup performance? _____

Question 3: Positional Encoding Trade-offs

You implemented both learned and sinusoidal positional encodings.

• Learned PE for max_seq_len=2048, embed_dim=512 adds how many parameters? _____
• What happens if you try to process a sequence longer than max_seq_len with learned PE? _____
• Which type of PE can handle sequences longer than seen during training? _____

Question 4: Production Implications

Your complete EmbeddingLayer combines token and positional embeddings.

• In GPT-3 (vocab_size≈50K, embed_dim≈12K), approximately what percentage of total parameters are in the embedding table? _____%
• If you wanted to reduce memory usage by 50%, which would be more effective: halving vocab_size or halving embed_dim? _____
• Why might sinusoidal PE be preferred for models that need to handle variable sequence lengths? _____

🎯 MODULE SUMMARY: Embeddings

Congratulations! You've built a complete embedding system that transforms discrete tokens into learnable representations!

Key Accomplishments

• Built Embedding class with efficient token-to-vector lookup (10M+ token support)
• Implemented PositionalEncoding for learnable position awareness (unlimited sequence patterns)
• Created create_sinusoidal_embeddings with mathematical position encoding (extrapolates beyond training)
• Developed EmbeddingLayer integrating both token and positional embeddings (production-ready)
• Analyzed embedding memory scaling and lookup performance trade-offs
• All tests pass ✅ (validated by test_module())

Technical Achievements

• Memory Efficiency: Optimized embedding table storage and lookup patterns
• Flexible Architecture: Support for learned, sinusoidal, and no positional encoding
• Batch Processing: Efficient handling of variable-length sequences with padding
• Systems Analysis: Deep understanding of memory vs performance trade-offs

Ready for Next Steps

Your embeddings implementation enables attention mechanisms and transformer architectures! The combination of token and positional embeddings provides the foundation for sequence-to-sequence models.

Next: Module 12 will add attention mechanisms for context-aware representations!

Production Context

You've built the exact embedding patterns used in:

• GPT models: Token embeddings + learned positional encoding
• BERT models: Token embeddings + sinusoidal positional encoding
• T5 models: Relative positional embeddings (variant of your implementations)

Export with: tito module complete 11