Eli Bendersky's website - Machine Learning

Rewriting pycparser with the help of an LLM

2026-02-04T19:35:00-08:00

pycparser is my most widely used open source project (with ~20M daily downloads from PyPI [1]). It's a pure-Python parser for the C programming language, producing ASTs inspired by Python's own. Until very recently, it's been using PLY: Python Lex-Yacc for the core parsing.

In this post, I'll describe how I collaborated with an LLM coding agent (Codex) to help me rewrite pycparser to use a hand-written recursive-descent parser and remove the dependency on PLY. This has been an interesting experience and the post contains lots of information and is therefore quite long; if you're just interested in the final result, check out the latest code of pycparser - the main branch already has the new implementation.

The issues with the existing parser implementation

While pycparser has been working well overall, there were a number of nagging issues that persisted over years.

Parsing strategy: YACC vs. hand-written recursive descent

I began working on pycparser in 2008, and back then using a YACC-based approach for parsing a whole language like C seemed like a no-brainer to me. Isn't this what everyone does when writing a serious parser? Besides, the K&R2 book famously carries the entire grammar of the C99 language in an appendix - so it seemed like a simple matter of translating that to PLY-yacc syntax.

And indeed, it wasn't too hard, though there definitely were some complications in building the ASTs for declarations (C's gnarliest part).

Shortly after completing pycparser, I got more and more interested in compilation and started learning about the different kinds of parsers more seriously. Over time, I grew convinced that recursive descent is the way to go - producing parsers that are easier to understand and maintain (and are often faster!).

It all ties in to the benefits of dependencies in software projects as a function of effort. Using parser generators is a heavy conceptual dependency: it's really nice when you have to churn out many parsers for small languages. But when you have to maintain a single, very complex parser, as part of a large project - the benefits quickly dissipate and you're left with a substantial dependency that you constantly grapple with.

The other issue with dependencies

And then there are the usual problems with dependencies; dependencies get abandoned, and they may also develop security issues. Sometimes, both of these become true.

Many years ago, pycparser forked and started vendoring its own version of PLY. This was part of transitioning pycparser to a dual Python 2/3 code base when PLY was slower to adapt. I believe this was the right decision, since PLY "just worked" and I didn't have to deal with active (and very tedious in the Python ecosystem, where packaging tools are replaced faster than dirty socks) dependency management.

A couple of weeks ago this issue was opened for pycparser. It turns out the some old PLY code triggers security checks used by some Linux distributions; while this code was fixed in a later commit of PLY, PLY itself was apparently abandoned and archived in late 2025. And guess what? That happened in the middle of a large rewrite of the package, so re-vendoring the pre-archiving commit seemed like a risky proposition.

On the issue it was suggested that "hopefully the dependent packages move on to a non-abandoned parser or implement their own"; I originally laughed this idea off, but then it got me thinking... which is what this post is all about.

Growing complexity of parsing a messy language

The original K&R2 grammar for C99 had - famously - a single shift-reduce conflict having to do with dangling elses belonging to the most recent if statement. And indeed, other than the famous lexer hack used to deal with C's type name / ID ambiguity, pycparser only had this single shift-reduce conflict.

But things got more complicated. Over the years, features were added that weren't strictly in the standard but were supported by all the industrial compilers. The more advanced C11 and C23 standards weren't beholden to the promises of conflict-free YACC parsing (since almost no industrial-strength compilers use YACC at this point), so all caution went out of the window.

The latest (PLY-based) release of pycparser has many reduce-reduce conflicts [2]; these are a severe maintenance hazard because it means the parsing rules essentially have to be tie-broken by order of appearance in the code. This is very brittle; pycparser has only managed to maintain its stability and quality through its comprehensive test suite. Over time, it became harder and harder to extend, because YACC parsing rules have all kinds of spooky-action-at-a-distance effects. The straw that broke the camel's back was this PR which again proposed to increase the number of reduce-reduce conflicts [3].

This - again - prompted me to think "what if I just dump YACC and switch to a hand-written recursive descent parser", and here we are.

The mental roadblock

None of the challenges described above are new; I've been pondering them for many years now, and yet biting the bullet and rewriting the parser didn't feel like something I'd like to get into. By my private estimates it'd take at least a week of deep heads-down work to port the gritty 2000 lines of YACC grammar rules to a recursive descent parser [4]. Moreover, it wouldn't be a particularly fun project either - I didn't feel like I'd learn much new and my interests have shifted away from this project. In short, the Potential well was just too deep.

Why would this even work? Tests

I've definitely noticed the improvement in capabilities of LLM coding agents in the past few months, and many reputable people online rave about using them for increasingly larger projects. That said, would an LLM agent really be able to accomplish such a complex project on its own? This isn't just a toy, it's thousands of lines of dense parsing code.

What gave me hope is the concept of conformance suites mentioned by Simon Willison. Agents seem to do well when there's a very clear and rigid goal function - such as a large, high-coverage conformance test suite.

And pycparser has an very extensive one. Over 2500 lines of test code parsing various C snippets to ASTs with expected results, grown over a decade and a half of real issues and bugs reported by users.

I figured the LLM can either succeed or fail and throw its hands up in despair, but it's quite unlikely to produce a wrong port that would still pass all the tests. So I set it to run.

The initial port

I fired up Codex in pycparser's repository, and wrote this prompt just to make sure it understands me and can run the tests:

This repository implements a Python parser producing an AST. It has a very
extensive test suite; can you first verify you successfully run all the tests?
(with python3 -m unittest discover, from the root dir)

Codex figured it out (I gave it the exact command, after all!); my next prompt was the real thing [5]:

pycparser uses the PLY library for its lexer and parser, and its own AST
package for the ASTs created by the parser. I want to explore replacing the
Parser component (leaving the Lexer alone for now) with a hand written
recursive descent parser that will pass all the same tests. Can you try that?

Here Codex went to work and churned for over an hour. Having never observed an agent work for nearly this long, I kind of assumed it went off the rails and will fail sooner or later. So I was rather surprised and skeptical when it eventually came back with:

Replaced the PLY parse path with a handwritten recursive‑descent parser while
keeping the existing lexer and AST helpers intact; CParser now points to the
new RD implementation, with the old PLY version preserved as PLYCParser in
pycparser/c_parser.py. The RD parser includes a token stream, declarator/type
parsing, expression precedence handling, pragmas, and statement parsing to
match the previous AST output and coordinate behavior.

It took me a while to poke around the code and run it until I was convinced - it had actually done it! It wrote a new recursive descent parser with only ancillary dependencies on PLY, and that parser passed the test suite. After a few more prompts, we've removed the ancillary dependencies and made the structure clearer. I hadn't looked too deeply into code quality at this point, but at least on the functional level - it succeeded. This was very impressive!

A quick note on reviews and branches

A change like the one described above is impossible to code-review as one PR in any meaningful way; so I used a different strategy. Before embarking on this path, I created a new branch and once Codex finished the initial rewrite, I committed this change, knowing that I will review it in detail, piece-by-piece later on.

Even though coding agents have their own notion of history and can "revert" certain changes, I felt much safer relying on Git. In the worst case if all of this goes south, I can nuke the branch and it's as if nothing ever happened. I was determined to only merge this branch onto main once I was fully satisfied with the code. In what follows, I had to git reset several times when I didn't like the direction in which Codex was going. In hindsight, doing this work in a branch was absolutely the right choice.

The long tail of goofs

Once I've sufficiently convinced myself that the new parser is actually working, I used Codex to similarly rewrite the lexer and get rid of the PLY dependency entirely, deleting it from the repository. Then, I started looking more deeply into code quality - reading the code created by Codex and trying to wrap my head around it.

