i-machine-think / diagnnose

diagNNose is a Python library that facilitates a broad set of tools for analysing hidden activations of neural models.

diagNNose ·

Paper: https://arxiv.org/abs/2011.06819

Demo:

Documentation: https://diagnnose.readthedocs.io

This library contains a set of modules that can be used to analyse the activations of neural networks, with a focus on NLP architectures such as LSTMs and Transformers. In particular, it contains functionality for :

• Extracting activations from different types of (language) models and providing quick access to these stored activations.
• Training diagnostic classifiers (Hupkes et al., 2018) on extracted activations.
• Training control tasks (Hewitt & Liang, 2019) parallel to these diagnostic classifiers.
• Performing model-agnostic feature attributions (Murdoch et al., 2018) on a model.
• Running a broad linguistic suite of targeted syntactic evaluations on a language model.

🎉 diagNNose has been presented at BlackboxNLP 2020! The paper can be found here.

Documentation can be found at diagnnose.readthedocs.io.

Our library is officially registered with pip and can be installed by running pip install diagnnose. The preferred version of Python is ≥3.6. The required packages are stated in requirements.txt.

Quick Tour

The workflow of diagNNose is divided into several building blocks, that can be combined for various experiments.

We provide a few examples that demonstrate the library. An interactive and more extensive interface for these scripts is also provided in the form of a notebook:

Activation Extraction

The activations of a model can be extracted using an Extractor that takes care of batching and selecting activations of interest.

Fine-grained activation selection is possible by defining a selection_func, that selects an activation based on the current sentence index and corpus item.

from torchtext.data import Example

from diagnnose.config import create_config_dict
from diagnnose.corpus import Corpus
from diagnnose.extract import Extractor
from diagnnose.models import LanguageModel, import_model
from diagnnose.tokenizer.create import create_tokenizer

if __name__ == "__main__":
    config_dict = create_config_dict()

    tokenizer = create_tokenizer(**config_dict["tokenizer"])
    corpus: Corpus = Corpus.create(tokenizer=tokenizer, **config_dict["corpus"])
    model: LanguageModel = import_model(**config_dict["model"])

    def selection_func(w_idx: int, item: Example) -> bool:
        return w_idx == item.extraction_idx

    extractor = Extractor(
        model, corpus, selection_func=selection_func, **config_dict["extract"]
    )
    activation_reader = extractor.extract()

Research using diagNNose

• Jumelet, Zuidema & Hupkes (2019): Analysing Neural Language Models: Contextual Decomposition Reveals Default Reasoning in Number and Gender Assignment
• Weber, Jumelet, Bruni & Hupkes (2021): Language Modeling as a Multi-Task Problem

Citing

If you intend on using diagNNose for your research, please cite us as follows. Feel free to reach out, we'd love to help!

@inproceedings{jumelet-2020-diagnnose,
    title = "diag{NN}ose: A Library for Neural Activation Analysis",
    author = "Jumelet, Jaap",
    booktitle = "Proceedings of the Third BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP",
    month = nov,
    year = "2020",
    address = "Online",
    publisher = "Association for Computational Linguistics",
    url = "https://www.aclweb.org/anthology/2020.blackboxnlp-1.32",
    pages = "342--350",
}

diagnnose's Issues

Create array dim hints throughout code

Quite some functions don't have array dims as comments, this could become confusing very quickly

Investigate the "some weights were not initialized" warning on HF models

