PolySAE: Modeling Feature Interactions in Sparse Autoencoders via Polynomial Decoding

• SAEs support mechanistic interpretability and safety steering, but additive reconstruction cannot distinguish “Starbucks” composition from co-occurrence or reveal its constituent-level composition.This conflation obscures how networks build complex representations from simpler parts.Source: §1 · pp. 1–2
• PolySAE preserves a linear encoder and adds quadratic and cubic decoder interactions via low-rank tensor factorization on a shared projection subspace, with 3% overhead on GPT2 small.Source: §1 · pp. 2–3
• Setting interaction coefficients to zero recovers vanilla SAE behavior, and PolySAE applies to TopK (Gao et al., 2025), BatchTopK (Bussmann et al., 2024), and Matryoshka (Bussmann et al., 2025).Source: §1 · p. 3
• Across four language models and three sparsification strategies, PolySAE achieves an average 8% F1 improvement while maintaining comparable reconstruction error.Source: §1 · p. 3
• PolySAE produces 2–10× larger Wasserstein distances between class-conditional feature distributions, indicating more separated semantic structure.Source: §1 · p. 3
• Interaction weights show negligible correlation with co-occurrence frequency (r = 0.06 vs. r = 0.82 for SAE feature covariance), and interaction directions causally steer outputs toward compositional semantics.Qualitative examples cover morphological binding, phrasal composition, and contextual disambiguation.Source: §1 · p. 3

Sparse dictionary learning

• SAEs have become a leading LLM feature-discovery paradigm, but their architectural variants still assume additive linear reconstruction.The field has scaled to millions of features and developed variants including BatchTopK, Matryoshka, Gated, and JumpReLU SAEs.Source: §2 · p. 3

Modeling feature interactions

• Prior interpretability work models pairwise feature interactions with Bilinear MLPs (Pearce et al., 2025) or tensor-network architectures; PolySAE extends interaction modeling to SAEs.Multiplicative interactions also have a broader history in deep learning, including bilinear models and gating mechanisms.Source: §2 · p. 3

Polynomials

• Unlike BAE (Dooms & Gauderis, 2025), which models pairwise input-neuron interactions, PolySAE models pairwise and higher-order interactions among learned sparse features while preserving linear SAE latents.This explicitly allocates capacity to non-additive feature composition.Source: §2 · p. 3

3.1. Preliminaries

• Because additive SAE reconstruction cannot capture joint effects that require a distinct direction, PolySAE extends the decoder while preserving a linear, interpretable encoder.The passage illustrates this limitation with “coffee” and “star” features whose joint activation may need a separate “Starbucks” reconstruction direction.Source: §3.1 · p. 4

3.2. Design Principles for Feature Interactions

• PolySAE’s design keeps encoding linear, uses polynomial reconstruction for composition, aligns interactions to a shared low-dimensional feature subspace, and constrains higher-order capacity and directions.The latent interaction subspace is orthonormal to avoid rotational ambiguity and redundant capacity in correlated interaction directions.Source: §3.2 · p. 4

3.3. PolySAE: Polynomial Sparse Autoencoder

• Quadratic and cubic decoder terms extend the standard SAE, which PolySAE recovers when both interaction coefficients are set to zero.The polynomial terms operate on sparse codes, with learnable scalar coefficients controlling each order.Source: §3.3 · p. 4
• The implementation shares a projection U across interaction orders and uses nested ranks R1 ≥ R2 ≥ R3; orthonormality is enforced by QR retraction after each gradient step.In practice, R2 = R3 ≪ R1 (for example, R2 = R3 = 64) is reported to suffice for capturing most interaction structure.Source: §3.3 · p. 5

3.4. Discussion