And - oh my - this was quite the journey. Much has been written about the code produced by agents, and much of it seems to be true. Maybe it's a setting I'm missing (I'm not using my own custom AGENTS.md yet, for instance), but Codex seems to be that eager programmer that wants to get from A to B whatever the cost. Readability, minimalism and code clarity are very much secondary goals.

Using raise...except for control flow? Yep. Abusing Python's weak typing (like having None, false and other values all mean different things for a given variable)? For sure. Spreading the logic of a complex function all over the place instead of putting all the key parts in a single switch statement? You bet.

Moreover, the agent is hilariously lazy. More than once I had to convince it to do something it initially said is impossible, and even insisted again in follow-up messages. The anthropomorphization here is mildly concerning, to be honest. I could never imagine I would be writing something like the following to a computer, and yet - here we are: "Remember how we moved X to Y before? You can do it again for Z, definitely. Just try".

My process was to see how I can instruct Codex to fix things, and intervene myself (by rewriting code) as little as possible. I've mostly succeeded in this, and did maybe 20% of the work myself.

My branch grew dozens of commits, falling into roughly these categories:

The code in X is too complex; why can't we do Y instead?
The use of X is needlessly convoluted; change Y to Z, and T to V in all instances.
The code in X is unclear; please add a detailed comment - with examples - to explain what it does.

Interestingly, after doing (3), the agent was often more effective in giving the code a "fresh look" and succeeding in either (1) or (2).

The end result

Eventually, after many hours spent in this process, I was reasonably pleased with the code. It's far from perfect, of course, but taking the essential complexities into account, it's something I could see myself maintaining (with or without the help of an agent). I'm sure I'll find more ways to improve it in the future, but I have a reasonable degree of confidence that this will be doable.

It passes all the tests, so I've been able to release a new version (3.00) without major issues so far. The only issue I've discovered is that some of CFFI's tests are overly precise about the phrasing of errors reported by pycparser; this was an easy fix.

The new parser is also faster, by about 30% based on my benchmarks! This is typical of recursive descent when compared with YACC-generated parsers, in my experience. After reviewing the initial rewrite of the lexer, I've spent a while instructing Codex on how to make it faster, and it worked reasonably well.

Followup - static typing

While working on this, it became quite obvious that static typing would make the process easier. LLM coding agents really benefit from closed loops with strict guardrails (e.g. a test suite to pass), and type-annotations act as such. For example, had pycparser already been type annotated, Codex would probably not have overloaded values to multiple types (like None vs. False vs. others).

In a followup, I asked Codex to type-annotate pycparser (running checks using ty), and this was also a back-and-forth because the process exposed some issues that needed to be refactored. Time will tell, but hopefully it will make further changes in the project simpler for the agent.

Based on this experience, I'd bet that coding agents will be somewhat more effective in strongly typed languages like Go, TypeScript and especially Rust.

Conclusions

Overall, this project has been a really good experience, and I'm impressed with what modern LLM coding agents can do! While there's no reason to expect that progress in this domain will stop, even if it does - these are already very useful tools that can significantly improve programmer productivity.

Could I have done this myself, without an agent's help? Sure. But it would have taken me much longer, assuming that I could even muster the will and concentration to engage in this project. I estimate it would take me at least a week of full-time work (so 30-40 hours) spread over who knows how long to accomplish. With Codex, I put in an order of magnitude less work into this (around 4-5 hours, I'd estimate) and I'm happy with the result.

It was also fun. At least in one sense, my professional life can be described as the pursuit of focus, deep work and flow. It's not easy for me to get into this state, but when I do I'm highly productive and find it very enjoyable. Agents really help me here. When I know I need to write some code and it's hard to get started, asking an agent to write a prototype is a great catalyst for my motivation. Hence the meme at the beginning of the post.

Does code quality even matter?

One can't avoid a nagging question - does the quality of the code produced by agents even matter? Clearly, the agents themselves can understand it (if not today's agent, then at least next year's). Why worry about future maintainability if the agent can maintain it? In other words, does it make sense to just go full vibe-coding?

This is a fair question, and one I don't have an answer to. Right now, for projects I maintain and stand behind, it seems obvious to me that the code should be fully understandable and accepted by me, and the agent is just a tool helping me get to that state more efficiently. It's hard to say what the future holds here; it's going to interesting, for sure.

[1]	pycparser has a fair number of direct dependents, but the majority of downloads comes through CFFI, which itself is a major building block for much of the Python ecosystem.

[2]	The table-building report says 177, but that's certainly an over-dramatization because it's common for a single conflict to manifest in several ways.

[3]	It didn't help the PR's case that it was almost certainly vibe coded.

[4]

There was also the lexer to consider, but this seemed like a much simpler job. My impression is that in the early days of computing, lex gained prominence because of strong regexp support which wasn't very common yet. These days, with excellent regexp libraries existing for pretty much every language, the added value of lex over a custom regexp-based lexer isn't very high.

That said, it wouldn't make much sense to embark on a journey to rewrite just the lexer; the dependency on PLY would still remain, and besides, PLY's lexer and parser are designed to work well together. So it wouldn't help me much without tackling the parser beast.

[5]	I've decided to ask it to the port the parser first, leaving the lexer alone. This was to split the work into reasonable chunks. Besides, I figured that the parser is the hard job anyway - if it succeeds in that, the lexer should be easy. That assumption turned out to be correct.

Sparsely-gated Mixture Of Experts (MoE)

2025-04-18T09:33:00-07:00

In transformer models, the attention block is typically followed by a feed forward layer (FF), which is a simple fully-connected NN with a hidden layer and nonlinearity. Here's the code for such a block that uses ReLU:

def feed_forward_relu(x, W1, W2):
    """Feed-forward layer with ReLU activation.

    Args:
        x: Input tensor (B, N, D).
        Wh: Weights for the hidden layer (D, DH).
        Wo: Weights for the output layer (DH, D).

    Returns:
        Output tensor (B, N, D).
    """
    x = x @ W1  # hidden layer (B, N, DH)
    x = np.maximum(0, x)  # ReLU activation (B, N, DH)
    x = x @ W2  # output layer (B, N, D)
    return x

This layer typically holds most of the weights in the transformer, because the hidden dimension (DH in this post, hidden_dim in some papers) is large - 4x the embedding depth D is common. Intuitively, this makes sense because this layer does the majority of the heavy lifting; while the attention block mixes the embeddings of tokens together to express their relationships to one another, the FF block does the actual "reasoning" on these tokens.

Transformer blocks are repeated dozens of times in a model, so the total size of these layers becomes problematic. One approach for improving efficiency that became very popular is called sparsely-gated mixture of experts (paper).

Mixture of Experts architecture (MoE)

The basic idea of MoE is this:

The large FF layer is split into a number (NEXP) of blocks called "experts". Each expert is still a FF layer. It takes a vector of size D and transforms it to another vector of size D.
There's an additional piece called "router" or "gate". This is just a fully-connected layer (D, NEXP) that takes a token and produces a score for each expert. The router is learned by the model, along with the experts themselves.
K experts with the highest scores are selected for each token, and the token is only fed through these experts.
The scores are also used to calculate a weighted average from the experts' outputs, eventually producing an answer of size D.

Here's a diagram for a single token, assuming NEXP=8 and TOPK=2 (the two highest scoring experts are selected for each token, out of a total of eight):

Notes:

Experts #1 and #5 are selected because the router produced the highest scores for them among all the experts. The input token is routed to these experts, but not to the others.
The output of each expert is element-wise multiplied by a corresponding weight, calculated from the scores of the selected experts using a softmax function (to ensure balanced weighting across multiple tokens and experts).
The weighted expert outputs are added up for the final output of the layer.

The key point to understand about this architecture is: the experts that were not among the top K for a token aren't used at all - the computation required to propagate the token through these experts is eschewed (both on the forward and backward passes).

This is the goal of the MoE architecture - we increase the overall model size, but keep the computational cost in check by only using a portion of the parameters for every single token. This is also reflected in the models' names; for example, the Mixtral model has size 8x7B; it has 8 experts, and it would be incorrect to just multiply the size of each expert by 8 because not all these parameters participate in the calculation of every token [1]. According to the Mixtral paper, the model only uses 13B active parameters for each token.

A summary of the idea behind MoE is:

MoE increases the model's capacity without proportionally increasing its computational cost.

Numpy implementation

Here's a well-commented implementation of the MoE layer using pure Numpy. First, some parameters:

# Parameters for a feed-forward layer with a fixed activation function.
@dataclass
class FFParams:
    Wh: np.ndarray
    Wo: np.ndarray


# Parameters for a Mixture of Experts (MoE) layer.
@dataclass
class MoEParams:
    # Embedding dimension of each token (a.k.a. model dimension, Dmodel)
    D: int

    # Hidden dimension in FF layers
    DH: int

    # Total number of experts
    NEXP: int

    # K in the top-k selection of top experts per token
    TOPK: int

    # List of experts: each expert is a forward layer with FFParams.
    ff_weights: List[FFParams]

    # Router weights: a linear layer (D, NEXP) that maps input to expert scores.
    router_weights: np.ndarray

And now the implementation. Note that it takes a general (B, N, D) input, assuming batch dimension D and sequence length N:

def moe(x: np.ndarray, params: MoEParams):
    """Mixture of Experts (MoE) layer.

    Args:
        x: Input tensor (B, N, D).
        params: MoEParams.

    Returns:
        Output tensor (B, N, D).
    """
    # Run input through router to get expert scores for each token.
    expert_scores = x @ params.router_weights  # (B, N, NEXP)

    # Select the top-k expert scores and their indices for each token.
    top_scores, top_experts = topk_lastdim(expert_scores, params.TOPK)  # (B, N, TOPK)

    # Apply softmax to the top scores to get weights that sum to 1.
    weights = softmax_lastdim(top_scores)  # (B, N, TOPK)

    out = np.zeros_like(x)
    for b in range(x.shape[0]):
        for n in range(x.shape[1]):
            # Unvectorized implementation: for each token in the batch and
            # sequence, select the top-k experts and apply them with the
            # calculated weights.
            for expert_idx, weight in zip(top_experts[b, n], weights[b, n]):
                expert = params.ff_weights[expert_idx]
                out[b, n] += weight * feed_forward_relu(x[b, n], expert.Wh, expert.Wo)

    return out

Calculating the experts themselves is not vectorized here - it is done token by token. MoE is inherently sparse: different tokens in the same sequence (and batch) may go through different sets of experts. Vectorizing this efficiently is tricky in general and depends on the HW we run the model on [2]. For a popular approach on GPUs, see the MegaBlocks paper from 2022. This remains an active area of research.

All that's left are some helper functions:

def topk_lastdim(x, k):
    """Get the top k elements and their indices.

    x is an arbitrary array with at least two dimensions. The returned
    array has the same shape as x, but its elements are the top k elements
    across the last dimension. The indices of the top k elements are also
    returned.
    """
    idx = np.argpartition(x, -k, axis=-1)[..., -k:]
    return np.take_along_axis(x, idx, axis=-1), idx

def softmax_lastdim(x):
    """Compute softmax across last dimension of x.

    x is an arbitrary array with at least two dimensions. The returned array has
    the same shape as x, but its elements sum up to 1 across the last dimension.
    """
    # Subtract the max for numerical stability
    ex = np.exp(x - np.max(x, axis=-1, keepdims=True))
    # Divide by sums across last dimension
    return ex / np.sum(ex, axis=-1, keepdims=True)

Additional considerations

A major area of focus with MoE architectures is load balancing among experts. Without special provisions, the model may learn to prefer certain experts over others and this leads to inefficient utilization of the model's weights. There are various approaches to tackle this, for example:

Adding noise to the top-k selection process to inject randomness
Defining a special loss function during training that encourages experts to receive a roughly equal number of training samples

Code

The full code for this post is available on GitHub.

[1]

Another way to think about MoE is that each "expert" specializes in a certain area of the model's capability. For example, one expert would be good at math, another at prose, etc. This is a very rough approximation, though, because transformer models consist of dozens of repeating blocks, and all these different experts end up thoroughly intermixed as tokens flow through the entire model.

[2]

In the sparsely-gated mixture of experts paper, this is referred to as The Shrinking Batch Problem:

"In modern CPUs and GPUs, large batch sizes are necessary for computational efficiency, so as to amortize the overhead of parameter loads and updates. If the gating network chooses k out of n experts for each example, then for a batch of b examples, each expert receives a much smaller batch of approximately kb/n << b examples. This causes a naive MoE implementation to become very inefficient as the number of experts increases"

Cross-entropy and KL divergence

2025-04-12T06:54:00-07:00

Cross-entropy is widely used in modern ML to compute the loss for classification tasks. This post is a brief overview of the math behind it and a related concept called Kullback-Leibler (KL) divergence.

Information content of a single random event

We'll start with a single event (E) that has probability p. The information content (or "degree of surprise") of this event occurring is defined as:

\[I(E) = \log_2 \left (\frac{1}{p} \right )\]

The base 2 here is used so that we can count the information in units of bits. Thinking about this definition intuitively, imagine an event with probability p=1; using the formula, the information we gain by observing this event occurring is 0, which makes sense. On the other extreme, as the probability p approaches 0, the information we gain is huge. An equivalent way to write the formula is:

\[I(E) = -\log_2 p\]

Some numeric examples: suppose we flip a fair coin and it comes out heads. The probability of this event happening is 1/2, therefore:

\[I(E_{heads})=-\log_2 \frac{1}{2} = 1\]

Now suppose we roll a fair die and it lands on 4. The probability of this event happening is 1/6, therefore:

\[I(E_4)=-\log_2 \frac{1}{6} = 2.58\]

In other words, the degree of surprise for rolling a 4 is higher than the degree of surprise for flipping to heads - which makes sense, given the probabilities involved.

Other than behaving correctly for boundary values, the logarithm function makes sense for calculating the degree of surprise for another important reason: the way it behaves for a combination of events.

Consider this: we flip a fair coin and roll a fair die; the coin comes out heads, and the die lands on 4. What is the probability of this event happening? Because the two events are independent, the probability is the product of the probabilities of the individual events, so 1/12, and then:

\[I(E_{heads}\cap E_{4})=-\log_2 \frac{1}{12} = 3.58\]

Note that the entropy is the precise sum of the entropies of individual events. This is to be expected - we need so many bits for one of the events, and so many for the other; the total of the bits adds up. The logarithm function gives us exactly this behavior for probabilities:

\[\log(p_1 \cap p_2) = \log(p_1 \cdot p_2) = \log(p_1) + \log(p_2)\]

Entropy

Given a random variable X with values x_1\dots x_n and associated probabilities p_1\dots p_n, the entropy of X is defined as the expected value of information for X:

\[H(X)=-\sum_{j=1}^{n}p_j \log_2 p_j\]

High entropy means high uncertainty, while low entropy means low uncertainty. Let's look at a couple of examples:

