7.3.4 Vector Embeddings of Simple Text Features

Why dense vectors?

In TF-IDF, each SMS message became a vector with 7,618 entries, almost all of them zero. These sparse vectors work really well with linear models, as we saw. But they have one basic limit: every word is its own column, with no connection to any other column. To a TF-IDF vector, prize, prizes and cash are just as different from each other as prize and lunch. Nothing in it says that some words are alike.

That is the gap this module fills. An embedding is a short, dense vector of real numbers, say 16 or 100 of them, built so that similar words or documents get similar vectors. It is also how every modern language model represents text: the first layer of a transformer is a table of embeddings. We will first learn to measure how close two vectors are, with cosine similarity, and use it to find similar messages. Then we will train word embeddings ourselves in PyTorch on the SMS data and compare them honestly with a bag of words.

What you need to know

Here is what I would like you to take away from this module. We will meet each idea below, so don't worry if some are new.

Learn it properly

This page only touches lightly on where good word vectors come from. The lecture below fills in that background.

Pick one main course

Our pick: Lena Voita's Word Embeddings lecture

Read the first parts of Word Embeddings from Lena Voita's free NLP Course | For You: Represent as Discrete Symbols: One-hot Vectors, then Distributional Semantics. Skip Count-Based Methods, and read Word2Vec: a Prediction-Based Method, at least up to how it is trained. It is visual and careful, and it covers exactly the step from sparse counts to dense vectors that we take here. The later sections are optional.

If you prefer something else

• The textbook: Chapter 5 of Speech and Language Processing, Embeddings, covers count-based vectors, cosine similarity (section 5.4) and word2vec, with exercises.
• In PyTorch: the official tutorial Word Embeddings: Encoding Lexical Semantics explains nn.Embedding and trains embeddings with a small n-gram language model.

A quick tour

Go through your chosen course first, then open a notebook and run each cell as you read. Training takes only a few seconds on a laptop CPU.

One-hot vectors

Let us start with the simplest possible vector for a word, and see exactly what is missing. A one-hot vector is as long as the vocabulary, with a 1 at the word's index and 0 everywhere else. With a vocabulary of five words, word 1 and word 3 look like this:

import numpy as np

one_hot = np.eye(5)          # row i is the one-hot vector of word i
print(one_hot[1], one_hot[3])
print(one_hot[1] @ one_hot[3], one_hot[1] @ one_hot[1])

[0. 1. 0. 0. 0.] [0. 0. 0. 1. 0.]
0.0 1.0

np.eye(5) is the 5 × 5 identity matrix, whose rows are exactly the one-hot vectors. The dot product, written @, of two different one-hot vectors is always 0, and of a vector with itself always 1. So one-hot vectors can say "same word" or "different word", and nothing more. A count vector from Bag of Words is just the sum of the one-hot vectors of a document's words, so it has the same blind spot.

Dense vectors and cosine similarity

A dense vector does better by spreading the meaning of a word over a few real numbers. To get a feel for it, here are three made-up vectors of three numbers each. I chose them by hand so that two of them point in a similar direction:

prize = np.array([0.9, 0.8, 0.1])     # made-up vectors
cash = np.array([0.8, 0.9, 0.2])
lunch = np.array([0.1, 0.2, 0.9])

def cosine_similarity(a, b):
    return (a @ b) / (np.linalg.norm(a) * np.linalg.norm(b))

print(round(cosine_similarity(prize, cash), 3))
print(round(cosine_similarity(prize, lunch), 3))

0.99
0.303

Cosine similarity is the dot product of two vectors divided by the product of their lengths ($\lVert a \rVert$, computed by np.linalg.norm):

$$\cos(a, b) = \frac{a \cdot b}{\lVert a \rVert \, \lVert b \rVert}$$

It is the cosine of the angle between the two vectors: 1 when they point the same way, 0 when they have nothing in common, and −1 when they point in opposite directions. Dividing by the lengths means only the direction counts, so a message and the same message written twice get exactly 1. Here, prize and cash come out very similar (0.99), and prize and lunch do not (0.303).

Finding similar messages

Cosine similarity works on sparse vectors too, and that turns TF-IDF into a small search engine. Let us load the SMS data with the split from Bag of Words (download sms.csv; SMS Spam Collection, CC BY 4.0) and look for the training messages closest to a made-up one:

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity as all_pairs

