Embeddings: things as vectors

A label number would invent an order that is not there

Say a network must read one of four kinds of object: cat, dog, car and truck. The simplest idea is to number them 0, 1, 2 and 3 and feed in that number. This misleads the network. A linear layer multiplies its input by weights, so it would treat "truck" as three times "dog", and "car" as halfway between "dog" and "truck". None of that is true.

An embedding avoids the problem. Each item gets its own vector, and no item's vector depends on another item's number. The network can then place cat near dog and car near truck, if that helps it.

An embedding table holds one learned row per item

An embedding table has one row for each item in the list and d columns. The number d is the embedding size. To embed an item, the network looks up its row. Real tables are large: for example, 50,000 words with d = 512. The figure uses four items and d = 2, so each row is a point you can draw.

Predict first. Is "cat" closer to "dog" or to "truck"? Guess from the picture, then check with numbers.

Looking up a row is the same as a linear layer without a bias applied to a one-hot vector: a vector with a 1 in the item's place and 0 everywhere else. For dog, the one-hot vector is [0, 1, 0, 0]. Each output is a weighted sum in which only dog's weight survives, so the result is dog's row. Libraries skip the multiplication and look the row up directly; PyTorch calls the table torch.nn.Embedding.

A point already has numbers, and a linear layer lifts them to d numbers

A 3D point p = (x, y, z) is already a vector, but only of 3 numbers. Networks usually want longer vectors, such as 64 or 256 numbers, so that later layers have room to store what they learn about the point. A linear layer with weights W (d rows, 3 columns) and a bias b makes the embedding Wp + b.

Once every item has a vector of the same size d, the network treats them all alike. Each such vector is called a token. A model of a point cloud often makes one token per point, or one per small group of nearby points.

A linear layer gives nearby points nearby embeddings. The layers after it then find fine detail, such as a sharp change between two close points, slow to learn. Many models therefore first expand each coordinate into waves of several frequencies, called Fourier features, and embed those instead.

Training moves the vectors to where they help

At the start of training, the table and the weights hold small random numbers, so cat and truck are as close as cat and dog. Each training step nudges every weight that took part, table rows included, in the direction that lowers the error (see gradient descent). Items that the task treats alike tend to drift together.

embed.py

import numpy as np

E = np.array([[ 0.9, 0.2],    # cat
              [ 0.8, 0.3],    # dog
              [-0.7, 0.6],    # car
              [-0.8, 0.5]])   # truck

one_hot = np.array([0, 1, 0, 0])
print(E[1], one_hot @ E)      # look up dog's row, two ways

def dist(a, b):
    return np.linalg.norm(E[a] - E[b])
print(round(dist(0, 1), 3), round(dist(0, 3), 3))

W = np.array([[0.5, -0.2, 0.1], [0.0, 0.3, 0.8],
              [-0.4, 0.6, 0.2], [0.7, 0.1, -0.5]])
b = np.array([0.1, 0.0, -0.2, 0.3])
print(W @ np.array([1.0, 2.0, -1.0]) + b)   # embed a 3D point

[0.8 0.3] [0.8 0.3]
0.141 1.726
[ 0.1 -0.2  0.4  1.7]