🎓 A1Learner
🔍

PyTorch Why Embeddings

An embedding is a short list of numbers that stands for one word. Embeddings replace one-hot rows, which grow with the vocabulary and say nothing about which words are alike.

One word as a one-hot row

A one-hot row is a row of zeros with a single 1 that marks the word's position. The row for word 4213 in a 10,000-word vocabulary holds 10,000 numbers.

row = F.one_hot(torch.tensor([4213]), num_classes=10_000)
print(row.shape)
print(int(torch.count_nonzero(row)))

Two positive words, zero similarity

great and lovely are both positive words. Their one-hot rows still multiply to a dot product of zero.

great = F.one_hot(torch.tensor([4213]), num_classes=10_000).float()
lovely = F.one_hot(torch.tensor([87]), num_classes=10_000).float()
print(float(great @ lovely.T))

The embedding side of the bill

An embedding table gives each word one short row. A word then costs 16 numbers instead of 10,000.

table = torch.empty(10_000, 16)
print(tuple(table.shape))
print(table.numel())

Example

import torch
import torch.nn.functional as F

vocab_size = 10_000
dim = 16

one_word = F.one_hot(torch.tensor([4213]), num_classes=vocab_size)
print("Numbers in one one-hot row:", one_word.numel())
print("Non-zero numbers in that row:", int(torch.count_nonzero(one_word)))

review = torch.tensor([4213, 87, 999, 12, 500, 77, 3000, 42, 9, 100, 2024, 555])
review_rows = F.one_hot(review, num_classes=vocab_size)
print("Numbers in a 12-word review as one-hot rows:", review_rows.numel())
print("Non-zero numbers in that review:", int(torch.count_nonzero(review_rows)))

great = F.one_hot(torch.tensor([4213]), num_classes=vocab_size).float()
lovely = F.one_hot(torch.tensor([87]), num_classes=vocab_size).float()
print("Dot product of the great and lovely rows:", float(great @ lovely.T))

table = torch.empty(vocab_size, dim)
print("Embedding table shape:", tuple(table.shape))
print("Numbers in the whole table:", table.numel())
print("Numbers per word from the table:", dim)
print("Numbers in the same review as embedding rows:", len(review) * dim)

Key points

  • One-hot row: a vocabulary-sized row of zeros with one 1 at the word's position
  • A one-hot row costs one number per vocabulary word
  • Any two different words have one-hot rows with a dot product of 0
  • Embedding: a short fixed-length list of numbers per word
  • An embedding table for 10,000 words at width 16 holds 160,000 numbers

In this example, F.one_hot builds the row for word 4213 in a 10,000-word vocabulary. The row holds 10,000 numbers and one nonzero entry. A 12-word review as one-hot rows holds 120,000 numbers with 12 nonzero entries. The dot product of the great and lovely rows prints as 0.0, because distinct one-hot rows never overlap. An embedding table of width 16 holds 160,000 numbers in total, shared by all words. The same review as embedding rows costs 12 times 16, which is 192 numbers.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What is wrong with one-hot rows for a large vocabulary?