sms = pd.read_csv("sms.csv")
train, test = train_test_split(sms, test_size=0.2,
                               stratify=sms["label"],
                               random_state=0)

sms_tfidf = TfidfVectorizer()
train_vectors = sms_tfidf.fit_transform(train["text"])

query = "You have won a free prize, call now to claim"
query_vector = sms_tfidf.transform([query])
similarity = all_pairs(query_vector, train_vectors)[0]
for i in similarity.argsort()[::-1][:3]:
    print(round(similarity[i], 3), train["label"].iloc[i],
          train["text"].iloc[i][:60])

0.496 spam You have WON a guaranteed £1000 cash or a £2000 prize.To cla
0.437 spam You have won ?1,000 cash or a ?2,000 prize! To claim, call09
0.416 spam You have WON a guaranteed £1000 cash or a £2000 prize. To cl

all_pairs compares the query with every training message at once (scikit-learn's TF-IDF rows already have length 1). argsort() sorts the positions from least to most similar, [::-1] reverses that, and [:3] keeps the top three. All three neighbours are spam. If you predicted a label from the labels of the most similar messages, you would have rebuilt K-nearest neighbours with cosine similarity.

Word embeddings in PyTorch: nn.Embedding

So far, I picked the dense vectors by hand. Where do real ones come from? One answer is to let a neural network learn a vector for every word, together with the classifier. In PyTorch, the table of word vectors is nn.Embedding: a matrix with one row per token id, random at first and then trained like any other weights. Looking up a batch of ids gives back their rows:

import torch
from torch import nn

torch.manual_seed(0)
embedding = nn.Embedding(num_embeddings=6, embedding_dim=4,
                         padding_idx=0)
ids = torch.tensor([[3, 5, 1, 0, 0]])     # one padded message
vectors = embedding(ids)
print(vectors.shape)
print(embedding.weight[0])

torch.Size([1, 5, 4])
tensor([0., 0., 0., 0.], grad_fn=<SelectBackward0>)

The table has 6 rows (a vocabulary of 6 ids) of 4 numbers each. A batch of shape (1, 5), one message of 5 ids, becomes a tensor of shape (1, 5, 4): every id is replaced by its vector. padding_idx=0 tells PyTorch that id 0 is <pad>, so its row stays zero and is never updated.

To classify a whole message, though, we need one vector per message, not one per token. The simplest choice is the mean of its token vectors, counting real tokens only, which is what the mask from Tokenization is for. nn.EmbeddingBag with mode="mean" does the lookup and this masked average in one step, so that is what our model will use.

The SMS messages as token ids

The network needs the messages as padded id tensors, built as in Tokenization, and a validation set to choose the best epoch, as in A Simple MLP in PyTorch. So we split 20% off the training messages and leave the test messages alone until the end:

import re
from collections import Counter

fit_part, val_part = train_test_split(train, test_size=0.2,
                                      stratify=train["label"],
                                      random_state=0)

def tokenize(text):
    return re.findall(r"\w+", text.lower())

word_counts = Counter()
for text in fit_part["text"]:
    word_counts.update(tokenize(text))
vocabulary = {"<pad>": 0, "<unk>": 1}
for word, count in word_counts.most_common():
    if count >= 2:
        vocabulary[word] = len(vocabulary)

print(len(fit_part), len(val_part), len(vocabulary))

3308 828 2923

The vocabulary comes from the 3,308 fitting messages only, as Tokenization recommends: 2,921 words seen at least twice, plus <pad> and <unk>. Now each message becomes 40 ids:

MAX_LENGTH = 40

def encode(text):
    ids = [vocabulary.get(token, 1) for token in tokenize(text)]
    ids = ids[:MAX_LENGTH]
    return ids + [0] * (MAX_LENGTH - len(ids))

def to_tensors(part):
    ids = torch.tensor([encode(text) for text in part["text"]])
    is_spam = (part["label"] == "spam").values
    labels = torch.tensor(is_spam, dtype=torch.float32)
    return ids, labels

fit_ids, fit_labels = to_tensors(fit_part)
val_ids, val_labels = to_tensors(val_part)
test_ids, test_labels = to_tensors(test)
print(fit_ids.shape, val_ids.shape, test_ids.shape)

torch.Size([3308, 40]) torch.Size([828, 40]) torch.Size([1035, 40])

The labels are float32, as BCEWithLogitsLoss requires for a yes/no label (PyTorch Pipeline for a Simple Machine Learning Problem explains why).

A model that learns word vectors

The model is small. It has two layers: an EmbeddingBag that turns a message into the mean of its word vectors, and an nn.Linear layer that turns that mean into one logit, positive for spam:

from torch.utils.data import TensorDataset, DataLoader

class SpamEmbeddingModel(nn.Module):
    def __init__(self, vocabulary_size, embedding_dim):
        super().__init__()
        self.bag = nn.EmbeddingBag(vocabulary_size, embedding_dim,
                                   mode="mean", padding_idx=0)
        self.out = nn.Linear(embedding_dim, 1)

    def forward(self, ids):
        message_vectors = self.bag(ids)          # (batch, dim)
        return self.out(message_vectors).squeeze(1)

torch.manual_seed(0)
train_loader = DataLoader(TensorDataset(fit_ids, fit_labels),
                          batch_size=64, shuffle=True)
model = SpamEmbeddingModel(len(vocabulary), embedding_dim=16)
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
sum(p.numel() for p in model.parameters())

46785

Almost all of the 46,785 weights sit in the embedding table: $2923 \times 16 = 46{,}768$, plus 16 weights and a bias in the output layer. The model is close to the logistic regression of Bag of Words, except that each word is first squeezed into 16 numbers, so words that behave alike can share what they learn.

The training and evaluation functions are the ones from PyTorch Pipeline for a Simple Machine Learning Problem, except that the evaluation reports F1, our metric for spam, and takes whole tensors, since 828 messages fit in one batch:

import copy
from sklearn.metrics import f1_score

def train_one_epoch(model, loader, loss_fn, optimizer):
    model.train()
    total_loss = 0.0
    for id_batch, label_batch in loader:
        loss = loss_fn(model(id_batch), label_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * len(label_batch)
    return total_loss / len(loader.dataset)

def evaluate(model, ids, labels, loss_fn):
    model.eval()
    with torch.no_grad():
        logits = model(ids)
        loss = loss_fn(logits, labels).item()
    predictions = (logits > 0).int()
    return loss, f1_score(labels.int(), predictions)

best_val_loss = float("inf")
for epoch in range(1, 21):
    train_loss = train_one_epoch(model, train_loader, loss_fn,
                                 optimizer)
    val_loss, val_f1 = evaluate(model, val_ids, val_labels, loss_fn)
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_epoch = epoch
        best_state = copy.deepcopy(model.state_dict())
    if epoch % 4 == 0:
        print(epoch, round(train_loss, 3), round(val_loss, 3),
              round(val_f1, 3))

model.load_state_dict(best_state)
print("best epoch:", best_epoch)

4 0.054 0.07 0.929
8 0.018 0.055 0.928
12 0.007 0.055 0.937
16 0.004 0.056 0.932
20 0.002 0.059 0.932
best epoch: 10

The columns are the epoch, the training loss, the validation loss and the validation F1. The training loss keeps falling, but the validation loss is lowest at epoch 10 and then creeps up: the model is starting to memorise the 3,308 messages, as in A Simple MLP in PyTorch. So we keep the weights from epoch 10.

An honest comparison

Now the test messages, used exactly once. To make the comparison fair, we also train the bag-of-words model from Bag of Words on the same 3,308 messages:

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.metrics import precision_score, recall_score

y_test = (test["label"] == "spam").astype(int)

def report(name, y_true, y_pred):
    precision = precision_score(y_true, y_pred)
    recall = recall_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)
    print(f"{name:<14} precision {precision:.3f}  "
          f"recall {recall:.3f}  F1 {f1:.3f}")

fit_labels_01 = (fit_part["label"] == "spam").astype(int)
bag_of_words = make_pipeline(CountVectorizer(),
                             LogisticRegression(max_iter=1000))
bag_of_words.fit(fit_part["text"], fit_labels_01)
report("bag of words", y_test, bag_of_words.predict(test["text"]))

model.eval()
with torch.no_grad():
    test_predictions = (model(test_ids) > 0).int()
report("embeddings", y_test, test_predictions)

bag of words   precision 0.975  recall 0.885  F1 0.928
embeddings     precision 0.928  recall 0.885  F1 0.906

Both catch 116 of the 131 spam messages, but the embedding model raises 9 false alarms against 3, so its F1 is lower: 0.906 against 0.928. And it is not bad luck with the seed: with seeds 1 to 4 instead of 0, the same code gave test F1 scores between 0.898 and 0.919, all below the bag of words.

So on this task, learned embeddings do not beat a plain bag of words. That is a common result for short texts and a few thousand examples: the linear model needs one weight per word, while the network must learn 16 numbers per word from the same data. Embeddings pay off with much more data, with models that use word order, or when the vectors come pretrained on a huge corpus. In a contest, always check the simple baseline before reaching for a network.

What the word vectors learned

What did the word vectors actually learn? The rows of model.bag.weight are the 16-number vectors of the vocabulary words. We scale every row to length 1, so that a single matrix product gives the cosine similarities of one word with all the others at once:

word_vectors = model.bag.weight.detach()
lengths = word_vectors.norm(dim=1, keepdim=True).clamp(min=1e-8)
unit_vectors = word_vectors / lengths
id_to_word = {i: word for word, i in vocabulary.items()}

def nearest_words(word, k=5):
    similarity = unit_vectors @ unit_vectors[vocabulary[word]]
    closest = similarity.argsort(descending=True)[1:k + 1]
    return [id_to_word[int(i)] for i in closest]

for word in ["prize", "claim", "home", "sorry"]:
    print(word, nearest_words(word))

prize ['recent', '16', 'uk', 'fri', 'asp']
claim ['2004', 'min', 'txt', '16', 'reward']
home ['peaceful', 'borrow', 'ability', 'accept', 'daily']
sorry ['can', 'continue', 'bslvyl', 'remove', 'scold']

.clamp(min=1e-8) avoids dividing by zero for the all-zero <pad> row, and [1:k + 1] skips the word itself. The neighbours of prize and claim are other spam words (uk, txt, reward, numbers), but not synonyms, and those of home and sorry look random. That makes sense: the vectors were trained only to answer "spam or ham?", so what they encode reliably is how spam-like a word is. We can check this by projecting every word vector onto the direction the output layer uses:

spam_direction = model.out.weight.detach()[0]
spam_scores = word_vectors @ spam_direction
top = spam_scores.argsort(descending=True)[:8]
print([id_to_word[int(i)] for i in top])

['150p', 'uk', 'claim', 'service', 'text', '18', 'win', 'txt']

These are much the same words as the largest logistic regression weights in Bag of Words. Learned embeddings capture what their task needs, and nothing more. Vectors in which prize lies near reward need far larger corpora and a task about meaning, such as predicting a word from its neighbours (word2vec).

Pretrained embeddings

That is why, in practice, you often start from vectors someone else trained on billions of words. Three libraries for this are on the IOAI 2026 list of available libraries:

• gensim loads and trains classic word vectors such as word2vec: one fixed vector per word.
• fastText builds a word's vector from its character n-grams, like those in N-grams, so it can handle words it never saw. It also has a fast built-in text classifier.
• sentence-transformers turns a whole sentence into one vector: model.encode(list_of_texts) returns one row per text, ready for logistic regression or cosine-similarity search.

One catch: all three need model files, from tens of megabytes to several gigabytes, normally downloaded on first use. A contest has no free internet; the organisers publish a list of pre-downloaded models that you load from a local folder. So check that list before planning around pretrained vectors, and practise loading a model from a local path at home.

Common mistakes

How to practise

The best way to get a feel for embeddings is to experiment with the model you just trained. Keep the same SMS split for everything, so your results stay comparable:

• Change embedding_dim to 2, train, and plot every word as a point using its two numbers. Where do claim, txt and prize end up?
• Replace the mean with a sum (mode="sum") or a maximum (mode="max"). Which works best on the validation set?
• Train the embedding model with five different seeds and report the mean and spread of the validation F1, as you would for any small neural network.
• Use K-nearest neighbours with metric="cosine" on TF-IDF vectors, and compare it with logistic regression on the same split.

Check yourself

What comes next

That wraps up Part 7: from raw characters to word vectors learned by a neural network, and a reminder that the simple baseline is often hard to beat. Here is where to go next:

• Waveforms, Sampling Rate, Duration, Amplitude, and Channels: the next modality, sound, which also becomes a sequence of numbers.
• Model Evaluation Metrics and Cross-Validation: the tools behind every honest comparison in Part 7.
• A Simple MLP in PyTorch: revisit the training loop used on this page, and try adding a hidden layer between the embedding and the output.