This is a random variable with 5 distinct values; the probability of x_1 is 1, and the rest is 0. The entropy here is 0, because 1\cdot \log 1 = 0 and also 0\cdot \log 0 = 0 [1]. We gain no information by observing an event sampled from this distribution, because we knew ahead of time what would happen.

Another example is a uniform distribution for the 5 possible outcomes:

The entropy for this distribution is:

\[H(X)=-\sum_{j=1}^{5}0.2 \log_2 0.2 = 2.32\]

Intuitively: we have 5 different values with equal probabilities, so we'll need \log_{2} 5=2.32 bits to represent that. Note that entropy is always non-negative, because 0\leq p_j \leq 1 and therefore \log_2 p_j \leq 0 for all j in a proper probability distribution.

It's not hard to show that the maximum possible entropy for a random variable occurs for a uniform distribution. In all other distributions, some values are more represented than others which makes the result somewhat less surprising.

Cross-entropy

Cross-entropy is an extension of the concept of entropy, when two different probability distributions are present. The typical formulation useful for machine learning is:

\[H(P,Q)=-\sum_{j=1}^{n}p_j \log_2 q_j\]

Where:

P is the actual observed data distribution
Q is the predicted data distribution

Similarly to entropy, cross-entropy is non-negative; in fact, it collapses to the entropy formula when P and Q are the same:

\[H(P,P)=-\sum_{j=1}^{n}p_j \log_2 p_j=H(P)\]

An information-theoretic interpretation of cross-entropy is: the average number of bits required to encode an actual probability distribution P, when we assumed the data follows Q instead.

Here's a numeric example:

p = [0.1, 0.2, 0.4, 0.2, 0.1]
q = [0.2, 0.2, 0.2, 0.2, 0.2]

Plotted:

The cross-entropy of these two distributions is 2.32

Now let's try a Q that's slightly closer to P:

p = [0.1, 0.2, 0.4, 0.2, 0.1]
q = [0.15, 0.175, 0.35, 0.175, 0.15]

The cross-entropy in these distributions is somewhat lower, 2.16; this is expected, because they're more similar. In other words, the outcome of measuring P when our model predicted Q is less surprising.

KL divergence

Cross-entropy is useful for tracking the training loss of a model (more on this in the next section), but it has some mathematical properties that make it less than ideal as a statistical tool to compare two probability distributions. Specifically, H(P,P)=H(P), which isn't (usually) zero; this is the lowest value possible for cross-entropy. In other words, cross-entropy always retains the inherent uncertainty of P.

The KL divergence fixes this by subtracting H(P) from cross-entropy:

\[D_{KL}(P,Q)=H(P,Q)-H(P)=-\left (\sum_{j=1}^{n}p_j \log_2 q_j - \sum_{j=1}^{n}p_j \log_2 p_j \right )\]

Manipulating the logarithms, we can also get these alternative formulations:

\[D_{KL}(P,Q)=-\sum_{j=1}^{n}p_j \log_2 \frac{q_j}{p_j}=\sum_{j=1}^{n}p_j \log_2 \frac{p_j}{q_j}\]

Thus, the KL divergence is more useful as a measure of divergence between two probability distributions, since D_{KL}(P,P)=0. Note, however, that it's not a true distance metric because it's not symmetric:

\[D_{KL}(P,Q)\ne D_{KL}(Q,P)\]

Uses in machine learning

In ML, we often have a model that makes a prediction and a set of training data which defines a real-world probability distribution. It's natural to define a loss function in terms of the difference between the two distributions (the model's prediction and the real data).

Cross-entropy is very useful as a loss function because it's non-negative and provides a single scalar number that's lower for similar distributions and higher for dissimilar distributions. Moreover, if we think of cross-entropy in terms of KL divergence:

\[H(P,Q)=D_{KL}(P,Q)+H(P)\]

We'll notice that H(P) - the entropy of the real-world distribution - does not depend on the model at all. Therefore, optimizing cross-entropy is equivalent to optimizing the KL divergence. I wrote about concrete uses of cross-entropy as a loss function in previous posts:

That said, the KL divergence is also sometimes useful more directly; for example in the evidence lower bound used for Variational autoencoders.

Relation to Maximum Likelihood Estimation

There's an interesting relation between the concepts discussed in this post and Maximum Likelihood Estimation.

Suppose we have a true probability distribution P, and a parameterized model that predicts the probability distribution Q_\theta. stands for all the parameters of our model (e.g. all the weights of a deep learning network).

The likelihood of observing a set of samples x_1\cdots x_n drawn from P is:

\[L=\prod ^{n}_{i=1}P(x_i)\]

However, we don't really know P; what we do know is Q_\theta, so we can calculate:

\[L(\theta)=\prod ^{n}_{i=1}Q_\theta(x_i)\]

The idea is to find an optimal set of parameters \widehat{\theta} such that this likelihood is maximized; in other words:

\[\widehat{\theta}=\underset{\theta}{argmax}\ L(\theta)=\underset{\theta}{argmax}\ \prod ^{n}_{i=1}Q_\theta(x_i)\]

Working with products is inconvenient, however, so a logarithm is used instead to convert a product to a sum (since log(f(x)) is a monotonically increasing function, maximizing it is akin to maximizing itself):

\[\widehat{\theta}=\underset{\theta}{argmax}\ \log L(\theta)=\underset{\theta}{argmax}\ \sum ^{n}_{i=1}\log Q_\theta(x_i)\]

This is the maximal log-likelihood.

Now a clever statistical trick is employed; first, we multiply the function we're maximizing by the constant \frac{1}{n} - this doesn't affect the maxima, of course:

\[\widehat{\theta}=\underset{\theta}{argmax}\ \frac{1}{n}\sum ^{n}_{i=1}\log Q_\theta(x_i)\]

The function inside the argmax is now the average across n samples obtained from the true probability distribution P. The Law of Large numbers states that with a large enough n, this average converges to the expected value of drawing from this distribution:

\[\widehat{\theta}=\underset{\theta}{argmax}\ \sum ^{n}_{i=1}P(x_i)\log Q_\theta(x_i)\]

This should start looking familiar; all that's left is to negate the sum and minimize the negative instead:

\[\widehat{\theta}=\underset{\theta}{argmin}\ -\sum ^{n}_{i=1}P(x_i)\log Q_\theta(x_i)\]

The function we're now minimizing is the cross-entropy between P and Q_\theta. We've shown that maximum likelihood estimation is equivalent to minimizing the cross-entropy between the true and and predicted data distributions.

[1]	This can be proven by taking the limit \lim_{p\to 0} p \log p and using L'Hopital's rule to show that it goes to 0.

Reproducing word2vec with JAX

2025-04-05T06:19:00-07:00

The word2vec model was proposed in a 2013 paper by Google researchers called "Efficient Estimation of Word Representations in Vector Space", and was further refined by additional papers from the same team. It kick-started the modern use of embeddings - dense vector representation of words (and later tokens) for language models.

Also, the code - with some instructions - was made available openly. This post reproduces the word2vec results using JAX, and also talks about reproducing it using the original C code (see the Original word2vec code section for that).

Embeddings

First, a brief introduction to embeddings. Wikipedia has a good definition:

In natural language processing, a word embedding is a representation of a word. The embedding is used in text analysis. Typically, the representation is a real-valued vector that encodes the meaning of the word in such a way that the words that are closer in the vector space are expected to be similar in meaning

Here's a framework that made sense to me when I was first learning about embeddings many years ago:

