pyvene: A Library for Understanding and Improving P y T orch Models via Interventions (2024)

Abstract: 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 pyvene , an open-source Python library that supports customizable interventions on a range of different PyTorch modules. pyvene supports complex intervention schemes with an intuitive configuration format, and its interventions can be static or include trainable parameters. We show how 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‘.

## introduction when we intervene on a neural network, we make an in-place change to its activations, putting the model in a counterfactual state. this fundamental operation has emerged as a powerful tool for both understanding and improving models; interventions of various kinds are key to recent efforts in model robustness @xcite , model editing @xcite and steering @xcite , causal abstraction @xcite @xcite or activation patching @xcite , circuit finding @xcite , and knowledge tracing @xcite . as intervention-based techniques have matured, the need has arisen to run ever more complex interventions on ever larger models. currently, there is no unified and generic intervention-oriented library to support such research. existing libraries are often project-based (see implementations for wang figure 1 : an inference-time intervention @xcite on tinystories-33m. the model is prompted with "once upon a time there was a", and is asked to complete the story. we add a static word embedding (for "happy" or "sad") into the mlp output at each decoding step for all layers with a coefficient of 0.3. pyvene's complete implementation is provided. the original and intervened generations use greedy decoding. @xcite as examples) that lack extensibility and are hard to maintain and share, and current toolkits focus on single or non-nested interventions (e.g., ablation neurons in a single forward pass) and are often limited to interventions on transformers @xcite without natively supporting other neural architectures. some of these existing libraries @xcite @xcite can support complex interventions such as exchanging activations across multiple forward passes yet they require sophisticated knowledge and heavy implementations. to address these limitations, we introduce pyvene, an open-source python library that supports customizable interventions on different neural architectures implemented in pytorch. different from previous libraries @xcite @xcite @xcite , pyvene is interventionoriented. it supports complex interventions by manipulating or exchanging activations across multiple model forward runs while allowing these interventions to be shared with a serialization configuration file. specifically, pyvene has a number of advantages: all pyvene intervention schemes and models are serializable objects that can be shared through a public model hub such as hugging-face. in this paper, we provide two detailed case studies using pyvene as well: (1) we fully reproduce @xcite 's locating factual associations in gpt2-xl (figure 1 in the original paper) in about 20 lines of code, and (2) we show intervention and probe training with pyvene to localize gender in pythia-6.9b. pyvene is published through the python package index (pypi), foot_0 and the project site foot_1 hosts more than 20 tutorials that cover interventions at different levels of complexity with various model architectures from simple feed-foward models to multi-modal models. ## system design and architecture two primary components of pyvene are the intervenable configuration, which outlines which model components will be intervened upon, and the intervenable model, which decorates the original torch model with hooks that allow activations to be collected and overwritten. 3 here is a setup for performing a zero-out intervention (often called a zero ablation; @xcite ) on the 10th, 11th, and 12th dimensions of the mlp output for 3rd token embedding of layer 0 in gpt-2: import torch import pyvene as pv # built -in helper to get a huggingface model _ , tokenizer , @xmath0. create_gpt2 () # create with dict -based config @xmath1. intervenableconfig ({ " layer ": 0, " component ": " mlp_output " , " intervention_type ": pv . vanillaintervention }) # initialize model @xmath2. intervenablemodel ( pv_config , @xmath3, @xmath4, # the location to intervene at (3 rd token ) @xmath5: 3} , # the individual dimensions targetted @xmath6, # the intervention values @xmath7. zeros ( gpt2 . config . n_embd ) ) # sharing pv_gpt2 . save ( " ./ tmp / " , @xmath8 the model takes a tensor input base and runs through the model's computation graph modifying activations in place to be other values source. in this code, we specified source in the forward call. when source is a constant, it can alternatively be specified in the intervenableconfig. to target complete mlp output representations, one simply leaves out the subspaces argument. the final line of the code block shows how to serialize and share an intervened model remotely through a model hub such as huggingface. ## case study i: locating factual associations in gpt2-xl we replicate the main result in @xcite 's locating factual associations in gpt2-xl with pyvene. the task is to trace facts via interventions on fact-related datasets. following meng et al.'s setup, we first intervene on input embeddings by adding gaussian noise. we then restore individual states to identify the information that restores the results. specifically, we restore the transformer block output, mlp activation, and attention output for each token at each layer. for mlp activation and attention output, we restore 10 sites centered around the intervening layer (clipping on the edges). our with this configuration, we corrupt the subject token and then restore selected internal activations to their clean value. our main experiment is implemented with about 20 lines of code with pyvene: @xmath9, tokenizer , @xmath10. create_gpt2 (" gpt2 -xl ") @xmath11, @xmath12, " mlp_activation " , " attention_output " ]: for l in range ( gpt . config . n_layer ) : for p in range ( 7 ) : @xmath13, @xmath14, s , w ) @xmath15. intervenablemodel ( t_config , gpt ) _ , @xmath16, [ none ] + [ base ]* n_r , { " sources -> base ": ([ none ] + [[[ p ]]]* n_r , [[[0 , 1 , 2 , 3]]] + [[[ p ]]]* n_r ) }) @xmath17. embed_to_distrib ( gpt , outs . last_hidden_state , @xmath18. append ( { " stream ": s , " layer ": l , " pos ": p , " prob ": dist [0][ -1][7312]}) 4 case study ii: intervention and probe training with pythia-6.9b we showcase intervention and probe training with pyvene using a simple gendered pronoun prediction task in which we try to localize gender in hidden representations. for trainable intervention, we use a one-dimensional distributed alignment search (das; @xcite , that is, we seek to learn a 1d subspace representing gender. to localize gender, we feed prompts constructed from a template of the form "[john/sarah] walked because [he/she]" (a fixed length of 4) where the name is sampled from a vocabulary of 47 typically male and 10 typically female names followed by the associated gendered pronoun as the output token. we use pythia-6.9b @xcite in this experiment, which achieves 100% accuracy on the task. we then train our interventions and probes at the transformer block output at each layer and token position. for intervention training, we construct pairs of examples and train the intervention to match the desired counterfactual output (i.e., if we swap activations from an example with a male name into another example with a female name, the desired counterfactual output should be "he"). for linear probe training, we use activation collection intervention to retrieve activations to predict the pronoun gender with a linear layer. as shown in figure 3 , a trainable intervention finds sparser gender representations across layers and positions, whereas a linear probe achieves 100% classification accuracy for almost all components. this shows that a probe may achieve high performance even on representations that are not causally relevant for the task. ## limitations and future work we are currently focused on two main areas: ## conclusion we introduce pyvene, an open-source python library that supports intervention-based research on neural models. pyvene supports customizable interventions with complex intervention schemes as well as different families of model architectures, and intervened models are shareable with others through online model hubs such as huggingface. our hope is that pyvene can be a powerful tool for discovering new ways in which interventions can help us explain and improve models.