Ground-Up VLA: Embedding and Positional Encoding

From GPU basics to robot foundation models

In the previous post, we wrapped up our introduction to the transformer block as it is commonly used in state-of-the-art large language models. However, transformers on their own are not enough. We need to figure out how they can operate on inputs and produce meaningful outputs.

That is, what needs to happen for “My brain has too many tabs open” to turn into something that a transformer block can take as an input? This is the process of turning raw data into a vector embedding (often abbreviated to just embedding). We’ll be investigating what this process looks like, in addition to a related concept called positional encoding, which provides a means by which we can glean information about where a particular token falls within a sequence of input tokens.

The way in which we produce embeddings and perform positional encoding is typically fixed either before or during model training. So, when we want to perform inference, we can call lookup functions to tell us how to represent a token given its value and its position in a sequence of tokens.

Vector Embedding

The real world gives us messy data of many types. That data might be a drawing, photograph, audio clip, video, or text spanning many different languages, among other things. Machines, however, have trouble with this messiness, and strongly prefer to work with hard numbers. For deep learning pipelines, we need to produce embeddings that represent pieces of the “thing” that a deep learning pipeline is examining.

So, for instance, if my total vocabulary is “yes” or “no”, then I can have a 1-bit encoding, where 1 is “yes” and 0 is “no”. If we wanted to convey more ideas, then we might have a 3-bit encoding and expand our vocabulary to eight words, perhaps [yes, no, please, thanks, hello, bye, good, bad], and have a specific 3-bit number apply to each word.

In reality, languages are somewhat more complex, so we need something a bit more sophisticated. We’ll therefore move from bit encodings to embeddings — i.e., what real-world models use. First, instead of each element being a single bit, each element of an embedding is a real number. This vastly expands the expressiveness of each element in the embedding. Second, we expand the embedding size dramatically. For example, OpenAI’s text-embedding-3-small embedding model represents a token in its model with 1,536 separate 32-bit floating point numbers, which translates to 49,152 bits = 6,144 bytes. In other words, regardless of whether the word we create an embedding for is “bag” or “pterodactyl” or “defenestrate”, the embedding will consume 6,144 bytes.

So how is the exact embedding for a word like ‘bag’ determined? The short answer is: we don’t. The values of unique tokens are determined by training over a large corpus of data. Training is guided by the goal of making embeddings for similar text end up close to one another in vector space. Finding the distance between two embeddings can be as simple as computing the Euclidean distance between them. For instance, the distance between [1.0, 2.2] and [10.0, 2.2] is ((2.2 - 2.2)² + (10.0 - 1.0)²)^0.5 = 9. The embeddings, in modern large language models, get trained along with the weights of the deep learning pipeline.

Since we haven’t really touched on the training process yet, we won’t get into that much more detail, but this should hopefully complete the picture of how we go from a word like “bag” to something numeric that can be used as the input to a transformer block.

Positional Encoding

Producing an embedding works for individual tokens, but in reality tokens don’t really show up in isolation. Typically, we will have a sequence of tokens, and the position of a token in a sequence conveys useful information. So, before we send an embedding into a transformer block, we would also capture this positional information.

Broadly speaking, there are three families of positional encoding that are used:

1. Sinusoidal: Given a position p in a sequence (i.e., third token) and a dimension index i (i.e., the fourth dimension out of d_model = 512), compute a fixed PE(p, i). In the original paper that introduced modern transformer architecture (Attention is All You Need), this translates into a sine function if the dimension i is even, and a cosine function if i is odd. Below is the formula, but for now, we won’t worry about why the equation is the way it is.

   \(PE(p, i) = \begin{cases} \sin\left(\frac{p}{10000^{2\lfloor i/2 \rfloor / d_{model}}}\right) & \text{if } i \% 2 = 0,\\ \cos\left(\frac{p}{10000^{2\lfloor i/2 \rfloor / d_{model}}}\right) & \text{if } i \% 2 = 1 \end{cases}\)

2. Learned: Under the assumption that p never exceeds some value p_max, we can actually learn how much position should matter alongside learning all the weights that are needed in the transformer blocks. That is, we learn what value to add to the first token of a given input sequence given that it is the first token.

3. Relative: The position of a token is useful insofar as it describes relationships. Most (but perhaps not all) meaning derives from knowing where a token is relative to the other tokens in a sequence. For instance, in “I will come to your office”, it’s not particularly important that “I” is the first word, but it is important that it precedes “will”, thereby indicating who is taking an action in the sentence. Relative encodings are also learned; the model learns what it means for two tokens to be ‘one token to the right’ vs ‘two tokens to the right’, etc.

We don’t need to worry about the exact math quite yet. The reason for including it in this post is to develop an intuitive understanding of what advantages we get by having any positional encoding in our deep learning pipeline.

Recap

In this post, we took a look at how we can take raw data, like the word “yes”, and turn it into a vector embedding, which is a large vector of floating point numbers. These embeddings are usable by the transformer blocks that make up a deep learning pipeline. Before passing them into the first transformer block, we typically add a positional encoding to each token’s embedding.

Suppose we are given an input sequence S, where S₁ is the first token, S₂ is the second token, and suppose we can produce a lookup of what the embedding of a token is with the function Embed(S₁). If PE(S₁) gives us the positional encoding of S₁, then the vector we will pass to the first transformer block will be Embed(S₁) + PE(S₁).

Now that we’ve taken a look at the input side of a deep learning pipeline, the next post will examine the opposite side of the pipeline — i.e., how a deep learning pipeline produces a usable output.