AlignSAE: Concept-Aligned Sparse Autoencoders

In plain words

A language model stores facts inside its hidden numbers, and those numbers mix many ideas together, so it is hard to point at the spot where, say, "this person's birth city" lives. A sparse autoencoder is a tool that unpacks those numbers into a larger set of separate slots, each meant to switch on for one idea. The usual version learns the slots without any labels, so nothing says which slot should mean which idea.

This paper adds a teaching step after the usual one. First the autoencoder learns to rebuild the hidden numbers. Then it gets labelled examples and is trained so that each named concept gets its own dedicated slot, while the rest of the slots stay free to carry everything else. Once that is done, you can read off which slot is on, and you can turn one slot up to nudge the model toward a different answer. The authors test this on a small GPT-2 with made-up biographies and with a harder two-step question type.

TL;DR

• A standard sparse autoencoder is trained on one layer of a frozen model; a second, supervised phase binds each ontology relation to a fixed latent slot.
• Losses: reconstruction, L1 sparsity on free slots, cross-entropy binding on the relation slots, a decorrelation penalty between relation and free slots, and an auxiliary answer-prediction loss.
• Author-reported gains are on synthetic data with GPT-2 only; the official code is MIT-licensed but differs from the paper in several places (pooling, decorrelation form, swap steering scale, two-hop training settings).

Problem

Standard sparse autoencoders are trained without labels, so they have no reason to put a given human-defined concept into one feature. Concepts end up split across several features, and one feature can carry several concepts §1, §3. The authors argue this makes feature-level steering and knowledge editing unreliable §1. The method assumes a predefined ontology of concepts with labelled inputs; the paper's experiments use relation types as concepts §3.

Method

Setup. A frozen base model's hidden state h (dimension d) is read at one layer ℓ. An encoder maps h to a sparse code z of size R + K, where R is the number of ontology concepts and K the number of free features. The code is split into concept slots and free features. A decoder reconstructs h from z §4.1 app. C.

Training is two phases §3, app. C.2:

1. Pre-training (Stage 1). Only reconstruction is optimised, so the decoder forms a dictionary. The appendix and the code use reconstruction alone here; the main text's description of this phase also mentions sparsity §3. See Paper vs code.
2. Post-training (Stage 2). All six loss terms below are optimised jointly app. C.3. The main text describes the binding and value losses as strengthened and the orthogonality penalty as switched on at this point §3.

Loss terms (Stage 2, as stated in appendix C.3; the main text gives an equivalent set with different forms, see Paper vs code):

\mathcal{L}_{recon} = \mathrm{MSE}(\hat{h}, h) = \frac{1}{d}\sum_i (\hat{h}_i - h_i)^2 \quad [app.\ C.3]

\mathcal{L}_{sparse} = \frac{1}{B\, n_{free}} \sum_{b}\sum_{j \le n_{free}} |z_{b,j}| \quad [app.\ C.3]

\mathcal{L}_{align} = \mathrm{CE}\big(\mathrm{softmax}(z_{concept}), y_{rel}\big) \quad [app.\ C.3]

\mathcal{L}_{ortho} = \sum_{r}\sum_{j} \Big(\tfrac{1}{B}\sum_b (z_{b,r}-\bar z_r)(z_{b,j}-\bar z_j)\Big)^2 \quad [app.\ C.3]

\mathcal{L}_{value} = \frac{1}{B}\sum_b \mathrm{CE}\big(V_{r_b}(z_{b,r_b}), t_b\big) \quad [app.\ C.3]

where t_b is the first token of the gold answer and V_r is a two-layer MLP from the scalar activation of relation slot r to vocabulary logits. The value heads are not used during interventions app. C.3.

Total: L = L_recon + 1e-3·L_sparse + 1.0·L_align + 1e-2·L_ortho + 0.5·L_value app. C.3, weights in app. C.4.

The main text writes the same objective as L_SAE + λ_align·L_align + λ_⊥·L_⊥ + λ_val·L_val, where L_⊥ is a decorrelation penalty between concept and rest slots and L_val uses a softmax over the value head output §3, Objectives.

Two-hop variant §4.2.2, app. M. Each example has a start entity and two relations. The target is the intermediate entity followed by the final entity. The relation-1 slot is supervised at the first entity token and the relation-2 slot at the second. The layer is 6 for all two-hop experiments §7.

Controllability (swap test) §4.4, app. D. To make the model answer for a different relation, the decoded direction of the target slot is added to the residual stream at layer ℓ, at the token position of the question, and the model generates from there. A swap succeeds when the answer equals the gold value of the target relation §4.4.

# Stage 1: reconstruction only  [app. C.2] [§3]
# Input:  h (B, d) hidden activations; d = 768 for GPT-2

def stage1_loss(h):
    # Encode into free features and concept slots  [§4.1]
    z = relu(encoder(h))                  # (B, K + R)
    h_hat = decoder(z)                    # (B, d)
    # Mean squared error over batch and dimensions  [app. C.3]
    return mse(h_hat, h)

