Publications
(*) denotes equal contribution
Accepted
2026
- Effectiveness of Curriculum Learning Depends on Reward Sparsity and Competing OptimaJohn J. Vastola, Ann Huang, Satpreet H. Singh, Samuel J. Gershman, and Kanaka RajanIn Advances in Neural Information Processing Systems 2026
While animals and people tend to learn a task more quickly and reliably when they first train on simpler versions of it, curriculum learning’s effectiveness in artificial settings varies, and appears substantially greater in reinforcement learning than in supervised learning. In this work, we construct and analyze a minimal model of policy learning to better understand why. We consider an open-loop, episodic target interception task whose parameters—especially how close the agent must be to the target to receive an informative reward signal—can be chosen so that the agent either receives informative rewards frequently (’dense’ rewards), or rarely (’sparse’ rewards). We mathematically and numerically analyze the dynamics of REINFORCE-based policy learning in both cases. In the dense reward case, curricula have only marginal benefits; in the sparse case, curricula can be required for learning to occur at all. We find that this is especially true when the agent has the conflicting goals of minimizing effort and accumulating high reward, since this can produce distinct reward landscape optima which compete with one another. Our characterization of useful curricula in this setting indicates that they address two issues: they (i) make rewards less sparse to speed up learning, and (ii) shape the reward landscape to steer agents away from bad local optima.
- Temporal Smoothness Constraints on Efficient Neurobiological Codes Imply Temporal SpecializationJohn J. Vastola, and Samuel J. GershmanIn Advances in Neural Information Processing Systems 2026
Neurobiological circuits may be energy efficient in part because there is a division of labor: different subsystems compute or represent different things. In light of this, it is interesting that a major difference between subsystems is the time scale on which they typically vary. Here, we argue that temporal smoothness constraints—the idea that it may be energetically costly for systems to vary much more quickly or slowly than their typical time scale—imply a certain division of labor with respect to temporal stimulus features. In particular, slow subsystems ought to represent slowly varying stimulus features, and fast subsystems ought to represent quickly varying features. We formulate a novel efficient coding model that we use to investigate this claim, and exactly solve for the optimal division of labor. In addition to finding that slower subsystems ought to represent more slowly varying temporal features, we find surprising structure of mathematical interest: optimal codes have subsystems which represent orthogonal temporal features, and which correspond to the eigenfunctions of a Sturm-Liouville problem.
- Quantifying the failure of symmetry learning in linear diffusion models due to initialization and finite-sample effectsJohn J. Vastola, and Kanaka RajanIn NeurIPS 2026 Workshop on Symmetry and Geometry in Neural Representations 2026
Data distributions often possess one or more symmetries: for example, all orientations of any specific dog are included in the distribution of all possible dog pictures. Despite the omnipresence of such symmetries, modern generative modeling approaches may fail to incorporate or learn them, which can lead to lower sample efficiency and poorer generalization. To what extent do diffusion models fail to learn data distribution symmetries, and what sources of symmetry breaking are most responsible? Here, we answer this question in the theoretically tractable setting of linear diffusion models. We show that of three potential sources of symmetry breaking—a non-symmetric initial weight matrix, a noisy learning target, and finite-sample noise—target noise produces the most robust symmetry breaking, since it persists even in the sample-rich regime. Using local laws from random matrix theory, we derive formulas for average symmetry error in fairly general cases involving symmetry within a subspace and show that smaller symmetries are harder to learn.
- Solvable models of learning to pursue a moving targetJohn J. Vastola, and Kanaka RajanReinforcement Learning Journal 2026
Agents tasked with intercepting a moving target must learn how to reach it, ideally as efficiently as possible. As observers of animal pursuit behavior have noticed, different strategies are possible: on one extreme, the agent reactively moves toward the target’s current location; on the other, the agent predicts the target’s future location and moves directly there. Motivated by the desire to understand how such strategies might be learned, we introduce a continuous-time linear-quadratic model of open-loop target pursuit whose optimal control strategies interpolate between these possibilities. Usefully, a small number of interpretable parameters control which type of strategy is optimal, and it is possible to derive closed-form solutions for optimal strategies, policy learning dynamics, and value learning dynamics. Exploiting our model’s linear structure, we find that the time scales of learning precisely correspond to the eigenvalues of certain matrices, and that relevant eigenvalue spectra generically have a gap. We show that this gap indicates that agents tend to learn how to reach their target before optimizing their movement. Our results provide a detailed mathematical characterization of pursuit behavior and its learning dynamics, which can both serve as a benchmark for empirical work, and as a foundation for more elaborate theoretical treatments of pursuit.
Publications
2026
- Implications of hierarchical Markov models of behavior: on irreversibility, predictability, and dimensionalityJohn J. Vastola, and Kanaka RajanIn 9th Annual Conference on Cognitive Computational Neuroscience 2026
The maturation of quantitative tools for studying the high-level structure of animal behavior, and especially tools which represent spontaneous behavior as a sequence of stereotyped and neurally well-defined ’syllables’, demands that the field revisit a fundamental theoretical question: if the coarse structure of behavior can be accurately described by Markov models, what do these models really tell us about behavior? In this work, we explore the theoretical implications of these models and discuss how they allow us to quantitatively formulate questions about the sequence-like nature and effective dimensionality of behavior. One important insight is that the eigenvalues and eigenvectors of various model-associated matrices furnish interpretable time scales and modifications of behavior that occur on those time scales. We illustrate our points using both toy examples and Markov models fit to real data. By analyzing the consequences of Markov representations, we clarify the theoretical meaning of progress in quantifying behavior.
2025
- Gradient Descent as Loss Landscape Navigation: a Normative Framework for Deriving Learning RulesJohn J. Vastola, Samuel J Gershman, and Kanaka RajanIn Advances in Neural Information Processing Systems 2025
Learning rules—prescriptions for updating model parameters to improve performance—are typically assumed rather than derived. Why do some learning rules work better than others, and under what assumptions can a given rule be considered optimal? We propose a theoretical framework that casts learning rules as policies for navigating (partially observable) loss landscapes, and identifies optimal rules as solutions to an associated optimal control problem. A range of well-known rules emerge naturally within this framework under different assumptions: gradient descent from short-horizon optimization, momentum from longer-horizon planning, natural gradients from accounting for parameter space geometry, non-gradient rules from partial controllability, and adaptive optimizers like Adam from online Bayesian inference of loss landscape shape. We further show that continual learning strategies like weight resetting can be understood as optimal responses to task uncertainty. By unifying these phenomena under a single objective, our framework clarifies the computational structure of learning and offers a principled foundation for designing adaptive algorithms.
- A Variational Manifold Embedding Framework for Nonlinear Dimensionality ReductionJohn J. Vastola, Samuel J. Gershman, and Kanaka RajanIn NeurIPS 2025 Workshop on Symmetry and Geometry in Neural Representations 2025
Dimensionality reduction algorithms like principal component analysis (PCA) are workhorses of machine learning and neuroscience, but each has well-known limitations. Variants of PCA are simple and interpretable, but not flexible enough to capture nonlinear data manifold structure. More flexible approaches have other problems: autoencoders are generally difficult to interpret, and graph-embedding-based methods can produce pathological distortions in manifold geometry. Motivated by these shortcomings, we propose a variational framework that casts dimensionality reduction algorithms as solutions to an optimal manifold embedding problem. By construction, this framework permits nonlinear embeddings, allowing its solutions to be more flexible than PCA. Moreover, the variational nature of the framework has useful consequences for interpretability: each solution satisfies a set of partial differential equations, and can be shown to reflect symmetries of the embedding objective. We discuss these features in detail and show that solutions can be analytically characterized in some cases. Interestingly, one special case exactly recovers PCA.
- Deep RL Needs Deep Behavior Analysis: Exploring Implicit Planning by Model-Free Agents in Open-Ended EnvironmentsRiley Simmons-Edler, Ryan P. Badman, Felix Berg, Raymond Chua, John J. Vastola, Joshua Lunger, William Qian, and Kanaka RajanIn Advances in Neural Information Processing Systems 2025
Understanding the behavior of deep reinforcement learning (DRL) agents—particularly as task and agent sophistication increase—requires more than simple comparison of reward curves, yet standard methods for behavioral analysis remain underdeveloped in DRL. We apply tools from neuroscience and ethology to study DRL agents in a novel, complex, partially observable environment, ForageWorld, designed to capture key aspects of real-world animal foraging—including sparse, depleting resource patches, predator threats, and spatially extended arenas. We use this environment as a platform for applying joint behavioral and neural analysis to agents, revealing detailed, quantitatively grounded insights into agent strategies, memory, and planning. Contrary to common assumptions, we find that model-free RNN-based DRL agents can exhibit structured, planning-like behavior purely through emergent dynamics—without requiring explicit memory modules or world models. Our results show that studying DRL agents like animals—analyzing them with neuroethology-inspired tools that reveal structure in both behavior and neural dynamics—uncovers rich structure in their learning dynamics that would otherwise remain invisible. We distill these tools into a general analysis framework linking core behavioral and representational features to diagnostic methods, which can be reused for a wide range of tasks and agents. As agents grow more complex and autonomous, bridging neuroscience, cognitive science, and AI will be essential—not just for understanding their behavior, but for ensuring safe alignment and maximizing desirable behaviors that are hard to measure via reward. We show how this can be done by drawing on lessons from how biological intelligence is studied.
- Monod: model-based discovery and integration through fitting stochastic transcriptional dynamics to single-cell sequencing dataGennady Gorin, Tara Chari, Maria Carilli, John J. Vastola, and Lior PachterNature Methods Nov 2025
Single-cell RNA sequencing analysis centers on illuminating cell diversity and understanding the transcriptional mechanisms underlying cellular function. These datasets are large, noisy and complex. Current analyses prioritize noise removal and dimensionality reduction to tackle these challenges and extract biological insight. We propose an alternative, physical approach to leverage the stochasticity, size and multimodal nature of these data to explicitly distinguish their biological and technical facets while revealing the underlying regulatory processes. With the Python package Monod, we demonstrate how nascent and mature RNA counts, present in most published datasets, can be meaningfully ‘integrated’ under biophysical models of transcription. By using variation in these modalities, we can identify transcriptional modulation not discernible through changes in average gene expression, quantitatively compare mechanistic hypotheses of gene regulation, analyze transcriptional data from different technologies within a common framework and minimize the use of opaque or distortive normalization and transformation techniques.
- Generalization through variance: how noise shapes inductive biases in diffusion modelsJohn J. VastolaIn The Thirteenth International Conference on Learning Representations Nov 2025
How diffusion models generalize beyond their training set is not known, and is somewhat mysterious given two facts: the optimum of the denoising score matching (DSM) objective usually used to train diffusion models is the score function of the training distribution; and the networks usually used to learn the score function are expressive enough to learn this score to high accuracy. We claim that a certain feature of the DSM objective—the fact that its target is not the training distribution’s score, but a noisy quantity only equal to it in expectation—strongly impacts whether and to what extent diffusion models generalize. In this paper, we develop a mathematical theory that partly explains this ’generalization through variance’ phenomenon. Our theoretical analysis exploits a physics-inspired path integral approach to compute the distributions typically learned by a few paradigmatic under- and overparameterized diffusion models. We find that the distributions diffusion models effectively learn to sample from resemble their training distributions, but with ‘gaps’ filled in, and that this inductive bias is due to the covariance structure of the noisy target used during training. We also characterize how this inductive bias interacts with feature-related inductive biases.
- Dynamical symmetries in the fluctuation-driven regime: an application of Noether’s theorem to noisy dynamical systemsJohn J. VastolaIn Proceedings of the 3rd NeurIPS Workshop on Symmetry and Geometry in Neural Representations Nov 2025
Noether’s theorem provides a powerful link between continuous symmetries and conserved quantities for systems governed by some variational principle. Perhaps unfortunately, most dynamical systems of interest in neuroscience and artificial intelligence cannot be described by any such principle. On the other hand, nonequilibrium physics provides a variational principle that describes how fairly generic noisy dynamical systems are most likely to transition between two states; in this work, we exploit this principle to apply Noether’s theorem, and hence learn about how the continuous symmetries of dynamical systems constrain their most likely trajectories. We identify analogues of the conservation of energy, momentum, and angular momentum, and briefly discuss examples of each in the context of models of decision-making, recurrent neural networks, and diffusion generative models.
2024
- The Unreasonable Effectiveness of Gaussian Score Approximation for Diffusion Models and its ApplicationsBinxu Wang, and John J. VastolaTransactions on Machine Learning Research Nov 2024
Diffusion models have achieved remarkable results in multiple domains of generative modeling. By learning the gradient of smoothed data distributions, they can iteratively generate samples from complex distributions, e.g., of natural images. The learned score function enables their generalization capabilities, but how the learned score relates to the score of the underlying data manifold remains largely unclear. Here, we aim to elucidate this relationship by comparing the learned scores of neural-network-based models to the scores of two kinds of analytically tractable distributions: Gaussians and Gaussian mixtures. The simplicity of the Gaussian model makes it particularly attractive from a theoretical point of view, and we show that it admits a closed-form solution and predicts many qualitative aspects of sample generation dynamics. We claim that the learned neural score is dominated by its linear (Gaussian) approximation for moderate to high noise scales, and supply both theoretical and empirical arguments to support this claim. Moreover, the Gaussian approximation empirically works for a larger range of noise scales than naive theory suggests it should, and is preferentially learned by networks early in training. At smaller noise scales, we observe that learned scores are better described by a coarse-grained (Gaussian mixture) approximation of training data than by the score of the training distribution, a finding consistent with generalization. Our findings enable us to precisely predict the initial phase of trained models’ sampling trajectories through their Gaussian approximations. We show that this allows one to leverage the Gaussian analytical solution to skip the first 15-30% of sampling steps while maintaining high sample quality (with a near state-of-the-art FID score of 1.93 on CIFAR-10 unconditional generation). This forms the foundation of a novel hybrid sampling method, termed \textitanalytical teleportation, which can seamlessly integrate with and accelerate existing samplers, including DPM-Solver-v3 and UniPC. Our findings strengthen the field’s theoretical understanding of how diffusion models work and suggest ways to improve the design and training of diffusion models.
- Optimal packing of attractor states in neural representationsJohn J. VastolaIn Proceedings of the 2nd NeurIPS Workshop on Symmetry and Geometry in Neural Representations Nov 2024
Animals’ internal states reflect variables like their position in space, orientation, decisions, and motor actions—but how should these internal states be arranged? Internal states which frequently transition between one another should be close enough that transitions can happen quickly, but not so close that neural noise significantly impacts the stability of those states, and how reliably they can be encoded and decoded. In this paper, we study the problem of striking a balance between these two concerns, which we call an ‘optimal packing’ problem since it resembles mathematical problems like sphere packing. While this problem is generally extremely difficult, we show that symmetries in environmental transition statistics imply certain symmetries of the optimal neural representations, which allows us in some cases to exactly solve for the optimal state arrangement. We focus on two toy cases: uniform transition statistics, and cyclic transition statistics.
2023
- Studying stochastic systems biology of the cell with single-cell genomics dataGennady Gorin, John J. Vastola, and Lior PachterCell Systems Oct 2023
Recent experimental developments in genome-wide RNA quantification hold considerable promise for systems biology. However, rigorously probing the biology of living cells requires a unified mathematical framework that accounts for single-molecule biological stochasticity in the context of technical variation associated with genomics assays. We review models for a variety of RNA transcription processes, as well as the encapsulation and library construction steps of microfluidics-based single-cell RNA sequencing, and present a framework to integrate these phenomena by the manipulation of generating functions. Finally, we use simulated scenarios and biological data to illustrate the implications and applications of the approach.
- Causal inference during closed-loop navigation: parsing of self- and object-motionJean-Paul Noel, Johannes Bill, Haoran Ding, John J. Vastola, Gregory C. DeAngelis, Dora E. Angelaki, and Jan DrugowitschPhilosophical Transactions of the Royal Society B: Biological Sciences Oct 2023
A key computation in building adaptive internal models of the external world is to ascribe sensory signals to their likely cause(s), a process of causal inference (CI). CI is well studied within the framework of two-alternative forced-choice tasks, but less well understood within the cadre of naturalistic action–perception loops. Here, we examine the process of disambiguating retinal motion caused by self- and/or object-motion during closed-loop navigation. First, we derive a normative account specifying how observers ought to intercept hidden and moving targets given their belief about (i) whether retinal motion was caused by the target moving, and (ii) if so, with what velocity. Next, in line with the modelling results, we show that humans report targets as stationary and steer towards their initial rather than final position more often when they are themselves moving, suggesting a putative misattribution of object-motion to the self. Further, we predict that observers should misattribute retinal motion more often: (i) during passive rather than active self-motion (given the lack of an efference copy informing self-motion estimates in the former), and (ii) when targets are presented eccentrically rather than centrally (given that lateral self-motion flow vectors are larger at eccentric locations during forward self-motion). Results support both of these predictions. Lastly, analysis of eye movements show that, while initial saccades toward targets were largely accurate regardless of the self-motion condition, subsequent gaze pursuit was modulated by target velocity during object-only motion, but not during concurrent object- and self-motion. These results demonstrate CI within action–perception loops, and suggest a protracted temporal unfolding of the computations characterizing CI. This article is part of the theme issue ‘Decision and control processes in multisensory perception’.
2022
- Interpretable and tractable models of transcriptional noise for the rational design of single-molecule quantification experimentsGennady Gorin*, John J. Vastola*, Meichen Fang, and Lior PachterNature Communications Dec 2022
The question of how cell-to-cell differences in transcription rate affect RNA count distributions is fundamental for understanding biological processes underlying transcription. Answering this question requires quantitative models that are both interpretable (describing concrete biophysical phenomena) and tractable (amenable to mathematical analysis). This enables the identification of experiments which best discriminate between competing hypotheses. As a proof of principle, we introduce a simple but flexible class of models involving a continuous stochastic transcription rate driving a discrete RNA transcription and splicing process, and compare and contrast two biologically plausible hypotheses about transcription rate variation. One assumes variation is due to DNA experiencing mechanical strain, while the other assumes it is due to regulator number fluctuations. We introduce a framework for numerically and analytically studying such models, and apply Bayesian model selection to identify candidate genes that show signatures of each model in single-cell transcriptomic data from mouse glutamatergic neurons.
- Is the information geometry of probabilistic population codes learnable?John J. Vastola, Zach Cohen, and Jan DrugowitschIn NeurIPS 2022 Workshop on Symmetry and Geometry in Neural Representations Dec 2022
One reason learning the geometry of latent neural manifolds from neural activity data is difficult is that the ground truth is generally not known, which can make manifold learning methods hard to evaluate. Probabilistic population codes (PPCs), a class of biologically plausible and self-consistent models of neural populations that encode parametric probability distributions, may offer a theoretical setting where it is possible to rigorously study manifold learning. It is natural to define the neural manifold of a PPC as the statistical manifold of the encoded distribution, and we derive a mathematical result that the information geometry of the statistical manifold is directly related to measurable covariance matrices. This suggests a simple but rigorously justified decoding strategy based on principal component analysis, which we illustrate using an analytically tractable PPC.
2021
- Solving the chemical master equation for monomolecular reaction systems and beyond: a Doi-Peliti path integral viewJohn J. VastolaJournal of Mathematical Biology Oct 2021
The chemical master equation (CME) is a fundamental description of interacting molecules commonly used to model chemical kinetics and noisy gene regulatory networks. Exact time-dependent solutions of the CME—which typically consists of infinitely many coupled differential equations—are rare, and are valuable for numerical benchmarking and getting intuition for the behavior of more complicated systems. Jahnke and Huisinga’s landmark calculation of the exact time-dependent solution of the CME for monomolecular reaction systems is one of the most general analytic results known; however, it is hard to generalize, because it relies crucially on special properties of monomolecular reactions. In this paper, we rederive Jahnke and Huisinga’s result on the time-dependent probability distribution and moments of monomolecular reaction systems using the Doi-Peliti path integral approach, which reduces solving the CME to evaluating many integrals. While the Doi-Peliti approach is less intuitive, it is also more mechanical, and hence easier to generalize. To illustrate how the Doi-Peliti approach can go beyond the method of Jahnke and Huisinga, we also find an explicit and exact time-dependent solution to a problem involving an autocatalytic reaction that Jahnke and Huisinga identified as not solvable using their method. Most interestingly, we are able to find a formal exact time-dependent solution for any CME whose list of reactions involves only zero and first order reactions, which may be the most general result currently known. This formal solution also yields a useful algorithm for efficiently computing numerical solutions to CMEs of this type.
2020
- Chemical Langevin equation: A path-integral view of Gillespie’s derivationJohn J. Vastola, and William R. HolmesPhys. Rev. E Mar 2020
In 2000, Gillespie rehabilitated the chemical Langevin equation (CLE) by describing two conditions that must be satisfied for it to yield a valid approximation of the chemical master equation (CME). In this work, we construct an original path-integral description of the CME and show how applying Gillespie’s two conditions to it directly leads to a path-integral equivalent to the CLE. We compare this approach to the path-integral equivalent of a large system size derivation and show that they are qualitatively different. In particular, both approaches involve converting many sums into many integrals, and the difference between the two methods is essentially the difference between using the Euler-Maclaurin formula and using Riemann sums. Our results shed light on how path integrals can be used to conceptualize coarse-graining biochemical systems and are readily generalizable.
Submitted
2026
- A theory of noisy linear recurrent neural network learning dynamicsJohn J. Vastola, and Kanaka Rajan2026
- On the stability of learningJohn J. Vastola2026
- Field-theoretic diffusion models and sampling as spatial pattern formationJohn J. Vastola, and Binxu Wang2026
- Spatial Reward Sparsity Drives Slow Policy Learning While Temporal Reward Sparsity Delays RefinementJohn J. Vastola, and Kanaka Rajan2026
- Curriculum learning in brains and machinesJohn J. Vastola, Andrew Saxe, Christopher Summerfield, and Kanaka RajanSubmitted to Nature Reviews Neuroscience 2026
Teaching students and non-human animals to perform complex tasks usually involves structured curricula. Teachers may give students a series of tasks that increase in difficulty or complexity, or perhaps which encourage them to learn specific skills that can be reused to solve future problems. But although the idea of using structured training to improve learning is obvious and widely used, curriculum learning (CL) remains surprisingly poorly understood. In machine learning, where optimizing learning has received sustained empirical and theoretical attention, curricula have only sometimes proved useful. Fundamental questions regarding when curricula are useful, whether good curricula should go from easy to hard, and how to define and determine optimal curricula lack general-purpose answers. Here, we synthesize evidence from across neuroscience, psychology, and artificial intelligence to propose a tentative framework for thinking about when and why curricula are effective. Motivated especially by a geometric picture of learning dynamics, we argue that good curricula should involve informative learning signals and performance-sensitive difficulty modulation, and should reorganize the representational geometry of learners (e.g., to support compositionality). We illustrate our points using concrete examples from animal and machine learning. Taken seriously, this perspective suggests a route towards practical and effective curriculum learning.
2025
- Storing long-lived memories via molecular error correction: a minimal mathematical model of Crick’s memory switchJohn J. Vastola, Tejas Ramdas, and Samuel J. GershmanbioRxiv 2025
Molecular memory storage mechanisms face two serious obstacles: they must be robust to the intrinsic noise characteristic of molecular-scale processes, and they must rely on components (like proteins) whose effective lifetimes are limited for biophysical reasons. For example, if one stores a memory by attaching a post-translational ‘mark’ to a protein, naively the lifetime of that memory is limited by protein turnover and ambient demarking activity. In this work, we investigate the principles of persistent and noise-robust molecular memory by studying simple mathematical models of protein-marking-state-based memory switches. We pay special attention to two possible strategies for improving memory storage: effectively averaging over noise by using a large number of proteins, and an enzyme-based ‘error correction’ mechanism first proposed by Francis Crick. Through an exhaustive mathematical analysis of these models, which includes solving most of them exactly, we find that error correction provides an effective mechanism for storing long-lived memories, and that using large populations of proteins only slightly improves the lifetime of stored memories. But we also find that error correction is not always useful, since too much of it can make the probability of false positives excessively high. Our findings lend credence to Crick’s proposal, and suggest motifs that may be exploited by potential intracellular memory storage mechanisms. Finally, in addressing these issues, we also provide general insight about how to define memory performance in the molecular setting, as well as associated notions like lifetime and capacity.