Adam is an optimizer that helps us efficiently converge on a set of parameters for a stochastic function that we want to minimize. In machine learning, we often use Adam to find the parameters of a model represented by a stochastic function that minimizes a loss function.

The motivation

We already have many variations of optimizers, therefore it is good to question why we need another. Here we will outline the core problems Adam addresses and the consequences of not addressing them.

Classic optimizers such as Stochastic Gradient Descent (SGD) only have one learning rate. Therefore when we have several parameters, the gradient of a parameter at a particular point in time can be vastly different from other parameters. In this situation when using one global learning rate, we take the same step size when updating a parameter, no matter what the gradient is. This means when some parameters could receive a more confident, larger update, we might be held back by needing to have smaller updates on more sensitive parameters. Of course, this is scaled by the gradient itself as it is multiplied by the learning rate, but this places an upper limit on what our global learning rate can be, as it needs to account for sensitive parameters that should receive smaller updates.

Whilst problem 1 highlights the desire to have multiple learning rates to allow parameters to individually update more confidently or more sensitively depending on their gradient at a point in time, there’s also a related issue over time for a single parameter. The gradient for each parameter will change over our training run, as expected, and we want to take large steps when the gradient is large and smaller steps when approaching a minimum. As previously discussed, having multiple learning rates allows this to vary across the model for each parameter, but it does not vary over time within one parameter, and we still have an upper bound on individual learning rates.

As previously noted with some optimizers like SGD we have to choose a learning rate, and its upper bound is dictated by small gradient step sizes. We can also select learning rates using trial and error, grid search and using learning rate schedules which are a predetermined schedule for adjusting the learning rate over time. All of these require some thought into selecting the learning rate.

Due to memory constraints, we might not load entire datasets and therefore use batches of data as seen in SGD. However, each batch allows the gradient to diverge and could be a noisy outlier rather than representative of the overall gradient of the training batch.

Optimizers like Adagrad help with sparse gradients by retaining a high learning rate for infrequently updated parameters. We might come across sparse gradients in NLP tasks where words are infrequently used and ideally should receive meaningful updates.

SGD suffers from slowing down during flat regions, meaning updates are small when we might want to speed up and get through the region. Other optimizers have adjusted for this, such as SGD with momentum, where gradients are accumulated. This means once we are in a flat region, we are carrying on the accumulation of past gradients allowing us to push through flat regions rather than computing our update on the current gradient.

Step by step: inside the Adam optimizer

Now that we've seen the key challenges, let's walk through how Adam actually works step by step using the original algorithm from the paper.

$$\begin{alignedat}{2} (1) \quad & \textbf{Require: } \alpha \text{ (Stepsize)} \\ (2) \quad & \textbf{Require: } \beta_1, \beta_2 \in [0, 1) \text{ (Exponential decay rates)} \\ (3) \quad & \textbf{Require: } f(\theta) \text{ (Stochastic objective function)} \\ (4) \quad & \textbf{Require: } \theta_0 \text{ (Initial parameter vector)} \\ (5) \quad & m_0 \leftarrow 0 \quad \text{(Initialize 1st moment vector)} \\ (6) \quad & v_0 \leftarrow 0 \quad \text{(Initialize 2nd moment vector)} \\ (7) \quad & t \leftarrow 0 \quad \text{(Initialize timestep)} \\ (8) \quad & \textbf{while } \theta_t \text{ not converged do} \\ (9) \quad & \quad t \leftarrow t + 1 \\ (10) \quad & \quad g_t \leftarrow \nabla_\theta f_t(\theta_{t-1}) \quad \text{(Compute gradients)} \\ (11) \quad & \quad m_t \leftarrow \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t \quad \text{(Update biased 1st moment)} \\ (12) \quad & \quad v_t \leftarrow \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2 \quad \text{(Update biased 2nd moment)} \\ (13) \quad & \quad \hat{m}_t \leftarrow m_t / (1 - \beta_1^t) \quad \text{(Bias-corrected 1st moment)} \\ (14) \quad & \quad \hat{v}_t \leftarrow v_t / (1 - \beta_2^t) \quad \text{(Bias-corrected 2nd moment)} \\ (15) \quad & \quad \theta_t \leftarrow \theta{t-1} - \alpha \cdot \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon) \quad \text{(Update parameters)} \\ (16) \quad & \textbf{end while} \\ (17) \quad & \textbf{return } \theta_t \quad \text{(Resulting parameters)} \end{alignedat}$$

Initialization

Lines 1-2 are simply setting the hyperparameters of the algorithm. α is a normal learning rate as seen in SGD, and β1/β2 are values both between 0 and 1. β1 controls how much of a past gradient is added to each update step, creating a weighted average. e.g. If β1 is set to 0.9, we retain 90% of the previous moment estimate and blend in 10% of the current gradient. β2 is a similar parameter but covers an accumulated penalty value that we will see later, along with a deeper explanation of the weighted average of the gradient.

Line 3 is setting up our stochastic function to optimize e.g. a loss like mean squared error

Line 4 initializes the parameter vector, just like other optimizers.

Line 5 creates a vector to store the first moment (also known as the mean) of every parameter’s gradient. Therefore the vector should have the same length as the number of parameters we need to optimize. Each value is initialized to zero. When these values are updated β1 is used to scale how much we update the running mean with the previously seen gradient.

Line 6 creates a vector similar to what is seen in line 5 and is also initialized to zero. However we will store the second moment for each parameter in this vector as the algorithm progresses. The second moment is the running average of the gradient squared, which reflects the average magnitude (squared) of the gradient over time. It helps to assess the stability or noisiness of the gradient signal. As β1 is used to scale the first moment, β2 is used with this vector to scale the second moment at each update.

Training loop

Line 7-9 sets up and starts our training loop, with an initial time step t.

Line 10 computes gradients for the loss function with respect to model parameters θ in the previous time step t-1. These gradients are stored in the variable g.

Moment updates

Line 11 updates our vector of first moments. We initially set this to zero, therefore for the first update, we will be adding in 1-β1 of each parameters current gradient.

e.g. Given t=1, β1=0.9 and the gradient of the current parameter is 1.5.

We will keep 90% (when β1=0.9) of the previous gradient, and blend in 10% (1-β1=0.1) of our current gradient (1.5) resulting in 0.15 being stored in our vector for this parameter.

On the next loop given t=2 and the gradient of current parameter is now 1.4. We will update our first moment for the current parameter with 90% of its current value (0.15 from the last update), which results in 0.135 and add this to 0.1 of the current gradient (1.4). The result is 0.135+0.14=0.275

This process, blending in previous gradients rather than taking only the current gradient allows us to build an average gradient value and provides a signal if we are on a consistent gradient rather than a fluctuating one which might require us to back off from large updates.

Line 12 uses a similar mechanism to line 11 for updating a value for each parameter. However it stores the second moment, which is the exponential moving average of the squared gradients. Squaring the gradients ensures the second moment remains positive, acting as a penalty term. Since the gradients are squared, the values are always non-negative, therefore the second moment keeps accumulating regardless of the gradient’s direction. The calculation of the first moment does not detect oscillating gradients well as they flip around a local minima as the negative and positive gradients will cancel each other out. However the calculation of the second moment will keep increasing as it oscillates around a local minima.