ML models and NNs specifically are all about vector math.
Words in a human language (like English) are just sequences of characters with no semantic meaning (there's nothing in the word "dog" that conveys dog-ness any more than the same concept in other human languages). Also, words have different lengths which isn't convenient.
To represent words as vectors, we typically use indices into a vocabulary; equivalently, this can be seen as a one-hot vector with the value at the correct vocabulary index being 1, and the rest 0.
This latter vector representation has no semantic meaning either, because "Paris" and "France" will be as different from each other as "Paris" and "Armadillo". Also, these vectors are huge (a typical vocabulary can have tens of thousands of words, just for a single language!)
Therefore, we need some magic to convert words into vectors that carry meaning.

Embeddings are that magic. They are dense vectors of floats - with typically hundreds or thousands of elements, and serve as representations of these words in high-dimensional space.

The word2vec CBOW architecture

The word2vec paper proposed two related architectures: CBOW (Continuous Bag Of Words) and Continuous Skip Gram. The two are fairly similar, and in this post I'm going to focus on CBOW.

The idea of the CBOW approach is to teach the model to predict a word from its surrounding words. Here's an example with window size of four [1]:

The goal here is to have the model predict that "liberty" should be the word in the middle, given the context words in peach-colored boxes. This is an unsupervised model - it learns by consuming text, sliding its window word by word over arbitrary amounts of (properly formatted and sanitized) input.

Concretely, the following diagram shows the model architecture; here are the dimensions involved:

B: batch (for computational efficiency, whole batches are processed together)
V: vocabulary size (the number of unique words in our vocabulary)
D: model depth (the size of the dense embedding vectors we're trying to learn)
W: window size

Here's the flow of data in the forward pass:

context is the context words for a given position. For example, in the sample diagram above the context would be of length 8. Each element is an integer representation of a word (its index into the vocabulary). Since we're processing batches, the shape of this array is (B,2W).
The context indexes into a projection matrix P, which has the learned embedding per row - one for each word in the vocabulary. The result is projection with shape (B,2W,D). The first two dimensions remain the same (because we still have the same batch and window size), but every integer is replaced with the word's embedding - so an extra dimension is added.
Next, a mean (arithmetic average) is taken across the window dimension. The embeddings of all the words in the window are averaged together. The result is (B,D) where each row is the average of the embeddings of 2W words.
Finally, the hidden layer matrix H is used to map the dense representation back into a sparse one [2] - this is the prediction of the middle word. Recall that this tries to predict a one-hot encoding of the word's vocabulary index.

For training, the loss is calculated by comparing out to the one-hot encoding of the actual target word for this window, and the calculated gradient is propagated backwards to train the model.

JAX implementation

The JAX implementation of the model described above is clean and compact:

@jax.jit
def word2vec_forward(params, context):
    """Forward pass of the word2Vec model.

    context is a (batch_size, 2*window_size) array of word IDs.

    V is the vocabulary size, D is the embedding dimension.
    params["projection"] is a (V, D) matrix of word embeddings.
    params["hidden"] is a (D, V) matrix of weights for the hidden layer.
    """
    # Indexing into (V, D) matrix with a batch of IDs. The output shape
    # is (batch_size, 2*window_size, D).
    projection = params["projection"][context]

    # Compute average across the context word. The output shape is
    # (batch_size, D).
    avg_projection = jnp.mean(projection, axis=1)

    # (batch_size, D) @ (D, V) -> (batch_size, V)
    hidden = jnp.dot(avg_projection, params["hidden"])
    return hidden


@jax.jit
def word2vec_loss(params, target, context):
    """Compute the loss of the word2Vec model."""
    logits = word2vec_forward(params, context)  # (batch_size, V)

    target_onehot = jax.nn.one_hot(target, logits.shape[1])  # (batch_size, V)
    loss = optax.losses.softmax_cross_entropy(logits, target_onehot).mean()
    return loss

Training

For training, I've been relying on the same dataset used by the original word2vec code - a 100MB text file downloaded from http://mattmahoney.net/dc/text8.zip

This file contains all-lowercase text with no punctuation, so it requires very little cleaning and processing. What it does require for higher-quality training is subsampling: throwing away some of the most common words (e.g. "and", "is", "not" in English), since they appear so much in the text. Here's my code for this:

def subsample(words, threshold=1e-4):
    """Subsample frequent words, return a new list of words.

    Follows the subsampling procedure described in the paper "Distributed
    Representations of Words and Phrases and their Compositionality" by
    Mikolov et al. (2013).
    """
    word_counts = Counter(words)
    total_count = len(words)
    freqs = {word: count / total_count for word, count in word_counts.items()}

    # Common words (freq(word) > threshold) are kept with a computed
    # probability, while rare words are always kept.
    p_keep = {
        word: math.sqrt(threshold / freqs[word]) if freqs[word] > threshold else 1
        for word in word_counts
    }
    return [word for word in words if random.random() < p_keep[word]]

We also have to create a vocabulary with some limited size:

def make_vocabulary(words, top_k=20000):
    """Creates a vocabulary from a list of words.

    Keeps the top_k most common words and assigns an index to each word. The
    index 0 is reserved for the "<unk>" token.
    """
    word_counts = Counter(words)
    vocab = {"<unk>": 0}
    for word, _ in word_counts.most_common(top_k - 1):
        vocab[word] = len(vocab)
    return vocab

The preprocessing step generates the list of subsampled words and the vocabulary, and stores them in a pickle file for future reference. The training loop uses these data to train a model from a random initialization. Pay special attention to the hyper-parameters defined at the top of the train function. I set these to be as close as possible to the original word2vec code:

def train(train_data, vocab):
    V = len(vocab)
    D = 200
    LEARNING_RATE = 1e-3
    WINDOW_SIZE = 8
    BATCH_SIZE = 1024
    EPOCHS = 25

    initializer = jax.nn.initializers.glorot_uniform()
    params = {
        "projection": initializer(jax.random.PRNGKey(501337), (V, D)),
        "hidden": initializer(jax.random.PRNGKey(501337), (D, V)),
    }

    optimizer = optax.adam(LEARNING_RATE)
    opt_state = optimizer.init(params)

    print("Approximate number of batches:", len(train_data) // BATCH_SIZE)

    for epoch in range(EPOCHS):
        print(f"=== Epoch {epoch + 1}")
        epoch_loss = []
        for n, (target_batch, context_batch) in enumerate(
            generate_train_vectors(
                train_data, vocab, window_size=WINDOW_SIZE, batch_size=BATCH_SIZE
            )
        ):
            # Shuffle the batch.
            indices = np.random.permutation(len(target_batch))
            target_batch = target_batch[indices]
            context_batch = context_batch[indices]

            # Compute the loss and gradients; optimize.
            loss, grads = jax.value_and_grad(word2vec_loss)(
                params, target_batch, context_batch
            )
            updates, opt_state = optimizer.update(grads, opt_state)
            params = optax.apply_updates(params, updates)

            epoch_loss.append(loss)
            if n > 0 and n % 1000 == 0:
                print(f"Batch {n}")

        print(f"Epoch loss: {np.mean(epoch_loss):.2f}")
        checkpoint_filename = f"checkpoint-{epoch:03}.pickle"
        print("Saving checkpoint to", checkpoint_filename)
        with open(checkpoint_filename, "wb") as file:
            pickle.dump(params, file)

The only thing I'm not showing here is the generate_train_vectors function, as it's not particularly interesting; you can find it in the full code.

I don't have a particularly powerful GPU, so on my machine training this model for 25 epochs takes 20-30 minutes.

Extracting embeddings and finding word similarities

The result of the training is the P and H arrays with trained weights; P is exactly the embedding matrix we need! It maps vocabulary words to their dense embedding representation. Using P, we can create the fun word demos that made word2vec famous. The full code has a script named similar-words.py that does this. Some examples:

$ uv run similar-words.py -word paris \
      -checkpoint checkpoint.pickle \
      -traindata train-data.pickle
Words similar to 'paris':
paris           1.00
france          0.50
french          0.49
la              0.42
le              0.41
henri           0.40
toulouse        0.38
brussels        0.38
petit           0.38
les             0.38

And:

$ uv run similar-words.py -analogy berlin,germany,tokyo \
      -checkpoint checkpoint.pickle \
      -traindata train-data.pickle
Analogies for 'berlin is to germany as tokyo is to ?':
tokyo           0.70
japan           0.45
japanese        0.44
osaka           0.40
china           0.36
germany         0.35
singapore       0.32
han             0.31
gu              0.31
kyushu          0.31

This brings us to the intuition for how word2vec works: the basic idea is that semantically similar words will appear in the vicinity of roughly similar context words, but also that words are generally related to words in the context their appear in. This lets the model learn that some words are more related than others; for example:

$ uv run similar-words.py -sims soccer,basketball,chess,cat,bomb \
      -checkpoint checkpoint.pickle \
      -traindata train-data.pickle
Similarities for 'soccer' with context words ['basketball', 'chess', 'cat', 'bomb']:
basketball      0.40
chess           0.22
cat             0.14
bomb            0.13

Optimizations

The word2vec model can be optimized in several ways, many of which are focused on avoiding the giant matrix multiplication by H at the very end. The word2vec authors have a followup paper called "Distributed Representations of Words and Phrases and their Compositionality" where these are described; I'm leaving them out of my implementation, for simplicity.

Implementing these optimizations could help us improve the model's quality considerably, by increasing the model depth (it's currently 200, which is very low by modern LLM standards) and the amount of data we train on. That said, these days word2vec is mostly of historical interest anyway; the Modern text embeddings section will have more to say on how embeddings are trained as part of modern LLMs.

Original word2vec code

As mentioned above, the original website for the word2vec model is available on an archived version of Google Code. That page is still useful reading, but the Subversion instructions to obtain the actual code no longer work.

I was able to find a GitHub mirror with a code export here: https://github.com/tmikolov/word2vec (the username certainly checks out, though it's hard to know for sure!)

The awesome thing is that this code still builds and runs perfectly, many years later. Hurray to self-contained C programs with no dependencies; all I needed was to run make, and then use the included shell scripts to download the data and run training. This code uses the CPU for training; it takes a while, but I was able to reproduce the similarity / analogy results fairly easily.

Modern text embeddings

The word2vec model trains an embedding matrix; this pre-trained matrix can then be used as part of other ML models. This approach was used for a while, but it's no longer popular.

These days, an embedding matrix is trained as part of a larger model. For example, GPT-type transformer-based LLMs have an embedding matrix as the first layer in the model. This is basically just the P matrix from the diagram above [3]. LLMs learn both the embeddings and their specific task (generating tokens from a given context) at the same time. This makes some sense because:

LLMs process enormous amounts of data, and consuming this data multiple times to train embeddings separately is wasteful.
Embeddings trained together with the LLM are inherently tuned to the LLM's specific task and hyper-parameters (i.e. the kind of tokenizer used, the model depth etc.)

Specifically, modern embedding matrices differ from word2vec in two important aspects:

Instead of being word embeddings, they are token embeddings. I wrote much more on tokens for LLMs here.
The model depth (D) is much larger; GPT-3 has D=12288, and in newer models it's probably even larger. Deep embedding vectors help the models capture more nuance and semantic meaning about tokens. Naturally, they also require much more data to be trained effectively.

Full code

The full code for this post is available here. If you want to reproduce the my word2vec results, check out the README file - it contains full instructions on which scripts to run and in which order.

[1]	The window size is how many words to the left and right of the target word to take into account, and it's a configurable hyper-parameter during training.

[2]

The terms dense and sparse are used in the post in the following sense:

A sparse array is one where almost all entries are 0. This is true for one-hot vectors representing vocabulary words (all entries are 0 except a single one that has the value 1).

A dense array is filled with arbitrary floating-point values. An embedding vector is dense in this sense - it's typically short compared to the sparse vector (in the word2vec example used in this post D=200, while V=20000), but full of data (hence "dense"). An embedding matrix is dense since it consists of dense vectors (one per word index).

[3]	The rest (mean calculation, hidden layer) isn't needed since it's only there to train the word2vec CBOW model.

Notes on implementing Attention

2025-03-26T17:15:00-07:00

Some notes on implementing attention blocks in pure Python + Numpy. The focus here is on the exact implementation in code, explaining all the shapes throughout the process. The motivation for why attention works is not covered here too deeply - there are plenty of excellent online resources explaining it.

Several papers are mentioned throughout the code; they are:

AIAYN - Attention Is All You Need by Vaswani et al.
GPT-3 - Language Models are Few-Shot Learners by Brown et al.

Basic scaled self-attention

We'll start with the most basic scaled dot product self-attention, working on a single sequence of tokens, without masking.

The input is a 2D array of shape (N, D). N is the length of the sequence (how many tokens it contains) and D is the embedding depth - the length of the embedding vector representing each token [1]. D could be something like 512, or more, depending on the model.

A self-attention module is parameterized with three weight matrices, Wk, Wq and Wv. Some variants also have accompanying bias vectors, but the AIAYN paper doesn't use them, so I'll skip them here. In the general case, the shape of each weight matrix is (D, HS), where HS is some fraction of D. HS stands for "head size" and we'll see what this means soon. This is a diagram of a self-attention module (the diagram assumes N=6, D is some large number and so is HS). In the diagram, @ stands for matrix multiplication (Python/Numpy syntax):

Here's a basic Numpy implementation of this:

# self_attention the way it happens in the Transformer model. No bias.
# D = model dimension/depth (length of embedding)
# N = input sequence length
# HS = head size
#
# x is the input (N, D), each token in a row.
# Each of W* is a weight matrix of shape (D, HS)
# The result is (N, HS)
def self_attention(x, Wk, Wq, Wv):
    # Each of these is (N, D) @ (D, HS) = (N, HS)
    q = x @ Wq
    k = x @ Wk
    v = x @ Wv

    # kq: (N, N) matrix of dot products between each pair of q and k vectors.
    # The division by sqrt(HS) is the scaling.
    kq = q @ k.T / np.sqrt(k.shape[1])

    # att: (N, N) attention matrix. The rows become the weights that sum
    # to 1 for each output vector.
    att = softmax_lastdim(kq)
    return att @ v  # (N, HS)

The "scaled" part is just dividing kq by the square root of HS, which is done to keep the values of the dot products manageable (otherwise they would grow with the size of the contracted dimension).

The only dependency is a function for calculating Softmax across the last dimension of an input array:

def softmax_lastdim(x):
    """Compute softmax across last dimension of x.

    x is an arbitrary array with at least two dimensions. The returned array has
    the same shape as x, but its elements sum up to 1 across the last dimension.
    """
    # Subtract the max for numerical stability
    ex = np.exp(x - np.max(x, axis=-1, keepdims=True))
    # Divide by sums across last dimension
    return ex / np.sum(ex, axis=-1, keepdims=True)

When the input is 2D, the "last dimension" is the columns. Colloquially, this Softmax function acts on each row of x separately; it applies the Softmax formula to the elements (columns) of the row, ending up with a row of numbers in the range [0,1] that all sum up to 1.

Another note on the dimensions: it's possible for the Wv matrix to have a different second dimension from Wq and Wk. If you look at the diagram, you can see this will work out, since the softmax produces (N, N), and whatever the second dimension of V is, will be the second dimension of the output. The AIAYN paper designates these dimensions as d_k and d_v, but in practice d_k=d_v in all the variants it lists. I found that these dimensions are typically the same in other papers as well. Therefore, for simplicity I just made them all equal to D in this post; if desired, a variant with different d_k and d_v is a fairly trivial modification to this code.

Batched self-attention

In the real world, the input array is unlikely to be 2D because models are trained on batches of input sequences. To leverage the parallelism of modern hardware, whole batches are typically processed in the same operation.

The batched version of scaled self-attention is very similar to the non-batched one, due to the magic of Numpy matrix multiplication and broadcasts. Now the input shape is (B, N, D), where B is the batch dimension. The W* matrices are still (D, HS); multiplying a (B, N, D) array by (D, HS) performs contraction between the last axis of the first array and the first axis of the second array, resulting in (B, N, HS). Here's the code, with the dimensions annotated for each operation:

# self_attention with inputs that have a batch dimension.
# x has shape (B, N, D)
# Each of W* has shape (D, D)
def self_attention_batched(x, Wk, Wq, Wv):
    q = x @ Wq  # (B, N, HS)
    k = x @ Wk  # (B, N, HS)
    v = x @ Wv  # (B, N, HS)

    kq = q @ k.swapaxes(-2, -1) / np.sqrt(k.shape[-1])  # (B, N, N)

    att = softmax_lastdim(kq)  # (B, N, N)
    return att @ v  # (B, N, HS)

Note that the only difference between this and the non-batched version is the line calculating kq:

Since k is no longer 2D, the notion of "transpose" is ambiguous so we explicitly ask to swap the last and the penultimate axis, leaving the first axis (B) intact.
When calculating the scaling factor we use k.shape[-1] to select the last dimension of k, instead of k.shape[1] which only selects the last dimension for 2D arrays.

In fact, this function could also calculate the non-batched version! From now on, we'll assume that all inputs are batched, and all operations are implicitly batched. I'm not going to be using the "batched" prefix or suffix on functions any more.

The basic underlying idea of the attention module is to shift around the multi-dimensional representations of tokens in the sequence towards a better representation of the entire sequence. The tokens attend to each other. Specifically, the matrix produced by the Softmax operation is called the attention matrix. It's (N, N); for each token it specifies how much information from every other token in the sequence should be taken into account. For example, a higher number in cell (R, C) means that there's a stronger relation of token at index R in the sequence to the token at index C.

Here's a nice example from the AIAYN paper, showing a word sequence and the weights produced by two attention heads (purple and brown) for a given position in the input sequence:

This shows how the model is learning to resolve what the word "its" refers to in the sentence. Let's take just the purple head as an example. The index of token "its" in the sequence is 8, and the index of "Law" is 1. In the attention matrix for this head, the value at index (8, 1) will be very high (close to 1), with other values in the same row much lower.

While this intuitive explanation isn't critical to understand how attention is implemented, it will become more important when we talk about masked self-attention later on.

Multi-head attention

The attention mechanism we've seen so far has a single set of K, Q and V matrices. This is called one "head" of attention. In today's models, there are typically multiple heads. Each head does its attention job separately, and in the end all these results are concatenated and feed through a linear layer.

In what follows, NH is the number of heads and HS is the head size. Typically, NH times HS would be D; for example, the AIAYN paper mentions several configurations for D=512: NH=8 and HS=64, NH=32 and HS=16, and so on [2]. However, the math works out even if this isn't the case, because the final linear ("projection") layer maps the output back to (N, D).

Assuming the previous diagram showing a self-attention module is a single head with input (N, D) and output (N, HS), this is how multiple heads are combined:

Each of the (NH) heads has its own parameter weights for Q, K and V. Each attention head outputs a (N, HS) matrix; these are concatenated along the last dimension to (N, NH * HS), which is passed through a final linear projection.

Here's a function implementing (batched) multi-head attention; for now, please ignore the code inside do_mask conditions:

# x has shape (B, N, D)
# In what follows:
#   NH = number of heads
#   HS = head size
# Each W*s is a list of NH weight matrices of shape (D, HS).
# Wp is a weight matrix for the final linear projection, of shape (NH * HS, D)
# The result is (B, N, D)
# If do_mask is True, each attention head is masked from attending to future
# tokens.
def multihead_attention_list(x, Wqs, Wks, Wvs, Wp, do_mask=False):
    # Check shapes.
    NH = len(Wks)
    HS = Wks[0].shape[1]
    assert len(Wks) == len(Wqs) == len(Wvs)
    for W in Wqs + Wks + Wvs:
        assert W.shape[1] == HS
    assert Wp.shape[0] == NH * HS

    # List of head outputs
    head_outs = []

    if do_mask:
        # mask is a lower-triangular (N, N) matrix, with zeros above
        # the diagonal and ones on the diagonal and below.
        N = x.shape[1]
        mask = np.tril(np.ones((N, N)))

    for Wk, Wq, Wv in zip(Wks, Wqs, Wvs):
        # Calculate self attention for each head separately
        q = x @ Wq  # (B, N, HS)
        k = x @ Wk  # (B, N, HS)
        v = x @ Wv  # (B, N, HS)

        kq = q @ k.swapaxes(-2, -1) / np.sqrt(k.shape[-1])  # (B, N, N)

        if do_mask:
            # Set the masked positions to -inf, to ensure that a token isn't
            # affected by tokens that come after it in the softmax.
            kq = np.where(mask == 0, -np.inf, kq)

        att = softmax_lastdim(kq)  # (B, N, N)
        head_outs.append(att @ v)  # (B, N, HS)

    # Concatenate the head outputs and apply the final linear projection
    all_heads = np.concatenate(head_outs, axis=-1)  # (B, N, NH * HS)
    return all_heads @ Wp  # (B, N, D)

It is possible to vectorize this code even further; you'll sometimes see the heads laid out in a separate (4th) dimension instead of being a list. See the Vectorizing across the heads dimension section.

Masked (or Causal) self-attention

Attention modules can be used in both encoder and decoder blocks. Encoder blocks are useful for things like language understanding or translation; for these, it makes sense for each token to attend to all the other tokens in the sequence.

However, for generative models this presents a problem: if during training a word attends to future words, the model will just "cheat" and not really learn how to generate the next word from only past words. This is done in a decoder block, and for this we need to add masking to attention.

Conceptually, masking is very simple. Consider the sentence:

People like watching funny cat videos

When our attention code generates the att matrix, it's a square (N, N) matrix with attention weights from each token to each other token in the sequence:

What we want is for all the gray cells in this matrix to be zero, to ensure that a token doesn't attend to future tokens. The blue cells in the matrix add up to 1 in each row, after the softmax operation.

Now take a look at the previous code sample and see what happens when do_mask=True:

First, a (N, N) lower-triangular array is prepared with zeros above the diagonal and ones on the diagonal and below.
Then, before we pass the scaled QK^T to softmax, we set its values to -\infty wherever the mask matrix is 0. This ensures that the softmax function will assign zeros to outputs at these indices, while still producing the proper values in the rest of the row.

Another name for masked self-attention is causal self-attention. This is a very good name that comes from causal systems in control theory.

Intuition - what attention does

What does the attention block try to accomplish? To think about it intuitively, let's focus on a single token in the input (ignoring batch) - x[i]. For this token, the attention block produces an output token out[i] that blends x[i]'s embedding (multi-dimensional dense vector representation) with contextual information from all the tokens preceding it in the sequence, i.e. x[:i].

The way this is done is first calculating the query vector q for x[i] (using Wq). This query can be thought of as "what attributes does this token care about in its context tokens".

Then, for each of the context tokens (including x[i] itself) we calculate:

Key (using Wk): these are the attributes of the token that queries may refer to.
Value (using Wv): these are the associated values tokens carry.

When attention calculates q @ K.T for each token, the result is - for each context token - the weights to use for mixing in the token's value. Then, when this is multiplied by V, the values are properly weighted.

So this is a very general approach for the model to learn what kind of information each token "cares" about in its context tokens, and how to blend the token's embedding with those of the preceding context tokens, to properly encode the context the token is encountered in.

Our implementation, starting with the basic scaled self-attention, implements this for all tokens in the input sequence simultaneously; hence, we don't just take a single x[i], calculate its q and then multiply that by K.T. Rather, we calculate Q from all x, and continue using matrix multiplications to vectorize these calculations across the entire sequence.

It's important to keep in mind that this intuitive explanation suffers from anthropomorphism. We try to explain what the model does intuitively, but in reality this is only a very abstract approximation of what's happening (consider that attention has multiple heads, and also that that LLMs typically have dozens of repeating transformer layers with self-attention blocks, applying the same mechanism over and over again).

Cross-attention

So far we've been working with self-attention blocks, where the self suggests that elements in the input sequence attend to other elements in the same input sequence.

Another variant of attention is cross-attention, where elements of one sequence attend to elements in another sequence. This variant exists in the decoder block of the AIAYN paper. This is a single head of cross-attention:

Here we have two sequences with potentially different lengths: xq and xv. xq is used for the query part of attention, while xv is used for the key and value parts. The rest of the dimensions remain as before. The output of such a block is shaped (Nq, HS).

This is an implementation of multi-head cross-attention; it doesn't include masking, since masking is not typically necessary in cross attention - it's OK for elements of xq to attend to all elements of xv [3]:

# Cross attention between two input sequences that can have different lengths.
# xq has shape (B, Nq, D)
# xv has shape (B, Nv, D)
# In what follows:
#   NH = number of heads
#   HS = head size
# Each W*s is a list of NH weight matrices of shape (D, HS).
# Wp is a weight matrix for the final linear projection, of shape (NH * HS, D)
# The result is (B, Nq, D)
def multihead_cross_attention_list(xq, xv, Wqs, Wks, Wvs, Wp):
    # Check shapes.
    NH = len(Wks)
    HS = Wks[0].shape[1]
    assert len(Wks) == len(Wqs) == len(Wvs)
    for W in Wqs + Wks + Wvs:
        assert W.shape[1] == HS
    assert Wp.shape[0] == NH * HS

    # List of head outputs
    head_outs = []

    for Wk, Wq, Wv in zip(Wks, Wqs, Wvs):
        q = xq @ Wq  # (B, Nq, HS)
        k = xv @ Wk  # (B, Nv, HS)
        v = xv @ Wv  # (B, Nv, HS)

        kq = q @ k.swapaxes(-2, -1) / np.sqrt(k.shape[-1])  # (B, Nq, Nv)

        att = softmax_lastdim(kq)  # (B, Nq, Nv)
        head_outs.append(att @ v)  # (B, Nq, HS)

    # Concatenate the head outputs and apply the final linear projection
    all_heads = np.concatenate(head_outs, axis=-1)  # (B, Nq, NH * HS)
    return all_heads @ Wp  # (B, Nq, D)

Vectorizing across the heads dimension

The multihead_attention_list implementation shown above uses lists of weight matrices as input. While this makes the code clearer, it's not a particularly friendly format for an optimized implementation - especially on accelerators like GPUs and TPUs. We can vectorize it further by creating a new dimension for attention heads.

To understand the trick being used, consider a basic matmul of (8, 6) by (6, 2):

Now suppose we want to multiply our LHS by another (6, 2) matrix. We can do it all in the same operation by concatenating the two RHS matrices along columns:

If the yellow RHS block in both diagrams is identical, the green block of the result will be as well. And the violet block is just the matmul of the LHS by the red block of the RHS. This stems from the semantics of matrix multiplication, and is easy to verify on paper.

Now back to our multi-head attention. Note that we multiply the input x by a whole list of weight matrices - in fact, by three lists (one list for Q, one for K, and another for V). We can use the same vectorization technique by concatenating all these weight matrices into a single one. Assuming that NH * HS = D, the shape of the combined matrix is (D, 3 * D). Here's the vectorized implementation:

# x has shape (B, N, D)
# In what follows:
#   NH = number of heads
#   HS = head size
#   NH * HS = D
# W is expected to have shape (D, 3 * D), with all the weight matrices for
# Qs, Ks, and Vs concatenated along the last dimension, in this order.
# Wp is a weight matrix for the final linear projection, of shape (D, D).
# The result is (B, N, D).
# If do_mask is True, each attention head is masked from attending to future
# tokens.
def multihead_attention_vec(x, W, NH, Wp, do_mask=False):
    B, N, D = x.shape
    assert W.shape == (D, 3 * D)
    qkv = x @ W  # (B, N, 3 * D)
    q, k, v = np.split(qkv, 3, axis=-1)  # (B, N, D) each

    if do_mask:
        # mask is a lower-triangular (N, N) matrix, with zeros above
        # the diagonal and ones on the diagonal and below.
        mask = np.tril(np.ones((N, N)))

    HS = D // NH
    q = q.reshape(B, N, NH, HS).transpose(0, 2, 1, 3)  # (B, NH, N, HS)
    k = k.reshape(B, N, NH, HS).transpose(0, 2, 1, 3)  # (B, NH, N, HS)
    v = v.reshape(B, N, NH, HS).transpose(0, 2, 1, 3)  # (B, NH, N, HS)

    kq = q @ k.swapaxes(-1, -2) / np.sqrt(k.shape[-1])  # (B, NH, N, N)

    if do_mask:
        # Set the masked positions to -inf, to ensure that a token isn't
        # affected by tokens that come after it in the softmax.
        kq = np.where(mask == 0, -np.inf, kq)

    att = softmax_lastdim(kq)  # (B, NH, N, N)
    out = att @ v  # (B, NH, N, HS)
    return out.transpose(0, 2, 1, 3).reshape(B, N, D) @ Wp  # (B, N, D)

This code computes Q, K and V in a single matmul, and then splits them into separate arrays (note that on accelerators these splits and later transposes may be very cheap or even free as they represent a different access pattern into the same data).

Each of Q, K and V is initially (B, N, D), so they are reshaped into a more convenient shape by first splitting the D into (NH, HS), and finally changing the order of dimensions to get (B, NH, N, HS). In this format, both B and NH are considered batch dimensions that are fully parallelizable. The QK^T computation can then proceed as before, and Numpy will automatically perform the matmul over all the batch dimensions.

Sometimes you'll see an alternative notation used in papers for these matrix multiplications: numpy.einsum. For example, in our last code sample the computation of kq could also be written as:

kq = np.einsum("bhqd,bhkd->bhqk", q, k) / np.sqrt(k.shape[-1])

See this post for my detailed notes on this notation.

Code

The full code for these samples, with tests, is available in this repository.

[1]	In LLM papers, D is often called d_{model}.

[2]	In the GPT-3 paper, this is also true for all model variants. For example, the largest 175B model has NH=96, HS=128 and D=12288.

[3]	It's also not as easy to define mathematically: how do we make a non-square matrix triangular? And what does it mean when the lengths of the two inputs are different?