Chapter 2: Models

2.1 Neural Networks

Linear layers, matrix multiplication, and activation functions as the compute primitives behind generative models.

Generations of research into neural networks form the theoretical foundation for generative AI.

To be a productive inference engineer, you need a basic intuition for essential concepts in neural networks. This section provides a high-level introduction; Appendix B offers recommendations for further reading.

The fundamental unit of a neural network is a node (a.k.a., neuron). A node is a short program that takes an input, multiplies it by some weights, adds some bias, and returns the result.

A group of nodes forms a layer. Nodes within a layer are independent of each other – they do their own calculations. The connection between nodes, or the “network” in a neural network, is between layers, where the nodes in a layer receive the output of the previous layer.

The neural networks behind LLMs contain dozens to hundreds of layers. There are three types of layers:

  • Input layer: The first layer, which accepts and processes the input to the neural network.
  • Hidden layers: Every layer between the first and last, which iteratively transform the input to arrive at an output.
  • Output layer: The final layer, which returns the prediction from the network.

Each layer produces an output that the next layer reads as input. For the hidden layers, these outputs are called hidden states.

Hidden states are one kind of internal representation for data within a neural network. A key aspect of internal representation is its dimensionality, or the actual size of the vectors used.

Figure 2.1: Multi-layer neural networks have one input layer, many hidden layers, and one output layer.
Figure 2.1: Multi-layer neural networks have one input layer, many hidden layers, and one output layer.

Internal representations for text input increase the dimensionality, encoding text chunks into vectors of hundreds or thousands of numbers to capture semantic meaning. But internal representations for image models reduce the dimensionality from millions of pixels down to a manageable size.

There are neural networks for creating these internal representations, and there are neural networks for using them:

  • Encoder: Takes an input like text or an image and creates an internal representation of the input that includes additional information and semantic meaning.
  • Decoder: Uses the internal representation to generate an output like text or an image.

Neural networks are composable. You can combine multiple neural networks together into a single model or use them sequentially to build a pipeline.

Modern LLMs are decoder-only, while encoder-only models are somewhat rare today, with old-school text embedding models from the BERT family as a prominent example.

Many models in other modalities use an encoder-decoder architecture. Whisper, a popular open model for audio transcription, uses an encoder to process audio input and a decoder to generate text tokens.

2.1.1 Linear Layers and Matmul

The most essential operation within a neural network is a matrix multiplication, or matmul. A matmul takes an input vector (a list of numbers) and a matrix (a grid of numbers) and multiplies the vector through the matrix to produce an output vector.

Within a neural network, a linear layer is the simplest form of matmul. Given an input vector, the linear layer applies a weight matrix and adds a bias vector:

Figure 2.2: In a matmul, the output vector y is the product of an input vector x and a weights matrix W plus a bias vector b.
Figure 2.2: In a matmul, the output vector y is the product of an input vector x and a weights matrix W plus a bias vector b.

The weights of any given linear layer are a small part of a generative AI model’s total weights, and the individual values within the weights matrix are set during training.

2.1.2 Activation Functions

Matrix multiplication is composable, meaning that multiplying a vector by two matrices is equivalent to multiplying that vector by the product of those matrices.

Figure 2.3: Two matmul equations representing separate layers collapse due to composition of linearity.
Figure 2.3: Two matmul equations representing separate layers collapse due to composition of linearity.

This is a problem for multi-layer neural networks because a series of linear layers, each one a matmul, would collapse into a single layer with all of the matrices multiplied together.

Deep multi-layer neural networks are useful because more layers use more parameters effectively and encode more meaning in hidden states.

Neural networks separate layers by breaking linearity with an activation function. Activation functions are non-linear to prevent composable matmul from collapsing layers, and are differentiable or mostly-differentiable to support back propagation.

One of the most basic activation functions in inference is ReLU, which stands for Rectified Linear Unit. ReLU is a simple function: if X is greater than zero, return X, else return zero. There are dozens of activation functions – including one named “Swish” thanks to its resemblance to the Nike logo – but most follow the same general pattern of mapping negative values to zero or near-zero, while keeping positive values unchanged.

Figure 2.4: Activation functions like ReLU are used to break linearity in multi-layer neural networks.
Figure 2.4: Activation functions like ReLU are used to break linearity in multi-layer neural networks.

Activation functions like ReLU, SiLU, Swish, and SwiGLU are fast to run, easy to train on (as they are mostly differentiable, they have a gradient at least for most values), and break linearity to support multi-layer neural networks.