🤖 machine learning

Code Correctness Signals in LLM Hidden States: Pre-Generation Probing and Repair Geometry

This paper demonstrates that code correctness is linearly decodable from the hidden states of a Qwen3-4B model before generation, while revealing through rigorous residualization controls that apparent "repair" signals are actually confounded by context rather than representing isolated comprehension features.

Original authors: Carlo Di Cicco

Published 2026-06-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Carlo Di Cicco

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 a large language model (LLM) as a brilliant, but slightly secretive, student taking a coding exam. This paper is like a detective story where researchers try to peek inside the student's "brain" (its hidden states) to see what it knows before it writes a single word of code, and to see if its brain changes in a specific way when it tries to fix a mistake.

The researchers used a specific student (a model called Qwen3-4B) and gave it 444 coding problems. Here is what they found, explained through simple analogies.

1. The "Pre-Game" Crystal Ball

The Question: Before the student even starts writing the answer, does their brain already "know" if they are going to get the problem right or wrong?

The Analogy: Imagine you are about to take a math test. You haven't written a single number yet. The researchers looked at the student's brain activity the exact moment they finished reading the question. They asked: "Can we predict a passing grade just by looking at how the brain is buzzing right now?"

The Finding: Yes.
They built a simple "detector" (a linear probe) that looked at the brain's final state after reading the prompt. It was incredibly accurate (about 93% accuracy) at predicting whether the code would pass the tests.

The Twist: They worried the detector was just cheating by noticing that harder problems have longer questions. So, they "subtracted" the effect of question length from the brain data. Even after removing the "length" clue, the detector still worked well (dropping slightly to 91% accuracy, but still far better than random guessing).
The Takeaway: The model's brain contains a clear signal about whether it can solve the problem before it even tries to generate the solution.

2. The "Fix-It" Direction

The Question: When the student fails a problem and tries to fix it, does their brain shift in a specific, consistent direction that signals "I am about to succeed"?

The Analogy: Imagine the student gets a question wrong. They get a hint (the error message) and try again. The researchers looked at the difference between the brain state of the failed attempt and the brain state of the repair attempt. They asked: "Is there a specific 'arrow' or direction in the brain's activity that points toward a successful fix?"

The Finding: It's complicated.

First Look: At first glance, yes! There was a clear "arrow" in the brain data. When the model successfully fixed a bug, the brain moved in a specific direction. When it failed to fix it, the brain moved differently. This looked like a "success signature."
The Second Look (The Reality Check): The researchers then asked, "Wait a minute. Is this 'arrow' actually about the model understanding the fix, or is it just reacting to the circumstances of the fix?"
- They noticed that successful fixes often happened when the original mistake was a specific type of error (like a runtime error) or when the code was shorter.
- When they mathematically "subtracted" these external circumstances (the context) from the brain data, the "success arrow" disappeared.
The Takeaway: The brain didn't have a special "I understand the fix" signal. Instead, the brain was just reacting to the context of the repair (e.g., "Oh, this is a short code with a runtime error, I know how to handle this"). The "success direction" was just a side effect of the situation, not a deep internal realization.

3. The "Magic Eraser" Tool

The paper introduces a method called Residualization. Think of this as a "Magic Eraser" for data.

How it works: If you think a signal (like "I know the answer") is real, but you suspect it's actually just caused by something else (like "the question was short"), you use the Magic Eraser to wipe away the "short question" part.
The Result:
- When they used the eraser on the Pre-Game signal, the signal stayed strong. (It was real).
- When they used the eraser on the Fix-It signal, the signal vanished. (It was an illusion caused by the context).

Summary

The paper teaches us two main lessons about how these AI models think:

They know before they speak: The model's brain holds a clear map of whether it will succeed or fail the moment it finishes reading the instructions.
Don't trust the "Fix" signal blindly: When a model tries to repair its own code, the changes in its brain aren't necessarily a sign of "learning" or "understanding." Often, they are just the model reacting to the specific details of the error message and the code length.

The authors emphasize that their most important contribution isn't just the results, but the honesty of their method. They used the same "Magic Eraser" tool to prove one thing was real and another was an illusion, showing that we must be careful not to mistake the context of a problem for the model's internal understanding.

Technical Summary: Code Correctness Signals in LLM Hidden States

Problem Statement

This work investigates the legibility of code correctness within the hidden states of Large Language Models (LLMs), specifically focusing on the Qwen3-4B-Instruct-2507 model. The study addresses two distinct temporal questions regarding code generation on 444 LiveCodeBench tasks:

Pre-Generation: Is the correctness of a model's first-attempt code already decodable from the hidden state at the final prompt token, before any output token is generated?
Repair Geometry: When a model attempts to repair a failed code generation, does the shift (delta) in its hidden state between the failing attempt and the repair attempt carry a stable, isolable geometric signature of whether the repair will succeed?

The study operates within the framework of mechanistic interpretability, utilizing linear probes and contrastive direction analysis to determine if internal representations encode verifiable properties like code correctness or repair success.

Methodology

Model and Data

Model: Qwen/Qwen3-4B-Instruct-2507 (36-layer decoder-only transformer, 2560-dimensional hidden state).
Dataset: 444 LiveCodeBench tasks (easy, medium, hard).
Hardware: Single NVIDIA RTX 4050 (6 GB VRAM).
Data Capture: Hidden states are extracted from the final prompt token of a single forward pass (before autoregressive generation) for the pre-generation analysis. For repair analysis, hidden states are captured for the first attempt (attempt 0) and the subsequent repair attempt (attempt 1).

