Learned Position Embeddings › Implementation in PyTorch
Language AI HandbookAll chapters Learned Position Embeddings › Implementation in PyTorchOne section · Sections in this chapter15 / 19 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
In practice, you would implement learned position embeddings using PyTorch's nn.Embedding layer. Here's what a real implementation looks like:
import torch
import torch.nn as nn
class TransformerEmbedding(nn.Module):
"""
Transformer embedding layer with learned token and position embeddings.
"""
def __init__(self, vocab_size, max_seq_len, embed_dim, dropout=0.1):
super().__init__()
# Token embedding table
self.token_embed = nn.Embedding(vocab_size, embed_dim)
# Position embedding table
self.pos_embed = nn.Embedding(max_seq_len, embed_dim)
# Dropout for regularization
self.dropout = nn.Dropout(dropout)
# Store for position indexing
self.max_seq_len = max_seq_len
def forward(self, token_ids):
"""
Args:
token_ids: LongTensor of shape (batch_size, seq_len)
Returns:
Embeddings of shape (batch_size, seq_len, embed_dim)
"""
batch_size, seq_len = token_ids.shape
# Create position indices: [0, 1, 2, ..., seq_len-1]
positions = torch.arange(seq_len, device=token_ids.device)
positions = positions.unsqueeze(0).expand(batch_size, -1)
# Look up embeddings and add
token_embeds = self.token_embed(token_ids)
pos_embeds = self.pos_embed(positions)
return self.dropout(token_embeds + pos_embeds)The PyTorch implementation is clean and efficient. The nn.Embedding layer handles the lookup table and gradient computation automatically. Position indices are created on-the-fly based on sequence length, and the same position embeddings are shared across all examples in a batch.
Notice the positions.unsqueeze(0).expand(batch_size, -1) idiom. This creates a position index tensor of shape (batch_size, seq_len) where every row is [0, 1, 2, ..., seq_len-1]. The expansion is memory-efficient because it shares the underlying storage rather than creating copies. When the embedding lookup runs, every example in the batch uses identical position indices (positions 0 through seq_len-1), so they all get the same position embeddings. This is correct: position 0 in sequence A should use the same position embedding as position 0 in sequence B, because position is an absolute index that does not depend on the content of the sequence.
The dropout applied after combining embeddings is regularization. By randomly zeroing out dimensions of the combined embedding during training, dropout prevents the model from relying too heavily on any single dimension of the positional or token representation. This tends to improve generalization, especially for smaller datasets.