nnsight-remote-interpretability

nnsight: ニューラルネットワーク内部への透過的アクセス

nnsight（/ɛn.saɪt/）は、研究者があらゆる PyTorch モデルの内部を解釈・操作できるようにするものです。小規模なモデルはローカルで、70B 以上の大規模モデルは NDIF を介してリモートで、同じコードを実行できるという独自の機能を備えています。

GitHub: ndif-team/nnsight（730+ stars） 論文: NNsight and NDIF: Democratizing Access to Foundation Model Internals（ICLR 2025）

主な価値提案

一度書けばどこでも実行可能: 同じ解釈可能性のコードが、GPT-2 のローカル実行でも Llama-3.1-405B のリモート実行でも動作します。remote=True を切り替えるだけです。

# Local execution (small model)
with model.trace("Hello world"):
    hidden = model.transformer.h[5].output[0].save()

# Remote execution (massive model) - same code!
with model.trace("Hello world", remote=True):
    hidden = model.model.layers[40].output[0].save()

nnsight を使うべき場合

次のような場合に nnsight を使用してください:

• ローカルの GPU には大きすぎるモデル（70B、405B）で解釈可能性の実験を行う
• 任意の PyTorch アーキテクチャ（transformers、Mamba、カスタムモデル）を扱う
• 複数トークンの生成に対する介入を行う
• 異なるプロンプト間でアクティベーションを共有する
• 再実装なしにモデルの内部へ完全にアクセスする

次のような場合は代替手段を検討してください:

• モデル間で一貫した API が必要 → TransformerLens を使用
• 宣言的で共有可能な介入が必要 → pyvene を使用
• SAE を学習している → SAELens を使用
• 小規模なモデルだけをローカルで扱う → TransformerLens の方がシンプルな場合があります

インストール

# Basic installation
pip install nnsight

# For vLLM support
pip install "nnsight[vllm]"

リモートの NDIF 実行には、API キーを取得するために login.ndif.us でサインアップしてください。

基本概念

LanguageModel ラッパー

from nnsight import LanguageModel

# Load model (uses HuggingFace under the hood)
model = LanguageModel("openai-community/gpt2", device_map="auto")

# For larger models
model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto")

トレースコンテキスト

trace コンテキストマネージャーは遅延実行を可能にします。操作は計算グラフに収集されます:

from nnsight import LanguageModel

model = LanguageModel("gpt2", device_map="auto")

with model.trace("The Eiffel Tower is in") as tracer:
    # Access any module's output
    hidden_states = model.transformer.h[5].output[0].save()

    # Access attention patterns
    attn = model.transformer.h[5].attn.attn_dropout.input[0][0].save()

    # Modify activations
    model.transformer.h[8].output[0][:] = 0  # Zero out layer 8

    # Get final output
    logits = model.output.save()

# After context exits, access saved values
print(hidden_states.shape)  # [batch, seq, hidden]

プロキシオブジェクト

trace の内部では、モジュールへのアクセスは操作を記録するプロキシオブジェクトを返します:

with model.trace("Hello"):
    # These are all Proxy objects - operations are deferred
    h5_out = model.transformer.h[5].output[0]  # Proxy
    h5_mean = h5_out.mean(dim=-1)              # Proxy
    h5_saved = h5_mean.save()                   # Save for later access

ワークフロー 1: アクティベーション分析

ステップごとの手順

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

prompt = "The capital of France is"

with model.trace(prompt) as tracer:
    # 1. Collect activations from multiple layers
    layer_outputs = []
    for i in range(12):  # GPT-2 has 12 layers
        layer_out = model.transformer.h[i].output[0].save()
        layer_outputs.append(layer_out)

    # 2. Get attention patterns
    attn_patterns = []
    for i in range(12):
        # Access attention weights (after softmax)
        attn = model.transformer.h[i].attn.attn_dropout.input[0][0].save()
        attn_patterns.append(attn)

    # 3. Get final logits
    logits = model.output.save()

# 4. Analyze outside context
for i, layer_out in enumerate(layer_outputs):
    print(f"Layer {i} output shape: {layer_out.shape}")
    print(f"Layer {i} norm: {layer_out.norm().item():.3f}")

