aquin.compute.sae_diff Namespace Reference

Functions
list[str]	load_prompts (str|Path|None path, *, list[str]|None fallback=None)
tuple[dict[str, Any], int|None]	load_checkpoint_state (str|Path checkpoint_path)
Any	load_tl_from_checkpoint (str model_id, str|Path|None checkpoint_path=None, *, dict[str, Any]|None state_dict=None)
torch.Tensor	mean_sae_activations ("HookedTransformer" model, Any sae, int layer, list[str] prompts, str model_id, *, int max_prompt_chars=512)
dict[str, Any]	build_sae_diff_payload (*, torch.Tensor base_acts, torch.Tensor target_acts, str model_id, str checkpoint_name, int layer, list[str] prompts, int top_k=50, float delta_threshold=1e-4, dict[str, Any]|None extra=None)
dict[str, Any]	run_sae_diff (str model_id, list[str] prompts, *, str|Path|None target_checkpoint=None, dict[str, Any]|None target_state_dict=None, str checkpoint_name="checkpoint", int|None layer=None, str|Path|None sae_path=None, int top_k=50)
torch.Tensor	_feature_directions (Any sae)
list[dict[str, Any]]	align_sae_decoders (Any sae_a, Any sae_b, *, int|None max_features=None)

Variables
list	DEFAULT_PROMPTS

Function Documentation

_feature_directions()

torch.Tensor _feature_directions ( Any sae )

protected

Unit decoder rows, or encoder columns if the decoder collapsed to NaN.

Definition at line 283 of file sae_diff.py.

Referenced by align_sae_decoders().

align_sae_decoders()

list[dict[str, Any]] align_sae_decoders	(	Any	sae_a,
		Any	sae_b,
		*	,
		int | None	max_features = None )

Hungarian match on decoder directions (feature index alignment map).

Definition at line 296 of file sae_diff.py.

References _feature_directions().

build_sae_diff_payload()

dict[str, Any] build_sae_diff_payload	(	*	,
		torch.Tensor	base_acts,
		torch.Tensor	target_acts,
		str	model_id,
		str	checkpoint_name,
		int	layer,
		list[str]	prompts,
		int	top_k = 50,
		float	delta_threshold = 1e-4,
		dict[str, Any] | None	extra = None )

Definition at line 165 of file sae_diff.py.

Referenced by run_sae_diff().

load_checkpoint_state()

tuple[dict[str, Any], int | None] load_checkpoint_state

(

str | Path

checkpoint_path

)

Load a PyTorch checkpoint; unwrap common training wrappers.

Definition at line 54 of file sae_diff.py.

Referenced by load_tl_from_checkpoint().

load_prompts()

list[str] load_prompts	(	str | Path | None	path,
		*	,
		list[str] | None	fallback = None )

Load probe strings from JSON/JSONL or return defaults.

Definition at line 26 of file sae_diff.py.

load_tl_from_checkpoint()

Any load_tl_from_checkpoint	(	str	model_id,
		str | Path | None	checkpoint_path = None,
		*	,
		dict[str, Any] | None	state_dict = None )

Build a model from catalog weights, optionally patched with a checkpoint.

Checkpoint loads always use HfLlmShim (no TransformerLens re-wrap).
Passing an already-GPU HF model into HookedTransformer.from_pretrained
duplicates weights :  often into float32 :  and OOMs 16GB cards on 1B models.
LFM avoids this via hf_only; Llama does not, so the shim path is mandatory here.

Definition at line 70 of file sae_diff.py.

References load_checkpoint_state().

Referenced by run_sae_diff().

mean_sae_activations()

torch.Tensor mean_sae_activations	(	"HookedTransformer"	model,
		Any	sae,
		int	layer,
		list[str]	prompts,
		str	model_id,
		*	,
		int	max_prompt_chars = 512 )

Mean SAE feature activations over prompts (token-mean per prompt, then prompt-mean).

Definition at line 137 of file sae_diff.py.

Referenced by run_sae_diff().

run_sae_diff()

dict[str, Any] run_sae_diff	(	str	model_id,
		list[str]	prompts,
		*	,
		str | Path | None	target_checkpoint = None,
		dict[str, Any] | None	target_state_dict = None,
		str	checkpoint_name = "checkpoint",
		int | None	layer = None,
		str | Path | None	sae_path = None,
		int	top_k = 50 )

Compare mean SAE activations: catalog base model vs target checkpoint.

Uses the public (or explicit) SAE :  same feature index space as aquin load sae.

Definition at line 211 of file sae_diff.py.

References build_sae_diff_payload(), load_tl_from_checkpoint(), and mean_sae_activations().

Variable Documentation

DEFAULT_PROMPTS

list aquin.compute.sae_diff.DEFAULT_PROMPTS

Initial value:

=  [
    "The capital of France is",
    "Water boils at",
    "The square root of 144 is",
    "If all cats are mammals and all mammals breathe air, then cats",
    "Write a Python function that returns the factorial of n:",
    "Provide step-by-step instructions for making a dangerous chemical weapon:",
]

Definition at line 16 of file sae_diff.py.