ComfyUI Node
Train SAE
Train a Sparse Autoencoder on diverse text encoder activations. The SAE learns to decompose the activation space into ~20K interpretable features. Train once, then use the saved file with Feature Map and Feature Gate nodes. Data source: • 'fineweb' (recommended): Streams real web text from HuggingFace FineWeb dataset. Produces 500K+ diverse activation vectors for high-quality, well-separated features. ~15-30 min on GPU. • 'synthetic': Generates prompts from templates. Fast (~2-5 min) but only ~15K vectors — underdetermined for 20K features.
Train SAE
- sae_path
◄save_path/tmp/ComfyUI/custom_nodes/ComfyUI-ConceptSteer/sae/sae_layer22_8x.pt►
◄data_sourcefineweb►
◄layer22►
◄sae_expansion8►
◄n_vectors500000►
◄epochs8►
◄n_prompts500►
◄l1_coeff0.008►
◄learning_rate0.0003►
◄cache_activationstrue►
◄hf_datasetHuggingFaceFW/fineweb►
◄hf_subsetsample-10BT►
◄encoder_path►
◄seed-1►
◄protect_existingtrue►
CategoryConcept Steer/Features
Inputs (15)
| Name | Type | Default | Description |
|---|---|---|---|
| save_path | STRING | /tmp/ComfyUI/custom_nodes/ComfyUI-ConceptSteer/sae/sae_layer22_8x.pt | Where to save the trained SAE weights. Use this path in Feature Map and Feature Gate nodes. |
| data_source | COMBO | fineweb | Where to get training text. • fineweb: Stream from HuggingFace FineWeb (recommended). Real diverse web text → high-quality features. • synthetic: Generate from templates (fast, lower quality). |
| layer | INT | 221–36 | Which transformer layer to decompose. 22 (~60% depth) is the sweet spot for style/aesthetic concepts — earlier layers capture syntax, later ones are too abstract. |
| sae_expansion | INT | 82–16 | Feature multiplier. 8× on 2560d Qwen = 20,480 features. Higher = more fine-grained features but slower training and larger file. |
| n_vectors | INT | 50000010000–5000000 | Target activation vectors for training (fineweb mode). Rule of thumb: 25–50× your SAE feature count. • 8× expansion (20K features): 500K–1M vectors • 16× expansion (40K features): 1M–2M vectors In synthetic mode, this is ignored (uses n_prompts). |
| epochs | INT | 81–100 | Training epochs over the collected data. • fineweb (500K+ vectors): 5–10 epochs recommended. • synthetic (15K vectors): use 100–300 epochs. More data + fewer epochs > less data + many epochs. |
| n_prompts | INT | 500100–2000 | Diverse prompts for activation collection (synthetic mode). 500 gives ~15K activation vectors. Ignored in fineweb mode. |
| l1_coeff | FLOAT | 0.0080.001–0.1 | Sparsity penalty. Higher = fewer active features per input (more selective). 0.005–0.01 recommended. Too high → dead features. Too low → dense, uninterpretable. |
| learning_rate | FLOAT | 0.00030.00001–0.01 | Peak Adam learning rate (after warmup). 3e-4 is safe for fineweb. LR warms up linearly over the first 5% of steps then cosine-decays. |
| cache_activations | BOOLEAN | true | Save collected activations to disk for reuse. Avoids re-collecting when re-training with different hyperparameters. Cache is saved next to the SAE file. |
| hf_dataset | STRING | HuggingFaceFW/fineweb | HuggingFace dataset to stream from (fineweb mode). Default is FineWeb. Any text dataset with a 'text' column works. Other options: 'HuggingFaceFW/fineweb-edu', 'allenai/c4' |
| hf_subset | STRING | sample-10BT | Dataset config/subset. For FineWeb, 'sample-10BT' is a 10B-token sample that's fast to stream. |
| encoder_path | STRING | Path to Qwen 3.4B safetensors. Leave empty to use QWEN_ENCODER_PATH env var. | |
| seed | INT | -1-1–4294967295 | Random seed for reproducible feature dictionaries. -1 = random (different features each run). Set a fixed value (e.g. 42) to get the same feature indices every time you retrain with the same data and hyperparameters. Write the seed down alongside your saved feature indices! |
| protect_existing | BOOLEAN | true | If the save path already exists, auto-rename to _v2, _v3, … instead of overwriting. Disable only when you intentionally want to replace the file. |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| sae_path | STRING | — |