Relearning LLMs and RL from First Principles

English · 中文 · Français

Writing note: The ideas and original draft material are mine. ChatGPT assisted with structuring, editing, and translation.

I have used many of the ideas behind language models and reinforcement learning before. Some of them I first encountered during my PhD; others became part of the normal vocabulary of working with modern AI: Softmax, Cross-Entropy, Attention, KL Divergence, Value Functions, Policy Gradients.

I knew the formulas. I could use them.

But at some point, especially as large language models became something I was using almost every day, I started noticing a gap between recognizing a mathematical object and really understanding why it had that form.

Recognizing a formula is not the same as understanding why it has that form.

Take Softmax:

The two-class case makes one part of the structure especially visible. If

then the probability of the first class can be written as

Move the slider below and watch what happens as the relative difference between the logits changes.

The visualization runs in your browser with Marimo. Source notebook.

I had seen the Softmax equation countless times. But why the exponential? Why use the natural exponential base? Why not divide the scores by their sum? Why should adding the same constant to every logit leave the probabilities unchanged?

Once I started asking questions at that level, other familiar formulas became strange again.

When we write

where did that probability actually come from?

The dataset gives us observations, not a table of probabilities.

When we maximize likelihood, why should the fact that something was observed mean that a model ought to assign it higher probability?

When policy-gradient methods rely on terms such as

how can that help when the action itself was sampled discretely?

These are basic questions. In some sense, they are much more basic than the research problems I had been working on. That was exactly why I found them uncomfortable.

I had developed enough familiarity with the machinery that I could move past these equations without stopping at them.

So I decided to stop.

What I mean by “first principles”

The goal of this project is not to collect a list of equations that everyone working on LLMs is supposed to know.

It is also not to pretend that every successful method can be derived uniquely from some small set of mathematical axioms.

Many choices in machine learning are exactly that: choices.

A probability identity, a modeling assumption, a useful parameterization, an optimization convention, and an engineering decision do not have the same status.

Part of understanding a formula is knowing what kind of thing it is.

For each mathematical object, I want to come back to questions such as:

• What problem are we trying to solve?
• What quantities do we actually observe?
• What has to be assumed?
• Which steps are mathematical identities?
• Which steps are modeling choices?
• Why is this particular choice useful?
• What would break if we changed it?

The test I have in mind is simple:

If I forget the equation, can I reconstruct the reasoning that led to it?

Softmax is a good example. Remembering

is easy. Understanding becomes more interesting when the equation can be reduced to

or equivalently,

Now the logits are no longer just mysterious numbers waiting to be normalized. Their differences describe log probability ratios.

That is the kind of transition I want to preserve in these notes: not only the final equation, but the question that made the equation worth understanding.

Learning by explaining

There is another reason I am writing this series.

I want to get better at explaining technical ideas.

It is surprisingly easy to convince myself that I understand something when the explanation remains inside my head. Writing forces a different standard. If I cannot explain why a step follows from the previous one, or if I quietly replace a modeling choice with “this is just how it works,” the gap becomes visible.

So explanation is not something I do after learning.

It is part of the test of whether I have learned.

The process I have gradually settled into looks roughly like this:

The order is not always clean. Sometimes an implementation exposes a misunderstanding. Sometimes a small numerical example does more than a derivation. Sometimes an explanation that sounded convincing falls apart when I try to write it precisely.

Those failures are useful. I do not want to remove all of them from the final articles and replace them with a textbook-style path in which every idea appears obvious in retrospect.

The question “why does Softmax use an exponential?” matters to me partly because I really asked it.

Learning LLMs with an LLM

There is an obvious recursive aspect to this project: a large part of this learning process happens through conversations with ChatGPT.

I use it as a tutor and a thinking partner. I ask questions, challenge explanations, propose interpretations, work through derivations, and often return to the same idea several times until I can express it in my own terms.

That does not remove the need to understand the material. If anything, it makes that distinction more important.

An LLM can produce a plausible explanation very quickly. Plausibility is not the standard I want here. I still need to check the algebra, inspect the assumptions, test the intuition against examples, and decide whether the explanation actually makes sense to me.

The conversations are therefore part of the source material for this series. I want to preserve the real questions and changes of mind that happened during them rather than reconstructing a cleaner learning history afterward.

There is something appropriately circular about using a language model to investigate the mathematics underneath language models. I am curious to see where that circularity is productive, and where it becomes a shortcut that hides understanding.

From equations to code

In parallel, I have been building a repository called llm-from-first-principles.

The idea is complementary to the blog.

Writing asks:

Can I explain this?

Implementation asks:

Can I build it without hiding the mechanism from myself?

The repository is where I work through the components directly. The blog is where I try to reconstruct the reasoning behind them.

Softmax is also a small example of why I do not want to stop at the equation itself.

On paper, we write

Mathematically, that definition is fine. But a literal implementation can be numerically unstable: if a logit is a large positive number, its exponential may overflow.

A stable implementation first finds

and shifts every logit by the same value:

It then computes

This produces exactly the same probabilities, because Softmax is invariant to adding or subtracting the same constant from every logit:

Here is the implementation I worked through in PyTorch:

import torch

def softmax_stable(x: torch.Tensor) -> torch.Tensor:
    max_x = x.max(
        dim=-1, keepdim=True
    ).values  # .values => get values, here [0] the same.
    shifted_x = x - max_x

    exp_x = torch.exp(shifted_x)
    row_sums = exp_x.sum(dim=-1, keepdim=True)

    return exp_x / row_sums

Subtracting the maximum does not define a different Softmax. It is a numerically safer way to compute the same mathematical function.

Understanding a mathematical formula is not yet the same thing as understanding how it lives in code.

There is often another layer between the expression written on paper and a reliable implementation: numerical stability, tensor shapes, masking, broadcasting, memory layout, gradient flow, or some other constraint that the compact mathematical notation does not show.

That gap is one of the reasons I started the repository in the first place.

Together, I want the process to look something like

I do not expect this to end with a complete theory of why large language models work. That question is much larger than any collection of formulas, and many parts of it remain open.

The goal is smaller and, for me, more useful:

To make fewer parts of the machinery feel like magic.

The first place I want to start is probability.

We train language models by assigning probabilities to tokens and sequences, but the data we begin with are observations.

So where, exactly, does the probability come from?

That is the question for the first article.