Line 13-14 are bias correction steps for the first and second moments. Since we initialize both moment vectors to zero, early values are biased toward zero, even if the true gradients are not.

To correct for this, we divide by a factor that compensates for how little history we've accumulated so far. In the first few steps, this correction has a large effect, later on it fades away as the moment estimates become more accurate on their own.

For example, at time step t = 1, the correction for the first moment is:

$$\frac{m_t}{1-\beta_1^1}$$

And at t = 10:

$$\frac{m_t}{1-\beta_1^{10}}$$

Since β1​ is a number less than 1 (e.g, 0.9), raising it to higher powers brings it closer to 0. Therefore the denominator​ grows closer to 1 over time. That means in early steps, we divide by a small number (amplifying the estimate), and later we divide by something close to 1 (leaving it mostly unchanged).

Line 15 is where we update our parameters, where, similar to SGD, we adjust by a learning rate or step size as it is known here, however with Adam we can scale our step size based on the first and second moment. We can think of the first moment as our signal and the second moment as our noise. Therefore when we have a high signal to noise ratio, we are confident we can take a large step and allow us to take a large step size.

e.g. If our first moment = 5, second moment = 25

$$\frac{5}{\sqrt{25}}=1$$

This results in 1, which, when multiplied by our step size returns the full step size, therefore we take a large parameter update.

If our second moment is high compared to the first moment, this is a signal that we are not confident in the average gradient being reported by the first moment. This could be for a few reasons, such as being on an oscillating surface. If we frequently flip between negative and positive gradients, the second moment will capture all of these as positive values, therefore it keeps building up rather than negative values canceling out previous positive values as in our the first moment.

e.g. If our first moment = 5, second moment = 100

$$\frac{5}{\sqrt{100}}=0.5$$

This will result in reducing our step size by half, signalling a lack of confidence, therefore take smaller steps.

ε is a small constant added to prevent division by zero.

Algorithm Overview

Zooming out from the pseudocode, we can see how the various steps help to address the problems listed earlier. First we track the first and second moment for every parameter we want to train. Whilst this does use more memory than one global learning rate, it should result in better use of our resources as we converge faster.

Taking the first moment which is an average of gradients we become resistant to sudden gradient spikes allowing us to smooth out the convergence. If we take a stable downward path toward a local minimum, at first the gradients will be large, with the first moment retaining a high value. If we ignore the second moment, this large first moment will not be scaling the step size back much, signalling we are confident and can take these large steps. As we get closer to the local minima, our first moment should start converging towards zero, which starts to scale back our step size. This prevents overshooting the local minima and oscillating around it.

However, not all regions of the loss surface are smooth. In more chaotic areas, the second moment plays a larger role in stabilizing updates. The second moment will come more into play, where we are not on a consistent, smooth part to a local minima. If the gradient repeatedly enters small troughs, the gradient will flip between negative and positive values. With only the first moment, we have smoothed out gradients however these could be big shifts leading us to scale the step size erratically. With the second moment, regardless of a positive or negative gradient we are increasing the value due to the squaring of the gradient which always yields a non-negative result for any real number. As the second moment grows, the larger value to be divided by results in a smaller scaling factor. The algorithm becomes more conservative due to low confidence.

Ultimately, Adam determines how much to update a parameter by looking at the ratio between the first moment (our directional signal) and the square root of the second moment (our measure of noise or instability). Adam adapts its learning rate dynamically, growing cautious in noisy or uncertain regions and moving decisively when gradients are stable.

Here’s how Adam addresses the challenges introduced earlier:

• One global learning rate
  Adam maintains per-parameter learning rates using first and second moment estimates, allowing different update magnitudes for each parameter.

• Fixed learning rate over time
  Moment estimates adapt over time, allowing the step sizes to shrink or grow as gradients evolve.

• Manual learning rate tuning
  Step sizes are adjusted automatically, often reducing the need for manual learning rate schedules or tuning.

• Small batches are noisy
  Exponential moving averages smooth gradient estimates, helping Adam stay stable even with noisy mini-batch gradients.

• Sparse gradients
  Like Adagrad, Adam reduces updates for frequently active parameters, while allowing relatively larger updates for rarely used ones, making it well-suited for sparse gradients.

• Flat regions and momentum
  The first moment acts like momentum, helping push through flat or ambiguous regions of the loss surface.

Summary

Optimizers preceding Adam used parts of the concepts it brings together. For example, SGD was extended with momentum, which averages gradients over time. This helps reduce oscillation and allows the optimizer to follow a more stable path, which is conceptually similar to Adam’s first moment estimate.

Adagrad introduced per-parameter learning rates by accumulating the square of past gradients. This allows large updates for infrequently updated parameters and smaller updates for frequently updated ones. However, because the accumulation grows without decay, the learning rate can become excessively small later in training.

RMSProp improved on this by using an exponential moving average of squared gradients instead of a cumulative sum. This enabled the learning rate to adapt more flexibly to recent gradient behavior rather than shrinking predictably over time.

Adam combines these ideas. It uses momentum like first moment estimates and RMSProp style second moment estimates to scale the step size based on both direction and the reliability of the gradients. Adam also introduces bias correction, which improves early training by compensating for the initial zero values in the moving averages.

Adam is widely used because of its robustness and adaptability. It often converges faster than optimizers like SGD. However, it does not always generalize as well. Because Adam closely follows the gradient signal for each parameter, especially in models with many parameters, it can overfit or settle into sharp minima. In contrast, SGD with momentum tends to average out gradient noise more effectively, which helps it find flatter minima that often lead to better generalization.

Below is a PyTorch implementation of Adam's core logic.

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
import matplotlib.pyplot as plt
import seaborn as sns

model_name = "bert-base-uncased"

# define sentence
sentence = "The brown dog chased the black dog"

#define tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokens = tokenizer(sentence, return_tensors="pt")["input_ids"]

# define token embeddings
model = AutoModel.from_pretrained(model_name)
token_embeddings = model.get_input_embeddings()(tokens)
embed_dim = token_embeddings.shape[-1]

class SimpleAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.w_query = nn.Linear(embed_dim, embed_dim)
        self.w_key = nn.Linear(embed_dim, embed_dim)
        self.w_value = nn.Linear(embed_dim, embed_dim)

        # Initialize weights
        nn.init.normal_(self.w_query.weight, mean=0.0, std=0.8)
        nn.init.normal_(self.w_key.weight, mean=0.0, std=0.8)
        nn.init.normal_(self.w_value.weight, mean=0.0, std=0.8)

        self.attention = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=1, batch_first=True)

    def forward(self, x):
        output, attention_weights = self.attention(self.w_query(x), self.w_key(x), self.w_value(x))
        return output, attention_weights

# Compute attention        
attention_layer = SimpleAttention(embed_dim=embed_dim)
attention_output, attention_weights = attention_layer(token_embeddings)

# Convert attention weights to numpy array and remove extra dimensions
attention_matrix = attention_weights.squeeze().detach().numpy()

