paper
active
2024
ⓒ 1
paper:doi-10-48550-arxiv-2403-07809

pyvene: A Library for Understanding and Improving PyTorch Models via Interventions

ByZhengxuan Wuⓘ·Atticus Geigerⓘ·Aryaman Arora·Jing Huangⓘ·Zheng Wangⓘ·Noah D. Goodmanⓘ+2 moreGoodfire, Harvard University + 2 more
DOI 10.48550/arxiv.2403.07809arXiv 2403.07809OpenAlex W4392781392
Causal TracingBackpack Language ModelsAmnesic ProbingCounterfactual Statetransformer architectureBoundless DASGender Representation in LLMsCausal ScrubbingGetter and Setter HooksCausal Structural ProbeInterchange Intervention Accuracy (IIA)Distributed Alignment SearchIntervenable ConfigurationInference-Time Intervention (ITI)+12 more

TL;DR

pyvene is an open-source Python library that unifies intervention-based research on PyTorch neural models by treating the intervention itself—rather than model surgery code—as the primitive abstraction, expressed in a serializable dict-based configuration that can be shared via HuggingFace. Prior libraries (BauKit, TransformerLens, nnsight, graphpatch, Transformer Debugger) either lack extensibility to recurrent and convolutional architectures or require sophisticated custom code for multi-source, cross-forward-pass interventions; pyvene resolves both limitations with Getter/Setter hooks that track state variables enabling intervention at arbitrary time steps in GRU and other recurrent models. The library ships with trainable intervention types including RotatedSpaceIntervention (Distributed Alignment Search), LowRankRotatedSpaceIntervention, and BoundlessRotatedSpaceIntervention, and reproduces Meng et al.'s factual-association localization result in GPT2-XL in approximately 20 lines of code. A second case study on Pythia-6.9B demonstrates that a 1D DAS intervention finds sparse, causally localized gender representations across layers, whereas a linear probe achieves near-100% classification accuracy almost everywhere—implying that high probe accuracy is insufficient evidence of causal relevance, and that trainable interventions provide a strictly more diagnostic test of whether a representation is mechanistically load-bearing for a behavior.

What to take away

  1. 1. pyvene introduces the IntervenableConfig and IntervenableModel abstractions, which express interventions as serializable dict-based configuration objects rather than runtime code hooks, enabling sharing through HuggingFace or any model hub.
  2. 2. Meng et al.'s factual-association localization result in GPT2-XL (a 48-layer model) is fully reproduced in approximately 20 lines of pyvene code, demonstrating the library's conciseness for complex causal tracing experiments.
  3. 3. On Pythia-6.9B, a trainable 1D Distributed Alignment Search (DAS) intervention finds sparse gender-representing subspaces across layers, while a linear probe achieves near-100% accuracy at almost every layer and token position—showing the two methods disagree sharply on which representations are causally relevant.
  4. 4. pyvene's Getter/Setter hook system records a state variable per hook so that interventions can be triggered at a specific time step in recurrent models (e.g., GRU with h_dim=32), a capability absent from all prior hook-based libraries including BauKit, nnsight, graphpatch, and Transformer Debugger.
  5. 5. The library supports parallel multi-source interchange interventions in which activations from multiple source forward passes are simultaneously patched into a base computation graph; in a GPT-2 experiment mixing tokens from 'The language of Spain' and 'The capital of Italy', 'Italian' appears in the top-5 output logits.
  6. 6. Serial multi-source interventions are also supported, where an activation patched into an intermediate source forward pass propagates into a subsequent base forward pass, and 'Italian' again appears in the top-5 logits after chaining two such interventions through GPT-2's residual stream.
  7. 7. To replicate the inference-time intervention of Li et al. (2023a) on TinyStories-33M, pyvene adds a static word embedding for 'happy' or 'sad' to the MLP output at every decoding step across all layers with a coefficient of 0.3, demonstrating per-decoding-step intervention support for generative LMs.
  8. 8. An open question raised by the gender-localization case study is whether the divergence between DAS intervention accuracy (IIA) and linear probe accuracy is a general phenomenon across tasks and model scales, or specific to simple grammatical features in large models like Pythia-6.9B.
  9. 9. To replicate the DAS training setup, a researcher should construct counterfactual pairs from a template ('John/Sarah walked because he/she') with 47 male and 10 female names, train a LowRankRotatedSpaceIntervention with low_rank_dimension=1 at each Transformer block output layer using cross-entropy loss against the gold counterfactual pronoun token.
  10. 10. pyvene is published on PyPI (pip install pyvene) and provides more than 20 tutorials spanning simple feed-forward networks, Transformers, recurrent models, and multi-modal models, with all code runnable on Google Colab.

