PaperlayerView exact revision on arXiv

Source-linked AI summary

PolySAE: Modeling Feature Interactions in Sparse Autoencoders via Polynomial Decoding

Panagiotis Koromilas, Andreas D. Demou, James Oldfield, Yannis Panagakis, Mihalis Nicolaou

arXiv:2602.01322v2cs.LGcs.CL

TL;DR

Standard SAEs rely on additive reconstruction, which limits their ability to represent non-linear compositional structure. PolySAE keeps the encoder linear but adds low-rank polynomial decoder interactions, improving probing and semantic separation while preserving comparable reconstruction. The method is intended as an interpretability analysis tool rather than a deployment component, and its studied model scale is limited.

  • Problem

    Standard SAEs assume features combine additively through linear reconstruction, limiting their ability to distinguish compositional structure from co-occurrence.

  • Method

    PolySAE preserves a linear SAE encoder and uses low-rank factorized quadratic and cubic decoder terms to model pairwise and triple feature interactions.

  • Results

    Across four language models and three SAE variants, PolySAE improves probing F1 by 8% on average, maintains comparable reconstruction, and yields 2–10× larger Wasserstein distances.

  • Takeaways & Limitations

    Interaction weights show negligible correlation with co-occurrence frequency (r = 0.06), while learned interaction directions steer outputs toward corresponding compositional semantics.

  • Takeaways & Limitations

    PolySAE is studied in models up to 2B parameters and is intended as an analysis tool rather than a deployment-facing component.

Abstract

from arXiv · show

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.

1. Introduction

SAEs support mechanistic interpretability but their additive reconstruction cannot represent non-linear composition or distinguish compositional phrases from accidental co-occurrence. PolySAE addresses this gap with a polynomial decoder, retaining a linear encoder while improving semantic probing and feature separation.

  • SAEs decompose activations into sparse, interpretable features that support monitoring, debugging, and steering model behavior.
  • Standard SAE reconstruction assumes features combine additively, limiting representation of non-linear morphology, phrasal meaning, and emergent concepts.
  • PolySAE preserves a linear encoder while adding quadratic and cubic decoder terms for pairwise and triple feature interactions.
  • 3% parameter overhead for GPT2 small makes PolySAE a tractable extension applicable to TopK, BatchTopK, and Matryoshka SAEs.
  • 8% average probing F1 improvement across four language models and three sparsification strategies accompanies comparable reconstruction error.
  • 2–10× larger Wasserstein distances and r = 0.06 interaction-weight correlation indicate more separated semantic structure with weaker dependence on co-occurrence frequency.

2. Related Work

Related work establishes additive sparse dictionary learning and multiplicative interaction modeling as relevant foundations. PolySAE differs from prior bilinear autoencoders by modeling interactions between learned sparse features, including higher-order terms.

  • SAEs and their variants use sparse dictionary learning to discover features in large language models, but assume additive linear reconstruction.
  • Multiplicative interactions have been used in bilinear models, gating mechanisms, Hadamard products, and conditioning systems.
  • Polynomial models provide a natural framework for representing higher-order interactions across applications including classification and interpretability.
  • PolySAE models pairwise and higher-order interactions between learned sparse features, unlike BAE, which models pairwise interactions between input neurons.

3. Sparse Polynomial Decoding

PolySAE preserves the linear SAE encoder while extending decoding with factorized quadratic and cubic feature interactions. Shared low-rank projections make these interactions expressive, context-dependent, and parameter-efficient.

  • Linear Encoding: PolySAE keeps sparse codes interpretable by deriving them through a linear encoder and sparsification.This follows the design principle that features should remain directions in activation space.
  • Polynomial Reconstruction: PolySAE reconstructs activations with linear, quadratic, and cubic terms that model pairwise and triple feature interactions.Setting λ2 = λ3 = 0 recovers a standard linear sparse autoencoder.
  • Factorized Interaction Structure: Low-rank interaction subspaces restrict pairwise and higher-order dictionaries, imposing a reusable inductive bias and reducing overfitting risk.Explicitly modeling all pairwise combinations would require O(d_sae^2) parameters; the factorized form limits interaction dictionaries to rank at most R_k.
  • Factorized Interaction Structure: Shared projection U forms interactions from zU, (zU) ∗(zU), and (zU) ∗(zU) ∗(zU), aligning them with the linear feature basis.This avoids explicitly materializing all feature combinations and keeps interaction effects interpretable as compositions of the same features.
  • Context-Dependent Dictionary Structure: Higher-order terms make each feature’s effective dictionary contribution context-dependent, varying with which other features are active.The pairwise dictionary modifies reconstruction through z_i z_j, while the triple dictionary contributes through z_i z_j z_k.
  • Compositional Capacity: PolySAE supports many semantic compositions through shared interaction directions without increasing the number of base features.Its shared low-rank modes provide reusable directions for feature combinations rather than independent parameters for every combination.
  • Parameter Efficiency: Approximately 2.5–3% additional parameters suffice for GPT-2 small under the empirically optimal low-rank configuration.The reported setting uses d = 768, d_sae = 16,384, with R2 = R3 and R2 ∈ [0.06R1, 0.11R1].