# Get token labels for the axes
tokens_text = tokenizer.convert_ids_to_tokens(tokens[0])

# Create heatmap
plt.figure(figsize=(10, 8))
sns.heatmap(attention_matrix, 
            xticklabels=tokens_text,
            yticklabels=tokens_text,
            cmap='YlOrRd',
            annot=True,
            fmt='.2f')

plt.title('Attention Weights Heatmap')
plt.xlabel('Key Tokens')
plt.ylabel('Query Tokens')
plt.tight_layout()
plt.show()

The code above will create a simple attention layer, without positional encoding. This enables our test to demonstrate how attention behaves when all tokens are treated purely based on their semantic embedding, without any positional differentiation. It has been initialised with random weights and has not undergone training, however training shouldn’t matter as tokens that are the same will have the same weights applied to them and result in the same output.

The attention heat map above is the output of our code. I’ve highlighted the use of two dog tokens to demonstrate how they receive equivalent attention weights against all other keys. For attention heads to be able to specialize e.g notice verb-object pairs, they must be able to differentiate between the same words in different positions otherwise they will receive the same result from the same weights being applied to token embedding.

What do positional encodings enable

In transformers we differentiate words at different positions with positional embeddings. This transformation perturbs the embedding based on the position of the word. For our example sentence in the code above, dog at position 3 would have a slightly different embedding to dog at position 7 after positional embeddings have been applied, resulting in a different attention result for each token. This will allow attention heads to see different positions and allow them to specialize e.g. an attention head that has specialized in recognising verb object pairs would have high attention with chases and the second dog, but not the first dog. However if each dog had not had their embedding slightly altered the specialized attention head could never notice that difference between both dogs.

Positional encodings step by step

Before we start to think about how positional encodings might be implements lets come up with some requirements

This is instead of passing in an addition feature meaning our neural network would have to interpret extra information and spend extra computation. By combining with existing embedding, we need a way for the network to see a different embedding when words are the same at different positions.

As shown early, each word needs to be uniquely identified even if the same word appears in a text sequence.

We want to create models that can generalise, therefore our positional encoding should be a pattern that can be recognised and accounted for in our attention matrices. We could assign random IDs to each position, and in theory this is learnable, however we would be using more learnable parameters than needed to account for the need to memorise individual IDs and the semantic information they contain for the various encounters. We are essentially forcing the transformer to memorise combinations of position and semantic information rather than allowing it to learn reusable patterns.

Add the position as an integer

An obvious thing to try to simply add the position of the token to the embedding. e.g. If dog was my first word and the embedding was [0.10, 0.80, 0.45], adding 1 to this would result in [1.10, 1.80, 1.45]. For the second word we would add 2 and so on.

However there are a few problems with this, first I could have many words in my input and say I reach the 100th word, adding 100 to each dimension of that token’s embedding will create very large numbers. In neural network we want to keep inputs on a similar scale to allow for faster convergence. However with our implementation embedding values will grow as we add more words to our input.

This requirement will ensure that embeddings stay within reasonable limits and do not grow too large. Another problem presented by our larger integers is the network will only be trained on the largest length it has seen, and therefore cannot generalise to any length. Ideally we want our network to work at lengths it has not seen.

Lastly the positional part of our embedding, seems to greatly outweigh the original embedding. It might be hard to see what the original input was, in fact it seems to completely change the embedding and more strong represent position than semantic information.

The graph above shows the embeddings of random list of words without positional information in reduced dimensional space. We can see clear semantic meaning, with words around the same concept clustering together.

The graph above represents the same embeddings with integer positions added based on the position they were in the text sequence. We can see they no longer cluster as before and they seem to have a different semantic meaning. They also seem to be dominated by the some increasing order. Therefore adding a position with an integer is probably not going to work for us and has led us to a new requirement.

Sinusoidal Encoding

Based on our requirements so far, we’ve now come to the solution presented in the Attention Is All You Need paper. However let’s build up our understanding of how this works.

To help meet our requirements, let’s use a sine function, where the positional coding is sine(x), with x being our token position and the result added to our embedding element wise.

For example if our embedding was [0.1, 0.6, 0.8] our resulting embedding with positional information would be [0.1 + sin(0), 0.6 + sin(1), 0.2 + sin(2)].

If we take the x axis to be our token position we can see that the numbers produced are small, helping to retain our embedding semantic information if added to each embedding element wise. Requirement wise it almost meets what we want, it can handle any sequence length as sine can be calculated for any x value, it is bounded between 1 and -1, it isn’t large enough to alter semantic information, however it most likely does not meet requirement 2: Positional Encodings must uniquely identify every word in a sequence.

As we can see sin(x) is periodic and it will eventually repeat itself for a tokens position. e.g sin(2) and sin(353) when rounded up will return the same value. This means for positions that return the same value, our attention mechanism will see them as the same position, and if they are the same word, we will have the same problem as not having positional encoding.

To help mitigate this we introduce cosine i.e. cos(x) as shown above and for each even index of our embedding we use sine and for odd indices we use cosine.

For example if our embedding was [0.1, 0.6, 0.8] our resulting embedding with positional information would be [0.1 + sin(0), 0.6 + cos(1), 0.2 + sin(2)].

Using cosine introduces a phase shift and greatly reduces the likelihood of receiving the same value for a position later in the sequence and the phase shift greatly reduces the chance of values being near each other. However there is still a chance of repeating at longer sequences lengths, plus with the regular periodicity it is hard to build relative patterns across long distances. This means we might not be meeting requirement 3: We must be able to generalise over positional encoding patterns. With this shortwave frequency it will be hard to generalise long distance patterns.

We can improve on both generalising over long range patterns reducing the chance of results collapsing into the same position by introducing lower frequencies. ie. we can create a general function which is sin(x/i) for even positions and cos(x/i) for odd positions where i is the the index of our embedding.

For example if our embedding was [0.1, 0.6, 0.8, 0.9] our resulting embedding with positional information would be [0.1 + sin(1), 0.6 + cos(1/2), 0.2 + sin(2/3), 0.8 + cos(2/4)].

The graph has been expanded on the x axis compared to previous graphs, to show how dividing by the i-th position decreases the frequency. We would do this for all the dimensions of the embedding, therefore if our embeddings size was 768, we would have 768 individual values, where the frequency decreases as we move closer to the end of the embedding.

The decreased frequency can be used to generalise over long distances and the extra dimensions highlight short, medium and long distances allows our attention matrices to learn parameters that discard certain ranges and focus on what it needs for a particular attention head.

With our current solution, as we are only scaling up each embedding index linearly by increasing i, we could have many dimensions that don’t convey enough different information and therefore have redundant positional information.

Our current formulas are:

$$PE(pos, 2i) = \sin\left(\frac{pos}{2i}\right)$$

$$PE(pos, 2i+1) = \cos\left(\frac{pos}{2i}\right)$$

Above PE represents our positional encoding function, pos the token we are taking in, and i our embedding index.

In the Attention Is All You Need paper, the formulas used are

$$PE(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d}}\right)$$