Some weights of BertForMaskedLM were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['cls.predictions.decoder.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.```

Remove indexing by key in ActivationReader

Indexing by "pos" should be sufficient, sentence "key" can't diverge from its position in the Corpus.

Allow labels to be added to corpus afterwards

Pass corpus to InitStates and use eos activations of that corpus

The "use eos activation" setup is rather cumbersome, and could be made much easier to use.

Remove config overwrite option

I think it is confusing and not that useful: requiring arguments to always be set in a json file is not that demanding.

Decrease coupling between config_dict and base classes

E.g.:

• https://softwareengineering.stackexchange.com/questions/40313/is-a-single-config-object-a-bad-idea
• https://softwareengineering.stackexchange.com/questions/335005/is-there-a-name-for-the-anti-pattern-of-passing-parameters-that-will-only-be
• https://sourcemaking.com/refactoring/replace-method-with-method-object

Create small suite of select functions

• every final token
• every n'th token
• every sen_id part of a provided list
• etc.

Create visualiser for ShapleyTensors

Add functionality for downloading Downstream corpora

Remove overrides

Doesn't add much and isn't a standard lib

Apply Factory pattern to LM creation

Port MaskedArray's to PackedSequences

Numpy arrays should only be used when working with sklearn classes.

Add Apache license

Investigate influence of using float32 vs. float64

Difference seems minor, but might be important for more finegrained classification.

Activations are stored as float32 as of now.

Skip extraction of sentences with no items from selection_func

Batchify CD

Labels no longer part of extraction

I decided that label dumping should not be inherently tied to activation extraction, and pulled it out of those classes.

In most cases this is no issue: labels can easily be created outside the extraction phase.

The only downside of this could perhaps be that in case a special selection_func is used during extraction it should match with the labels that are used later on, but this of course could also be part of the creation of those labels.

Perhaps it would be useful to create a label creation class, although it wouldn't involve much coding and could also be presumed not to be a part of this library.

Incorporate Structural Probes

Create subword averaging for transformer LMs

Improve lazy loading functionality

Add numpy compatibility next to torch tensors

Batchify LSTM and Corpus

As far as I am aware, the LSTM model implementation doesn't work with batches so far. Furthermore, the corpus might ne wrapped in an extra class as well to allow easy batch processing.

Reconsider GoogleLM support

Invokes a series of edge cases that muddy code clarity (such as a list of vocab_paths)

Make lib py3.6 able

Unify specification for activations

Throughout the project, different ways to specify the kind of activation are used, which sometimes makes awkward conversions necessary. See e.g.

• {type}_l{layer_num} for pickled activations and DCs, like e.g. hx_l2.pickle
• {type}{layer_num} for the extraction config, e.g. hx0
• {layer_num} -> {type} -> activation for activations dict

Allow decompositions to be created for a sentence directly

Right now the flow expects it to go from corpus extraction -> decomposition

This is annoying if you want to be able to quickly test a decomposition for a single sentence without having to extract the activations from an entire corpus

Incorprate HF transformer compatability

Allow downstream corpora to be downloaded automatically

Make init states part of model

Init state creation is now done separately from model initialisation, even though this is inherently tied to its respective model.

Activation indexing to tensor

Right now if multiple sentences are indexed the activations are concatenated instead of stored in separate dimensions. This means you'd still need the activation ranges to look up the exact activations of a certain sentence, which might be very useful.

If the indexing would return a tensor of the shape (num_sens, max_sen_len, hidden_size), with either nan or 0 as default values for sentences smaller than the max_sen_len it would still be easy to convert that back to a concatenated tensor.

Always save pickle dump sequence as 1 array

The dynamic dumping functionality is great for not clogging up RAM, but the data format we end up with can only be read in with our own ActivationReader.

It might be much more convenient to always write each activation type to just one array at the end of extraction. This way a user can simply load in the activations without having to use the utils of our library.

Add preproc_linzen code to downstream

Current setup expects format that does not follow from publicly available corpus

Recreate ContextualDecomposition from ShapleyTensor

Creating mul_contributions and assigning the relevance partitions there should in principle be sufficient, and serves as a good check-up for the quality of the contributions.

Make suppress_print toggleable

At least for simple_extract this might be useful

More control over activation extraction

For some experiments (including mine ;-)) it is useful or even necessary to have more control over which activations are being extracted. We could therefore try to add a filter function which returns true if an activation should be extracted based on its value, time step and arbitrary additional information.

Allow labels attribute name to be set for DCTrainer

Right now it always expects it to be labels

Consider using torchtext, torch.DataLoader, etc.

Unify vocab_path configuration

tokenizer / vocab_path, and its place in the config json

Use os.path.join instead of custom trim

No need to redefine this util

Input unk for GoogleLM

Should still be mapped to the char ids of the token itself! Unk is only relevant for the output softmax layer there

Check test.txt with uneven sen lens

Batchification not working properly yet

Add feature_mask to ShapleyTensor

Move example_dir stuff outside scripts folder

They serve a more important purpose than just being example scripts, as the argparsers are defined there as well.

example_dir would be better suited for actual example use cases with small corpora.

Consider using h5py instead of pickle

http://docs.h5py.org/en/stable/quick.html

pickle dumping leaves some nasty memory leaks. This is most notable in ActivationWriter.concat_pickle_dumps, where for some reason the activations variable is still stored in memory.

Use ApacheArrow for activation storing

Inspired by recent improvements to huggingface's transformer lib: https://twitter.com/Thom_Wolf/status/1272512974935203841

Create a model-agnostic attribution method

__torch_function__!

Incorporate Transformer LM tasks

Don't use defaultdict for w2i

Defaultdict's are kinda neat, but their major downside in my opinion is that unknown tokens are added to the dictionary. So d[unk_token] actually creates an entry for unk_token in d.

I therefore propose the following:

class W2I(dict):
    """ Provides vocab functionality mapping words to indices.

    Non-existing tokens are mapped to the id of an unk token that should
    be present in the vocab file.
    """
    def __init__(self, vocab_lines: List[str], unk_token: str = '<unk>') -> None:
        w2i: Dict[str, int] = {w.strip(): i for i, w in enumerate(vocab_lines)}
        self.unk_idx = w2i[unk_token]
        super().__init__(w2i)

    def __missing__(self, key: Any) -> Any:
        return self.unk_idx

This enables the w2i dictionary to be created at once from a list of tokens, and is not prone to the unk_token storage issue of defaultdicts.

Consider structuring arg config as structured dict

{
model: {
...
}
} etc