LLMs: Neural Networks

Mar 28
9 min read

Updated: Apr 3

I hope you have read my previous blogs covering the basic intuition and tokenization aspects of LLMs. If not, I highly recommend reading them.

LLMs: Build Intuition

LLMs: Tokenization

So, as you remember, the whole “magic” behind LLMs eventually boils down to a state-of-the-art mathematical equation.

Today, we are going to dig a little deeper into this mathematical equation, a.k.a. Neural Networks.

Most probably, you would have seen an image like this showcasing a complex neural net.

By the end of this blog, I promise that you will build the intuition around them.

Neural Network

A Neural Network is basically a multi-layer perceptron… ah, too many jargons. What is a perceptron?

Okay, wait—but we were discussing a mathematical equation, right? Where does a Neural Network come into play?

The image above represents a single neuron (perceptron), which denotes a small linear equation.

For example:

f(x) = w1 * i1 + w2 * i2 + b

f(x) = ax + by + c

Here, the coefficients of these equations are referred to as the weights of a neuron.

When you stack multiple such layers together to create more complex non-linearities (instead of a simple linear equation), we call it a Neural Network.

Let’s Build It

Okay, enough theory. Let’s build this in practice.

Note: torch is a Python neural network library. It’s completely fine if you are not familiar with it—we are focusing on building intuition.

import torch
import torch.nn as nn
input_size = 3 
model = nn.Sequential(
    nn.Linear(input_size, 2),
)
model = model.to(device)

In the above code snippet:

Input size is 3, i.e., our equation will have 3 variables
It has a single layer of size 2, i.e., 2 neurons
You can stack more layers with more neurons as needed

So, the above neural network actually looks like a single-layer neural network.

And if we convert this to a mathematical representation

And there is a more interesting way to represent this equation as matrices

For now, just remember:

Neural networks (mathematical representations) eventually boil down to matrix operations, and matrix operations are GPU-friendly.

(I will cover this more in a future blog.)

Our current neural network has 8 parameters (coefficients).

So now, we have built a basic understanding around neural networks, let's try to put them into a usecase

A Simple Use Case

We have a digital lock:

Input size = 3
Each input can be 0 or 1
Lock opens only if exactly one ‘1’ is present

Examples:

0 0 0 -> locked
0 1 1 -> locked
0 0 1 -> unlocked

Mapping:

locked = 0
unlocked = 1

Now, let's try supplying the input to our current "untrained" model

X_V = torch.tensor([[0.0, 0.0, 1.0]]).to(device)

logits = model(X_V.to(device))
prediction = torch.argmax(torch.softmax(logits, dim=1), dim=1).item()
print("Prediction:", prediction)

When I run the above code, it returns '0' (i.e., locked), which is incorrect. But this is kind of expected, right, since we haven't trained the model for our usecase, it's just a mathematical function with random values as coefficients. And it could probably have given a correct answer for this input, and an incorrect one for some other input

So, our job is to train this model for our usecase. And by train, I mean optimizing the values of these weights or coefficients instead of being random.

Training

That's where Gradient Descent comes into the picture. But before we dig deeper, let me use a real-life analogy, so that it becomes easier to understand.

Let's say you have an exam, Now what do you typically do to prepare for it?

You will have a course material or books that you learn from, right? For models, this is equivalent to a "training dataset."
You revise the material multiple times, right? Let's say you revise 3 times. For models, this is defined by "epochs."
For every epoch or revision cycle, you can either focus on mugging up completely or you can try building intuition gradually as and when with each cycle. So basically, there is some concrete amount by which you learn or remember. And if you try mugging up completely, you often forget what you previously read. The same quantity for models is defined as "learning rate". And a large learning rate is typically not better, same as humans :p
After each iteration of the revision cycle or epoch, you measure yourself by checking answers to quantify the level of preparation you currently have. For ex- by answering some questions and checking the number of correct ones. For ex- You answered 10 questions, and 6 were correct. So, currently you have a "loss" of 4, and we target to minimize this loss as much as we can by revising again. The same happens with models as well, and is defined by "loss" and is computed via "loss function".
- For an exam scenario, Loss Function = Total Questions - Correct Answers

Okay, so now we have some understanding of some of the parameters that are useful to control model training. We are ready for Gradient Descent

