🤖 AI

Inference Time Causal Probing in LLMs

This paper proposes Hidden-state Driven Margin Intervention (HDMI), a probe-free, gradient-based method that directly steers LLM hidden states using the model's native output to achieve more reliable causal interventions than existing classifier-dependent approaches, while also introducing a lookahead variant for text editing.

Original authors: Sadegh Khorasani, Saber Salehkaleybar, Negar Kiyavash, Matthias Grossglauser

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sadegh Khorasani, Saber Salehkaleybar, Negar Kiyavash, Matthias Grossglauser

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

The Big Picture: Fixing the "Black Box"

Imagine a Large Language Model (LLM) as a super-smart, but slightly opaque, chef in a kitchen. This chef can write amazing stories, but if you ask, "Why did you add salt to this soup?" the chef might not have a clear answer. They just "know" it tastes right.

Scientists want to understand how this chef decides things. Specifically, they want to know: Does the chef actually use the concept of "plural" (like "cats") to decide on the verb "run," or is that just a coincidence?

To test this, researchers use a technique called Causal Probing. It's like a "what-if" experiment. They want to say, "Okay, let's pretend the subject is plural instead of singular. Does the chef change the verb from 'runs' to 'run'?"

The Problem with Old Methods: The "Translator" Issue

Previously, to do this experiment, researchers had to hire a translator (called a "probe").

They would train this translator to read the chef's secret notes (the hidden states) and guess if the subject was singular or plural.
Then, they would use the translator's feedback to nudge the chef's notes to change the meaning.

The Flaw: This translator might not speak the chef's language perfectly. The translator might have its own rules for what "plural" looks like, which don't match how the chef actually thinks. It's like trying to fix a car engine using a manual written for a different brand of car. You might push the right button, but you could break something else because you don't understand the engine's true layout.

The New Solution: HDMI (The "Direct Nudge")

The authors propose a new method called HDMI (Hidden-state Driven Margin Intervention).

The Analogy: Instead of hiring a translator, HDMI talks directly to the chef's head.

The Goal: The chef is about to say "runs" (singular). You want them to say "run" (plural).
The Old Way: Ask a translator to find the "plural" button and push it.
The HDMI Way: HDMI looks at the chef's final decision-making process (the output). It calculates the exact mathematical "nudge" needed to make the word "run" slightly more likely and "runs" slightly less likely. It does this by looking at the difference (the margin) between the two words.

Why it's better:

No Translator Needed: It doesn't need to train a separate AI to understand the concept. It uses the model's own brain to figure out how to change the output.
Precision: Because it uses the model's own "geometry" (how it naturally arranges words), it aligns perfectly with how the model actually thinks.
Efficiency: It's a simple, fast calculation, like a quick tap on the steering wheel rather than a long detour.

The "Lookahead" Upgrade: LA-HDMI

The paper also introduces a fancy version called LA-HDMI (Lookahead HDMI) for text editing.

The Scenario: Imagine you are writing a story: "The cat was sleeping." You want to change it to "The cats were sleeping."

The Problem: If you just change "was" to "were," the sentence might break because the word before it ("The") or the word after it might need to change too to keep the grammar flowing.
The LA-HDMI Solution: This method doesn't just look at the current word; it looks ahead. It simulates the next few steps of the sentence while making the change.
- It sees that if it changes "was" to "were," it needs to change "The" to "The" (no change) but maybe adjust the article before the noun if the noun changes.
- It gently steers the hidden thoughts of the model before the word is even written, ensuring the whole sentence stays smooth and fluent, not just the one word you changed.

Think of it like a GPS that doesn't just tell you to turn left now, but calculates the whole route ahead so you don't get stuck in a dead end later.

Did It Work? (The Results)

The authors tested this on two major "gymnasiums" (datasets) designed to test grammar and logic:

LGD Agreement: Checking if subjects and verbs match (e.g., "he is" vs. "they are").
CausalGym: A complex set of grammar puzzles.

The Scorecard:
They measured two things:

Completeness: Did we successfully change the meaning? (e.g., Did "runs" become "run"?)
Selectivity: Did we accidentally break other things? (e.g., Did we change the tense or the meaning of the whole sentence by mistake?)

The Verdict:
HDMI consistently scored higher than the old methods. It was better at changing exactly what it was supposed to change without messing up the rest of the sentence. It worked well on different types of AI models (like Llama and Pythia), proving it's a robust tool.

Summary

Old Way: Use a separate tool to guess how to change the AI's mind (often inaccurate).
HDMI: Use the AI's own output to calculate the perfect nudge to change its mind (accurate and efficient).
LA-HDMI: Use HDMI with a "crystal ball" to edit text smoothly, ensuring the whole sentence flows naturally, not just the edited word.

The paper concludes that this "direct nudge" approach is a more reliable way to understand and control how AI models think and write.

Technical Summary: Inference Time Causal Probing in LLMs (HDMI)

1. Problem Statement

The paper addresses the limitations of current causal probing methods used to understand how Large Language Models (LLMs) utilize internal representations for prediction.

