Neural Networks

We don’t know what the desired mapping from context embedding vectors to feature scores actually looks like, so we will need to define the main part of our parametric function in such a way that it is flexible enough to represent a vast range of different relationships for different parameter values and hope that, given the right parameter values, it can approximate the desired "awesome feature extraction" function $f_\text{main}$ well enough.

For our toy model with 1D input and output, we said an intuitive way to achieve this would be to define a function made up of a bunch of continuously connected linear segments, whose boundaries, slopes and offsets are parameters we seek to optimize. Unfortunately, generalizing this to multidimensional inputs seems very non-trivial — even for a 2D input, it already isn’t clear how exactly we should define the boundaries of our linear regions, which would now be curves instead of points, let alone how to define our parameterized function in such a way that it is always continuous along all these boundary curves.

So while we could take on these challenges and proceed down a more mathematical route, let’s see if we can make our life easier and get inspiration from somewhere else instead.

The human brain

Ignoring the LLMs we’re reinventing, the only system that we know can perform sophisticated next-token prediction is the human brain, which means whatever type of function the human brain is implementing must definitely be powerful enough to do well.

Of course, we are far from understanding the human brain in detail and trying to recreate exactly what it is doing would likely be hopeless anyway, but just as we used some key observations about human development in coming up with our imitate + supervise + reinforce plan, we might be able to identify some basic patterns in how the brain functions and leverage them to tackle our problem.

The brain is largely comprised of two types of cells — neurons, which transmit electric signals and are the computational cells we are interested in, and glial cells, which support and protect neurons and serve all sorts of crucial auxiliary functions. Let's take a closer look at how neurons work.

A single neuron

Neurons are very diverse and can differ from each other in a million ways both big and small, but a typical neuron functions somewhat like this:

1. Inputs: It receives input (voltage) signals via terminals called synapses. These voltage signals can be both excitatory (promoting output) or inhibitory (suppressing output) and either stem from other cells (usually other neurons) or sometimes directly from the outside world.

2. Signal aggregation: Each of these input signals locally modifies the voltage near the synapse in question. Then, a complicated mess of biophysical processes determines how, if at all, these voltage changes at the synapses impact the voltage in the cell body, the so-called soma.

3. Thresholding: If the voltage in the cell body (soma) rises above a certain threshold, a biophysical chain reaction is triggered that causes a voltage signal to be sent to other neurons.

4. Output signal: For most neurons, this output signal is a so-called action-potential signal, which is binary — it is either sent or not sent, and the strength of the signal is always the same. However, the rate of these action-potential firings can vary and effectively correspond to different output signal strengths, since output signals can be integrated over small time windows by downstream neuronal circuitry.

So far, so good. But how are all these neurons wired together?

A biological neural network

We are far from understanding the brain in detail and no full map of all the neurons and their connections in a human brain has ever been made†The current state of the art for full-connectome reconstruction seems to be about $1\,\text{mm}^3$ of mouse brain, which is the sort of thing that sounds rather unimpressive if you don't think about it much and completely mindblowing if you do.[note], but we aren't trying to recreate the brain here, we just want to derive some inspiration from it to construct a powerful type of function that can learn to do many things.

So let's focus on just the simplest possible structural observations and see how far that gets us:

1. Many connections per neuron: How many other cells any given neuron connects to varies widely across the brain, but it's often many, varying from just $5\text{–}10$ input and $\sim\!10^2\text{–}10^3$ output connections to other cells (for the most common type of neuron) all the way to neurons receiving inputs from $\sim\!10^5$ or sending outputs to $\sim\!10^5\text{–}10^6$ other cells. The type of neurons which mostly comprise the neocortex (where high-level reasoning is performed) tend to connect to $\sim\!10^3\text{–}10^4$ other cells on both the input and the output side.

2. Layered structures: A recurring theme across different parts of the brain is that distinct layers of neurons can be identified. For example, both the neocortex (high-level reasoning) and the cerebellar cortex (learning and prediction) have clearly discernible layers of neurons and light impinging on our retinas gets processed into the brain by clearly distinguishable layers of cells as well. Even though tons of feedback loops via so-called recurrent connections and endless other complexities exist, a predominant direction of information flow can usually be discerned across these layers, with the outputs from one layer forming the inputs of another and so on.

Good enough. Let's see if we can leverage these insights to build up a simple but powerful function we can use for language modeling.

Artificial neural networks

If we would like to define a function inspired by the structure of the brain, we will first have to define some sort of mini-function that can act as a simplified, Artificial Neuron and then figure out how to make a bigger, more powerful function out of many such mini-functions, essentially wiring many artificial neurons together to form — drum roll — a Neural Network.

An artificial neuron

As mentioned above, biological neurons effectively receive inputs in a weighted fashion, allowing for both inhibitory and excitatory effects as well as varying sensitivity depending on the biophysical details of the synapses in questions.

The absolute simplest way to model this given some inputs $x_1, \,\dots, x_n$ would be to feed our artificial neuron a weighted sum of these inputs, i.e. to weight $x_1$ by some input weight $w_1$, $x_2$ by some input weight $w_2$, and so on. If we think of the various inputs and weights as elements of vectors, this weighted sum can simply be thought of as the dot product of a weight vector and the input vector, just like the weighted sum of the context-feature scores in the previous chapter:

$$\underbrace{\raisebox{-0.2em}{\vphantom{\mu}}w_1 x_1 \;+\; \dots \;+\; w_n x_n}_{\raisebox{-1em}{$\textstyle \vct{w} \cdot \vct{x}$}}$$

where the weights $\vct{w}$ are real-numbered parameters of our little artificial neuron model.

We also said that biological neurons will usually output a signal only if those weighted inputs add up to exceed some threshold $t$, so let’s say the same should be true for our artificial neuron, meaning its output, which is usually referred to in machine learning as the neuron’s Activation $a$, should be zero whenever the weighted input sum fails to exceed the neuron’s activation threshold:

$$a = 0\hspace{15px}\text{if}\hspace{15px} \vct{w} \cdot \vct{x} \,\leq\, t$$