As I have told multiple times now, eventually, we are dealing with a mathematical equation

 f(x1, x2) = w1 * x1 + w2 * x2

here x1 and x2 are the inputs and w1 and w2 are weights or coefficients

For each different value of x1 and x2, our equation will give out some value, right? Which we compare with the expected answer

Let's say Y represents the correct answer, so our loss will be

loss = f(L) = Y - f(x1, x2)

similar to our analogy "total questions - correct answers" -> "expected answer - actual answer."

f(L) = Y - f(x1, x2)

replacing f(x1, x2)

f(L) = Y - (w1 * x1 + w2 * x2)
f(L) = Y - w1 * x1 - w2 * x2

Our target is to minimize f(L); this is where gradient descent comes into play.

If I go back to the basics of derivatives that we studied in high school

f = a * x + b * y

d(f) / d(x) = a 
d(f) / d(y) = b

what does derivate of d(x) and d(y) represent here?

d(f) / d(x) = a 
It shows that a unit change in "x" will increase "f" by a magnitude of "a" 

d(f) / d(y) = b
It shows that a unit change in "y" will increase "f" by a magnitude of "b" 

so if x = x + delta
then 
f will change by a magintude of "a" in the direction of "delta"

Note- I used the term "magnitude," and not "factor," as the relation is not linear

Okay, enough theory, let's try this out with a neural network

Let's say we have a single-layer neural network with 2 parameters

f(y) = f(x1, x2) = w1 * x1 + w2 * x2

let's assume these values-

w1 = 2
w2 = 3

x1 = 4
x2 = 7

also let's assume for given x1 and x2, the expected answer is 32

f(y) = f(4, 7) = w1 * 4 + w2 * 7

expected answer is 32

actual answer is f(4, 7) = 2 * 4 + 3 * 7 = 29

loss = expected - actual = 32 - 29 = 3

so our target is to minimize the loss

loss = f(L) = 32 - f(4, 7)
f(L) = 32 - (w1 * 4 + w2 * 7)
f(L) = 32 - 4 * w1 - 7 * w2

df(L)/d(w1) = -4
df(L)/d(w2) = -7

What do these derived values denote?

df(L)/d(w1) = -4

it denotes that for every "unit" increase in "w1", the loss "increases" by a magniute of "-4" 


df(L)/d(w2) = -7

it denotes that for every "unit" increase in "w2", the loss "increase" by a magniute of "-7"

So what exactly do these 'gradients' or 'derivatives' for a "loss" function denote? And why did we study them?

The main intuition is "if we travel in the direction of the gradient, then it will result in an increase in the value of the function and vice versa."

Here, our target is to minimize the loss function, so here we will have to move in the 'opposite' direction of the gradient to achieve our target

w(new) = w(old) - gradient

Typically, we multiply the gradient by some factor so that weights are not changed rapidly, and this factor is known as "learning rate."

w(new) = w(old) - gradient * learning_rate

Now, coming back to our example, let's say our learning rate is 0.01

learning_rate = 0.01

df(L)/d(w1) = -4
df(L)/d(w2) = -7

old values of w1 and w2:
w1 = 2
w2 = 3

new values of w1 and w2 will be:

w1 = 2 - (-4 * 0.01) = 2.04
w2 = 3 - (-7 * 0.01) = 3.07

Now let's calculate the loss with new values:

loss = f(L) = 32 - f(4, 7)
f(L) = 32 - (w1 * 4 + w2 * 7)
     = 32 - (2.04 * 4 + 3.07 * 7)
     = 2.35

So, we were able to reduce our loss from 3 to 2.35

Let's iterate it again to reduce the loss further:

learning_rate = 0.01

df(L)/d(w1) = -4
df(L)/d(w2) = -7

old values of w1 and w2:
w1 = 2.04
w2 = 3.07

new values of w1 and w2 will be:

w1 = 2.04 - (-4 * 0.01) = 2.08
w2 = 3.07 - (-7 * 0.01) = 3.14

loss = f(L) = 32 - f(4, 7)
f(L) = 32 - (w1 * 4 + w2 * 7)
     = 32 - (2.08 * 4 + 3.14 * 7)
     = 1.7

