Mitchell Ostrow

To generate coherent responses, language models infer unobserved meaning from their input text sequence. One potential explanation for this capability arises from theories of delay embeddings in dynamical systems, which prove that unobserved variables can be recovered from the history of only a handful of observed variables. To test whether language models are effectively constructing delay embeddings, we measure the capacities of sequence models to reconstruct unobserved dynamics. We trained 1-layer transformer decoders and state-space sequence models on next-step prediction from noisy, partially-observed time series data. We found that each sequence layer can learn a viable embedding of the underlying system. However, state-space models have a stronger inductive bias than transformers-in particular, they more effectively reconstruct unobserved information at initialization, leading to more parameter-efficient models and lower error on dynamics tasks. Our work thus forges a novel connection between dynamical systems and deep learning sequence models via delay embedding theory.

The use of an internal model to infer and predict others‘ mental states and actions, broadly referred to as Theory of Mind (ToM), is a fundamental aspect of human social intelligence. Nevertheless, it remains unknown how these models are used during social interactions, and how they help an agent generalize to new contexts. We investigated a putative neural mechanism of ToM in a recurrent circuit through the lens of an artificial neural network trained with reinforcement learning (RL) to play a competitive matching pennies game against many algorithmic opponents. The network showed near-optimal performance against unseen opponents, indicating that it had acquired the capacity to adapt against new strategies online. Analysis of recurrent states during play against out-of-training-distribution (OOD) opponents in relation to those of within-training-distribution (WD) opponents revealed two similarity-based mechanisms by which the network might generalize: mapping to a known strategy (template matching) or known opponent category (interpolation). Even when the network’s strategy cannot be explained by template-matching or interpolation, the recurrent activity fell upon the low-dimensional manifold of the WD neural activity, suggesting the contribution of prior experience with WD opponents. Furthermore, these states occupied low-density edges of the WD-manifold, suggesting that the network can extrapolate beyond any learned strategy or category. Our results suggest that a neural implementation for ToM may be a reservoir of learned representations that provide the capacity for generalization via flexible access and reuse of these stored features.

How can we tell whether two neural networks are utilizing the same internal processes for a particular computation? This question is pertinent for multiple subfields of both neuroscience and machine learning, including neuroAI, mechanistic interpretability, and brain-machine interfaces. Standard approaches for comparing neural networks focus on the spatial geometry of latent states. Yet in recurrent networks, computations are implemented at the level of neural dynamics, which do not have a simple one-to-one mapping with geometry. To bridge this gap, we introduce a novel similarity metric that compares two systems at the level of their dynamics. Our method incorporates two components: Using recent advances in data-driven dynamical systems theory, we learn a high-dimensional linear system that accurately captures core features of the original nonlinear dynamics. Next, we compare these linear approximations via a novel extension of Procrustes Analysis that accounts for how vector fields change under orthogonal transformation. Via four case studies, we demonstrate that our method effectively identifies and distinguishes dynamic structure in recurrent neural networks (RNNs), whereas geometric methods fall short. We additionally show that our method can distinguish learning rules in an unsupervised manner. Our method therefore opens the door to novel data-driven analyses of the temporal structure of neural computation, and to more rigorous testing of RNNs as models of the brain.