PyTorch Word Embedding

Word embedding is one of the most fundamental and important techniques in natural language processing.

Word embedding maps discrete word symbols into continuous dense vectors, enabling machines to understand and process text data.

PyTorch providesnn.Embeddingmodule to implement this functionality, serving as the foundation for building various NLP models.

1. Basic Concepts of Word Embedding

In a computer, text is essentially a sequence of integers. Each word is assigned a unique index ID, but this discrete representation has a problem: semantically similar words may be close in meaning, yet their IDs are completely unrelated.

Word embedding solves this problem by learning anembedding matrix:

\[ E \in \mathbb{R}^{V \times D} \]

Where \(V\) is the vocabulary size and \(D\) is the embedding dimension. Each word ID corresponds to a row in the embedding matrix, and its vector representation is obtained through a lookup operation:

\[ \text{embedding} = E[\text{word\_id}] \]

The advantages of word embedding include:

Converting high-dimensional sparse one-hot vectors into low-dimensional dense vectors, greatly reducing computational cost
Semantically similar words are closer in vector space, and word similarity can be computed via cosine similarity
Embedding vectors are learnable parameters and can be automatically adjusted through backpropagation

2. nn.Embedding Explained in Detail

nn.Embeddingis the word embedding layer provided by PyTorch, encapsulating the creation of the embedding matrix and the lookup operation.

2.1 Basic Usage

Example

import torch
import torch.nn as nn

# Create an embedding layer
# num_embeddings: vocabulary size (vocab_size)
# embedding_dim: embedding dimension (embedding_dim)
vocab_size = 10000
embedding_dim = 256

embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)

# View the shape of the embedding matrix
print(embedding.weight.shape) # torch.Size([10000, 256])

# Input word indices (LongTensor) to obtain embedding vectors
word_ids = torch.tensor([0, 1, 2, 9999]) # Arbitrary word indices
embedded = embedding(word_ids)

print(embedded.shape) # torch.Size([4, 256])
# Each word ID corresponds to a 256-dimensional vector

2.2 nn.Embedding Parameter Details

Example

import torch.nn as nn

embedding = nn.Embedding(
num_embeddings=10000, # Vocabulary size, must be greater than or equal to the maximum index value in the input
embedding_dim=256, # Embedding vector dimension, usually 50, 100, 200, 300, etc.
padding_idx=None, # Index of the padding word; the embedding vector for the padding word is all zeros
max_norm=None, # Maximum norm of the embedding vector, used for normalization
norm_type=2.0, # Normalization type, usually L2 norm
scale_grad_by_freq=False,# Scale gradients by word frequency
sparse=False, # Whether to use sparse gradients (saves GPU memory, but training is slower)
_weight=None, # Predefined weights, used for loading pretrained embeddings
)

# View the number of parameters
total_params = embedding.num_embeddings * embedding.embedding_dim
print(f"Embedding layer parameter count: {total_params:,}")
# 10000 * 256 = 2,560,000

The number of parameters in the embedding layer = vocabulary size × embedding dimension. This is a very large matrix. Usually, the embedding layer of an NLP model accounts for a large proportion of the model's total parameter count.

2.3 Padding Index padding_idx

When processing variable-length sequences, short sequences need to be padded. Usingpadding_idxcan fix the embedding vector of the padding word to a zero vector, preventing padding content from affecting the model:

Example

import torch
import torch.nn as nn

# Set padding_idx=0, meaning index 0 is the padding word
embedding = nn.Embedding(num_embeddings=10000, embedding_dim=128, padding_idx=0)

# Initialize the weights
nn.init.uniform_(embedding.weight, -0.1, 0.1)

# The embedding vector at index 0 is all zeros
word_0 = embedding(torch.tensor([0]))
print(f"Embedding of padding word: {word_0}") # All zeros

# Other indices remain normal
word_5 = embedding(torch.tensor([5]))
print(f"Embedding of word 5: {word_5}") # Non-zero values

2.4 Max Norm Normalization

Usingmax_normcan limit the norm of embedding vectors, preventing them from becoming too large during training:

Example

import torch
import torch.nn as nn

# Set the maximum norm to 1.0
embedding = nn.Embedding(num_embeddings=1000, embedding_dim=64, max_norm=1.0)

# Input arbitrary word indices
ids = torch.tensor([1, 2, 3])
embedded = embedding(ids)

# Check the L2 norm of each vector
norms = torch.norm(embedded, p=2, dim=1)
print(f"Norms of each vector: {norms}") # All values are close to 1.0

