🤖 machine learning

Kronecker Factorization Improves Efficiency and Interpretability of Sparse Autoencoders

The paper introduces KronSAE, a sparse autoencoder architecture that factorizes the latent space into heads and employs a differentiable AND-like interaction to enforce compositional co-activation, thereby improving interpretability, capturing feature correlations, and reducing computational costs without sacrificing performance.

Original authors: Vadim Kurochkin, Yaroslav Aksenov, Daniil Laptev, Daniil Gavrilov, Nikita Balagansky

Published 2026-09-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Vadim Kurochkin, Yaroslav Aksenov, Daniil Laptev, Daniil Gavrilov, Nikita Balagansky

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are vast, complex systems that process human language by transforming text into streams of numbers. These numbers, known as activations, flow through the model's internal layers, carrying the meaning of words and the logic of sentences. For researchers trying to understand how these machines think, these hidden streams are a black box. To open it, scientists use a tool called a sparse autoencoder. Think of this tool as a translator that takes the dense, messy flow of numbers and breaks it down into a long list of simple, distinct concepts. Ideally, each item on this list represents a single, clear idea—like "mathematics," "legal contracts," or "the color blue"—that the model activates when it encounters that specific concept. This process, called decomposition, allows researchers to see the individual gears turning inside the machine.

However, standard translators have a limitation. They treat every concept as an independent, flat item on a list, ignoring the fact that human language is deeply structured. Words and ideas often appear together in predictable patterns; the concept of "geography" frequently co-occurs with "maps" or "directions." When a model learns these patterns, the standard tool struggles to capture the relationship, often forcing the system to learn the connection implicitly or, worse, to merge distinct ideas into a single, confusing feature. This makes the resulting list of concepts harder to interpret and less efficient to compute.

A team of researchers at T-Tech has proposed a new way to build these translators, called KronSAE, which respects the natural structure of language from the start. Instead of treating every concept as a separate, flat coordinate, their design organizes the internal dictionary into groups. Within each group, the system builds complex features by combining two simpler, pre-existing components. It is a bit like how a chef might create a specific dish by combining a base ingredient with a specific spice; the dish only exists if both the base and the spice are present. In this new architecture, a feature activates only when two underlying "parent" signals are active at the same time. This creates a built-in rule that encourages the system to learn features that are combinations of simpler parts, mirroring how language actually works.

The researchers tested this approach on two popular language models, Qwen2.5 and Gemma-2, using a massive dataset of educational web pages. They found that this structured approach did more than just mimic human language; it actually made the model's internal representation clearer. When they compared the new system to the standard method, they discovered that the new design produced features that were more specific and less likely to be confused with one another. In the standard method, a single feature often tries to do too much, absorbing the meaning of several related concepts and becoming vague. The new system, by forcing features to be built from specific combinations, reduced this "absorption." The resulting features were sharper, describing narrower, more precise ideas, which made them easier for researchers to understand and label.

Beyond clarity, the new design offered a significant boost in efficiency. Because the system builds complex features by combining simpler ones, it does not need to calculate every single possibility from scratch. The researchers found that they could reduce the amount of computing power required to run the encoder by more than forty percent while maintaining the same level of accuracy in reconstructing the original data. This means the system can learn the same amount of information using far fewer resources. In their experiments, even with this massive reduction in computational cost, the drop in performance was less than one percent, a negligible trade-off for such a large gain in speed and efficiency.

The study also explored how well this system could learn the hidden relationships between concepts. In a controlled experiment using synthetic data where the relationships between features were known in advance, the new system successfully recovered these patterns much better than the standard method. It learned to group related concepts together within its internal structure, whereas the standard method scattered them. When looking at real-world data, the researchers observed that features which naturally appeared together in the text were indeed mapped to the same internal groups in the new system. This suggests that the architecture successfully captures the statistical regularities of language, organizing the model's knowledge in a way that reflects how ideas are actually connected.

The researchers also examined the nature of the features themselves. They found that the simpler "parent" components within the system were often broad and abstract, capable of representing multiple different ideas. However, when these parents were combined, the resulting feature became highly specific. For example, a broad component related to "directions" might combine with another related to "medical terms" to create a feature specifically about "medical directions." This hierarchical structure allowed the system to reuse its basic components to build a vast array of specific concepts without needing a unique, separate neuron for every single idea. This compositional approach not only made the features more interpretable but also provided a more efficient way to store and retrieve knowledge.

Ultimately, this work suggests that the way we design these interpretability tools matters as much as the data they process. By introducing a simple structural bias that mimics the way concepts combine in language, the researchers created a system that is both more efficient and more transparent. The findings indicate that we do not need to rely on the model to accidentally discover these relationships during training; we can build them into the architecture itself. This approach offers a new path forward for understanding artificial intelligence, providing a clearer window into the machine's mind while simultaneously making the system faster and cheaper to run. The code for this new method is now available for others to use and build upon, inviting further exploration into how structured learning can improve our understanding of complex systems.

Technical Summary: KronSAE

Problem Statement

