The Non-Linear Representation Dilemma: Is Causal Abstraction Enough for Mechanistic Interpretability?

In plain words

Researchers want to know what a neural network is really doing inside. One approach is causal abstraction: find a simpler, step-by-step program that seems to do the same job, then match parts of the network to the program's intermediate steps. The test is interventional. If you change one step in the program, the matched part of the network should change in the same way.

Most published work makes the matching step a straight-line map, because many people assume that the concepts a network uses sit along straight directions in its internal activations. This paper asks what happens if that restriction is removed and the matching can be any invertible function. The authors show that, under mild conditions, a match of this kind can be found for any network and any program, so the test passes even when it tells you nothing.

They back this up with experiments on a small network trained on an equality puzzle and on language models that solve a name-tracking task. Flexible matches fit even networks that were never trained, while straight-line matches fit less well. The lesson is that causal abstraction only says something useful when you also commit to how information is encoded.

TL;DR

The paper argues that causal abstraction, used with unrestricted alignment maps, is trivial: under countability, layerwise injectivity and output surjectivity, any network can be aligned with any algorithm. Experiments show that non-linear RevNet alignment maps reach near-perfect interchange intervention accuracy even on randomly initialised models, while linear rotations are weaker. The authors conclude that causal abstraction needs an explicit assumption about information encoding before its results mean anything. Results are author-reported.

Problem

Causal abstraction asks whether a neural network implements an algorithm: a directed graph of named intermediate variables that computes the task output §2.1. The network is abstracted by an algorithm if a map from hidden states to the algorithm's variables makes interventions on the network correspond to interventions on the algorithm §3. Distributed alignment search (DAS) estimates such a map from data, by training it so that the intervened network produces the output the algorithm would produce §3.3.

Most work on causal abstraction implicitly assumes the linear representation hypothesis, so the map is restricted to linear or orthogonal transformations §1, §5. The definition of causal abstraction does not require linearity §1, §3. The paper's concern is that once the map is unrestricted, the notion becomes vacuous. The authors argue that the usual evidence (a low interchange-intervention loss) then no longer separates a real implementation from a coincidence §1, §4.

Method

Definitions

• Algorithm. A deterministic causal model: a directed acyclic graph with input nodes, inner nodes and output nodes. Each node's value is a function of its parents' values, and an intervention overwrites a node with a constant §2.1.
• Network. A stack of layer functions with real-valued hidden states and a softmax output. An intervention fixes a set of neurons to given values §2.2.
• Abstraction map τ. Maps the network's hidden states to the algorithm's node values. The intervention map ω_τ sends a network intervention to an algorithm intervention when both induce the same set of states §3.
• τ-abstraction. The algorithm is a τ-abstraction of the network if τ is surjective, the algorithm's allowed interventions are exactly those induced by the network, and τ commutes with the interventions for all inputs [§3, definition after Beckers et al.]. A strong abstraction allows every intervention.
• Constructive abstraction. τ is block-wise: each node reads its own disjoint set of neurons, and some neurons are unused §3.
• Distributed abstraction. Before the block-wise step, a bijective layerwise alignment map φ rotates the hidden states into latent coordinates. Each node then reads its own disjoint set of latent coordinates §3.2. Interventions are applied in latent space: map to latent, overwrite the chosen coordinates, map back §3.2, app. B.
• Input-restricted abstraction. Interventions are limited to values that other input-restricted interventions can produce. This is the notion used in practice, and it is the one the paper's theorem is about §3.2.

Main theorem and its assumptions §4, app. F

The theorem: given a DNN and an algorithm that satisfy the five assumptions below, the algorithm is an input-restricted distributed abstraction of the DNN §4.

1. Countable inputs. The input set is countable. This holds for language (finite strings) and, approximately, for images at a fixed resolution [§4, app. F].
2. Input-injectivity. Every prefix of the network, up to each layer, is injective §4, app. F.
3. Strict output-surjectivity. For every layer, the suffix of the network followed by the output map is strictly surjective onto the outputs. The softmax bottleneck may break this in theory §4, app. F.
4. Matchable partial orders. Some assignment of single neurons to algorithm nodes respects the algorithm's order. Each layer keeps at least one neuron unused app. F.
5. The DNN solves the task. For every input, the network's argmax equals the task output §4, app. F.