# 5. Find top predictions
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
for token, prob in zip(top_tokens.indices, top_tokens.values):
    print(f"{model.tokenizer.decode(token)}: {prob.item():.3f}")

チェックリスト

• LanguageModel ラッパーでモデルを読み込む
• 操作には trace コンテキストを使用する
• コンテキスト後に必要な値に対して .save() を呼び出す
• コンテキストの外で保存された値にアクセスする
• 分析には .shape、.norm() などを使用する

ワークフロー 2: アクティベーションパッチング

ステップごとの手順

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

clean_prompt = "The Eiffel Tower is in"
corrupted_prompt = "The Colosseum is in"

# 1. Get clean activations
with model.trace(clean_prompt) as tracer:
    clean_hidden = model.transformer.h[8].output[0].save()

# 2. Patch clean into corrupted run
with model.trace(corrupted_prompt) as tracer:
    # Replace layer 8 output with clean activations
    model.transformer.h[8].output[0][:] = clean_hidden

    patched_logits = model.output.save()

# 3. Compare predictions
paris_token = model.tokenizer.encode(" Paris")[0]
rome_token = model.tokenizer.encode(" Rome")[0]

patched_probs = torch.softmax(patched_logits[0, -1], dim=-1)
print(f"Paris prob: {patched_probs[paris_token].item():.3f}")
print(f"Rome prob: {patched_probs[rome_token].item():.3f}")

系統的なパッチングスイープ

def patch_layer_position(layer, position, clean_cache, corrupted_prompt):
    """Patch single layer/position from clean to corrupted."""
    with model.trace(corrupted_prompt) as tracer:
        # Get current activation
        current = model.transformer.h[layer].output[0]

        # Patch only specific position
        current[:, position, :] = clean_cache[layer][:, position, :]

        logits = model.output.save()

    return logits

# Sweep over all layers and positions
results = torch.zeros(12, seq_len)
for layer in range(12):
    for pos in range(seq_len):
        logits = patch_layer_position(layer, pos, clean_hidden, corrupted)
        results[layer, pos] = compute_metric(logits)

ワークフロー 3: NDIF によるリモート実行

ローカルの GPU を使わずに、同じ実験を大規模モデルで実行します。

ステップごとの手順

from nnsight import LanguageModel

# 1. Load large model (will run remotely)
model = LanguageModel("meta-llama/Llama-3.1-70B")

# 2. Same code, just add remote=True
with model.trace("The meaning of life is", remote=True) as tracer:
    # Access internals of 70B model!
    layer_40_out = model.model.layers[40].output[0].save()
    logits = model.output.save()

# 3. Results returned from NDIF
print(f"Layer 40 shape: {layer_40_out.shape}")

# 4. Generation with interventions
with model.trace(remote=True) as tracer:
    with tracer.invoke("What is 2+2?"):
        # Intervene during generation
        model.model.layers[20].output[0][:, -1, :] *= 1.5

    output = model.generate(max_new_tokens=50)

NDIF のセットアップ

1. login.ndif.us でサインアップする
2. API キーを取得する
3. 環境変数を設定するか、nnsight に渡します:

import os
os.environ["NDIF_API_KEY"] = "your_key"

# Or configure directly
from nnsight import CONFIG
CONFIG.API_KEY = "your_key"

NDIF で利用可能なモデル

• Llama-3.1-8B、70B、405B
• DeepSeek-R1 モデル
• 各種のオープンウェイトモデル（最新の一覧は ndif.us を確認してください）

ワークフロー 4: プロンプト間のアクティベーション共有

1 つのトレース内で、異なる入力の間でアクティベーションを共有します。

from nnsight import LanguageModel

model = LanguageModel("gpt2", device_map="auto")

with model.trace() as tracer:
    # First prompt
    with tracer.invoke("The cat sat on the"):
        cat_hidden = model.transformer.h[6].output[0].save()

    # Second prompt - inject cat's activations
    with tracer.invoke("The dog ran through the"):
        # Replace with cat's activations at layer 6
        model.transformer.h[6].output[0][:] = cat_hidden
        dog_with_cat = model.output.save()

