abelrguezr/token-embeddings

Token Embeddings Skill

This skill helps you create, understand, and work with token embeddings for large language models.

What This Skill Does

• Explains token embedding concepts and initialization
• Creates PyTorch embedding layers for your vocabulary
• Adds positional embeddings (absolute, relative, RoPE)
• Helps debug embedding-related issues
• Provides code templates for common embedding tasks

When to Use This Skill

Use this skill when you:

• Need to create token embeddings for a new vocabulary
• Want to understand how embeddings work in your model
• Need to add positional information to your embeddings
• Are debugging embedding dimension mismatches
• Want to extend context windows in RoPE-based models
• Need to implement or understand different positional encoding strategies

Core Concepts

Token Embeddings

Token embeddings convert discrete tokens into continuous vectors. Each token in your vocabulary gets a unique vector of fixed dimensions.

Key parameters:

• vocab_size: Number of unique tokens (e.g., 50257 for BPE)
• embedding_dim: Vector dimensions (e.g., 256, 512, 768)

Example:

Vocabulary: [1, 2, 3, 4, 5, 6] (6 tokens)
Embedding dim: 3
Token 3 → [-0.4015, 0.9666, -1.1481]

Positional Embeddings

Positional embeddings encode token positions in sequences. Without them, the model treats tokens as a "bag of words."

Types:

1. Absolute: Fixed position vectors (GPT-style)
2. Relative: Distance-based encoding (Transformer-XL, BERT variants)
3. RoPE: Rotary embeddings (modern decoder-only LLMs)

Quick Start

Create Basic Token Embeddings

import torch

vocab_size = 50257  # BPE vocabulary
embedding_dim = 256

token_embedding = torch.nn.Embedding(vocab_size, embedding_dim)

Add Absolute Positional Embeddings

context_length = 512
pos_embedding = torch.nn.Embedding(context_length, embedding_dim)

# Combine embeddings
token_emb = token_embedding(token_ids)  # [batch, seq_len, dim]
pos_emb = pos_embedding(torch.arange(seq_len))  # [seq_len, dim]
combined = token_emb + pos_emb  # [batch, seq_len, dim]

RoPE (Rotary Positional Embeddings)

For modern LLMs, RoPE is preferred:

def apply_rope(q, k, cos, sin):
    """Apply rotary positional embeddings to query/key vectors."""
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

Common Tasks

Task 1: Initialize Embedding Layer

import torch

def create_token_embeddings(vocab_size: int, embedding_dim: int) -> torch.nn.Embedding:
    """Create a token embedding layer."""
    return torch.nn.Embedding(vocab_size, embedding_dim)

# Usage
embedding_layer = create_token_embeddings(50257, 256)
print(embedding_layer.weight.shape)  # torch.Size([50257, 256])

Task 2: Create Positional Embeddings

def create_positional_embeddings(context_length: int, embedding_dim: int) -> torch.nn.Embedding:
    """Create absolute positional embeddings."""
    return torch.nn.Embedding(context_length, embedding_dim)

# Usage
pos_layer = create_positional_embeddings(512, 256)
pos_embeddings = pos_layer(torch.arange(512))
print(pos_embeddings.shape)  # torch.Size([512, 256])

Task 3: Combine Token and Positional Embeddings

def combine_embeddings(
    token_ids: torch.Tensor,
    token_embedding: torch.nn.Embedding,
    pos_embedding: torch.nn.Embedding
) -> torch.Tensor:
    """Combine token and positional embeddings.
    
    Args:
        token_ids: [batch_size, seq_len]
        token_embedding: Token embedding layer
        pos_embedding: Positional embedding layer
    
    Returns:
        Combined embeddings: [batch_size, seq_len, embedding_dim]
    """
    batch_size, seq_len = token_ids.shape
    
    # Get token embeddings
    token_emb = token_embedding(token_ids)  # [batch, seq_len, dim]
    
    # Get positional embeddings
    positions = torch.arange(seq_len).expand(batch_size, -1)
    pos_emb = pos_embedding(positions)  # [batch, seq_len, dim]
    
    # Combine
    return token_emb + pos_emb

Task 4: Position Interpolation for Extended Context

def position_interpolation(
    pos_ids: torch.Tensor,
    original_context: int,
    new_context: int
) -> torch.Tensor:
    """Scale position indices for context window extension.
    
    Args:
        pos_ids: Original position indices
        original_context: Training context length (e.g., 2048)
        new_context: Target context length (e.g., 8192)
    
    Returns:
        Scaled position indices
    """
    scale = original_context / new_context
    scaled_pos = (pos_ids * scale).long()
    return scaled_pos

# Usage
original_ctx = 2048
new_ctx = 8192
positions = torch.arange(8192)
scaled_positions = position_interpolation(positions, original_ctx, new_ctx)

Debugging Checklist

When embeddings aren't working correctly, check:

1. Dimension Mismatches

   # Verify shapes match
   assert token_emb.shape == pos_emb.shape, "Embedding dimensions must match"
   

2. Vocabulary Size

   # Ensure vocab_size matches your tokenizer
   max_token_id = token_ids.max()
   assert max_token_id < vocab_size, f"Token {max_token_id} exceeds vocab_size {vocab_size}"
   

3. Context Length

   # Ensure sequence doesn't exceed context length
   seq_len = token_ids.shape[1]
   assert seq_len <= context_length, f"Sequence {seq_len} exceeds context {context_length}"
   

4. Gradient Flow

   # Verify embeddings are trainable
   assert token_embedding.weight.requires_grad, "Embeddings should be trainable"
   

Best Practices

1. Embedding Dimensions: Use powers of 2 (256, 512, 768, 1024) for efficiency
2. Initialization: PyTorch's default Xavier initialization works well
3. Positional Encoding: Use RoPE for decoder-only models, absolute for encoder-only
4. Context Extension: Use position interpolation before fine-tuning for longer contexts
5. Batch Processing: Always process in batches for efficiency

Example: Complete Embedding Setup

import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, embedding_dim: int, context_length: int):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, embedding_dim)
        self.pos_embedding = nn.Embedding(context_length, embedding_dim)
        self.context_length = context_length
    
    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len = token_ids.shape
        
        # Token embeddings
        token_emb = self.token_embedding(token_ids)
        
        # Positional embeddings
        positions = torch.arange(seq_len).expand(batch_size, -1)
        pos_emb = self.pos_embedding(positions)
        
        # Combine
        return token_emb + pos_emb

# Usage
vocab_size = 50257
embedding_dim = 256
context_length = 512

embedding_model = TokenEmbedding(vocab_size, embedding_dim, context_length)

# Test with sample input
batch_size = 8
seq_len = 4
token_ids = torch.randint(0, vocab_size, (batch_size, seq_len))

output = embedding_model(token_ids)
print(output.shape)  # torch.Size([8, 4, 256])

References

• Build a Large Language Model from Scratch
• RoPE Paper
• YaRN: Context Extension
• Position Interpolation