Building LLMs from Scratch: A Deep Dive into Forward Computation and Feature Space Transformation

Introduction: Decomposing Large Models into Tensor Operations

Nanjing University's School of Computer Science offers a course series called Foundations of Large Language Models: A Zero-to-One Implementation Journey. It takes a highly condensed, engineering-first approach to help students understand how LLMs work from the ground up. In Lecture 2 (L2), the instructor zeroes in on a central question: How does a model perform forward computation? The subtitle — "Transformations in Feature Space" — hints at the mathematical essence underlying this process.

The course has a clear focus: it's not a conventional deep learning theory class. Instead, it moves at a brisk pace to fill in the most fundamental yet critical concepts, building intuitive understanding for students without prior background. One defining characteristic: the model is always analyzed through the lens of how tensors change shape, rather than through abstract architectural terminology.

Models as Functions: The Origin of Forward and Backward Passes

At the most fundamental level, any deep learning model can be understood as a function f(x). You feed in an x, internal computation happens, and you get an output y. In the early days, this process had no special name — it was simply a function call, an inference run.

So where did the term "forward" come from? The key lies in the learnable parameter W inside the model. The classical learning algorithm is Back Propagation — starting from output y, it traces back along the computation path toward input x to update parameters. Once "backward" became established terminology, the x-to-y computation naturally came to be called the "forward pass."

In PyTorch, this concept is codified as an engineering convention: every module that inherits from nn.Module must implement a forward method. When you pass x into the model object, the framework automatically triggers forward to execute the computation. Backpropagation is handled by PyTorch's automatic differentiation engine — at the application level, developers only need to focus on implementing the forward pass.

The Model Skeleton: init and forward

A standard PyTorch model class typically contains two core methods:

• __init__ (initialization): Handles the "static" part — defining which layers the model contains and how parameters are initialized, like snapping building blocks together.
• forward: Handles the "dynamic" part — describing how data flows through the layers and how computation is performed when input x arrives.

Whether you're reading the Transformer source code on HuggingFace or any complex LLM implementation, this init + forward structure is consistent throughout. The code is just larger and the module nesting deeper.

Parsing Model Inputs: 3D and 4D Tensor Structures

The model's input x is essentially a tensor. For natural language models, the input is typically a 3D tensor with three dimensions:

• B (Batch): The number of requests or samples. If 10 users send prompts simultaneously, B = 10.
• S (Sequence): The length of each prompt — the number of tokens.
• D (Dimension): The feature vector dimension for each token, such as 4096 (4K), 8192, or 128.

Here's a key insight: the first two dimensions are just "quantity" indicators — the semantic information is carried entirely by the last dimension's vector. Each token is represented as a high-dimensional vector of floating-point numbers. These values are learned during training and aren't directly human-interpretable, but the model uses them to understand meaning.

As for 4D tensors, they typically appear in Multi-Head Attention — the 4K dimension is split evenly across multiple heads (e.g., 8 heads), allowing the data to be arranged and processed as a 4D tensor.

Feature Space Transformation: The Mathematical Essence of Linear Layers

The course's central concept — "transformation in feature space" — gets a concrete explanation here. Deep learning models are often described as doing Representation Learning: their core purpose is to teach the model to "represent" a token or sample as a high-dimensional vector.

The instructor uses a four-class cat-and-dog classification example to illustrate this vividly. At input time, differently colored cats and dogs are mixed together, indistinguishable. But after multiple transformations through a deep model, similar categories cluster together in the final high-dimensional feature space, becoming linearly separable — a single line can now split different classes apart.

This is the point of feature space transformation: no matter what the input's spatial representation looks like, a single matrix operation projects it into a new space better suited for downstream tasks.

From Linear Regression to Deep Networks

The classic linear transformation formula is y = xW + b. Applying this just once to directly produce an output is a "shallow model." The essence of deep learning is stacking: the output of one layer becomes the input of the next, with repeated transformations and nonlinear activation functions inserted between them.

