🤖 AI

Bucketing the Good Apples: A Method for Diagnosing and Improving Causal Abstraction

This paper introduces a diagnostic method for causal abstraction in neural networks that partitions the input space into well-interpreted and under-interpreted regions based on interchange interventions, enabling researchers to pinpoint where interpretations fail and providing practical heuristics to refine and improve high-level causal hypotheses.

Original authors: Li Puyin, Jiyuan Tan, Ahmad Jabbar, Thomas Icard, Atticus Geiger

Published 2026-05-06

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Li Puyin, Jiyuan Tan, Ahmad Jabbar, Thomas Icard, Atticus Geiger

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 have a very smart, but slightly mysterious, robot chef. You want to understand exactly how it decides whether a dish is "delicious" or "not delicious." You have a theory (a hypothesis) about how it works: maybe you think it just checks if the salt is right.

To test your theory, you try a standard method called Causal Abstraction. This is like swapping ingredients between two different dishes to see if the robot's decision changes exactly as your theory predicts. If your theory is perfect, swapping the salt should always change the decision in a predictable way.

The problem is that in the real world, your theory is rarely perfect 100% of the time. Usually, you get a single score, like "78% accuracy." This score tells you the theory is "okay," but it doesn't tell you why it's failing or where it's failing. Is the robot confused about spicy food? Does it only work on soups but fail on salads? A single number hides all those details.

This paper introduces a new method called "Bucketing the Good Apples" to fix this. Instead of looking at the whole pile of fruit (all the dishes) and giving it one average score, the authors suggest sorting the fruit into different buckets based on how well your theory works for each one.

Here is how the method works, step-by-step, using simple analogies:

1. The Problem: The "Average" Score is Misleading

Imagine you are testing a theory that says, "The robot decides a dish is good if it has cheese."

Scenario A: You test it on 100 pizzas. Your theory works perfectly.
Scenario B: You test it on 100 salads. Your theory fails completely because salads don't usually have cheese.
The Old Way: You mix them all together and say, "The theory is 50% accurate." This is useless. It doesn't tell you that the theory is actually great for pizza and terrible for salad.

2. The Solution: Sorting into Buckets

The authors propose a four-step recipe to sort the inputs (the dishes) into "Good Buckets" (where the theory works) and "Bad Buckets" (where it fails).

Step 1: Pick the Reliable Dishes
First, only look at the dishes the robot actually got right. If the robot made a mistake on a pizza, we don't care about that pizza for this test; we only care about the ones where the robot was confident and correct. This ensures we are testing our theory, not the robot's basic ability to cook.

Step 2: Find the "Swap" Connection
Next, the researchers look at pairs of dishes. They ask: "If I take the 'cheese' part from Pizza A and put it into Salad B, does the robot's decision change in the exact way my theory says it should?"

If the swap works perfectly for a pair, they are "interchange-consistent."
If the swap fails, they are not.

Step 3: The "Bucketing" (The Core Idea)
Now, they draw a map. They connect dishes that work well together with lines.

If you have a group of dishes where every single pair works well when swapped, that group forms a tight cluster, or a "Bucket."
In our example, all the Pizzas would form one tight bucket (because swapping cheese between pizzas always works).
The Salads might form a messy, disconnected group where nothing swaps well.
The Discovery: The researchers found that the "Good Bucket" often reveals a hidden rule. For example, they might realize, "Oh! The theory 'Check for Cheese' only works when the dish is a Pizza. When the dish is a Salad, the robot is actually looking for something else (like 'freshness')."

Step 4: Teaching a Detective
Finally, they train a simple computer program (a classifier) to look at a new dish and guess which bucket it belongs to.

Does it belong in the "Pizza Bucket" (where the cheese theory works)?
Or the "Salad Bucket" (where the cheese theory fails)?
If the program can do this accurately, it proves that the difference between the buckets is real and structural, not just a random fluke.

Real-World Examples from the Paper

The authors tested this on three different "kitchens" (tasks):

