PaperlayerView exact revision on arXiv

Source-linked AI summary

OrtSAE: Orthogonal Sparse Autoencoders Uncover Atomic Features

Anton Korznikov, Andrey Galichin, Alexey Dontsov, Oleg Rogov, Elena Tutubalina, Ivan Oseledets

arXiv:2509.22033v1cs.LG

TL;DR

SAEs can suffer from feature absorption and composition, limiting the interpretability of their latent representations. OrtSAE penalizes high cosine similarity between features using a chunk-wise strategy. It reports more distinct features, reduced absorption and composition, and improved spurious-correlation removal while remaining on par on other downstream tasks.

  • Problem

    Feature absorption and composition can undermine the interpretability and downstream applicability of SAE latents.

    Source: §1 · p. 1

  • Method

    OrtSAE adds a chunk-wise penalty on high cosine similarity between SAE latents to encourage orthogonal features with linear scaling in latent count.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    OrtSAE yields 9% more distinct features, reduces absorption by 65% and composition by 15%, and improves spurious-correlation removal by 6% while matching other downstream tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    OrtSAE improves feature atomicity while preserving reconstruction fidelity comparable to traditional SAEs.

    Source: §5 · p. 9

Abstract

from arXiv · show

Sparse autoencoders (SAEs) are a technique for sparse decomposition of neural network activations into human-interpretable features. However, current SAEs suffer from feature absorption, where specialized features capture instances of general features creating representation holes, and feature composition, where independent features merge into composite representations. In this work, we introduce Orthogonal SAE (OrtSAE), a novel approach aimed to mitigate these issues by enforcing orthogonality between the learned features. By implementing a new training procedure that penalizes high pairwise cosine similarity between SAE features, OrtSAE promotes the development of disentangled features while scaling linearly with the SAE size, avoiding significant computational overhead. We train OrtSAE across different models and layers and compare it with other methods. We find that OrtSAE discovers 9% more distinct features, reduces feature absorption (by 65%) and composition (by 15%), improves performance on spurious correlation removal (+6%), and achieves on-par performance for other downstream tasks compared to traditional SAEs.

1 INTRODUCTION

SAEs aim to represent model activations as interpretable features, but absorption and composition can undermine that goal. OrtSAE penalizes similarity between latents; experiments report more distinct features, lower absorption and composition, and improved spurious-correlation removal.

1 INTRODUCTION

  • SAEs decompose activations into sparse latent features, but absorption and composition can undermine interpretability and downstream applicability.Absorption occurs when broad features are captured by more specific, token-aligned latents; composition merges independent features into a composite.Source: §1 · p. 1
  • Figure 1 compares explained variance, absorption, composition, and nearest-feature cosine similarity at L0=70; OrtSAEs show lower absorption, composition, and similarity than the traditional SAE variants.Explained variance is marginally lower than BatchTopK SAE.Source: Figure 1 · p. 2
  • OrtSAE penalizes high cosine similarity between SAE latents, using chunk-wise computation to reduce complexity from quadratic to linear in the number of latents.The penalty is designed to promote atomic features without changing the core SAE architecture or significant computational overhead.Source: §1 · p. 2
  • OrtSAE finds 9% more distinct features, reduces absorption by 65% and composition by 15%, and improves spurious-correlation removal by 6% while matching other downstream tasks.At L0=70, the feature results are compared with traditional SAEs; the SAEBench results report on-par performance on other downstream tasks.Source: §1 · p. 2

2 RELATED WORK

Related work develops SAEs to decompose activations into sparse, interpretable features and explores architectures for reconstruction and sparsity. Prior work also identifies absorption and composition, while hierarchical Matryoshka SAEs bring overhead and feature hedging concerns.

2 RELATED WORK

  • SAEs use sparse latent spaces to decompose activations into interpretable features, with later variants targeting reconstruction quality and sparsity.Examples include JumpReLU's learned threshold, TopK selection, and BatchTopK's batch-level constraint.Source: §2 · p. 3
  • Feature absorption and composition remain challenges; MetaSAE is introduced as a promising approach to identify them.Absorption involves broad features being absorbed into specific ones, while composition merges independent features.Source: §2 · p. 3
  • Matryoshka SAE uses nested feature spaces to learn broad and specific features, but its hierarchical design is associated with feature hedging and substantial overhead.The passage also reports degraded reconstruction performance and says its interpretability remains poorly explored.Source: §2 · p. 3
  • OrtSAE explores decorrelation by directly enforcing orthogonality between SAE latents, with performance similar to Matryoshka SAE in mitigating absorption and composition.The passage contrasts this approach with Matryoshka's hierarchical representation.Source: §2 · p. 3

