The Endgame for Mechanistic Interpretability

The endgame for mechanistic interpretability is formal methods.

Mechanistic interpretability is currently pulled between competing visions. On one side, Neel Nanda argues for pragmatic interpretability: grounding work in real models and judging progress by empirical feedback, even when the underlying understanding is partial. On the other, Leo Gao defends ambitious interpretability: a long-term bet on building circuits that are necessary and sufficient for behavior, on the view that deeper mechanistic understanding is what will generalize across model changes. This disagreement is treated as a question of methods, but the deeper divide lies elsewhere.

The disagreement persists because mechanistic interpretability lacks an agreed-upon end goal. What, exactly, would count as success?

The Telos of Mechanistic Interpretability

Today, feature labeling, circuit discovery, probing, activation patching, and causal interventions coexist as productive methods, yet they are only loosely coordinated. The field lacks a shared ideal for what interpretability methods are ultimately meant to deliver.

One possible telos is legibility: producing explanations that are intelligible to humans. On this view, interpretability succeeds when it tells a coherent story about why a model behaves as it does. But explanations that fail under counterfactual intervention may sound plausible, yet provide no reliable handle for control. Even advocates of curiosity-driven interpretability typically hope their findings will eventually support some downstream use, and legibility alone does not guarantee this.

A second telos is scientific understanding. Intervention is used to test hypotheses, identify causal structure, and build general explanations. Mechanistic interpretability often operates successfully under this ideal. But LLMs are not natural objects. They are engineered software artifacts. Focusing on understanding alone leaves a powerful affordance unused. Interventions need not merely reveal structure, they can permanently modify the system while preserving formally specified properties. A purely scientific telos does not demand patchability, certification, or correctness under modification.

A third telos is capability enhancement. From this perspective, interpretability is valuable only insofar as it accelerates optimization. The natural equilibrium of this ideal favors systems that are maximally effective and maximally opaque.

These ideals are not mutually exclusive, but no single one provides a stable foundation for the field. From this perspective, mechanistic interpretability ought to orient itself toward debuggability: the ability to localize failures to specific mechanisms, intervene on those mechanisms predictably, and certify that the intervention preserves desired behavior on bounded domains. This telos subsumes legibility and scientific understanding, while resisting the drift toward opacity implicit in a purely capability-driven view.

In what follows, I make this notion precise.

Desired Goals for LLM Debuggability

In an idealized setting, debugging an LLM would proceed from localization, to intervention, to certification. Each stage places strictly stronger demands on our mechanistic understanding, and each rules out large classes of explanations.

Localization: Identifying the Responsible Mechanism

The first requirement of debuggability is localization: the ability to identify which internal mechanisms are responsible for a given behavior, and to distinguish mechanisms that generalize from those that merely correlate with it.

In the strongest form, a debuggable localization supports counterexample search. For a bounded LLM input domain D, this means being able to determine whether the behavior can occur without the mechanism being active, or whether the mechanism can be active without producing the behavior (and, when possible, to surface concrete inputs that witness such cases).

Intervention: Surgical, Mechanism-Level Debugging

Localization is only meaningful if it admits intervention. Once a failure is traced to a mechanism, we’d like to modify that mechanism in a way that is predictable and targeted:

1. The responsible head, MLP, or subspace can be modified or constrained.

2. The intervention removes the undesired behavior on a specified domain.

3. The intervention does not induce collateral damage elsewhere in that domain.

Certification: Domain-Bounded Safety Guarantees

The final goal of debuggability is certification: the ability to make exhaustive, falsifiable claims about model behavior on bounded domains.

For a formally specified domain D, this can mean proving that no harmful token is produced for any input in D, a claim that is in principle achievable for sufficiently constrained settings. Certification may also take the form of subcircuit-level bounds, for example structural invariants that rule out entire classes of behavior by construction, such as proving that a circuit cannot bypass a guard layer unless a specific feature is active.

What a Debuggable Explanation Can (and Can’t) Promise

First, some anti-goals:

1. Debuggability does not imply that a trained Transformer can be cleanly de-compiled into a single symbolic program. Transformer models are not limited to discrete algorithmic control flow. They exploit continuous geometry in high-dimensional embedding spaces, superposition, and distributed representations. Any realistic abstraction must preserve this expressive freedom rather than erase it.

2. Debuggability does not entail identifying a single, privileged “cause” of an output or phenomenon. Model behavior is mediated by deep, branching causal pathways with redundancy, overlap, and compensatory mechanisms. A debuggable explanation need not be unique. What matters is that the identified mechanisms are enough to explain and control the behavior within scope, and that alternative bypasses would be surfaced by counterexample search rather than hidden by storytelling.

3. Debuggability does not aim at global safety proofs of the form “this model will never produce harmful output.” Even if “harmful” were formally defined, the unconstrained input space of frontier LLMs is beyond the reach of any existing or foreseeable verification technique. Any agenda that predicates success on global guarantees is doomed to either vacuity or false confidence.

Instead, debuggability is about constructing a family of verified, compositional abstractions that faithfully reproduce model behavior on bounded domains and support predictable intervention. These abstractions are partial, local, and plural, but they are exact where they apply.

The relevant analogy is debugging a large, safety-critical software system. One cannot prove “Chrome will never crash.” But one can prove that specific routines are memory-safe, that sandboxing prevents certain classes of process escape, that critical invariants are preserved across refactors, and that a given patch eliminates a vulnerability without introducing regressions.

The same logic applies to LLMs. Meaningful debuggability consists in guarantees such as:

• This subcircuit cannot activate a forbidden feature on domain D.

• This intervention removes a failure mode while preserving all other behaviors in scope.

• This pathway is structurally incapable of bypassing a guard unless a specific internal condition is met.

The Necessity of Formal Methods

Note that the above are universal claims over bounded domains. They are not distributional or probabilistic, and they do not rest on sampling. This is why a debuggability-oriented interpretability agenda is necessarily coupled to formal methods. SMT solvers, abstract interpretation, and neural verification frameworks are not optional add-ons. They are the only frameworks in which claims of impossibility, preservation, or closure under intervention can be made precise.

