Technical Summary: VASAE – Naming SAE Dictionary Directions with Vocabulary-Aligned Anchoring

1. Problem Statement

Sparse Autoencoders (SAEs) are increasingly used to decompose the high-dimensional residual streams of Transformer models into sparse codes over learned dictionaries, aiding in mechanistic interpretability. However, a significant limitation of standard SAEs is that their learned features (dictionary directions) are not intrinsically connected to the model's token vocabulary.

In standard practice, feature naming is a post hoc process. Researchers must inspect contexts where specific features activate or run separate automated interpretation procedures to assign semantic labels. This disconnect arises because the geometry of the learned dictionary and the geometry of the token embeddings are typically unrelated; a feature can effectively reconstruct activations while lying far from any token embedding vector, rendering it without an "intrinsic" vocabulary-level name.

The paper asks: Can we train SAE features such that they remain flexible enough for high-quality reconstruction but are geometrically anchored close enough to token embeddings to acquire intrinsic token names during training, without sacrificing reconstruction performance?

2. Methodology: Vocabulary-Aligned Sparse Autoencoder (VASAE)

The authors propose VASAE, a modification to the standard SAE training objective that introduces a vocabulary-aligned anchoring mechanism.

Core Architecture

• Input: VASAE operates on the post-residual stream vectors ($h^{(\ell)}$) of a Transformer layer $\ell$.
• Encoder: Uses a standard top-k sparse encoder (linear map + ReLU + Top-K selection) to produce a sparse code $z$.
• Decoder: Uses a learnable dictionary matrix $W_D$ (where columns are feature vectors $d_i$) to reconstruct the residual stream.
• Anchors: The method utilizes the model's fixed input token embeddings ($W_E$) as reference directions. These embeddings share the same dimensionality as the residual stream (or are projected to do so).

The Anchoring Objective

The key innovation is the addition of an anchor loss term to the standard reconstruction loss.

1. Nearest-Token Alignment: For each learned feature $d_i$, the method computes the cosine similarity with all fixed token embeddings $w_v$.
   $$s_i = \max_{v \in V} \cos(d_i, w_v)$$
   The token $v^*_i$ achieving this maximum defines the intrinsic token name for feature $i$.
2. Anchor Loss: The training objective encourages features to stay close to their nearest token embedding:
   $$L_{anchor} = -\frac{1}{d_{sparse}} \sum_{i=1}^{d_{sparse}} s_i$$
   Minimizing this loss maximizes the cosine similarity between features and their nearest vocabulary tokens.
3. Total Loss: The combined objective is:
   $$L_{vasae} = L_{recon} + \lambda_{anchor} L_{anchor}$$
   where $\lambda_{anchor}$ is a hyperparameter controlling the strength of the alignment constraint.

Distinction from Hard-Tying

The authors explicitly distinguish VASAE from a "hard-tied" decoder baseline, where the decoder weights are fixed to exactly equal the token embedding matrix ($W_D = W_E$). Hard-tying forces features to be exact token embeddings, which the paper hypothesizes is insufficient for high-quality reconstruction. VASAE keeps the dictionary learnable but pulls it toward the vocabulary geometry via the anchor loss.

3. Experimental Setup

The method was evaluated on two decoder-only Transformer models:

• GPT-2-small: 12 layers, $d_{model}=768$, vocabulary size $|V|=50,257$.
• Llama-3.1-8B: 32 layers, $d_{model}=4096$, vocabulary size $|V|=128,256$.

Baselines:

1. Standard SAE: Trained with only reconstruction loss.
2. Hard-Tied Decoder: Decoder weights fixed to token embeddings; only the encoder is trained.

Metrics:

• Reconstruction: Mean Squared Error (MSE), Variance Explained (VE), and Cross-Entropy (CE) recovery after residual stream substitution.
• Alignment: Geometric feature-alignment rate (percentage of features with alignment score $s_i \ge 0.8$) and vocabulary coverage.
• Case Studies: Qualitative inspection of intrinsic token names assigned to active features in specific contexts.

4. Key Results

Reconstruction Quality

VASAE achieves reconstruction metrics comparable to the Standard SAE, significantly outperforming the Hard-Tied Decoder baseline.

• GPT-2-small: VASAE and Standard SAE both achieve a Variance Explained (VE) of 0.965 and CE recovery of 0.975.
• Llama-3.1-8B: VASAE and Standard SAE both achieve a VE of 0.931 and CE recovery of 0.906.
• Hard-Tied Baseline: Performs poorly (e.g., negative VE in GPT-2-small), confirming that forcing features to be exact token embeddings destroys reconstruction capability.

Geometric Alignment

• GPT-2-small: With $\lambda_{anchor} = 10^{-4}$, 89–94% of features in layers 0–10 exceed the alignment cutoff ($s_i \ge 0.8$). The alignment is strong and consistent across shallow and middle layers.
• Llama-3.1-8B: Alignment is depth-dependent. At a higher coefficient ($\lambda_{anchor} = 5 \times 10^{-3}$), the shallow layer (L0) achieves 92.8% alignment. However, alignment drops significantly in deeper layers (e.g., L31), suggesting that vocabulary anchoring is more compatible with residual states closer to the embedding space.

Case Studies

Qualitative analysis of intrinsic token names reveals semantic coherence:

• In GPT-2-small, features active in a sentence about "Baker Street" are assigned intrinsic names like located, Street, corner, and around.
• Additional examples show features aligning to adjectives/adverbs (e.g., perfect, feasible), award-related terms (e.g., winner, praise), and self-introduction terms (e.g., fan, name).
• In Llama-3.1-8B, the shallow layer shows similar location-related alignment, while deeper layers show less consistency, mirroring the quantitative alignment metrics.

5. Contributions and Significance

The paper makes three primary contributions:

1. Introduction of VASAE: A soft vocabulary-anchoring objective that assigns learned SAE decoder directions nearest-token names during training, rather than post-hoc.
2. Evaluation of Anchoring: Demonstrating that vocabulary-aligned anchoring preserves reconstruction quality comparable to standard SAEs, while proving that hard-tying the decoder to token embeddings is insufficient for this task.
3. Geometric Analysis: Analyzing embedding-space feature alignment and vocabulary coverage, providing evidence that intrinsic token names act as geometry-derived identifiers.

Significance:
The authors argue that VASAE complements existing post-hoc interpretation methods by connecting learned features to intrinsic token names during training. This provides a "vocabulary-indexed handle" for features, offering a starting point for inspection.

Crucially, the paper maintains a modest stance on the implications of these names:

• Not Explanations: An intrinsic token name is not a full feature explanation. It does not explain why a feature activates or its causal role in the model's computation.
• Indexing Tool: The names serve as preliminary indexing tools. True understanding requires further analysis of how named features participate in circuits, how they connect to other components, and how interventions on them affect behavior.
• Limitations: The alignment is most robust in GPT-2-small and shallow/middle layers of Llama-3.1-8B. The method currently anchors only to input embeddings; for models with untied input/output embeddings, the choice of anchor matrix remains an open question.

In summary, VASAE successfully bridges the gap between the geometric space of learned features and the vocabulary space of the model, enabling features to carry intrinsic token names without compromising their utility for reconstruction.