3 ORTHOGONAL SPARSE AUTOENCODERS

The section frames feature absorption and composition as overlapping latent representations, then introduces OrtSAE’s cosine-similarity penalty to encourage more atomic features. It describes chunk-wise computation and an every-fifth-iteration option to reduce computational overhead.

3.1 TRADITIONAL SPARSE AUTOENCODERS

  • Traditional SAEs encode activations into sparse overcomplete latent codes and decode them as sparse combinations of feature vectors.ReLU/L1, TopK, and BatchTopK are variants that use different rules to encourage sparsity.Source: §3.1 · p. 4

3.2 CHALLENGES IN TRAINING SAES

  • Feature absorption and composition arise as distinct sparsity-related failure modes: latents absorb parts of a feature or merge independent features into a composite.Both can produce overlapping representations with higher-than-expected cosine similarities between decoder vectors.Source: §3.2 · p. 4

3.3 ORTSAE TRAINING PROCEDURE

  • OrtSAE adds an orthogonality penalty that penalizes high cosine similarities between SAE latents to encourage more atomic features.The penalty extends the traditional SAE objective.Source: §3.2 · p. 4, §3.3 · p. 4
  • Chunk-wise computation reduces orthogonality-loss complexity from O(m^2) to O(m) by averaging squared maximum pairwise similarities within randomly formed latent chunks.The chunks have equal size proportional to m.Source: §3.3 · p. 5
  • Computing the orthogonality loss every fifth training iteration, with γ multiplied by 5, yields comparable performance with significantly reduced computational overhead.The authors describe this as a further efficiency option.Source: §3.3 · p. 5

4 EXPERIMENTS

Across five experimental sections, the paper evaluates OrtSAE’s setup, core performance, feature atomicity, and downstream utility. OrtSAE improves atomicity and spurious-correlation removal while remaining comparable to baselines on other reported measures.

4.1 EXPERIMENTAL SETUP.

  • The experiments compare OrtSAE with ReLU, BatchTopK, and Matryoshka SAEs, training mainly on layer 12 of Gemma-2-2B and also testing other model layers.The main configuration uses a 65,536-feature latent space, eight L0 sparsity levels, and 500 million OpenWebText tokens.Source: §4.1 · p. 5

4.2 FOUNDATIONAL PERFORMANCE ANALYSIS

  • OrtSAE matches BatchTopK SAE in KL-divergence and interpretability, and has explained variance comparable to BatchTopK while exceeding Matryoshka SAE by 2%.The explained-variance comparison measures reconstruction quality; the Autointerp Score assesses interpretability.Source: §4.2 · pp. 6–7, Figure 3 · p. 6
  • OrtSAE’s MeanCosSim is 2.7 times lower than BatchTopK’s and 1.5 times lower than Matryoshka’s, approaching random-initialization levels.MeanCosSim measures each decoder feature’s cosine similarity to its closest neighboring feature.Source: §4.2 · p. 6

4.3 ATOMICITY ANALYSIS

  • At L0 of 70, OrtSAE’s composition rate is 0.06 lower and its absorption rate is 0.17 lower than BatchTopK SAE; both measures approach Matryoshka SAE.The composition comparison uses MetaSAE-based explained variance, where lower values indicate higher atomicity.Source: §4.3 · p. 7, Figure 4 · p. 7
  • OrtSAE has lower feature clustering than BatchTopK SAE and matches Matryoshka SAE, while retaining 9% unique features versus BatchTopK’s 1.5%.A feature counts as unique when all cross-model cosine similarities are below 0.2.Source: Figure 4 · p. 7, §4.3 · p. 8
  • In a qualitative decomposition, one BatchTopK “Queen terms” feature separates into OrtSAE features for “Queen terms” and for “titles and royal concepts”.The examples were selected for cosine similarity above 0.95 and coefficients of at least 0.1.Source: §4.3 · p. 8, Figure 5 · p. 8

4.4 DOWNSTREAM BENCHMARKS

  • OrtSAE outperforms traditional SAEs on spurious correlation removal, improves modestly on targeted probe perturbation, and is comparable on sparse probing and RAVEL.Matryoshka SAE has the highest absolute scores on spurious correlation removal and targeted probe perturbation.Source: Figure 6 · p. 9, §4.4 · p. 9

5 CONCLUSION

The conclusion presents OrtSAE as an orthogonality-constrained SAE training approach that improves atomicity while preserving downstream utility. It reports reduced absorption and composition, more distinct features, and improved spurious correlation removal with minimal computational overhead.

