docs/methods/sae/API.md

SAE API

SAE hub · Complete runnable example · Pretrained loading without fitting

Use existing weights

The snippets in this reference assume probe is a Prober bound to your loaded backbone and processor, and inputs is a prepared input batch. The linked complete loading example constructs all three from scratch.

# Native dictionary saved by this library; no fitting is needed.
method = probe.sae.load("outputs/dictionary", device="cpu")
result = method.run(inputs)

# Published GPT-2 layer-0 TopK recipe; downloads and converts the author weights.
# The already-loaded backbone/tokenizer must match this recipe exactly.
method = probe.sae.from_pretrained(
    "sae-6783018355728933dae6d812", cache_dir="outputs/sae_cache", device="cpu",
)
result = method.run(inputs)

The complete inference-only example includes the matching backbone, tokenizer, fresh input and visualization. from_pretrained accepts only implemented conversion recipes. A catalog record alone does not imply automatic source-format or backbone compatibility.

Loading parameter	Default	Meaning
card_id	Required	Published model ID with an implemented automatic loading recipe.
cache_dir	None	Source checkpoint download cache and converted native-artifact cache; defaults to the Hugging Face cache root.
device	"cpu"	Device for the loaded SAE dictionary; backbone placement remains caller-owned.

Configure and fit a new dictionary

Here probe, train_inputs and validation_inputs are constructed in a method's complete example, such as ReLU SAE.

from vlmprobe.sae import SAETrainingConfig

method = probe.sae.topk(layer=0, tokens="visual", d_sae=1024, k=32)
history = method.fit(
    train_inputs, validation_inputs, output_dir="outputs/fit",
    training=SAETrainingConfig(max_steps=200, batch_size=128),
    preprocessing_id="my-independent-training-split-v1",
)
method.save("outputs/dictionary", model_card=None, registry_dir=None)
loaded = probe.sae.load("outputs/dictionary", tokens=None, binding=None,
                        registry_dir=None, artifact_dir=None, device="cpu")
result = loaded.run(validation_inputs[0])

Every named factory is also available through probe.sae.method("topk", layer=0, tokens="visual", d_sae=1024, k=32).

.fit(...) trains the SAE encoder/decoder and any method-specific learned parameters on activations from the frozen backbone. It does not fine-tune that backbone. Loading a trained dictionary and calling .run(...) performs inference without optimizer updates. A pretrained backbone and a pretrained SAE are separate sets of weights; one does not supply the other. See the inference-only guide for complete loading examples.

layer is one zero-based residual layer from describe(). tokens accepts the common token selector (all, visual, text, last_prompt, positions or a boolean mask). A saved batch-specific selector needs matching coordinates at inference; it is not a dataset-independent image region description. d_sae defaults to four times the inferred activation width; expansion_factor=4 controls that default. An omitted k is capped at the inferred dictionary width; an explicitly invalid k raises. train_inputs and validation_inputs are finite native model input batches or ProbeInputs. Actual model forwards capture detached CPU float32 activations. Each split has a configurable max_cache_bytes payload cap. It is not a hard total-RAM limit: concatenation temporarily retains both chunks and their joined tensor, plus the other split; the trainer also uses in-memory tensors. train_ids and validation_ids, when provided, must be disjoint source groups. Exact repeated complete batches are independently rejected; this does not infer patient/slide overlap for you.

Fit hashes loaded model tensors and processor configuration, records the exact adapter hook and token selection, and derives a calibration fingerprint from inputs/activations/source IDs. preprocessing_id must describe your actual external transforms; its default explicitly records that they are unverified. Load verifies actual model tensors, architecture/configuration, numerical-stack versions and processor/tokenizer behavior when fingerprints are present. The frozen backbone is assumed unchanged between method calls; verification is performed at fit/load, rather than rehashing a large backbone on every run. Source artifacts without automatic fingerprints require a caller-supplied verified binding; width matching alone is rejected as an automatic compatibility claim.

Loading creates an inference method. Calling .fit() on that loaded object is rejected: it would otherwise silently discard weights or recipe settings. To resume training, recreate the original factory/configuration and call fit with its matching output directory; warm-start fine-tuning is not yet provided. Resaving modified dictionary weights invalidates inherited numerical/semantic evidence and marks earlier training metrics as stale.