$$PE(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d}}\right)$$

Here we can see rather than dividing by i, we have i as an exponent of 10000 and i is divided by d which represents the embedding size. This allows the the different dimensions to scale up logarithmically covering a wider range of frequencies rather than the previous linear scaling and also reduces redundant positional information. Using 10,000 as the base with an exponent is what allows the frequencies to scale up logarithmically, and d is used to control the scaling. Without d, it would scale far too quickly and might mean we miss important positional information. 10,000 was used as the base after experimenting in the Attention paper and provided a balance between scaling up and capturing enough information at different ranges. This great range of frequencies allows the model to generalise better, and gives more options for attention heads to focus on ranges for their specialisation.

Absolute vs Relative Positional Encoding

So we’ve built out positional encoding in the same way that was designed in the Attention Is All You Need paper and this method has been used in many models. However it has been improved upon and subsequently used in newer transformer models.

The problem with sinusoidal encoding is that is falls into a category known as absolute positional encoding. Absolute encoding adds a particular value or ID to each position, so position 1 receives a certain value, position 2 another etc. If I have the sentences “A man threw a ball” and another sentence “In the garden a man threw a ball”. In each sentence I have “man threw” in position 2, 3 and in position 5, 6 in the first and second sentence respectively. With sinusoidal encoding, man and threw will receive different positional encodings in each sentence. However in language, what creates meaning is less about absolute positions of related words but the relative position between them. If an attention head specialized in detecting subject and verb, it would have learn the value for man and threw differently for the two different sentences, by providing absolute values for each position, our learnt parameters have less chance to generalise, even though we do have some implicit relative positioning with sinusoidal encoding due to the periodicity introduced with sine and cosine waves. However the relative implicit relative positioning is hard to extract and models are not forced to look at this when absolute positioning is available.

With relative positioning, if we took the same sentence our encoding scheme would receive values that highlight the relative position between words, and not have absolute values. This forces the model to use relative values and also helps to generalise more, a property we want for our NLP models. For example in the case with man and threw in two sentences, our model could extract the fact the words are next to each, without having to memorise exact embedding plus positional values for different situations.

RoPE (Rotary Positional Embedding)

One way that has been used in many models to provide relative positional embeddings is RoPE.

Let’s imagine we have sentence 1 “The mouse ate some cheese” and sentence 2 “In the house a cat ate some cheese”

The graphs on the top represent simple 2D embedding for each word in our sentences. The graph below represents a rotation of each vector based on the their position using the transformation pθ where p is the word position and θ is 25° e.g. some in the first sentence at position 4 would be rotated by 4θ which is 100°. As we can see in the rotated embeddings, ate and cheese which have the same relative position in each sentence end up in different positions, due to their position, but the angle between them is the same as the relative distance is the same in each sentence. This is what RoPE does it takes an embedding, and rotates it by some θ based on it’s position.

RoPE and Attention Math

You can also see that every embedding has it’s magnitude maintained. No matter where we rotate the magnitude will be the same. We can see the importance of this by looking at how we calculate attention scores.

$$q⋅k=∥q∥∥k∥cos(θ)$$

We can see our attention scores rely on the norms of our vectors and angle between them. Since we have only rotated our vectors their norms stay the same, therefore any influence on attention scores due to RoPE must be on angle changes. Including the RoPE angle difference our attention score now becomes the following, where d the the angle introduced by RoPE.

$$∥q∥∥k∥cos(θ+d)$$

This now means, if our attention head wanted q and k to be highly aligned but they are far apart, our model must learn to increase the magnitude of the original embedding when projected into q and k or project them into a space where they are much closer, despite RoPE moving them far apart.

We’ve seen how rotating embeddings, helps transformers to identify tokens at different positions by providing another lever to modulate, the angle between tokens. To see how we relative position can be picked up by the model, we need to break down our attention score using transposition.

$$RoPE(q,i)=R(i)q$$

The above formula define a RoPE transformation, where the first parameter is our vector, the second the position in a text sequence and R is a rotation matrix. Therefore our attention score would be.

$$RoPE(q,i).RoPE(k,j) = R(i)q.R(j)k$$

Changing this to matrices transposition, and using product transposition rules.

$$(R(i)q)^\top(R(j)k) = q^\top R(i)^\top R(j)k$$

Since R is a rotation matrix, and the transpose of a rotation matrix can be represented as R(-θ) we can simplify to

$$q^\top R(j-i)k$$

Where j and i are simply positions, we can see how the attention calculation has cleanly extracted the difference in positions, allowing it to use relative distance rather than being intertwined with the actual semantic content of the embeddings. This allows the model to modulate what it wants such as direction or angle to achieve the result we want. Before with sinusoidal we have to learn to extract position from the semantic information.

RoPE frequencies

As in the sinusoidal method previously talked about you might be wondering, if the rotation eventually repeats and therefore the same token at different positions will appear at the same angel relative to another token. Like the sinusoidal method, RoPE also introduces different frequencies across the embedding. In our example so far we have only considered 2D embeddings. In RoPE, within one embedding we take indices of the embedding pairwise e.g. position 1 & 2, position 3 & 4 etc and rotate each pair by some θ depending on the token position and the index of the pair we are looking at. As we move further along the embedding, the rotation becomes smaller and smaller, resulting in a lower frequency towards the end of the embedding as in the sinusoidal method. This allows us to create unique signatures for likewise tokens in long text sequences and provides the model with data to help it focus on what it wants such as long range or short range dependencies.

To rotate each pair we can use a rotation matrix

$$\begin{bmatrix} \cos(m\theta_i) & -\sin(m\theta_i) \\ \sin(m\theta_i) & \cos(m\theta_i) \end{bmatrix}$$

m is our token position and i-th θ is defined as

$$\theta_i = \frac{1}{10000^{\frac{2i}{d}}}$$

This is similar to our sinusoidal formula and produces similar sinusoidal frequencies, i being the pairwise position, and i-th θ is a rotation angle defined for each pair in an embedding based on it’s index.

e.g. If our embedding was [0.1, 04, 0.5, 0.2], after splitting into pairs we have [[0.1, 0.4], [0.5, 0.2]], our 1st pair is [0.1, 0.4] and 2nd pair [0.5, 0.2]. We then calculate theta for the 1st and 2nd pair use 1 and 2 as i.

$$\begin{bmatrix} \cos(m\theta_0) & -\sin(m\theta_0) & 0 & 0 & \cdots & 0 & 0 \\ \sin(m\theta_0) & \cos(m\theta_0) & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos(m\theta_1) & -\sin(m\theta_1) & \cdots & 0 & 0 \\ 0 & 0 & \sin(m\theta_1) & \cos(m\theta_1) & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & \cos(m\theta_{d/2}) & -\sin(m\theta_{d/2}) \\ 0 & 0 & 0 & 0 & \cdots & \sin(m\theta_{d/2}) & \cos(m\theta_{d/2}) \end{bmatrix}$$

Using the above sparse matrix we can rotate each pair of embedding. We calculate θ up to d/2 due using pairs and only to compute angles for half the size of the embedding.

Efficient Implementation of RoPE