Core Technique: Residualization

The paper employs residualization as a primary confound-control tool. This involves:

Fitting an Ordinary Least Squares (OLS) regression for each hidden state dimension against a candidate covariate (e.g., prompt length, repair context features).
Subtracting the fitted prediction from the hidden state dimensions to remove the linear contribution of that covariate.
Re-evaluating the signal (probe performance or geometric direction) on the residualized data.

This approach is used to distinguish between signals that are intrinsic to the model's representation versus those that are artifacts of input structure or context differences.

Experimental Pipelines

Pre-Generation Probe (Section 3):
- A linear probe (StandardScaler $\to$ PCA $\to$ Logistic Regression) is trained to predict pass/fail labels from the prompt-final hidden state.
- Layer Selection: Nested cross-validation (50 outer splits, 5-fold inner CV) is used to select the optimal layer and regularization strength ( $C$ ) to prevent test-set leakage.
- Baselines: Comparison against raw hidden states, residualized hidden states (removing prompt length), and a prompt-length-only baseline.
Repair Direction Analysis (Section 4):
- Delta Construction: For tasks where attempt 0 fails and attempt 1 exists, the hidden state delta is defined as $\delta = h_{\text{attempt 1}} - h_{\text{attempt 0}}$ .
- Contrastive Direction: The vector $v = \text{mean}(\delta_{\text{success}}) - \text{mean}(\delta_{\text{failure}})$ is computed.
- Validation: Two tests are applied to $v$ $v$ :
  - Magnitude: Norm of $v$ compared against label-shuffled nulls.
  - Inclination: Split-half cosine similarity between directions recovered from two halves of the data.
- Covariate Check: Before residualizing, statistical tests (Welch's t-test, Chi-squared) verify if covariates differ between success and failure groups. Residualization is only applied to covariates that show significant group differences.
- Covariates: $\Delta$ prompt_tokens, code_length_attempt_0, and failure_type.

Key Results

1. Pre-Generation Correctness is Linearly Decodable

Raw Signal: A linear probe trained on the prompt-final hidden state (Layer 29 selected via nested CV) achieves a leakage-free AUC of 0.931 ± 0.008 across 50 random splits.
Robustness to Prompt Length: After residualizing hidden state dimensions against prompt length, the probe performance remains high at 0.911 ± 0.010.
Baseline Comparison: A prompt-length-only baseline achieves an AUC of 0.754 ± 0.014. Non-linear models (Random Forest, Gradient Boosting, MLP) trained on prompt length alone do not exceed this baseline.
Conclusion: The hidden state contains a correctness signal that is robust and not reducible to prompt length, existing before a single output token is generated.

2. Repair Direction is a Contextual Correlate, Not an Isolated Feature

Raw Signal: In the raw deltas (236 clean transitions: 128 success, 108 failure), a contrastive direction exists. It passes magnitude ( $p < 0.001$ ) and split-half cosine ( $p < 0.001$ ) tests against label-shuffled nulls.
Covariate Imbalance: Statistical tests reveal that success and failure groups differ significantly in $\Delta$ prompt_tokens, initial code length, and failure type distribution ( $p < 0.001$ for all).
Residualization Effect: When the deltas are residualized against these three covariates:
- The magnitude of the direction drops from 9.76 to 3.72 (38% of original).
- The split-half cosine drops from 0.767 to 0.140.
- Neither test remains statistically significant (Magnitude $p=0.126$ , Split-half $p=0.942$ ).
Conclusion: The geometric direction observed in raw deltas is a correlate of repair success driven by the structure of the repair context (e.g., the specific error message and code length differences) rather than an isolated "repair-comprehension" feature in activation space.

3. Limitations on Multi-Step Repairs

The study finds that the pass rate collapses sharply after the first repair attempt (Attempt 1: 36.4%, Attempt 2: 6.7%).
Data for transitions beyond $0 \to 1$ (e.g., $1 \to 2$ ) is too sparse (only 13 successes for $1 \to 2$ ) to support meaningful contrastive direction analysis.

Contributions and Significance

The paper claims its contribution is as much methodological as empirical:

Diagnostic Honesty: The study demonstrates the utility of a single confound-control procedure (residualization) in two opposing directions.
- In the pre-generation case, residualization upholds the claim by showing the signal survives the removal of a major confound (prompt length).
- In the repair case, the same procedure overturns the initial interpretation by showing the signal collapses when context covariates are removed.
Methodological Rigor: The work emphasizes the necessity of checking covariate distributions before residualization. It argues that residualizing against covariates with equal group means is unnecessary and potentially misleading, whereas residualizing against covariates with differing means is essential to isolate latent features.
Empirical Findings:
- Confirms that code correctness is legible in the hidden state before generation begins.
- Refutes the hypothesis that a stable, isolated geometric direction for "repair success" exists in the raw hidden state shifts, attributing the observed signal to the differing contexts of the repair attempts.

The paper concludes that while LLMs encode rich information about code correctness, interpreting geometric shifts in hidden states requires careful disentanglement of the input context, as apparent "features" may simply be reflections of the prompt structure.