Inference Time Causal Probing in LLMs

arXiv (Cornell University)

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.

Authors 4

1. Sadegh Khorasani

   École Polytechnique Fédérale de Lausanne

   Affiliation as printed

   School of Computer and Communication Sciences , EPFL" Lausanne , Switzerland

2. Saber Salehkaleybar Aachen Leiden Institute of Advanced Computer Science

   Leiden University · University of Applied Sciences Leiden

   Affiliation as printed

   Leiden Institute of Advanced Computer Science (LIACS) , Leiden University , , Leiden , The Netherlands

3. Negar Kiyavash

   École Polytechnique Fédérale de Lausanne

   Affiliation as printed

   College of Management of Technology , EPFL , Lausanne , Switzerland

4. Matthias Grossglauser

   École Polytechnique Fédérale de Lausanne

   Affiliation as printed

   School of Computer and Communication Sciences , EPFL" Lausanne , Switzerland