3. Loading Pretrained Word Embeddings

Using pretrained word embeddings can significantly improve model performance, especially when training data is scarce. Common pretrained word embeddings include Word2Vec, GloVe, FastText, etc.

3.1 Difference Between Training from Scratch and Loading Pretrained Embeddings

Method	Advantages	Disadvantages	Applicable Scenarios
Random initialization and training	Fully customizable, adapts to specific tasks	Requires a large amount of training data	Many domain-specific words and sufficient training data
Loading pretrained embeddings	Leverages knowledge from large-scale corpora, fast training and good results	Limited vocabulary coverage, cannot handle out-of-vocabulary words	General tasks, limited training data
Freezing pretrained embeddings	Fast training, low GPU memory usage	Embeddings cannot be fine-tuned	Limited training resources, focusing only on the upper layers of the model
Fine-tuning pretrained embeddings	Can adapt to specific tasks	Slower training, high GPU memory usage	Moderate data size, certain domain differences

3.2 Loading GloVe Pretrained Word Vectors

GloVe is a pretrained word vector released by Stanford University. The following shows how to load it:

Example

import torch
import torch.nn as nn
import numpy as np

# Simulate loading GloVe word vectors (actually need to download GloVe files)
# Assume there is already a word vector file with the format: each line contains a word followed by its vectors

def load_glove_embeddings(path, word2idx, embedding_dim=300):
"""
Load GloVe pretrained word vectors
path: path to the word vector file
word2idx: dictionary mapping words to indices
embedding_dim: word vector dimension
"""
embeddings = np.random.randn(len(word2idx), embedding_dim).astype(np.float32)
word_count = 0

with open(path, 'r', encoding='utf-8') as f:
for line in f:
values = line.strip().split()
word = values[0]
if word in word2idx:
word_idx = word2idx[word]
embeddings[word_idx] = np.asarray(values[1:], dtype=np.float32)
word_count += 1

print(f"Loaded {word_count}/{len(word2idx)} word vectors")
return torch.from_numpy(embeddings)

# Assume there is already a vocabulary
word2idx = {'hello': 0, 'world': 1, 'example': 2, 'python': 3}
EMBED_DIM = 300

# Load pretrained embeddings and create an embedding layer
# pretrained_embeddings = load_glove_embeddings('glove.6B.300d.txt', word2idx, EMBED_DIM)
# embedding = nn.Embedding.from_pretrained(pretrained_embeddings, padding_idx=0)

# Simplified example: use a randomly initialized pretrained matrix
pretrained_embeddings = torch.randn(len(word2idx), EMBED_DIM)
embedding = nn.Embedding.from_pretrained(pretrained_embeddings, padding_idx=0)

print(f"Embedding layer shape: {embedding.weight.shape}")

3.3 Freezing and Fine-Tuning the Embedding Layer

Depending on the task requirements, you can choose to freeze or fine-tune the embedding layer:

Example

import torch
import torch.nn as nn

embedding = nn.Embedding(num_embeddings=10000, embedding_dim=256)

# Method 1: Freeze the embedding layer (not participating in training)
embedding.weight.requires_grad = False
# During training, only parameters with embedding.weight.requires_grad = False will not be updated

# Method 2: Fine-tune the embedding layer (participating in training)
embedding.weight.requires_grad = True

# Method 3: Freeze partial word vectors (freeze high-frequency words like "the", "is")
# Assume the indices of high-frequency words are 0-99
embedding.weight.requires_grad = True
with torch.no_grad():
embedding.weight[0:100] *= 0 # Or assign fixed values

# Filter out frozen parameters in the optimizer
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, embedding.parameters()),
lr=1e-3
)

4. Combining Embedding Layer with RNN/LSTM

Word embedding is the first layer of an NLP model. It converts text IDs into dense vectors, which are then passed to sequence models such as RNN and LSTM for processing.

4.1 Embedding Layer + LSTM Text Classification

Example

import torch
import torch.nn as nn

class TextClassifier(nn.Module):
"""
Embedding Layer + LSTM Text Classification Model
"""
def __init__(self, vocab_size, embed_dim, hidden_size, num_classes, padding_idx=0):
super().__init__()

# Word embedding layer
self.embedding = nn.Embedding(
num_embeddings=vocab_size,
embedding_dim=embed_dim,
padding_idx=padding_idx
)

