VASAE: Naming SAE Dictionary Directions with Vocabulary-Aligned Anchoring
VASAE introduces vocabulary-aligned anchoring to train SAE features with intrinsic token names, maintaining reconstruction quality and achieving over 90% feature-token alignment.
Key Findings
Methodology
VASAE extends standard SAE by incorporating a vocabulary-aligned loss that encourages each learned feature to be close to its nearest token embedding in cosine similarity. The training jointly optimizes reconstruction error and geometric alignment, using a combined loss function with hyperparameter λanchor. The process involves encoding residual streams via a top-k sparse encoder, reconstructing with a learnable dictionary, and updating features to maximize similarity with fixed token embeddings. Experiments on GPT-2-small and Llama-3.1-8B validate high alignment rates (>90%) across layers, with features automatically named by their nearest tokens, providing interpretable, geometry-based feature labels.
Key Results
- In GPT-2-small, 0-10 layers exhibit approximately 90% of features with cosine similarity scores above 0.8, indicating strong geometric alignment. The variance explained (VE) remains at 0.965, comparable to standard SAE, with minimal performance loss. Llama-3.1-8B shows 92.8% alignment in shallow layers, though deeper layers display less stability. Case studies confirm that features assigned to tokens are contextually relevant, and the alignment improves interpretability without sacrificing reconstruction quality.
- Compared to baseline SAE and fixed token embedding methods, VASAE significantly enhances feature-to-token alignment while maintaining similar reconstruction metrics. Ablation studies demonstrate that the geometric alignment loss is crucial for high-quality token mapping. The results highlight the method’s robustness across different model sizes and layers, especially in early and middle layers where residual states are closer to embedding space.
- Qualitative analysis reveals that features with high alignment scores correspond to meaningful tokens in context, and subtracting sentence-level mean sparse codes further refines token relevance. These findings suggest that embedding-space geometry can serve as a reliable basis for intrinsic feature naming, facilitating interpretability and analysis of internal model states.
Significance
This work advances the interpretability of Transformer models by bridging the gap between learned residual features and the intrinsic vocabulary. By embedding features within the vocabulary’s geometric space, VASAE enables automatic, meaningful naming of internal features, which enhances transparency and trustworthiness. The approach addresses a key challenge in model explainability—linking high-dimensional internal states to human-understandable concepts—thus opening new avenues for debugging, bias detection, and controlled model behavior. Its compatibility with existing models and minimal impact on reconstruction performance make it a practical tool for both research and applied AI systems, fostering deeper insights into how language models encode knowledge and semantics.
Technical Contribution
VASAE’s main technical innovation lies in integrating a vocabulary-aligned loss into SAE training, leveraging fixed token embeddings as geometric anchors. This approach ensures features are not only sparse and reconstructive but also geometrically close to specific tokens, enabling automatic naming. The method combines sparse coding, cosine similarity-based alignment, and joint optimization, resulting in a dictionary that is both interpretable and effective. Unlike prior post-hoc labeling or hard tying, VASAE’s soft alignment preserves flexibility and improves the semantic coherence of features, offering a new paradigm for internal model interpretability grounded in geometric principles.
Novelty
This paper is the first to embed vocabulary-based geometric constraints directly into the training of SAE features, enabling automatic, intrinsic token naming during learning. Unlike previous post-hoc interpretability methods, VASAE integrates the alignment into the optimization process, resulting in features that are both reconstructively effective and semantically meaningful. The use of cosine similarity as a geometric measure for feature-to-token alignment in residual streams is novel, providing a scalable and principled approach to interpretability that aligns with the intrinsic structure of language embeddings.
Limitations
- Deep-layer features tend to show less stable alignment, possibly due to the increasing semantic abstraction and divergence from initial embedding space. This limits the method’s effectiveness in the deepest layers.
- Reliance on fixed token embeddings assumes their geometric structure accurately reflects semantic relationships, which may not hold in cases of embedding bias or domain shift.
- Computational cost increases with vocabulary size, as similarity calculations are expensive, especially for large vocabularies like Llama-3.1-8B. Further optimization is needed for real-time or large-scale deployment.
Future Work
Future research will focus on extending geometric alignment techniques to multimodal models, incorporating context-aware embeddings, and improving stability in deeper layers. Exploring dynamic or adaptive anchor sets, integrating with other interpretability tools, and reducing computational overhead are promising directions. Additionally, applying this framework to tasks like bias detection, controllable generation, and model editing could significantly enhance AI transparency and safety.
AI Executive Summary
Transformer models, especially in language processing, rely heavily on internal residual streams that encode complex semantic information. Despite their success, understanding these high-dimensional states remains challenging. Traditional interpretability methods often rely on post-hoc labeling, which lacks direct connection to the model’s vocabulary. This gap limits the transparency and controllability of language models. To address this, the paper introduces VASAE, a novel approach that incorporates vocabulary-aligned anchoring into the training of sparse autoencoders. By leveraging fixed token embeddings as geometric anchors, VASAE encourages each learned feature to align with a specific token in the vocabulary space, enabling automatic and meaningful feature naming during training.
The core idea is to optimize a combined loss function that balances reconstruction accuracy with geometric proximity to token embeddings. This is achieved through a sparse coding framework, where each residual stream is represented by a small set of features, and each feature is encouraged to be close to its nearest token embedding in cosine similarity. The training process involves separate models for each layer, ensuring layer-specific interpretability. Extensive experiments on GPT-2-small and Llama-3.1-8B demonstrate that over 90% of features in early and middle layers are strongly aligned with their nearest tokens, with minimal performance trade-offs.
Case studies confirm that features assigned to tokens are contextually relevant, and subtracting sentence-level mean sparse codes refines their semantic relevance. The results highlight that vocabulary-aligned anchoring not only preserves the model’s reconstructive ability but also provides a geometric basis for interpretability. This approach offers a scalable, automatic method to connect learned features with human-understandable concepts, advancing the field of model transparency. Despite some instability in deep layers, the method’s effectiveness in shallow and middle layers marks a significant step toward interpretable AI systems, with promising directions for future multimodal and large-scale applications.
Deep Dive
Abstract
Sparse autoencoders (SAEs) provide useful decompositions of Transformer residual streams, but their learned features are usually named post hoc rather than directly connected to the Transformer's token vocabulary. We introduce Vocabulary-Aligned Sparse Autoencoder (VASAE), a method that trains SAE features under vocabulary-aligned anchoring and assigns each feature an intrinsic token name: the token string whose embedding is nearest to that feature. Without reducing reconstruction quality compared with a standard SAE, VASAE produces dictionaries with vocabulary-aligned features. Using a 0.8 cutoff on the nearest-token alignment score, dictionaries trained on GPT-2-small post-residual streams align about 90% of features in layers 0--10. In Llama-3.1-8B, representative shallow and middle-layer dictionaries contain strongly aligned features, including 92.8% in the shallow layer, while the representative final-layer dictionary shows limited alignment. After subtracting the sentence-level mean sparse code, case studies show that many remaining intrinsic token names are relevant to nearby input tokens. These results suggest that vocabulary-aligned anchoring can connect learned features to intrinsic token names during training, complementing post hoc interpretation of learned dictionaries.