• PolySAE makes feature contributions context-dependent: its linear, quadratic, and cubic terms define dictionaries over individual features, pairs, and triples, respectively.The low-rank computational form avoids explicitly materializing the d_sae^2-dimensional interaction dictionaries.Source: §3.4 · p. 5
• Using the same d_sae base features as an SAE, shared R2 and R3 modes support many feature-pair and feature-triple compositions without increasing the number of features.Interactions are expressed through reusable shared modes rather than independent parameters for each combination.Source: §3.4 · p. 5
• For GPT-2 small, PolySAE adds ∼2.5–3% of the full SAE’s parameters; with R2 = R3 and R2 ∈ [0.06R1, 0.11R1], ∆P = (1.12–1.22)d^2 up to constants.With a full-rank linear term, the stated parameter increment is ∆P = d^2 + d(R2 + R3) + 2.Source: §3.4 · p. 5

4.1. Experimental Setup

• The evaluation uses K=64 and 16,384 latents on residual-stream activations from four pretrained language models, comparing standard SAEs with TopK, BatchTopK, and Matryoshka sparsification.Training uses 500M tokens (300M for GPT-2 Small) with context length 128; evaluation includes reconstruction metrics and six classification tasks.Source: §4.1 · p. 5

4.2. Reconstruction and Semantic Modeling

• 8% average probing gains across Pythia-410M, Pythia-1.4B, and Gemma2-2B, with gains of more than 10% on GPT-2, while PolySAE maintains comparable MSE to SAE.The comparison covers four language models and three sparsification strategies.Source: Table 1 · p. 6, §4.2 · p. 6
• 2–10× higher Wasserstein distances show stronger class-conditional separation, while CE recovery stays within 0.003 of vanilla SAE across all 12 configurations.The Wasserstein improvements are reported across models other than GPT-2; the CE comparison indicates no functional degradation.Source: Table 1 · p. 6, §4.2 · p. 6
• Evaluation combines decoder-output MSE and cross-entropy recovery with feature-level logistic-regression probing and class-conditional 1-Wasserstein distance.Probing uses the feature with the largest mean activation difference between positive and negative classes; Wasserstein distance measures global distributional separation.Source: §4.2 · p. 6

4.3. Ablations

• Low-rank factorization cuts parameters by 65% at a 1pp F1 cost, while orthogonality recovers the gap and exceeds the original polynomial by 2.9pp at zero parameter cost.These architectural ablations were evaluated on GPT-2 Small.Source: §4.3 · p. 6, Table 2 · p. 7
• Across sparsity levels, PolySAE outperforms standard SAEs, with a wider gap at higher k and lower variance across dictionary widths, enabling competitive performance with fewer active features.Figure 3 shows probing F1 across active-feature counts k and dictionary widths of 2k–16k.Source: Figure 3 · p. 7, §4.3 · p. 7
• PolySAE has smaller K=1-to-K=5 F1 gains in 9 of 12 model–sparsifier configurations, including gaps of −7 to −8 for all three GPT-2 Small sparsifiers.The authors suggest higher-order interactions absorb contextual variability, leaving linear features more semantically focused.Source: §4.3 · p. 7

5.1. Making Sense of Learned Interactions

• r = 0.06 versus r = 0.82: PolySAE’s learned interactions had negligible correlation with co-occurrence frequency, unlike SAE activation covariance.The analysis compares each learned quadratic interaction capacity with empirical feature-pair co-occurrence frequency.Source: §5.1 · p. 7
• Second-order interactions capture compositions such as coffee × star yielding Starbucks, while third-order bindings use added context to disambiguate meanings.Examples include distinguishing financial investing from unrelated “-ing” usages; higher-order terms absorb contextual variation that would otherwise fragment linear features into context-specific atoms.Source: §5.1 · pp. 7–8, Table 4 · p. 8, Table 5 · p. 8
• 12%: 8,550 of the 70,000 evaluated feature pairs scored above 0.9 for compositional interpretability, adding at least 8,550 second-order concepts beyond the linear dictionary.The evaluated pairs were selected from those with non-negligible interaction strength.Source: §5.1 · p. 8

5.2. Steering with Learned Interactions

