Cross-seed explainability using Procrustes-conditioned Joint End-to-end Top-K Sparse Autoencoders
Procrustes-conditioned joint end-to-end Top-K sparse autoencoder extracts universal features across independent BERT seeds, achieving Pearson r≥0.70.
Key Findings
Methodology
This paper introduces a Procrustes-conditioned joint end-to-end Top-K sparse autoencoder (SAE) designed to extract cross-seed universal features from independently trained BERT models. By computing an orthogonal Procrustes rotation matrix before joint training, the method aligns the activation spaces of different seeds, addressing the non-convexity of dictionary learning that causes feature misalignment. The model employs Top-K sparsity constraints, combined with end-to-end optimization and an auxiliary dead-feature revival loss, to enhance interpretability and feature universality. Experiments on five seed pairs across SST-2, Stanford Politeness, and TweetEval datasets demonstrate that this approach yields features with high cross-seed correlation (Pearson r≥0.70), outperforming post-hoc alignment baselines, and capturing interpretable sociolinguistic patterns.
Key Results
- The Procrustes-conditioned joint SAE achieved an average top-10 feature Pearson correlation of 0.74 across datasets, significantly higher than independent SAEs (≈0.005), indicating robust cross-model feature alignment.
- The percentage of universal features (r≥0.70) increased to 57.9% with the full pipeline, surpassing joint training without rotation (44.9%), confirming the effectiveness of spatial alignment.
- Qualitative analysis revealed that high-universality features encode interpretable social linguistic patterns, such as politeness and modality, demonstrating the model’s capacity for mechanistic interpretability.
Significance
This work addresses a fundamental challenge in mechanistic interpretability: the misalignment of features learned by independently trained models. By integrating Procrustes spatial alignment into joint SAE training, it provides a scalable, effective solution for extracting consistent, interpretable features across different model seeds. This advances our understanding of large language models’ internal representations, facilitating research in model transparency, knowledge transfer, and fairness. The methodology’s simplicity and effectiveness make it a promising tool for future mechanistic studies, bridging the gap between black-box models and human-understandable concepts.
Technical Contribution
The paper introduces a novel combination of Procrustes orthogonal rotation with joint end-to-end Top-K sparse autoencoder training. Unlike prior post-hoc alignment methods, this approach performs spatial calibration during training, ensuring feature space consistency. It leverages the sparsity constraint to prevent feature splitting, while the orthogonal Procrustes rotation aligns activation spaces without task-distorting penalties. The integration of multiple loss components—KL divergence, local reconstruction, cross-model sparse code discrepancy, and dead neuron revival—enables robust extraction of universal, interpretable features. The approach simplifies the pipeline by avoiding multiple SAE training and complex regularizations, offering both theoretical and practical advantages.
Novelty
This study is the first to incorporate Procrustes orthogonal rotation directly into the joint training of autoencoders for cross-seed feature alignment. Unlike existing post-hoc or multi-model regularization methods, it performs spatial calibration during training, ensuring high feature correlation and interpretability. The method effectively combines spatial alignment with sparse, end-to-end optimization, setting a new standard for cross-model mechanistic interpretability in large language models.
Limitations
- The validation is limited to BERT models; applicability to other architectures like GPT or T5 remains to be tested, as different model structures may exhibit different activation space properties.
- Feature interpretability relies on qualitative analysis; automated semantic validation or quantitative metrics for interpretability are still under development.
- Computational overhead increases with the number of model pairs and layers, especially during Procrustes alignment, which may limit scalability for very large models or datasets.
Future Work
Future research will extend this framework to diverse model architectures and larger datasets, aiming to automate semantic validation of features. Developing more efficient spatial alignment algorithms and integrating with other interpretability tools will further enhance robustness. Additionally, applying this method to real-world tasks such as model debugging, knowledge transfer, and fairness auditing could significantly impact AI transparency and trustworthiness.
AI Executive Summary
Understanding the internal mechanisms of large language models remains a significant challenge in AI research. Despite their impressive performance, models like BERT are often regarded as black boxes, with their learned representations difficult to interpret consistently across different training runs. Variability introduced by random initialization leads to misaligned feature spaces, hampering efforts to extract universal, human-understandable concepts. Existing alignment techniques, such as post-hoc Procrustes rotation or multi-model regularization, offer partial solutions but often fall short of providing a seamless, end-to-end approach.
This paper proposes a novel solution: a Procrustes-conditioned joint end-to-end Top-K sparse autoencoder (SAE). The core idea is to compute an orthogonal Procrustes rotation matrix before training, aligning the activation spaces of independently trained models. This spatial calibration ensures that the features learned are directly comparable, significantly improving the correlation of extracted features across seeds. The method employs a Top-K sparsity constraint to enforce interpretability, coupled with a multi-objective loss function that includes KL divergence, local reconstruction, cross-model sparse code discrepancy, and dead neuron revival. This comprehensive approach allows the model to learn highly universal features that encode meaningful sociolinguistic patterns.
Extensive experiments on SST-2, Stanford Politeness, and TweetEval datasets demonstrate that the proposed pipeline outperforms baseline methods, achieving Pearson correlations above 0.70 for cross-seed features and increasing the proportion of universal features. Qualitative analyses confirm that these features encode interpretable social cues, such as politeness and modality, providing insights into the model’s internal semantics. The approach’s simplicity, combined with its effectiveness, marks a significant step forward in mechanistic interpretability, offering a scalable and robust tool for understanding large language models.
Looking ahead, future work will explore broader model architectures, optimize computational efficiency, and develop automated semantic validation tools. The ultimate goal is to facilitate transparent, trustworthy AI systems capable of explaining their internal reasoning in human-understandable terms, thereby fostering greater trust and utility in AI applications.
Deep Dive
Glossary
Procrustes旋转 (Procrustes rotation)
一种线性变换，用于最大程度地对齐两个空间的点集，确保它们的形状一致。
用于校正不同模型激活空间的空间错位。
Top-K稀疏性 (Top-K sparsity)
在每个位置只保留最大的K个激活值，减少冗余，提高可解释性。
自编码器中的稀疏约束策略。
机制解释 (Mechanistic interpretability)
理解模型内部表示和操作机制的方法，揭示模型“思考”的具体过程。
本文旨在通过特征对齐实现机制解释。
Pearson相关系数 (Pearson correlation coefficient)
衡量两个变量线性相关程度的统计指标，值范围-1到1。
用于评估不同模型中特征的相似性。
稀疏自编码器 (Sparse Autoencoder)
通过稀疏约束学习稀疏表示的自编码网络，用于特征提取和机制理解。
本文采用Top-K稀疏SAE进行特征提取。
Open Questions Unanswered questions from this research
- 1 如何将Procrustes空间校正推广到多层、多任务模型中，保持特征一致性和语义解释性仍是挑战。
- 2 自动化机制验证工具尚未成熟，定量评估机制特征的可解释性仍需突破。
- 3 在大规模模型和多任务场景下，空间校正的计算成本和效率问题亟待优化。
Applications
Immediate Applications
模型机制调试
帮助研究者理解模型内部表示，识别潜在偏差或错误，提升模型透明度。
知识迁移与模型对齐
实现不同模型之间的知识迁移，提升多模型系统的协同能力。
Long-term Vision
公平性与偏差检测
通过机制特征对比，识别模型中的偏差和不公平现象，推动公平AI发展。
Abstract
We present a Procrustes-conditioned Joint End-to-end Top-K Sparse Autoencoder (SAE) for extracting cross-seed universal features from independently trained BERT models. Cross-seed feature universality is a fundamental challenge in mechanistic interpretability: because dictionary learning is non-convex, independently trained networks learn misaligned feature spaces, so apparently identical features may differ by random initialization. We address this by computing an orthogonal Procrustes rotation between seeds' activation spaces before joint SAE training, combining Top-K sparsity, end-to-end downstream optimization, and an auxiliary dead-feature revival loss based on previous SAE literature. Evaluating on five independent seed pairs (ten BERT models) across three benchmark datasets (SST-2, Stanford Politeness, TweetEval Emotion), our full pipeline produces more universal features (Pearson r $\geq$ 0.70 across seeds) than post-hoc alignment baselines on all three datasets. A minimal qualitative analysis confirms that high-universality features encode interpretable sociolinguistic patterns.