kai wen ng

Posted on Aug 12

Building an Inference Engine from Scratch - Part 2: Basics of Attention Mechanism

#ai #deeplearning #learning #machinelearning

Before building an LLM inference engine, I need to understand the fine-grained details behind how modern LLMs work.

The book "Build a Large Language Model from Scratch" has been an excellent companion throughout this journey:

https://www.manning.com/books/build-a-large-language-model-from-scratch

It not only explains the theory behind Transformers, but also provides a practical understanding of how each component connects together to form a working language model.

Why Transformers?

Before Transformers, many sequence models were based on RNN architectures.
However, RNNs have several limitations:

Sequential computation makes training and inference slow.
Long sequences suffer from information loss due to limited hidden-state memory.
Parallelization is difficult.

What is Attention?

At its core, attention allows each token to determine which other tokens are important when building its representation.
Starting from token embeddings:

Token Embeddings
        |
        |
Linear Projection
        |
   --------------
   |     |      |
   Q     K      V

Each embedding is transformed into:

Query (Q): What information is this token looking for?
Key (K): What information does this token contain?
Value (V): What information should be passed forward?

For each token:

The query vector is compared with all key vectors using dot product:

Attention Score = Q × Kᵀ

The scores are normalized using softmax:

Attention Weight = Softmax(Attention Score)

The weights are applied to the value vectors:

Context Vector = Attention Weight × V

The result is a context-aware representation where each token contains information from other relevant tokens.

Multi-Head Attention

Instead of learning only one attention pattern, Transformers use multiple attention heads.
Each head can learn different relationships:

Syntax relationships
Long-distance dependencies
Semantic similarity
Position-based patterns

The outputs from all heads are combined to create a richer representation.

Causal Attention

For autoregressive language models, the model should only use previous tokens when predicting the next token.

Example:

The cat sat on the ...

When predicting the next word, the model should not have access to future tokens.
This is achieved using a causal mask:

Before Softmax:

[0    -inf  -inf]
[0     0    -inf]
[0     0     0 ]

The -inf values become zero probability after applying softmax.

Dropout

Dropout is a regularization technique used during training. During training, some activations are randomly disabled so the model does not depend too heavily on specific neurons or pathways. This encourages the network to learn more robust representations and reduces overfitting. During inference, dropout is disabled.

Layer Normalization

Layer normalization stabilizes activations throughout the network. It normalizes hidden representations and applies learnable:

Scale parameters
Shift parameters

This helps maintain stable signal propagation through deep Transformer layers.

Activation Function: GELU

Modern Transformers commonly use GELU instead of ReLU.
ReLU:

ReLU(x) = max(0, x)

completely removes negative values.
GELU provides a smoother activation function, allowing small negative inputs to contribute instead of being completely discarded which in turn improved representation learning in deep networks.

Hidden Dimension / Neurons

The hidden dimension determines the size of the representation space.
A larger hidden dimension allows the model to capture more complex patterns and relationships.

However:

More parameters require more memory.
More computation is required.
Larger models require more training data.

Residual Connections (Shortcut Connections)

Transformer blocks contain residual connections:

Output = Layer(x) + x

These shortcut paths:

Improve gradient flow.
Allow deeper networks to train effectively.
Preserve information between layers.

This concept was introduced in residual networks and became a fundamental component of modern deep learning architectures.

Transformer Block

A Transformer is not only an attention mechanism.
A typical Transformer block contains:

Input
 |
Layer Normalization
 |
Multi-Head Attention
 |
Residual Connection
 |
Layer Normalization
 |
Feed Forward Network
 |
Residual Connection
 |
Output

Keeping input and output dimensions consistent allows multiple Transformer blocks to be stacked together.

Model Parameters

To understand the size of an LLM, we need to know how many trainable parameters it contains.
During model implementation, the number of parameters can be calculated using numel, which counts the total number of elements inside a tensor.
For example, each weight matrix and bias vector contributes to the overall parameter count:

Total Parameters = numel(all model tensors)

Logits and Token Generation

The final layer outputs logits.
Logits are raw scores before converting them into probabilities.

Applying softmax:

Logits → Softmax → Probability Distribution

gives the probability of each possible next token.
The simplest decoding method:

argmax(probability)

selects the highest probability token.
To generate more diverse outputs, modern LLMs often use sampling strategies on the probability distribution

Top-k sampling
Top-p sampling

Autoregressive Generation Loop

LLMs generate text one token at a time.

The process:

Input tokens are passed through the Transformer.
The model predicts the next token.
The predicted token is appended to the input.
The process repeats. This simple loop is the foundation behind text generation in modern LLMs.

Building an inference engine from scratch is giving me a deeper understanding of what happens behind every API call.
The goal is not only to run models, but to understand every matrix multiplication, tensor transformation, and design decision that makes modern LLMs possible.

DEV Community