So after running 2 iterations (epochs) with a learning rate of 0.01, we reduced our loss from 3 to 1.7. (almost half, remarkable isn't it), And as we run it for more epochs, our loss will reduce further.

So in a nutshell:

Gradient Descent is a process of updating the weights based on gradients so that our loss is minimized.

And in the training phase, we are running multiple iterations (epochs) of gradient descent, where you do a "forward pass" to calculate the actual value and"backpropagate" the loss to update the weights.

And after "n" forward and backward passes, we get our desired set of weights or parameters.

Example Contd.

So, since now we have a deep understanding of training a neural network, let's put them into action in our "digital lock" example:

The first step is to prepare a training dataset so that our model can learn patterns from it.

For our usecase, the below will serve as the training set since it contains some of the possible (not all) cases for our usecase

    [0.0, 0.0, 0.0] -> 0
    [1.0, 0.0, 0.0] -> 1
    [0.0, 1.0, 0.0] -> 1
    [1.0, 1.0, 0.0] -> 0
    [0.0, 1.0, 1.0] -> 0
    [1.0, 0.0, 1.0] -> 0
    [1.0, 1.0, 1.0] -> 0

Performing Gradient Descent

import torch.optim as optim

X = torch.tensor([
    [0.0, 0.0, 0.0],
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0],
    [1.0, 1.0, 0.0],
    [0.0, 1.0, 1.0],
    [1.0, 0.0, 1.0],
    [1.0, 1.0, 1.0]
]).to(device)

y = torch.tensor([0, 0, 0, 1, 1, 1, 0])

loss_function = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Training loop (1000 epochs)
for epoch in range(1000):
    # Forward pass
    outputs = model(X).to(device)

    # Calculating Loss
    loss = loss_function(outputs, y.to(device))
    print(f"Epoch {epoch+1} | Loss: {loss.item():.4f}")

    # Backward pass
    optimizer.zero_grad()
    loss.backward()

    # Weights Update
    optimizer.step()

Note:

In actual problems, we don't use a linear loss function (expected - output); instead, they are more complex functions, since the pattern is not linear in real life
SGD - Stochastic Gradient Descent, runs gradient descent in batches instead of running individually for each dataset for efficient hardware use

Output:

Notice that our loss has reduced significantly.

Epoch 1 | Loss: 0.8561
Epoch 2 | Loss: 0.8427
Epoch 3 | Loss: 0.8307
Epoch 4 | Loss: 0.8199
......................
......................
......................
Epoch 999 | Loss: 0.5786
Epoch 1000 | Loss: 0.5786

If I just plot the equation again, to see the updated weights:

Now, if I run it for the same input [0.0, 0.0, 1.0], I get the correct answer. Notice that it is not present in our training dataset, and our model has learned the pattern.

Inferencing:

X_V = torch.tensor([[0.0, 0.0, 1.0]]).to(device)

logits = model(X_V.to(device))
print(logits)

# output:
[ 0.1156, -0.0326]

As you would have noticed, the output from our neural network is not classifiable currently, i.e., what does "-0.0326" denote?

We want these raw numerics to be converted into some representation that's easy to interpret

This is where the last step comes into picture, i.e., Generating the Probabilistic Distribution.

Typically, we use the "softmax" function for this conversion.

normalized = torch.softmax(logits, dim=1)
print(normalized)

# output:
[0.4630, 0.5370]

Here, if you notice, the final output represents the probabilities of each label, so our current model is claiming it to be "locked" with a probability of 46% and "unlocked" with a confidence of 53%

Activation Function and Bias:

Just to mention at the end, the neural network we built was relatively a very simple one and not enough to handle complex problems, where we need more non-linearities. Typically, the output of each neuron is passed through a mathematical function (activation) to add more non-linearities.

Example of activation functions could be tanh(), sigmoid(), relu() etc

We also add "bias" as an additional parameter that can be trained and independent of inputs to have more control if needed

without activation function and bias:

f(x1, x2) = x1 * w1 + x2 * w2

-------------------------------------------------------------

with bias (c):
f(x1, x2) = x1 * w1 + x2 * w2 + c

-------------------------------------------------------------

with activattion function (tanh):

F(x1, x2) = tanh(f(x1, x2))

So I hope you were able to build some intuition around Neural Networks and all the "jargon" surrounding it. I believe this understanding is more than enough to understand any complex neural architecture/workflow independently.

In the next blog, we will be exploring "Token Embeddings" as we progress through the journey of understanding LLMs. Stay tuned

Aditya x Mittal