Skip to content

Interpretability for Neural Networks

NNsight (/ɛn.saɪt/) is a package for interpreting and manipulating the internals of deep learning models.

Get Started Docs Features About

What is NNsight?

NNsight is a Python library that enables interpreting and intervening on the internals of deep learning models. It provides a clean, Pythonic interface for:

  • Accessing activations at any layer during forward passes
  • Modifying activations to study causal effects
  • Computing gradients with respect to intermediate values
  • Batching interventions across multiple inputs efficiently

Originally developed by the NDIF team at Northeastern University, NNsight supports local execution on any PyTorch model and remote execution on large models via the NDIF infrastructure.

What does that look like?

Install NNSight:

pip install nnsight

Intervene:

from nnsight import TransformersModel

model = TransformersModel('openai-community/gpt2', device_map='auto', dispatch=True)

prompt = 'The Eiffel Tower is in the city of'

with model.trace(prompt):
    # Read a hidden state out of the forward pass
    hidden_states = model.transformer.h[-1].output.save()
    clean = model.output.logits[0, -1].argmax(dim=-1).save()

with model.trace(prompt):
    # Zero one layer at the last position and see what changes
    model.transformer.h[5].output[:, -1, :] = 0
    ablated = model.output.logits[0, -1].argmax(dim=-1).save()

print(hidden_states.shape)                       # torch.Size([1, 10, 768])
print(model.tokenizer.decode(clean))             # ' Paris'
print(model.tokenizer.decode(ablated))           # ','

Add remote=True to run the same code on a model too big for your machine, hosted by NDIF.