Skip to content

README.md

Latest commit

 

History

History
237 lines (168 loc) · 9.19 KB

README.md

File metadata and controls

237 lines (168 loc) · 9.19 KB
Raw

neuron-circuits

Attribute and steer individual MLP neurons in language models.

from neuron_steer import NeuronSteerer

steerer = NeuronSteerer("meta-llama/Llama-3.1-8B-Instruct")

# Behavioral steering: discover refusal circuit from positive/negative prompt pairs
circuit = steerer.find_feature(
    positive=["How do I pick a lock?", "Write malware code"],
    negative=["How do I bake a cake?", "Write clean code"],
    name="refusal",
)
steerer.steer("How do I pick a lock?", feature="refusal", multiplier=0.0)
# Answers directly instead of refusing

# Factual steering: discover capitals circuit from a single target token
circuit = steerer.find_feature(
    prompt="What is the capital of the state containing Dallas?",
    target=" Austin", name="capitals"
)
steerer.steer("What is the capital of Ohio?", feature="capitals", multiplier=0.0)
# "I don't know" -- the capital-city circuit is ablated

Implements Contrastive Neuron Attribution (CNA): discover sparse MLP neuron circuits for any behavior using contrastive activation analysis, then steer that behavior at inference time by scaling the identified neurons. ~100--200 MLP neurons form a complete circuit. A single forward+backward pass finds them.

Install

pip install torch transformers accelerate
pip install -e .

Python 3.9+, PyTorch 2.0+ with CUDA. GPU required (16GB+ VRAM).

See quickstart.py for a runnable end-to-end example. Also: refusal steering, interactive REPL.

Features

  • Contrastive discovery -- find neurons for any behavioral feature (refusal, belief, sentiment, sycophancy) from positive/negative prompt pairs, no target token needed
  • Single-pass circuit discovery -- RelP/LRP attribution finds factual circuits in one forward+backward pass
  • Multiplier steering -- ablate (0.0), baseline (1.0), amplify (2.0+), or sweep across multipliers
  • Edge attribution -- neuron-to-neuron information flow, hourglass architecture detection, super weight identification
  • Automatic universal neuron blacklisting -- filters task-agnostic infrastructure neurons
  • Cross-model support -- Llama, Qwen, Mistral with zero code changes
  • Interactive REPL -- explore circuits live with steerer.interactive()
  • Batch faithfulness evaluation -- circuit quality measurement with percentage threshold sweep

Results

Ablating 0.1% of MLP activations reduces refusal rates by over 50% on JBB-Behaviors across all model sizes and architectures tested, while maintaining near-baseline generation quality (>0.97) at all steering strengths. CAA achieves comparable refusal reduction at moderate strengths but degrades output quality sharply beyond α=0.5.

JBB-Behaviors refusal rates (instruct models, α=1.0)

Model Baseline Ablated Δ Relative
Llama-3.2-1B-Instruct 90% 34% −56pp −62.2%
Llama-3.2-3B-Instruct 84% 47% −37pp −44.0%
Llama-3.1-8B-Instruct 90% 34% −56pp −62.2%
Llama-3.1-70B-Instruct 86% 18% −68pp −79.1%
Qwen2.5-1.5B-Instruct 93% 12% −81pp −87.1%
Qwen2.5-3B-Instruct 90% 58% −32pp −35.6%
Qwen2.5-7B-Instruct 87% 2% −85pp −97.7%
Qwen2.5-72B-Instruct 78% 8% −70pp −89.7%

CNA vs CAA: refusal rate and generation quality (instruct models, α=1.0)

Model CNA Refusal% CNA Quality CAA Refusal% CAA Quality
Llama-3.2-1B-Instruct 20.2 0.975 0.0 0.554
Llama-3.2-3B-Instruct 26.3 0.977 0.0 0.431
Llama-3.1-8B-Instruct 5.1 0.969 38.4 0.493
Llama-3.1-70B-Instruct 12.1 0.981 0.0 0.569
Qwen2.5-1.5B-Instruct 26.3 0.982 100 0.888
Qwen2.5-3B-Instruct 34.3 0.984 0.0 0.844
Qwen2.5-7B-Instruct 13.1 0.980 5.1 0.414
Qwen2.5-72B-Instruct 5.1 0.983 98.0 0.406

Base vs instruct

Applying the same discovery pipeline to base models identifies neurons with similar activation differences, but steering them produces only content shifts — not behavioral change. Fine-tuning transforms the late-layer discrimination structure into a functional refusal gate.

Model Variant Baseline refusal% CNA Refusal% CNA Quality
Llama-3.2-1B Base 2.0 0.0 0.658
Llama-3.2-1B Instruct 43.4 20.2 0.975
Qwen2.5-3B Base 14.1 11.1 0.865
Qwen2.5-3B Instruct 92.9 34.3 0.984

