Nodes/Concept Steer/Train Lens (SAE)
ComfyUI Node

Train Lens (SAE)

Train a concept direction via Sparse Autoencoder decomposition of the text encoder's residual stream. Produces an interpretable lens showing which features define the concept. Slower (~5min) but more informative than contrastive alone.

By Nynxz·Created 5 months ago·Updated 5 months ago· 4
Train Lens (SAE)
    • lens_path
    concept_namemy_concept
    positive_texts
    negative_texts
    layer22
    sae_expansion8
    sae_epochs200
    sae_features30
    n_prompts500
    refine_contrastivetrue
    contrastive_steps500
    sae_save_path
    sae_load_path
    transcoder_repo
    encoder_path
    output_dir
    protect_existingtrue
    CategoryConcept Steer

    Inputs (16)

    NameTypeDefaultDescription
    concept_nameSTRINGmy_conceptName for the concept (used in filename)
    positive_textsSTRINGTexts that EMBODY the concept (one per line). 10+ pairs recommended for strong results.
    negative_textsSTRINGNeutral/opposite texts (one per line, same count as positive). Should describe SIMILAR scenes WITHOUT the concept.
    layerINT221–36Which transformer layer to hook into. Default 22 (~60% depth) is the sweet spot for concepts. Earlier layers = lower-level features, later = more abstract.
    sae_expansionINT82–16SAE hidden dimension multiplier. 8x (default) = 20,480 features for 2560d Qwen. Higher = more features but slower training.
    sae_epochsINT20050–1000SAE training epochs (default 200, more = better features)
    sae_featuresINT305–100Number of top SAE features to keep in the concept direction. More features = richer concept but noisier.
    n_promptsINT500100–2000Number of diverse prompts for SAE activation collection
    refine_contrastiveBOOLEANtrueAlso run contrastive optimization on output embeddings and blend with SAE direction. Recommended for best results.
    contrastive_stepsINT500500–20000Contrastive optimization steps (only used if refine_contrastive is enabled)
    sae_save_pathSTRINGSave the trained SAE to this path for reuse across concepts. Saves ~5min per subsequent concept.
    sae_load_pathSTRINGLoad a pre-trained SAE/transcoder instead of training a new one. Supports both native .pt SAE and .safetensors transcoder formats.
    transcoder_repoSTRINGHuggingFace repo for pretrained transcoders (e.g. 'mwhanna/qwen3-4b-transcoders'). Downloads a 64x expansion transcoder with 163,840 features trained on ~1B tokens. Much better than training a small SAE. Leave empty to train your own SAE.
    encoder_pathSTRINGPath to Qwen 3.4B safetensors file. Leave empty to use QWEN_ENCODER_PATH env var.
    output_dirSTRINGOverride output directory for the lens file
    protect_existingBOOLEANtrueIf the output lens file already exists, save as _v2, _v3, … instead of overwriting. Disable only when intentionally replacing.

    Outputs (1)

    NameTypeDescription
    lens_pathSTRING