In this framing, the threshold $t$ is a parameter of the so-called Activation Function determining the neuron’s output activation for a given sum of inputs. People usually prefer to think of the activation function itself as parameter-free and absorb the threshold parameter in the weighted sum instead, yielding an adjusted input sum known as the Preactivation $z$:

$$z \;=\; \vct{w} \cdot \vct{x} - t$$

But since writing things as simple sums is nicer than having mixed signs, this is more commonly written as

$$z \;=\; \vct{w} \cdot \vct{x} + b$$

Here, we use the letter $b$ just like we did in the previous chapter, since we can think of the real-numbered parameter $b$ as a Bias added to the input of the artificial neuron, in the sense that a neuron with larger (more positive) bias is more likely to be active and a neuron with lower (more negative) bias is less likely to be active.

In this framing, in which the threshold is absorved into the preactivation itself, enforcing the thresholding requirement simply translates to enforcing that the activation of the neuron should be $0$ whenever its preactivation is $0$ or less:

$$a = 0\hspace{20px}\text{if}\hspace{20px} z \leq 0$$

But what if $z \geq 0$? In other words, what should the output look like if the artificial neuron is actually active?

For biological neurons, we said the strength of the output tends to reflect the strength of the overall input, and that said output strength might be encoded in the frequency of output rather than the voltage of the signal. But encoding stronger signals in more frequently produced output and then integrating over time again seems completely inappropriate for our purposes — after all, we are not interested in simulating the brain’s dynamics, we merely want to define a simplified toy function roughly modeling the “net” computation implemented by the neuron.

So let’s avoid this frequency-encoding and simply define the output of our artificial neuron such that it is larger when the neuron's preactivation is larger. In other words, let's make sure our activation function is an increasing function. Of course, there are infinitely many such functions to choose from, but keeping with our philosophy of starting simple and making adjustments as needed, let’s just use what is arguably the simplest possible option — a linear function with slope one. Combining this with the thresholding requirement, we get

$$a(z) \;=\; \begin{cases} 0 & \text{if } z \leq 0 \\ z & \text{if } z > 0 \end{cases}$$

This can be written more compactly as simply choosing the larger value from $0$ and $z$, and is known in machine learning as the ReLU function:

$$a(z) \;=\; \underbrace{\raisebox{-0.1em}{\vphantom{\mu}}\max(0, z)}_{\raisebox{-1.2em}{$\textstyle \relu(z)$}}$$

If we plot it, it looks like this:

The name $\relu$ derives from “Rectified Linear Unit”, with the unit in question being the artificial neuron and the word “rectified” indicating the thresholded behavior. Despite its simplicity, the $\relu$ function turns out to work remarkably well for machine learning.

So far, so good! Having decided to define our artificial-neuron mini-function this way, let’s return to the brain to get some more guidance on how we might want to wire these artificial neurons together to hopefully get one big parameterized function that is sufficiently powerful to serve as the core of our language model.

An artificial neural network

Our main takeaway from our (incredibly superficial) look at the structure of the brain was that there was a persistent theme of densely connected layers of many neurons processing information.

Let's adopt this simple structure by constructing our overall function from many artificial neurons, each with its own input-weight and bias parameters, and by grouping those neurons in layers as well. In principle, these layers could vary in size and number, but to keep things as simple as possible, let's give each layer the same number of neurons $N$ for now and consider both that number $N$ and the number of layers $L$ hyperparameters we will have to choose values for when building an actual model.

Naturally, the inputs of the first layer need to be the actual inputs. In our LM scenario, these will be the feature scores of our context embedding vector, but let's just consider them elements of a general real-number input vector $\vct{x}$ for now, as we already did for the individual neuron above. To honor the general notion of "many connections per neuron", let's simply connect each input to all the neurons in the first layer. If we visualize this, it looks something like this:

Algebraically, the output activation of the $i$-th neuron in the first layer will then be

$$a_i^{(1)} \;=\; \relu\left(\vct{w}_i^{(1)} \cdot \vct{x} \;+\; b_i^{(1)}\right)$$

But writing this for every single neuron individually is very unwieldy, so let's condense our notation. First, let's pack all the input weights in one big first-layer weight matrix $W^{(1)}$, with the $i$-th row representing the input weights $\vct{w}_i^{(1)}$ of the $i$-th neuron:

$$W^{(1)} \;=\; \begin{pmatrix} \;\rule[0.5ex]{4em}{0.4pt} & \vct{w}_1^{(1)} & \rule[0.5ex]{4em}{0.4pt}\; \\[1em] & \smash{\vdots} & \\[0.8em] \;\rule[0.5ex]{4em}{0.4pt} & \vct{w}_N^{(1)} & \rule[0.5ex]{4em}{0.4pt}\; \end{pmatrix}$$

Secondly, let's pack all the biases $b_i$ of the first layer neurons into a single vector $\vct{b}^{(1)}$ as well. We can then write the output of the entire first layer as a vector $\vct{a}^{(1)}$ which we can compute from the input vector $\vct{x}$ like this:

$$\vct{a}^{(1)} \;=\; \relu\left(W^{(1)} \,\vct{x} \;+\; \vct{b}^{(1)}\right)$$

Sweet. Let's keep building our network.

The whole idea of a layered structure is that the outputs from one layer serve as inputs to the next layer. If we once again connect everything to everything, meaning every neuron in a given layer to every neuron in the subsequent layer†Since each connection in our model has a weight parameter associated with it that can equal $0$ in principle, this everything-to-everything connection scheme is not just a simple choice but actually the most general layer-to-layer connectivity pattern possible.[note], it looks something like this:

Each layer here is simply implementing the same type of operation as the first layer, except the inputs are the $N$ activations from the previous layer instead of the original inputs $\vct{x}$. Algebraically, the activations $\vct{a}^{(l+1)}$ of the $(l\!+\!1)$-th layer neurons are then given by