Sparse Autoencoders (SAEs) are a standard tool for decomposing dense language model activations into sparse, interpretable features. However, standard SAE architectures typically treat the latent dictionary as a flat set of independent coordinates. This "flat-latent" default fails to explicitly encode higher-level structural relationships found in natural language, such as hierarchy, compositionality, and systematic co-activation. Consequently, these relationships must emerge implicitly during training or be recovered through post-hoc analysis. Furthermore, the standard encoder requires a dense projection to $F$ independent latents, resulting in a computational cost of $O(Fd)$ per token, which becomes expensive for large dictionary sizes.

Methodology: KronSAE

The authors propose KronSAE, an encoder-side modification that introduces a hierarchical, compositional prior to the latent space. Instead of projecting directly to $F$ independent features, KronSAE factorizes the latent space into $h$ independent "heads."

Architecture

Factorization: For each head $k$ , the input activation $x \in \mathbb{R}^d$ is projected into two lower-dimensional "pre-latents":
- A base vector $p^k \in \mathbb{R}^m$ via $u^k = P^k x$ .
- An extension vector $q^k \in \mathbb{R}^n$ via $v^k = Q^k x$ .
- Here, $m, n \ll d$ , and the total dictionary size is $F = h \times m \times n$ .
Compositional Interaction (mAND): The post-latent features for a head are formed by the pairwise composition of the base and extension pre-latents using a differentiable interaction kernel called mAND:
$z^k_{i,j} = \text{mAND}(u^k_i, v^k_j) = \begin{cases} \sqrt{u^k_i v^k_j} & \text{if } u^k_i > 0 \text{ and } v^k_j > 0 \\ 0 & \text{otherwise} \end{cases}$
This operation approximates a logical AND, ensuring a post-latent is active only when both parent pre-latents are active. The resulting matrix $z^k \in \mathbb{R}^{m \times n}$ is flattened to form the post-latent vector.
Sparsity: The flattened vectors from all heads are concatenated, and a standard TopK selection is applied to produce the final sparse latent vector $f \in \mathbb{R}^F$ .

Compatibility and Efficiency

Compatibility: KronSAE is compatible with standard SAE objectives and variants, including TopK, Matryoshka (nested dictionaries), Switch SAEs (Mixture-of-Experts), and different activation functions like JumpReLU.
Computational Gain: By replacing a dense projection to $F$ latents with projections to $h(m+n)$ pre-latents, KronSAE reduces the encoder's parameter count and floating-point operations (FLOPs). The encoder cost scales as $O(d \cdot h(m+n))$ rather than $O(d \cdot F)$ .

Key Contributions

Identification of Encoder Hierarchy: The paper identifies the lack of explicit encoder-side hierarchy as a missing design axis in SAEs, motivated by the structural inductive biases inherent in feature organization.
KronSAE Design: Introduction of a Kronecker-factorized encoder equipped with the mAND interaction, which constructs post-latent features from reusable pre-latent factors.
Improved Interpretability: Empirical evidence showing that KronSAE reduces "feature absorption" (where one feature becomes a strict subset of another) and yields more specific, interpretable latents compared to flat TopK SAEs.
Reconstruction-Compute Trade-off: Demonstration that KronSAE can reduce encoder FLOPs by more than 40% while incurring less than a 1% drop in Explained Variance (EV).

Experimental Results

The authors trained SAEs on the residual streams of Qwen2.5-1.5B and Gemma-2-2B models using the FineWeb-Edu dataset.

Reconstruction Fidelity: Under an equal token budget (250M tokens), KronSAE achieves Explained Variance (EV) comparable to standard TopK SAEs across various dictionary sizes ( $2^{14}$ to $2^{16}$ ) and layers.
Interpretability: Using an automated pipeline (detection and fuzzing scores), KronSAE variants (specifically with $m \ge 2$ ) consistently outperform TopK baselines. The features exhibit lower token entropy and lower multitoken ratios, indicating higher specificity.
Correlation Learning:
- Real Data: Correlated features in standard TopK SAEs tend to map to features within the same KronSAE head, suggesting the architecture successfully groups related concepts.
- Synthetic Data: In a controlled toy experiment with block-structured covariance, KronSAE recovered the ground-truth correlation structure significantly better (higher RV coefficient) than TopK SAE.
Feature Absorption: KronSAE configurations showed reduced mean absorption fractions and full-absorption scores compared to baselines, indicating a reduction in the "subsumption" of features.
Efficiency: The method reduces encoder FLOPs by over 40% while maintaining reconstruction quality, as shown in the Pareto frontier analysis.

Significance and Claims

The paper claims that KronSAE offers a new perspective on how sparsity and factorization can work together in representation learning. By imposing an architectural co-activation prior (via the mAND interaction), the model makes structured feature interactions easier to learn without compromising reconstruction fidelity.

The authors position KronSAE as a method that:

Enhances Interpretability: By reducing feature absorption and creating more specific, monosemantic features.
Improves Efficiency: By drastically reducing the computational cost of the encoder, making large-dictionary SAEs more feasible to train and run.
Aligns with Linguistic Structure: By explicitly modeling the compositional nature of language (e.g., a concept being the intersection of a base category and a specific modifier) rather than treating features as independent flat coordinates.

The authors note limitations, including that the evaluation is currently limited to mid-sized transformer models and that the engineering overhead of implementing the factorized encoder is moderate but non-zero. They also acknowledge that while the toy model validates the inductive bias, real-world correlation structures may differ.