Finding Manifolds With Bilinear Autoencoders

1 Introduction

• Latents lack compositional interpretation, so their isolated behaviour may not explain the whole representation and can be skewed by composition.Their domains of applicability can shift through splitting and absorption [7].Source: §1 · p. 2
• Polynomials are proposed as nonlinear primitives that can be analysed through their coefficients and geometrically describe manifolds.They are presented as basis-independent and suitable for studying nonlinear computation and representations [27].Source: §1 · p. 2
• Bilinear autoencoders decompose neural representations into polynomial latents, supporting automated manifold discovery and exact similarity and distance calculations without input-dependent correlation analysis.Source: §1 · p. 2
• The paper proposes extensions that induce clustering, importance ordering, and activation sparsity without post-hoc regularisation penalties.Source: §1 · p. 2
• The analysis uses bilinear autoencoders to uncover nonlinear manifolds and quantify latent consistency through compositionality.Source: §1 · p. 2

2.1 Motivation and architecture

• Low-rank constraints can be too restrictive because neural representations typically span all dimensions; sparse autoencoders add nonlinearity to a linear decoder, recovering linearly separable latents traceable to original representations.Both methods aim to reveal easily extractable structures but may produce latents reflecting imposed priors rather than underlying structure [7].Source: §2.1 · p. 2
• Each bilinear latent is a rank-1 bilinear form over X = x ⊗x, where the encoder stacks these forms across latents.Removing σ yields second-order polynomials that remain linearly analysable in the lifted interaction space [27] [30].Source: §2.1 · p. 3
• Bilinear autoencoders reconstruct the quadratic product representation X rather than the input x, using a bottleneck to capture linearly extractable structure in X.Their decoder is the encoder transpose, and the product-space representation is quadratic in x.Source: §2.1 · p. 3, Figure 2 · p. 3

2.2 Bilinear latents and their geometry

• Bilinear latents capture linear directions in product space, and their closed-form compositions yield quadratic manifolds in input space.Linear combinations remain bilinear forms whose activation landscapes define quadratic surfaces and intricate activation regions.Source: §2.2 · pp. 3–4
• Closure under linear combinations makes single latents and composites analysable within the same vector space, including after rotations or mixtures.Figure 3 illustrates selective activation on superposed directions and more complex manifolds such as oriented cones and paraboloids.Source: §2.2 · p. 4, Figure 3 · p. 4

2.3 Efficient training with the kernel trick

• Rewriting reconstruction error in terms of latent activations and BB⊤ avoids materializing the quadratic reconstruction.The kernel-style computation uses inner products; the method scales quadratically in time and linearly in memory with latent dimension, without materializing the kernel matrix.Source: §2.3 · pp. 4–5

3.1 Avoiding dead latents with scale-invariant sparsity

• Hoyer density penalizes relative activation density without pushing absolute activation magnitudes toward zero, helping avoid dead latents while promoting selective representations.It is computed per latent across batch and sequence, and added to the SSE loss with coefficient α.Source: §3.1 · p. 5

3.2 Ordering latents by importance through cumulative reconstruction

• Cumulative reconstruction loss trains every latent prefix to reconstruct well, providing an ordering intended to support truncation while retaining reconstruction quality.An upper-triangular mask selects the prefixes, and uniform weights make them equally important; the ordered objective is recast as an efficient kernel-matrix computation.Source: §3.2 · pp. 5–6, Figure 5 · p. 6

3.3 Mixing latents through a linear bottleneck

• A lower-dimensional latent bottleneck produces mixture latents whose bilinear forms can attain higher rank, highlighting interacting spaces that may contain complex manifolds.The bottleneck encourages identification of latent anti-correlations as well as correlations to improve reconstruction [13].Source: §3.3 · p. 6

3.4 Combining extensions

• The extensions are modular, but cumulative ordering boosts early latents’ reconstruction contribution, so activation-sparsity penalties should be scaled across the ordered basis to preserve balance.The paper applies ordering and sparsity to the latent basis; sparsity could instead be applied within the bottleneck to promote sparse latent mixtures.Source: §3.4 · p. 6

4.1 Measuring sparsity and reconstruction

