PolySAE: Modeling Feature Interactions in Sparse Autoencoders via Polynomial Decoding

TL;DR

PolySAE addresses the limitation of standard sparse autoencoders, whose linear decoding prevents faithful modeling of compositional structure in language representations. It introduces a polynomial decoder that adds second- and third-order feature interactions, implemented via a low-rank, shared projection to maintain efficiency and interpretability of the linear encoder. Across four language models and three SAE variants, PolySAE improves probing F1 by about 8% on average while preserving reconstruction quality, and yields 2–10× larger Wasserstein distances, indicating more semantically structured representations; importantly, learned interaction weights show near-zero correlation with co-occurrence, suggesting true compositional binding beyond surface statistics. The work advances mechanistic interpretability by providing a scalable tool to analyze and intervene on compositional structure within large language models, with potential applications in safety, debugging, and model auditing.

Abstract

Sparse autoencoders (SAEs) have emerged as a promising method for interpreting neural network representations by decomposing activations into sparse combinations of dictionary atoms. However, SAEs assume that 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 approximately 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, such as morphological binding and phrasal composition, largely independent of surface statistics.

Figures (4)

• Figure 1: Semantic Dimension Expansion via Feature Interaction. Consider two semantic directions---Famous and Beverage---and their associated learned features Star and Coffee. (a) Additive interactions yield co-occurrence semantics that remain in the original feature span. (b) Multiplicative interactions enable representations to escape this subspace via $z_i \cdot z_j$, lifting into orthogonal dimensions (Brand) to capture emergent concepts like Starbucks."Starbucks" example from \ref{['tab:polysae_second_order']}.
• Figure 2: An overview of PolySAE: (1) sparse latent features are first extracted with a standard SAE encoder. (2) Activations in the residual stream are then reconstructed by modeling 2nd- and 3rd-order interactions in addition to the standard linear component. The example "Investing.com --- Philippines stocks were higher after" comes from \ref{['tab:polysae_third_order']}.
• Figure 3: Probing Mean F1 vs. sparsity $k$. Shaded regions show range across widths (2k--16k). PolySAE consistently outperforms SAE with significant separation at higher $k$.
• Figure 4: Reconstruction MSE for different $R_2$ and $R_3$ values, with $R_1=768$, using activations from GPT-2 Small.