$$\vct{a}^{(l+1)} \;=\; \relu\left(W^{(l+1)} \, \vct{a}^{(l)} \;+\; \vct{b}^{(l+1)}\right)$$

In contrast to the calculations performed by the brain, the main function $f_\text{main}$ we are trying to define for our LM scenario has a clear output, namely the vector of context feature scores our model subsequently uses to calculate the individual next-token scores.

So as a final step, we should map the outputs of our last layer of neurons onto this feature vector. Arguably the simplest way to do this is to once again use our simple weighted-sum-plus-bias recipe, so let's do that:

In matrix lingo, this corresponds to performing the same type of operation we performed in each layer, except that we are not applying the $\relu$ activation functino here since we would like to be able to have negative outputs:

$$\vct{y} \;=\; W^{(\text{out})} \; \vct{a}^{(L)} \;+\; \vct{b}^{(\text{out})}$$

Et voilà! Using this simple layer-wise wiring, we have built a type of parameterized function that takes in a vector of real-numbered inputs (such as our context embedding vector) and produces a vector of real-numbered outputs (such as our context feature scores), with the parameters being all the individual input weights and biases of all the neurons and final outputs.

Functions like this, which are composed of many small units that can be loosely thought of as artificial neurons, are called Neural Networks, and the specific type we just invented is referred to as a Fully Connected Feedforward Neural Network, with the term “fully connected” indicating that each neuron is connected to all neurons in the previous layer, and the term “feedforward” indicating that no recurrent connections are present, meaning information flows through the network only in the forward direction, from input to output.

The layers of artificial neurons are usually referred to as the Hidden Layers of the network, since they can be thought of as being "hidden" between the inputs and outputs interfacing with the external world. Even though they do not contain any activation functions mimicking neurons, the inputs and outputs are also commonly framed as forming layers, which are simply referred to as the Input Layer and the Output Layer.

Neural networks are instrumental to modern machine learning and form a key part of most state-of-the-art AI models, including large language models like ChatGPT, but so far, our hope that they might be able to approximate a broad range of powerful functions is resting solely on the fact that we drew some loose inspiration from the brain in inventing them. So let's take a closer look at what all this corresponds to mathematically and see if this hope feels warranted.

Understanding $\relu$ networks

If we were to use a neural network as the core part of our model, it would ultimately have to take in a large context embedding vector as its input and, given the complexity of the task, would likely have to process that input through a number of intermediate layers with many neurons. But thinking about high-dimensional input spaces is tricky and visualizing them is virtually impossible, plus considering a bunch of layers at once is very overwhelming.

So let's start small.

1 input, 1 layer, 1 output

For now, let's consider only the simplest $\relu$ network we could think of — a network with just 1 input, 1 output, and 1 hidden layer of $\relu$ neurons in between.

The preactivation of the $i$-th hidden-layer neuron in such a network is simply

$$z_i^{(1)} \;=\; w_i^{(1)} \, x \,+\, b_i^{(1)}$$

where, as before, $w_i^{(1)}$ is the input weight for that neuron and $b_i^{(1)}$ is its bias. As a result, the activation of the neuron equals

$$a_i^{(1)} \;=\; \relu\left(w_i^{(1)} \, x \,+\, b_i^{(1)}\right)$$

As a function of the input $x$, this corresponds to a linear function with slope $w_i^{(1)}$ and y-axis intercept $b_i^{(1)}$ whose negative part has been zeroed out. Let's play around with these parameter values a bit to get a feel for this:

Since all the weights and biases of the individual hidden-layer neurons are independent parameters, each neuron can represent an arbitrary function of this type, and the network as a whole can represent any function that can be written as a linear combination $\vct{w}^{(\text{out})} \cdot \vct{a}^{(1)}$ of these $\relu$ outputs (plus an optional bias $b^{(\text{out})}$):

$$y \;=\; \vct{w}^{(\text{out})} \cdot \vct{a}^{(1)} \;+\; b^{(\text{out})}$$

So let’s figure out what such linear combinations look like.

Picturing the output

A useful way of thinking about the single-neuron activation graph above is that each individual neuron output is made up of two pieces — the zero part, and the linear part — with each neuron having its own transition point $x_{0,i} = -b_i^{(1)}/w_i^{(1)}$, where the neuron's preactivation equals exactly zero and the neuron transitions from being inactive to being active or vice versa:

If we consider all the hidden-layer neurons simultaneously, these individual transition points partition the input space into a bunch of Regions, with each region corresponding to a specific and unique subset of neurons being active:

Inactive neurons have zero output and might as well not exist, so for any one such region, the overall output of the network will simply be the weighted sum of the outputs of the active neurons. But each of those active neurons simply represents a linear mapping of the input, and adding a bunch of linear mappings together just gives another linear mapping!

This last fact is intuitive graphically since adding a bunch of lines of differing slopes together simply results in a new line whose slope is the sum of the individual ones, but it's also easy to show algebraically:

What this means is that the graph of the overall function $y(x)$ that the network represents takes the form of a straight line for each region. But since the slope of that line depends on the weights tied to the active neurons, and since a different subset of neurons is active in each region, the slope of this line can vary from region to region!

What’s more, since crossing the boundary between two regions corresponds to an individual neuron becoming active or inactive, and since any neuron’s output is continuous at that transition, the output of the network is guaranteed to be continuous at the boundaries of the regions (and is, of course, continuous within the regions already, as it corresponds to the output of a simple linear function in each region).

This is quite amazing. It means that a single-hidden-layer $\relu$ network with one input and one output represents a continuous, piecewise-linear function whose parameter values determine the widths and locations of the various linear regions as well as the slopes and offsets of the corresponding line segments. In other words, it represents the exact type of function we said would be an intuitive choice for the purpose of function approximation when discussing our simplified toy model with a single input and output in the previous chapter!

The reason we considered such a function desirable is that we can easily envision approximating any continuous function on some finite interval by using a bunch of linear segments if we can freely choose the segments’ slopes, offsets and boundaries.