• 71.0%: PolySAE lowered target rank versus no steering in 230/324 pairs, with 3 degradations; it improved over no steering in 27/27 concepts and vanilla SAE in 21/27, with +41.5 mean rank improvement versus vanilla.For example, steering with [canada]×[oil] produced “Keystone XL” under PolySAE; Table 6 shows selected GPT-2 Small examples.Source: §5.2 · p. 9, Table 6 · p. 9
• 0.372 ± 0.093 versus 0.311 ± 0.158: PolySAE’s mean cosine similarity to ground-truth compositional directions was higher than vanilla SAE’s, a 19.7% relative improvement.PolySAE won on 11/27 concepts, vanilla SAE on 2/27, and 14 were ties.Source: §5.2 · p. 9

6. Conclusion

• Across four LLMs and three SAE variants, PolySAE improves probing F1 by 8% on average, maintains comparable reconstruction, and achieves 2–10× larger Wasserstein distances between class-conditional distributions.Interaction weights correlate weakly with co-occurrence (r = 0.06), and learned interaction directions causally steer outputs toward corresponding compositional semantics.Source: §6 · p. 9
• PolySAE uses low-rank tensor factorization on a shared projection subspace to model pairwise and triple feature interactions with small parameter overhead while preserving a linear encoder.Its higher-order decoder terms are intended to capture compositional structure rather than co-occurrence.Source: §6 · p. 9
• The study covers models up to 2B parameters and restricts experiments to forced-sparsity SAE variants.Source: §6 · p. 9

Impact Statement

• PolySAE models non-additive feature interactions while preserving linear, human-interpretable features, supporting mechanistic analysis and the long-term goal of more transparent, controllable, and trustworthy systems.The authors characterize the work's societal impact as net positive.Source: Impact Statement · p. 10
• The method may support model auditing, debugging, and safety research by helping researchers identify, analyze, and intervene on meaningful internal representations.The authors describe these as potential downstream uses of improved interpretability.Source: Impact Statement · p. 10
• PolySAE is intended as an analysis tool, not a deployment-facing component, and does not directly increase language-model capabilities.The paper identifies methodological and scientific benefits as its primary anticipated benefits.Source: Impact Statement · p. 10
• Insights into internal representations could be misused to manipulate model behavior, but the authors identify no novel or unique risks beyond those in interpretability research.Source: Impact Statement · p. 10

A. PolySAE Algorithm and Codebase

• Algorithm 1 provides the full training algorithm, detailing encoding, polynomial decoding, and optimization steps used throughout the experiments.Source: A. PolySAE Algorithm and Codebase · p. 15
• The PolySAE codebase is publicly available on GitHub.The paper gives the repository at github.com/pakoromilas/PolySAE.Source: A. PolySAE Algorithm and Codebase · p. 15

B. Implementation Details

• SAEs and PolySAEs use 16,384 latent dimensions, K = 64 sparsity, and Top-K, BatchTopK, or Matryoshka encoders on residual-stream activations.Top-K and BatchTopK retain the K largest activations per token or batch, respectively.Source: B. Implementation Details · p. 15
• Evaluation covers GPT-2 Small, Pythia-410M and Pythia-1.4B (Biderman et al., 2023), and Gemma-2-2B (Gemma Team, 2024).For each model, activations come from a single layer near the network center.Source: B. Implementation Details · p. 15
• SAEBench evaluation reports held-out reconstruction error and sparse probing on six classification tasks.Gemma-2-2B and GPT-2 Small use OpenWebText, while Pythia models use an uncopyrighted deduplicated Pile variant; reconstruction is evaluated on held-out data from the training distribution. Tasks are Bias in Bios, AG News, EuroParl, GitHub programming languages, Amazon Sentiment, and Amazon-15.Source: B. Implementation Details · p. 16
• PolySAE changes only the decoder; the encoder, sparsification, optimizer, and evaluation pipeline remain shared through an extended SAELens interface.The training pipeline preserves the standard SAE training interface.Source: B. Implementation Details · p. 16

C. Extended Qualitative Analysis