Proof idea (for orientation, not a full derivation): build the alignment map layer by layer by induction. Three invariants are kept: each node's value can be read from the layer's latent coordinates, the output still matches the algorithm, and earlier node information survives in the unused coordinates app. F. Because the input set is countable and the reals are uncountable, injective per-coordinate maps exist, and uncountably many hidden states share any given output class, so the output can be kept app. F. The final map is made a bijection by a separate construction on the unused coordinates app. F. The existence proof does not say the map is learnable from finite data §7.

Alignment maps used in experiments §5

• Identity. φ is the identity, and the neuron partition is found greedily, adding one neuron at a time to each node's set §5.
• Linear. φ(h) = Q h with Q orthogonal §5. The code parametrises an orthogonal matrix, and its inverse is the transpose [code: das/Rotation_Model.py#L6-L50].
• Non-linear (RevNet). φ is a reversible residual network with K blocks and hidden size d_h §5. Each block splits the vector in half, adds a one-hidden-layer ReLU MLP F to the second half and a second MLP G to the first, and has an exact inverse RevNet.py.

Evaluation

The metric is interchange intervention accuracy (IIA): on held-out interventions, how often the intervened network's argmax equals the algorithm's output §5. Algorithm outputs on the intervention set are computed by running the algorithm with the same intervention §3.3.

Pseudocode: interchange intervention app. B, §3.2

Descriptive pseudocode of the intervention, mirroring the code's hook logic DAS_MLP.py:

# Interchange intervention on one layer  [§3.2] [app. B]
# base_hidden, source_hidden: (B, d) activations at the chosen layer
# mask: (B, d) bool, True where the latent coordinate comes from the source
def interchange(phi, phi_inverse, base_hidden, source_hidden, mask):
    # Map both activations into the latent space defined by phi
    base_latent = phi(base_hidden)                     # (B, d)
    source_latent = phi(source_hidden)                 # (B, d)
    # Copy the chosen latent coordinates from the source
    latent = where(mask, source_latent, base_latent)   # (B, d)
    # Return to the neuron basis; the network continues from here
    return phi_inverse(latent)                         # (B, d)

Pseudocode: DAS training step §3.3, app. I.1

The algorithm sets the target. The network is frozen and only φ is trained. The loss is cross-entropy between the intervened network output and the algorithm's output §3.3, the DAS loss.

# One DAS training step; only phi receives gradients  [§3.3]
# [code: das/DAS_MLP.py#L84-L117]
def das_step(dnn, layer, phi, phi_inverse, batch):
    # Hidden states at the chosen layer for base and source inputs
    base_hidden = hidden_at(dnn, layer, batch.base)     # (B, d)
    # Each intervened variable takes its columns from its own source
    source_hidden = hidden_at(dnn, layer, batch.sources)  # (B, d)
    patched = interchange(phi, phi_inverse, base_hidden,
                          source_hidden, batch.mask)
    # Finish the frozen forward pass from the patched layer  [§2.2]
    logits = forward_from(dnn, layer, patched)          # (B, classes)
    # Target: the algorithm's output under the same intervention
    return cross_entropy(logits, batch.algorithm_label)

Pseudocode: RevNet alignment map §5

# One RevNet block (additive coupling)  [§5] [code: das/RevNet.py#L63-L81]
# x: (B, d); F and G are one-hidden-layer ReLU MLPs on d/2 inputs
def revnet_block(x):
    # Split the vector into two halves
    x1, x2 = split_half(x)                              # (B, d/2) each
    y1 = x1 + F(x2)                                     # (B, d/2)
    y2 = x2 + G(y1)                                     # (B, d/2)
    return concat(y1, y2)                               # (B, d)

# Exact inverse, used to map patched latents back to neurons
def revnet_block_inverse(y):
    y1, y2 = split_half(y)                              # (B, d/2) each
    x2 = y2 - G(y1)                                     # (B, d/2)
    x1 = y1 - F(x2)                                     # (B, d/2)
    return concat(x1, x2)                               # (B, d)

Pseudocode: fitting loop with early stopping app. I.1, code

This follows train_test in the code, including how the best map is kept DAS.py. The code's exact bookkeeping is in Pitfalls.

# Fit the alignment map with the network frozen  [app. I.1]
def fit_alignment(dnn, layer, phi, train_set, eval_set):
    best_phi = copy_of(phi)
    best_eval = inf           # lowest eval loss seen (keeps the phi)
    counted_eval = inf        # eval loss counted for patience
    epochs_without_gain = 0
    for epoch in range(50):   # at most 50 epochs  [app. I.1]
        train_losses = []
        for batch in batches(train_set, size=6400):     # [app. I.1]
            train_losses.append(das_step(dnn, layer, phi,
                                         phi.inverse, batch))
        # The learning-rate schedule sees the mean training loss
        scheduler_step(mean(train_losses))              # [code: das/DAS.py#L291]
        eval_loss = mean_loss(dnn, layer, phi, eval_set)
        if eval_loss < best_eval:
            best_eval = eval_loss
            best_phi = copy_of(phi)
            if eval_loss + 0.001 < counted_eval:        # [app. I.1]
                counted_eval = eval_loss
                epochs_without_gain = 0
            else:
                epochs_without_gain += 1
        else:
            epochs_without_gain += 1
        if epochs_without_gain >= 5:                    # patience 5 [app. I.1]
            break
    return best_phi

Hyperparameters

Hierarchical equality task and MLP (paper values first; code values are noted where they differ or add detail).

Distributive law task (app. I.3): input 24-d, MLP width 24, three layers, Adam 1e-3, batch 1024, 1,048,576 training samples, 20 max epochs, early stopping after 3 epochs app. I.3. Alignment maps: Adam 1e-3, batch 6400, up to 50 epochs, early stopping after 5 epochs app. I.3.

Indirect object identification with Pythia (app. I.2; code values differ, see Paper vs code):

Results

Author-reported. Setups are the hierarchical equality task with a 3-layer MLP of width 16 (dataset and task in §5), and indirect object identification (IOI) with Pythia models §5, §6.

Qualitative results, from figures only: for the both-equality and left-equality algorithms, rotation IIA drops in the third layer of the equality MLP, and RevNet removes this drop fig. 2 §6. For Pythia-410m across training, 8-layer RevNet is near perfect at all checkpoints, including random initialisation, while rotation tracks the network's own accuracy more closely fig. 4 §6. Smaller fully trained Pythia models show slightly lower RevNet alignment §6.

Injectivity check: no collisions among 1,280,000 hidden states, repeated over 10 seeds app. H. The minimal distances are smaller for pairs that share equality-variable values, which the authors read as harder separation tab. 1 app. H.

Paper vs code

Official repository: densutter/non-linear-representation-dilemma, linked from the paper §1, app. A. The repository has no LICENSE file; the code is cited by path and line only. Commit checked: 7a8d538. Results.zip is committed and was not inspected.

1. IOI learning rate. Paper: 1e-4 app. I.2. The argparse default is 1e-4 das_llm.py. The run scripts override it: 5e-4 for the 410M main and name-split runs pythia410m.sh [code: run/pythia410m_name_split.sh#L8], and 1e-3 for the size sweep and seed runs pythia_size.sh pythia410m_seeds.sh.
2. IOI epochs. Paper: 2 epochs app. I.2. The size sweep runs 4 epochs pythia_size.sh; the seed runs run 3 [code: run/pythia410m_seeds.sh#L9].
3. IOI precision. Paper: float64 and lr 1e-3 for the 31M and 70M models only app. I.2. The size script uses float64 and lr 1e-3 for all three sizes it loops over, including 160M pythia_size.sh.
4. IOI warmup and early stopping. The paper states neither. The code uses a 500-step linear warmup pythia410m.sh. Early stopping counts validation checks run every 100 steps, not epochs, with patience 15 to 100 depending on the script das_llm.py [code: das/DAS.py#L239-L267]. The default improvement threshold is 0.01, not 0.001 das_llm.py.
5. MLP early stopping. Paper: stop after 3 epochs without improvement app. I.1. The code does not pass a threshold to the MLP training function, so the default of 1024 applies. The counter increments on each batch step, not each epoch, and the model kept is the one with the best eval accuracy Classification_Model.py [code: das/Classification_Model.py#L120-L129] [code: das/Classification_Model.py#L144-L149]. With batch 1024 and 1,048,576 samples, one epoch is 1024 steps, so the effective patience is about one epoch.
6. Which alignment map is kept. Paper: silent. The code returns the map with the lowest eval loss, not the last one DAS.py [code: das/DAS.py#L304-L306] DAS.py.
7. LR schedule input. Paper: silent. The schedule steps on the mean training loss, not the eval loss DAS.py.
8. LM loss. Paper: cross-entropy of the intervened output, with the output set given as the two name tokens §3.3 app. I.2. The default loss is cross-entropy over the full vocabulary at the final position [code: das/LLM_Model.py#L81-L85] DAS_LLM.py. Two alternative losses exist behind flags das_llm.py. IIA is a two-way comparison of the two name-token logits [code: das/DAS_LLM.py#L106].
9. LM intervention position. Paper: intervene at the middle layer fig. 4. The code intervenes only at the final token position of that layer's output, for base and source DAS_LLM.py [code: das/DAS_LLM.py#L91-L96].
10. Launch script. run/pythia410m.sh#L12 joins two flags into one token (--model_init_mode 0--layer_index 12). As written, argparse should reject the merged token as a non-integer, so the layer flag is never read. I did not run it.
11. Script name. The README points to run/pythia_seeds.sh; the file is run/pythia410m_seeds.sh.
12. Equality MLP setup. The code matches the paper on width, batch size, learning rate, sample counts, intervention variables and RevNet blocks and hidden size Standard_DAS.py [code: notebooks/Standard_DAS.py#L124-L126]. No discrepancy found there.

Pitfalls

• The theorem is existential. It proves that some alignment map exists. It does not say the map is learnable from finite data or that it generalises §7. Its proof uses uncountably many hidden states per output class to keep the output correct app. F, which is a form of extreme fitting §7.
• Assumptions that can fail. Injectivity holds almost surely for transformers at initialisation with real weights app. G. Finite precision and neural collapse may break it in practice §4. Output-surjectivity may fail because of the softmax bottleneck §4 §8.
• Low IIA can be capacity, not absence. For identity-of-first-argument on a trained MLP, RevNet reaches over 0.99 train IIA with 1,000 samples but under 0.55 with 1.28M samples. The authors read this as RevNet lacking the capacity to fit the required function app. H.
• Generalisation depends on names. Randomly initialised Pythia models reach high IIA only when train and eval share names; with disjoint names, no alignment was found even with complex maps app. I.2.
• Seed variance. Two of five Pythia-410m seeds failed, one with poor IIA and one that did not learn the task. The authors attribute the failures to suboptimal alignment training and did not tune further app. I.2.
• Plateaus can trigger early stopping. For randomly initialised Pythia-31M, the loss plateaus between about 4k and 6k steps before dropping app. I.2. The code's patience of 15 validation checks is 1,500 steps das_llm.py pythia410m.sh, so a plateau this long can stop the run before it converges.
• Linear maps can exceed network accuracy. The rotation map sometimes reaches higher IIA than the network's own accuracy. The authors note this may mean the map is learning the task itself §7.
• Label construction for equality data. For each sample, the label reuses the equality flag of the intervened variable, so the label reflects the source, not the base Dataset_Generation.py. This matches the paper's intent, but a re-implementation must do the same.
• Possible label inversion in the IOI data (unverified). In Preprocess_IOI_Data, when the final name appears before the subject, the code takes the subject as the true answer [code: das/Dataset_Generation.py#L257-L263]. That looks inverted for ABBA sentences, where the indirect object is the first introduced name. I did not fetch the dataset (fahamu/ioi), so I could not check the sentence format. Verify before using the LM labels.
• Compute. The equality MLP experiments took about 3 weeks on 10 CPU computers, because DAS on small MLPs ran faster on CPU app. J. The Pythia runs took about 30 GPU hours on one A100 80 GB, including tuning app. J.

Lineage

• Builds on: causal abstraction [arXiv:2301.04709] [arXiv:2303.02536] (DAS, Geiger et al.), and RevNet (Gomez et al., 2017) for the non-linear map. The paper cites them in its bibliography.
• Compared against: Boundless DAS and HyperDAS (named in app. C; no card here) and probing-based analyses §7.
• Superseded by: not known from this wiki.