4. Empirical Evaluation

PolySAE is evaluated for reconstruction fidelity and semantic modeling across language models, sparsification strategies, and sparsity levels. It preserves comparable reconstruction while improving probing and distributional semantic separation, with architectural ablations supporting tractable design choices.

  • Evaluation setup: The evaluation uses MSE and cross-entropy recovery for reconstruction, plus probing and 1-Wasserstein distance for semantic structure.Probing uses logistic regression on SAE activations, while Wasserstein distance measures global class-conditional distributional separation.
  • Reconstruction and semantic modeling: 8% average F1 improvement is achieved across four language models and three sparsification strategies while maintaining comparable reconstruction error.PolySAE also gains more than 10% on GPT-2, while cross-entropy recovery remains within 0.003 of vanilla SAE across all 12 configurations.
  • Reconstruction and semantic modeling: 2–10× larger Wasserstein distances are observed across the other models, indicating more separated class-conditional representations.The result supports probing gains that reflect broader distributional separation rather than only favorable decision thresholds.
  • Ablations: 65% fewer parameters incur only a 1pp F1 cost under low-rank factorization, while orthogonality recovers the gap and adds +2.9pp at zero parameter cost.The ablation evaluates parameter count, MSE, and F1 on GPT-2 Small.
  • Sparser codes: PolySAE consistently outperforms standard SAEs across sparsity levels, with a widening gap and lower width variance at higher active-feature counts.Figure 3 varies active features k and shows ranges across dictionary widths from 2k to 16k.
  • Semantic concentration: PolySAE shows smaller F1 gains when expanding from K=1 to K=5 in 9 of 12 configurations, consistent with stronger semantic concentration into fewer features.All three GPT-2 Small sparsifiers show gaps of −7 to −8; the paper attributes this behavior probably to higher-order interactions absorbing contextual variability.

5. Understanding and Utilizing Interactions

The paper tests whether PolySAE interactions reflect compositional structure rather than co-occurrence, whether they are interpretable, and whether they can steer generation. The analyses find frequency-independent interactions, interpretable pairwise and triple compositions, and reliable causal steering toward target concepts.

  • Surface statistics versus composition: PolySAE interaction capacity has negligible co-occurrence correlation (r = 0.06), unlike SAE activation covariance (r = 0.82).The comparison tests whether learned quadratic interactions merely replicate feature co-occurrence statistics.
  • Interpretable interactions: Quadratic interactions capture context-dependent phrase compositions such as coffee × star yielding Starbucks, which isolated SAE features do not recover.Third-order terms add contextual conditioning that refines pairwise compositions and disambiguates meanings.
  • Interpretable interactions: Third-order interactions integrate morphological and semantic cues to distinguish investing from unrelated -ing usages and disambiguate generic entities.These higher-order terms absorb contextual variation that would otherwise fragment linear features.
  • Interpretability at scale: 12% of evaluated feature pairs score above 0.9 for highly interpretable compositional interactions.Among 70,000 evaluated pairs, 8,550 exceed the 0.9 interpretability score.
  • Causal steering: PolySAE improves target rank over no steering in 230 of 324 prompt-concept pairs (71.0%), with only 3 degradations.At the concept level, it improves over no steering in 27/27 concepts and over vanilla SAE in 21/27, with a mean rank improvement of +41.5 against vanilla.
  • Direction alignment: Steering vectors are evaluated against difference-in-means ground-truth compositional directions using cosine similarity.The directions combine decoder directions for the constituent features and are injected into GPT-2’s layer-8 residual stream.

6. Conclusion

PolySAE extends SAE decoding with higher-order feature interactions while retaining a linear encoder, improving semantic separation and enabling compositional steering without sacrificing comparable reconstruction. Its interaction capacity is largely independent of co-occurrence statistics, though the study is limited to models up to 2B parameters and forced-sparsity SAE variants.

  • Conclusion: PolySAE models pairwise and triple feature interactions through higher-order decoding and low-rank tensor factorization while preserving a linear encoder.The method adds small parameter overhead and targets compositional structure in learned representations.
  • Conclusion: 8% average probing F1 improvement and 2–10× larger Wasserstein distances are achieved across four LLMs and three SAE variants with comparable reconstruction.The conclusion summarizes improvements in both probing performance and class-conditional distributional separation.
  • Conclusion: Interaction capacity shows compositional structure largely independent of co-occurrence, with r = 0.06.The learned interaction directions also causally steer outputs toward corresponding compositional semantics.
  • Limitations: The study covers models up to 2B parameters and restricts experiments to forced-sparsity SAE variants.These are the stated scope limitations of the evaluation.

Impact Statement

PolySAE is presented as a methodological and scientific analysis tool for representing compositional structure in interpretable language-model features. It is intended to support analysis, auditing, debugging, and safety research rather than deployment or capability gains.

  • PolySAE models non-additive feature interactions while preserving linear, human-interpretable features.The approach is intended to advance mechanistic analysis of large language models.
  • Improved interpretability may make meaningful internal representations easier to identify, analyze, and intervene on for auditing, debugging, and safety research.
  • PolySAE is intended as an analysis tool rather than a deployment-facing component and does not directly increase language-model capabilities.
  • The authors identify possible misuse of interpretability insights but no novel or unique risks beyond those already present in the interpretability literature.