• Activation density falls from 0.6 to 0.2 as α rises from 0 to 0.1 with no reconstruction cost; further reductions increase error.Latent mixtures improve this trade-off, while ordering degrades it; density penalties sometimes improve reconstruction, suggesting naturally sparse features.Source: §4.1 · p. 7, Figure 6 · p. 7
• About 2% of latents remain dense despite sparsification, and ordering moves these dense latents to the front independently of sparsity regularization [31].The vanilla and ordered variants show one prominent peak of highly sparse latents and a smaller peak of dense latents.Source: §4.1 · p. 7, Figure 7 · p. 8

4.2 Discovering manifolds through latent interactions

• The mixed autoencoder uses latent interactions to identify candidate clusters, combines top candidates into composite latents, and extracts their salient input-space subspaces from weights.It projects the 64k highest-activating inputs out of 256k onto each subspace; unlike PCA, this reveals manifolds without input-dependent statistics.Source: §4.2 · p. 8
• Among about 50 inspected candidates, roughly half show nonlinear structure, with visualizations spanning quadratic surfaces, separated clusters, and linear directions.Most concepts are linearly represented, often as a cloud of near-zero activations with one or a few high-activation directions.Source: §4.2 · p. 8

4.3 Computing exact reconstruction similarity between autoencoders

• The Frobenius comparison measures exact reconstruction globally and independent of basis, while a separate latent-mapping approach investigates how similarly individual latents are represented.Empirical results for latent consistency are provided in Appendix I.Source: §4.3 · pp. 8–9
• Comparable-sparsity bilinear autoencoders are near-identical at 98% similarity for vanilla and ordered variants, while varying seeds yields over 99% similarity for vanilla models.Similarity drops significantly between the densest and sparsest instances, indicating that sparsity penalties affect whole-autoencoder behavior as well as latents.Source: §4.3 · pp. 8–9, Figure 8 · p. 9

• The paper uses bilinear autoencoders to decompose representations into polynomial latents, induce sparsity, importance ordering, and clustering, and automate manifold discovery.Source: §5 · p. 9
• Future work includes using discovered manifolds for causal interventions and steering, and scaling to larger models and latent spaces to investigate circuits, failures, and safety properties.Source: §5 · p. 9
• Other proposed directions are multi-layer interpretability, composite-manifold representations, and architectural biases for known geometries.Source: §5 · p. 9
• Polynomial computation and complexity can be analysed in closed form, potentially facilitating description-length quantification and study of composition and interactions.Source: §5 · p. 9
• An ordering-aware similarity can be formed by including M ⊤M, but the paper does not explore this cumulative-similarity variant.Source: §5 · p. 9

• TD and WG jointly developed the theory on sparse bilinear autoencoders.Source: §6 · p. 10
• TD conducted experiments, developed extensions, wrote the code, and composed the initial manuscript.Source: §6 · p. 10
• WG provided theoretical guidance for formalisation and collaborated on manuscript development.Source: §6 · p. 10

A Experimental setup

• All discussed autoencoders use the same hyperparameters, although the authors consider this likely suboptimal because they performed few ablations.Source: A Experimental setup · p. 12
• Experiments primarily use layer 18 of Qwen3-0.6B-Base [34], with the vanilla autoencoder training in about 10 minutes on an RTX 4080; the setup was also replicated on other models.Source: A Experimental setup · p. 12
• Muon is the optimiser; its convergence advantage over Adam is setup-dependent and roughly 2× to 10×, with a trapezoidal schedule and α warmup.Source: A Experimental setup · p. 12
• The architecture uses orthogonal initialisation, an expansion factor of 16, mixing expansion of 2, and default α = 0.1.Source: A Experimental setup · p. 12
• Inputs are L2-normalised and trained with sum of squares error for convenience; the linear autoencoder's objective is otherwise unchanged by this scaling.Source: A Experimental setup · p. 12

B Efficient evaluation of the loss

• The bilinear autoencoder kernel can be computed as an element-wise product of matrices, with contraction order and splitting used to avoid memory allocation.The kernel remains quadratic in latent count in memory, so it can be evaluated in blocks rather than fully materialised.Source: B Efficient evaluation of the loss · p. 13, Figure 10 · p. 13
• Computing only the contributing symmetric portion of BB⊤ halves the computation, and blockwise evaluation extends to mixed and ordered variants.Source: B Efficient evaluation of the loss · p. 13

