Skip to content

arushisharma17/NeuroX

 
 

Repository files navigation

NeuroX Toolkit

NeuroX-Workflow

1. Clone NeuroX

cd <your project dir>    #eg. cd /work/LAS/jannesar-lab/arushi/LatentConceptAnalysis

# Clone the NeuroX repository
git clone https://github.com/arushisharma17/NeuroX.git
cd NeuroX

2. Creating the environment

a. For pronto users only: Set Up NeuroX on Pronto using micromamba

sbatch pronto_setup.sh /path/to/project/dir

Important: Please not that if you ran setup.sh while installing CodeConceptNet repo, you do not need to install NeuroX again. The setup script already does that for you. It should be located at your $PROJECTDIR/ . Its an editable install, so you can cd into $PROJECTDIR/NeuroX and git pull to get latest changes.

For further information about pronto: Pronto Documentation

b. For other (non-pronto) users: Create a python virtual environment

python3 -m venv neurox-env
source neurox-env/bin/activate
pip install --upgrade pip

# Create a virtual environment
python3 -m venv neurox-env

# Activate the environment
source neurox-env/bin/activate

# Upgrade pip
pip install --upgrade pip

pip install -e . #Installs package in editable mode

3. Explore functionality of this toolkit: (For Testing and Educational Purposes)

We have created a temp directory where we store all relevant files for testing purposes.

a. If you are on pronto: you can use interactive mode to run scripts.

srun --time=01:00:00 --nodes=1 --cpus-per-task=8 --partition=gpu-interactive --gres=gpu:1 --pty /usr/bin/bash
cd temp/
./test_run_interactive.sh /path/to/project/dir   #eg. /work/LAS/jannesar-lab/arushi/LatentConceptAnalysis

Now you can run the code that you normally put in your shell script here directly. You will need to activate the environment and load modules too. Pronto Interactive Computing Guide

b. If you are not on pronto: can use following colab notebook as a reference.

You can run code interactively in colab or on your own machine after setting up the environment.

Add link to notebook

After running extract_activations.py

/path/to/your/project_directory/
└── NeuroX/
    └── temp/
        ├── config.json
        ├── input_data/
        │   ├── <dataset1>/                # e.g., test, java, cuda
        │   │   └── <input files>          # e.g., test.in, test.label
        │   ├── <dataset2>/                # e.g., cuda
        │   │   └── <input files>          # e.g., cuda.in cuda.label
        │   └── <dataset3>/                # e.g., java
        │       └── <input files>          # e.g., java.in java.label 
        ├── outputs/
        │   ├── <dataset1>/                # e.g., test
        │   │   ├── <model1>/              # e.g., microsoft-codebert-base
        │   │   │   └── Activations/
        │   │   │       ├── layer0/
        │   │   │       ├── layer1/
        │   │   │       └── layerN/
        │   │   └── <model2>/              # e.g., bert-base-cased
        │   │       └── Activations/
        │   │           ├── layer0/
        │   │           ├── layer1/
        │   │           └── layerN/
        │   └── <dataset2>/                # e.g., java
        │       ├── <model1>/              # e.g., microsoft-codebert-base
        │       │   └── Activations/
        │       │       ├── layer0/
        │       │       ├── layer1/
        │       │       └── layerN/
        │       └── <model2>/              # e.g., bert-base-cased
        │           └── Activations/
        │               ├── layer0/
        │               ├── layer1/
        │               └── layerN/
        └── test_run_interactive.sh

Tests Runner API Reference

NeuroX provide all the necessary tooling to perform Interpretation and Analysis of (Deep) Neural Networks centered around Probing. Specifically, the toolkit provides:

  • Support for extraction of activation from popular models including the entirety of transformers, with extended support for other models like OpenNMT-py planned in the near future
  • Support for training linear probes on top of these activations, on the entire activation space of a model, on specific layers, or even on specific set of neurons.
  • Support for neuron extraction related to specific concepts, using the Linear Correlation Analysis method (What is one Grain of Sand in the Desert? Analyzing Individual Neurons in Deep NLP Models.). The toolkit can extract either a local ranking of neurons important to a particular target class, or a global ranking of neurons important to all the target classes.
  • Support for ablation analysis by either removing or zeroing out specific neurons to determine their function and importance.
  • Support for subword and character level aggregation across a variety of tokenizers, including BPE and all tokenizers in the transformers library.
  • Support for activation visualization over regular text, to generate qualitative samples of neuron activity over particular sentences.

A demo using a lot of functionality provided by this toolkit is available.

demo picture

Getting Started

This toolkit requires and is tested on Python versions 3.6 and above. It may work with older Python versions with some fiddling, but is currently not tested nor supported. The easiest way to get started is to use the published pip package:

pip install neurox

Manual Installation

If you wish to install this package manually (e.g. to modify or contribute to the code base), you can clone this repository into a directory of your choice:

git clone https://github.com/fdalvi/NeuroX.git

Create and activate a new virtual environment for the toolkit (This step can be skipped if you manage your environment in other ways like Conda or System-level installations):

python -m venv .neurox-env
source .neurox-env/bin/activate

Install the dependencies required to run the toolkit:

pip install -e .

Sample Code

A Jupyter notebook with a complete example of extracting activations from BERT, training a toy task, extracting neurons and visualizing them is available in the examples directory for a quick introduction to the main functionality provided by this toolkit.

Documentation

API Reference contains an API reference for all of the functions exposed by this toolkit. Primarily, the toolkit's functionality is separated into several high-level components:

  • Extraction
  • Data Preprocessing
  • Linear Probing
  • Neuron extraction and interpretation
  • Neuron cluster analysis
  • Visualization

Citation

Please cite our paper published at AAAI'19 if you use this toolkit.

@article{dalvi2019neurox,
  title={NeuroX: A Toolkit for Analyzing Individual Neurons in Neural Networks},
  author={Dalvi, Fahim
    and Nortonsmith, Avery
    and Bau, D Anthony
    and Belinkov, Yonatan
    and Sajjad, Hassan
    and Durrani, Nadir
    and Glass, James},
  journal={Proceedings of the AAAI Conference on Artificial Intelligence (AAAI)},
  year={2019}
}

Planned features

  • Pip package
  • Support for OpenNMT-py models
  • Support for control tasks and computing metrics like selectivity
  • Support for attention and other module analysis

Publications

About

A Python library that encapsulates various methods for neuron interpretation and analysis in Deep NLP models.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 98.2%
  • Shell 1.8%