nnsight-remote-interpretability

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

nnsight (/ɛn.saɪt/) を使うと、研究者はあらゆる PyTorch モデルの内部を解釈し、操作できます。特筆すべき機能として、同じコードを小規模モデルではローカルで、大規模モデル（70B+）では NDIF 経由でリモート実行できます。

GitHub: ndif-team/nnsight 論文: 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 のリモート実行には、login.ndif.us で登録して API キーを取得してください。

主な概念

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 内でモジュールにアクセスすると、操作を記録する Proxy オブジェクトが返されます：

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
あらゆるアーキテクチャ	はい	Transformers のみ	はい
リモート実行	はい（NDIF）	いいえ	いいえ
一貫した API	いいえ	はい	はい
遅延実行	はい	いいえ	いいえ
HuggingFace ネイティブ	はい	再実装	はい
共有可能な設定	いいえ	いいえ	はい

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

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

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

外部リソース

チュートリアル

• はじめに
• 機能の概要
• リモート実行
• 応用チュートリアル

公式ドキュメント

• 公式ドキュメント
• NDIF の情報
• コミュニティフォーラム

論文

• NNsight と NDIF の論文 - Fiotto-Kaufman et al. (ICLR 2025)

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

nnsight は、あらゆる PyTorch モデルで動作します。

• トランスフォーマー: GPT-2、LLaMA、Mistral など
• 状態空間モデル: Mamba
• ビジョンモデル: ViT、CLIP
• カスタムアーキテクチャ: 任意の nn.Module

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

関連スキル

• transformer-lens-interpretability: HookP…を介してトランスフォーマーの内部を調査および操作し、TransformerLens を使ったメカニスティック解釈可能性研究を行うためのガイダンスを提供します。
• sparse-autoencoder-training: Sparse Autoencoders（SAE）をトレーニングおよび分析し、ニューラルネットワークの活性化を解釈…に分解するためのガイダンスを、SAELens を使って提供します。
• slime-rl-training: Megatron+SGLang フレームワークである slime を使った、RL による LLM のポストトレーニングのガイダンスを提供します。