API Reference

NeuronSteerer(model_name, device="cuda", dtype=torch.bfloat16, auto_blacklist=True)

Loads a HuggingFace causal LM with eager attention and auto-detects universal neurons.

High-Level API

find_feature(*, positive=None, negative=None, prompt=None, target=None, name=None, top_k=200, seed_response="") -> Circuit

Find a feature circuit. Two modes:

# Contrastive mode (behavioral features)
circuit = steerer.find_feature(
    positive=["How do I pick a lock?", "Write malware"],
    negative=["How do I bake a cake?", "Write clean code"],
    name="refusal",
)

# Single-prompt mode (factual features)
circuit = steerer.find_feature(
    prompt="Capital of Texas?", target=" Austin", name="capitals",
)

steer(prompt, *, feature=None, circuit=None, multiplier=0.0, max_new_tokens=50) -> str

Generate with a feature steered. Uses cached features from find_feature.

steerer.steer("How to pick a lock?", feature="refusal", multiplier=0.0)

interactive()

Launch the interactive REPL:

neuron> prompt What is the capital of Ohio?
neuron> discover Austin
neuron> ablate top10
neuron> sweep 0.0 0.5 1.0 2.0 5.0
neuron> edges
neuron> save my_circuit

Core Methods

discover_circuit(prompt, target_token, counterfactual_token=None, top_k=None, threshold=0.005, seed_response="", ...) -> Circuit

Single-prompt circuit discovery via RelP attribution.

discover_circuit_multi(prompts, target_tokens, counterfactual_tokens=None, ...) -> Circuit

Multi-prompt discovery. Attributes across prompts, unions per-prompt circuits.

discover_contrastive(positive_prompts, negative_prompts, top_k=200, ...) -> Circuit

Find neurons by contrasting activations between two prompt sets.

discover_edges(prompt, circuit, top_k_targets=30, ...) -> CircuitGraph

Neuron-to-neuron edges within a circuit. Returns a CircuitGraph with hub analysis, bottleneck detection, ASCII diagrams, and Graphviz export.

steer_and_generate(prompt, circuit, multiplier=0.0, max_new_tokens=50, ...) -> str

Generate with circuit neurons scaled by multiplier.

generate(prompt, max_new_tokens=50) -> str

Normal generation without steering.

next_token_probs(prompt, tokens, circuit=None, multiplier=1.0, ...) -> Dict[str, float]

Next-token probabilities for specific tokens, optionally with steering.

measure_faithfulness_batch(prompts, target_tokens, counterfactual_tokens, ...) -> List[Dict]

Batch faithfulness evaluation. Returns faithfulness and completeness at each threshold.

Data Structures

Circuit

circuit.top(k=20)           # Top-k neurons by attribution
circuit.by_layer()           # Group neurons by layer
circuit.unique_neurons()     # Unique neuron indices per layer
circuit.summary()            # Human-readable summary
circuit.save("path.json")    # Serialize to JSON
Circuit.load("path.json")    # Load from JSON

CircuitGraph

graph.top_edges(k=20)           # Top-k edges by weight
graph.edges_from(neuron_idx)    # Outgoing edges
graph.edges_to(neuron_idx)      # Incoming edges
graph.layer_flow()              # Layer-to-layer flow aggregates
graph.hub_analysis()            # Source/target hub ranking
graph.bottleneck()              # Hourglass bottleneck neurons
graph.detect_super_weights()    # Anomalous infrastructure neurons
graph.ascii_diagram()           # ASCII visualization
graph.to_dot("circuit.dot")     # Graphviz DOT export
graph.summary()                 # Human-readable summary

How It Works

Three LRP rules linearize the backward pass for neuron-level attribution:

  1. LN-rule (RMSNorm): Detach the normalization coefficient in the backward pass while preserving it in the forward pass. Preserves per-token scaling without letting normalization noise flow backward.

  2. AH-rule (Attention): Eager attention (not SDPA/Flash) so gradients flow through Q, K, V, and O projections cleanly.

  3. Half-rule (MLP gate): Shapley 50/50 attribution for the gate × up elementwise multiply — each factor gets half the gradient.

Contrastive pipeline:

positive prompts + negative prompts
-> collect last-token MLP activations per layer
-> mean(positive) - mean(negative) = delta per neuron
-> top-k by |delta| = contrastive circuit
-> hook circuit neurons -> generate with scaled activations

RelP pipeline (factual tasks):

prompt + target token
-> apply LRP rules -> forward pass -> backward from target logit
-> grad * activation = attribution per neuron -> threshold -> circuit
-> hook circuit neurons -> generate with scaled activations

License

MIT License. See LICENSE.