The natural world is richly structured over space and time. Much of this structure arises from the interplay between spatial geometry and motion. However, most existing world models ignore this structure, leading to an inability to generalize in dynamic environments. In this work, we show that enforcing equivariance between an agent’s representations and the world’s dynamics necessarily induces an efficient, structured memory. Concretely, we introduce Flow Equivariant World Modeling, a framework in which both self-motion and external object motion are unified as one-parameter Lie-group ``flows’’ acting on a latent world memory; and models are built to be equivariant with respect to these transformations. On 2D and 3D partially observed video world modeling benchmarks, we demonstrate that Flow Equivariant World Models significantly outperform comparable state-of-the-art diffusion-based and memory-augmented world modeling architectures in their ability to track and predict the locations of moving objects over long horizons.

H. Lillemark, B. Huang, F. Zhan, Y. Du, T. Anderson Keller

Paper: https://arxiv.org/abs/2601.01075
Project page: https://flowequivariantworldmodels.github.io
Accepted at ICML ‘26

https://github.com/AnonFloWM/Flow-Equivariant-World-Modeling

Data arrives at our senses as a continuous stream, smoothly transforming from one instant to the next. These smooth transformations can be viewed as continuous symmetries of the environment that we inhabit, defining equivalence relations between stimuli over time. In machine learning, neural network architectures that respect symmetries of their data are called equivariant and have provable benefits in terms of generalization ability and sample efficiency. To date, however, equivariance has been considered only for static transformations and feed-forward networks, limiting its applicability to sequence models, such as recurrent neural networks (RNNs), and corresponding time-parameterized sequence transformations. In this work, we extend equivariant network theory to this regime of ‘flows’ – one-parameter Lie subgroups capturing natural transformations over time, such as visual motion. We begin by showing that standard RNNs are generally not flow equivariant: their hidden states fail to transform in a geometrically structured manner for moving stimuli. We then show how flow equivariance can be introduced, and demonstrate that these models significantly outperform their non-equivariant counterparts in terms of training speed, length generalization, and velocity generalization, on both next step prediction and sequence classification. We present this work as a first step towards building sequence models that respect the time-parameterized symmetries which govern the world around us.

T. Anderson Keller

Accepted at NeurIPS 2025 (Spotlight, Top 13% accepted)
Paper: https://openreview.net/forum?id=N1KPOlcN6P
arXiv: https://arxiv.org/abs/2507.14793

https://neurips.cc/virtual/2025/loc/san-diego/poster/118412

https://kempnerinstitute.harvard.edu/research/deeper-learning/flow-equivariant-recurrent-neural-networks/

https://github.com/akandykeller/FERNN

• IAIFAI Colloquium Lecture: “Flow Equivariance”

IAIFAI Colloquium Lecture: “Flow Equivariance”

Orientation-rich images, such as fingerprints and textures, often exhibit coherent angular directional patterns that are challenging to model using standard generative approaches based on isotropic Euclidean diffusion. Motivated by the role of phase synchronization in biological systems, we propose a score-based generative model built on periodic domains by leveraging stochastic Kuramoto dynamics in the diffusion process. In neural and physical systems, Kuramoto models capture synchronization phenomena across coupled oscillators – a behavior that we re-purpose here as an inductive bias for structured image generation. In our framework, the forward process performs \textit{synchronization} among phase variables through globally or locally coupled oscillator interactions and attraction to a global reference phase, gradually collapsing the data into a low-entropy von Mises distribution. The reverse process then performs \textit{desynchronization}, generating diverse patterns by reversing the dynamics with a learned score function. This approach enables structured destruction during forward diffusion and a hierarchical generation process that progressively refines global coherence into fine-scale details. We implement wrapped Gaussian transition kernels and periodicity-aware networks to account for the circular geometry. Our method achieves competitive results on general image benchmarks and significantly improves generation quality on orientation-dense datasets like fingerprints and textures. Ultimately, this work demonstrates the promise of biologically inspired synchronization dynamics as structured priors in generative modeling.

Yue Song, T. Anderson Keller, Sevan Brodjian, Takeru Miyato, Yisong Yue, Pietro Perona, Max Welling

Accepted at NeurIPS 2025 (Poster)
Paper: https://arxiv.org/abs/2509.15328
OpenReview: https://openreview.net/forum?id=dxK2QgEKvz

https://github.com/KingJamesSong/OrientationDiffusion

Recent work has revealed that state space models (SSMs), while efficient for long-sequence processing, are fundamentally limited in their ability to represent formal languages—particularly due to time-invariant and real-valued recurrence structures. In this work, we draw inspiration from adaptive and structured dynamics observed in biological neural systems and introduce the Adaptive Unitary State Space Model (AUSSM): a novel class of SSMs that leverages skew-symmetric, input-dependent recurrence to achieve unitary evolution and high expressive power. Using algebraic automata theory, we prove that AUSSM can perform modulo counting and simulate solvable group automata at finite precision, enabling SSMs to model a broad class of regular languages out of reach for other SSM architectures. To overcome the practical inefficiencies of adaptive recurrence, we develop a separable convolution formulation and a CUDA implementation that enables scalable parallel training. Empirically, we show that AUSSM and its hybrid variant—interleaved with Mamba—outperform prior SSMs on formal algorithmic tasks such as parity and modular arithmetic, and achieve competent performance on real-world long time-series classification benchmarks. Our results demonstrate that adaptive unitary recurrence provides a powerful and efficient inductive bias for both symbolic and continuous sequence modeling.