• The extended qualitative analysis examines pairwise interactions first, then triplet compositions, comparing PolySAE with a vanilla Top-K SAE trained under identical conditions.Source: C. Extended Qualitative Analysis · p. 16

C.1. Second-Order Analysis

• PolySAE interaction strength correlates with co-occurrence at r = 0.06, versus r = 0.82 for vanilla SAE feature covariance.The analysis used 1M OpenWebText documents and the top 10,000 features, yielding approximately 5 × 10^7 candidate pairs.Source: §C.1 · p. 16
• High-strength, low-co-occurrence pairs reveal latent interactions that often encode meaningful compositional patterns not recoverable from frequency alone.Table 7 contrasts PolySAE compositions with vanilla SAE features, which fail to capture those compositions.Source: §C.1 · p. 16, Table 7 · p. 17

C.2. Third-Order Analysis

• Third features in selected latent pairs consistently modulate pair meaning, forming interpretable entity–attribute–domain or subject–object–context structures.For each high-strength, low-co-occurrence pair, the analysis scores co-active third features with the cubic decoder and retains the highest-scoring feature after filtering stopword-like features.Source: §C.2 · p. 17

C.3. Additional Second-Order Interaction Examples

• Across additional second-order examples, PolySAE features are typically more specific than vanilla SAE features, which often capture only a broad or morphological feature.The comparison covers Tables 7–11 and highlights tokens where two PolySAE features activate together.Source: §C.3 · p. 17

C.4. Additional Third-Order Interaction Examples

• In additional third-order examples, three PolySAE features vary with context and are more specific than vanilla SAE activations, which often capture one component or generic features.The examples are shown in Tables 12 and 13.Source: §C.4 · p. 17

D. LLM-as-Judge Prompt for Interaction Interpretability

• Of 70K evaluated interaction pairs, 8,550 (12%) scored above 0.9 for matching the composite concept implied by their constituent features.GPT-4o-mini independently scored contexts for pairs with non-negligible interaction strength, using feature tokens and representative co-activation contexts; 70K of 292,361 eligible pairs were evaluated.Source: D. LLM-as-Judge Prompt for Interaction… · p. 18

E. Activation Steering: Additional Examples

• +41.5 mean rank improvement is illustrated by GPT-2 Small cases where PolySAE steering changes output semantics while vanilla SAE directions do not shift the baseline.Table 14 covers selected examples across 27 compositional concepts; Table 15 isolates these failure cases for vanilla SAE.Source: E. Activation Steering: Additional… · p. 19
• Three PolySAE features can bind in context to represent third-order compositional concepts that vanilla SAE often misses.The table caption describes features Fi, Fj, and Fk as capturing compositional structure beyond individual components.Source: Table 12
• Figure 4 provides the full LLM-as-judge prompt used to score the interpretability of PolySAE feature interactions in Section 5.2.Source: Figure 4 · p. 20

F. Interaction Rank Ablation

• With R1 = 768 on GPT-2 Small, PolySAE achieves competitive reconstruction at R2 = R3 = 64, and increasing ranks beyond 64 does not improve reconstruction.The result suggests that additional interaction capacity is unnecessary for capturing interaction structure in this setting.Source: F. Interaction Rank Ablation · p. 20

G. Computational Overhead

• Under identical GPT-2 Small training conditions, the evaluation tracks reconstruction MSE over training and records GPU memory, system RAM, and CPU utilization.Figure 6 shows reconstruction MSE over 3 × 108 training tokens; Figure 7 reports system-resource usage.Source: G. Computational Overhead · p. 20, Figure 6 · p. 24, Figure 7 · p. 24
• +1.19 GB of GPU memory and +18.2 min of wall-clock time are PolySAE’s reported overheads, with comparable system RAM and CPU utilization.Table 16 summarizes median resource use and total wall-clock time; measurements used 3 × 108 training tokens, context length 128, batch size 4096 tokens per step, and a single GPU.Source: G. Computational Overhead · pp. 20–21, Table 16 · p. 24, Figure 7 · p. 24