The Logic Puzzle (Toy Task):
They had a model solving a math-like logic problem. Their initial theory was too simple. By bucketing the inputs, they realized the model was actually using a two-step process: first checking one condition, then another. The bucketing helped them "reverse engineer" the full, correct logic step-by-step, turning a vague guess into a precise map of how the model thinks.
The "Who Said What?" Game (Entity Binding):
Imagine a story with many characters. The model needs to remember who did what. The researchers thought the model just looked at the "position" of the name in the sentence.
- The Result: The bucketing showed the theory only worked for names at the very beginning or end of the story. For names in the middle, the model got confused. The "Bad Bucket" revealed that the model struggles with the middle of long lists, a nuance the old "average score" missed.
The Language Mixer (Factual Recall):
They tried to teach the model to separate "Language" (e.g., English vs. Spanish) from other facts (like "Country").
- The Result: The method showed that the model's "Language detector" was actually just a "English/Spanish detector." It worked great for English and Spanish inputs but failed completely for other languages. The bucketing revealed that the model hadn't learned a general "Language" concept; it had just memorized two specific languages.

The Big Takeaway

The paper argues that we shouldn't just accept a single "good/bad" score for how well we understand an AI. Instead, we should partition the AI's world.

By sorting the inputs into "Good Apples" (where the theory holds) and "Bad Apples" (where it breaks), we can:

Diagnose exactly where a theory fails.
Discover hidden variables or steps the model is using that we didn't know about.
Improve our theories by combining the insights from different buckets.

It turns the process of understanding AI from a vague "it's about 70% right" into a precise, constructive map of exactly how the machine thinks.

Technical Summary: Bucketing the Good Apples

Problem Statement

Causal abstraction is a rigorous framework for mechanistic interpretability that evaluates whether a high-level causal model ( $H$ ) accurately abstracts a low-level neural network ( $L$ ) under counterfactual interventions. This evaluation typically relies on Interchange Intervention Accuracy (IIA), a scalar metric summarizing the agreement between the low-level model's behavior under intervention and the high-level model's predictions.

The paper identifies a critical limitation in the current state of the art: IIA is a global summary statistic. While it indicates whether a hypothesis is generally faithful, it fails to reveal where in the input space the hypothesis succeeds or fails. Intermediate IIA scores are particularly difficult to interpret, offering little guidance on how to refine a hypothesis that is partially correct. Furthermore, existing workflows are largely evaluative rather than constructive; they can diagnose that a hypothesis is imperfect but do not provide a mechanism to improve it or discover missing intermediate variables.

Methodology

The authors propose a four-step diagnostic pipeline that shifts the unit of analysis from a global score to the structure of the input space. The core idea is to partition the input space into "well-interpreted" regions (target buckets) and "under-interpreted" regions based on pairwise interchange-intervention behavior.

1. Formalization: Interchange-Consistent Subsets

The method formalizes the concept of an Interchange-Consistent Pair: two inputs $i_1, i_2$ are consistent if swapping their internal representations yields outputs that match the high-level counterfactual predictions in both directions.

Interchangeability Graph: The input space is modeled as a graph where vertices are inputs and edges connect interchange-consistent pairs.
Target Buckets: A subset of inputs is considered a "target bucket" if it forms a clique (or a high-density $\gamma$ -quasi-clique) in this graph. Within these buckets, the abstraction is nearly perfect; between buckets, it fails.

2. The Four-Step Diagnosis Pipeline

The authors instantiate this framework into a practical recipe:

Specify Task and Input Space: Define a task where the low-level model $L$ performs reliably. The analysis is restricted to the subset of correct predictions ( $I_{correct}$ ) to isolate interpretability failures from task failures.
Find a Candidate Alignment: Identify a candidate alignment $\Pi$ between a high-level variable and a low-level feature using standard methods (e.g., Distributed Alignment Search (DAS), full-vector patching). This step often yields intermediate IIA scores that motivate further diagnosis.
Bucket the Input Space: Construct the interchangeability graph over $I_{correct}$ . Search for dense quasi-cliques (target buckets) where the abstraction holds. The remainder of the space forms the "other bucket." This step localizes specific failure modes.
Generalize and Characterize: Train a classifier to predict bucket membership for unseen inputs. The classifier uses either hand-labeled features or model-internal features (e.g., Sparse Autoencoder (SAE) activations) to determine if the discovered partition reflects a stable structural boundary rather than an artifact of the finite sample.

Key Contributions

Diagnostic Shift: The paper moves causal abstraction from a purely global evaluation to a diagnostic tool that reveals where an interpretation works and where it fails.
Constructive Refinement: By analyzing the structure of well-interpreted vs. under-interpreted regions, the method provides heuristics for improving interpretations. It can identify missing distinctions in a high-level hypothesis, discover unmodeled intermediate variables, and combine partial interpretations.
Recursive Hypothesis Discovery: The pipeline supports an iterative process where a coarse hypothesis is refined into a finer-grained causal model by recursively applying the bucketing procedure to the newly discovered failure modes.

Experimental Results

The authors validate the method across three settings of increasing complexity:

Toy Logic Task:
- Setup: A fine-tuned GPT-2-small model solves a logic task $o_5 = (o_1 \land o_2) \lor o_3$ .
- Result: Starting with a hypothesis containing only the output variable $o_5$ , the pipeline identified an alignment with intermediate IIA. Bucketing revealed that the alignment was faithful only when $o_1 \land o_2$ was false.
- Refinement: By promoting the missing intermediate variable $o_4 = o_1 \land o_2$ and re-running the diagnosis, the authors recursively recovered the full high-level causal hierarchy ( $o_1, o_2, o_3 \to o_4 \to o_5$ ), demonstrating that the model uses the intended factorization rather than an extensionally equivalent alternative.
Entity Binding Task:
- Setup: Using Gemma-2-2B-Instruct, the task involves retrieving entities from templated groups based on a queried group position ( $q_{group}$ ).
- Result: A positional hypothesis showed a U-shaped alignment curve (high faithfulness at sequence edges, low in the middle). Bucketing automatically isolated these regions without manual inspection.
- Generalization: Classifiers trained on internal SAE features successfully generalized this boundary to unseen inputs, achieving 98% IIA on the target bucket, confirming the distinction is encoded in the model's representations.
Entangled Factual Recall (RAVEL Benchmark):
- Setup: Isolating the LANGUAGE attribute from entangled attributes (e.g., COUNTRY, TIMEZONE) in Llama-3.1-8B using Multi-task DAS (MDAS).
- Result: Despite MDAS being designed for disentanglement, the global IIA was low (14.3%). Bucketing revealed that the learned subspace was not a general abstraction of "language" but rather collapsed the input space into coarse, language-specific clusters (e.g., English and Spanish buckets with high IIA, but low cross-language IIA).
- Insight: This exposed a structural mismatch where the representation preserved within-language structure but failed at cross-language interchangeability.

Significance and Claims

The paper claims that partitioning the input space is a crucial step toward more precise, constructive, and scalable mechanistic interpretability.

From Evaluation to Construction: The method transforms "imperfect" abstractions into actionable diagnostic objects, allowing researchers to move from post-hoc evaluation toward the constructive discovery of high-level hypotheses.
Scalability: The approach is shown to be broadly useful across different model sizes (fine-tuned and pretrained) and alignment methods (DAS, MDAS, patching).
Limitations: The authors modestly note that the method depends on a prior task construal and candidate abstraction (it is a structured program search, not unconstrained automatic interpretation) and currently focuses on single-variable hypotheses. They suggest future work should address multi-variable interactions and larger input spaces.

The paper concludes that by "bucketing the good apples," researchers can identify the specific conditions under which causal abstractions hold, thereby turning abstraction failures into evidence for refining and discovering the underlying mechanisms of neural networks.