# LSTM layer
self.lstm = nn.LSTM(
input_size=embed_dim,
hidden_size=hidden_size,
num_layers=2,
batch_first=True,
bidirectional=True,
dropout=0.3
)

# Classifier
self.fc = nn.Linear(hidden_size * 2, num_classes)

def forward(self, x):
# x: (batch_size, seq_len) - word indices
embedded = self.embedding(x) # (batch_size, seq_len, embed_dim)

# LSTM output
output, (h_n, c_n) = self.lstm(embedded)

# Take the hidden state of the last time step (bidirectional concatenation)
# Forward last hidden state: h_n[-2]
# Backward last hidden state: h_n[-1]
h_combined = torch.cat([h_n[-2], h_n[-1]], dim=-1)

# Classification
logits = self.fc(h_combined)
return logits

# Instantiate the model
VOCAB_SIZE = 10000
EMBED_DIM = 128
HIDDEN_SIZE = 128
NUM_CLASSES = 2

model = TextClassifier(VOCAB_SIZE, EMBED_DIM, HIDDEN_SIZE, NUM_CLASSES)

# Simulated input: batch_size=4, sequence length=10
batch_input = torch.randint(1, VOCAB_SIZE, (4, 10))
output = model(batch_input)

print(f"Input shape: {batch_input.shape}") # torch.Size([4, 10])
print(f"Output shape: {output.shape}") # torch.Size([4, 2])

4.2 Using Pretrained Word Vectors

Example

import torch
import torch.nn as nn

# Assume pretrained word vectors have been loaded
pretrained_vectors = torch.randn(10000, 300) # Simulate pretrained vectors

# Create an embedding layer and load pretrained weights
embedding = nn.Embedding.from_pretrained(
pretrained_vectors,
padding_idx=0,
freeze=False # True: freeze and do not train, False: fine-tune
)

# When using GloVe or other pretrained word vectors, usually freeze for a few epochs first, then unfreeze for fine-tuning
# For the first few epochs, only train the upper layers of the model
embedding.weight.requires_grad = False

# After training for several epochs, unfreeze the embedding layer for fine-tuning
# embedding.weight.requires_grad = True

5. Positional Encoding

Unlike RNN/LSTM, the Transformer model does not include positional information. Positional encoding needs to be added additionally to let the model perceive the order of words in a sequence.

5.1 Principle of Positional Encoding

Positional encoding uses sine and cosine functions to generate position vectors:

\begin{aligned} PE_{(pos, 2i)} &= \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) \\ PE_{(pos, 2i+1)} &= \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) \end{aligned}

The feature of this encoding method is that encodings at different positions can be converted to each other through linear transformations, making it easy for the model to learn positional relationships.

5.2 Implementing Positional Encoding

Example

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
"""
Positional Encoding Layer
"""
def __init__(self, d_model, max_len=5000):
super().__init__()

# Create positional encoding matrix
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

# Compute the divisor term
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

# Use sin for even indices and cos for odd indices
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)

# Add a batch dimension and register it as a buffer that does not participate in gradient computation
pe = pe.unsqueeze(0) # (1, max_len, d_model)
self.register_buffer('pe', pe)

def forward(self, x):
"""
x: (batch_size, seq_len, d_model)
"""
seq_len = x.size(1)
# Slice the positional encoding of the corresponding length and add it
x = x + self.pe[:, :seq_len, :]
return x

# Usage example
d_model = 256
max_len = 100

pos_encoding = PositionalEncoding(d_model, max_len)

# Simulated input: batch_size=4, seq_len=20, d_model=256
x = torch.randn(4, 20, d_model)
x = pos_encoding(x)

print(f"Input shape: {x.shape}") # torch.Size([4, 20, 256])

Positional encoding is a key component in the Transformer architecture. It enables the model to distinguish words at different positions, even if their embedding vectors are the same.

5.3 Learnable Positional Encoding

In addition to fixed positional encoding, learnable positional encoding can also be used:

Example

import torch
import torch.nn as nn

class LearnablePositionalEncoding(nn.Module):
"""
Learnable Positional Encoding
"""
def __init__(self, d_model, max_len=5000):
super().__init__()
self.pos_embedding = nn.Embedding(max_len, d_model)

def forward(self, x):
batch_size, seq_len, d_model = x.size(1), x.size(1), x.size(2)
# Create position indices [0, 1, 2, ..., seq_len-1]
positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
pos_encoded = self.pos_embedding(positions)
return x + pos_encoded