# Stage 2: supervised post-training  [§3] [app. C.3]
# Inputs: h (B, d); rel (B,) gold relation index in 0..R-1;
#         ans (B,) first token of gold answer

def stage2_loss(h, rel, ans):
    # Split the code: first K entries are free, last R are concept slots
    z = relu(encoder(h))                  # (B, K + R)
    z_free = z[:, :K]                     # (B, K)
    z_concept = z[:, K:]                  # (B, R)
    h_hat = decoder(z)                    # (B, d)
    # Reconstruction term  [app. C.3]
    recon = mse(h_hat, h)
    # L1 penalty on free features only  [app. C.3]
    sparse = mean(abs(z_free))
    # Binding: cross-entropy of concept slots against the gold relation  [§3]
    align = cross_entropy(z_concept, rel)
    # Decorrelate concept slots from free features  [app. C.3]
    ortho = sum(cross_cov(z_concept, z_free) ** 2)
    # Value head of the gold slot predicts the first answer token  [app. C.3]
    value = cross_entropy(value_head(z_concept, rel), ans)
    return (1.0 * recon + 1e-3 * sparse + 1.0 * align
            + 1e-2 * ortho + 0.5 * value)

# Swap intervention  [§4.4] [app. D]
# Inputs: residual stream h at layer l, target slot index j, strength alpha

def swap(h, target_slot, alpha):
    # Decoded direction of the target slot  [§4.4]
    direction = W_dec[:, target_slot]     # (d,)
    # Add it at the question's token position  [§4.4]
    return h + alpha * direction

Hyperparameters

Values in the paper's appendix unless marked otherwise. "Not stated" means the paper does not say.

Base model (GPT-2, 124M, frozen during SAE training)

SAE

Datasets and evaluation

Results

Author-reported. Setup: synthetic biography QA with GPT-2 124M, 1-hop (6 relations) and 2-hop (20 relations) §4.2, §5, §7. No hardware or training time is stated.

Layer 0 vs layer 6, 1-hop §5.1, layer comparison table

Ablation of post-training terms, 1-hop, layer 6 §6

Stage 1 ablation §6, stage-1 subsection: starting post-training from scratch, without Stage 1, gives a higher total loss (3.34 vs 1.86), higher reconstruction loss and far higher sparsity loss. Diagonal accuracy on train is 0.970 without Stage 1 and 1.000 with it; on the held-out split it is 0.970 and 0.954. The authors conclude that Stage 1 is needed for a stable basis.

Disentanglement §5.2 app. G

• Mutual Information Gap, 1-hop: 0.173 for AlignSAE, 0.000 for the standard SAE.
• Mutual Information Gap, 2-hop: 0.035 for AlignSAE, 0.004 for the standard SAE (reported as 9.2×).
• Concept fragmentation and concentration curves per layer are in a figure §5.2, fig. in §5.2; the text reports fragmentation near 1 for AlignSAE from mid layers on, and hundreds to thousands of features for the standard SAE.

Swap controllability, 1-hop §5.3: success rises from 0.54 at alpha = 1 to 0.85 at alpha = 2 at layer 6, and falls to 0.23 at alpha = 10. Failed swaps keep the right answer category in 75.3% of cases at alpha = 2 and 83.0% at alpha = 10 §5.4.

Two-hop §7

• Step-wise diagonal accuracy: 100% at both hop steps, after training.
• Grokking-style transition: diagonal accuracy at the first step is 38.4% at epoch 10 and 100% at epoch 40; validation accuracy rises later §7.2.
• Swap success peaks near alpha ≈ 50; AlignSAE reaches 7.86%, a standard SAE 1.5–2.7% across alpha app. K. The main text describes this as 4× higher swap success §1, §7; the appendix gives 2.9× against the standard SAE app. K.
• The authors name a gap between binding and causal control in two-hop settings: diagonal binding is perfect while swap success stays low app. K.

Abstract headline numbers: the abstract reports 94.8% swap success in 1-hop and a 9–∞× disentanglement gain. The 94.8% figure does not appear in the results tables; the table value is 0.847 at alpha = 2 §5.3, §6. See Pitfalls.

Free-form generation (LLM-judge, 1–5 scale, judged by Claude Haiku 4.5) app. I: steering toward one relation changes the first generated token to the target type while keeping fluency; the judge rates target-concept success at 3.9 for 1-hop and 2.5 for 2-hop at alpha = 10. Only 70 generations were rated.

Paper vs code

Code at commit 4fd6da0ccb25d05fb1729c7062eb039fef610d3a of github.com/yminglai/AlignSAE, MIT licence, owner is the first author's GitHub account as linked from the paper header §front matter.

Searched: links in the paper header, which point to this repository. No separate project page was needed. The repository is the authors' own.

Differences, each with both citations:

