Koopman Autoencoders Learn Neural Representation Dynamics

TL;DR

This work reframes neural network layerwise representations as nonlinear dynamical systems and introduces Koopman autoencoders (KAEs) to learn a surrogate linear evolution in an observable space. By coupling an encoder $\phi$, a decoder $\phi^{-1}$, and a linear operator $\mathcal{K}$, KAEs interpolate and edit neural representation trajectories with a loss that enforces reconstruction, linearity in the observable space, state-space consistency, and topology preservation via encoder isometry. The authors demonstrate that KAEs yield intermediate representations whose topology progressively simplifies, akin to the original network, and show practical model editing via EMMET, achieving targeted unlearning on MNIST but with caveats on datasets lacking neural collapse. The framework leverages topology (Betti numbers) and RSA-inspired metrics to quantify and interpret representation dynamics, offering a tool for analysis and rapid, targeted interventions in deep models with potential applications in safety and alignment.

Abstract

This paper explores a simple question: can we model the internal transformations of a neural network using dynamical systems theory? We introduce Koopman autoencoders to capture how neural representations evolve through network layers, treating these representations as states in a dynamical system. Our approach learns a surrogate model that predicts how neural representations transform from input to output, with two key advantages. First, by way of lifting the original states via an autoencoder, it operates in a linear space, making editing the dynamics straightforward. Second, it preserves the topologies of the original representations by regularizing the autoencoding objective. We demonstrate that these surrogate models naturally replicate the progressive topological simplification observed in neural networks. As a practical application, we show how our approach enables targeted class unlearning in the Yin-Yang and MNIST classification tasks.

Figures (6)

• Figure 1: The top three principal components of the neural representations from the first layer (L0) and all residual blocks (L1-5) of a multi-layer perceptron (MLP) with a ResNet-style architecture. Each plot contains $2\times10^3$ points and undergoes the preprocessing steps outlined in Section \ref{['sec:preprocessing']} before PCA for plotting. The model is trained on the Yin-Yang dataset kriener2022, a three-way classification task. See Appendix \ref{['app:dataset']} for details on architecture and dataset.
• Figure 2: A summary of our framework presented in Section \ref{['sec:framework']}. We gather neural representations from a trained, residual network and preprocess them to bring them into the same space. Afterwards, we train a Koopman autoencoder on a pair of the representations, resulting in predictive autoencoder with manipulable and visualizable observabe space.
• Figure 3: (Top left) The $\beta_0$ Betti numbers of the neural representations from each residual block of a residual MLP trained on MNIST. The Betti numbers are computed using the Vietoris-Rips complex at a filtration $\epsilon=0.166$. (Top right) The average $\beta_0$ Betti numbers of intermediate outputs, projected into state space, for five KAEs trained on the first and penultimate layer representations of the residual MLP. The Betti numbers are computed using the Vietoris-Rips complex at a filtration $\epsilon=0.14$. (Bottom) Select intermediate outputs from an MNIST KAE, projected into the state space. At each successive iteration, the topology is simplified until it arrives at the penultimate layer representations.
• Figure 4: (A) Each scatter plot displays $2\times10^3$ points projected onto the top three principal components (PCs) derived from representations in the penultimate layer. The leftmost plot shows PCs from the original MLP representations, while the remaining show PCs computed after embedding the representations into observable space via different KAEs. All PCs are aligned via the orthogonal Procrustes problem. (B) Betti curves, for $\beta0 \text{ and } \beta1$, across a filtration threshold of $\epsilon=4$ for the penultimate layer representations of the original model (black) and the observable space representations via different KAEs.
• Figure 5: $10^4$ points projected on the top three principal components of the neural representations produced by the Koopman operator in observable space before editing (left) and after editing (right). The KAE is trained on the first and penultimate-layer representations of a MNIST classifier. The operator is edited to forget class 4 (violet) by merging the outputs of that class with those of class 9 (light blue). The result of the merge is visible on the top right corner.