Package: codebook-features
Codebook FeaturesAlex Tamkin, Mohammad Taufeeque and Noah D. Goodman: "Codebook Features: Sparse and Discrete Interpretability for Neural Networks", 2023.[arXiv]Codebook Features is a method for training neural networks with a set of learned sparse and discrete hidden states, enabling interpretability and control of the resulting model.Codebook features work by inserting vector quantization bottlenecks calledcodebooksinto each layer of a neural network. The library provides a range of features to train and interpret codebook models, including by analyzing the activations of codes, searching for codes that activate on a pattern, and performing code interventions to verify the causal effect of a code on the output of a model. Many of these features are also available through an easy-to-use webapp that helps in analyzing and experimenting with the codebook models.InstallationPyPIInstall from PyPI to directly use the library:pip install codebook-featuresSource CodeInstall from source code if you plan to modify part of the code or contribute to the library:git clone https://github.com/taufeeque9/codebook-features
cd codebook-features
pip install -e .For development mode, we recommend using Poetry:poetry installUsageTraining a codebook modelWe adapt therun_clm.pyscript from HuggingFace to train/finetune conventional models or the codebook models. We also use thehydralibrary for configuration management of the training scripts. The default config for training codebooks is available incodebook_features/config/main.yaml. The hydra syntax can be used to override any of the default config values, which includes arguments for the codebook model and arguments inherited from Huggingface'sTrainingArguments. For example, to train a codebook model using gpt2-small on the wikitext dataset, run:python -m codebook_features.train_codebook model_args.model_name_or_path=roneneldan/TinyStories-1M 'data_args.dataset_name=roneneldan/TinyStories'Interpretability WebApp for Codebook ModelsOnce a codebook model has been trained and saved on disk, we can use the interpretability webapp to visualize the codebook. First, we need to generate the relevant cache files for the codebook model that is required for the webapp. This can be done by running the scriptcodebook_features/code_search_cache.py:python -m codebook_features.code_search_cache --orig_model_name <orig name/path of model> --pretrained_path <path to codebook model> --dataset_name <dataset name> --dataset_config_name <dataset config name> --output_base_dir <path to output directory>Once the cache files have been generated, we can run the webapp using the following command with the base output directory used in the above command:python -m streamlit run codebook_features/webapp/Code_Browser.py -- --cache_dir <path to the base cache directory>Code InterventionTo control a network, one caninterveneon codes by causing them to always be activated during the forward pass. This can be useful to influence the sampled generations, e.g., to cause the network to discuss certain topics. For a general tutorial on using codebook models and seeing how you can perform code intervention, please see theCode Intervention Tutorial.Guide to the codebase [click to expand]Codebook Modelcodebook_features/modelsis the main module used to define codebooks. It has the following classes:CodebookLayer: defines atorch.nn.Modulethat implements the codebook layer. It takes in arguments likenum_codes,dim,snap_fnkcodesthat define the codebook. It provides various functionalities including logging methods, hook function that can disable specific codes during inference, etc.GroupCodebookLayer: defines atorch.nn.Modulethat implements a group of codebook layer each of which are applied to a different part of the input vector. This is useful for applying a group of codebooks on the attention head outputs of a transformer model.CodebookWrapper: is an abstract class to wrap a codebook around anytorch.nn.Module. It takes in themodule_layer,codebook_cls, and arguments for the codebook class to instantiate the codebook layer. The wrapper provides asnapboolean field that can be used to enable/disable the codebook layer.TransformerLayerWrapper: subclassesCodebookWrapperto wrap a codebook around a transformer layer, i.e. a codebook is applied on the output of the a whole transformer block.MLPWrapper: subclassesCodebookWrapperto wrap a codebook around an MLP layer, i.e. a codebook is applied on the output of the MLP block.CodebookModelConfig: defines the config to be used by a codebook model. It contains important parameters likecodebook_type,num_codes,num_codebooks,layers_to_snap,similarity_metric,codebook_at, etc.CodebookModel: defines the abstract base class for a codebook model. It takes in a neural network model through themodelargument and the config through theconfigargument and return a codebook model.GPT2CodebookModel: subclassesCodebookModelto define a codebook model specifically for GPT2.GPTNeoCodebookModel: subclassesCodebookModelto define a codebook model specifically for GPTNeo.GPTNeoXCodebookModel: subclassesCodebookModelto define a codebook model specifically for GPTNeoX.HookedTransformerCodebookModel: subclassesCodebookModelto define a codebook model for any transformer model defined using theHookedTransformerclass oftransformer_lens. This is mostly while interpreting the codebooks while the other classes are used for training the codebook models. Theconvert_to_hooked_model()function can be used to convert a trained codebook model to aHookedTransformerCodebookModel.Codebook TrainingThecodebook_features/train_codebook.pyscript is used to train a codebook model based on a causal language model. We use therun_clm.pyscript provided by the transformers library for training. It can take in a dataset name available in thedatasetslibrary or a custom dataset. The default arguments for the training script is available incodebook_features/config/main.yaml. The hydra syntax can be used to override any of the default config values.TokFSM ExperimentThecodebook_features/train_fsm_model.pyscript provides an algorithmic sequence modeling task to analyse the codebook models. The task is to predict the next element in a sequence of numbers generated using a Finite State Machine (FSM). Thetrain_fsm_model/FSMclass defines the FSM by taking in the number of states throughN, the number of outbound edges from each state throughedges, and the base in which to represent the state usingrepresentation_base. Thetrain_fsm_model/TokFSMDatasetclass defines an iterable torch dataset using the FSM that generates the dataset on the fly. Thetrain_fsm_model/TokFSMModelTrainerprovides additional logging feature specific to the fsm models like logging the transition accuracy of a model.Thecodebook_features/train_fsm_model.pyscript can be used to train a codebook model on the TokFSM dataset. The syntax for the arguments and training procedure is similar to thetrain_codebook.pyscript. The default arguments for the training script is available incodebook_features/config/fsm_main.yaml.For tutorials on how to use the library, please see theCodebook Features Tutorials.Citation@misc{tamkin2023codebook,title={Codebook Features: Sparse and Discrete Interpretability for Neural Networks},author={Alex Tamkin and Mohammad Taufeeque and Noah D. Goodman},year={2023},eprint={2310.17230},archivePrefix={arXiv},primaryClass={cs.LG}}