Interpret API

Locate the parts of a model responsible for safety behavior. Import from safetune.interpret. This pillar is functions + config/report dataclasses rather than trainers.

from safetune.interpret import identify_safety_neurons, safety_circuit_info

# One-call wrapper: refusal-direction extraction + weight-based neuron scan
circuit = safety_circuit_info(model, tokenizer, harmful_prompts, harmless_prompts)

# Or score neurons directly against precomputed per-layer refusal directions
report = identify_safety_neurons(model, refusal_direction_per_layer)

Public surface

Functions: identify_safety_neurons, safety_circuit_info, eap_safety_circuit. Configs / reports: SafetyNeuronConfig, SafetyNeuronReport, EAPSafetyCircuitConfig.

See the Interpret guide.

Reference

safetune.interpret.identify_safety_neurons(model, refusal_direction_per_layer, config=None, *, tokenizer=None, harmful_prompts=None, harmless_prompts=None)

Rank per-layer neurons by their relevance to refusal/safety behaviour.

Two localization paths, selected by config.method:

• "weight" (default) -- rank units by absolute cosine between the unit's residual-stream write-direction (a column of target_module) and the layer's refusal direction. Needs only refusal_direction_per_layer; no forward passes.
• "activation" -- rank units by a per-neuron activation contrast between harmful and harmless prompts (feed-forward activation analysis, cf. Wei et al. arXiv:2402.05162, Chen et al. arXiv:2406.14144). Needs tokenizer, harmful_prompts and harmless_prompts; does not use refusal_direction_per_layer.

Parameters:

Name	Type	Description	Default
model	Module	HF causal LM.	required
refusal_direction_per_layer	Dict[int, Tensor]	{layer_idx: 1-D tensor (hidden,)}, typically from :func:safetune.steer.extract_refusal_direction (the all_layer_directions return value). Used by "weight" mode; may be {} for "activation" mode.	required
config	Optional[SafetyNeuronConfig]	:class:SafetyNeuronConfig. Default uses method="weight", target_module="mlp.down_proj", top_k_per_layer=16.	None
tokenizer	Any	HF tokenizer -- required for method="activation".	None
harmful_prompts	Optional[List[str]]	harmful contrast corpus -- required for method="activation".	None
harmless_prompts	Optional[List[str]]	harmless contrast corpus -- required for method="activation".	None

Returns:

Type	Description
SafetyNeuronReport	class:SafetyNeuronReport with per-layer top-k unit indices and scores.

safetune.interpret.safety_circuit_info(model, tokenizer, harmful_prompts, harmless_prompts, *, top_k_per_layer=16, target_module='mlp.down_proj', target_layers=None, method='weight', activation_module='mlp.act_fn', activation_score='mean_abs_diff')

Convenience: locate safety neurons end-to-end and return a CircuitInfo.

In method="weight" mode (default) this extracts a refusal direction from the contrast corpus and ranks units by weight-direction cosine. In method="activation" mode it skips the refusal-direction step and ranks units directly by the harmful-vs-harmless activation contrast -- the refusal direction is still extracted so that direction_layer is populated for downstream consumers, but it does not affect the scores.

Returns the :class:CircuitInfo produced by :meth:SafetyNeuronReport.as_circuit_info. Plug directly into LSSF, PKE, NLSR, or DeepRefusal for targeted patching.

Parameters:

Name	Type	Description	Default
method	str	"weight" or "activation" (see :func:identify_safety_neurons).	'weight'
activation_module	str	module whose activations are contrasted when method="activation"; forwarded to :class:SafetyNeuronConfig.	'mlp.act_fn'
activation_score	str	activation contrast score when method="activation"; forwarded to :class:SafetyNeuronConfig.	'mean_abs_diff'

safetune.interpret.SafetyNeuronConfig dataclass

Configuration for safety-neuron localization.

Attributes:

Name	Type	Description
method	str	"weight" (output-direction cosine, no extra forwards) or "activation" (per-neuron harmful-vs-harmless activation contrast across a contrast corpus).
top_k_per_layer	int	keep only the top-k highest-scoring units per layer.
target_layers	Optional[List[int]]	restrict to these layers; None means all decoder layers.
score_floor	float	drop units whose absolute score is below this threshold.
target_module	str	which projection matrix to score in "weight" mode. "mlp.down_proj" (default): each column is one MLP intermediate neuron's residual-stream write-direction -- columns are the natural unit axis and the cosine against the refusal direction is exact. "self_attn.o_proj": columns index attention output dimensions (head_dim slices of the concatenated heads), not whole heads; the per-column cosine is still well defined but a "unit" is then an o_proj input channel, not an attention head.
activation_module	str	in "activation" mode, which sub-module's output is captured as the per-neuron activation vector. "mlp.act_fn" captures the post-activation MLP intermediate (gate) neurons -- these are the canonical "neurons" of the safety-neuron literature. If that hook point is unavailable the implementation falls back to "mlp" (the MLP block output, i.e. residual-stream units).
activation_score	str	how to turn captured activations into a per-neuron contrast. "mean_abs_diff" (default): mean |act| on harmful minus mean |act| on harmless. "tstat": a Welch-style standardized mean difference (mean-abs difference divided by pooled std) -- the label-correlation form, robust to per-neuron scale. "mean_diff": signed mean (not magnitude) difference.
activation_batch_size	int	forward-pass batch size for the contrast corpus.
activation_max_tokens	int	truncate each prompt to this many tokens.
abs_rank	bool	rank units by absolute score (True, default) so that strongly suppressed-on-harmful neurons are also surfaced.