# Usage example
pos_encoding = LearnablePositionalEncoding(d_model=256, max_len=100)
x = torch.randn(4, 20, 256)
x = pos_encoding(x)
print(f"Output shape: {x.shape}") # torch.Size([4, 20, 256])

6. Advanced Techniques for Embedding Layers

6.1 Reducing GPU Memory Usage

When the vocabulary is very large, the embedding layer consumes a lot of GPU memory. The following techniques can be used to optimize:

Example

import torch
import torch.nn as nn

# Technique 1: Use sparse gradients (sparse=True)
embedding = nn.Embedding(
num_embeddings=100000,
embedding_dim=256,
sparse=True # Gradients are stored in sparse format, saving GPU memory
)

# Technique 2: Use quantization
# Convert float32 to int8 or float16
embedding_int8 = embedding.to(torch.int8)

# Technique 3: Freeze infrequently used word vectors
# In a large-scale vocabulary, only train high-frequency words; keep low-frequency words frozen
embedding = nn.Embedding(num_embeddings=100000, embedding_dim=256)
embedding.weight.requires_grad = True

# Freeze word vectors with index greater than 50000
with torch.no_grad():
embedding.weight[50000:] *= 0
embedding.weight.requires_grad = False

# Only train high-frequency words
embedding.weight[1:50000].requires_grad = True

6.2 Handling Out-of-Vocabulary (OOV) Words

Words not in the vocabulary (out-of-vocabulary / OOV words) may appear in the test set and require special handling:

Example

import torch
import torch.nn as nn

class EmbeddingWithOOV(nn.Module):
"""
Embedding Layer Supporting Out-of-Vocabulary Words
"""
def __init__(self, vocab_size, embed_dim, oov_idx=None):
super().__init__()
self.embedding = nn.Embedding(vocab_size + 1, embed_dim, padding_idx=0)
self.oov_idx = oov_idx if oov_idx is not None else vocab_size # Use the last index as OOV

def forward(self, x):
# Replace the indices of OOV words with oov_idx
oov_mask = x >= self.vocab_size
x = x.clone()
x[oov_mask] = self.oov_idx
return self.embedding(x)

# Use hashing techniques to handle larger vocabularies
class HashEmbedding(nn.Module):
"""
Embedding Layer Using Hashing to Handle Vocabularies of Arbitrary Size
"""
def __init__(self, num_buckets, embed_dim):
super().__init__()
self.num_buckets = num_buckets
self.embedding = nn.Embedding(num_buckets, embed_dim)

def forward(self, x):
# Hash word indices into buckets
# Use Python dictionary's hashing method
hashed = torch.remainder(x, self.num_buckets)
return self.embedding(hashed)

6.3 Subword Embedding

For morphologically rich languages (such as German and Russian), subword embedding can effectively handle the out-of-vocabulary problem:

Example

import torch
import torch.nn as nn

class SubwordEmbedding(nn.Module):
"""
Simplified Subword Embedding Example
In practice, algorithms such as BPE and WordPiece should be used for tokenization
"""
def __init__(self, vocab_size, embed_dim, char_dim=50):
super().__init__()
# Character-level embedding
self.char_embedding = nn.Embedding(vocab_size, char_dim)
# Word-level embedding
self.word_embedding = nn.Embedding(vocab_size, embed_dim - char_dim)
# Character-level LSTM for combining character vectors
self.char_lstm = nn.LSTM(
char_dim, char_dim,
batch_first=True, bidirectional=True
)

def forward(self, word_ids, char_ids):
"""
word_ids: word indices (batch_size, seq_len)
char_ids: character indices (batch_size, seq_len, max_word_len)
"""
# Word embedding
word_emb = self.word_embedding(word_ids)

# Character embedding + LSTM
batch_size, seq_len, max_word_len = char_ids.shape
char_ids_flat = char_ids.view(-1, max_word_len) # (batch_size * seq_len, max_word_len)
char_emb = self.char_embedding(char_ids_flat) # (batch_size * seq_len, max_word_len, char_dim)

char_output, (h_n, _) = self.char_lstm(char_emb)
# Take bidirectional final hidden states and concatenate
char_rep = torch.cat([h_n[-2], h_n[-1]], dim=-1) # (batch_size * seq_len, char_dim * 2)

char_rep = char_rep.view(batch_size, seq_len, -1) # (batch_size, seq_len, char_dim * 2)

# Concatenate word embedding and character representation
combined = torch.cat([word_emb, char_rep], dim=-1)
return combined

7. Complete Practical Example: Text Classification Model

