🤖 machine learning

Markovian Circuit Tracing for Transformer State Dynamic

This paper introduces Markovian Circuit Tracing (MCT), a diagnostic framework that demonstrates transformer activations on synthetic Hidden Markov Model tasks encode coarse state-transition structures, enabling state abstractions that significantly improve counterfactual prediction accuracy through activation patching.

Original authors: Abdullah X

Published 2026-05-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Abdullah X

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are watching a magician perform a trick. You see the cards they shuffle and the words they say (the visible emissions), but you can't see the secret order they are keeping in their head (the hidden states).

For a long time, researchers have tried to figure out how AI models (like Transformers) "think" by looking at their internal gears and wires. This paper, titled "Markovian Circuit Tracing" (MCT), proposes a new way to peek inside the magician's mind, specifically when the AI is trying to predict what happens next in a sequence.

Here is the breakdown of their experiment and findings, using simple analogies.

1. The Setup: A Controlled Magic Show

To test the AI, the researchers didn't use real-world data (like Shakespeare or news articles). Instead, they built a perfectly controlled, fake world based on a mathematical concept called a Hidden Markov Model (HMM).

The Game: Imagine a board game where a token moves around a circle of hidden rooms. You can't see which room the token is in, but you see a colored light flash every time the token moves.
The Rules: The researchers know the exact rules: how the token moves between rooms, how the lights flash, and exactly what the token should predict next.
The AI: They trained a tiny AI (a "Transformer") to play this game. It only sees the colored lights and has to guess the next light.

2. The Problem: Is the AI Actually "Thinking"?

When the AI guesses correctly, is it just memorizing patterns, or is it actually building an internal map of the hidden rooms?

The researchers wanted to know: Does the AI's internal "brain activity" (activations) contain a map of these hidden rooms?

3. The Method: "Markovian Circuit Tracing" (MCT)

They created a diagnostic pipeline called MCT. Think of it as a translator that tries to convert the AI's messy, high-dimensional brain signals into a simple list of "rooms" (states).

They tested this translator in four ways:

Belief Check: Can we read the AI's mind to see if it knows the probability of being in each room? (Like asking, "Am I 80% sure I'm in the Red Room?")
Map Check: If we force the AI to think it is in a specific room, does it act like it's actually in that room?
Transition Check: Does the AI's internal state change in a way that matches the rules of the game?
The "Patching" Test (The Magic Trick): This is the most important part. They took the AI's internal "state" (the translator's guess of which room it's in) and swapped it with a different state mid-game.
- Analogy: Imagine the AI is driving a car. Mid-drive, you reach in and swap the driver's mental map from "I am on the Highway" to "I am in the City." If the AI suddenly starts driving like it's in the city (slowing down, turning), then the internal map was real and useful.

4. The Results: Partial Success

The findings were "partial but consistent," meaning the AI did some things well and others poorly.

The AI is a Good Student: The AI learned the game very well. It predicted the next light almost as well as a perfect mathematician could.
The "Map" is Fuzzy: When the researchers tried to translate the AI's brain signals into specific "rooms," it wasn't a perfect 1-to-1 match.
- In easy games (where the lights clearly show which room you are in), the AI's internal map was quite clear.
- In hard games (where the lights are confusing or there are too many rooms), the AI's internal map was blurry.
The "Patching" Proof: This was the biggest win. When they forced the AI to use a specific internal state (a "centroid"), the AI's behavior changed exactly as the math predicted.
- The Result: The AI's behavior became much closer to the "perfect" mathematical target than when they used random states or wrong states. This proves the AI is using a specific internal structure to make decisions, even if that structure isn't a perfect copy of the game's rules.

5. The Big Takeaway

The paper concludes that Transformers do build internal "state" summaries when solving sequence problems, but these summaries are more like probabilistic beliefs (e.g., "I think I'm likely in Room A") rather than exact labels (e.g., "I am definitely in Room A").

Key Limitations Mentioned:

This was tested on tiny, synthetic models playing a fake game.
The "translator" they used was simple.
They cannot claim this works for complex real-world tasks like writing poetry or diagnosing diseases yet. The paper strictly limits its claims to this controlled, mathematical benchmark.

Summary Analogy

Imagine you are trying to figure out how a GPS works by watching a car drive.

Old way: You look at the wires and try to guess how the map is stored.
This paper's way: You build a tiny, perfect maze. You watch the car navigate it. Then, you reach into the car and swap the driver's mental map with a different one. If the car suddenly turns left instead of right, you know for a fact the driver was using a map, and you've successfully identified how that map influences the driving.

The paper says: "We built a perfect maze, swapped the map, and proved the driver was using a map. But we don't know if this works for driving on real highways yet."

Technical Summary: Markovian Circuit Tracing for Transformer State Dynamics

Problem Statement