This brings up the question whether our single-layer $\relu$ network allows for this freedom — is it merely the case that any such network represents some piecewise-linear, continuous function, or is it true that any such function can be implemented using an appropriately constructed network, and that we can therefore approximate virtually any function $f(x)$ (on some finite interval) using a single-layer $\relu$ network?

Let’s find out.

Universal approximation

We have already established that the boundaries of the linear regions are located at the individual-neuron activation thresholds $x_{0,i} = -b_i^{(1)}/w_i^{(1)}$ where the preactivations $z_i$ equal zero and the neuron in question turns on or off. This means we can always create a new region simply by adding another neuron, and we can freely choose the location of the corresponding boundary in input space by choosing the neuron's input weight and bias accordingly.

But what about the slope of the output in that new linear region?

We already determined above that the algebraic expression for the slope of any given linear region is

$$\sum\nolimits_{\text{active neurons}}\; w_i^{(\text{out})}\; w_i^{(1)}$$

This means that whenever a newly added neuron becomes active, the slope will change by $\,w_{\text{new}}^{(\text{out})} \, w_\text{new}^{(1)}\,$, which we can easily make arbitrarily small or large, for example by keeping $\,w_\text{new}^{(1)}\,$ fixed and choosing whatever value we desire for $w_{\text{new}}^{(\text{out})}$.

This implies we can tune the slope of the network output in the new region to be whatever we want, and suggests a simple recipe for approximating any continuous 1D function (on some finite interval) using a single-layer $\relu$ network: Simply add a new neuron with input weight $w_i^{(1)}\!=\!1$ whenever you would like to change the slope of the network output, using $b_i^{(1)}$ to ensure the neuron becomes active at the given point in input space, and using $w_i^{(\text{out})}$ to adjust the slope of the overall output to your liking.†Not convinced? Simply play around with the biases and output weights in the widget above and convince yourself you can adjust all the segments at will.[note]

For example, here is what this would look like if we chose to approximate one period of a sine wave with five neurons:

In math speak, this means single-layer $\relu$ networks are Universal Approximators of 1D functions — at least in theory, we can approximate any continuous function $f(x)$ on some finite interval arbitrarily well using a single-layer $\relu$ network.

Given that we started this neural-network journey looking for a powerful function architecture that would be flexible enough to model whatever unknown functional relationship underlies next-token prediction, this is a highly encouraging result.

That being said, it isn’t more than that either — the mere fact that there must exist some hypothetical network approximating a function well neither means that it can be done with a reasonable amount of neurons nor that we will actually be able to find the right parameter values in practice, when the function we wish to approximate is unknown and we can’t simply look at a graph and eye-ball our slopes and boundaries. Plus, this is all assuming we are dealing with only 1 input and 1 output so far, whereas the inputs and outputs of the network we hope to use as the core part of our language model will have to be be many-dimensional.

So before we worry about anything else, let’s see if any of this promise even survives if we go beyond one-dimensional inputs, or if we have to make further adjustments to our function architecture.

$N_i$ inputs, 1 layer, 1 output

Let's say we have a context embedding dimension $N_i$ and our network must therefore deal with $N_i$ inputs. So for $N_i\!=\!2$, for example, the network would look something like this:

The preactivation of the $i$-th hidden-layer neuron is then given by

$$z_i^{(1)} \;=\; \vct{w}_i^{(1)} \cdot \vct{x} \;+\; b_i^{(1)}$$

where, as before, $\,\vct{x} \!\in\! \mathbb{R}^{N_i}\,$ is the vector of inputs, $\,\vct{w}_i^{(1)} \!\!\in\! \mathbb{R}^{N_i}\,$ is the vector of input weights for that neuron and $b_i^{(1)}$ is the neuron's bias. As a result, the activation of the neuron equals

$$a_i^{(1)} \;=\; \relu\left( \vct{w}_i^{(1)} \cdot \vct{x} \;+\; b_i^{(1)} \right)$$

Picturing the output

As before, the neuron will become active/inactive whenever the preactivation is exactly zero, meaning in this case whenever

$$0 \;=\; \vct{w}_i^{(1)} \cdot \vct{x} \;+\; b_i^{(1)} \quad \implies \quad \vct{w}_i^{(1)} \cdot \vct{x} \;=\; -b_i^{(1)}$$

With a single input $x$, this condition simply yielded the individual boundary points $x_{0,i} = -b_i^{(1)}/w_i^{(1)}$. Let's see what it corresponds to in input space for the more general vectorial case.

Recall that the dot product of two vectors always equals the product of their magnitudes times the cosine of the angle between them:

$$\vct{a} \cdot \vct{b} \;=\; \|\vct{a}\| \, \|\vct{b}\| \, \cos\theta$$

A useful way of thinking about this equation is that $\|b\| \cos\theta$ is simply the signed scalar component $b_{\text{along }\vct{a}}$ of the vector $\vct{b}$ along the vector $\vct{a}$:

So another way of thinking of the dot product of two vectors is that it equals the magnitude of one of the vectors times the scalar projection of the second onto the first:

$$\vct{a} \cdot \vct{b} \;=\; \|\vct{a}\| \, b_{\text{along }\vct{a}}$$

Using this intuition, our above dot-product condition simply translates to

$$\| \vct{w}_i^{(1)} \| \; x_{\text{along }\vct{w}_i^{(1)}} \;\,=\;\, -b_i^{(1)}$$

In other words, the preactivation will equal exactly $0$ and the neuron will become active/inactive whenever the scalar component of the input vector $\vct{x}$ along the vector of the neuron's input weights $\vct{w}_i^{(1)}$ equals a specific value:

$$x_{\text{along }\vct{w}_i^{(1)}} \;=\; -\frac{b_i^{(1)}}{\| \vct{w}_i^{(1)} \|}$$

But this only constraints the input vector in the direction of the weight vector. If a given input $\vct{x}$ yields $z\!=\!0$, taking a step in any direction in input space that is orthogonal to $\vct{w}_i^{(1)}$ is not going to move the needle on the preactivation. So generally, the set of points in input space yielding $z\!=\!0$ for a given neuron will not just consist of isolated points but be of dimension $N_i\!-\!1$.