Arjun Karuvally, Franz Nowak, T. Anderson Keller, Carmen Amo Alonso, Terrence J. Sejnowski, Hava T. Siegelmann

Accepted at NeurIPS 2025 (Poster)
Paper: https://openreview.net/forum?id=s4zitEu2R8

https://github.com/arjunkaruvally/AUSSM

When trained on data with missing symmetries, diffusion models face a fundamental challenge: how can they generate samples respecting symmetries they have never observed? We prove that this failure stems from the structure of the learning problem itself. Unconditional models must satisfy a global equivariance constraint, coupling all group elements into a single optimization that requires high-dimensional data extrapolation across gaps. In contrast, conditioning on group elements factorizes this into independent problems, transforming the task into low-dimensional function generalization. Our theory predicts—and experiments confirm—that this simple change yields 5-10× error reduction on held-out symmetries. On synthetic 2D rotation tasks, conditional models maintain low error even with 300° gaps while unconditional models collapse catastrophically. We further suggest that topology-aware group embeddings may help improve this generalization by ensuring smoother functions over the group manifold.

Sid Bharthulwar, T. Anderson Keller, Manos Theodosis, Demba E. Ba

Accepted at Symmetry and Geometry in Neural Representations (NeurReps) @ NeurIPS 2025 (Poster)
Paper: https://openreview.net/forum?id=UI82R3lwar

Foundation models have shown remarkable success in fitting biological visual systems; however, their black-box nature inherently limits their utility for under- standing brain function. Here, we peek inside a SOTA foundation model of neural activity (Wang et al., 2025) as a physiologist might, characterizing each ‘neuron’ based on its temporal response properties to parametric stimuli. We analyze how different stimuli are represented in neural activity space by building decoding man- ifolds, and we analyze how different neurons are represented in stimulus-response space by building neural encoding manifolds. We find that the different processing stages of the model (i.e., the feedforward encoder, recurrent, and readout modules) each exhibit qualitatively different representational structures in these manifolds. The recurrent module shows a jump in capabilities over the encoder module by “pushing apart” the representations of different temporal stimulus patterns. Our “tubularity” metric quantifies this stimulus-dependent development of neural activ- ity as biologically plausible. The readout module achieves high fidelity by using numerous specialized feature maps rather than biologically plausible mechanisms. Overall, this study provides a window into the inner workings of a prominent neural foundation model, gaining insights into the biological relevance of its internals through the novel analysis of its neurons’ joint temporal response patterns. Our findings suggest design changes that could bring neural foundation models into closer alignment with biological systems: introducing recurrence in early encoder stages, and constraining features in the readout module.

J. Bertram, L. Dyballa, T. Anderson Keller, S. Kinger, S. W. Zucker

Accepted at Data on the Brain and Mind @ NeurIPS 2025 (Workshop)
Paper: https://openreview.net/forum?id=jUy7vFgoZf Under Review at ICML ‘26

As Vision Transformers (ViTs) become standard backbones across vision, a mechanistic account of their computational phenomenology is now essential. Despite architectural cues that hint at dynamical structure, there is no settled framework that interprets Transformer depth as a well-characterized flow. In this work, we introduce the Block-Recurrent Hypothesis (BRH), arguing that trained ViTs admit a block-recurrent depth structure such that the computation of the original L blocks can be accurately rewritten using only k ≪ L distinct blocks applied recurrently. Across diverse ViTs, between-layer representational similarity matrices suggest few contiguous phases. To determine whether this reflects reusable computation, we operationalize our hypothesis in the form of block recurrent surrogates of pretrained ViTs, which we call Recurrent Approximations to Phase-structured TransfORmers (Raptor). Using small-scale ViTs, we demonstrate that phase-structure metrics correlate with our ability to accurately fit Raptor and identify the role of stochastic depth in promoting the recurrent block structure. We then provide an empirical existence proof for BRH in foundation models by showing that we can train aRaptor model to recover 94% of DINOv2 ImageNet-1k linear probe accuracy in only 2 blocks. To provide a mechanistic account of these observations, we leverage our hypothesis to develop a program of Dynamical Interpretability. We find (i) directional convergence into class-dependent angular basins with self-correcting trajectories under small perturbations (ii) token-specific dynamics, where clsexecutes sharp late reorientations while patch tokens exhibit strong late-stage coherence reminiscent of a mean-field effect and converge rapidly toward their mean direction and (iii) a collapse of the update field to low rank in late depth, consistent with convergence to low-dimensional attractors. Altogether, we find that a compact recurrent program emerges along the depth of ViTs, pointing to a low-complexity normative solution that enables these models to be studied through principled dynamical systems analysis.

