PolySAE: Modeling Feature Interactions in Sparse Autoencoders via Polynomial Decoding

AI summary

Overview

Research area: Mechanistic interpretability of neural networks, specifically sparse autoencoders (SAEs) and compositional feature representation. Technical level: Advanced. Scope: This paper introduces PolySAE, a polynomial-decoding extension to SAEs that models higher-order feature interactions to capture compositional structure while preserving an interpretable linear encoder.

What This Paper Is About

Sparse autoencoders explain neural network activations by decomposing them into sparse, additive dictionary features. This additive linear assumption cannot represent compositional meaning: for example, a model cannot tell whether "Starbucks" comes from combining "star" and "coffee" features or from their mere co-occurrence. PolySAE aims to fix this by adding polynomial interaction terms to the SAE decoder so compound concepts can be decomposed into interacting constituent features rather than monolithic units.

Key Contributions

1. Introduces PolySAE, an SAE variant whose decoder includes higher-order polynomial terms for pairwise and triple feature interactions while keeping the encoder linear for interpretability.
2. Uses low-rank tensor factorization on a shared projection subspace to capture these interactions with minimal parameter overhead, reported as 3% on GPT-2.
3. Evaluates across four language models and three SAE variants, showing improved probing F1 and class-conditional feature separation without worse reconstruction.
4. Demonstrates that learned interaction weights are largely independent of co-occurrence frequency and that interaction directions causally steer model outputs toward compositional semantics.

Main Findings

• Improved probing performance: PolySAE improves probing F1 by about 8% on average across four language models and three SAE variants, while maintaining comparable reconstruction error.
• Better class separation: It produces 2–10× larger Wasserstein distances between class-conditional feature distributions, indicating more distinct feature representations for different concepts.
• Not just co-occurrence: Learned interaction weights correlate only weakly with co-occurrence frequency (r = 0.06), compared with a strong correlation for standard SAE feature covariance (r = 0.82). This suggests polynomial terms capture compositional structure beyond surface statistics.
• Causal steering: The learned interaction directions can causally steer model outputs toward the corresponding compositional semantics.
• Efficiency: Low-rank tensor factorization limits the extra parameters, making interaction modeling practical for large language models.

Methodology in Plain English

Standard SAEs learn a dictionary of features and reconstruct an activation by adding up active feature vectors. PolySAE keeps this feature detection step linear but changes the reconstruction step. Instead of only adding individual features, it also includes products of feature activations to represent two-feature and three-feature interactions. For example, it can represent a compound concept as the combination of "star" and "coffee" features, not just their separate presence. To avoid an explosion in parameters, the model compresses interaction tensors using low-rank factorization and a shared projection subspace. The researchers then train and compare PolySAE against standard SAE variants on several language models, measuring probing accuracy, reconstruction quality, distribution separation, correlation with co-occurrence, and causal effects of interaction directions.

Why This Matters

This work challenges a core simplifying assumption in SAE-based interpretability: that features combine additively. If SAEs can model interactions, they may reveal more faithful compositional features and circuits in neural networks, improving both scientific understanding and practical control of models.

Real-world applications:

• LLM auditing and debugging: Better identification of how compound concepts are represented can help diagnose unwanted associations or failures.
• Controllable text generation: Interaction directions offer a way to steer model outputs toward specific compositional meanings.
• Concept discovery and knowledge editing: Decomposing compound concepts into constituents could support more precise editing of model knowledge.
• Safety and monitoring: More separable class-conditional feature distributions can help detect harmful or biased concept representations.

Industry relevance: The method targets AI safety, model interpretability tooling, and LLM product teams that need reliable ways to inspect, control, and audit large models. Its low parameter overhead makes it more practical for deployment in existing interpretability pipelines.

Future Directions

• Scaling PolySAE to larger models, longer contexts, and more features while controlling parameter and compute costs for higher-order interactions.
• Extending polynomial decoding beyond language to vision, multimodal, and reinforcement learning representations.
• Developing stronger benchmarks and theoretical guarantees for distinguishing true compositional structure from co-occurrence and other surface statistics.
• Integrating learned interaction directions into practical model editing, steering, and safety intervention workflows.

Target Audience

This paper is best suited for mechanistic interpretability researchers, AI safety researchers, and machine learning practitioners working with sparse autoencoders or LLM representation analysis. It also benefits graduate students and advanced undergraduates with a background in neural networks, representation learning, or natural language processing who want to understand compositional feature modeling.

Authors’ abstract

Sparse autoencoders (SAEs) interpret neural network representations by decomposing activations into sparse combinations of dictionary atoms. However, SAEs assume features combine additively through linear reconstruction, an assumption that cannot capture compositional structure: linear models cannot distinguish whether ''Starbucks'' arises from the composition of ''star'' and ''coffee'' features or merely their co-occurrence. This forces SAEs to allocate monolithic features for compound concepts rather than decomposing them into interpretable constituents. We introduce PolySAE, which extends the SAE decoder with higher-order terms to model feature interactions while preserving the linear encoder essential for interpretability. Through low-rank tensor factorization on a shared projection subspace, PolySAE captures pairwise and triple feature interactions with small parameter overhead (3% on GPT2). Across four language models and three SAE variants, PolySAE achieves an average improvement of $\sim$8% in probing F1 while maintaining comparable reconstruction error, and produces 2--10$\times$ larger Wasserstein distances between class-conditional feature distributions. Critically, learned interaction weights exhibit negligible correlation with co-occurrence frequency ($r = 0.06$ vs $r = 0.82$ for SAE feature covariance), suggesting that polynomial terms capture compositional structure largely independent of surface statistics. Finally, the learned interaction directions causally steer model outputs toward the corresponding compositional semantics.