This vision does not require that today’s frontier LLMs be fully verifiable end-to-end. What matters is that the debuggability of Transformer models has precedent:

• Sparse circuit extraction shows that models contain relatively isolated, algorithmic subcircuits that remain stable under targeted intervention.

• Symbolic Circuit Distillation is an early example of automated extraction, where neural mechanisms can be proven formally equivalent to symbolic programs.

• Neural verification work (e.g. Reluplex and Marabou) establishes that exhaustive reasoning is possible once models are reduced to verification-friendly components on bounded domains.

• Alternative attention mechanisms suggest that standard attention (the dominant barrier to SMT verification in Transformers) is an architectural choice rather than a theoretical necessity, opening the door to verification-aware model design with comparable performance.

Taken together, these results shift the problem from conceptual impossibility to engineering integration and scale. Debuggability is about being able to say, with confidence: “This mechanism, on this domain, behaves this way, and if it didn’t, we would know.”

What De-compiling Actually Looks Like

Here’s what a possible “de-compilation” pipeline might look like:

1. Identify stable linear regions (local programs)

The smallest unit of analysis is a particular mechanism that has a stable branch structure on a bounded domain. Many verification-friendly components (affine maps, threshold gates, max/Top-K selection) behave like ordinary programs. Once you know which branch you’re in (e.g. which segment of a piecewise function is active, which items win a Top-K) the remaining computation is just affine arithmetic.

A “local program” is a region of inputs defined by explicit guard conditions (linear inequalities) together with the affine map executed under those guards. The stability part matters, because you want margins on the guards (e.g. thresholds not near zero, Top-K winners separated from runners-up) so that small perturbations or permissible interventions don’t flip the branch decisions and invalidate the explanation.

Here’s what a concrete example might look like:

Head 31.2: Helps break text into paragraphs

Empirically verified: attention weight exceeds ε when a token from set {‘\n’, ‘\n\n’} or a discourse marker token occurs at position t−1.

2. Factor into meaningful subspaces

Factoring into meaningful subspaces is the step where you decompose a mechanism’s activations into low-dimensional directions that have stable semantics across inputs and contexts, such as syntactic markers, sentiment, or safety-relevant features. A single local program may operate over several subspaces, and the same subspace may participate in many different local programs.

Without subspaces, interventions are blunt (ablating whole heads or MLPs). With them, interventions can be surgical (editing or bounding specific directions while leaving others untouched).

Ideally, these subspaces exhibit “functional coherence,” where moving along the subspace produces predictable, monotonic changes in model behavior on a bounded domain.

3. Extract formally verifiable causal circuits

In this step, we compose local programs and subspaces into a single object that supports global, counterfactual claims about behavior on a bounded domain. Formally, this means specifying an interface, a domain, and a set of admissible interventions, and then proving that the neural subcircuit is equivalent to (or soundly approximated by) a symbolic specification on that domain.

My project Symbolic Circuit Distillation builds in this direction by providing formally verified functional abstractions on bounded domains. Achieving this level of debuggability places strong pressure to redesign core Transformer components that are poorly suited to formal verification.

Multiple circuits may implement overlapping or reconstructed features elsewhere in the model. What matters is that, within scope, the abstraction is correct and closed under counterfactuals. If a bypass exists, formal search will find a concrete counterexample. This is the point where mechanistic interpretability becomes robust to patching and refactoring, and where safety-relevant guarantees become possible. Circuits stop being explanatory stories and become objects you can edit, reason about, and certify.

How these verified control abstractions are surfaced to human operators (whether through explicit query languages, automated tooling, or learned interfaces) is an important but orthogonal problem, and not required for the core claim of debuggability.

Interpretability as Control

The central question is whether mechanistic insight can support reliable, bounded, and verifiable control over systems that matter.

The result is a patchwork of verified abstractions: local programs, meaningful subspaces, and formally specified circuits. Many such decompositions may exist. Verification removes arbitrariness not by enforcing uniqueness, but by enforcing sufficiency.

I am curious to hear thoughts from other researchers. You can reach me on X: @neelsomani

Thanks to Maaz for giving feedback prior to posting.

Intro to Routing: Mixture-of-Experts and Expert Choice

I derive MoE and EC from first-principles.

In this post, I’ll cover routing mechanisms for large language models. People often talk about Mixture-of-Experts and Expert Choice, so my goal is to give a first-principles walkthrough that explains how these methods arise naturally. You can think of this as how I would have derived them myself, or as an ex post facto explanation that clarifies the logic behind their design.

Mixture-of-Experts (MoE)

Historical Roots of MoE

The engineering motivation behind MoE is straightforward. The goal is to take N expert functions f_i and compute a weighted average of their outputs based on how confident the model is in each expert.

The basic idea is simple. First, compute logits for each expert:

\(z_i = (W_g x + b)_i\)

Then convert these logits into a probability distribution:

\(g_i(x) = \text{softmax}(z(x))_i\)

Finally, compute the convex combination of expert outputs:

\(y = \sum_i g_i(x) * f_i(x)\)

Training proceeds exactly as it does for a standard feedforward layer. This formulation matches the original work of Jordan and Jacobs (1991).

The drawback is that this approach requires evaluating every expert for every token, including experts with very low probability. This becomes expensive as N grows.

Top-1 Gating

Let’s say we want to perform “Top-1 gating,” where we select only the highest scoring expert. Specifically, we want to compute all g_i(x), pick the top expert s, run only that expert, and set:

\(y = g_s(x) * f_s(x)\)

This has a property you might find unexpected. Even if all f_i point in roughly the same direction, the magnitude of y is smaller than the raw output of f_s(x), since f_s is scaled by g_s. You can argue that this is sometimes desirable, because lower confidence often corresponds to smaller updates in many other ML architectures. But this is a post hoc justification, and the real reason is just that attempts to renormalize y tend to perform worse in practice.

