Ground-Up VLA: Deep learning and transformers

From GPU basics to robot foundation models

In the last post, we talked about self-attention and said that we’d talk about how they fit into the larger picture of transformers next. In this post, we’ll first see how transformers fit into the larger “deep learning” picture, and then take a look at transformer architecture. We won’t get into linear algebra and other math in this post, and instead keep the focus on architecture.

Deep learning, at its base, is the same as building a complex neural network. The reason it gets a special name is because it contains many layers of neurons stacked one atop another. Recall that a neural network is broken down into a few different components:

• the input layer,

• one or more hidden layers (since neither the input nor the output give you visibility on what the hidden layers are doing), and

• the output layer

A layer is a set of neurons, where each neuron follows the design we laid out in the neural networks post. Two adjacent layers (let’s call them parent and child layers) of a neural network are connected because the outputs of the neurons from a parent layer serve as the inputs to the child layer. A network becomes deep when we stack many such layers one after another, so that we have many hidden layers.

There are a bunch of different architectures that people have proposed on what math should be run inside of the layers, and how we should connect the neurons of a given parent layer to the neurons of a child layer. That’s beyond the scope of this post, however, so we’ll treat the exact connectivity as given and focus on a specific choice: a deep learning network built from transformer blocks.

So what does a modern deep learning network look like? Let’s break it down.

1. Input

   1. Raw data (text, image, etc.)

   2. Convert raw data to embeddings (vectors). I.e. produce

      \(X \in ℝ^{n\times d_{model}}\)

   3. Add positional encoding, so that a we retain knowledge of where a particular embedding falls in the sequence of embeddings representing the original input. Let’s call this P = X + Position(X). Here, Position(…) is the function that encodes positional information about the tokens of X (i.e., whether a word is the first, fifth, or hundredth in a sequence of words).

2. Transformer Block 1: T₁ = Transformer(P)

3. Transformer Block 2: T₂ = Transformer(T₁)

4. …

5. Transformer Block N: T_N = Transformer(T_N-1)

6. Output: Apply normalization and post-processing to the output of the final transformer block.

Let’s now zoom into Transformer Block 1 to see what it contains (the other blocks use the exact same architecture). Here, the input P is a given from the Input portion of the deep-learning architecture above, and we want to produce the output that we called T₁ above.

1. Input: P

2. Self-Attention Sublayer

   1. L₁ = LayerNorm(P): A normalization layer

   2. A₁ = SelfAttention(L₁): The self-attention block that we’ve already discussed

   3. R₁ = ResidualAdd(P, A₁) = P + A₁: Re-add P so that we retain an essence of what the original input to this sublayer was

3. Feed-Forward Sublayer

   1. L₂ = LayerNorm(R₁)

   2. F₁ = FeedForward(L₂): Position-wise transformations

   3. R₂ = ResidualAdd(R₁, F₁) = R₁ + F₁

4. Output T₁ = R₂

In the future posts, we’ll take a look at the pieces of this transformer block that we haven’t discussed yet: LayerNorm normalization, residual addition, the feed-forward sublayer, and output transformations.

We’ll also delve into how we actually produce the embeddings, do positional encoding, and what needs to be done to make the final output usable.