1. Stage 1 content. Main text describes pre-training with reconstruction and sparsity §3, Model Training. Appendix and code use reconstruction only app. C.2 04_train_sae.py. Stage 2 is gated by stage >= 2 04_train_sae.py.
2. Reconstruction loss. Main text uses a squared L2 norm summed over dimensions §3, Objectives. Appendix and code use mean squared error over batch and dimensions app. C.3 04_train_sae.py.
3. Decorrelation penalty. Main text names it as a correlation Frobenius norm §3, Objectives. Appendix and code use un-normalised centred cross-covariance, summed over pairs app. C.3 04_train_sae.py.
4. Value loss. Main text takes a softmax over the value head of the concept vector §3, Objectives. Code uses only the gold slot's scalar, passed through that relation's head 04_train_sae.py 04_train_sae.py 04_train_sae.py. The appendix matches the code app. C.3.
5. Activation position. Main text gives mean pooling as an example §4. Appendix and code use the last prompt token app. B 03_collect_activations.py 01_collect_activations_2hop.py.
6. Swap steering. Paper: h + alpha · W_dec · e_j §4.4 app. D. The released hook-based script defaults to decoder-column steering scaled by beta = ||h|| / ||d_j||, so the injected vector has norm alpha · ||h|| 06_swap_evaluate_alignsae_hook.py and hook registration 06_swap_evaluate_alignsae_hook.py. Whether the 1-hop tables were produced by this script is not stated.
7. Two-hop base training. Paper: 100,000 steps, batch 96, LR 5e-4 with warmup of 2,000, minimum 1e-5, max length 512 app. M. Repository run script: batch 128, max length 128, 100 epochs, LR 1e-4, weight decay 0.01, warmup ratio 0.1, mixed mode with 30% QA training examples run_train_two_hop.sh train_two_hop.py. Note that the repo script's comments refer to "paper setup" for the epoch count.
8. Two-hop SAE training. The 2-hop SAE script is one loop over all terms with no Stage 1 02_train_sae_2hop.py, uses Adam rather than AdamW 02_train_sae_2hop.py, LR 1e-3, batch 128, 100 epochs 02_train_sae_2hop.py. The paper does not state a two-hop SAE schedule.
9. Two-hop layer. The two-hop pipeline sets LAYER=6 run_full_pipeline.sh, but the collection script's default is the last layer 01_collect_activations_2hop.py.
10. Unchanged. The 1-hop base-model training settings in appendix B match the script defaults 02_sft_base_model.py. The SAE loss weights match the script defaults 04_train_sae.py. The SAE learning rate is constant, with no clipping, as stated 04_train_sae.py.
11. Excluded independence term. The paper's appendix excludes the free-free independence term app. C.3. The code computes it only when its weight is positive, and its CLI default is 0.0 04_train_sae.py.

Pitfalls

• Stage 1 scope is ambiguous. The main text and the appendix disagree on whether sparsity is active in Stage 1. Follow the appendix and code: reconstruction only.
• Pooling is ambiguous. Use the final prompt token as the code does. Mean pooling is only given as an example and is not used in the released code.
• Layer numbering. "Layer 6" is not defined in the paper. The collection code uses hidden_states[6], which counts the embedding output as index 0, so layer 6 is the output after six transformer blocks. Confirm this matches your framework.
• Ortho term scale. The decorrelation penalty sums squared cross-covariances over R × K pairs with no normalisation by K app. C.3. Because the sum grows with K, the stated weight 1e-2 is only meaningful at the K used. The paper does not report this term's magnitude against the others, so log it during training.
• Swap success headline. The 94.8% number in the abstract is not in any table the paper reports. The reported table value is 0.847 at alpha = 2 §5.3. Do not implement 94.8% as a target.
• Swap steering form. The released code scales the decoder column by the residual norm, which changes the effective alpha by the activation norm. Pick one form and state it.
• Two-hop claims. Binding is 100% in two-hop, but swap success stays at 7.86% at the best alpha. Do not read binding as control in multi-hop settings app. K.
• Scope. All results are GPT-2 on templated synthetic biographies with 6 or 20 relations. The authors say the method is model-agnostic but evaluate only GPT-2 [Limitations].
• Missing details. Hardware, training time, and the exact 2-hop SAE schedule are not stated.
• Versions. The card follows arXiv v4 (updated 2026-10-15). Its errata renamed the binding metric to "diagonal accuracy" and removed two earlier definitions, so older TMLR-era text may use other names §4.3.
• Ablation text. The ablation prose gives the full model's unseen diagonal accuracy as 0.891, but the table gives 0.912 §6, app. ablation. Use the table.

Lineage

• Builds on: sparse autoencoders for language models §2. Its supervised value term follows the concept-bottleneck line of work the authors cite §3.
• Compared against: a standard SAE trained with the same architecture on reconstruction and sparsity only §5.2, app. K; a supervised linear probe used for swap steering in the two-hop setting §7.
• Superseded by: not stated.