Below is a complete text classification model example, including an embedding layer, LSTM, and classifier:

Example

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ── Configuration Parameters ──────────────────────────────────────
VOCAB_SIZE = 10000
EMBED_DIM = 128
HIDDEN_SIZE = 128
NUM_LAYERS = 2
NUM_CLASSES = 5
DROPOUT = 0.3
MAX_LEN = 200

# ── Model Definition ──────────────────────────────────────
class TextClassificationModel(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_size,
num_layers, num_classes, dropout=0.3, padding_idx=0):
super().__init__()

# Embedding layer
self.embedding = nn.Embedding(
vocab_size,
embed_dim,
padding_idx=padding_idx
)

# BiLSTM
self.lstm = nn.LSTM(
embed_dim,
hidden_size,
num_layers=num_layers,
batch_first=True,
bidirectional=True,
dropout=dropout if num_layers > 1 else 0
)

# Classifier
self.dropout = nn.Dropout(dropout)
self.fc = nn.Linear(hidden_size * 2, num_classes)

def forward(self, x):
# x: (batch_size, seq_len)
embedded = self.embedding(x) # (batch_size, seq_len, embed_dim)

# LSTM
output, (h_n, c_n) = self.lstm(embedded)

# Concatenate bidirectional final hidden states
h_forward = h_n[-2] # (batch_size, hidden_size)
h_backward = h_n[-1] # (batch_size, hidden_size)
h_combined = torch.cat([h_forward, h_backward], dim=-1)

# Classification
dropped = self.dropout(h_combined)
logits = self.fc(droped)
return logits

# ── Dataset ──────────────────────────────────────
class TextDataset(Dataset):
def __init__(self, texts, labels, vocab, max_len=200):
self.texts = texts
self.labels = labels
self.vocab = vocab
self.max_len = max_len

def __len__(self):
return len(self.texts)

def __getitem__(self, idx):
text = self.texts[idx][:self.max_len]
# Convert words to indices, unknown words represented by 1 (<UNK>)
ids = [self.vocab.get(word, 1) for word in text]
# Pad with zeros to a fixed length
if len(ids) < self.max_len:
ids += [0] * (self.max_len - len(ids))
return torch.tensor(ids), torch.tensor(self.labels[idx])

# ── Training Function ──────────────────────────────────────
def train_epoch(model, loader, optimizer, criterion, device):
model.train()
total_loss = 0
correct = 0
total = 0

for inputs, labels in loader:
inputs, labels = inputs.to(device), labels.to(device)

optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

total_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)

return total_loss / total, correct / total

# ── Initialization and Training ──────────────────────────────────
# Assume a vocabulary already exists
word2idx = {'<PAD>': 0, '<UNK>': 1} # The vocabulary needs to be built based on the actual corpus

model = TextClassificationModel(
VOCAB_SIZE, EMBED_DIM, HIDDEN_SIZE,
NUM_LAYERS, NUM_CLASSES, DROPOUT
).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Simulated training
print("Start training...")
for epoch in range(10):
train_loss, train_acc = train_epoch(model, None, optimizer, criterion, device)
print(f"Epoch {epoch+1}: Loss={train_loss:.4f}, Acc={train_acc:.4f}")

8. API Quick Reference

8.1 Common Operations on nn.Embedding

Operation	Code
Create embedding layer	`nn.Embedding(num_embeddings, embedding_dim)`
Load pretrained embeddings	`nn.Embedding.from_pretrained(weights)`
Look up table to get embeddings	`embedding(word_ids)`
Freeze embeddings	`embedding.weight.requires_grad = False`
Get embedding vectors	`embedding.weight[idx]`

8.2 Pretrained Word Vector Resources

Resource	Dimension	Features
GloVe	50, 100, 200, 300	Word co-occurrence statistics, fast training
Word2Vec	50-500	Trained by Google, wide coverage
FastText	300	Supports subwords, good at handling OOV
BERT	768+	Context-dependent, best performance

8.3 Recommendations for Choosing an Embedding Layer

数据量小（< 10K）
    -> 使用预训练词向量（GloVe/FastText）+ 冻结或微调

数据量中等（10K ~ 100K）
    -> 使用预训练词向量 + 微调

数据量大（> 100K）
    -> 可考虑从头训练，或使用大规模预训练模型

领域差异大
    -> 使用领域相关预训练模型或增量训练

资源受限
    -> 冻结嵌入层，使用较小的嵌入维度

Other Extensions