Backpropagation follows from the product rule. For the expert parameters θ_f:

\(\frac{dL}{d\theta_{f_s}} = \frac{dL}{dy} g_s(x) \frac{df_s(x)}{d\theta_{f_s}}\)

By symmetry, the gradient for the router parameters θ_g is:

\(\frac{dL}{d\theta_{g_s}} = \frac{dL}{dy} f_s(x) \frac{dg_s(x)}{d\theta_{g_s}}\)

For i != s, we have dL/dθ_{f_i} = 0, since those experts do not run. But dL/dθ_{g_i} is not zero, because the softmax couples all logits z_i, so each g_i influences the scaling of f_s.

The gradients are undefined at the exact boundaries where the identity of the top expert changes. But that’s typically fine, just like it’s not an issue for ReLU or other piecewise differentiable functions.

The major problem is that unused experts never improve. Training collapses to a solution where g_s ~= 1 for whichever expert happened to win early in training. In an ideal setting, we would want all experts f_j for j != s to receive at least some tokens so they continue to learn.

So we define the proportion of tokens routed to expert i in a batch:

\(p_i = \frac{1}{|B|} \sum_{t \in B} \mathbf{1}[s_t = i]\)

You might attempt to regularize this by optimizing something like:

\(L’ = L + \lambda * \mathrm{KL}(p || U)\)

where U is a uniform distribution. This would flatten the token allocation, and in principle could prevent collapse.

But since p_i above depends on an argmax (through s_t), the gradient of p_i is zero almost everywhere. The model receives no useful signal from this penalty.

As a result, we need some other differentiable penalty that discourages any single expert from dominating the routing. The goal is to reduce the confidence g_i for experts that receive too many tokens and increase it for experts that receive too few.

Flattening the Argmax Distribution

Since p_i isn’t useful for differentiation, we try using the Gumbel max trick, which provides a way to make the argmax behave like a soft, differentiable sampling process. As a general rule, if z_j are logits and ε_j ~ Gumbel(0, 1), then:

\(\arg\max_j { z_j + \varepsilon_j } \sim \text{Categorical}(\text{softmax}(z))\)

This distribution gives us Pr[s_t = i] for each expert i. (In principle, we could even use the noisy sampling in the forward pass.) More importantly, this lets us compute an expected load for each expert and attempt to flatten it.

In our case, let z_j be the logit that produces the gating probability g_j(x_t) = softmax(z(x_t))_j. Then we assume:

\(s_t = \arg\max_j { z_j(x_t) + \varepsilon_j }\)

This implies:

\(\Pr[s_t = i] = \text{softmax}(z(x_t))_i\)

and therefore the expected load for expert i is:

\(\text{E}[\text{load}_i] = \sum_{t \in B} \text{softmax}(z(x_t))_i = \sum_{t \in B} g_i(x_t)\)

With this expected load vector, we can now try to flatten it. Possible approaches include:

• KL(E[load] || U)

• L2 distance to uniform

• Entropy maximization

A common alternative is to minimize the coefficient of variation (or a similar quantity), which flattens the distribution without the numerical instability of KL or entropy when some components get small:

\(\text{CV}(\text{E}[\text{load}]) = \frac{\text{std_dev}(\text{E}[\text{load}])}{\text{mean}(\text{E}[\text{load}])}\)

We add an auxiliary term:

\(L_{\text{aux}} = \text{CV}(\text{E}[\text{load}])^2\)

where CV is squared for optimization convenience. So the full loss becomes:

\(L’ = L + \lambda * L_{\text{aux}}\)

Practical Implementation Today

In modern MoE implementations, a simpler surrogate auxiliary loss is used. It’s not statistically derived or theoretically clean, but it works well in practice:

\(L_{\text{aux}} = \sum_i \frac{1}{|B|} \left[ p_i(B) \sum_{t \in B} g_i(x_t) \right]\)

In theory, you might reach this form by starting from the objective:

\(\sum_i {p_i}^2\)

Since the p_i sum to 1, minimizing this objective encourages a uniform allocation via the method of Lagrange multipliers.

Then, assuming we really are sampling with Gumbel noise, then given a sufficiently large batch, the empirical load fraction can be approximated via the soft probabilities g_i:

\(p_i \approx q_i = \left( \frac{1}{|B|} \sum_{t \in B} g_i(x_t) \right)\)

where q_i is a differentiable surrogate for the observed load proportion. Using this approximation on one of the factors:

\(\sum_i {p_i(B)}^2 \approx \sum_i p_i(B) * \left( \frac{1}{|B|} \sum_{t \in B} g_i(x_t) \right)\)

which is the surrogate used above.

The effect is straightforward. If an expert receives too many tokens (that is, if p_i is too large), the loss increases, which pushes the model to reduce g_i.

The key detail is that the derivative of p_i(B) is locally zero, since p_i depends on an argmax that does not change in a small neighborhood. For this reason, implementations mark p_i(B) as a “no-grad” quantity.

Compared to this heuristic approach, the formulation based on Gumbel noise and the coefficient of variation is mathematically cleaner. The stochastic forward pass aligns naturally with the probabilistic reasoning used in the backward pass, and the load penalties follow from that framework without ad hoc constructions. Despite this conceptual clarity, the surrogate loss above remains more widely used.

Generalizing to Top-K

If we want to use the Top-K experts rather than Top-1, first we need a constructive definition of the statistical process. As before, we compute g_i(x_t) for each expert i. But this time, we select the Top-K experts instead of a single one.

We’ll try to take the same approach as before. To compute the expected load for expert i, we need

\(q_i(x_t) = \Pr[i \in S]\)

where S is the set of the K selected experts. The expected load is:

\(\text{E}[\text{load}_i] = \sum_t q_i(x_t)\)

Just like in the Top-1 case, we can minimize CV(E[load_i])². In theory, E[load_i] is differentiable because Top-K sampling follows a Plackett-Luce distribution. But in practice, the gradient is computationally intractable. So we end up using the same surrogate as in the previous section.