B. Implementation Details

The experiments train PolySAEs and standard SAEs with shared sparse-encoding infrastructure across four language models, three sparsification strategies, and matched evaluation procedures. PolySAE changes only the decoder and is analyzed through pairwise and triplet interactions.

  • Architecture and Sparsification: Models use latent width 16,384 and sparsity K = 64 with Top-K, BatchTopK, or Matryoshka encoders.
  • Implementation Details: Experiments use GPT-2 Small, Pythia-410M, Pythia-1.4B, and Gemma-2-2B, extracting residual-stream activations from central transformer layers.
  • Evaluation: Evaluation measures held-out reconstruction error and sparse probing across six classification tasks using SAEBench.
  • Implementation: PolySAE differs from standard SAEs only in the decoder; encoders, sparsification, optimizers, and evaluation pipelines are shared.
  • Interaction Analysis: Qualitative analysis examines second-order pairwise interactions and then third-order triplet compositions against a vanilla Top-K SAE trained under identical conditions.

C.1. Second-Order Analysis

The second-order analysis tests whether PolySAE interactions capture compositional structure beyond surface co-occurrence, then examines whether third-order terms refine those pairwise compositions. Across qualitative examples, PolySAE features are more specific and its triplets contextualize pairwise meanings.

  • Setup: The analysis ranks features by activation mass, retains the top 10,000, and evaluates approximately 5 × 10^7 candidate pairs from 1M OpenWebText documents.
  • Relation to Co-occurrence: PolySAE interaction strengths correlate negligibly with co-occurrence (r = 0.06), whereas vanilla SAE covariance correlates strongly (r = 0.82).The contrast suggests PolySAE interactions reflect structure beyond surface statistics.
  • Qualitative Regimes: Strong interactions with low co-occurrence identify latent pairs that often correspond to compositional patterns unrecoverable from frequency alone.
  • Third-Order Analysis: Third-order triplets are consistently interpretable, with the third feature modulating pair meaning in entity–attribute–domain and subject–object–context structures.
  • Qualitative Examples: Across second-order examples, PolySAE features are typically more specific than vanilla SAE features, which often capture broad or high-level patterns.
  • Third-Order Analysis: Third-order activations vary with context and are more specific than vanilla SAE activations, which often capture one component or generic features.

E. Activation Steering: Additional Examples

Additional GPT-2 Small examples examine activation steering for compositional concepts and compare PolySAE directions with vanilla SAE directions. The examples include cases where PolySAE produces semantic changes that vanilla directions do not.

  • Activation Steering: Additional GPT-2 Small examples cover 27 compositional concepts and extend the activation-steering experiment from the main paper.
  • Activation Steering: Some vanilla SAE directions fail to shift output from the unsteered baseline, while PolySAE directions produce a meaningful semantic change.
  • Activation Steering: The examples illustrate the gap represented by the aggregate +41.5 mean rank improvement reported in Section 5.2.
  • Interaction Scoring: The interpretability scoring prompt asks for 0–1 judgments of how strongly examples match the composite concept implied by two constituent features.

F. Interaction Rank Ablation

PolySAE reaches competitive reconstruction with modest interaction ranks, while larger ranks provide no further benefit in this setting.

  • F. Interaction Rank Ablation: R2 = R3 = 64 achieves competitive reconstruction with R1 = 768 on GPT-2 Small.Increasing interaction ranks beyond 64 does not improve reconstruction.

G. Computational Overhead

PolySAE adds modest training-resource overhead under identical GPT-2 Small conditions, with comparable system RAM and CPU utilization. The supplied passages also describe reconstruction tracking, steering evaluations, and additional interaction examples, but do not report their outcomes here.

  • G. Computational Overhead: +1.19 GB of GPU memory and +18.2 min of wall-clock time are the reported PolySAE overheads.These measurements use 3 × 10^8 training tokens, context length 128, batch size 4096 tokens per step, and a single GPU.
  • G. Computational Overhead: Figure 6 shows reconstruction MSE over training, while Figure 7 reports per-step system resource usage.Table 16 summarizes median resource use and total wall-clock time.
  • G. Computational Overhead: PolySAE has comparable system RAM and CPU utilization to a vanilla SAE under identical training conditions.The passage characterizes these overheads as modest relative to the reported gains.
  • G. Computational Overhead: Table 14 evaluates activation steering across 27 concepts and 12 prompts under No Steering, SAE, and PolySAE conditions.Bold marks cases where PolySAE steers generation toward the compositional target.
  • G. Computational Overhead: Additional steering examples show PolySAE producing semantic shifts where vanilla SAE directions leave outputs unchanged from the unsteered baseline.The passage describes these as cases where vanilla SAE directions fail to shift the output.
Loading 2602.01322v2…