Uncovering Neural Geometry in Vision Models With Block-Sparse Featurizers

Picture a tree.

Now make it small, a sapling; then tall, an oak; then enormous, a redwood; then tiny again, a bonsai. Turn its leaves summer green, then autumn red, then strip them bare. Change its species: birch, pine, willow. Each one is still a tree, but they differ along many axes of variation. The concept “tree” is not a one-dimensional idea, but a multidimensional space of variation.

1

Even before the emergence of LLMs, many have argued we should think of concepts as regions in structured geometric spaces, e.g. Peter Gärdenfors in Conceptual Spaces: The Geometry of Thought.

Trees are just one example here; our world is full of intricate structure. So are the internals of neural networks! But the most popular methods we use to understand what’s happening inside neural networks, like sparse autoencoders (SAEs) and transcoders, treat each internal concept as a single straight line – an assumption which is surprisingly powerful, but which nonetheless doesn’t effectively capture the naturally curved, multidimensional structures we know exist in model representations.

In this post, we explain our new paper which introduces a family of techniques – Block-Sparse Featurizers (BSF) – that we think provide a simple, viable alternative. BSFs decompose model activations into subspaces, which empirically contain multidimensional shapes.

We believe this is a significant step forward in understanding neural networks on their own terms. Applied to vision models, we find that BSFs find interpretable, multidimensional features which better explain model activations; that those features enable fine-grained steering; and that most concepts in the models are multidimensional.

Introduction: A toy model of manifolds in superposition

As part of interpretability’s broader goal of understanding what goes on inside neural networks, we want to know what a model is “thinking” at any given moment. The primary way we do this

2

But not the only way! E.g., see our work on decomposing weights. is by reading the activations flowing through the model and mapping them to meaningful concepts. In their raw form, though, those activations are just large vectors with no convenient labels. So how do we disentangle them into meaningful patterns?

That question can be reframed as a different one: how are a model’s representations structured? If we know that, then we can choose the best method to disentangle them.

If you’ve been following this series of posts, you can probably guess where this is going. We’ve discussed the growing body of research, from Goodfire as well as many others, showing that concepts are often represented inside neural networks as manifolds – multidimensional, curved shapes. We’ve also discussed the mismatch between these complex representations and the tools we use to find them: the most popular methods, like SAEs and transcoders, look only for straight lines.

3

This straight-line assumption has gotten SAEs quite far! We want the simplest possible building blocks for what a network computes, and a straight line is about as simple as a shape can be. Straight lines are also easy to work with computationally. But we think it’s a fundamentally limiting assumption, if we want to fully understand model activations.

To get an intuition for the consequences of this mismatch, consider the toy model above. We constructed clean artificial activations by summing samples from predefined manifolds – implicitly stacking the manifolds in superposition. This forms a simplified model of how we think real neural network activations are structured. Then we use two different featurizer methods to try to recover the known underlying structure from the activations.

On the left are the ground-truth manifolds, which are what we’re aiming to recover. In the middle column, we see that a standard SAE does a poor job of finding the underlying manifolds!

4

In our previous work, we showed SAEs don’t lose manifold geometry so much as fragment it by tiling manifolds across many redundant features that only post-hoc analysis (with an Ising-model and extra compute) can regroup. In contrast, BSFs recover these subspaces coherently during training, even when the SAE is given the advantage of selecting its best top-g directions against a single BSF block of size g.

But if we instead begin with the assumption that model activations are made of manifolds, that implies a different strategy, whose results are on the right. A Block-Sparse Featurizer (BSF) – the method this post describes – finds the ground-truth structures with much higher fidelity!

5

In the full paper, we test not only an SAE and vanilla BSF, but also the two other variants of BSF we propose, as well as two concurrent featurizer methods (MFA and SMixAE). Applying the block-sparsity principle allows us to greatly improve the ability of MFA and SMixAE to faithfully recover our toy manifolds!

So: how do BSFs work, and why are they effective at finding manifolds?

Block-Sparse Featurizers (BSFs)

We want a method that respects neural geometry. So instead of looking for one-dimensional lines as the fundamental “atoms” of neural computations, we look for multidimensional subspaces. Since each subspace is multidimensional, it can far better capture the shape taken by a concept’s representations.

How do we disentangle these subspaces? The key insight is that a block – the set of directions that span a subspace – should tend to “fire” together on the same input. Unlike SAEs, which assume that a small (sparse) number of 1D features are firing at any given time, BSFs assume that a small (sparse) number of blocks are firing at any given time. That is, we enforce sparsity at the block level.

This principle is what we call block-sparsity, and it is not new: it belongs to the older family of structured sparsity, a long-studied topic in dictionary learning

6

E.g., Yuan & Lin, 2006; Eldar & Bolksei, 2009; Jenatton et al., 2009..

In the course of this research, we found that the principle of block-sparsity matters more than the particular implementation. We built several methods that follow the general principle, yet differ in how they impose that structure,

7

Though removing the ReLU turns out to matter across all of them! See the paper for more details. and every one of them recovers concept manifolds. The same principle even lets us improve several existing methods, as we show in the full paper.

The diagram below shows the architecture of a generic BSF:

(Excerpt: introduction, toy model of manifolds in superposition, and the block-sparse method; full text at source URL; paper arXiv:2606.25234.)