Mechanistic interpretability has largely focused on analyzing local circuits, features, and attention heads to explain isolated behaviors in transformers. However, partially observed sequence tasks require a different computational paradigm: the model must maintain and update an internal summary (state) over time to inform future predictions. While Hidden Markov Models (HMMs) provide a mathematically rigorous framework for such tasks—where the optimal solution is a Bayesian belief vector over latent states—there is a lack of controlled benchmarks to test whether transformer activations actually encode these state-transition dynamics. Existing work often studies components (probing, patching, state extraction) in isolation without access to exact ground truth for latent states, belief vectors, and counterfactual targets simultaneously.

Methodology: Markovian Circuit Tracing (MCT)

The paper introduces Markovian Circuit Tracing (MCT), a diagnostic pipeline designed to test whether transformer activations contain coarse state-transition structure. The methodology operates within a controlled synthetic benchmark using six distinct HMM families with varying transition and emission complexities.

The Benchmark

The benchmark trains tiny causal transformers (2 layers, 128 width, ~275k parameters) on visible emissions from HMMs where the latent states, transition matrices, and emission probabilities are known exactly. This setup provides:

Exact Ground Truth: Latent state sequences, Bayesian belief vectors, and transition matrices.
Exact Targets: Bayes-optimal next-token predictions and forced-state counterfactual targets (the distribution of the next token if the latent state were forced to a specific value).

The MCT Pipeline

Given a transformer trained on HMM emissions, MCT performs the following steps:

State Extraction: A state extractor $g$ maps activation vectors $a_{\ell,t}$ to discrete candidate states $z_t$ . The reference extractor uses K-means clustering on residual activations.
Belief Recovery: A linear ridge probe maps activations to the exact Bayesian belief vector to test if belief information is present.
Transition Recovery: The extracted state sequence is used to estimate a transition matrix, which is compared against the true HMM transition matrix using row-wise KL divergence.
Markov Order Testing: Empirical models of order 0, 1, and 2 are fitted to the extracted states to evaluate the strength of the first-order Markov signal.
State Forcing (Intervention): This is the core causal test. The centroid of a recovered state cluster is patched into the model at a specific token position. The model's output is then compared to the exact HMM counterfactual target (the distribution resulting from forcing the latent state). The performance is measured by the KL divergence between the patched model's output and the exact target.

Key Results

Task Learning

Across 18 runs (6 HMM families $\times$ 3 seeds), transformers learned the prediction task effectively, achieving a mean excess loss over the Bayes-optimal predictor of only 0.0138.

State and Transition Recovery

Belief Information: Linear probes confirmed that belief information is recoverable from activations across all families, though the quality varies significantly by HMM family.
Transition Structure: Residual and PCA-based state abstractions recovered partial transition structure. The signal was strongest in persistent (sticky transitions) and easy-separable (low emission entropy) regimes.
Limitations: Recovery was weaker in ambiguous-emission and six-state regimes. Exact belief clustering remained a stronger upper bound than activation-derived clustering, suggesting the model represents a continuous belief-like object rather than a direct copy of the discrete latent label.

State Forcing (Causal Evidence)

The state-forcing intervention provided the clearest evidence of causal utility:

Performance: Patching a recovered-state centroid reduced the KL divergence to the exact counterfactual target from 0.1957 (unpatched) to 0.0532 on average.
Controls: This improvement significantly outperformed controls including wrong-state centroids, mean activations, random activations, shuffled-label centroids, and even true-state centroids in some aggregated metrics.
Implication: The recovered state abstraction is causally meaningful and closer to the model's internal predictive state than the sampled latent label in certain contexts.

Contributions

Controlled Benchmark: A unified framework providing exact latent labels, Bayesian beliefs, transition matrices, and counterfactual targets for testing transformer interpretability methods.
Diagnostic Pipeline (MCT): A reference pipeline for mapping activations to state abstractions, testing Markov order, and validating causal relevance via state forcing.
Empirical Findings: Evidence that small transformers learn near-Bayes predictors and that their residual activations contain partial, coarse-grained transition structure, particularly in regimes with strong transition signals.

Significance and Claims

The paper positions MCT not as a complete algorithm for recovering hidden states, but as a diagnostic framework for transformer state dynamics. The authors claim that:

MCT successfully detects partial predictive transition structure in synthetic settings.
It identifies where interpretability signals are strong (persistent/easy regimes) and where they are weak (ambiguous/high-entropy regimes).
State-forcing tests using exact HMM counterfactuals provide a rigorous method for evaluating the causal relevance of extracted state abstractions.

The work is modest in scope, acknowledging that the tasks are synthetic, models are tiny, and the reference extractor is simple. It does not claim that natural language reasoning possesses identical state structures, but rather establishes a baseline for testing whether activation-derived abstractions can expose coarse predictive dynamics against exact ground truth.