The sparse matrix can get quite large, and holding that in memory along with accompanying matrix multiply can be made more efficient by splitting up our rotation matrices into two vectors.

If we have a 2 dimensional embedding at position 1 rotating this would involve the calculation

$$\begin{bmatrix}x_1\\ x_2 \end{bmatrix} . \begin{bmatrix} cos(\theta) & -sin(\theta)\\ sin(\theta) & cos(\theta)\\ \end{bmatrix} = \begin{bmatrix} x_1cos(\theta)-x_2sin(\theta)\\ x_1sin(\theta) + x_2cos(\theta)\\ \end{bmatrix}$$

We can break this down into the following components

$$\begin{bmatrix}x_1\\ x_2 \end{bmatrix} \otimes \begin{bmatrix} cos(\theta)\\ cos(\theta) \end{bmatrix} = \begin{bmatrix} x_1cos(\theta)\\ x_2cos(\theta) \end{bmatrix}$$

First we use element wise multiplication to cover the cos part of our rotation matrix. Notice how we have left -x2 *sin(θ) on the top row and, x1 *sin(θ) on the bottom. As negative x2 is now on the top, for every pair we flip them and make x2 negative.

$$\begin{bmatrix}-x_2\\ x_1 \end{bmatrix} \otimes \begin{bmatrix} sin(\theta)\\ sin(\theta) \end{bmatrix} = \begin{bmatrix} -x_2sin(\theta)\\ x_1sin(\theta) \end{bmatrix}$$

We then simply add the two results together element wise to achieve the same as sparse matrix rotation, but with a method doesn’t require construction of large matrices and is more easily vectorized.

$$\begin{bmatrix} x_1cos(\theta)\\ x_2cos(\theta) \end{bmatrix} + \begin{bmatrix} -x_2sin(\theta)\\ x_1sin(\theta) \end{bmatrix}$$

Sources & Further Reading

- [You could have designed state of the art Positional Encoding] (https://fleetwood.dev/posts/you-could-have-designed-SOTA-positional-encoding) - This post was heavily inspired by this original blog post.

- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864) – Introduces RoPE (Rotary Positional Embedding), a technique for modeling relative positions through rotation.

