Embeddings: from IDs to vectors

An embedding represents an item as a vector of numbers. In a language model, a token ID selects one row from a learned embedding table. The ID is an index; its numerical size does not describe the token’s meaning.

The table learns with the model

During training, gradients update the selected rows along with the other model parameters. The vectors become useful because of the prediction task the network learns. Individual coordinates usually do not have simple names such as “animal” or “past tense.” Information can be spread across many coordinates.

A token starts with the same table row wherever it occurs. Later layers combine it with position and surrounding context. The representation of “bank” can therefore develop differently in a sentence about a river and one about money.

Feed the vectors into a model

With torch and transformers installed, this example makes the lookup explicit, then passes the vectors into GPT-2:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")
inputs = tokenizer("The river bank", return_tensors="pt")

with torch.inference_mode():
    vectors = model.get_input_embeddings()(inputs.input_ids)
    scores = model(inputs_embeds=vectors).logits[0, -1]

next_id = scores.argmax().item()
continuation = tokenizer.decode([*inputs.input_ids[0].tolist(), next_id])

vectors has a batch dimension, one position per input token and one coordinate per embedding dimension. Supplying inputs_embeds lets the model continue from these vectors. GPT-2 adds its positional information internally before processing them through its layers.

The resulting scores still predict the next token. The embedding is the numerical starting representation that makes those computations possible.

Text-retrieval systems also use embeddings, but typically represent whole passages with a model trained for that purpose. A raw token-table row is not automatically a useful document-search representation.

Related: tokens, attention, vector database.