Torch and explicit loader/sampler RNG states are preserved during collection. Arbitrary Python/NumPy transforms and persistent worker processes need their own deterministic data or caching protocol.

Dictionary hyperparameters

The model-bound factory forwards keyword arguments to SAEConfig after inferring d_in. Methods share these fields, but only relevant algorithm options affect the selected method. Explicit recipe defaults and mathematical scope appear in the method inventory.

Parameter	Default	Meaning
architecture	required	Native algorithm name; see the method inventory.
d_in	required	Captured activation width; inferred by the model-bound fit workflow.
d_sae	required	Dictionary feature count; fixed for the artifact.
k	32	Positive per-vector TopK count, or batch budget as selected below; pursuit defaults to this many iterations.
l1_coefficient	0.01	Coefficient for L1 sparsity in ReLU/Gated or explicit JumpReLU L1 recipe.
seed	0	Local dictionary initialization seed; does not reset caller RNG.
normalization	None	Optional named normalization override: rms, none, layer_norm, constant_norm_rescale.
input_normalization	'scalar'	Internal coordinate recipe: scalar, none, layer_norm or constant_norm.
normalization_eps	1e-05	Positive stability constant for per-input normalization.
normalization_std_correction	0	Standard-deviation correction: 0 (population) or 1 (sample).
apply_decoder_bias_to_input	True	Subtract decoder bias before encoding normalized inputs.
apply_b_dec_to_input	None	Optional source-format alias of apply_decoder_bias_to_input.
use_decoder_bias	True	Whether reconstruction includes the learned decoder bias.
normalize_decoder	True	Apply per-atom decoder normalization after optimizer steps.
rescale_acts_by_decoder_norm	False	Explicit source recipe that rescales feature values using decoder norms.
auxk_coefficient	0.0	Optional dead-feature auxiliary reconstruction weight; 0 disables it.
auxk_k	None	Number of dead features considered by AuxK; None uses the documented method default.
dead_feature_steps	100	Number of training steps without firing before a feature is dead for AuxK.
jumprelu_threshold	0.1	Positive starting threshold for each JumpReLU feature.
jumprelu_bandwidth	0.05	Rectangle surrogate width for threshold derivatives.
l0_coefficient	0.01	Weight on JumpReLU L0 surrogate sparsity.
jumprelu_sparsity	'l0'	Explicit JumpReLU sparsity recipe: l0 or l1.
jumprelu_ste_to_input	False	Whether threshold surrogate gradients also reach encoder preactivations.
batch_budget	'average'	average: k times number of vectors; global: k across the entire minibatch.
threshold_ema	0.99	EMA decay for BatchTopK inference thresholds.
threshold_init	'first_batch'	Initialize inference threshold from first_batch or zero_ema.
matryoshka_widths	()	Increasing prefix widths ending at d_sae; empty selects half/full widths.
matryoshka_weights	()	Positive weights for each nested reconstruction loss; empty means equal weights.
dictionary	'learned'	learned or archetypal; archetypal method names select the latter.
archetypal_candidates	128	Maximum number of training-only anchor vectors.
archetypal_delta	1.0	Bound on the archetypal dictionary relaxation component.
pursuit_iterations	None	Maximum greedy support additions; None uses k.
residual_tolerance	0.01	Residual stopping tolerance in pursuit.
stop_on_duplicate_support	True	Stop matching pursuit when no new support is added.
omp_nnls_iterations	32	Projected-gradient refit steps for nonnegative OMP; finite approximation, not exact NNLS.
quantization_levels	4	Number of learned scalar code levels; QSAE is experimental.
hard_quantization	False	Use discrete quantization rather than the soft code path.

Training parameters

SAETrainingConfig is shared by all algorithms. The normalization field here chooses whether the training-only RMS scale is fitted for scalar-normalized models (rms) or disabled (none); per-input dictionary normalization takes precedence. These are not different source-specific trainers.