Correlational Probing Limitation: Traditional methods train auxiliary classifiers ("probes") on hidden states to detect linguistic properties (e.g., subject number). While these show what is encoded, they do not prove how the model uses that information to generate the next token.
Causal Probing Limitation: Existing causal approaches attempt to intervene on hidden states to alter properties (e.g., changing singular to plural) and observe changes in output. However, these methods typically rely on training a separate probe to determine the direction of intervention (e.g., using the probe's gradient to flip a label).
- Drawbacks: This introduces extra training costs, requires property-specific supervision, and risks misalignment. The probe's classification boundary may not coincide with the model's internal "predictive geometry," leading to interventions that do not effectively steer the model's generation or inadvertently alter unrelated properties.

The authors propose Inference-Time Causal Probing, a framework where interventions are applied entirely at inference time without retraining the model or training auxiliary probes.

2. Methodology: Hidden-state Driven Margin Intervention (HDMI)

The core contribution is HDMI, a probe-free, gradient-based technique that directly leverages the generative model's native output head to steer hidden states.

2.1 Core Mechanism

HDMI operates by defining a margin objective between a "source" token (the model's original preferred continuation) and a "target" token (the desired counterfactual continuation).

Objective: Maximize the logit margin $L(x) = \phi(x)_\tau - \phi(x)_\sigma$ , where $\phi(x)$ is the logit vector, $\tau$ is the target token index, and $\sigma$ is the source token index.
Intervention: The method performs gradient ascent on the hidden state $h_\ell(x)$ at a specific layer $\ell$ to increase the probability of the target token while decreasing the source token.
Computational Efficiency: For the final layer, the gradient is a closed-form matrix-vector product: $\nabla_{h_L} L(x) = W_U^\top (e_\tau - e_\sigma)$ . This avoids the need for iterative probe training.
Multi-token Extension: The method supports sets of tokens (e.g., multiple valid plural verbs) by summing logits over target sets ( $T^+$ ) and source sets ( $T^-$ ), defined as $L_{set}(x) = \sum \phi(x)_{\tau_i} - \sum \phi(x)_{\sigma_j}$ .

2.2 Lookahead HDMI (LA-HDMI) for Text Editing

To address text editing where a user provides a modified sequence, the authors introduce LA-HDMI.

Challenge: Standard autoregressive generation does not allow hidden states at step $t$ to be aware of future edits at $t' > t$ .
Solution: LA-HDMI introduces a "lookahead" mechanism that backpropagates gradients through the expected token embeddings.
- It decouples the forward pass (using low temperature $\beta_f$ to maintain fluency) from the gradient path (using high temperature $\beta_g$ to compute expected embeddings).
- The objective function $J_t(h_t)$ accumulates margin objectives for future edit positions ( $S_{max}$ steps ahead), adjusting the current hidden state $h_t$ to steer future generations toward the user's edits while preserving fluency in non-edited regions.

3. Evaluation Framework

The paper evaluates interventions using two primary metrics, following the framework of Canby et al. [2024]:

Completeness: Measures whether the intervention successfully shifts the targeted property (e.g., changing subject number) in the hidden state, verified by a separate validation probe.
Selectivity: Measures whether unrelated properties (nuisance variables, e.g., non-subject noun number) remain stable after intervention.
Reliability: The harmonic mean of Completeness and Selectivity, serving as the overall measure of intervention quality.

4. Key Results

The authors evaluated HDMI on Meta-Llama-3-8B-Instruct and Pythia-70M using the LGD agreement corpus and the CausalGym benchmark.

Performance: HDMI consistently achieves higher Reliability than probe-driven baselines (AlterRep, FGSM, PGD) across most tasks.
- In agreement suites (e.g., subject-verb agreement), HDMI often achieves perfect Completeness and high Selectivity.
- Unlike baselines, HDMI's performance is independent of the quality of an interventional probe, as it does not use one. Baselines like AlterRep and PGD showed significant performance drops on smaller models (Pythia-70M) when probe accuracy was lower.
Robustness: HDMI maintains consistent reliability across different model sizes and tasks, whereas probe-based methods showed mixed results sensitive to the specific property and model.
Text Editing: LA-HDMI successfully demonstrated the ability to steer generation toward user-specified edits (e.g., changing "girl" to "owl" or "was" to "were") while maintaining grammatical fluency and adjusting necessary context (e.g., changing "a" to "an").

5. Significance and Claims

The paper claims that HDMI represents a shift from probe-dependent to model-native causal probing.

Alignment: By using the model's own output head to define the intervention direction, HDMI ensures the perturbation aligns with the model's actual predictive geometry, avoiding the misalignment risks of external probes.
Efficiency: The method is computationally lightweight, requiring only gradient ascent on hidden states without the overhead of training auxiliary classifiers.
Generality: The approach naturally extends to multi-token properties and offers a mechanism for fluent text editing via lookahead gradients.

The authors acknowledge limitations, including sensitivity to hyperparameters (step size, temperature, horizon) and the potential attenuation of gradients in long-range text editing tasks. They suggest future work in adaptive layer selection and step-size scheduling to further improve selectivity without sacrificing completeness.