Setting up for failure: automatic discovery of the neural mechanisms of cognitive errors
Authors: Puria Radmard
Paul M. Bays
Máté Lengyel
Venue: ICLR 2026 Poster (2026)
Keywords: neuroscience
working memory
recurrent neural networks
diffusion models
behavioral modeling
Primary area: applications to neuroscience & cognitive science
TL;DR: Training RNNs to reproduce realistic error patterns (rather than optimal performance) produces networks that better mimic biological neural computation, demonstrated through a working memory task where networks were taught to make swap errors.

Abstract: Discovering the neural mechanisms underpinning cognition is one of the grand challenges of neuroscience. However, previous approaches for building models of recurrent neural network (RNN) dynamics that explain behaviour required iterative refinement of architectures and/or optimisation objectives, resulting in a piecemeal, and mostly heuristic, human-in-the-loop process. Here, we offer an alternative approach that automates the discovery of viable RNN mechanisms by explicitly training RNNs to reproduce behaviour, including the same characteristic errors and suboptimalities, that humans and animals produce in a cognitive task. Achieving this required two main innovations. First, as the amount of behavioural data that can be collected in experiments is often too limited to train RNNs, we use a non-parametric generative model of behavioural responses to produce surrogate data for training RNNs. Second, to capture all relevant statistical aspects of the data, rather than a limited number of hand-picked low-order moments as in previous moment-matching-based approaches, we developed a novel diffusion model-based approach for training RNNs. To showcase the potential of our approach, we chose a visual working memory task as our test-bed, as behaviour in this task is well known to produce response distributions that are patently multimodal (due to so-called swap errors). The resulting network dynamics correctly predicted previously reported qualitative features of neural data recorded in macaques. Importantly, these results were not possible to obtain with more traditional approaches, i.e., when only a limited set of behavioural signatures (rather than the full richness of behavioural response distributions) were fitted, or when RNNs were trained for task optimality (instead of reproducing behaviour). Our approach also yields novel predictions about the mechanism of swap errors, which can be readily tested in experiments. These results suggest that fitting RNNs to rich patterns of behaviour provides a powerful way to automatically discover the neural network dynamics supporting important cognitive functions.