A brief historical note: early models were simply called "neural networks." As layers were stacked deeper, they ran into the vanishing gradient problem, making training difficult to converge. Around 2010, a series of breakthroughs solved the challenge of training deep networks, and the word "Deep" was added — giving us Deep Neural Networks (DNNs). Today's large language models are fundamentally deep neural networks with tens or even hundreds of stacked layers.

Understanding Matrix Operations Through the Lens of Shape

A recurring methodological emphasis in this course is: understand each operation in the model by observing how the shapes of inputs, outputs, and parameters change.

Take a linear layer as an example: an input tensor of shape BHD, multiplied by a weight matrix W of shape D×E, produces an output of shape BHE. Recall the basic rule from linear algebra: an M×N matrix multiplied by an N×K matrix yields an M×K result — the middle dimension N is eliminated. When the input is a 3D tensor, the B dimension can be thought of as a "batch dimension" — B matrices each independently multiplied by the same W.

From the perspective of individual tokens, this is even more intuitive: there are B×H tokens in total, and each token is projected from a D-dimensional space to an E-dimensional space via a linear transformation. The same weight matrix W is shared across all tokens — this is parameter sharing in action.

Broadcasting and Element-wise Operations

In PyTorch, @ denotes matrix multiplication, while * denotes element-wise multiplication. There's also an important mechanism: Broadcasting.

When a 3D tensor needs to be added to a 1D bias vector, this wouldn't be valid by strict mathematical definition — but PyTorch follows rules to scan dimensions from right to left, automatically expanding the lower-dimensional tensor to align with the higher-dimensional one before performing the operation. It's worth noting: broadcasting doesn't happen arbitrarily. It follows strict rules, and if shapes can't be aligned, you'll get an error.

Parameter Management and Modular Model Construction

Unlike the constantly changing input x, W and b are the model's static parameters. In PyTorch, they are wrapped with nn.Parameter, which unlocks several conveniences:

• Parameter counting: Automatically aggregates the total number of model parameters (this is how 7B, 13B model scales are calculated).
• Device migration: A single call to model.to(device) moves all parameters from CPU memory to GPU memory.
• Save and load: Parameters can be located, swapped, and persisted via state_dict.

Building Blocks: From Linear Layers to a Full Feed-Forward Network

The course demonstrates how to build a miniature FFN (Feed-Forward Network) module using linear layers: an up layer expands the vector's dimensionality (e.g., 4K → 8K), followed by an activation function (such as SwiGLU gated activation), then a down layer projects it back to the original dimension.

The course also introduces Residual Connections — analogous to parallel circuits in electronics. The input x travels down two paths: one goes through the computation, the other is a direct copy of x, and the two are merged with addition at the end. This design is one of the key techniques that enables deeper models to train successfully.

One noteworthy engineering detail: in large language models, linear layers do not include a bias by default — this is disabled by setting bias=False at creation time. Traditional computer vision models (like ResNet), on the other hand, typically retain bias terms.

Course Highlights and Hands-on Practice

Another standout feature of this course is its strong practical orientation. The course has partnered with Huawei to integrate their Agent framework (Open9n), and plans to distribute API token credits for domestic LLMs for students to complete their assignments.

For the homework, Assignment A0 requires students to manually implement a linear layer, including explicitly saving the forward input and manually computing the backward pass — which requires knowledge of partial derivatives and the chain rule. Learners are advised to review partial differentiation from calculus in advance, as preparation for the upcoming backpropagation content.

Summary

L2 begins with "models as functions" and uses tensor shape transformations as the central analytical tool throughout, breaking down LLM forward computation into a series of understandable matrix operations and feature space projections. The key takeaways are:

1. The key to understanding a model is understanding how tensor shapes transform;
2. A linear layer's essence is projecting a token from one dimensional space to another;
3. Depth comes from stacking layers combined with nonlinear activations;
4. PyTorch's init/forward structure is the universal skeleton for all models.

This teaching philosophy — "start from implementation, anchor on shape" — is highly valuable for learners who want to genuinely understand how large models work, rather than remaining at the level of conceptual buzzwords.