# The dog prompt now has cat's internal representations

ワークフロー 5: 勾配に基づく分析

逆伝播中に勾配にアクセスします。

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

with model.trace("The quick brown fox") as tracer:
    # Save activations and enable gradient
    hidden = model.transformer.h[5].output[0].save()
    hidden.retain_grad()

    logits = model.output

    # Compute loss on specific token
    target_token = model.tokenizer.encode(" jumps")[0]
    loss = -logits[0, -1, target_token]

    # Backward pass
    loss.backward()

# Access gradients
grad = hidden.grad
print(f"Gradient shape: {grad.shape}")
print(f"Gradient norm: {grad.norm().item():.3f}")

注意: vLLM やリモート実行では勾配へのアクセスはサポートされていません。

よくある問題と解決策

問題: モデルによってモジュールパスが異なる

# GPT-2 structure
model.transformer.h[5].output[0]

# LLaMA structure
model.model.layers[5].output[0]

# Solution: Check model structure
print(model._model)  # See actual module names

問題: 保存を忘れる

# WRONG: Value not accessible outside trace
with model.trace("Hello"):
    hidden = model.transformer.h[5].output[0]  # Not saved!

print(hidden)  # Error or wrong value

# RIGHT: Call .save()
with model.trace("Hello"):
    hidden = model.transformer.h[5].output[0].save()

print(hidden)  # Works!

問題: リモートのタイムアウト

# For long operations, increase timeout
with model.trace("prompt", remote=True, timeout=300) as tracer:
    # Long operation...

問題: 保存するアクティベーションが多い場合のメモリ

# Only save what you need
with model.trace("prompt"):
    # Don't save everything
    for i in range(100):
        model.transformer.h[i].output[0].save()  # Memory heavy!

    # Better: save specific layers
    key_layers = [0, 5, 11]
    for i in key_layers:
        model.transformer.h[i].output[0].save()

問題: vLLM の勾配の制限

# vLLM doesn't support gradients
# Use standard execution for gradient analysis
model = LanguageModel("gpt2", device_map="auto")  # Not vLLM

主要 API リファレンス

メソッド/プロパティ	用途
model.trace(prompt, remote=False)	トレースコンテキストを開始
proxy.save()	トレース後にアクセスできるよう値を保存
proxy[:]	プロキシをスライス/インデックス指定（代入するとパッチされる）
tracer.invoke(prompt)	トレース内にプロンプトを追加
model.generate(...)	介入を伴って生成
model.output	モデルの最終出力ロジット
model._model	基盤となる HuggingFace モデル

他のツールとの比較

機能	nnsight	TransformerLens	pyvene
任意のアーキテクチャ	はい	Transformerのみ	はい
リモート実行	はい (NDIF)	いいえ	いいえ
一貫したAPI	いいえ	はい	はい
遅延実行	はい	いいえ	いいえ
HuggingFaceネイティブ	はい	再実装	はい
共有可能な設定	いいえ	いいえ	はい

リファレンスドキュメント

詳細なAPIドキュメント、チュートリアル、高度な使い方については、references/ フォルダーを参照してください。

ファイル	内容
references/README.md	概要とクイックスタートガイド
references/api.md	LanguageModel、トレース、プロキシオブジェクトの完全なAPIリファレンス
references/tutorials.md	ローカルおよびリモートの解釈可能性に関するステップバイステップのチュートリアル

外部リソース

チュートリアル

• Getting Started
• Features Overview
• Remote Execution
• Applied Tutorials

公式ドキュメント

• Official Docs
• NDIF Info
• Community Forum

論文

• NNsight and NDIF Paper - Fiotto-Kaufman et al. (ICLR 2025)

アーキテクチャのサポート

nnsight は任意の PyTorch モデルで動作します:

• Transformers: GPT-2、LLaMA、Mistral など
• State Space Models: Mamba
• Vision Models: ViT、CLIP
• カスタムアーキテクチャ: 任意の nn.Module

重要なのは、適切なコンポーネントにアクセスするためにモジュール構造を把握しておくことです。