Peer brief — for seminar discussion

Wu et al. introduce pyvene, an open-source Python library distributed via PyPI that reframes neural network intervention as a first-class, serializable abstraction rather than bespoke hook code. Where existing tools—BauKit, TransformerLens, nnsight, graphpatch, and OpenAI's Transformer Debugger—implement interventions as runtime code executed on a single forward pass, pyvene's IntervenableConfig and IntervenableModel classes allow multi-source, cross-forward-pass, parallel, and serial interventions to be specified declaratively, stored, and shared through HuggingFace. The library also handles recurrent architectures (GRU, convolutional models) by maintaining per-hook state variables that gate execution to a target time step—something vanilla hook approaches cannot do. Two case studies anchor the claims. First, Meng et al.'s causal tracing experiment localizing factual associations in GPT2-XL (48 layers) is reproduced in roughly 20 lines of pyvene code, validating the API's expressiveness against a well-known benchmark result. Second, gender localization in Pythia-6.9B (a 32-layer, ~7B-parameter model from the Pythia suite) using a 1D Distributed Alignment Search (DAS) intervention—pyvene's LowRankRotatedSpaceIntervention with low_rank_dimension=1—shows that DAS finds sparse causally relevant subspaces, whereas a linear probe achieves near-100% accuracy across virtually every layer and token position. The load-bearing implication is that probe accuracy is not a reliable proxy for causal relevance: a representation can be linearly decodable without being mechanistically load-bearing, and trainable causal interventions are needed to adjudicate. An alternative approach pyvene could have used for the gender localization study is Amnesic Probing (Elazar et al., 2020), which uses concept erasure rather than subspace rotation to test causal necessity. A critical reader would push back on the scope of the empirical validation: both case studies are replications or simple demonstrations rather than novel findings, so the library's claimed advantage—handling complex intervention schemes more easily than predecessors—is never stress-tested against a task that those predecessors genuinely fail on at scale. The paper implicitly predicts that the probe-versus-intervention divergence observed in Pythia-6.9B will generalize, motivating pyvene as infrastructure for systematically testing that hypothesis across architectures and tasks, but no such evidence is provided here.

Methods (8)

  • Amnesic Probing
    Behavioral explanation technique using amnesic counterfactuals by Elazar et al. 2020
  • Boundless DAS
    A variant of DAS implemented in pyvene via BoundlessRotatedSpaceIntervention, introduced by Wu et al. 2023
  • Causal Scrubbing
    Method by Chan et al. 2022 for rigorously testing interpretability hypotheses via interventions
  • Causal Structural Probe
    Probe method combining causal interventions and structural analysis, supported by pyvene's activation collection
  • Distributed Alignment Search
    The core method introduced in this paper: finds alignments between high-level causal variables and distributed neural representations via gradient descent.
  • Inference-Time Intervention (ITI)
    Method by Li et al. 2023a that adds static vectors to model activations at inference time to steer behavior
  • Interchange Intervention Training (IIT)
    Training technique that induces specific causal structures in neural networks by co-training with interchange interventions
  • Path Patching
    Method by Goldowsky-Dill et al. 2023 for localizing model behavior via targeted activation interventions

Frameworks (2)

  • Backpack Language Models
    LM architecture with sense vectors showing multiplication effects, illustrating custom intervention in pyvene
  • transformer architecture
    Neural network architecture based on attention, commonly used in large language models

Findings (5)

Claims (6)

Hypotheses (1)

Questions (2)

Original abstract (expand)

Interventions on model-internal states are fundamental operations in many areas of AI, including model editing, steering, robustness, and interpretability. To facilitate such research, we introduce $\textbf{pyvene}$, an open-source Python library that supports customizable interventions on a range of different PyTorch modules. $\textbf{pyvene}$ supports complex intervention schemes with an intuitive configuration format, and its interventions can be static or include trainable parameters. We show how $\textbf{pyvene}$ provides a unified and extensible framework for performing interventions on neural models and sharing the intervened upon models with others. We illustrate the power of the library via interpretability analyses using causal abstraction and knowledge localization. We publish our library through Python Package Index (PyPI) and provide code, documentation, and tutorials at https://github.com/stanfordnlp/pyvene.

Related work— refs + corpus + external arXiv

Cited / in-corpus / arXiv badges show which signals surfaced each row. Multi-source rows weighted higher.

+16 more

Similar preprints — Semantic Scholar

Cited by (3)