Inference Time Causal Probing in LLMs

2026-05-08

AI Summary

This work addresses the limitations of existing causal probing methods, which rely on auxiliary classifiers that often misalign with a model’s internal geometry and hinder generalizable, controllable interventions in large language models. The authors propose HDMI, a probe-free gradient-based intervention method that directly leverages the model’s native outputs to steer hidden states, employing a boundary optimization objective to enhance the probability of target continuations while suppressing source ones. They further extend this approach to text editing as LA-HDMI, which adjusts softmax embeddings via backpropagation to boost the generation probability of specified tokens without compromising fluency. HDMI constitutes the first framework enabling causal intervention grounded solely in the model’s own outputs and introduces a forward-looking hidden state optimization mechanism, achieving state-of-the-art performance on the LGD and CausalGym benchmarks while balancing intervention completeness and selectivity.

Abstract

Causal probing methods aim to test and control how internal representations influence the behavior of generative models. In causal probing, an intervention modifies hidden states so that a property takes on a different value. Most existing approaches define such interventions by training an auxiliary probe classifier, which ties the method to a specific task or model and risks misalignment with the model's predictive geometry. We propose Hidden-state Driven Margin Intervention (HDMI), a probe-free, gradient-based technique that directly steers hidden states using the model's native output. HDMI applies a margin objective that increases the probability of a target continuation while decreasing that of the source, without relying on probe classifiers. We further introduce a lookahead variant (LA-HDMI) for text editing that backpropagates through the softmax embeddings, modifying the current hidden state so that the likelihood of user-specified tokens increases in next token generations while preserving fluency. To evaluate interventions, we measure completeness (whether the targeted property changes as intended) and selectivity (whether unrelated properties are preserved), and report their harmonic mean as an overall measure of reliability. HDMI consistently achieves higher reliability than prior methods on the LGD agreement corpus and the CausalGym benchmark, across Meta-Llama-3-8B-Instruct, and Pythia-70M.

Problem

• causal probing
• hidden states
• probe classifier
• intervention
• large language models

Innovation

• causal probing
• probe-free intervention
• gradient-based steering
• hidden-state manipulation
• text editing