C Autoencoder pseudo code

• The vanilla autoencoder pseudo code documents initialization and the forward pass, including elementwise multiplication of two projected feature vectors.Source: C Autoencoder pseudo code · pp. 13–14
• The pseudo code computes the reconstruction loss using a kernel formed from the elementwise product of the left and right projection Gram matrices.Source: C Autoencoder pseudo code · p. 14

D Comparison of reconstruction error

• Mixed autoencoders outperform vanilla counterparts everywhere except in the first model layers.The added bottleneck is interpreted as indicating that latents are better described as composites than independently; the proposed explanation for the first-layer exception is speculative.Source: D Comparison of reconstruction error · p. 14
• Across most layers and variants, bilinear autoencoders outperform the TopK baseline in product-space reconstruction after its error is roughly doubled for a fair comparison.The comparison corrects for bilinear models reconstructing product inputs rather than the original inputs.Source: D Comparison of reconstruction error · p. 14, Figure 11 · p. 15

E Anecdata and miscellaneous ablations

• Training on 16M tokens took minutes and yielded interpretable features; a 32× mixed autoencoder trained on 256M tokens reached SSE 0.14.The authors report that training metrics continued to decrease and say they lack resources to train on billions of tokens.Source: E Anecdata and miscellaneous ablations · p. 15
• Bilinear autoencoders produced interpretable features on GPT2-small and Pythia-160M, though reconstruction error was about 20% higher on equivalent setups.More competitive models of similar size yielded lower reconstruction errors.Source: E Anecdata and miscellaneous ablations · p. 15
• Expansions up to 64× beat the random-input reconstruction baseline; encoder nonlinearities reduced sparsity slightly but worsened reconstruction and increased loss.TopK required k > 500 for competitive reconstruction when the activation followed element-wise multiplication; Muon often reduced training time by an order of magnitude versus Adam.Source: E Anecdata and miscellaneous ablations · pp. 15–16

F Reconstruction with ordered latents

• Ordered autoencoders reached lower reconstruction errors significantly faster than the baseline, while ordering also influenced final reconstruction.Figure 12 compares reconstruction errors over latent prefixes and includes a greedily reordered vanilla baseline.Source: F Reconstruction with ordered latents · p. 16, Figure 12 · p. 16
• Cumulative reconstruction errors can make the first latent’s reconstruction orders of magnitude more important than the last’s.A gentler penalty schedule for later latents is proposed, but its effect on latent behaviour has not been tested.Source: F Reconstruction with ordered latents · p. 16

G Zooming in on a latent

• For the first vanilla latent without a sparsity penalty, 99% of sampled activations were below 2 · 10^-5; its activations corresponded to possessive pronouns.The authors say this rough activation-distribution shape is representative of other latents, though tail density and length vary.Source: G Zooming in on a latent · p. 16

H Random Manifolds

• The random latent visualisations include a split world/global versus believe/think concept, conjunctions, conditional clauses, and the linear bigram direction “at the.”Only top-activating samples are shown, so a long tail of weak activations can create an illusory gap.Source: Figure 14 · p. 17, H Random Manifolds · p. 17

I Similarity between latent bases

• A bilinear latent is represented by B_i = l_i ⊗ r_i; the best permutation’s recoverable norm in B^⊤B′ measures how much two bases align.This similarity compares full bilinear forms, making it stricter than comparing decoder directions alone.Source: I Similarity between latent bases · p. 18
• Self-similarity ranges from 0.8 for densely activating vanilla autoencoders to 0.65 for sparse counterparts, while cross-similarities lie between 0.5 and 0.6.The diagonal reflects latent interference because decoders are encoder transposes, not inverses; off-diagonal entries compare two autoencoders.Source: I Similarity between latent bases · p. 18
• The authors argue that consistency of the whole matters more than latent identifiability because bilinear latents can be rotated or combined without leaving the analysable bilinear space.They say analyses remain equivalent regardless of the basis used.Source: I Similarity between latent bases · p. 18