5 CONCLUSION

  • OrtSAE’s chunk-wise orthogonalization penalty scales linearly with feature count and achieves atomicity gains with minimal computational overhead.Across model families, the conclusion reports 9% more distinct features, superior spurious correlation removal (+6%), and on-par performance on other SAEBench tasks.Source: §5 · p. 9

6 REPRODUCIBILITY STATEMENT

The reproducibility statement describes documentation, planned code release, public data, and supplementary experiments supporting replication of the work.

6 REPRODUCIBILITY STATEMENT

  • The authors document experimental details and plan to release code, hyperparameters, training and evaluation tools, and SAEBench integration publicly upon acceptance.OpenWebText is publicly available, and the statement points to supplementary experiments on other architectures and layers and hyperparameter ablations.Source: §6 · p. 10

A ADDITIONAL DETAILS OF SAE TRAINING SETUP

The training appendix details SAE configurations and shared training conditions across models and sparsity levels. OrtSAE uses eight chunks for its 65,536-dimensional latent space.

A ADDITIONAL DETAILS OF SAE TRAINING SETUP

  • Training uses 500 million OpenWebText tokens, a 1,024-token context, AdamW, and one H100 GPU; one SAE takes approximately 10 hours.The main setup uses Gemma-2-2B layer 12, with additional experiments on layer 20 of Gemma-2-2B and Llama-3-8B.Source: A ADDITIONAL DETAILS OF SAE TRAINING… · p. 12

B ADDITIONAL EXPERIMENTS ON GEMMA-2-2B AND LLAMA-3-8B

At L0=70, experiments on layer 20 of Gemma-2-2B and Llama-3-8B confirm reductions in feature absorption and composition and improved Spurious Correlation Removal; OrtSAE also has higher explained variance than Matryoshka SAE on Llama-3-8B.

B ADDITIONAL EXPERIMENTS ON GEMMA-2-2B AND LLAMA-3-8B

C EFFECT OF NUMBER OF CHUNKS ON ORTSAE PERFORMANCE

OrtSAE retains its reconstruction and atomicity performance across chunk counts, while computing orthogonality loss less often substantially reduces its overhead. The section also defines the KL-divergence score used to assess predictive preservation.

C EFFECT OF NUMBER OF CHUNKS ON ORTSAE PERFORMANCE

D KL-DIVERGENCE SCORE DEFINITION

  • The KL-divergence score compares next-token distributions from the original model, SAE reconstructions, and zeroed activations; it ranges from 0 for no improvement over zero ablation to 1 for perfect reconstruction.Source: D KL-DIVERGENCE SCORE DEFINITION · pp. 13–14

E LOGLOSS RESULTS

LogLoss evaluation on layer 12 of Gemma-2-2B shows that OrtSAE preserves predictive performance comparably to BatchTopK SAE, with both outperforming ReLU SAE and Matryoshka SAE.

E LOGLOSS RESULTS

  • Across L0 values 40, 70, 100, and 130, OrtSAE maintains predictive performance comparable to BatchTopK SAE, and both outperform ReLU SAE and Matryoshka SAE on LogLoss.Lower LogLoss indicates better preservation of the base language model’s predictive performance.Source: E LOGLOSS RESULTS · p. 14, Table 4 · p. 14

F FEATURE INTERPRETABILITY METRICS DETAILS

The section describes how feature interpretability and composition are evaluated, defines graph-based clustering measures, and gives qualitative decompositions of composite features into OrtSAE components.

F FEATURE INTERPRETABILITY METRICS DETAILS

  • Autointerp Score is the accuracy of a separate GPT-4o-mini model predicting feature activation from generated descriptions, tested on 100 sequences per feature.OrtSAE’s interpretability is reported as comparable to BatchTopK and Matryoshka SAEs across sparsity levels.Source: F FEATURE INTERPRETABILITY METRICS… · p. 14

G METASAE-BASED FEATURE COMPOSITION METRICS DETAILS

  • Lower MetaSAE explained variance indicates greater atomicity: the score measures how much of the primary SAE decoder-weight variance a MetaSAE reconstructs.The MetaSAE decomposes primary SAE feature vectors into sparse combinations of meta-features.Source: G METASAE-BASED FEATURE COMPOSITION… · p. 15

H CLUSTERING COEFFICIENT DEFINITION

  • In the feature graph, cosine-similar decoder-feature pairs above a threshold form edges; edge density varies with the threshold, and clustering coefficient ranges from 0 to 1.A triangle is three fully connected nodes; a connected triple has three nodes linked by at least two edges.Source: H CLUSTERING COEFFICIENT DEFINITION · p. 15

I ADDITIONAL QUALITATIVE EXAMPLES

Related papers

This paper cites

Loading 2509.22033v1…