Finding Manifolds With Bilinear Autoencoders

TL;DR

The paper introduces bilinear autoencoders that decompose neural representations into quadratic, polynomial latents by linearly reconstructing the quadratic space $X = \mathbf{x} \otimes \mathbf{x}$. It combines this with a kernel-trick training approach for efficiency and proposes extensions for scale-invariant sparsity, latent importance ordering via cumulative reconstruction, and latent mixing through a linear bottleneck to reveal interacting latent structures. Empirically, it demonstrates that these latents uncover interpretable manifolds, quantify latent consistency with exact Frobenius-based similarity, and show stable reconstruction across hyperparameters. The work argues for polynomials as algebraic primitives in interpretability, enabling principled, compositional analysis of nonlinear latents and offering a path toward nonlinear yet analysable latent representations with potential applications in model debugging and verification.

Abstract

Sparse autoencoders are a standard tool for uncovering interpretable latent representations in neural networks. Yet, their interpretation depends on the inputs, making their isolated study incomplete. Polynomials offer a solution; they serve as algebraic primitives that can be analysed without reference to input and can describe structures ranging from linear concepts to complicated manifolds. This work uses bilinear autoencoders to efficiently decompose representations into quadratic polynomials. We discuss improvements that induce importance ordering, clustering, and activation sparsity. This is an initial step toward nonlinear yet analysable latents through their algebraic properties.

Figures (15)

• Figure 1: Conceptual subspaces extracted directly and solely from the weights of a bilinear autoencoder. To visualise the structure of these spaces, inputs are projected linearly onto it, with colour indicating activation strength. This reveals learned manifolds: the top-left organises nuances of 'new' into an off-centre triangle, while the top-right forms an off-centre circle for the first digits in a year. The bottom-left linearly clusters subordinating conjunctions, and the bottom-right captures verb negation. The top two are nonlinear manifolds, whereas the bottom two suggest linear concepts, possibly in superposition elhage2022toymodelssuperposition. An interactive version is available at https://mi-manifold.netlify.app/.
• Figure 2: Most autoencoders reconstruct their inputs nonlinearly. Instead, bilinear autoencoders linearly reconstruct the product space $X$, which is quadratic in the input $\mathbf{x}$.
• Figure 3: Coloured areas illustrate high activations $\mathbf{f}_i(\mathbf{x}) > 1$ for two bilinear latents $B_i$. Isolated latents can selectively activate on superposed directions (red). Composite latents can activate on more complicated manifolds, such as oriented cones and paraboloids (blue).
• Figure 4: Diagrammatic formulation of \ref{['eq:vanilla-loss']}. Lines indicate tensor contractions over that index, while unconnected lines represent unmodified open indices. Thick lines represent a quadratic space.
• Figure 5: Diagrammatic formulation of the ordered loss. Here, $M$ is an upper triangular mask matrix with columns $\mathbf{m}_k$ and $\mathbf{w}$ are the weights of each cumulative reconstruction. We use a uniform distribution for $\mathbf{w}$ to make all prefix reconstructions ($\text{Pre}$) equally important. This encourages the latents to contribute to the reconstruction in an approximately arithmetic progression of importance.