Ground-Up VLA: Training Concepts

From GPU basics to robot foundation models

In this post, we’re going to take a look at the final piece needed to make deep learning pipelines produce useful outputs: training. We’ve covered all the pieces needed to build an end-to-end pipeline:

• Input

  • Convert raw data to embeddings

  • Create a token by summing the embedding of a “piece” of raw data and its positional encoding in the input sequence

• Transformer block (one or more of these)

  • Self-attention sublayer

    • LayerNorm normalization

    • Linear mapping, followed by a nonlinear activation function

    • Residual addition

  • Feed-forward sublayer

    • LayerNorm normalization

    • Linear mapping to expand dimensionality

    • Nonlinear activation function

    • Linear mapping to bring dimensionality back to the same value as the sublayer’s input

    • Residual addition

• Output

  • LayerNorm normalization

  • Linear mapping to logits

  • Optional softmax to produce probability distributions

For each linear mapping we need to perform, we have a weight matrix W and a bias vector b that must be known. For each LayerNorm, we again have a set of weights 𝛄 and biases 𝛃 that must be learned. In the training phase of a neural network (this is the same regardless of network depth), we are able to modify these parameters to better approximate what we see in the training data.

This post will examine the process of how we determine what the parameter values are. This, also, is how we define the size of a model. For instance, when we say that a model is of size “ten billion”, it means that the model has a total of ten billion weight and bias numbers that are used in the deep network of the model.

At its core, learning is an iterative loop. All the parameters in the pipeline, whether they are a linear mapping parameter, a LayerNorm parameter, or some other parameter, are initialized to random values. Then, using training data, we iteratively find better and better values for each parameter until the model starts to produce high quality outputs. We will take an input, use a forward pass through the pipeline to produce a prediction on what the output should be, compute the loss relative to what the model should have produced, and then perform a backward pass that updates all the parameters while moving from the output logit layer towards the input embedding layer.

Suppose we take the input as x and the output as y. If x and y were from a sentence example like “Greg Goggins got good at games,” then the goal of a training pipeline is to predict, at any point in time, what the next word in the sentence will be. So, for instance, suppose x = “Goggins”. The pipeline is supposed to produce the output of y = “got”.

Both “Goggins” and “got” can be represented as vectors where each element is a floating point number. This also means that we can compute the loss as the distance between the output the model actually predicted, and the output it was supposed to produce. The loss can be as simple as the Euclidean distance between two vectors.¹ So, suppose that we embedded “got” as [1.3, 1.5] and the model predicted the next word after “Goggins” was “described”, which can be represented with the embedding [3.4, 10.1]. The loss is computed as sqrt((10.1 - 1.5)² + (3.4 - 1.3)²) = 8.85. Instead, if the model predicted the next word as “received” with embedding [1.5, 1.1], then the loss would be sqrt((1.1 - 1.5)² + (1.5 - 1.3)²) = 0.45. Clearly, 0.45 is a smaller loss than 8.85. This also makes logical sense, because we know that while “received” is not quite what we would want, it is certainly a much better next-word prediction than “described”.

Using the loss, we can start the process of the backward pass, followed by optimization. We’ll come back to optimization, and for the moment focus on the backward pass (also called backpropagation). We move backward from the output to the input, hence the name. For each step we take backward, we ask the question “how should the parameters in this step change?” For instance, suppose the final output was the logits vector z in the vocabulary space. Then, the layer prior to this was the linear mapping z = W_logit y + b_logit, where y is the input to the linear mapping. We need to find how we should update the values in W_logit and b_logit. To simplify, let’s assume that we do not have any biases, so we only need to update W_logit.

We can compute the gradient for a weight w in W_logit as the partial derivative of L with respect to w.

\(\frac{\partial L}{\partial w}\)

If the gradient is positive, then increasing w would make the loss L increase, and if the gradient is negative, that would make the loss decrease. Consider the single equation: z = w y, where z and y are simple scalars. This means that the loss is L = sqrt((z_pred - z_true)²) = | z_pred - z_true |.

• If z_true = 4, y = 2, and the current value of w was 3, then the predicted value of z would be z_pred = 3 * 2 = 6. The loss would be L = | 6 - 4 | = 2.

• The partial derivative of L with respect to w would be sign(z_pred - 4) * y = sign(w y - 4) * y.

• Plugging in the numbers, we find the gradient as sign(w y - 4) * y = sign(3 * 2 - 4) * 2 = 2.

Now that we have a gradient, we need to do something with it. This is where we return to the concept of optimization. We are optimizing to minimize the loss, so the optimization step figures out how to change parameters such that the loss would have been smaller, given the previous observed loss. We now introduce the concept of the learning rate, which describes the rate at which we change w based on the loss we incurred. Typically, the learning rate is a pretty small value, such as η = 0.001. The update we perform is w ← w - η × (∂L / ∂w) = 3 - 0.001 × 2 = 2.998. Over many iterations, we would expect that w would eventually reach the value of w = 2, since that will minimize the loss L. This particular form of optimization is called gradient descent, but other optimizers exist (e.g., a popular one is the Adam optimizer).

The training process for a deep learning pipeline is quite similar to the toy scalar example we just examined. It applies the same principle, but at a much larger scale, and to matrices and vectors instead of to individual scalars.

And that’s training (conceptually)! We’ve now completed our fast pass through the concepts and some of the math used to put together a deep learning pipeline. Next, we’ll try to apply some of these concepts by building a not-very-deep neural network.

1. In reality, modern LLMs use cross-entropy loss. Euclidean distance just happens to be easier to understand, since most engineers have come across it before.