Eli Bendersky's website - Machine Learning

Thoughts on starting new projects with LLM agents

2026-06-06T17:38:00-07:00

A few months ago I wrote about using LLM agents to help restructuring one of my Python projects. It's worth beginning by saying that the rewrite has been successful by all reasonable measures; I've been able to continue maintaining that project since then without an issue.

In this post, I want to discuss another project I've recently completed with significant help from agents: watgo. In this project many things are different; most notably, it's a from-scratch project rather than a rewrite, and it uses a different programming language (Go). This post describes my experience working on the project, and some lessons learned along the way.

The process

This is a new project, so it required extensive design. I began by iterating on the design with the agent, with a sketch of the API. For this purpose, I recommend using a Markdown file committed into the repository for future reference.

After that, I started asking the agent to write CLs [1] in a logical order that made sense to me, keeping them small and reviewable (more on this in the next section). Sometimes it's not easy to have a small CL, and multiple rounds of revision may confuse the agent; in this case, I commit the CL and then go back and ask the agent to modify or refactor the code, as much as needed, with separate CLs. In the worst case, the whole sequence can be reverted if I feel we've taken the wrong direction (branches could also be helpful here for more complicated scenarios).

This point is worth reiterating: sometimes a single CL is a huge step forward, but requires lots of review, cleanup and refactoring to be viable. I've had multiple instances where an agent produced several days of work in a single CL, but I then spent hours instructing it to clean up and refactor. Overall, it's still a productivity gain, just not as much as some pundits would like us to believe.

Keeping the human in the loop

Given the current state of agent capabilities, I think it's worth splitting projects into two categories:

Low importance / prototype / throw away projects where deep code understanding is unnecessary. These can be "vibe-coded" (submitting agent code without even reviewing it).
High importance projects that I actually want to maintain; here, vibe-coding is ill advised and I insist on reviewing and guiding all code the agent writes before it's submitted (or shortly after, as discussed above).

The watgo projects is a clear example of (2): I certainly intend to maintain this project in the long term, so I insist on code that I understand. With very few exceptions, no code gets in without full review and often multiple rounds of revisions.

Even if the cost for writing code went down, maintaining a project is so much more than that. It's triaging and fixing bugs, it's thinking through what needs to be done rather than how to do it, it's keeping the code healthy over time, and so on. As Brian Kernighan said:

Everyone knows that debugging is twice as hard as writing a program in the first place. So if you're as clever as you can be when you write it, how will you ever debug it?

Maybe at some point agents will become good enough that projects in category (2) can be implemented and maintained completely autonomously. Maybe. But we're certainly not there yet. My hunch is that getting there will require crossing the AGI line [2], after which little in our world remains certain.

Practical workflow

If you're using an agent to send an actual PR and only review that, it's difficult to be disciplined enough to actually perform a thorough review. I find the following method to be more reliable:

I use a CLI agent running locally in my repository, and ask it to update the code there. In parallel, I have a VSCode window open in the same project, where I can:

Review the agent's changes using VSCode's diff view
Make my own tweaks and code changes if needed

Once I'm pleased with the change, I manually create a commit.

Keeping the CLs small

As mentioned above, it's imperative to keep making progress in small chunks, with small enough CLs that a human can fully understand in a single review. It's very tempting to sprint ahead submitting thousands of lines of code every day, but this temptation has to be avoided. Coding with an agent is like speed-reading; yes, you're making more progress, but comprehension suffers the faster you go.

Particularly for refactoring, agents still take the shortest route to destination. It's important to guide them to think about the "big picture" at all times, find all instances where X is better done as Y, not just a single place noticed during a review. This is why it's sometimes OK to have a CL submitted before you fully agree with everything, and go back to it later for several refactoring rounds. Source control works amazingly well when pair-coding with agents.

Testing strategy

It's a key point discussed in every "how to succeed with AI" article, but still critical enough to reiterate here: a solid testing strategy is absolutely crucial for success. Agents produce - by far - the best results when they have a solid test suite to test their code against.

With the pycparser rewrite, I had a large existing test suite. For watgo, the very first thing I did was think through how to adapt the test suites of the WASM spec and of the wabt project for my needs.

If your project doesn't have such tests to rely on, this should be your first order of business - finding one, or building one from scratch. Beware of self-reinforcing loops though; it's dangerous to trust agents for both the tests and the implementations tested against them.

Language choice - Go for agent-written projects

Go is a fantastic language for agents to write, because it's designed to be very readable by humans. The biggest strengths of Go are exactly what makes the experience of reviewing agent code so positive:

Go changes very infrequently, so you don't have to wonder "are we using the most modern / idiomatic approach" or "what the hell is this construct" as often as with other languages (looking at you, Python and TypeScript).
There are relatively few ways to accomplish the same thing in Go, further lowering the mental burden.
The standard library is rich and there's much less need to keep abreast of the package-everyone-uses du jour.
In general, Go is designed for readability, with a mild-but-still-strong type system, uniform formatting, explicit error propagation and opinionated choices already made for you.

Since most of the time spent by humans when using agents is reading rather than writing code, these effects compound and produce a great experience. Recall the discussion of how some languages are optimized for writability (Perl) while others are optimized for readability (Go)? Well, when working on a project with an agent we live in a world of 99% reading vs. 1% writing, so this really matters.

I find this aspect really crucial in light of the earlier points made in this post - namely, keeping the human in the loop by understanding and reviewing all of the agent's design choices and code.

Final thoughts

If you're working on a subject that's completely new to you, I would strongly recommend against the approach described in this post. To really learn something, you have to work through it from scratch, yourself, reading, designing, writing the code. Agents don't change this basic fact; even before agents, if you wanted to learn X, copying it from Stack Overflow or some other project clearly wasn't the right way to go. Similarly, while agents can be used as a prop for learning, they cannot learn for you.

As a corollary, junior engineers should exercise extreme caution when relying on LLMs. There's no replacement to hard-won experience and the sweat and tears of learning new, challenging topics. Learning is supposed to be hard; if it's too easy, you're probably not learning.

For senior engineers, agents are a boon; it's a great tool to increase productivity, avoid the boring stuff, and get unstuck from procrastination; but only when used judiciously.

[1]	CL stands for Changelist, also known as a "patch" or a "diff" - basically a standalone commit that touches one or more files. This term originates from the source control systems Perforce and Subversion.

[2]	Programming is the ultimate realization of thought; if machines can design, produce, maintain and understand code better than humans, it means they can start improving themselves, which is the definition of singularity.

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.