For example, for a two-dimensional input ($N_i\!=\!2$), the boundary across which a given neuron becomes active/inactive will take the form of a straight line in input space. If the neuron's bias $b_i^{(1)}$ equals $0$, this line will simply be the line through the origin that is perpendicular to the direction of $\vct{w}_i^{(1)}$, since that is the line for which $x_{\text{along }\vct{w}_i^{(1)}}\!=\!0$.

If the neuron's bias is not equal $0$, the activation boundary will still be perpendicular to the direction of $\vct{w}_i^{(1)}$, but offset from the origin along (or opposite) that direction to ensure that the scalar component of $\vct{x}$ along $\vct{w}_i^{(1)}$ equals $-b_i^{(1)}/\| \vct{w}_i^{(1)} \|$.

Of course, the weight vectors and bias parameters for the individual neurons can and will all be different, so the activation boundaries of all the individual neurons will be different as well. Collectively, they will partition the input space into a bunch of regions defined by straight lines. For $N_i\!=\!2$, this looks something like this:

For three-dimensional inputs, these region boundaries would correspond to planes instead of lines, and for higher dimensions they will correspond to hyperplanes, but nothing about the basic logic changes. It's all based on dot-product intuition that readily extends to higher dimensions.

As before, each individual region will correspond to a specific subset of all the neurons being active. To figure out what the network output looks like as a function of the input $\vct{x}$ for such a region, we can follow the same logic as before, the only differences being that our input is now a vector $\vct{x}$ and each neuron has multiple input weight vectors $\vct{w}_i^{(1)}$, instead of merely multiple scalar input weights $w_i^{(1)}$:

But by definition, the dot product of the "slope vector" $\vct{m}$ and the input $\vct{x}$ is just

$$\vct{m} \cdot \vct{x} \;=\; m_1 x_1 \;+\; \dots \;+\; m_{N_i} x_{N_i}$$

meaning the output of the network is simply a linear function of all the inputs $x_i$ (with different slopes for each input) plus some offset!

In other words, going from a single input to many just means the output for any given region goes from looking like a straight line to looking like a plane (or hyperplane, if we have more than two inputs). It's still straight and sloped along a certain direction, it just has more dimensions. As before, this slope can change from region to region.

For $N_i\!=\!2$ inputs and $N\!=\!5$ neurons, this looks something like this:

Universal approximation

So far, so good. But the thing we are really trying to figure out is whether the universal-approximation property we observed for a single input still holds now that we have many.

Based on the observation that we can slice up the input space into arbitrarily many regions by adding more and more neurons to the hidden layer and that each neuron will give us more parameters to control the output across these regions, we might be optimistic that it does.

But it's far from obvious whether this is indeed true, because the activation of each neuron can only vary along a single direction in input space, namely that of its input-weight vector $\vct{w}_i^{(1)}$. When we only had a single input, having full control along a single direction was enough, and we could simply splice our desired function together step-by-step by introducing new region boundaries via additional neurons and adjusting the slope of the output as we saw fit. With a multi-dimensional input, we can still always add another neuron and use it to adjust the slope of the output in the half of input space where the neuron is active, but it's not obvious how we can use all the resulting neuron outputs to locally sculpt an arbitrary function.

Fortunately, there is a default recipe in math for constructing continuous functions from components that vary only along a single dimension: the Fourier transform. We will not go down the rabbit hole of explaining the intuition underlying Fourier transforms here, but for our purposes, the key insight is that any continuous function can be approximated arbitrarily well on some finite interval by a sum of plane waves, i.e. of functions that vary sinusoidally along only a single direction.

Using our 1D approximation recipe, we can approximate plane waves of arbitrary wavelengths pointing along arbitrary directions even in higher dimensional space, simply by choosing our neuron parameters such that all the activation boundaries are parallel and stacked in the direction along which the plane wave exhibits variation. Here is what this would look like for a 2D input space, for example:

So given a sufficiently large number of neurons to construct all these individual plane waves, we can theoretically make our network output approximate any continuous function arbitrarily well on some finite interval.⁠[Don't like it?]

But by definition, the dot product of the "slope vector" $\vct{m}$ and the input $\vct{x}$ is just

$$\vct{m} \cdot \vct{x} \;=\; m_1 x_1 \;+\; \dots \;+\; m_{N_i} x_{N_i}$$

meaning the output of the network is simply a linear function of all the inputs $x_i$ (with different slopes for each input) plus some offset!

In other words, single-layer $\relu$ nets are universal approximators even for many inputs.

$N_i$ inputs, $L$ layers, 1 output

So in principle, if we could have arbitrarily many neurons and were able to find the perfect set of parameter values, a single-layer network would be enough to approximate pretty much any function arbitrarily well.

But of course, we cannot afford to have arbitrarily many parameters in practice. Seeing as we already observed that the brain is partially comprised of layered structures and already hypothesized our artificial network should likely contain multiple layers as well, let's see what happens if we do add additional layers, i.e. if $L\!>\!1$.

As always, let's start simple and hope general trends will reveal themselves. For $N_i\!=\!2$ and $L\!=\!2$, the network would look something like this:

The equations mapping the inputs of such a two-layer network to its output are given by

$$\begin{align*} \vct{a}^{(1)} \;&=\; \relu\left(W^{(1)} \,\vct{x} \;+\; \vct{b}^{(1)}\right) \\[1em] \vct{a}^{(2)} \;&=\; \relu\left(W^{(2)} \, \vct{a}^{(1)} \;+\; \vct{b}^{(2)}\right) \\[1em] y \;&=\; \vct{w}^{(\text{out})} \cdot \vct{a}^{(2)} \;+\; b^{(\text{out})} \end{align*}$$

Let's see what this does.

Picturing the output

As far as the first-layer neurons are concerned, absolutely nothing about our previous logic has changed. Each first-layer neuron will switch from being active to inactive (or vice versa) when its preactivation $z_i^{(1)}\!= \vct{w}_i^{(1)} \cdot \vct{x} + b_i^{(1)}$ equals zero. As explained above, this will be the case along $N_i\!-\!1$-dimensional activation boundaries. For our simple 2D-input example, this leaves us with the same straight-line boundaries as before:

But what about the second-layer neurons? Looking at our network structure, we see that the preactivation of a second-layer neuron is calculated the exact same way the output of our single-layer network was calculated, meaning we can do the same exact algebra as earlier:

Complicated as this might look, the intuition is simple: Only active neurons matter. Along each individual active path through the network ending up at our neuron of interest, a given input will get weighted by the product of weights along the path. If we sum over all these paths for each input and lump all the resulting effective weights for the individual inputs into a vector, we get the above $\vct{w}_{\text{eff}}$. Similarly, the overall bias $b_{\text{eff}}$ of our preactivation is just the sum of the bias of the given neuron plus a weighted sum of the biases of earlier neurons.

Even though the above algebra is for the two-layer example, this intuitive understanding readily extends to networks with more layers, simply because it's always true that only active neural pathways matter and that weights get multiplied along those paths.

So what does this mean for our activation boundaries and the resulting linear regions?

For a fixed activation pattern in the earlier layers, the preactivation $\vct{w}_{\text{eff}} \cdot \vct{x} + b_{\text{eff}}$ of a given neuron will equal zero along an $N_i\!-\!1$-dimensional hyperplane that is orthogonal to the weight vector $\vct{w}_{\text{eff}}$. This is simply due to the same dot-product logic we discussed earlier when figuring out the activation boundaries for a single-layer network. So for a two-dimensional input, the activation boundary for a fixed activation pattern in the earlier layers will still just be a straight line.

But crucially, the effective weight vector $\vct{w}_{\text{eff}}$ depends on which neurons are active in the earlier layers. So whenever we cross an activation boundary of an earlier-layer neuron in input space, the effective weight vector determining the orientation of our later-layer activation boundary will change! So for a 2D input space, the activation boundary of later-layer neuron is not just a straight line, but a straight line that changes direction whenever it crosses an earlier neuron's activation boundary:

If we plot this for a full two-layer network with $N\!=\!4$ neurons per layer and $L\!=\!2$ layers, we get four such kinked activation boundaries:

If we add more layers, we simply get more and more activation boundaries which can get redirected more and more frequently, namely wherever they cross any activation boundary belonging to an earlier-layer neuron. For example, if we simply add two more layers, the regions in input space might look something like this:

In higher dimensions, the same logic applies, it's just that instead of kinked lines, the activation boundaries of later-layer neurons will be kinked $(N_i\!-\!1)$-dimensional hyperplanes. We can't really visualize it but it's the same exact idea – each first-layer activation boundary cuts the input space in half along one direction, every later-layer activation boundary does so until it encounters an earlier-layer boundary at which point its direction might change. As a result, the input space is divvied up into lots of regions in a parameter-dependent fashion.

But what does the output look like for each of these regions?

Crucially, each of these regions corresponds to a unique activation pattern, meaning to a specific subset of the network's neurons being active, just as it did for our single-layer network. But if we restrict ourselves to such a fixed subset of neuron's being active to begin with, the $\relu$ function doesn't do anything at all and each layer simply outputs a linear combination of the activations from the previous layer plus some offset. Even if we chain many such layers and do this over and over again, the end result will still be a linear combination of the inputs plus some bias, exactly as it was for the single-layer network.

In other words, for any one region, we can still always write

$$y \;=\; \vct{m} \cdot \vct{x} \,+\, b_\text{eff}$$

meaning the output for each region still takes the form of an $N_i$-dimensional hyperplane. For a network with $N_i\!=\!2$ inputs, the output looks like this:

Universal approximation

So far, so good, but what does this imply for the network's ability to flexibly approximate arbitrary functions?

In some sense, it doesn't change much. Given that single-layer $\relu$ networks are universal approximators, we already know that multi-layer $\relu$ networks must be, too, because any multi-layer network could always be wired to simply pass on the outputs of its first-layer neurons and act as a single-layer network. In other words, for the same number of neurons $N$ per layer, adding more layers can only increase the representational capacity of a network, never decrease it.

But this universal-approximation property is just a theoretical guarantee for the limit of arbitrarily many neurons, and comparing multi-layer networks to single-layer networks with the same number of neurons per layer is comparing apples to oranges. In practice, more parameters will mean more complexity and more compute, so we will only be able to afford a certain number of them. So the question is what, if anything, adding more layers gets us, if we have to cut down the number of neurons $N$ per layer in turn.

Shallow vs. deep

Let's try to get some intuition for this by thinking about it in more mathematical terms.

A "wide" single-layer $\relu$ network with many neurons can be thought of as representing a parameterized function $f_\text{wide}\left(\vct{x}\right)$ with a relatively large number of parameters.

If we restructure such a network to instead have multiple "narrower" layers with fewer neurons, we are effectively chaining together many functions $f_\text{narrow}^{(l)}$ with fewer parameters.

So the question is what's more desirable: Having a single, more powerful parameterized function operating directly on input space, or being able to chain a number of less powerful parameterized functions together?

$$\smash{\underbrace{\raisebox{-0.7em}{\vphantom \mu} \hspace{2px} f_\text{wide}\left(\vct{x}\right) \hspace{1px}}_{\raisebox{-1.1em}{\small\text{single-layer net}}}} \qquad \text{vs.} \qquad \smash{\underbrace{\raisebox{-0.7em}{\vphantom \mu} \hspace{2px}f_\text{narrow}^{(L)} \left( f_\text{narrow}^{(L-1)} \left( {\scriptstyle \dots} \left( f_\text{narrow}^{(2)} \left( f_\text{narrow}^{(1)} \left( \vct{x} \right) \right) \right) {\scriptstyle \dots} \right) \right)}_{\raisebox{-1.1em}{\small\text{multi-layer net}}}} \\[2.7em]$$

If we assume both have the same amount of parameters, the answer to the question can only lie in how these parameters are used, so let's think about that.

For a single-layer network, the hidden-layer output activations $\vct{a}^{(1)}$ are calculated by forming weighted sums of the network inputs, shifting those weighted sums by the neuron biases, and then killing off all the negative parts via the $\relu$ activation:

$$\vct{a}^{(1)} \quad=\quad \overbrace{\raisebox{0.5em}{\vphantom A}\relu}^{\mathclap{\raisebox{0.3em}{\small\text{kill negative parts}}}} \;\left(\; \smash{\underbrace{\raisebox{-0.3em}{\vphantom \mu} W^{(1)} \, \vct{x}}_{\mathclap{{\raisebox{-1em}{\small\text{weighted sums}}}}}} \quad\smash{\overbrace{\raisebox{0.5em}{\vphantom A}+\; \vct{b}^{(1)}}^{\mathclap{\raisebox{0.3em}{\small\text{shift sums}}}}} \;\right) \vphantom{\underbrace{\raisebox{-0.3em}{\vphantom \mu}}_{\raisebox{-1em}{\small X}}}$$

Visually, the weighted sums of the inputs simply correspond to $N_i$-dimensional hyperplanes, meaning the resulting activations $\vct{a}^{(1)}$ (after the $\relu$ killed off negative parts) are a bunch of $N_i$-dimensional half-planes in input space.

The output of the network is then simply a weighted sum of these half-plane "building blocks", plus some shift:

$$y \quad=\quad \vphantom{\overbrace{X}^{\raisebox{-1em}{\small X}}} \smash{\underbrace{\raisebox{-0.3em}{\vphantom \mu} \vct{w}^{(\text{out})} \cdot \vct{a}^{(1)}}_{\mathclap{{\raisebox{-1em}{\small\text{weighted sum}}}}}} \quad\smash{\overbrace{\raisebox{0.5em}{\vphantom A}+\; \vct{b}^{(\text{out})}}^{\mathclap{\raisebox{0.3em}{\small\text{shift}}}}} \vphantom{\underbrace{\raisebox{-0.3em}{\vphantom \mu}}_{\raisebox{-1em}{\small X}}}$$

Since each of these half-plane building blocks is tied to one of the neurons, each parameter†The only exception being the output bias, which simply shifts the output up or down.[note] of the network affects a single half-plane contribution to the output, whether it's by affecting the activation boundary where the half-plane starts (in the case of the neuron bias parameters $\vct{b}^{(1)}$), by affecting the half-plane's slope (in the case of the output weights $\vct{w}^{(\text{out})}$), or by affecting both (in the case of the input weights $W^{(1)}$).

This means that, by and large, the network's parameters are on equal footing, and each of them has a semi-global impact on the output in this specific half-plane sense. While we have established that we can, in theory, approximate arbitrary functions if we combine enough such half-planes, this does not seem like a particularly efficient use of our parameters, both because most things aren't conveniently represented by sums of half-planes and because the fact that each parameter can only impact the output in a very specific way makes the network very inflexible.

For a multi-layer network, the situation is exactly the same in the beginning. The activations $\vct{a}^{(1)}$ of the first hidden layer can be visualized as half-planes in input space. We just have fewer such half-planes now, because our multi-layer network must have fewer neurons per layer.

After that, each subsequent layer does more of the same: build weighted sums of its inputs, apply shifts, kill of negative parts.

$$\vct{a}^{(l+1)} \quad=\quad \overbrace{\raisebox{0.5em}{\vphantom A}\relu}^{\mathclap{\raisebox{0.3em}{\small\text{kill negative parts}}}} \;\left(\; \smash{\underbrace{\raisebox{-0.3em}{\vphantom \mu} W^{(l+1)} \, \vct{a}^{(l)}}_{\mathclap{{\raisebox{-1em}{\small\text{weighted sums}}}}}} \quad\smash{\overbrace{\raisebox{0.5em}{\vphantom A}+\; \vct{b}^{(l+1)}}^{\raisebox{0.3em}{\small\text{shift sums}}}} \;\right) \vphantom{\underbrace{\raisebox{-0.3em}{\vphantom \mu}}_{\raisebox{-1em}{\small X}}}$$

But crucially, the inputs to each subsequent layers are the outputs of the previous layer. Meaning while the input-space building blocks for the second-layer activations are half-planes, the building blocks for the third layer are something built from half-planes, the building blocks for the fourth layer are something built from that, and so on.

From a parameter-usage point, this seems very intriguing. While parameters affiliated with first-layer neurons still determine the slopes and boundaries of a bunch of half-planes which are then fed to the second layer, the parameters of neurons in later layers are directly tied to far more complex functional landscapes across input space, namely whatever scalar functions $a_i^{(l)}(\vct{x})$ the output activations from the previous layer happen to represent.

As these functional landscapes themselves are essentially outputs of a smaller neural net (namely the network up to that layer), they should be able to take near-arbitrary shapes given enough neurons. As a result, multi-layer nets should, in principle, be able to employ their parameters in a vastly more flexible fashion and generally be able to approximate functions much more efficiently, meaning using fewer parameters, than a single-layer net.†Of course, the parameterization will only be more efficient if the scalar functions represented by earlier-layer activations are actually useful building blocks, meaning all of this hinges on our ability to have our model learn the appropriate parameter values. As promised, we will start tackling this issue in the next chapter.[note]

Intriguingly, this argument does not even rely on our choice of the $\relu$ activation function or other minutae of our network. It's simply an outcome of the fact that we are processing the input information layer-by-layer with each layer representing a powerful parameterized mapping. This automatically allows later layers to operate on more complex mappings of the inputs and rather generally suggests that Deep networks of many layers might make for flexible and efficient function approximators. Indeed, this lies at the heart of the success of Deep Learning, which essentially means "machine learning with deep, multi-layer networks".

The scalar functions across input space represented by the activations of individual neurons are sometimes referred to as (hidden or learned) Features, which is related to but different from our earlier notion of a feature on which our tokens are scored. Unfortunately, the term feature is used in a number of related but different ways in ML contexts and people are often very sloppy about what exactly they are referring to.⁠[more]

To visualize the notion of the network building increasingly intricate features from one layer to the next, as well as the basic idea of approximating arbitrary functions, here is a step-by-step animation of a simple $\relu$ net approximating Half Dome (a famous monolith in Yosemite National Park):

$N_i$ inputs, $L$ layers, $N_o$ outputs

Any network playing the role of $f_\text{main}$ in our model architecture would necessarily have to produce a large number of context feature scores, not just a single number.

So as a final step of our deep dive, let's examine what, if anything, changes if our $\relu$ network has more than one output. We already visualized this when we first wired our neurons together to form a full network earlier, but as a reminder, here is what it looks like:

Computationally, the only difference to the single-output scenario is that the activations of the final hidden layer are now mapped to a vector $\vct{y}$:

$$\vct{y} \;=\; W^{(\text{out})} \; \vct{a}^{(L)} \;+\; \vct{b}^{(\text{out})}$$

Picturing the output

For the single-output scenario, we determined that the network output could be visualized as a bunch of spliced-together plane segments, with each segment covering a region in input space corresponding to a specific subset of neurons being active.

Since we can think of each of the outputs of our multi-output network as the output of a single-output network, the same is still true now. We just have multiple such outputs, meaning we have to think of each of them as one such surface across input space.

Because the outputs are all computed as weighted sums of the same building blocks, namely the activations of the neurons in the final hidden layer, these individual output surfaces are not independent, however, and must all share the same activation boundaries.

Here is what this might look like for a small network with $N_i\!=\!2$ inputs, $N_o\!=\!4$ outputs and $N\!=\!5$ neurons per layer:

Universal approximation

The extension of our previous universal-approximation logic is also straightforward: We've already established that we can approximate any single output $y_1(\vct{x})$ arbitrarily well over some interval. If we would like to additionally approximate some other $y_i(\vct{x})$, we could simply add the required neurons to all the layers and avoid any cross-talk with the existing $y_1(\vct{x})$ subnetwork by choosing all the cross-connecting weights to be zero and essentially having independent subnetworks. In this case our outputs would still be built from the same overall set of final-layer activations, but all the individual outputs would simply rely on independent subsets of those activations.

That said, this is once again just a theoretical guarantee for the limit of arbitrarily many neurons, and having a bunch of independent subnetworks would be horribly inefficient in practice.

Recall that the very hope underlying our choice to define a single function $f_\text{main}$ extracting high-level features from the input contexts was that the same features might be relevant to many different next-token probabilities. Having this function shared amongst tokens will help with compression exactly to the degree that our learning problem does not require independent sub-networks for all the individual token scores.

Bringing it all together

This was a lot, so let's briefly summarize the key takeaways and then tie everything back to our quest of building a better language model.

A brief summary

Taking loose inspiration from the human brain, we can define artificial neural networks by defining an "artificial neuron" mini function and wiring layers of such mini functions together to map an input vector $\vct{x}$ to an output vector $\vct{y}$.

An intuitive and straightforward way to define such an artificial neuron is to have it form a weighted sum of its various inputs and enforce simple thresholding behavior, with the neuron becoming active only if said weighted sum is large enough. This is exactly what the $\relu$ function does.

Wiring a bunch of such $\relu$neurons together in a sequence of layers yields a feedforward $\relu$ network.

For a fixed subset of neurons being active, such a network simply represents a linear mapping and its ouput can be visualized as a hyperplane segment. But whenever an activation threshold is crossed and some neuron in the network becomes active or inactive, the orientation of that plane segment can change, and so the overall output can be thought of as a bunch of tiled-together segments.

A closer examination of these types of outputs reveals that such $\relu$networks are universal approximators, meaning given enough neurons, they could in principle approximate any continuous function on some finite interval arbitrarily well. Taken on its own, this basic theoretical fact implies little regarding their practical usefulness, but it does give us reason to hope they might be useful.

Another reason we might have that hope is that deep networks allow for later-layer parameters to be tied to highly intricate and tailored functional features across input space, which should be conducive to parameter-efficient approximations.

A neural-network-based architecture

This is all very promising, so let's give neural nets a try. For now, let's define our architecture's main function $f_\text{main}$ to be a feedforward $\relu$ net:

But wait a minute! The various $a_i$ in the above architecture are just simple "remixes" (weighted sums plus biases) of the final neuron activations, which we are then remixing again to obtain the token scores. This double remixing is completely redundant, so we might as well drop the intermediate step and replace the two-step mapping with a single step.†This is the same thing we already discussed earlier when analyzing the output of a $\relu$ network: Simply chaining multiple layers of weighted sums and biases together still results in a simple linear mapping, which is why for a fixed neuronal activation pattern, the output still just looks like a hyperplane even with many layers.[note]

This will remove the explicit stage of "context feature scores" able to take on any real-numbered value from our architecture, but we can still think of the output activations of the final neuron layer as representing non-negative context feature scores, and because the two formulations are equivalent, we are not losing any functionality at all:

Conceptually, this pretty much completes our quest for a (rudimentary) function-approximation architecture, but we are not quite ready to fully cross the model-definition bullet point from our list of urgent issues yet — defining our model does not just mean specifying the types of operations it should perform; it means specifying the exact mathematical function it represents. For our simple $\relu$ network, that means we need to decide how many layers to use and how many neurons to place in each layer.

This finding of suitable hyperparameter values is a different problem from the separate challenge of figuring out how to find suitable parameter values. We hope to find the latter via some sort of fitting procedure, but the former change the very structure of the function we are aiming to fit.

In the absence of any clear priors on what values our hyperparameters should take, this feels like an intensely empirical question, meaning we will simply have to figure out what model size and shape are best by training different models and comparing their performance down the road, once we have figured out how to perform such model training in the first place.

This only intensifies our pre-existing conviction that everything hinges on our ability to figure out a suitable method of finding appropriate parameter values though, so let's focus on that challenge in the next chapter.

Next Chapter