DEV Community

Karam Sayed
Karam Sayed

Posted on Originally published at karamsayed.substack.com

Why Do AI Models Need High-Dimensional Embeddings?

A language model transforms words into context-rich representations that capture the meaning of each word relative to the words around it.

The human mind can visualize up to 3 dimensions intuitively, so let's use a 3-dimensional example to understand why higher dimensions allow richer representations.

Take the sentence: "The dog with that man barked."

Imagine that after the Attention layers process this sentence, the final representation of "dog" lives in a hypothetical 3-dimensional space with axes: Pet (the animal is a pet, hence its association with a man), Barking (the vector points high in this direction, indicating that it is barking), and Attention-seeking (the vector points in this direction, representing a hypothetical association with seeking attention).

Now consider another sentence: "The lion is sleeping."

The final representation of "lion" would be different because its surrounding context and associated properties are different. In our hypothetical space, it might point lower along the Pet axis, lower along the Barking axis, and lower along the Attention-seeking axis.

Vector Embeddings

Note: These three axes are hypothetical and created only for visualization. Actual Transformer dimensions do not necessarily correspond to individually identifiable concepts such as being a pet or barking.

This is just a simplified 3-dimensional representation. Now, a machine can work with representations in far greater dimensions.

In the original Transformer base architecture, the model uses 512-dimensional representations (d_model = 512). More dimensions provide a larger representational space in which the model can encode information about tokens and their context.

We don't know what specifically each dimension represents. These representations are learned by the model during training on the dataset. The input embedding provides the initial 512-dimensional representation, and the Attention layers and Feed-Forward Networks progressively transform it to incorporate contextual information.

You can read more about distributed word representations here: Efficient Estimation of Word Representations in Vector Space.


The Superposition principle

The superposition principle makes high-dimensional representations even more interesting. It describes how neural networks can represent more features than their number of individual neurons by allowing multiple features to be encoded through overlapping directions in activation space.

This means that individual neurons do not necessarily represent individual concepts. Multiple features can coexist within the same neural network activations.

You can explore this concept in Anthropic's research: Toy Models of Superposition

References:
1) Attention Is All You Need - Vaswani et al.
2) Efficient Estimation of Word Representations in Vector Space - Mikolov et al.
3) How might LLMs store facts | YouTube - 3Blue1Brown

This post is based on my Substack article: Dissecting the Transformer: Dissecting the Transformer: Inside “Attention Is All You Need”

Top comments (0)