Finally, it’s worth noting that the original formulation in Shazeer et al. (2017) used a different approach. Instead of Gumbel noise, the authors added normal noise to the logits, and they renormalized the weights g_i after selecting the Top-K elements. The methodology and derivation differ from the conceptual framework presented above.

Expert Choice (EC)

Expert Choice (Zhou et al. 2022) observes that the biggest pitfall of MoE is that an expert can get overloaded. If too many tokens want the same expert, that expert overloads and gets a huge share of the gradient. If too few tokens go to an expert, that expert’s weights collapse and it fails to train. That is why we needed that hacky regularization term earlier.

Imagine you’re Google serving inference at massive scale. You don’t care about routing every token perfectly. You care about keeping all experts busy, avoiding hot spots, and guaranteeing predictable latency.

Rather than computing g_i(x_t) for each token and selecting argmax_i g_i (letting the tokens pick the experts), you can let the experts pick which tokens they want to serve. Each expert receives a fixed budget of M tokens and selects the M tokens for which it believes it is most useful.

You still evaluate g_i(x_t) for all experts i and all tokens t in the batch. For each expert i, you select:

\(B_i = \text{TopM}_{t \in B} (g_i(x_t))\)

Each expert receives exactly M tokens (or up to M if the batch is small). How this affects backpropagation depends on what you do when multiple experts select the same token. For simplicity, assume the output is the sum:

\(y = \sum_{i : x_t \in B_i} g_i(x_t)* f_i(x_t)\)

(Note that real EC implementations use more complicated aggregations.) In this case, backpropagation for the expert parameters is exactly the same as in MoE. If an expert is not selected, it receives no gradient. If it is selected, then

\(\frac{dL}{d\theta_{f_i}} = \frac{dL}{dy} g_i(x_t) \frac{\partial f_i(x_t)}{\partial \theta_{f_i}}\)

The gradient with respect to the router parameters is surprisingly simple:

\(\frac{dL}{d\theta_{g_i}} = \sum_{x_t \in B_i} \frac{dL}{dy} f_i(x_t) \frac{\partial g_i(x_t)}{\partial \theta_{g_i}}\)

Notice that we didn’t have to differentiate through the Top-M operator at all, since the gradients only flow through g_i for the tokens actually selected by expert i.

There’s a glaring pitfall here. What if a token isn’t selected by any of the experts? In practice, implementations handle this by increasing M or by routing those stray tokens to the expert with the largest g_i(x_t).

Future Directions

I hope this post was informative. In the future, I plan to cover other routing mechanisms such as Mixture-of-Depths (MoD) or Switch Transformers, a paper by authors I respect a lot.

Conceptually, MoD pushes sparsity in a more radical direction. Instead of selecting which expert network should process a token, the router selects which layers of the transformer the token should flow through. The resulting interactions make MoD a significantly harder routing problem than MoE or Expert Choice.

When the research landscape matures further, I plan to revisit this topic with a dedicated analysis.

A Minimal Route to Transformer Attention

Is attention inevitable?

In this post, I’ll show how a small set of reasonable assumptions can recover the Transformer attention mechanism. Some parts of attention are theoretically motivated, while others are arbitrary choices. I’ll explicitly call out which is which.

To see why attention exists, it helps to recall its predecessor: the recurrent neural network (RNN). Classic encoder-decoder RNNs process a sequence token by token. Each new hidden state incorporates the current token and the previous hidden state, producing a vector you can think of as an “accumulator” of everything seen so far. After ingesting the final token, that accumulated vector is repeatedly fed to the decoder, which predicts output tokens until it emits a STOP symbol.

The problem is long-range dependence. If an important token appeared far earlier in the sequence (say, the first of 10,000 tokens), its influence becomes diluted as the RNN processes additional tokens. The model simply forgets.

Ideally, the model should use all previously seen tokens to compute the information needed to predict the next token, weighting each earlier token by how relevant it is for that prediction. That suggests computing a relevance score between a position i and each other position j, and then combining some function of the embeddings accordingly.

Formally, define a scalar relevance function that takes the embeddings X along with indices i and j:

\(u(i, j | X_1, …, X_n)\)

We work in embedding space rather than raw token IDs to avoid meaningless geometric assumptions (e.g., token 9 is not inherently “closer” to token 10). One-hot encodings would work, but are much more sparse.

Then the model’s output vector at position i can be written as:

\(y_i = G(\{x_j, u(i, j | X)\}_j)\)

where G aggregates some function of the embeddings x_j alongside their relevance to position i. (During autoregressive generation, i corresponds to the most recently produced token.) We don’t yet know the form of G or u. Our goal is to characterize the simplest constraints that lead directly to Transformer-style attention.

Observation: Enforce permutation symmetry

We want to constrain the space of possible functions for G and u.

Once we have the relevance scores u(i, j), the output y_i should not depend on the order in which the pairs (x_j, u(i, j)) are provided. In other words, if we reorder the elements indexed by j, the result should remain the same. This requires G to be “permutation-invariant” over the set {(x_j, u(i, j))}_j.

The Deep Sets theorem (Zaheer et al., 2017) tells us that any such function can be written as:

\(G(\{x_j, u(i, j)\}_j) = ρ\left(\sum_j φ(x_j, u(i, j))\right)\)

Here ρ and φ are arbitrary differentiable functions. We fix the index i, since the invariance only applies over j. Differentiability ensures that the overall model can be trained with gradient-based methods.

At this point, ρ and φ are still completely general, and we also need to define u. We will impose further assumptions to narrow down their form.

Assumption 1: ρ is the identity function

ρ could output many different types of objects. For example:

• It could output a scalar, but that would discard most of the information from the embeddings.

• It could output an O(N)-dimensional vector, with one component per input element, but that would make the output scale with sequence length and defeat the purpose of summarizing information.

• It could output a vector in some intermediate dimension, or even map into a different space/manifold entirely.

All of these are technically possible. In practice, Transformers set ρ to be the identity function, so:

\(G(X) = \sum_j φ(x_j, u(i, j))\)

This simplifies the structure of G and lets us focus on constraining φ and u.

Assumption 2: Relevance-contribution proportionality

Even with ρ set to the identity, φ could be any function of the embedding x_j and the relevance score u(i, j). To simplify the form, we assume that if a token’s relevance is scaled by a constant k, its contribution scales by the same factor:

\(φ(x_j, k * u(i, j)) = k * φ(x_j, u(i, j))\)

This is not the only possible relationship. For example, we could have chosen a quadratic or some other monotonic transformation in u(i, j). The key requirement is simply that φ should separate into:

• A scalar measuring how important x_j is

• A vector capturing what x_j contributes

Under the linear version of this assumption, we get:

\(φ(x_j, u(i, j)) = u(i, j) * φ(x_j, 1)\)

Define v(x_j) = φ(x_j, 1), yielding:

\(φ(x_j, u(i, j)) = u(i, j) * v(x_j)\)

This makes φ explicitly separable, where u(i, j) purely controls magnitude (relevance), and v(x_j) determines the content being contributed.

Assumption 3: Linear change of coordinates

At this point, v(x_j) could be any function of x_j. To simplify the model and keep it efficient to compute, we assume v is a linear transformation of x_j:

\(v(x_j) = W_V x_j\)

Substituting this into the previous expression gives:

\(G(X) = \sum_j {u(i, j) * W_V x_j}\)

This means each token contributes a linearly transformed version of its embedding, weighted by its relevance score u(i, j).

Observation: Constrain u for efficient parallel computation

We want u(i, j) to be computable efficiently on hardware like GPUs. Here, “efficient” refers to low sequential depth in the computational graph, not necessarily a low number of arithmetic operations. GPUs can execute many multiplications in parallel, but long chains of dependent operations create bottlenecks. For example, a recurrent computation with O(N) sequential steps is slow for long sequences, but a matrix multiply has O(1) sequential depth and is highly parallelizable.

If we allowed a fully general relevance function such as:

\(u(i, j | X) = g_\theta(x_i, x_j, \text{context}(X))\)

where context(X) examines all tokens at once, we would need to evaluate this network O(N²) times for a single layer, which is too slow.

Alternatively, we could define a single model:

\(g_θ: X → ℝ^{n \times n}\)

that outputs all pairwise relevance values directly. But that would require storing and training parameters of size O(N²), which locks the model to a fixed input length and scales poorly.

To keep computation parallelizable and scalable, we restrict u to be built from tensor operations such as:

• Linear projections

• Element-wise functions

• Inner products

• Reductions like sums

and avoid control flow or long sequential recurrences.

Assumption 4: Dot product similarity for u

A simple way to score the interaction between x_i and x_j is with a dot product. However, we don’t necessarily want similarity in the embedding space - we want similarity in a space optimized for relevance.

So, as we did for v(x_j), we first apply learned linear projections:

\(\begin{aligned} q_i &= W_Q x_i \\ k_j &= W_K x_j \end{aligned}\)

Then we define the relevance score as:

\(u'(i, j | X) = ⟨q_i, k_j⟩\)

We denote this version as u’ because additional modifications will be applied later.

Assumption 5: Pick a normalization for u

Next, we want the relevance scores u’(i, j) to measure relative importance. If the same constant were added to all scores, or if they were scaled uniformly, the ranking of tokens should not change. This motivates applying a differentiable normalization function over j.

There are several possibilities (e.g., softmax, Gumbel-Softmax). In practice, Transformers use softmax.

One final issue: the dot product ⟨q_i, k_j⟩ tends to grow in magnitude with the key/query dimension d_k. To prevent extremely large values from dominating the softmax, we scale the logits:

\(u’(i, j) = \frac{⟨q_i, k_j⟩}{\sqrt{d_k}}\)

Applying softmax normalization over j then gives:

\(G(x_i) = \sum_j \operatorname{softmax}_j\!\left( \frac{\langle W_Q x_i, W_K x_j \rangle}{\sqrt{d_k}} \right) \, W_V x_j\)

This is exactly the scaled dot-product attention used in Transformers.

Does a better attention mechanism exist?

So there you have it. If we impose the following assumptions:

1. ρ is the identity function

2. Each token’s contribution scales proportionally with its relevance score

3. A linear transformation maps embeddings to the value, key, and query vectors

4. Relevance is based on a dot product

5. Relevance scores are normalized with a softmax

we obtain the exact scaled dot-product attention used in Transformers.

While some of the choices were forced, they weren’t all theoretically required. There may be better options for ρ, for the similarity measure, or for the normalization function. Even more fundamentally, the Deep Sets form at the beginning was from imposing permutation-invariance, but we end up reinjecting positional encodings in practice. Exploring these variations could reveal new attention mechanisms with different computational or modeling advantages.

Killing the GIL: How To Use Python 3.14's Free-Threading Upgrade

The global interpreter lock (GIL) has been interfering with true parallelism in Python. That ends with Python 3.14.

For almost three decades, Python’s Global Interpreter Lock (GIL) has been the single mechanism standing between your CPU cores and real parallelism.

That changes with Python 3.14.

The new free-threaded build removes the GIL, allowing multiple threads to execute Python bytecode simultaneously. No multiprocessing, no pickle, no hacks.

In this post, I’ll:

1. Explain why the GIL existed and what it was protecting

2. Compare Python’s old concurrency models (threading, multiprocessing, asyncio)

3. Build Python 3.14 with the GIL disabled

4. Run a short multithreaded benchmark that finally scales with cores

5. Explain the results

Why the GIL existed in the first place

The GIL is a global mutex that historically allowed only one thread to execute Python bytecode at a time. It was there to protect CPython, the C implementation of the interpreter.

Every Python object in CPython lives on the heap as a C struct with a reference count. Each assignment or function call increments and decrements those counters constantly. If two threads updated the same object’s reference count simultaneously, you could get memory corruption or premature frees that crash the interpreter.

Adding locks around every Python object would have been complex and slow, so early CPython took the simple route: wrap the entire interpreter in one global lock. That made single-threaded execution safe, but prevented true multithreading for CPU-bound workloads.

How concurrency previously worked

Before 3.14, Python offered three main concurrency models, each with trade-offs:

• threading (old): uses real OS threads, but only one can execute Python bytecode at a time because of the GIL. Good for I/O, useless for parallel CPU work.

• multiprocessing: spawns multiple processes, each with its own interpreter and GIL. True parallelism, but expensive, requiring separate memory, pickling overhead, and slower process startup.

• asyncio (green threads): runs everything cooperatively on one thread. Excellent for high-concurrency I/O, but it never uses more than one core.

With Python 3.14’s free-threaded build, threading becomes the best of all worlds: true parallelism across cores, shared memory without serialization, and minimal overhead.

Building Python 3.14 without the GIL

Compile it yourself:

git clone https://github.com/python/cpython
cd cpython
git checkout v3.14.0
./configure --prefix=$HOME/.py-314-ft --disable-gil
make -j && make install
$HOME/.py-314-ft/bin/python3 -V

Or with pyenv:

pyenv uninstall -f 3.14.0 || true
PYTHON_CONFIGURE_OPTS=”--disable-gil” pyenv install 3.14.0
pyenv local 3.14.0

Verify that you’re running a free-threaded build:

python3 - <<’PY’
import sys
print(”Free-threaded build:”, not sys._is_gil_enabled())
PY

You want: Free-threaded build: True.

Running a realistic multithreaded benchmark

Here’s a bit-optimized N-Queens solver. You can also download the repo on GitHub. It’s already efficient and CPU-bound, a good test of whether threads can finally scale.

# nqueens.py
import threading, time

def solve_row(n, cols=0, diags1=0, diags2=0, row=0):
    if row == n: return 1
    count = 0
    free = (~(cols | diags1 | diags2)) & ((1 << n) - 1)
    while free:
        bit = free & -free
        free -= bit
        count += solve_row(
            n, cols|bit, (diags1|bit)<<1, (diags2|bit)>>1, row+1
        )
    return count

def solve_threaded(n, n_threads):
    first_row = [(1 << c) for c in range(n)]
    chunks = [first_row[i::n_threads] for i in range(n_threads)]
    total = 0
    lock = threading.Lock()

    def work(chunk):
        nonlocal total
        local = 0
        for bit in chunk:
            local += solve_row(
                n, cols=bit, diags1=bit<<1, diags2=bit>>1, row=1
            )
        with lock:
            total += local

    threads = [threading.Thread(target=work, args=(c,)) for c in chunks]
    for t in threads: t.start()
    for t in threads: t.join()
    return total

if __name__ == “__main__”:
    for threads in (1, 2, 4, 8):
        t0 = time.perf_counter()
        solve_threaded(14, threads)
        dt = time.perf_counter() - t0
        print(f”threads={threads:<2}  time={dt:.2f}s”)

Run it once with standard CPython 3.14 (GIL on) and once with your free-threaded build. With the GIL, all runs take about the same time. With the free-threaded build, performance improves almost linearly with thread count.

Results

Example benchmark:

14-Queens on an M1 Pro, Python 3.14

That’s an ~8x speed-up without changing a single line of logic - just running under the free-threaded interpreter.

Caveats:

• C extensions: any binary package must be recompiled for free-threading, or it might quietly re-enable the GIL.

• Thread safety: without the GIL, race conditions are real. Protect shared state with locks, queues, or immutable data.

• Single-thread overhead: expect a 5-10% slowdown for purely single-threaded scripts due to atomic ops and internal locks.

Closing thoughts

The GIL made CPython simple and safe to implement, but it locked Python to a single core.

With Python 3.14, that trade-off is gone. For the first time, standard Python threads can run in true parallel on modern CPUs.

So go ahead and kill the GIL, and let me know how it works for you.

What You Didn't Learn in Berkeley CS 188 — Part 4

Is GRPO broken?

This is the fourth and final piece in my series on reinforcement learning. Previously, we covered classical RL, continuous control, and off-policy methods. The topic of LLM post-training is discussed all over X, so this primer should help anyone get up to speed.

Here’s how I like to think about post-training methodologies:

2x2 Quadrant of Post-Training Methods

SFT is simple. It’s just applying additional training iterations like the pre-training stage, but on a curated set of ideal (prompt, response) pairs. You might make this more efficient with a LoRA adapter.

In this post, we’ll focus on quadrant 2: DPO and offline GRPO. Along the way, I’ll point out how methods like online PPO and online GRPO fit in. Historically, online PPO came first, so understanding it helps explain DPO.

Theory of Relative Scoring

Before getting into the objective function of direct preference optimization (DPO), we need to motivate the idea of relative scoring.

We’re given lists of prompts x and pairwise responses a₊ and a_-:

\((x, a_{+}, a_{-})\)

All we know is that a₊ is preferred to a_-. A human may have rated them that way, or another signal might imply it (e.g. code that compiles > code that fails).

That setup doesn’t immediately lend itself to the methods we’ve seen so far. On-policy methods don’t work because neither response may be likely under the current policy. Off-policy methods still don’t work because we lack a defined reward.

You might try to make the model more likely to output a₊ than a_- by optimizing Pr[π(a₊ | x) > π(a_- | x)]. But that expression makes no sense. π(a | x) are constants given the model. Unless we add tunable parameters (say, some γ where f_γ(π) produces a new model), those probabilities don’t change.

You could define f_γ(π)(x, a) and optimize Pr[f_γ(π)(x, a₊) > f_γ(π)(x, a_-)], or build an even more general model f_γ(π, x, a₊, a_-) that directly outputs the likelihood that a₊ is better. But f is still abstract. It’s unclear how to parameterize it.

Instead, DPO (and the original online PPO post-training) take a simpler route by introducing a latent reward. The assumption is that if a human preferred a₊ to a_-, then there exists some implicit reward function r such that

\(r(a_{+} | x) + ε_{+} > r(a_{-} | x) + ε_{-}\)

where ε represents human noise or ambiguity. If we can learn that reward function, we can optimize the model accordingly.

Learning the Reward Function

One approach is maximum likelihood estimation. We denote a₊ ≻ a_- if a₊ is preferred. We’d like a function g such that:

\(g(r_\phi(a_{+} | x), r_\phi(a_{-} | x)) = Pr[a_{+} ≻ a_{-}]\)

and then optimize φ to maximize:

\(\ \prod Pr[a_{+} ≻ a_{-}]\ \)

Let’s try to define g. Notice:

\(\begin{align*} Pr[a_{+} ≻ a_{-}] &= Pr[r(a_{+} | x) + ε_{+} > r(a_{-} | x) + ε_{-}] \\ &= Pr[r(a_{+} | x) − r(a_{-} | x) > ε_{-} − ε_{+}] \end{align*}\)

So preference depends only on the difference between rewards. That implies translational invariance: g(u, v) = g(u + c, v + c). That property implies that g must be f(r(a₊ | x) - r(a_- | x)) for some function f, since g(u, v) = g(u − v, 0) = f(u − v), where the first equality follows by translational invariance.

Second, if r(a₊₊ | x) > r(a₊ | x) > r(a_- | x), the higher-reward response should never be less preferred. In other words, f must be non-decreasing: f’(x) >= 0

Finally, f(r(a₊ | x) - r(a_- | x)) + f(r(a_- | x) - r(a₊ | x)) = 1, which along with the previous condition, implies f(0) = ½, lim _t→∞ f(t) = 1, and lim _t→-∞ f(t) = 0.

Many functions f satisfy these conditions. The choice depends on what noise distribution you assume. In practice, DPO uses the logistic sigmoid, which assumes Gumbel noise:

\(\begin{aligned} σ(x) &:= \frac{1}{1 + e^{-x}} \vphantom{\frac{1}{1 + e^{-x}}} \\ U(a) &= r(a) + ε,\quad ε ∼ \text{Gumbel}(0, 1) \vphantom{\frac{1}{1 + e^{-x}}} \\ \implies Pr[U(a_{+}) > U(a_{-})] &= σ(r(a_{+}) − r(a_{-})) \vphantom{\frac{1}{1 + e^{-x}}} \end{aligned}\)

If noise were Gaussian, you’d recover the probit model instead.

The final objective to optimize r_φ is:

\(\begin{align*} \max J(\phi) &= \sum \log \sigma(r_{\phi}(a_{+}) − r_{\phi}(a_{-})) \\[4pt] \nabla J(\phi) &= \sum (1 − \sigma(r_{\phi}(a_{+}) − r_{\phi}(a_{-}))) [\nabla r_{\phi}(a_{+}) − \nabla r_{\phi}(a_{-})] \end{align*} \)

The KL Divergence Penalty & PPO

Now we have a reward function. Just like the traditional methodology for REINFORCE, you can optimize your policy with respect to the objective function:

\(\max_\theta J(\theta) = E_{\pi_{\theta}}[r(x, a)]\)

It’s a bit different from REINFORCE since there’s no discounted sum of rewards across a trajectory. Instead, it’s just a single-step reward that we’re optimizing with respect to. The problem with this approach is that it’s going to completely alter your model. The optimization will force the policy to output a₊ with very high probability, at the cost of everything else.

So the actual optimization for online PPO and DPO actually adds a constraint to prevent the policy from diverging too much from the original policy, π_ref:

\( \max_θ E_{\pi_{\theta}}[r(x, a)] \ \text{s.t. } KL(\pi_\theta | \pi_\text{ref}) < δ\)

That KL divergence constraint might make you think of PPO. But that similarity is completely superficial. Recall that the KL divergence constraint for PPO came from rewriting the objective function:

\(\begin{align*} \max_\theta J(θ) &= E_\tau[g(\tau)] = E_{x∼d^{π_{\text{old}}},a∼π_{\text{new}}}[A^{π_{\text{old}}}(x, a)] \end{align*}\)

We needed to constrain d^π_old ≈ d^π_new so we didn’t have to re-sample, and the best we could do was penalize KL(π_new || π_old) and establish an upper bound on the divergence of the state distributions.

The KL divergence constraint for online PPO and DPO is not fundamentally justified in the same way. It is simply the heuristic notion that we want π_new to be not too different from π_ref. You could theoretically derive this constraint if you think the true model follows a Boltzmann distribution, and you impose π_ref as a prior. This leads to the same objective function as above. But that’s not really where this KL divergence constraint comes from.

If you’re running online PPO, you’ll see the KL divergence penalty in the objective function to keep the policy π_new close to π_ref, and a clipping mechanism to keep the policy π_new within the trust region of π_old.

To finish the derivation for online PPO, we add the constraint that π_θ(a | x) must sum to 1:

\(\begin{align*} J(θ) &= \mathbb{E}_{\pi_\theta}[r(x, a)] - \frac{1}{β} \, KL(\pi_\theta \,\|\, \pi_\text{ref}) - λ \left[\sum_a \pi_\theta(a | x) - 1\right] \\[6pt] J(θ) &= \sum_a \pi_\theta(a | x) \left[r(x, a) - \tfrac{1}{β}\big(\log \pi_\theta(a | x) - \log \pi_\text{ref}(a | x)\big) - λ\right] - λ \end{align*} \)

Taking the gradient:

\(\begin{align*} \nabla J(θ) &= \sum_a \nabla \pi_\theta(a | x) \left[r(x, a) - \tfrac{1}{β}\big(\log \pi_\theta(a | x) - \log \pi_\text{ref}(a | x)\big) - λ\right] \\[6pt] &\quad - \tfrac{1}{β} \sum_a \pi_\theta(a | x) \, \nabla \log \pi_\theta(a | x) \end{align*} \)

You can go ahead and optimize π_θ using this gradient, and that’s exactly where methods like online PPO (or as we cover later, online GRPO) fit in.

Direct Preference Optimization (DPO)

DPO, on the other hand, attempts to turn this into a supervised learning problem, eliminating the need for rollouts or trajectories altogether. DPO starts by solving for the closed form solution of π_θ. Note that π_θ(a | x) * ∇log(π_θ(a | x)) = ∇π_θ(a | x) by the log-gradient trick, so that second summation sums to 1:

\( \sum ∇π_θ(a|x)[r(x,a) - \tfrac{1}{β}(\log π_θ(a|x) - \log π_{\text{ref}}(a|x)) - λ] - \tfrac{1}{β} = 0\)

Simplifying,

\(\beta r(x,a) - [\log π_θ(a|x) - \log π_{\text{ref}}(a|x)] - λ - 1 = 0\)

so:

\( \log π_θ(a|x) = \beta r(x,a) + \log π_{\text{ref}}(a|x) + λ - 1\)

and exponentiating gives:

\(π_θ(a|x) = \frac{π_{\text{ref}}(a|x) e^{\beta r(x,a)}}{C(x)}\)

where C(x) is the normalization constant ensuring probabilities sum to one.

Then, DPO moves in the reverse direction, substituting this definition of π_θ to express r:

\(\ r(x, a)=\frac{1}{β}[ \log C(x)+\log π_θ(a | x)−\log π_\text{ref}(a | x) ]\)

Previously we solved this expression, which we’ll use for MLE:

\(Pr[U(a_{+}) > U(a_{-})] = σ(r(a_{+}) − r(a_{-}))\)

Plugging in r:

\(\begin{align*} Pr[U(a_{+}) > U(a_{-})] &= \sigma\Big(\frac{1}{\beta}[(\log \pi_{\theta}(a_{+} | x) − \log \pi_{\text{ref}}(a_{+} | x)) − (\log \pi_{\theta}(a_{-} | x) − \log \pi_\text{ref}(a_{-} | x))]\Big) \end{align*}\)

Then we maximize likelihood:

\( \begin{align*} J(θ) &= ∑ \log \sigma\Big(\frac{1}{\beta}[(\log \pi_{\theta}(a_{+} | x) − \log \pi_\text{ref}(a_{+} | x)) − (\log \pi_{\theta}(a_{-} | x) − \log \pi_\text{ref}(a_{-} | x))]\Big) \end{align*}\)

That’s the final DPO objective. It can be optimized via standard supervised learning on your dataset. Choosing β controls the trade-off between imitation and divergence. But note that you no longer get additional signal beyond the dataset, unlike online PPO.

Offline Group Relative Policy Optimization (GRPO)

Now we reach the modern variant. GRPO was introduced by DeepSeek in 2024.

GRPO begins with the same pairwise setup as DPO. In fact, pairwise GRPO is mathematically identical to DPO, just rewritten.

To simplify notation, define:

\(z(x):=\frac{1}{β}[(\log \pi_{\theta}(a_{+} | x)−\log \pi_\text{ref}(a_{+} | x))−(\log \pi_{\theta}(a_{-} | x)−\log \pi_\text{ref}(a_{-} | x))]\)

Then the objective function for DPO becomes:

\(\begin{align*} J(θ) &= \sum \log σ(z(x)) \\[4pt] \nabla J(θ) &= \sum (1 - σ(z(x))) \, \nabla z(x) \\[4pt] \nabla z(x) &= \frac{1}{β}\left[\nabla \log \pi_{\theta}(a_{+} \mid x) - \nabla \log \pi_{\theta}(a_{-} \mid x)\right] \end{align*} \)

Define a shorthand:

\(w(x,a_{+},a_{-}):=\frac{1−σ(z(x))}{β}\)

Then:

\(∇J(θ)=∑ w(x,a_{+},a_{-})[∇ \log \pi_\theta(a_{+} | x) − ∇ \log \pi_\theta(a_{-} | x)]\)

Next, GRPO defines some synthetic reward function Ŕ:

\(\hat{R}(a) := \begin{cases} +w(x, a_{+}, a_{-}), & a = a_{+} \\[4pt] -w(x, a_{+}, a_{-}), & a = a_{-} \\[4pt] 0, & \text{otherwise} \end{cases}\)

Then we might rewrite the gradient as:

\(∇J(θ)=∑ \hat{R}(a) ∇ \log π_θ(a | x)\)

This is exactly the REINFORCE objective! This is the basic formulation for offline GRPO in the pairwise case. As you can see, all we did was make a few substitutions, but we didn’t fundamentally change the optimization. Thus, offline GRPO (pairwise) ≡ DPO ≡ REINFORCE in disguise.

Extending to Groups

So that’s the pairwise case. But the “group” in “group relative policy optimization” implies that you can have more than two responses. To be clear, with a₁ > a₂ > a₃, you could decompose that into pairs (a₁ > a₂, a₂ > a₃, …), but GRPO treats the group as a first-class citizen.

Here’s where theory gets shaky. In DPO, the weights w_i are strictly determined ±(1−σ(z))/β. GRPO merely observes that these weights satisfy ∑ w_i = 0 and generalizes: any set of scores with ∑ w_i = 0 is allowed.

The same supervised objective then applies:

\(∇J(θ)=∑ \hat{R}(a_{i}) ∇ \log π_θ(a_{i} | x)\)

where Ŕ uses the custom group weights.

To adapt this to online GRPO, we reuse the same idea as online PPO. After generating k responses for a prompt, compute their scores, center them (subtract the mean), and treat those as the rewards.

End of the Series

So is GRPO broken? Many people report that it works for them empirically. But it’s fair to say that GRPO’s theoretical foundations are weaker than many other methods. I’ll end this with a take I posted about GRPO:

Source: https://x.com/neelsomani/status/1976690361553895711

I hope to cover other RL/ML topics in future posts, but that concludes my blog series on reinforcement learning. Feedback is appreciated!