Package: causal-tracer
Causal TracerCausal trace plots for transformer language models.Demo:AboutThis library generates causal trace plots for transformer language models like Llama and GPT2, and should work with any decoder-only model on Huggingface. This library is based on causal tracing code fromROME, and broadly packages and improves on their excellent work. Thank you to these authors! There are some notable differences between the original ROME causal tracing code and this library, such as support for batch processing, automatic noise calculation, more processing options, and a slightly different API.Causal tracing is a technique to find which activations at which layers are causally important for the model to generate any given output. The way this works is by scrambing subject tokens, then slowly replacing activations in the scrambled computation graph and observe if replacing an activation gets the model closer to its original answer.For instance, if we prompt a languge model with "Rome is located in the country of", it will output "Italy". If we want to understand how the model generated that answer, we can scramble the tokens for "Rome" by adding gaussian noise so the model now sees gibberish instead, like "@#(* is located in the country of". Of course, after this scrambling, there's no way for the model to output "Italy" since the subject is just noise. However, we can take this corrupted computation graph and start replacing activations in it with the original uncorrupted activations, and see if the model starts outputting "Italy" again. If it does, we know that activation is important to the computation!For more info on causal tracing, check out the original ROME paper,Locating and Editing Factual Associations in GPT.Installationpip install causal-tracerBasic usageIf you're generating causal traces for a Llama-based model or GPT2, you don't need any further configuration.fromtransformersimportAutoModel,AutoTokenizerfromcausal_tracerimportCausalTracer,plot_hidden_flow_heatmapmodel=AutoModel.from_pretrained("gpt2-medium")tokenizer=AutoTokenizer.from-pretrained("gpt2-medium")tracer=CausalTracer(model,tokenizer)# perform causal tracing across hidden layers (residual stream) of the modelhidden_layer_flow=tracer.calculate_hidden_flow("The Space Needle is located in the city of",subject="The Space Needle",)# plot the resultplot_hidden_flow_heatmap(hidden_layer_flow)You can also generate causal traces of MLP layers or attention layers in the transformer.fromtransformersimportAutoModel,AutoTokenizerfromcausal_tracerimportCausalTracer,plot_hidden_flow_heatmapmodel=AutoModel.from_pretrained("gpt2-medium")tokenizer=AutoTokenizer.from-pretrained("gpt2-medium")tracer=CausalTracer(model,tokenizer)# perform causal tracing across MLP layers of the modelmlp_layer_flow=tracer.calculate_hidden_flow("The Space Needle is located in the city of",subject="The Space Needle",kind="mlp",window=10,)plot_hidden_flow_heatmap(mlp_layer_flow)# perform causal tracing across MLP layers of the modelattn_layer_flow=tracer.calculate_hidden_flow("The Space Needle is located in the city of",subject="The Space Needle",kind="attention",window=10,)plot_hidden_flow_heatmap(attn_layer_flow)When generating MLP or attention causal traces, it's you should typically set a window size. In the ROME paper, this is set to 10, which means the MLP or attention traces are replaced as a group and their impact is averaged to make it easier to see the impact of smaller changes.Batching and samplingBy default, causal traces will be calculated by scrambling the subject tokens with 10 different noise samples, and will run in batches of size 32. You can improve the quality of the causal trace by increasing the number of samples higher. Also, if you run out of RAM during processing, you can try decreasing the batch size.hidden_layer_flow=tracer.calculate_hidden_flow("The Space Needle is located in the city of",subject="The Space Needle",samples=50,batch_size=8,)Limiting patching for performanceRunning causal tracing can be slow as it requires a lot of passes through the model to generate a trace. You can get a speed-up by only calculating causal traces of certain layers, or only performing patching on subject tokens themselves. The results won't be complete if you do this, but depending on the use-case, that might be fine.hidden_layer_flow=tracer.calculate_hidden_flow("The Space Needle is located in the city of",subject="The Space Needle",start_layer=10,end_layer=15,patch_subject_tokens_only=True,)Custom layer configsIf you're using a model that isn't automatically detected by the library, you'll need to add aLayerConfigto tell CausalTracer where to findthe embeddings, attention, MLP, and hidden layers within the model. You can do this by creating aLayerConfigobject and passing it in when creating aCausalTracerobject.fromcausal_tracerimportCausalTracer,LayerConfigcustom_layer_config=LayerConfig(hidden_layers_matcher="h.{num}",attention_layers_matcher="h.{num}.attn",mlp_layers_matcher="h.{num}.mlp",embedding_layer="wte",)tracer=CausalTracer(model,tokenizer,layer_config=custom_layer_config)Note thathidden_layers_matcher,attention_layers_matcher, andmlp_layers_matcherare template strings, containg{num}in the middle. During processing,{num}will get replaced with the layer number. These strings correspond to the named modules of the transformer. You find all the named modules of a Pytorch model by runningmodel.named_modules().ContributingContributions are welcome! If you submit code, please make sure to add or update tests coverage along with your change. This repo uses Black for code formatting, MyPy for type checking, and Flake8 for linting.