Inference Time Causal Probing in LLMs

TL;DR

This paper introduces HDMI, a probe-free, gradient-based causal probing method for LLMs that directly manipulates hidden states to test and control internal representations without task-specific classifiers.

Contribution

HDMI and LA-HDMI are novel interventions that improve reliability in causal probing by avoiding auxiliary classifiers and enabling targeted text editing in language models.

Findings

HDMI outperforms prior methods on LGD agreement corpus and CausalGym benchmark.

HDMI achieves higher reliability in changing targeted properties while preserving unrelated ones.

LA-HDMI effectively modifies token likelihoods for text editing while maintaining fluency.

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.