Parameter	Default	Meaning
epochs	10	Maximum epochs unless a step/vector budget is chosen.
lr	0.001	Adam learning rate.
batch_size	256	Number of activation vectors per optimizer minibatch.
patience	3	Stop after this many epochs without min_delta validation-loss improvement.
min_delta	0.0001	Minimum validation-loss improvement for best-model selection.
seed	0	Local dictionary initialization seed; does not reset caller RNG.
normalization	'rms'	Fit train-only scalar RMS (rms) or disable it (none); per-input model normalization takes precedence.
max_steps	None	Optional exact optimizer-step budget; mutually exclusive with max_vectors.
max_vectors	None	Optional exact cumulative training-vector budget; mutually exclusive with max_steps.
device	'cpu'	Dictionary training device (cpu or cuda); the caller owns backbone placement.
resume	True	Resume matching optimizer/RNG/cursor/method state from output_dir.
grad_clip	5.0	Positive maximum gradient norm applied before Adam updates.
progress_every	20	Progress/checkpoint write interval in optimizer steps.

Probe operations and outputs

Call	Parameters	Tensors
method.run(inputs)	Native input batch	features[N,F], activation[N,D], reconstruction[N,D], reconstruction_error[N,D], l0[N], positions[N,2]
method.profile(top_k=10).run(inputs)	Positive top count per feature; ranks absolute values and retains signs	Feature mean/variance/firing frequency/min/max; top_activation_values[K,F], top_positions[K,F,2]
method.reconstruct().run(inputs, metric=None)	Optional task metric	Baseline/edited logits or representation outputs; optional baseline/edited score and score change
method.edit(feature_ids=[...], scale=0., mode="error_preserving").run(inputs, metric=None)	Unique valid IDs; finite scaling; error_preserving or reconstruction	Same output/score comparison as above, with explicit intervention metadata

N packs selected batch/token positions, F=d_sae, D=d_in; the two position columns are batch and native expanded-token index. Signed features count as active when nonzero. Profiles rank absolute activation, retaining negative values. Representation outputs remain vectors; they are not fabricated logits.

Default feature editing uses x + decode(z_edited) - decode(z) with the same per-input normalization context reused for both decodes. scale=1 returns an exact no-op, including native bf16 precision. Reconstruction is deliberately a separate control. Generic low-level operations remain available as probe.sae.features/profile/reconstruct/edit(layer=..., sae=..., binding=...).

Lower-level training, artifacts and discovery

from vlmprobe.sae import (SAEConfig, SparseAutoencoder, SAETrainingConfig,
                         fit_sae, list_methods, list_models, get_model_card,
                         load_artifact, save_artifact, import_checkpoint)

fit_sae(model, train_vectors, validation_vectors, output_dir, training=..., provenance=...) uses the same engine as model-bound fitting. Vectors must be finite floating [N,D] in the bound native coordinate system. Legacy train_sae(...) is only a compatibility adapter to this engine.

save_artifact(model, path, model_card=None, registry_dir=None) writes either config.json/weights.pt/model_card.json or the same schema in a safe .pt envelope. load_artifact(path, expected_binding=None, device="cpu") validates checksums, types, exact state, config and optional binding. Inference weights and optimizer resume state are separate. Browse list_models(...) and get_model_card(id, ...); see model cards for local registration, published pretrained cards and explicit source checkpoint conversion.

Model discovery parameters

list_models(*, modality=None, method=None, status=None, query=None,
            artifact_kind="sae", registry_dir=None, include_builtin=True, limit=None)
get_model_card(card_id, *, registry_dir=None)

Parameter	Meaning
modality	Filter language, vision or multimodal resources. Unreported source metadata remains unspecified.
method	Filter an algorithm name such as relu, topk or jumprelu.
query	Case-insensitive search over model, hook, release and other card metadata.
artifact_kind	"sae" selects dictionary weights; "backbone" selects activation-extraction models; None includes both.
registry_dir	Optional directory containing your saved artifact/model cards.
include_builtin	Include the packaged author checkpoint catalog; use False to browse only local runs.
status	Optional artifact metadata filter, for example published or trained.
limit	Optional positive maximum result count.

Cards provide source model/hook/configuration, dimensions when reported, direct checkpoint links, and a link to our model information page. load_model(...) restores a registered native artifact; source weights use the explicit checkpoint conversion interface before model-bound probing.