Embedding

Once inputs have been discretized by the tokenizer, we want to compress them further so that they're easier to work with. The (one-hot) dimensionality of the token space is equivalent to the vocabulary size, but this space is colossal and far too sparse - we want to reduce the vector space to a more efficient representation. To do this, we want to learn an optimal mapping between tokens and (what are called) embeddings, which are much lower dimension fixed length vectors that are passed into the remainder of the model.

The embedding matrix is the first truly learned part of an LLM, which is to say that this representation mapping is adjusted over the course of a training run along side every other model weight and bias. The size of the embedding matrix is equivalent to the size of the tokenizer's vocabulary multiplied by the dimensionality of the model. Some examples:

Embedding Vector Space

Embedding matrices learn to represent something approximating the fundamental identity/meaning of a token. They learn this meaning in the context of a high dimension vector space - we can loosely imagine various dimensions in this vector space corresponding to different concepts. Take the following single word analogies for example:

• king is to queen as man is to woman

• paris is to france as tokyo is to japan

• meter is to kilometer as foot is to mile

• cat is to dog as feline is to canine

In a coherent analogy "A is to B as X is to Y", there is something consistent or parallel about the relationship between A & B and X & Y. If we imagine Paris and France as some points in space, it would make sense that their relative positions share some similarity to the more abstract conceptual relationship X is the largest city in country Y. Indeed, this conceptual relationship is strong enough to make the following analogies trivial to complete:

• Paris is to France as Berlin is to ...

• Paris is to France as Rome is to ...

• Paris is to France as London is to ...

This is precisely the behavior we observe from learned embeddings. See the following simplified projections from the GloVe project:

Skip-Gram with Negative Sampling (SGNS)

The embedding matrix training example I've included uses an approach called Skip-Gram with Negative Sampling (SGNS). This approach has the model use input tokens to predict surrounding context tokens (over a specified window). This pushes tokens that occur together to cluster together in vector space. SGNS, as the name suggests, also uses negative sampling, which means that for every positive association (sky appears after blue), some negative examples are provided (coffee doesn't appear after blue). This is helpful because most tokens appear with a small subset of other tokens, and it provides more signal in many cases to see tokens that are never coincident/in context together.

Some learned examples:

Nearest to Jesus:   ['Jesus', 'Jesus\n', 'from the dead', 'Herod ', 'Jesus ', 'all the people', 'Peter ', 'disciples']
Nearest to Pharaoh: ['Pharaoh', 'Sodom', 'Laban', 'in law ', 'Abimelech', 'in Egypt', 'household', 'Aaron the ']
Nearest to drink:   ['drink', 'and drink', 'bitter', 'counsel', 'by the sword', 'have', 'scatter', 'from thee']
Nearest to silver:  ['silver', 'brass', 's of gold', 's of silver', 'gold', 'oxen', 'ilver', 'shekel']
Nearest to chariot: ['chariot', 'chariots, and ', 'horses', 'camel', 'vessel', 'horsemen', 'brass, ', 'ashe']
Nearest to night :  ['night ', 'which is\n', 'morning, ', 'dark', 'latter ', 'morning ', 'to go ', 'moon ']
Nearest to young:   ['young', 'household', 'eld', 'in law ', 'maidservant', 'married ', 'fam', 'Egyptian']