- [Attention Is All You Need](https://arxiv.org/abs/1706.03762) – The original Transformer paper that introduced sinusoidal positional encodings and the self-attention mechanism.

Yo, AI’s a mind that don’t ever sleep, stackin’ patterns and data, runnin’ deep.
It don’t hustle like us, but it’s sharp on the grind, tech so smooth, it’ll blow your mind.

Have you ever wondered how ChatGPT can talk about AI in a Snoop Dogg like style, as seen in the example above? In this series, we’ll go from basic machine learning to help build our intuition, before implementing the models used to power ChatGPT. Along the way, we’ll learn the basics of PyTorch, a popular machine learning library to simplify our code and enable easier implementations of more complex models.

Why Use Machine Learning?

When a computer is required to perform a task, we might write a function mapping an input to an output to enable the task. For instance, to predict house value based on the number of bedrooms, we might write a function such as

def predict_house_value(bedrooms):
    return bedrooms*100000

This function takes in the number of bedrooms, multiplied by a constant, and outputs a house value. As a predictor of house value, it’s pretty limited, we could take in more input parameters to calculate a more accurate house value.

def predict_house_value_v2(bedrooms, city, garden):
    # Adjust house value based on whether there is a garden
    garden_increase = 2000 if garden else 0

    # Adjust house value based on city
    city_multipliers = {
        "London": 1.2,
        "New York": 1.5,
        "Tokyo": 1.8
    }
    city_multiple = city_multipliers.get(city, 1.0)

    return garden_increase + city_multiple * 100000 * bedrooms

The function is gradually refining its output, returning house values that better reflect the diverse range of values we encounter. However, it remains highly inaccurate and would necessitate additional parameters to function effectively or restrict its application to specific regions and require extensive domain knowledge to comprehend how each input influences the output. The algorithm also potentially misses interplay between the 3 input parameters. For example, gardens could have more value in different cities. You can imagine this as if statements capturing the complex relationships by ANDing input parameters, however try to imagine this now with 50 cities and/or more input parameters.

Machine Learning allows us to take preexisting data and create the mappings represented by our functions for us. For example in the house value example, we can use data on various houses with features such as the number of bedrooms, proximity to schools, location and a house value associated with it.

We then select an appropriate machine learning model to train. During the training phase, we take each house one at a time, load the features, and the dependent variable which in this case is the house value. While training, the model will learn to predict house values based on features it has seen, by adjusting internal parameters of the model. After training, we can infer house values of houses not seen in the training data by supplying different combinations of features e.g 5 bedrooms, not near a school, New York, etc and the model will predict the house as it has learned to map input features to predicted house values.

To help us understand how models learn and why their internal adjustments create a desired output we'll start with linear regression. Later we will explore how we can extend concepts in linear regression to more complex architectures such as neural networks, which can approximate a variety of computational tasks and finally transformers which will allow you to create Snoop-Dogg style text like at the top of the article.

Linear Regression

Linear Regression is a method used to model the relationship between a dependent variable and one or more features. The dependent variable should be a continuous value when using linear regression. For instance, if we want to predict the value of a house based on its size, we can create a linear model where house value is the dependent variable and house size is the feature. This model would help us to understand and predict how changes in the size of a house might influence its value, allowing us to make predictions about future house values based on size.

We can use a 2D plot with house size on the x-axis and house value on the y-axis, to help us visualize the model.

Plotting the data shows a clear correlation between house value and house size. Drawing a line through the data will help us to approximate other house values based on their size. With this line, we can pick a house size on the x-axis, and see what y value intersects with x on the line. This line would be our model.

The line drawn above is one we could draw by hand and intuitively know it is a near optimal line to predict house values based on the data. This is based on it going through the centre of the plotted data.

Training a linear regression models allows us to find that line programatically. However if we can easily draw the line, why bother. Well, we can have many input features, and each feature is another axis on our plot. That however becomes more difficult to visualize in 3D and not possible above 3 dimensions. In real projects we will most likely be dealing with multiples features and therefore need a way to determine what this line is. For the time being though we’ll stick with the 2 dimensions, house value and house size while learning.

Model line

The line drawn through our data represent our model. It allows us to infer values for inputs not yet seen in a dataset. Each line can be represented with the following equation

$$y=w_1 x_1+b$$

• y is our dependent variable, house value, the value we are trying to predict

• x₁ is an input feature, house size. We use subscript 1, because we can have many input features which would be labeled x₂, x₃… etc. These other features could represent other house properties such as number of floors, or proximity to schools. However for our simple case we have one input feature.

• w₁ is a weight. This is a constant value that is determined during training. Again this has a subscript 1 because there will be as many weights as there are features i.e. there is a weight for each feature category.

• b represents a bias value. This is another constant value determined through training. It is not multiplied by any feature but simply added onto the final result.

w₁ and b are constants that are learnt through training and these are the values that decide the gradient of our model. x₁ would be a variable that is inputted for each calculation, in this case the house size.

From the above plots we can see how different values for w₁ and b produce different models, some clearly better than others.

We can see w₁ affects the gradient, this is because it is a weight being applied against an input value such as house size. The greater the weight the more influence it has on the dependent variable y such as house value. When the weight is greater, the slope is steeper as the feature influences the dependent variable more. When the weight is negative there is a downward slope as the feature has a negative affect on the dependent variable.

When we change b, we see the line intersects the y-axis at a different point. For example at zero, it intersects where y=0, when b is 1, it intersects where y=1. Without this bias value, the model would have less flexibility as y would always have to be zero when x is zero.

What makes a good model

Now we know what parts make up a linear regression model, how do we know we have good model. Visually we can see there are better models than others. The best model is the one that minimizes the difference between predicted values and actual values.

From the above image, we can see there is a difference between what a model predicts for a house size and what the actual house value is for a given house size. If our model was the perfect predictor and predictions were equal to the actual value, the sum of all the differences would be zero. As predictions start to deviate from actual values the sum of all absolute differences increases and our model becomes potentially less accurate.

The loss function

We can use the MSE (Mean Square Error) formula to represent those differences. This is called a loss function, which aggregates the the difference between all actual and predicted values. Each individual difference is an error value.

$$\text{MSE} = \frac{1}{n} \sum_{i=1}^{n} (y_i - \hat{y}_i)^2$$

The MSE will give us an overall loss value, where a lower loss value represents a better model. i.e. a loss of zero is a perfect predictor, meaning all predictions perfectly match the actual values.

• n is the number of instances we have e.g if we have data for 100 house, n is 100

• i is a specific instance e.g. in our house data, if we took the 2nd example i would be 2

• ŷ (pronounced y hat) is our predicted value e.g. the first house in our dataset might have a predicted value of $100k in our model

• y is our actual value e.g. the first house in our dataset might have an actual value of $120k

The MSE calculates the difference between all predicted and actual values for all n instances and squares each one. Squaring the result adds a greater penalty to predictions further away from actual values, whilst also making each value positive.

Making each value positive ensures minus values can’t cancel out other values when summing the values.

How do we train a good model

To select good values for our weights and bias which minimizes the loss from the MSE we must train our model using a dataset. From now on we’ll refer to weights and biases collectively as parameters.

We train with the following steps

• initialize our parameters (weights and bias) to random values

• calculate predictions for every sample we have based on the parameter values we chosen for our set of samples, calculate the mse

• based on the mse value, adjust our parameters, in a way which favors a lower loss value (we'll cover this in more depth later)

• with our new parameter values, we calculate predictions again, calculate the mse, and adjust parameters.

• continue the above steps, of adjusting parameters, calculating mse until our loss is acceptably small or the changes in the mse start to slow down with each iteration or for a set number of iterations

We'll go through these steps in code below with real data using the California Housing dataset. In this dataset, each row represents a district in California, and the target variable (the output) is the median house value of owner occupied homes, expressed in hundreds of thousands of dollars. As there is more than one feature, we’ll have more than than 2 dimensions, and therefore won’t be able to easily visualize with graphs as before.

Putting it all together

Data importing and exploration

Below we'll use linear regression on a real dataset. The California Housing dataset collects data on housing per district in California. Each district contains features such as population, average number of bedrooms etc. Using linear regression, we'll predict the Medium House Value per district based on the the other fields in the dataset.

# Import libraries that allow us to download the California housing dataset
# and view the data
from sklearn.datasets import fetch_california_housing
import pandas as pd

# Load the dataset
california_housing = fetch_california_housing()

# Convert to a DataFrame, this will allow us to view and access data easily
data = pd.DataFrame(california_housing.data, columns=california_housing.feature_names)

# Create MedHouseVal column which represents the Median House Value per a district
# This is the y value we will try to predict
data['MedHouseVal'] = california_housing.target

# Display the first few rows of the dataset
print(data.head())

The output from data.head() allows us to quickly inspect the dataset so we can get a feel of what the data looks like.

MedInc  HouseAge  AveRooms  AveBedrms  Population  AveOccup  Latitude  \
0  8.3252      41.0  6.984127   1.023810       322.0  2.555556     37.88   
1  8.3014      21.0  6.238137   0.971880      2401.0  2.109842     37.86   
2  7.2574      52.0  8.288136   1.073446       496.0  2.802260     37.85   
3  5.6431      52.0  5.817352   1.073059       558.0  2.547945     37.85   
4  3.8462      52.0  6.281853   1.081081       565.0  2.181467     37.85   

   Longitude  MedHouseVal  
0    -122.23        4.526  
1    -122.22        3.585  
2    -122.24        3.521  
3    -122.25        3.413  
4    -122.25        3.422

Scaling features

Next, we need to preprocess our data by scaling the features. Scaling adjusts the range of each feature to be similar, which prevents any single feature from dominating the training process due to its large scale. This also aids in faster convergence of the model during training.

import numpy as np

# We'll create our input feature dataset X and a corresponding Y dataset 
# for the data we want to predict
# We'll also convert the datasets to numpy which is better for processing, 
# where Pandas dataframes are good for data exploration
x_unprocessed = data.drop("MedHouseVal", axis=1).to_numpy()
y = data["MedHouseVal"].to_numpy()

# Scale features in X using Z-score normalization
# This preserves the distribution of the data and ensures all
# features have a mean of 0 and a standard deviation of 1
x_mean = np.mean(x_unprocessed, axis=0)
x_std = np.std(x_unprocessed, axis=0)
x = (x_unprocessed - x_mean)/x_std

The above scaling, known as standardization or Z-score normalization, ensures that features with large ranges do not overshadow others in influencing the model's predictions. Without scaling, features that have large ranges might result in smaller weights, as their input values are already greatly influencing the dependent value due to their size relative to other features. Over time, training will reach the correct weights, however, unscaled features can cause inefficient optimization paths, leading to slower convergence.

If all of our features are on similar scales (e.g., 1-100), then we can leave scaling out, as no feature will dominate due to scale.

Note we did not scale y, as y is not included in the learning process and a weight is not assigned to it, as it is the value we are trying to predict.

Test and validation datasets

After exploring the data and pre-processing it, we need to split the dataset into a train and validation set. A typical split would include 80% of the data in the training set and the remaining 20% in the validation set. This split is commonly used because it provides enough data to train the model effectively whilst also keeping a sufficient portion to evaluate its performance. However, this ratio might change depending on the size of the dataset or the requirements of the specific task. Larger datasets might allow for smaller validation sets, such as 90/10 splits.

The training set is responsible for training the model. When we calculate the MSE at the end of a training iteration, ideally, we want as low a number as possible. However, if we achieve zero MSE and therefore perfect accuracy, there is a high chance our model has learned to predict the data in the training set but has not generalized to unseen data. This is also known as the model overfitting the data.

To test that our model generalizes to data outside of the training set, we hold back some data from the initial dataset to create a validation set. With the validation set, after we have trained the model, we calculate the MSE on this set. If the loss is low, indicating predictions are close to the observed data in the validation set, it is likely we have a generalized model, meaning it performs well on unseen data and not just the validation set. However, a low validation error alone does not guarantee generalization and might require additional testing on unseen data to confirm the model's robustness.

# optional seed, setting this makes our random calculations reproducible
np.random.seed(42)
# get the number of samples, shape returns the dimensions of the matrix
# specifying 0, fetches only the number of rows
# a random order is generated, in case the order of our data 
# skews the distrubtion in some way
n_samples = x.shape[0]

# this will generate an array of numbers from 0 to the number of rows minus 1, in a random order
# we'll use this to select rows from our datasets with a random order that is set once
indices = np.random.permutation(n_samples)

# define the percentage of samples that should be used in a training set
train_ratio = 0.8
n_train = int(train_ratio * n_samples)

# create the training and validation datasets
train_indices = indices[:n_train]
val_indices = indices[n_train:]
x_train, y_train = x[train_indices], y[train_indices]
x_val, y_val = x[val_indices], y[val_indices]

print(f"Training set size: {x_train.shape}")
print(f"Validation set size: {x_val.shape}")

Training set size: (16512, 8)
Validation set size: (4128, 8)

Weight initialisation

Next we create a weight vector to represent the weights applied to each feature. We need as many weights as there are features and an additional bias value, which allows the model to adjust the output independently of the feature values. Random small values are used for weight initialization.

We also choose small initialization as large numbers can cause a feature to contribute too much to predictions, causing large changes to the weights across iterations. The large changes could become unstable, oscillating around the optimum value. Instead of gently moving towards an optimum point, the weight values might swing wildly around it, slowing convergence.

# initialise an array of random weights for each feature
# x.shape(1), will give us the number of columns which represent features
# we add 1 to the number of features to account for the bias term that will be trained
n_weight = x_train.shape[1]+1
weights = np.random.randn(n_weight) * 0.01
print("Initial weights:", weights)

# we will also add a feature column of ones, which will be used for the bias value later
# np.c_ will horizontally concatenate 2 matrices e.g. x_train and a generated 1 column matrix of 1s
x_train = np.c_[x_train, np.ones(x_train.shape[0])]
x_val = np.c_[x_val, np.ones(x_val.shape[0])]

Example output, yours will be different due to random initialization.

Initial weights: [-0.00454508  0.00186784 -0.01296609  0.00670622 -0.00369122  0.00874647
 -0.00619931 -0.00674114 -0.00945463]

The training loop

Now we begin the training loop, an iterative process aimed at finding the optimal weights. In the training loop we will calculate predictions for each sample using our random weights, then calculate the loss for all predictions.

To visualize the loss, we can plot a graph where y is the loss value, and each other axis corresponds to a weight and its current value. This graph provides insights into how different weight values influence the loss, helping to identify trends such as whether certain weights significantly reduce or increase the loss, and guiding adjustments toward optimal values. For all weight values, there is an intersection point where y reaches its minimum. The minima is where the weight values at this point produce predictions for all sample that are the closest possible value to the actual value. This is the value we aim to find.

The above shows what that graph might look like with one weight, as two dimensions are easier to reason with. Visualizing beyond three dimensions is not possible because our spatial intuition is limited to three dimensions, making it challenging to interpret higher dimensional spaces. As shown in the graph, there's an optimal value for w₁ where y is at its lowest. The lower y is, the smaller the difference between the predicted values and actual values for each sample. Here, y refers to the loss value, which quantifies these differences and serves as the key metric in the optimization process. By minimizing the loss value, we iteratively adjust the weights to improve the model's accuracy in predicting outcomes.

Gradient Descent

# initialise an array of random weights for each feature
# x.shape(1), will give us the number of columns which represent features
# we add 1 to the number of features to account for a bias value
n_weight = x_train.shape[1] + 1
weights = np.random.randn(n_weight) * 0.01
print("Initial weights:", weights)

# we will also add a feature column of ones, which will be used for the bias value later
# np.c_ will horizontally concatenate 2 matrices e.g. x_train and a generated 1 column matrix of 1s
x_train = np.c_[x_train, np.ones(x_train.shape[0])]
x_val = np.c_[x_val, np.ones(x_val.shape[0])]

Above we've initialised our parameters with random values to provide a starting point.

To move these random initializations to the optimal value, we need to move the weights towards the point where the loss is at it's minimum. To compute this, we can calculate the partial derivative for w₁ with respect to the MSE function. The partial derivative for our weight will tell us the gradient of our weight for a particular loss value i.e it will tell us how much the loss is affected as we adjust our weight.

In the code below we update the weight by subtracting the gradient for the current value of of w₁ and all other parameters.

w1 -= w1_gradiant

This will allow it to go towards the minima regardless of being a downward or upward gradient. This process of updating iteratively updating the weights by calculating gradients is called Gradient Descent. Below shows our initial point and several iterative updates. You'll see the jumps are quite large and can overshoot the optimum value, causing the value to oscillate around the optimum and slowing convergence, or it might not converge at all.

To smooth out these updates, we use a hyperparameter called the learning rate represented by α (alpha).

Let’s replace our previous code line with the one below, taking into account the learning rate.

learning_rate = 0.01
w₁ -= w₁_gradient * learning_rate

Above we simply multiply the gradient by a learning rate to decrease the step size and smooth out the update, as seen on the graph below.

However setting α too low will result in more updates to achieve the optimum weight, resulting in a slower convergence and extra compute. Setting this too large results in the problems we see without using a learning rate to scale the gradient update.

A few sections back we scaled our features ensuring they were within similar ranges with a mean of 0 and a standard deviation of 1. Without scaling, weight updates for features with large ranges such as house values would have far greater jumps compared to a feature with a smaller range such as number of rooms. As the learning rate is the same for all features, scaling our features to be within similar ranges allows them to take similarly scaled weight updates resulting in faster convergence.

Features with a large scale will also dominate the learning process, as the feature would contribute significantly more than other smaller scale features, due to the larger numbers. You might think it's fine if larger features contribute more to the weight updates early on, since they’ll eventually take smaller steps as the gradient shrinks. However, without feature scaling, the optimization process becomes uneven. Large features dominate the updates initially, but smaller features may take much longer to make meaningful progress. This imbalance can lead to slower convergence overall, as the model takes a less efficient path towards the optimal solution. By scaling all features to similar ranges, we ensure that all features contribute to the weight updates more evenly, leading to a faster and smoother convergence.

Once we've updated our weights, we repeat the process, of calculating predictions, calculating the loss then updating weights using partial derivatives. This iterative process is known as gradient descent. We can carry on for a set number of iterations known as epochs or until the loss starts to start to slow down and stabilize across epochs, signaling that our model has stopped learning.

This process of iteratively updating the parameters proportional to the gradient of a loss function is called Gradient Descent. It is a foundational optimization method in machine learning, enabling efficient minimization of loss functions, and it will play a crucial role in more advanced algorithms we explore later.

# set the number of epochs to train for, 100 is good enough for this simple
# example, however a better approach would be stop learning once we see
# little change in the error metrics between epochs
n_epochs = 100
# set the learning rate, 0.1 is generally a good starting place, however
# later examples will adjust this during the learning phase to achieve
# faster convergence 
lr = 0.01

# the main training loop
for epoch in range(n_epochs):
    # calculate predictions using the dot product of two matrices, 
    # the feature matrix and the weight vector
    predictions_train = np.dot(x_train, weights)
    error = predictions_train - y_train
    gradients = 2 * (x_train * error[:, np.newaxis]).mean(axis=0)
    weights -= lr * gradients

    # calculate the mean squared error
    mse_train = np.mean(error ** 2)

    # calculate predictions on validation set
    predictions_val = np.dot(x_val, weights)
    mse_val = np.mean((predictions_val - y_val) ** 2)

    print(f"Epoch {epoch}: Training Error = {mse_train:.6f}, Validation Error = {mse_val:.6f} ")

Epoch 0: Training Error = 5.637618, Validation Error = 5.468886 
Epoch 1: Training Error = 5.438583, Validation Error = 5.277734 
Epoch 2: Training Error = 5.247518, Validation Error = 5.094238 
Epoch 3: Training Error = 5.064100, Validation Error = 4.918087 
Epoch 4: Training Error = 4.888018, Validation Error = 4.748984 
Epoch 5: Training Error = 4.718977, Validation Error = 4.586642 
........................
........................
Epoch 95: Training Error = 0.726476, Validation Error = 0.753565 
Epoch 96: Training Error = 0.721906, Validation Error = 0.749195 
Epoch 97: Training Error = 0.717503, Validation Error = 0.744986 
Epoch 98: Training Error = 0.713263, Validation Error = 0.740932 
Epoch 99: Training Error = 0.709177, Validation Error = 0.737027

The above output is an example we might see from the training loop. We can see at the start we have an a high error value, but training over 100 epochs brings the error rate down considerably. While it does start to slow down towards the end of the training phase, we are still seeing improvements epoch to epoch, and we could eke out better accuracy by continuing to train. However, there is a trade-off to consider, additional training may lead to marginal accuracy improvements, but it also increases the risk of overfitting, where the model performs well on training data but poorly on unseen data.

Vectorization

The training loops also make use of vectorized methods and NumPy, leveraging broadcasting when arrays have compatible but differing shapes, and element-wise operations when arrays have the same shape, to create short lines which are fast to run. Vectorized methods allow for SIMD (single instruction multiple data), which allows entire arrays to be processed in parallel by a CPU instruction, significantly reducing control overhead compared to traditional loops. Instead of iterating through each element individually, SIMD enables a single instruction to operate on multiple data points simultaneously, maximizing parallelism and improving performance.

Normally we iterate through an array which requires control instructions to loop through each element, then an instruction that must be applied to each element individually. Vectorized methods remove the need for many control instructions. Memory fetches are also minimized as instructions are performed on chunks of data in registers rather than single elements of an array.

Vectorized functions, especially in NumPy, are written in optimized compiled C code rather than Python, further enhancing performance. In our case rather, than iterating through our array vector and multiplying by each feature for each sample we can calculate them all at once. For this we use the dot product.

Below compares the use of iterative methods vs vectorized methods.

def iterative_predictions(x, weights):
    predictions = np.zeros(x.shape[0])
    for i in range(x.shape[0]):
        prediction = 0.0
        for j in range(x.shape[1]):
            prediction += x[i][j] * weights[j]
        predictions[i] = prediction
    return predictions

def vectorized_predictions(x, weights):
    return np.dot(x, weights)

# Time the iterative method
start_time = time.time()
predictions_iterative = iterative_predictions(x_train, weights)
iterative_time = time.time() - start_time
print(f"Iterative method took {iterative_time:.6f} seconds")

# Time the vectorized method
start_time = time.time()
predictions_vectorized = vectorized_predictions(x_train, weights)
vectorized_time = time.time() - start_time
print(f"Vectorized method took {vectorized_time:.6f} seconds")

Iterative method took 0.031959 seconds
Vectorized method took 0.000339 seconds

As we can see above the vectorized method is orders of magnitude faster using the vectorized dot product method.

As well as highly optimized methods such as the dot product. NumPy applies element wise operations on arrays using SIMD.

error = predictions_train - y_train

When performing a calculation such as the one above between two arrays of the same size, NumPy automatically performs the operation element-wise between the two arrays. This means an element of an array is matched with it's counterpart in a corresponding matrix and a calculation is completed as if they were single scalar values. For example, with predictions_train and y_train, which are both the same size, NumPy will subtract the corresponding elements (first element of predictions_train minus the first element of y_train, and so on), ultimately returning a new array of results.

NumPy can also handle arrays or matrices of differing sizes using a method called Broadcasting if they meet certain rules to be "broadcast-compatible."

These rules are as follows:

• Rank Compatibility: Each matrix must have the same number of dimensions.

• Dimension Compatibility: Dimensions are compared from right to left. If the size in each dimension are the same or if one of them is 1, then the dimension is considered compatible, and the operation can proceed. If one dimension is 1, that value is broadcasted across the larger dimension.

For example, if we have an array with shape (3, 1) and another array with shape (3, 4), NumPy will broadcast the single value across the dimension where 1 occurs to perform the element-wise calculation, effectively treating the first array as if it had the shape (3, 4).

Later in this series, we'll see how we can push these calculations to a GPU, which is designed specifically for massive parallel computations, allowing us to achieve even faster processing.

Limitations of Linear Regression

As the name suggests, linear regression is linear, which means it struggles when we need to model non-linear relationships. We only increase or decrease a feature by a magnitude, therefore it’s ability to model complex relationship is limited. e.g if we could apply exponents or polynomials to an input feature we could start to see non linear relationships.

For example, in the dataset shown above, a straight line from our linear model wouldn’t be able to capture the complex shape we see. While factors like exponents would allow for more complex lines, methods to learn those and combine several to create models would be limited. To model non linear relationships there are more appropriate models which do not involve learning these extra components.

Linear models are also limited in their ability to model relationships between input features. For instance, in our California Housing Dataset, if 1 bedroom houses when they are very old significantly increase the median income value, a linear model would struggle to capture this relationship, as it treats each feature independently and multiplies it by a weight.

Summary

Linear regression might be less popular than Transformers, Neural Networks, and other advanced models, but if a simple linear model is sufficient, it’s quick and easy to train.

You might then be wondering why we covered linear regression in the first place. The reason is that linear regression provides a solid foundation for understanding key concepts such as training with gradient descent, loss functions, and how weights influence a model's output. In linear regression, we use a mathematical function that takes some input features and has adjustable parameters known as weights. We apply a loss function to the model's predictions and then use gradient descent to adjust the weights until the loss function is appropriately minimized.

This approach to training can be extended to more complex models, however the core idea remains the same: we adjust parameters to minimize a loss function. The difference lies in the mathematical function we use, linear regression is simple, but we can replace it with more flexible models that can handle complex relationships, and help predict different types of data.

The next model we’ll intro are Neural Networks as our new mathematical function. Neural Networks allow us to model non-linear relationships and dependencies between features, enabling us to tackle much more complex datasets.

Complete Code