Mozes Jacobs, Thomas Fel, R. Hakim, A. Brondetta, Demba E. Ba, T. Anderson Keller

ArXiv: https://arxiv.org/pdf/2512.19941
OpenReview: https://openreview.net/forum?id=gH3HhnfWLC
Accepted at ICLR ‘26

https://kempnerinstitute.github.io/raptor/

Diffusion models have had remarkable success over the last decade in generating a diverse set of visually plausible images. These models work by transforming the data to a centered Gaussian and then learning the reverse process by training a neural network to approximate the score of the underlying distribution. A variety of architectures from CNNs, to U-Nets, to transformers have been used as the score-approximation network in diffusion modeling. We propose an analytically solvable parameterization of the score function using an expansion in a wavelet basis. In particular, we derive interpretable optimal score functions in a 2D, orthogonal wavelet basis in terms of the moments of the data distribution. We use this parametrization to provide an architecture-agnostic, moment-based analysis that reveals which attributes of the data distribution tend to matter most for denoising. Our score machine is flexible enough to partially mimic the relevant inductive biases of multiple architectures, including U-Nets, and CNNs, taking a step towards understanding why different score architectures can exhibit distinct generative behavior. Since our score is solvable in terms of the moments of the data, we can begin to understand how the data distribution interacts with the score network to produce the behavior we observe in diffusion models.

Emma Lucia Byrnes Finn, Binxu Wang, T. Anderson Keller, Demba E. Ba

Accepted in Proceedings of AISTATS ‘26
Also Accepted at SPIGM Workshop @ NeurIPS 2025
Paper: https://openreview.net/forum?id=fmlzopxjxq

Traveling waves of neural activity are widely observed in the brain, but their precise computational function remains unclear. One prominent hypothesis is that they enable the transfer and integration of spatial information across neural populations. However, few computational models have explored how traveling waves might be harnessed to perform such integrative processing. Drawing inspiration from the famous “Can one hear the shape of a drum?” problem – which highlights how normal modes of wave dynamics encode geometric information – we investigate whether similar principles can be leveraged in artificial neural networks. Specifically, we introduce convolutional recurrent neural networks that learn to produce traveling waves in their hidden states in response to visual stimuli, enabling spatial integration. By then treating these wave-like activation sequences as visual representations themselves, we obtain a powerful representational space that outperforms local feed-forward networks on tasks requiring global spatial context. In particular, we observe that traveling waves effectively expand the receptive field of locally connected neurons, supporting long-range encoding and communication of information. We demonstrate that models equipped with this mechanism solve visual semantic segmentation tasks demanding global integration, significantly outperforming local feed-forward models and rivaling non-local U-Net models with fewer parameters. As a first step toward traveling-wave-based communication and visual representation in artificial networks, our findings suggest wave-dynamics may provide efficiency and training stability benefits, while simultaneously offering a new framework for connecting models to biological recordings of neural activity.

Mozes Jacobs, Roberto C. Budzinski, Lyle Muller, Demba E. Ba, T. Anderson Keller

Accepted at CCN 2025 (Oral presentation, Top 7%)
Paper: https://openreview.net/forum?id=QEzqo546V5

https://github.com/KempnerInstitute/traveling-waves-integrate

• Tweet-print

Tweet-print

In the physical world, almost all information is transmitted through traveling waves -- why should it be any different in your neural network?

Super excited to share recent work with the brilliant @mozesjacobs: "Traveling Waves Integrate Spatial Information Through Time"
1/14 pic.twitter.com/Bs4UKR7j21

— Andy Keller (@t_andy_keller) March 10, 2025

This book introduces approaches to generalize the benefits of equivariant deep learning to a broader set of learned structures through learned homomorphisms. In the field of machine learning, the idea of incorporating knowledge of data symmetries into artificial neural networks is known as equivariant deep learning and has led to the development of cutting edge architectures for image and physical data processing. The power of these models originates from data-specific structures ingrained in them through careful engineering. To-date however, the ability for practitioners to build such a structure into models is limited to situations where the data must exactly obey specific mathematical symmetries. The authors discuss naturally inspired inductive biases, specifically those which may provide types of efficiency and generalization benefits through what are known as homomorphic representations, a new general type of structured representation inspired from techniques in physics and neuroscience. A review of some of the first attempts at building models with learned homomorphic representations are introduced. The authors demonstrate that these inductive biases improve the ability of models to represent natural transformations and ultimately pave the way to the future of efficient and effective artificial neural networks.

Yue Song, T. Anderson Keller, Nicu Sebe, Max Welling

Book: https://link.springer.com/book/10.1007/978-3-031-88111-4
DOI: https://doi.org/10.1007/978-3-031-88111-4