Embeddings and sequence models

From tokens to vectors to predictions

Text reaches a model as token ids: integers produced by a tokenizer, one per word or sub-word. nn.Embedding(vocab_size, dim) is a learnable lookup table that maps each id to a dense vector. Its input must be an integer tensor; its output adds a final dimension of size dim.

Recurrent layers such as nn.LSTM and nn.GRU read a sequence step by step while carrying a hidden state. Pass batch_first=True to use the familiar (N, T, features) layout; the default is (T, N, features). An LSTM returns all per-step outputs plus the final hidden and cell states, and for classification the final hidden state is a convenient summary.

Transformers (nn.TransformerEncoderLayer, nn.MultiheadAttention) use self-attention to look at all positions at once, which parallelises well on GPUs; they need positional information added to the embeddings. In practice, most NLP work today fine-tunes a pretrained transformer, for example from the Hugging Face transformers library, whose models are ordinary PyTorch nn.Modules you can train with the same loop.

An LSTM text classifier

h_n[-1] is the final hidden state of the last LSTM layer, one vector per sentence. padding_idx=0 keeps the padding embedding fixed at zero.

import torch
from torch import nn

class TextClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim=64, hidden=128, num_classes=2):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden, batch_first=True)
        self.fc = nn.Linear(hidden, num_classes)

    def forward(self, token_ids):              # (N, T) int64
        x = self.embed(token_ids)              # (N, T, embed_dim)
        _, (h_n, _) = self.lstm(x)             # h_n: (num_layers, N, hidden)
        return self.fc(h_n[-1])                # (N, num_classes)

model = TextClassifier(vocab_size=5000)
ids = torch.randint(1, 5000, (8, 20))          # 8 sentences of 20 tokens
print(model(ids).shape)                        # torch.Size([8, 2])

Padding leaks into the final state

If sentences are padded to the same length, a plain LSTM also reads the padding, and its final state reflects those zeros. Wrap the input with nn.utils.rnn.pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False) so the LSTM stops at each sentence's real end.

Quick check: What must the input to `nn.Embedding` be?

• Float vectors of size `embed_dim`
• An integer tensor of token indices
• One-hot encoded float tensors
• Normalised values between 0 and 1