Published as a conference paper at ICLR 2026
SETTING UP FOR FAILURE: AUTOMATIC DISCOVERY OF THE NEURAL MECHANISMS OF COGNITIVE ERRORS
Puria Radmard Paul M. Bays
Máté Lengyel*:
"Department of Engineering, University of Cambridge
Department of Psychology, University of Cambridge *Department of Cognitive Science Central European University
(pr450, pmb20, ml (68) @cam. ac.
ABSTRACT
Discovering the neural mechanisms underpinning cognition is of the grand challenges of neuroscience. However, previous approaches for building models of recurrent neural network (RNN) dynamics that explain behaviour required iterative refinement of architectures and/or optimisation objectives, resulting in a piecemeal, and mostly heuristic, numan-in-the-llee process. Here, we offer an alternative approach that automates the discovery of viable RNN mechanisms by explicitly training RNNs reproduce behaviour, including the same characteristic errors and suboptimalities, that humans and animals produce in a cognitive task. Achieving this required two main innovations. First, as the amount of behavioural data that can be collected in experiments is often too limited to train RNNs, we use non- parametric generative model of behavioural responses to produce surrogate data for training RNNs. Second, to capture all relevant statistical aspects of the data, rather than a limited number hand-picked low-order moments as in previous moment- matching-basee approaches, for training RNNs. To showcase the potential of our approach, we chose visual working memory task as our test-bed, as behaviour in this task is well known produce response distributions that are patently multimodal (due to so-called swap errors). The resulting network dynamics correctly predicted previously reported qualitative features of neural data recorded in macaques. Importantly, these results were not possible obtain with more traditional approaches, i.e., when only a limited set of behavioural signatures (rather than the full richness of behavioural response distributions) were fitted, or when RNNs were trained for task optimality (instead of reproducing behaviour). Our approach also yields novel predictions about the mechanism of swap errors, which can be readily tested experiments. These results suggest that fitting RNNs to rich patterns of behaviour provides powerful way to automatically discover the neural network dynamics supporting important cognitive functions.
developed novel diffusion model-based approach
INTRODUCTION
An important goal for computational neuroscience use behavioural data to generate viable and testable hypotheses about the neural network mechanisms that underpin cognition. To achieve this goal, the following requirements need to be met: (1) hypotheses should be formulated as recurrent neural networks (RNNs) so that dynamical neural mechanisms can be studied (2) models should be quantitatively fit behaviour, such that they capture as many statistical properties of the data as possible, as important cognitive mechanisms have been shown to only reveal themselves in detailed patterns of response variability (rather than averages); (3) model fitting should be automated, avoiding subjective, piecemeal iterative model construction, neural network mechanisms can be accelerated at scale.
Previous approaches have fallen short of these desiderata. For example, cognitive models including drift-diffusion (Resulaj et al.) 2009; Pardo- Vazquez et al. 2019), Bayesian (Heald et al. 2021). symbolic (Castro et al. 2025). and large language model-based models (Binz al. 2025
have

---

Published as a conference paper at ICLR 2026
been successfully fit to detailed behaviour with impressive predictive power, but their architectures remained too abstract and divorced from RNNs to be useful as hypotheses about neural network dynamics (thus failing Requirement 1). Neural population coding models have suggested causal connections between neural response properties and particular patterns in behaviour (Matthey et al. 2013 Schneegans and Bays, 2017, McMaster et al., 2022), but remained mute about the neural network dynamics that give rise to the hypothesised neural responses in the first place (failing Requirement 1), have been rarely fit to behaviour quantitatively (failing Requirement 2), and their construction was not automated (failing Requirement 3). RNN models either with hand-crafted (Edin et al., 2009 Bouchacourt and Buschman 2019). or with aatk-opt architectures and parameters (Mnnte et al., 2013 3 Stroud et al. 20233 have been highly successful at suggesting testable hypotheses about neural dynamies underlying cognitive performance, but they rarely capture behaviour beyond generically competent performance, let alone being quantitatively fit to detailed behavioural data (failing Requirements 2, and for hand-crafted networks also Requirement 3). In some cases, non-trivial aspects of behaviour were successfully captured by such models, but at the expense of including purpose-built design choices or ablations (Xie et al. 2023). thus making these models essentially hand-crafted in this sense (failing Requirement 3).
In this work, we introduce an approach that satisfies all three requirements. We extend automated neural mechanism discovery by directly training biologically plausible RNNs to exhibit behavioural suboptimali- ties seen in experimental data. We tackle two major challenges for doing so. First, RNNs are data hungry, and behavioural data is scarce. To make training feasible, we generate synthetic data using descrip- tive generative model, and use this to train our RNN (Figure []. The second challenge is defining suitable training objective for the RNN so that they can generate complex (for example, multimodal (Bays et al. 2009) or skewed (Fritsche et al. 2020) Resulaj et al., 2009)) continuous response distributions often found in experiments. Previously used methods for training RNNs are not appropriate for this. Task optimization of RNNs typically used simple optimisation criteria, such as mean squared error or cross entropy (Yang et al. 2019), which encourage unimodal, typically normal response distributions. Even when training RNNs explicitly to generate specific distributions of responses, moment- or score matching-bait based criteria were used (Echeveste et al. 2020; Chen et al. 2023), which also do not scale well to capturing complex distributions (and require arbitrary choices as to which moments or statistics of the target distribution are important). Instead, by framing elicitation of behavioural output generating samples in Euclidean space. train the RNN using a training criterion inspired by diffusion models (Ho et al. 2020), state-of-the-an class of generative neural models for complex, continuous distributions. We adapt the training procedure used for diffusion models to work within an RNN, harnessing their flexibility in sampling from complex distributions. This enables us to flexibly fit to continuous response data, unlike prior methods of fitting directly to behaviour (Ji-An et al. 2025), which are limited to categorical distributions (55).
Task
Task
Behaadiuural
Synthoric
Goneratee
Train
Task opemal
Reslittic outputs
culputs
Network
Subject
Flexble gencrative
modiel
Network
Figure I: Left: typical procedure involves training an RNN to perform optimally in task, without relating to real behaviour. Right: our novel method can replicate behavviur, by training surrogate training data.
We demonstrate the power of approach for automatic mechanism discovery on of the most studied cognitive paradigms whose neural underpinning is still poorly understood: visual working memory (VWM). While there are many neural circuit models of working memory, they typically do not address the more challenging (and ecologically more relevant) situation when multiple pieces of information ( ('items') need to be maintained in memory Zhang (1996). Even RNNs that do perform multi-item VWM tasks fail to capture the characteristic patterns of behaviour that are specific to the multi-item situation (Yang et al., 20199] Driscoll al. (222). or only capture some select behavioural biases by using hand-crafted networks with purpose-built architectural motifs with limited biological plausibility (Bouchacourt and Buschman] 2019). particular, so-called swap errors are a major source of errors that arise when participants recall the wrong item from their working memory, instead of the item that has been cued - and thereby create strongly multimodal response distributions (Bays) et al. 2009.. Current RNN models do not capture swap errors, and the resulting multimodal response distributions, at all. Conversely, there exist population coding models that do capture swap errors, but do not make predictions about neural circuit mechanisms; (Schneegans and Bays] (207)). We address this gap by showing that our approach produces neural circuits that generate rich, realistic response
Dechanisms,Schneegaan
2

---

Published as a conference paper at ICLR 2026
distributions by directly training networks to perform swap errors. We use trained networks to generate hypotheses about the biological implementation of observer models from the psychophysics literature (Schneegans and Bays (2016); McMMsstrr et al. (2022) $4) previously inaccessible to RNNs trained for task performance.
2
BACKGROUND
RNNs in neuroscience Neural models of biologically plausible cortical circuits are described by their continuous-time "leaky" dynamics:
- s,rr,,,,7:0 8,
(I)
where r' is the vector of neural firing rates (hidden RNN activations), is the input from an external population (e.g. sensory input into the population), 77 is the population's 's biological process noise, and A is the diagonal matrix of neurons' membrane time constants (typically constant across cells, i.c. = >T). The non-linear computation F, parameterised by 0,. models the recurrent communication between and information integration within neurons. The network output is generically æ = (r: 0o), often interpreted as a behavioural response. The exact form of F and G vary across studies; this work, we use dendritic tree architecture for each neuron (Appendix B Lyo and Savin (2024)), and interpret the output of the network- a simple linear projection of the firing rates as a colour estimate on the colour circle. More details are provided in $ and discussion of our choice of architecture is provided in Appendix B namely, we chose this architecture to be sufficiently expressive for the task posed, while also maintaining biologically plausible connectivity structure.
Previous attempts to train RNNs that generate samples typically train RNNs on task-optimality or matthntg criteria (Echeveste et al. 2020). This is insufficient for as outlined in $ we are interested in reproducing behaviour that takes on a distinctly multimodal distribution. To do so, we employ methods from the diffusion model literature. We briefly introduce diffusion models here, but defer many details to previous foundational work (Ho et al. 2020).
Denoising diffusion probabilistic models (DDPMs) DDPMs are class of generative models that involve neural network learning to undo fixed noising process applied to data samples. During training. some Dr drawn from data, and is iteratively corrupted:
q<-+~|=,,) Na(x; v1 1,I,,,,1)
(2)
which admits a closed form posterior q ( [0,,, (T), with mean (3,, æT). The DDPM describes xy, T). and is trained by maximising a hicrarchical evidence lower bound (ELBO) of the one-step denoising of the corrupted data. By parameterising the transition kernel by its mean A (x+ 7) only, then adding isotropic Gaussian noise with the same variance as the noising posterior, this ELBO can be expressed as a sequence of square distances, rather than KL divergences. We equate timesteps within the trial to timesteps the DDPM denoising process, and use this objective to train the RNN to generate realistic samples from the response distribution. Training RNN-like diffusion models to perform cognitive tasks further requires a form of teacher forcing (Appendix D)
non-stationary transition kernel (x (
Synthetic data Neural networks are data-hungry. This is not an obstacle training for task-optimality task variables such as stimulus features and target responses be generated procedurally but it is when training directly on behaviour producing human or animal behavioural datasets is prohibitively costly the scales required for training networks. Therefore, we generate synthetic behavioural data to plug this gap. Many of the models listed in T which accurately describe behaviour using neural population codes (Schneegans and Bays 2017p. Bayesian cognitive (Heald al. (202I) and descriptive (Bruijns et al.) (2023) models, and/or large language models (Binz et al. 2025) can, by the same token, be used to faithfully generate synthetic behavioural data. These synthetie datasets can be used to train our dynamical neural model; sufficient flexibility in the generative model allows synthetic data that captures desired statistical facets and suboptimalities of behaviour. Many of these models, or their relatives, have also been used to fit behavioural data from working memory (WM) tasks, with varying degrees of satisfaction of our three requirements. For this reason, and because of the sophisticated response distributions elicited from even simple WM tasks, we choose this as the testbed for our method. In $ we will return to potential extensions to other aspects of WM, and other cognitive tasks in general.
3

---

Published as a conference paper at ICLR 2026
Working memory (WM) and swap errors While not consistently distinguished from short-term memory (Aben et al. 2012). WM is generally defined as the temporary maintenance, manipulation, and retrieval of sensory information over an order of seconds, for executive functions. A dominant experimental paradigm in investigating visual working memory (VWM) is the delayed estimation task, in which human or animal subject is presented with sensory information and asked, after short delay during which the stimulus is withdrawn, to reproduce some information about it. Figure shows an example of a multiitem delay estimation task, where the subject is presented with multiple (here, 2) items, then presented with probe dimension (here, location) feature value, called a cue, and asked to reproduce the corresponding report dimension (here, colour) feature value of the cued item. From here, we will interchangeably use 'location" and 'probe feature', and 'colour' and 'report feature' for ease of language, however the features used for each role open for experimental design. Most responses fall near the cued item's 's report feature value, and the distribution of estimation errors along the report feature (here, the colour circle) provides key evidence for measuring the capacity of VWM (Ma et al. 2014).
T
However, there is also a central tendency in responses towards the colour of uncued items, a.k.a. distractors (Bays et al.) 2009] Gorgoraptis et al.] 2011). This indicates the presence of swap errors, in which sub- jects mistakenly recall the colour of a distractor. Be- havioural and neural evidence largely implicates mis- selection of the correct item from memory at cueing time, rather than misbinding of feature values dur-
In-trial time
+
+
Pre-stim.
Stim.
Delay
Cue
Delay 2 idenatsing happens here!
Figure 2:
Two-item delayed estimation task minimal VWM task for swap errors.
ing memory encoding or storage swap errors are failure to extract the correct item from memory. rather than a misencoding into memory (but see |Radmard et al. (2025)) Macaque neural signatures reflects this misselection hypothesis Alleman et al. (2024). The key evidence for this is that swap errors are more likely to happen when a distractor's 's probe feature value is close in feature space (here, physical space) to the cued one (Emrich and Ferber| 2012: [Schneegans and Bays 2017 McMaster ct al. 2022). Describing the expected distribution of responses, including swap errors, requires a multimodal distribution, with modes placed over both target and uncued colours.
Bayesian Non-parametric model of Swap errors (BNS) We use BNS (Radmard et al. 2025) to generate synthetic multiitem delayed estimation data. This model predicts the likelihood of swap eπror as a function of the distance of the distractor from the cued stimulus in each feature dimension (Appendix A). In the present work, we do not directly fit BNS to a real behavioural dataset. Instead, we opt to hand-tune the generative model to capture archetypal dependencies of swap errors. Specifically, we train RNNs with synthetic data generated by instantiations of BNS that either i) predicts swap errors, ii) predicts the same frequency of swap errors for any distractor, and iii) predicts swap errors more frequently when its probe feature is more similar to that of the cued item, as per the extensive experimental evidence. In $3] we show that RNN representations resemble the real macaque neural geometry only in the third case, when maximal realism is captured in the training data, even in non-swap trials. We also fit BNS the behavioural output of the trained RNNs to quantify its success in reproducing the target dependence on probe difference ($4). Again, we emphasise that the use of VWM and BNS is one instantiation of our method, and that its ethos of fitting directly to suboptimal behaviour extends beyond this domain, which we discuss briefly in (5)
3 METHODS
Network architecture Our RNN models a population of n densely connected cortical pyramidal neuron Each neuron integrates inputs from sensory population and all other neurons in the population via a dendritic tree (Appendix B Lyo and Savin (2024)). Overall, JF in our network dynamics (equation I is the somatic integration of the final, most proximal, layer of dendritic nodes with additive Gaussian noise (Appendices G We train the RNN to perform the multittem delayed estimation task ($2) with 2 items, providing the minimum viable task in which swap errors could arise. The network is first provided with two stimuli, each parameterised by two features, dubbed probe and report (visualised as location and colour), both on circle. For index-cued networks, only the stimuli colours are provided as individual ordered scalars in -,, and the cue as an
Conntaants used in defining the architecture/traniigg processes are declared in Appendix
4

---

Published as a conference paper at ICLR 2026
A
RNN responses Ground truth
:
D
Swap probability
01
0.
E
10
13
2.0
25
301
Distractor probe distance
B
Angullr error colour recall
Time after offset
Figure 3: DDPM-style trained RNNs accurately captures swap errors in training data, unlike ablated task-optimal networks. A dotted lines the target swap used generate the synthetic training data for various RNNs, and solid lines are average swap rates inferred from fitting BNS to RNN- generated behaviour after they are trained on this synthetic data (Appendix [A] Radmard et al. (2025)). RNNs achieve fair success in replicating training data for most trials in both probe distance-dependent and -independent cases. Corresponding sample sets in IR2 are shown in C, where border colours match the line colours here. B Trajectory snapshots at different points in the second delay period, during which the network is trained to denoise behaviour. Trials are coloured by their argument at the end of denoising, which is interpreted as the colour estimation made by the network. Final samples successfully sample from their multimodal target distribution (Appendix (A). C Typical set of final timestep samples for trials with close (left) and far (right) probe feature values. In all cases the red stimulus is cued. Borders indicate generative RNNs, referencing lines in A. Top row (grey border) indicates task-optimal network. D,E Two typical ablations that may be applied to task optimal networks to indcue swap errors: decreasing probe distance beyond the minimum margin between items seen during training induce confusion, and increasing process noise variance beyond training conditions to increase misselection likelihood. Swap errors would be evidenced by a second mode at the colour the distractor, indicated by a blue marker. This does not arise in either case, and further attempts would require subjective ablation design. E inset: Conversely, removing process noise for network trained to perform swap errors recapitulates optimal behaviour.
index. This provides direct analogy to the tasks performed by macaques (Panichello and Buschman 202T Alleman et al.) 2024), which we make comparison ($4). Locations are not provided to these networks, and colours are ordered. For ceature-cued networks, the stimuli are provided to the network as the conjunctions of locations and colours. These are fed to the network as activations of palimpsest conjunctively tuned sensory neural population (Appendix B Matthey et al. (2015); Schneegans and Bays (2017). Importantly, in the latter case, items are order invariant the network must store the two conjunctions between the two features. This allows us extend predictions about the representation of a variable probe feature. The network's firing rates vector is projected the 2D plane with fixed, orthonormal output projection matrix æ G(r) WrrrW. E 2g2x= (i.e. WIW, I..) Network activity is therefore split into two fixed orthogonal subspaces: the behavioural subspace (the network output) æ Wrr, and the behavioural nullspace m W We interpret the argument of this output at the final timestep of the trial, arctan an ([wrr)) (Wrrr]2) as the colour recalled by the network on that trial. To train for task performance, mean squared error (MSE) loss on the discrepancy between this 2D estimate and the report feature of the cued item embedded on the circle in IR² suffices (Appendices A D).
Training Instead of training for task optimality, adapt the training procedure used for generative diffusion models ($2) Ho et al. (0020; Lyo and Savin (2024) to train for behavioural realism, accounting for the presence of swap errors. We apply the DDPM training objective only to æ, not the full activity vectofr r, equating timesteps after cue offset to timesteps in the DDPM denoising process. Besides the of leaky dynamics (see above), this is key difference to previous integrations of
5

---

Published as a conference paper at ICLR 2026
diffusion models with RNNs in cortical modeling (Lyo and Savin 2024). Additionally, the DDPM-
style criterion is only applied to the T timesteps following cue offset (Figures 2 BB) - not repeatedly
throughout the full duration of the trial - and the target distribution is trial- specihc.
For each task trial, the target distribution of responses is defined by a mixture of 2 Gaussians in IR? with mode means determined by colours of stimuli (on the circle), and weights determined by BNS prediction (Appendix D). This final behavioural sample is then iteratively noised (equation 2 to produce the target trajectory for the second delay. The network is trained undo this process each timestep of the second delay, with the transition kernel mean provided by discretising the continuous projected to the behavioural subspace:
noiseless RNN dynamics (equation I
dt A F(r,, 8x: (o, (re. t) = W (
df r
(3)
Note that because diffusion only occurs during the second delay, the sensory term St is empty (Appendix [B). and information about the full set of stimuli, and hence the target distribution, is contained in the behavioural nullspace activity m, This dependence on previous timesteps means training requires simulation of the full trial trajectory in sequence. Because valid application of the DDPM criterion ($2) requires training examples of noised data (æ,) ) to be drawn from the marginal distribution of the noising process, i.c. (( (xT). we employ teacher-forcing during training to ensure that the distribution of trajectories of æ, stay in distribution. The training algorithm is summarised in with full details provided in Appendix [
algorithm !
Algorithm 1 Summary of training with teacher-forcing full algorithm in Appendix D
repeat
Generate stimulus set and select cued item
Using BNS, generate a sample x from the desired behavioural distribution in IR?, and apply the noising process (equation 2 to generate noised data 00.r-
Initialise network activity from (0,
Run network dynamics for the prestimulus, stimulus, delay, and cue periods (Figure 2 with noisy discretised leaky dynamics (cquation 3 without projection), finally arriving at network activity at cue offset To
for = 0 to T (i.e. during second delay) do with x (Appendix D
Apply teacher-forcin
Calculate mean transition kernel ( 1) (equation 3
Run dynamics one step (equation and add scheduled noise (Appendix G
end for
Train on regularised DDPM criterion (Appendix D
until converged
4 RESULTS
Swap dependence reproduction We start by summarising the success in reproducing a variety of desired behaviours, characterised by the target swap function. Figure 3A shows the ability of the RNNs to replicate the dependence of swap errors on distractor distance used to generate their synthetic training data, for a variety forms of dependence (Appendix(A) (A). Here, dotted lines show the ground truth swapping rates as function of probe distance, and solid lines show feature-cued networks' abilities to capture them. This is inferred by fitting BNS the estimates generated by the trained RNN. Disparities in low swap probability regions is likely due to the compounded lack of training samples for both the RNN and BNS in low probability modes. Example sample sets, coloured by the associated estimate, are shown for some of these in Fig with the multimodality achievable with our method contrasting the unimodal distribution of estimates produced by task task-optimised networks (top row). Figure BP and E summarise naïve attempts to induce this multimodality after training for task-optimality, which one might expect researchers to attempt. Failure to do so would typically set off an iterative and targetted sequence of ablations based on heuristics/prior assumptions on the origin of swap errors. We achieve swap errors by directly training for them, and find that the resulting neural signatures underlying these errors match those found in biology, as we show below.
Neural representations of memoranda We now compare the neural representations of stimuli during each delay to those of macacque lateral prefrontal cortex (LPFC) during a similar task

---

Published as a conference paper at ICLR 2026
A
B
Pre-cue
upper
25
PC11
P1
C Loragio tho
Locsan-ctbor
Loodton
Cosine anglee between planes
2
Selected stimulus 1
Selected stimulus
0.8
Index cued swap prob.
0.6
0.4
Task optimised
8C3
PC
0.0
2
PCC
Time relative to cue onset (s)
2C
D
Cosine of angle between pianes
0.6-
0.4
0.2
cue
Figure 4: Behaviourally realistic, not task optimal, networks capture neural signatures of VWM. A my averaged over many accurate trials with distractor report feature varying, for trials with close (top) and far (bottom) stimulus distance, before (left, planes misaligned) and after (right, planes aligned) exposure. Colours represent report feature of cued (and accurately recalled) stimulus. Representations drawn from network E (right), which best matches neural data (see below). B This qualitatively matches equivalent representations in macaque cortex in this case there are only two possible cue values. We find a similar pre- and post-cue geometry for our index-cued networks, which are the closest analogue to the real task. C The real data (dots) further shows that this planar alignment (quantified cosine similarity between normals) increases before and during cue exposure (grey stripe), and remains high throughout the second delay until time of recall. This is matched for our index cued networks when trained with data-matched proportion of swap errors (5- 10%). Even training with our DDPM-base bbsed method with no-swap target distribution does not achieve this rapid increase during the period. A longer second delay period chosen to encompass all experimental conditions in the original study. Plane alignment over time for the feature-cued networks. Axes borders indicate the model used, in Figure 3 Importantly, only the model with distance-dependent swap probability matches the description of the biological data like its index-cued counterparts. This is quantified in Appendix C Error bars show mean + std across all different spatial configurations with the denoted stimulus location distance (see inset; purple is closest distractor, yellow is furthest). C generated with data from Panichello and Buschman (2021)
(Panichello and Buschman, 2021p. Here, two macacque monkeys performed a version of the multiitem delayed estimation task with fixed stimulus locations one coloured square in the top and one in the bottom half their field of vision. A binary visual cue at the point fixation indicated which item is to be recalled. This task is most analogous to the index-cued version of our networks. As such, following subsections will begin by reproducing biological phenomena with these networks, before going to make predictions about the role of the probe feature with our feature-cued networks. For the former, we match behavioural statistics; for the latter, we construct archetypal swap dependencies the target distribution, form predictions about neural representations.
7

---

Published as a conference paper at ICLR 2026
Their dataset of two monkeys for this retrospectively cued task (4.769/3.943 completed retro trials over 13/10 sessions) revealed: i) a swap error rate of roughly 10% when pooling the data for the two monkeys, and ii) a consistent low dimensional representation of the stimuli in memory before and after the onset of the cue during accurate trials (Figure 3). At both points in accurate trials (i.e. no swap error, and low angular error), the mean population response of LPFC associated with the colour of each item (upper and lower) was found by averaging neural activity across trials with that colour in that position, factoring out the colour of the distractor. These representations formed a 2D cycle when averaged across trials, reflecting the cyclicality of the behavioural responses provided by the monkey on the colour eircle. Prior to cue onset, the representation of each item occupied perpendicular 2D planes (Figure [left]) and after cue onset, these planes reoriented to become parallel (Figure B [right]), with the cosine of the angle between the planes increasing to a maximum during the cue remaining high (i.e., parallel) during the second delay (Figure CC).
We compare these findings to representations in the behavioural nullspace activity m of our RNN, for index- and feature- cued networks, well as task-optimised networks. For index-cued networks (Figure (C), which most closely analogise the real task, optimisation for task performance does not capture the full alignment of the stimulus representations during and after cue presentation. The sharp and sustained increase and alignment is only captured when using our DDPM-style training and including swap errors in the trial-wise target distributions.
In the feature-cued case, where a continuous location is used to load and cue items from memory, this pattern is only captured when the network is trained with a pection-dependent swapping rate (Figure rate of swap errors. As motivated above,
Ap [right]), and not with a flat, probe crature-independee while we matched behavioural statistics for the index-cued task, which bears direct resemblance to the real task provided to monkeys. we modelled the form of the swap dependence off more sophisticated tasks (Emrich and Ferber 2012] Schnnegann and Bays/ 2017 McMaster et al. (022) on are in this task in order to make predictions about their neural representations. Furthermore, this network bears higher pre-cue representational alignment for distractors with closer probe feature values. This provides a prediction for the neural substrate of delay-time processes leading to memory retricval error at cucing, previously described in the psychophysics literature as noising or drift in the abstract representations of stimuli along the probe dimension over time (Schneegans and Bays 2017] McMaster et al.) 2022). Concretely, we predict that trial-by-trial decomposed analysis of the alignment between stimulus representations will reveal that swap errors occur more often when pre-cue representational alignment of colour is higher. This could be due to probe similarity, as in the case presented, or other factors such as attention before cueing in the fixed item location case.
The previous subsection made the prediction that physically closer, hence more swap-prone, distractors cause higher report feature representation alignment before cue onset (Figure Ap [right]). We now consider the neural signatures of the retrieval errors at cueing time, which lead swap errors. For the same macaque task, experimentalists later showed that item misselection is the primary cause of swap errors (Figure 5) Alleman et al.] (024)). The misselection is operationalised as follows: (i) before cue onset, items are stored as conjunctions between their probe (location or just upper and lower in the experimental case) and report feature (colour) values (Figure 5A [left]); after cue onset, items are stored as conjunctions between their role (target vs. distractor) and report feature value (Figure [right)) These are dubbed the spetial and report subspaces respectively; item "misselection' is misprojection between these subspaces at cueing time. Swap errors can also arise at cucing time due to misinterpretattion of the cue variable. Alternatively, report and probe features can be mishound prior to cueing. Alleman et al. (2024) find the least evidence for this mechanism, instead asserting that swap errors arise during or after cueing time. We now complement existing evidence of this neural basis of swap errors, which depends on across-trial statistical modelling of neural responses Alleman et al. 22024), with the representational geometry learned by our network.
Item misselection
First, we apply the same linear decoding analysis as Alleman et al. (024)) to the memory subspace neural activity of our index cued network. These networks hold direct analogy to the task analysed by this work, so results can be directly compared. This method is based on training linear decoders on non-swap trials, then projecting neural activity onto the line connecting these nominal representations to representations that arise due to misbinding, cue misinterpretation, or item misselection. We defer details of this analysis to Alleeann et al. (2024). As with the real macaque data, this analysis suggests that swap errors arise in our network during the second delay, rather than misbound representations of memoranda Schneegans and Bays (017)); McMaster et al. (2022).
8

---

Published as a conference paper at ICLR 2026
Next, we investigate our probe-cued networks for the existence of report subspace, where colours bound to roles (target or distractor) with locations factored out, and the neural signature of swap errors during delay 2. Figure [SC shows low-dimensional projection of my for fixed set of locations and one fixed colour forming two parallel rings, coloured by the free stimulus. In the left column, when the fixed colour stimulus is cued, the bottom ring (negative PC3) is made up of the representation of swap trials, which are coloured by the uncued but mistakenly recalled colour, and the top ring (positive PC3) is coloured by the distractor colour when the fixed colour item accurately recalled. Vice versa when the fixed colour stimulus is uncued (right column). The normal to these rings, which discriminates accurate and swap trials F [bottom]), is independent of the stimulus locations, only depending on the fixed stimulus colour considered (Figure 5 [top]). This makes this two ring structure a candidate for the geometry of the posteue report subspace (Alleman et al., 2024), which binds stimulus colour (around the ring) to its recalled role (choice of ring), and is independent of the original colour-locatoo binding Figurr[5A)) We have not considered exactly how this report subspace varies across different fixed report stimuli, just that the normals of these rings is shared whenever one of the items has a fixed colour, regardless of the item locations or whether the fixed item was cued (Figure [5c [left)). However, their alignment depends on colour similarity (Appendix suggesting an intuitive toroidal formation.
5 DISCUSSION
In this work, we trained RNNs to exhibit swap errors during a multi-item delayed estimation VWM task (33) Its success was contrasted with naïve training methods and ablations applied to task-optimal networks (Figure 3). Capturing swap errors post hoc would require a series of directed network modifications, likely driven by leading assumptions or heuristics as to how swap errors arise, whereas our method displays them directly. Then, we showed that this training method captured a more accurate description of neural mechanisms when trained to adhere to the empirical variability around real swap errors (Figure(4). Namely, we trained one network to produce swap errors more frequently when there is higher probe similarity between the cued item and the distractor (Figure 4 [left]), previously theorised to increase the likelihood of the item misselection underlying swap errors. Separating trials by probe distance, we provided a prediction for the neural representation of probe similarity before the cue. Finally, we provided a candidate representational geometry for swap errors, in line with previous work on their neural basis.
Our approach bridges the gap between flexibly modelling neural circuits and producing rich descrip- tions of behaviour, that were previously only achieved separately by task-optimised networks and abstract cognitive models, respectively. We did so by reversing the typical sequence of normative modelling- by directly training networks to reproduce the full distribution of behavioural responses, including suboptimalities, we bypass the need for an iterative, potentially heuristic set of ablations. These principles are most tightly shared by tiny RNNs (Ji-An et al. 2025). However these small networks, like previous approaches fitting to behavioural data (Xue et al. 2024) have only been used for discrete behaviour, e.g. categorical choices, for which simple maximum likelihood fitting to behavioural datasets suffices. Our method provides a novel, principled approach to holistically fitting to all, continuous dimensions of data. The neural predictions we made by reverse engineering behaviour in this way await experimental verification, but demonstrate the value of training networks to reproduce behavioural errors rather than optimising for performance.
One limitation of our approach is its dependence on a generative behavioural model produce training data, which may not be available for all tasks of interest. However, for many well-studied paradigms such models already exist, and recent advances in descriptive modelling (Binz al. 2025) suggest that capturing the relevant statistical structure of behaviour, rather than normative or mechanistic correctness, is often sufficient and increasingly achievable.
The present work can be extended in multiple directions. We used this VWM task as a testbed for our novel methodology because of its inherent multimodality, but our method can be extended to any other cognitive task associated with complex behavioural suboptimalities, or other non-trivial response distributions, such as reaction times (Pardo- Vazquez et al. 2019). For example, we report our success in training a model to perform cue combination (multisensory integration; [Beierholm et al. (207)) task, where the network is trained to generate samples from the posterior of a Bayesian ideal observer. This task currently acts demonstration of the flexibility of our technique, rather than an extension to the mechanistic hypotheses it generates. Future work can extend in this direction.

---

Published as a conference paper at ICLR 2026
A
B
Normalized Projection
(a-nominal, 1
1.0
Correct trials Misainding Misselection Mis nttepprrrta
0.8
nominal trajectory
bottom calour
distrator colour
0.6
missolection from memory
misinterppetation oue
misbound
repc. in
0.0
-0.5
2.0 relative to onset (s)
2.0
top colour
target colour
D
alignment across probes when reports are same
Sharco
Shared
Empricall density
Nui distribution 8
fixed reports are apposites
20
30
-100
Starred
110
Shared
PC1
50
100
-20
Shared
PC1
100
-20
Correct tries final activity
Swapped trias final activty PC3
Shated
PC3
Stared
PC3
Empiiiiaal dendit
1000
Sharad
100
10
Sharod
100
Shared
60
100
Value along relevant PC3
Figure RNNs trained with swap errors provide prediction about the geometry of item misselection. A prior works (Alleman al. 2024) suggest 3 causes of swap errors. Only misbinding errors arise during (or before) delay 1. prior to cucing. Misselection and misinterpretation cause incorrect transfer from a probe-report binding in delay to report binding in delay 2. B We repeat linear decoder analyses from Alleman et al. (2024) for our index-cued networks. Like the real data, these provide most evidence for swap errors arising during delay ruling out misbinding errors. We defer full details to Alleman et al. (2024) C Triial-byttiall my, projected to a shared set of principal components (PCs). One item's s colour is fixed to purple, and the other item's 's colour is varied and used to colour the scatterplots. Left (right) axes show when the purple item is cued (a distractor), so colours indicate the uncued (target) colour. Probe features are fixed (bottom) row, causing more (fewer) swap errors. Neural activities form two rings (with gap in each due to the minimum margin between colours) see main text. To illustrate this better, we used a network that swaps more often than seen for macaques (highest orange line in Figure 3A). ). D Left: there is high pairwise alignment in the third PC across different stimulus spatial configurations (i.e. fixed locations) if the fixed colour used to generate the two ring geometry is the same between these configurations less so the fixed colour is different (here, the opposite) for all pairs spatial configurations. See Appendix for a more graded pairwise comparison. This also applies to the first two PCs, unshown. Right: this third PC discriminates swapped versus correct trials in cases where the fixed colour item (purple here) cued, across many different probe values.
close (far) spatial configuration in the top
Within swap error analysis. our analysis of the representational geometry of stimuli and responses (Figures 4 and 5) mostly considers stationary snapshots of the population activity. Besides this, Figures 4 4D,10 E aggregate information across many trials. Future work should include trial-by-trial dynamical analysis. For example, the misselection process described in Figure (5A may be due to a failure in information loading (Stroud et al.] 2023) during the cue period, which is made more likely due to the higher representational alignment we uncovered here. Better interpreting why RNNs capture neural phenomena only when trained explicitly to exhibit errors, beyond the descriptive summary we have provided here, can help refine this inquiry. Looking forward. there exists a much richer set of behavioural phenomena associated with VWM which can be modelled with our method. Further dependence of swap errors on set size (Bays et al. 20099 Emrich and Ferber 2012) and presentation order (van Ede et al. 2021 Gorgoraptis et al. 2011). and other, unimodal, response distributions which cannot be captured by moment matthing (Fritsche et al., (020) are all suitable next paradigms to benefit from our method.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Bart Aben, Sven Stapert, and Arjan Blokland. About the distinction between working memory and short-term memory. Frontiers in Psychology, 3, 2012. ISSN 1664-1078 doi: 10.3389ffssg.2012
00301. URL https: http://wwwwww.tteie.eeeee http://www..oooooooooooooraaaaa/
00301
Matteo Alleman, Matthew Panichello, Timothy J. Buschman, and W. Jeffrey Johnston. The neural
basis of swap errors working memory. Proceedings of the National Academy of Sciences, 121 (33), August 2024. ISSN 1091-6490. doi: 10.1073/pnas.24000 2401032111,
URL | /xx doi
org 10130....31p11
2401032121
Bays, R. F. G. Catalao, and M. Husain. The precision of visual working memory set by allocation of a shared resource. Journal of Vision, 9(10):7-7, September 2009. doi: 10.116779.10.7. URL
016///0...111/11 10.7
Ulrik Beierholm, Ladan Shams, Wei J., and Konrad Koerding. Comparing bayesian models for multisensory combination without mandatory integration. In J. Platt, D. Koller, Y. Singer. and S. Roweis, editors, Advances in Neural Information Processing Systems, volume 20. Curran Asso- ciates, Inc., 2007. URL Ftt. 2007 e 2atee//e000offfffffaaaaa22eeeee2
www.inng.....uueeeee cttesccc/peee ww/paper.ppep..... fi Feleppappr/ pape
Marcel Binz, Elif Akata, Matthias Bethge, Franziska Brändle, Fred Callaway, Julian Coda-Forno, Peter Dayan, Can Demircan, Maria Eckstein, Noémi Ééltetö, Thomas L. Griffiths, Susanne Haridi, Akshay K. Jagadish, Li Ji-An, Alexander Kipnis, Sreejan Kumar, Tobias Ludwig, Marvin Mathony, Marcelo Mattar, Alireza Modirshanechi, Surabhi Nath, Joshua C. Peterson, Milena Rmus, Evan M. Russek, Tankred Saanum, Johannes A. Schubert, Luca M. Schulze Buschoff, Nishad Singhi, Xin Sui, Mirko Thalmann, Fabian J. Theis, Vuong Truong, Vishaal Udandarao, Konstantinos Voudouris, Robert Wilson, Kristin Witte, Shuchen Wu, Dirk U. Wulff, Huadong Xiong, and Eric Schulz. A foundation model to predict and capture human cognition. Nature, 644480789:1002-11099, July 2025. ISSN 1476-4687. http://ax.doi. org/1 1tt.://...../...1.1///1111111 0100010000000111110000000101-330
doi: 10. 10110///415566005--022211
URL
025-0225-----0-
Flora Bouchacourt and Timothy J. Buschman. A flexible model of working memory. Neuron, 103 (1):147-160.e8, July 2019. doi: htt.://.neroo.20199903
URL htttps:://oi.org//////to
1016/5.
020
Sebastian A. Bruijns, Kcénia Bougrova, Inês C. Laranjeira, Petrina Y. P. Lau, Guido T. Meijer, Nathaniel J. Miska, Jean-Paul Noel, Alejandro Pan- Vazquez. Noam Roth, Karolina Z. Socha, Anne E. Urai, and Peter Dayan. Dissecting the complexities of learning with infinite hidden markov models. December 2023. doi: 10. 10.10012023.12.22.53300 URLhttp:: //tp:///xxx.. doi
org/10.
Pablo Samuel Castro, Nenad Tomasev, Ankit Anand, Navodita Sharma, Rishika Mohanta, Aparna Dev, Kuba Perlin, Siddhant Jain, Kyle Levin, Noemi Elteto, Will Dabney, Alexander Novikov, Glenn C Turner, Maria Eckstein, Nathaniel D. Daw, Kevin Miller, and Kim Stachenfeld. Discovering symbolic cognitive models from human and animal behavior. In Forty-second International Conference on Machine Learning, 2025. URL pettps://openrevieee wetpe/ewwwweeeeeeeeeeeeeeeeeeeee nwt. forum ?id= @hRXGMJ027
Shirui Chen, Linxing Preston Jiang, Rajesh P. N. Rao, and Eric Shea- Brown. Expressive probabilistic sampling in recurrent neural networks. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2023. Curran Associates Inc.
Christopher J. Cueva, A. Ardalan, Misha Tsodyks, and Ning Qian. Recurrent neural network models for working memory of continuous variables: activity manifolds, connectivity patterns, and dynamic codes. ArXiv, abs/2111.01275, 2021. URL htR.h:ttt::///ee .emantssemmmt..mmmm wtw.emmnnnicchho..oo lcschola
org/CozpusID: 240419711
Laura Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. August 2022. doi: 10.1101/2022.00.5555333000 URL fttps://dee www.ddi.org//// www.orggo..... 1101 000022222222 08. 18...5.5500 503870

---

doi: (.7755//clifc.55388 Published as a conference paper at ICLR 2026
Rodrigo Echeveste, Laurence Aitchison, Guillaume Hennequin, and Máté Lengyel. Cortical-like dynamics in recurrent circuits optimized for sampling-bais based probabilistic inference. Nature Neuroscience, 1339))113---1111 119,, August 2020. doi:
URL https
//doi.org/ 1038/ 44883/4419933 020 0670-1
Fredrik Edin, Torkel Klingberg, Pär Johansson, Fiona McNab, Jesper Tegnér, and Albert Compte. Mechanism for top-down control of working memory capacity. Proceedings of the National Academy of Sciences, 106(16)6800-680)) April 2009. doi: (007://pas.0001894400 URL
https: doi..ii.... org 10w. 0901894106
S. M. Emrich and S. Ferber. Competition increases binding errors in visual working memory. Journal of Vision 12(4): April 2012. ISSN 1534-7362. doi: 10.1167/12.4.12. URL
Matthias Fritsche, Felke Spaak, and Floris de Lange. A bayesian and efficient observer model explains concurrent attractive and repulsive history biases visual perception. eLife, 9, June 2020.
7554/e115e 55389
N. Gorgoraptis, R. F. G. Catalao, P. M. Bays, and M. Husain. Dynamic updating of working memory resources for visual objects. Journal of Neuroscience, 31(23):8502 851 June 2011. doi: 10. 10.1523/jneurosc..o222 1t1.000008112.00.11 11.2011. URL http:ps//ooo...rg/11113333333 wtw.:/oo////eeocc
0208 2011
James B. Heald, Máté Lengyel, and Daniel M. Wolpert. Contextual inference underlies the learning of sensorimotor repertoires. Nature, 600(7889):489 493, November 2021. ISSN 1476-4687. doi: 00.0088/41186-002-0441233 541586- 021-6-022-4112--2
URL hctp: /tt:://dxx.o doi wwt.:ooi..ggggi.... wtt.ooo.g.......................
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neu- ral Information Processing Systems, volume 33, pages 6840-6851. Curran Associates, Inc., 2020. URL httpp:://proceeddiii
nwwrips.r.p..ee ccpppppr.rppppppp.ppppppp cimeer_ffleeee files//aperr2220/ htlee4Accccceeeeaaafoo6661111112
Li Ji-An, Mareus Benna, and Marcelo G. Mattar, Discovering cognitive strategies with tiny recur- rent neural networks. Nature, 644480778,9933-100,,,, 2025. ISSN 1476-4687. doi: 10.1038/ s41586-025 1111866025-09422-1 09142-4. URL/htt:://k...ol.rrg///..
Benjamin S. H. Lyo and Cristina Savin. Complex priors and flexible inference in recurrent circuits with dendritic nonlinearities. In The Twelfih International Conferencee on Learning Representations, 2024. URL /ttttp:://ppereviiie
net forum?id- Saw.rummiii-SSSuuuaaaaaaaaa S5aUhpuyap
Wei Ji Ma. Masud Husain, and Paul M Bays. Changing concepts of working memory. Nature February 2014. ISSN 1546-1726. doi: 10.1038/mm.3655. URL
Neuroscience, 1733):377-356,
http:// doi www.oorg/...... 10t.//............. 3655
Valerio Mante, David Sussillo, Krishna V. Shenoy, and William T. Newsome, Context-dependent computation by recurrent dynamics in prefrontal cortex. Nature, 503(7474):78 84, November 2013. ISSN 1476-4687. doi: 10.1038/nature 12742. URL http::ttp//ate doi wttpo/o.ooor0/111033313333 nature12742
Loie Matthey, Paul Bays, and Peter Dayan. probabilistic palimpsesst model of visual short-term memory. PLOS Computational Biology, 11(1):e0044003, January 2015. ISSN 1553-7358, doi:
10.1371/journal.ppe www......pi....0....... 1004003. URL fottp://at doi http:/ppp/...aaaa.ooo......////0 pcbi
1004003
Jessica M.V. McMaster, Ivan Tomić, Sebastian Schneegans, and Paul M. Bays. Swap errors in visual working memory are fully explained by cue-feature variability. Cognitive Psychology, 137: 101493, September 2022, doi: 10. 10.111/6/..oom wwtp:/...oogsycc..222..000
URL
10. 70.1016/3. wow.://c.copsycc....m. 2022. 101493
Matthew F. Panichello and Timothy J. Buschman. Shared mechanisms underlie the control of working
memory and attention. Nature, 592(7855):601 605, March 2021. ISSN 1476-4687. doi: 10. 1038/
141586-001-033900
URL https///w....//tttt////dooooo Ut..ittee///xx..doooo////.
0339000-
12

---

Published as a conference paper at ICLR 2026
Jose L. Pardo- Vazquez, Juan R. Castineiras-de Saa, Mafalda Valente, Iris Damião, Tiago Costa, M. Inés Vicente, André G. Mendonça, Zachary F. Mainen, and Alfonso Renart. The mechanis- tic foundation of weber's law. Nature Neuroscience, 22(9):1493 1502. August 2019. ISSN 1546-1728 doi: 10. 10.10888411999---9940999999
URI. http:://kx doi www.l/.rg//.... /
Puria Radmard, Paul M. Bays, and Máté Lengyel. A flexible bayesian non-parametric mixture model reveals multiple dependencies of swap enrors in visual working memory, 2025. URL
https://arxiv. org. 2555 01178
Carl Edward Rasmussen and Christopher K. Williams. Gaussian Processes for Machine Learning. The MIT Press, November 2005. ISBN 9780262256834, doi: 10.755 / 1000000000000011.....11
URL httt:://ax.doi..rrg////..
001 00001
Arbora Resulaj, Roozbeh Kiani, Daniel M. Wolpert, and Michael N. Shadlen. Changes of mind in
decision- making. Nature, 461 266, August 2009. ISSN 1476-4687. doi: 10.
nature08275. http:p/p..a...........00///////.
Sebastian Schneegans and Paul Bays. Neural architecture for feature binding in visual working memory. The Journal of Neuroscience, 37(14)33913-3925, March 2017. ISSN 1529-2401. doi: (1523//eeurosc.3343--66621011110 URL httttt/p/ddd.oo org/ /00 1523/ JNEUROSCI 3493- 1773-16.2017
Sebastian Schneegans and Paul M. Bays. No fixed item limit in visuospatial working memory. Cortex, 83:181-192, October 2016. ISSN 0010-9452. doi: 10.1016/j..ortees cortex.20116.07.020 URL
http::/dx.. doi. otg. 101.00.111660000
2016.
Wayne W.M. Soo and Máté Lengyel. Training stochastic stabilized supralinear networks by dynamies- neutral growth. October 2022. doi: 10.1101/2022.1011.551222 URL httpps::://doioo
10. www.12111222... 10. 19. 51.22
Jake Stroud, Kei Watanabe, Takafumi Suzuki, Mark G. Stokes, and Máté Lengyel. Optimal information loading into working memory explains dynamic coding in the prefrontal cortex. Proceedings of the National Academy of Sciences, 120(48), November 2023. ISSN 1091-6490. doi: 10.1073/pnas. 2307990 120. URL htRphtp::// //dx doi. wwt.oi.org11.111.
.aiiaas.337999 2307991120
Freek van Ede, Jovana Deden, and Amna C. Nobre. Looking ahead in working memory to guide sequential behaviour. Current Biology, (1( June 2021. ISSN 0960-9822. doi: 10. 1016/j.cub 2021 1.0.066.. URL B
doi http::///wx...0.0001111111111110 00t.0//0.....///////.........//.
2021. 04 063
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, page 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860966
Christopher Versteeg, Jonathan D. McCart, Mitchell Ostrow, David M. Zoltowski, Clayton B. Wash- ington, Laura Driscoll, Olivier Codol, Jonathan A. Michaels, Scott W. Linderman, David Sussillo, and Chethan Pandarinath. Computation- wwwperioon-tthhhhhhhh through-dynamics benchmark: Simulated datasets and quality metrics for dynamical models of neural activity. February 2025. doi: 10. 10.1100/2025.0..07. 637062. / //x doi. org. org. 11ee11111111 0000001222222 02 07. 00777.63770 637062
Saurabh Vyas, Matthew Golub. David Sussillo, and Krishna V. Shenoy. Computation through
neural population dynamics. Annual Review of Neuroscience, 43(1):249-275, July 2020. ISSN 1545-4126 doi: 10. 1 new. 1146/ http:/anurev-nuuu--002525
URL httthttt:://ee doi. org/
Ronald J. Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural Computation, 1(2):270-280. June 1989. ISSN 1530-888X. doi: 10.1162/ neco. neco.1989.11 1989.1.2.270. UtLhttp::/
doi org 1162 new.. 1989
2.20
Yudi Xie, Yu Duan, Aohua Cheng, Pengcen Jiang, Christopher J. Cueva, and Guangyu Robert Yang. Natural constraints explain working memory capacity limitations in sensory-cogmitive models, March 2023. URLthtte htpp://xxx /xx. doi org 110100..1..1.22....
30.534982
13

---

Published as a conference paper at ICLR 2026
Cheng Xue, Sol K. Markman, Ruoyi Chen, Lily E. Kramer, and Marlene R. Cohen. Task interference as a neuronal basis for the cost of cognitive flexibility. March 2024. doi: 10.1101/202..3.00 583375. URLhttp://dae dx. doi 1ttp://0..11112222.11 03 04. 583375
Guangyu Robert Yang. Madhura Joglekar, H. Francis Song, William T. Newsome, and Xiao- Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature Neuroscience, 22(2):297 Jamuary 2019. doi: 00.10088/41193-018-0010000 URL
Kechen Zhang. Representation of spatial orientation by the intrinsic dynamics of the head-direction cell ensemble: a theory. Journal of Neuroscience, 16(6):2112-21266 1996.
14

---

Published as a conference paper at ICLR 2026
A BAYESIAN NON-PARAMETRIC MODEL OF SWAP ERRORS (BNS)
Stimulus array Z = [2(m) n= 1,2) is denoted by the vector probe (location) and report
(colour) feature values: (m)) (=)") 2 E S X S Item n" is cued after the first delay. The
probe-dependent BNS Radmard et al. (2025) generative process defines a mixture distribution over estimates (recalled colour, y), based on the circular distance of the distractor from the cued item (i.e. in the probe dimension). The so-called swap function f dictates the form of this dependence. For two items is:
x (n) = (()) (n')
2200) 1
f [~ps(s)
#(n) Teoorett
n n # n²
{f (n)
π Softmax(
3 Cat(T))
y 2((0)
E E~ p~pe(t0)
The associated graphical model is shown in Figure A
To generate synthetic data, we fix one target swap function f and perform the generative process up to component 3, before embedding the samples into 2² and adding noise there, rather than using circular emission distribution P.- We also generate both feature values with minimum margin of T/ /4, akin to previous multiitem tasks Schneegans and Bays 2017 McMaster et al. (2022):
Generate with minimum feature value margins
2. Generate with BNS, using target function f*
3. Embed sample and add noise 2 ~ . (AWI
(oo(t)) sin(y)
The samples of æ. forms the target distribution for the DDPM training procedure (Appendix D
W
z (n)
n
(n)
1/
0
m
Z. y
Figure 6: Graphical model of BNS, adapted from Radmard et al. (2025). contains parameters for emission distribution Pe- w contains GP prior parameters and Tecorestt
To fit BNS to the outputs the trained RNN, we interpret the colour estimate as the argument of the output at the final trial timestep:
[ Wrrth
y aretan
Werrh
and use variational inference to fit (f). which approximates the dataset posterior over swap function f. which is equipped with Gaussian process Rasmussen and Williams/ (2005) prior py. We defer details to Radmard et al. (2025). We do not consider the unitorm component used in the introductory work, given the rarity of RNN behavioural samples away from the target modes.
15

---

Published as a conference paper at ICLR 2026
In both cases, for the single distractor case, BNS admits tractable probability of swapping given probe distance:
Azp E + Eoourea
p(3 + n2 A0)))
n",
(5)
p(f)
p(f)
For the target swap function, this can be evaluated directly (p(J)= -(( M. forming the dotted lines in VigureeaA, For the fitted behaviour, this can be estimated using i.i.d. samples from the inferred approximate posterior (p(ff) = {()).
B
ARCHITECTURE AND TASK DETAILS
A
B Aa
+
Task
T p
T
d'
T
T
8t
p(æ,)
C
b.
Lt.
x
Leaky
dynamis
Teacher
x
forcing
Figure 7: Schematic of architecture and training. A (top) task with durations; (middle) by sized conjunctive tuning curve representations of sensory imputs, and representation of time along trial (constant before denoising starts); (bottom) denoising to a mixture of Gaussians target distribution. B dendritic tree with sensory and time inputs. Here, 8 in matrix form has report features (colours) for each row, and probe features (locations) for each column, hence why at cucing time, strip column provided, with no colour information. C teacher- teacher-forcing, applied to the behavioural subspace throughout the second delay period during training.
Our RNN models population of densely connected cortical pyramidal neurons Lyo and Savin (2024). Each neuron integrates inputs from all other neurons in the population via a dendritic tree. Distal tree nodes are provided with sensory information, and all nodes provided with representation of time along the trial, necessary for training it as diffusion model.
Dendritic tree nodes are organised into layers, with each node integrating activity from nodes one layer more distal from the soma than itself. Each presynaptic neuron's axons synapse onto all leaf nodes of this tree, i.c. the layer most distal to the postsynaptic neuron, including its own. This comprises a weighted sum followed by a ReLU non-lincarity (Figure 7 B):
h = Σ Eiai + (+(())
where è indexes over neurons, indexes a particular path of nodes in the dendritic tree, and enumerates distal nodes connecting The population's n neuron receivee inputs from dendritic tree with L layers, each with branching factor of B, from the previous layer. This results in each neuron synapsing onto I Bi dendritic leaf nodes. The rate vector is projected via one set of weights W, E Bn.B,... By X then used as inputs the most distal layer, instead of h. As
Undeclared numerical constants defined and provided values in Appendix G
16

---

Published as a conference paper at ICLR 2026
mentioned in the main text, this architecture choice was partly motivated by minimum requirements of expressivity. We started with standard architectures employing linear projections followed by point non-linearities (e.g.. Yang al. (2019),, but consistently struggled to achieve multimodal dynamics. A two-layer network F could produce multimodality, but this is biologically implausible as it implies intermediary neural population with instantaneous responses (zero time constant). We therefore adopted dendritic tree architectures (Lyo and Savin, 2024), which constrain deep network's weights to block-diagonal form, producing tree-like connectivity that reflects cortical pyramidal neuron morphology whilst enabling the required distributional flexibility. While a dendritic structure not necessary for our method, a proven (Lyo and Savin 2024), expressive structure that maintains plausibility in this regard.
Each node is also provided with a bias term. Before the cue offset (t. < 0), when the network still accumulating information, this is constant bias b. After cue offset, when the network is denoising during the WM second delay, this becomes a smoothly changing ny-dimensional representation of time Ho al. (2020); Vaswani et al. (0177) projected to scalar with a unique set of weights for each dendritic node.
To provide sensory input to the 'index-cued' network, we separate channels for stimuli and cue:
[>() 2/ OTJT
t in stim. presentation
[0, 0, T [0,
I in cue presentation, n²
t cue presentation,
where čn E Re are learned embeddings of each cue case.
To provide sensory input the 'feature-cued' network, we represent these stimuli with activations of palimpsest conjunctively tuned neural population (Figure 7 [second row]; Matthey et al. (2015); Schneegans and Bays 2017)) then combine the tuning curve information before projection the neural population using two sets of weights Wp. W,:
Σ exp [cos(2))
(2p (n)
={²) + cos(2!") =
I in stim. presentation
S, RR
Stij
exp cos(2p
(n") =())
t in cue presentation otherwise
0
(8)
= We,(W,,,, E IR²
(9)
where Ep, Σy are fixed and evenly spaced preferred stimulus values for the probe and feature di- mensions. In both cases, a projection W,s, E R BB.....B is provided to the most distal dendritic nodes.
The final recurrent output of the network is the somatic activity, with added noise:
F(re: 81, ny; 0,) = he + of11 1 g~~N N I,))
(10)
with 0, containing all the dendritic tree weights, sensory input parameters, time embedding projection
weights, etc.. hg = [ .... A is. the somatic activity. As with time representation b, of is held at constant maximum value before cue offset (a = of), then quenched according to the diffusion noise schedule's 's posterior variance (Appendix G.. The overall time-discretised network dynamics are therefore:
r [[1- dt -
di
F(.,) + ore I~,( X
(11)
Note that additive noise is applied to the rates, rather than to the membrane activity, meaning that rate values can fall below zero. Therefore, neural rates vector should be interpreted as deviations of the population firing rates from some baseline firing rate.
C SUPPLEMENTARY RESULTS FOR FEATURE-CUED NETWORKS
As mentioned in S while index-cued networks enjoyed a plane alignment trajectory that resembled the data with just the introduction of swap errors, probe-cued networks required higher swap errors
17

---

Published as a conference paper at ICLR 2026
Trsk optimal
Binwaon with swazs
Difucco with conclrt 8.05 pms
Diflusen smooth samp funzitin
1.0
Deta 8:
-05
3.8
6.5
2.8
15
20
-05
06
05
20
-2.5
8.6
0.
300
15
2.6
-35
3.1
65
2.0
15 24
Figure 8: Sigmoid and flat line fits to the alignment trajectories of different probe-cued networks.
for proximal distractors to do so. Here, we quantify this improvement in matching trends found in the data. We fitted two models to the cosine similarity trajectories data from each network variant: (1) four-parameter sigmoid function (slope, x-offset, y-offiet, y-scale) and (2) a flat line representing the temporal mean (Figure 8). We fit these to all cosine similarity data in Figure 4D, aggregating across location distance (colours of lines). This approach was motivated by the experimental methodology in Panichello and Buschman 2)22), where the authors fitted sigmoid functions the cosine similarity over time. We adopted this as quantitative benchmark for evaluating whether our networks produce neural dynamics that match the temporal profile observed in biological data. The comparison to flat line serves as a control. While the constant swap rate network produces temporally stable alignments that are fit well by a sigmoid with near-zero slope, this does not match the characteristic temporal dynamics reported in the experimental results. By quantifying the mean squared error difference between sigmoid and flat line fits, we can assess which networks exhibit the dynamic temporal evolution of representational alignment seen in the biological data. These results are reported in Table I in the order of subplots in Figure
Task optimal
0.009734
0.005479
0.004255
Diffusion with no swaps
0.052094
0.039117
0.012977
Diffusion with constant 0.05 swap prob
0.000152
0.000000
0.000152
Diffusion with smooth swap function
0.065647
0.000340
0.065307
Table 1: Quantitative comparison of plane alignment trajectories for different feature-cued networks
D TRAINING DETAILS
The network trained to produce swap errors are optimised on a joint loss function. The primary terms of this loss function are the stepwise denoising
(see below: [o et al. (2020)). with the transition kernel mean provided by discretising the continuous RNN dynamics (equation 3} projected to the behavioural subspace. Unlike traditional DDPMs, the transition kernel takes in dimensional rate vector, and outputs 2 dimensional denoised mean in the behavioural subspace. Without projection Wa. this dictates the full network dynamics throughout the diffusion period (equation IT]. but is followed by a teacher forcing step in the diffusion period during training (algorithm 2p. As per the advice of seminal work on DDPMs, we weight each term of this loss equally, even though this violates its original formulation as an ELBO Ho et al. (2020).
The loss function also penalises the deviation of the distribution of the behavioural subspace activity at cue offset Do from the base (unit normal) distribution across many trials. This is a necessary starting point for the denoising process Ho et al. (2020) because the trajectory not reacher-forced prior to offset, must explicitly regularise it. We achieve this by applying L2/Frobenius penalty to the deviation of the first two moments of activity at the start of denoising (i.c. the first timestep of delay 2) from this target distribution [ooo and Lengyel (2022). Finally, we apply L2 regulariser on neural activity across the trial, including prior to the denoising period] Yang et al. (2019) Stroud et al. (2023).
Normally, training DDPM parameters across different timesteps, i.c. Po 7) can be done in parallel across data samples ær are noised to valid sample of g with the noising process
18

---

Published as a conference paper at ICLR 2026
accumulated over timesteps, then the DDPM objective can be calculated independently of the noised state at other timesteps. This is because DDPMs are Markovian trained parameters can denoise a sample independently of previous timesteps. This is not possible when training RNNs to perform memory tasks. The necessary computations memory maintenance during the delay period and implemented via non-linear dynamical motifs (Vyas al. 2020 Driscoll et al.) 2022; Versteeg et al. 2025) which are also driven by the remainder of the neural state space which not projected to the response space. As such, training RNNs requires simulation of full task trials in sequence. As neural activity trajectories will likely veer out of the sequence of distributions defined by the noising process q. this sequential training risks detrimentally biasing the training distribution. To solve this, we use leacher-forcing ensure the correct distribution of training targets satisfied (Appendix D Williams and Zipser 1989)
retrieval during the cueing period
Algorithm 2 Training with teacher forcing. Here, F, Σ, are the empirical mean and covariance of r across many trials with the same experimental conditions presented.
repeat
Generate stimulus set and cue (, ,2(2),n²) for each independent trial, indexed by do
From BNS, generate a sample from the desired behavioural distribution: x7 and apply the * noising process (equation 2 to generate noised data f0:7-1
Initialise network activity
Run network dynamics for the prestimulus, stimulus, delay, and cue periods (Figure 77,, Appendix B with noisy discretised leaky dynamics (equation 3 wittoout projection), finally
arriving at network activity at cue offset ro
for = 0 to T do
Replace behavioural subspace: v. <- r x. WI rr
Calculate mean transition kernel Pag (r', t) (equation 3) and stash
Run dynamics one step (equation IT] and add scheduled process noise (Appendix G).
end for
Train on joint loss function
|gg (ag * .x iso. (r,, 72
Σ
n (115ol12 + | 11,
timesteps
end for until converged
D.1 CHANGING OPTIMISATION WINDOW FOR TASK-OPTIMAL NETWORKS
Task optimal networks trained simple MSE loss function:
T
(xe") cos 2f
AWT xT
W WT
(12)
sin Z7
iee they are trained to minimise distance to the point on a circle in TR.2 with the same argument as the task-optimal colour estimate 2((x') Training with similar, output-based MSE loss function against multimodal target distribution would lead avee-aver,ggng, warranting the process-basee based MSE loss we have developed in this work. This is illustrated in Figure which shows samples generated from a network trained in this way.
Note that only the final network rate vector is behaviourally relevant, as it is used to evaluate the colour estimate of the network for that trial (Appendix [A). We refer to the number of timesteps T* before this critical timestep which the optimisation kicks in as the optimisation window. We maintained the same duration and neural noise schedules of each trial period across all networks (Appendix G). For different. simpler VWM task with single, variable delay duration, it has previously been shown that adjusting the optimisation window alters the neural coding schemes al. (2023). Specifically, it affected the cross-temporal decodability of task
used for VWM Stroud /
19

---

Published as a conference paper at ICLR 2026
+
+
+
+
Figure 9: Network trained on a MSE loss (equation 12] against samples from a multimodal target distribution of estimates. The samples are coloured by colour estimate which they are interpreted as, and the modes of the target distribution for each trial is coloured by the nominal colour estimate at which it centered. As the probe feature values get closer, the blue distractor is more likely to cause a swap error, as reflected in the target distributions. The target multimodality is not achieved, instead the network interpolates between the two distinct modes.
variables in storage for the duration of the delay before the estimate is produced. In the absence of neural evidence compare for the neural coding scheme across this delay Panichello and Buschman (2021). we presented the extreme case of T = T Figure 4E (left), i.e. the MSE loss function in equation 12 was applied all timesteps of the second delay. Figure I0 shows the equivalent plot for multiple task-optimal networks with different optimisation windows. In all cases, All estimates were tightly bound to the circle in the behavioural subspace, and no swap errors were observed (Figure BF [top]). Again, we see that the hallmarks of neural data not captured, although in networks which are afforded time before being penalised for behavioural deviation, there is a shared alignment in the darget-ditrractor colour planes during the optimisation window. Regardless, without swap errors, it is difficult to make connections to behavioural data, as was the original aim of our novel training method. Investigating the link between this property of neural activity and optimal information loading is key to understanding both the mechanism of memory retrieval, and the dynamical source of swap errors during this time (55) and is left to future works.
D.2 TRAINING CURRICULUM
Training feature-cued networks directly as described above, be it on task-optimality or behavioural realism, proved inefficient and time-costly. To speed up training, feature-cued networks were first trained much shorter version of the task, with second delay of T = T 5 timesteps. This ensured that networks were initialised with the ability to, at least, retrieve cued items from memory. From this checkpoint, each network type was then trained independently to prevent cross- contamination of the specialised mechanisms we sought to develop and study in each model variant.
D.3 COMPUTATIONAL RESOURCES
All models were trained using a single NVIDIA RTX A5000 24GB GPU at a time, All optimisation is implemented in PyTorch, utilising the Adam optimiser. A learning rate of 0,001 was found satisfactory across all experiments presented, and other optimiser hyperparameters were maintained the library defaults. A batch size of 32, each with 512 independently simulated trials (indexed by k: in Algorithm 2). was used throughout.
Index-cued networks trained for task optimality (behavioural realism) typically took 35000 (300000)
batches at 3.0 (3.0) batches per second, resulting in 3.2 (28) hours per run. These experiments typically occupied 1. 15650 MiB of VRAM. Feature-cued networks trained for task optimality (behavioural realism) typically took 2000 (300000) batches at 2.4 (2.3) batches per second, resulting in 0.23 (36)
'Compared to the training preamble (see later in main text) this was just a matter of extending time horizons, rather than learning new memory retrieval mechanism from scratch
20

---

Published as a conference paper at ICLR 2026
Cosine arange betwen pare
0.4
0.2
0.0
Time relative to cue onset (5)
Figure 10: Plane alignment for different task-optimal networks (see appendix text). Optimisation window highlighted in red
hours per run. These experiments typically occupied I 15870 MiB of VRAM. The training preamble with shortened trial duration was typically took 500000 batches at 8.6 batches per second, resulting in an additional 16 hours per initialisation; but this was inherited by multiple downstream networks with different final objectives. These experiments typically occupied 4560 MiB of VRAM. The reduction in memory is due to the shorter time per trial in this preamble.
E NEURAL SIGNATURES OF SWAP ERRORS
Alignment of PC3
Fired value
Figure 11: Continuing from Figure 5C cosine between PC3 for response time activity geometries
generated by different fixed colours, averaged over spatial configurations (spatial pairs).
Figure 11 provides an extension to the alignment results of Figure 5C (left). There, the item locations kept fixed, as well as one item's colour. The response-time hidden activity mr. when projected to its first three principal components (PCs), forms a characteristic 2 ring structure. These PCs are shared across initial item locations, as long as the fixed colour the same (blue histogram). This not the case when the colours are dissimilar (red histogram). Figurel 11 provides the mean value of this histogram for a wider range of fixed colour pairs. We see a cyclical structure of PC3 alignment with respect to the fixed stimulus colour: This provides preliminary evidence of toroidal structure, which would have to verified with further topological data analysis Cueva et al. (021)
F EXTENSION TO CUE COMBINATION TASKS
F.1 TASK AND GENERATIVE MODEL.
To demonstrate the generality of our approach beyond visual working memory, we trained networks causal inference task for multisensory cue combination. In this task, observer receives visual
21

---

Published as a conference paper at ICLR 2026
(s.) and auditory (sa) cues, and must estimate the location of the auditory source 2n (Figure H3.2 The cues may or may not originate from a common source. When they do share a common cause, integrating both cues improves estimation accuracy, but when they do not, the visual cue acts distractor that should be ignored.
We generated synthetic behavioural data using the Bayesian causal inference model from [Beierholm ett al. (2007), which has been successfully fit to human psychophysical data in this task. The generative model assumes that visual and auditory sources (Zy and =.) are drawn independently from Gaussian prior with standard deviation pp. On each trial, these sources either share a common cause (C = 1) with probability Prommon- or are independent (C 0). When
1, both cues are generated from the same underlying location; when C = 0, they are generated independently. The cues are corrupted by Gaussian sensory noise with standard deviations d, and da respectively (Figure F.3.2p). The ideal observer computes the posterior distribution over 8a by marginalising over the latent causal structure:
[ xa) - pCCC = | Ta (s)a((a To C= + p(C = 0| I (o))l(8 C= (13)
Following the original work, we used parameter values or = 2.14, 0a = 9.2, op = 12.3, and Poommon 0.28. Spatial standard deviations were scaled by factor of 1/12.3 to match the spatial scales used in our network training. The resulting posterior distributions are typically bimodal when the cues are separated (see Figure [.3.2p for examples with fixed H. -), with one mode corresponding to integration (visual cue has a common cause) and another to segregation (visual cue has an independent cause and should be ignored).
In summary, the generative model assumes:
Za ~ ~N(0, of))
(14)
C c=f0 10 w.p. Peommon
(15)
11 w.p.
1 1 Poommon
za
if C =1 if 0
(16)
N(0,0 of))
8a
N(22; o2)
(17)
~ Ni(z,,
(18)
F.2
NETWORK TRAINING
Unlike the circular feature spaces used in the VWM tasks (3. the cue combination task involves linear spatial locations. We therefore adapted our architecture to receive 2D sensory inputs (s,, sal during the stimulus presentation epoch, and to output ID estimates of auditory location &a via a single output dimension (Figure H).22). The trial structure comprised prestimulus baseline period, followed by simultaneous presentation of both cues, then a delay period during which the network performs diffusion-based denoising to generate estimate.
We used the same dendritic tree RNN architecture as in the feature-cued VWM networks, with identical regularisation and training hyperparameters. The key difference was in the dimensionality of the behavioural subspace - here, a ID projection x = w rather than the 2D projection used for circular colour reports. As in the VWM case, the network was trained using feacher-foreing during the delay period to denoise from the corrpption process to synthetic behavioural samples. In this case, these samples were drawn from the trial-specific posterior distribution p(2,) Se; 8a), as infered by an ideal observer (i.e., with access to the real generative model parameters). This posterior was discretised into 200 bins spanning [-4ap; +-4rp)), and samples were drawn by first selecting a bin according the posterior probability mass, then adding uniform noise within that bin to generate continuou--valuee targets.
F.3 RESULTS
F MODEL CAPTURES IDEAL OBSERVER POSTERIORS
We verified that our trained network successfully reproduces the ideal Bayesian observer posteriors across different cue configurations. Figure F3.2 shows results for the case where the auditory cue
22

---

Published as a conference paper at ICLR 2026
is fixed its prior mean (s., 0) while the visual cue location S, varies from -3 to +3. The ideal observer posterior (Figure|F.3.2)) exhibits characteristic structure: a strong horizontal streak at 2, 0 (corresponding to the auditory cue location) and diagonal streak interpolating between Fa = and 2n = xv (corresponding to the visual cue location when integrated under the common cause hypothesis). The relative strength of these two modes is determined by the posterior probability of common cause, p(C = 1 Ta), which decreases as the cues become more separated (Figure F..22)
When Tn and Ha are similar (left side of heatmap), the posterior is dominated by the common cause hypothesis, producing single integrated estimate that combines both cues with reliability- weighting. As the cues diverge (right side of heatmap), the posterior becomes increasingly bimodal, with the segregation mode (at 8a = Ia 0) gaining probability mass while the integration mode (along the diagonal) weakens. This bimodality reflects the observer's uncertainty about whether the cues share common cause.
The network successfully reproduces this structure (Figure F.3.2 2 2p), evidenced by histograms of samples generated by the network. Example slices (Figure F.3.2 show close correspondence between the target posteriors (black lines) and empirical sample distributions (colored histograms) across a range of visual cue locations, capturing both the unimodal regime when cues are close and the bimodal regime when they are separated.
F.3.2 QUANTIFYING ACCURACY ACROSS THE FULL CUE SPACE
To comprehensively evaluate the network's ' performance, we computed the Kolmogorov-Smiroov (KS) statistic comparing the ideal and empirical posterior distributions across a 20 20 grid spanning all combinations of E [-3, 3] and E 1-3, 3] (Figure F.3.2). The KS statistic quantifies the maximum vertical distance between two cumulative distribution functions (CDFs), with values near zero indicating close agreement.
The heatmap reveals that the network achieves high accuracy (low KS statistics, shown blue) across most of the space. Example CDF comparisons (Figure H3.2p) illustrate this agreement for representative conditions spanning the range of KS values. For each comparison, the ideal CDF (black line) closely tracks the empirical CDF computed from network samples (coloured line), with the maximum discrepancy (red vertical line) typically small. Cases with slightly higher KS statistics tend occur when cues are moderately separated regime where the posterior most strongly bimodal and thus most challenging to sample accurately.
23

---

Published as a conference paper at ICLR 2026
A
Shared source
Different sources
z (unseen)
Z
Z
Sg
Sa
Sa
B
Pre-stim.
In-trial time
Peommoo
Delay hore Stim.
Q T
Causal inference for multisensory combination. A Task schematic: the observer receives visual (s,) and auditory (s,) cues and must estimate the auditory source location 2, Cues may share a common cause. B Trial timeline for the RNN. After a prestimulus phase, both cues presented simultaneously. followed by a delay period during which the network denoises to produce an estimate. Unlike the VWM tasks, sensory inputs are 2D [Sp, so] and output is ID (estimate of auditory location fit to be a sample of the ideal observer posterior). C Bayesian graphical model. Sources 2n and za are drawn from common prior, and the causal structure C' determines whether they are bound together or independent. Cues and sa are noisy observations of the sources.
Figure 12:
24

---

Published as a conference paper at ICLR 2026
A
RNN Output Posterior O Ideal Observer Posterior
2
Possible auditory source locations za
Visual cue location Sy
Figure 13: Network reproduces ideal observer posteriors for fixed auditory cue. Ideal observer posterior p( sul 8a) across varying visual cue locations 8, with fixed auditory cue s. 0 (dashed white line). The diagonal streak (dashed white line) corresponds to the shared source hypothesis, while the horizontal streak corresponds the separate sources hypothesis. B Empirical posterior computed from network samples shows close correspondence to ideal observer. C Example slices selected visual cue locations (colored vertical lines in panels a.b) showing target posterior (black line) and sample histograms (colored). Network successfully captures both unimodal and bimodal regimes.
25

---

Published as a conference paper at ICLR 2026
A
KS Statistic across (sy, Sa) grid
B
0.25
location Sv
KS Statistic
Visual
0.30
0.05
Auditory cue location 5
Za
Figure 14: Kolmogorov-Smine Smirnov analysis across full cue space. A KS statistic quantifying agreement between ideal and empirical posteriors across all combinations of visual (I.,) and auditory (x.,) cue locations. Low values (blue) indicate close agreement. Colored dots mark locations of example comparisons shown panel B Example CDF comparisons for representative conditions spanning the range of KS values. Each panel shows the target posterior CDF (black), the empirical sample CDF from network samples (colored), and the maximum discrepancy (red vertical line). Network achieves high accuracy across diverse cue configurations.
26

---

Published as a conference paper at ICLR 2026
G
TABLE OF CONSTANTS AND NOISE SCHEDULE
Name
Symbol
Value
Minimum angular margin between item feature values
x/4
Radial magnitude of behavioural target
2.5
Variance of behavioural target models
0.2²
Variance of rate vector initialisation target models
o
1.02
Network population size
n
16
Number of dendritic tree layers
L
2
Branching factor of each dendritic tree layer
B1, B2
10, 10
Dimensionality of global time representation
n4
Shared time representation vector size Number of tuning curves for each feature dimension Feature-cued sensory projection size
n1
16
n
62
Index-cued cue embedding size
ne 4
Time discretisation step
di 0.05
Neural membrane time constant
0.5
Prestimulus duration
T.
0.15s, 3 timesteps
Stimulus exposure duration
T.
0.25s, 5 timesteps
Delay 1 duration
Ta
0.75s, 15 timesteps
Cue exporsure
Te
0.25s, 5 timesteps
Delay 2 (denoising) duration
Tg
2.0s, 40 timesteps
Base distribution regularisation weight
71
0.01
L2 rate regularisation weight
72
0.0001
Table 2: Constants and hyperparameters used in main text
27

---

Published as a conference paper at ICLR 2026
Noising process variance 8
Resulting process noise
0.3
0.2
0.1
0.0
5
30
40
Timestops after cue offset
Figure 15: Schedule for corruption noise variance (3;; equation equation 15) and corresponding generative noise standard deviation (ar,) The former linear interpolation between 0.3 and 0.01 across the 40 timesteps.
28


Reviews:
[{"review_id": "RBei83HacL", "reviewer": "Reviewer_LJQh", "summary": "In this submission, the authors propose an alternative method to train models of biological visual processing that don't rely on task-optimization. The proposed method aims to synthetically generate a large amount of behavioral data and introduces a diffusion-based method to train RNNs to fit the generated behavioral data, hence automating the development of brain-aligned dynamical system models without a human-in-the-loop. RNN models of a multi-item delayed estimation problem are trained to display 'swap errors' which are characteristic of human behavior on such visual working memory problems, and a representational geometry analysis is shown to highlight how the trained RNN models capture key indicators of visual working memory behavior.", "questions": "NA. Please refer to my review above.", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Training biologically plausible dynamical systems models using deep neural networks has been saturated for a while with task-optimized neural networks. The proposed approach in this paper to automate the process of generating ecologically valid synthetic behavioral data and a new paradigm to train RNNs on the above generated data with denoising diffusion is clearly a fresh take that is well distinguished from task-optimized network training.\n- While the proposed method is shown to model behavior in a visual working memory task, I believe the general framework proposed here has the potential to scale to, and help interpret other dynamic biological processes in perception and decision making.", "weaknesses": "- While task-optimized networks are notoriously difficult to design and train, they provide a direct way to compute representational similarity relative to biological brains on naturalistic stimuli (images, videos, audio, etc). The proposed methods here, however, train RNNs that operate at a different level of abstraction of the input if I understand correctly. This is expected, as the proposed model is one of biological decision making and not necessarily perception. But if one were to apply the proposed method to model perceptual processes, could the authors please comment on how one could bridge this gap in input complexity consumed by their proposed RNN models? Would BNS style synthetic data generation even be feasible for generating coherent, high-dimensional sensory behavioral data?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hipr4Jbfcf", "reviewer": "Reviewer_eTxk", "summary": "The authors train RNNs to fit distributions of behavioral output. This is done by modifying a diffusion based procedure. The algorithm is tested on a synthetic model of behavior. The resulting network’s hidden state statistics has similarities to those observed experimentally.", "questions": "Line 120 – Why dendritic tree?\nLines 208-214 The description of index and feature based is not written clearly. The overuse of the word feature does not help.\nLine 253: Clarity. If I understand correctly, the network activity r is split. But the sentence reads as if it is split into x and m. The next sentence reads “this output”, but the word output does not appear in the previous sentence.\nFigure 7A: could be useful to explain what the horizontal and vertical axes of s_t are", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Behavior is often probabilistic, and not necessarily unimodal. Most models add some noise on top of a deterministic backbone. Using flexible distributions directly is an important contribution.", "weaknesses": "Despite the introduction, there is no fit to actual behavioral data. Is there any guarantee that this can work?\n\nClarity: There are several places where the writing is not clear. For instance, the paper jumps between index and feature based codes, without explaining the rationale.  \n\nThe basic premise in line 41: RNNs are not the only option for testable hypotheses. They are a specific example of a latent variable model. \n\nLine 110 – biologically plausible is a broad and ill-defined term. Some people only consider spiking networks to be plausible. Others require conductance based models. Others are satisfied with rate models, but insist on Dale’s law.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dMQ5Wlgazl", "reviewer": "Reviewer_96uo", "summary": "This work trained RNNs to exhibit swap errors during a multi-item delayed estimation VWM task. Its success was contrasted with naive training methods applied to task-optimal networks. The proposed method captured a more accurate description of neural mechanisms when trained to adhere to the empirical variability around real swap errors.", "questions": "Please check my concerns and questions in Weaknesses section.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The research problem itself is interesting and promising, i.e. discovering neural mechanisms of cognitive errors. However, it is doubted whether such simulated errors by RNNs can be really called 'neural mechanisms' of human cognitive errors, since the mechanisms were revealed in RNNs rather than human neural activities.\n\n2. The background section makes it easier for the reader to understand the context.\n \n3. The proposed model was evaluated in a working memory task dataset and showed meaningful patterns and signatures of cognitive behaviors.", "weaknesses": "1. Most evaluations are qualitative results, rather than quantitative results. Although the discovered signatures are promising, the lack of quantitative measurements significantly reduces the rigor of this work.\n\n2. Only one task (working memory) is used for evaluation. Although I agree that this is a representative task, the small task domin limits the generalization ability of this proposed model and it is not convincing that it can be effective in other cognitive task domains.\n\n3. Limited baseline or ablation studies to show the unique advantages of the proposed model to discover mechanisms of cognitive errors compared with prior models in existing work. I can only find one naive training method as the baseline.\n\n4. Lack of details in the dataset used for evaluation, including working memory task settings/design, sample size of participants and trials in the task, etc. This makes it hard to evaluate the significance of this work and reliability of the findings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KCW30RJOLZ", "reviewer": "Reviewer_YEvN", "summary": "The paper proposes training RNNs to reproduce behavioral distributions including characteristic error modes than optimizing task performance. They do so by first using a generative model to produce synthetic data for training RNNs, then adapting diffusion training procedures for RNNs to produce realistic behavioral distributions. The authors tested their method on the visual working memory task, and found that the DDPM-trained RNNs, instead of the task-optimized RNNs, produces realistic swap errors and neural representations.", "questions": "1. Does well does the method generalize to a more complex versions of the task, for example, with more items?\n2. How does the method apply to tasks where an output is required at every time step of the trial? \n3. Fig 3C: for the brown-border model, shouldn’t close probe feature values have more swap error? (It looks like it's the opposite in the figure).", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Training RNNs to produce realistic and complex behavioral distributions is a very important and timely research direction. \n2. This paper is very well-written and easy to follow. \n3. The experiments are well designed and the authors provided analysis at both the behavioral and representational level.", "weaknesses": "1. The BNS model used to generate synthetic data still requires hand-tuning, so “automatic discovery” feels overstated. Could the authors comment on this limitation?\n2. In Figure 3A, the correspondence between border colors and RNN models is unclear. Please state this explicitly in the caption (or add a legend).\n3. On line 203, (F) has not been introduced, and the network equation is missing. It would help to present Equation 3 (or at least the network update equation within it) earlier.\n4. On line 252, I think the intended expression is (x = G(r) = W_x r) (the r appears to be missing).\n5. The approach depends on a generative model to create synthetic behavioral data for training the RNNs, which raises several concerns:\n\n   * For more complex tasks, well-founded hypotheses or statistical/generative models of the behavioral process may be unknown or unavailable.\n   * The method assumes the generative model captures key characteristics of behavior, yet the paper provides no quantitative evidence for this beyond showing that swap error increases when the probe is closer to the cued item.\n   * Lines 146–149 claim that models listed in §1 “can be used to faithfully generate synthetic behavioural data,” but this is not supported with empirical evidence or citations to prior work.\n\nIf these concerns are properly addressed, I may consider increasing my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)