What seems obvious in one context can take on an entirely different meaning if that context shifts. While context-dependent inference has been widely studied, a fundamental question remains: how does the brain simultaneously infer both the meaning of sensory input and the underlying context itself, especially when the context is changing? Here, we study flexible perceptual inference—the ability to adapt rapidly to implicit contextual shifts without trial and error. We introduce a novel change-detection task in dynamic environments that requires tracking latent state and context. We find that mice exhibit first-trial behavioral adaptation to latent context shifts driven by inference rather than reward feedback. By deriving the Bayes-optimal policy under a partially observable Markov decision process, we show that rapid adaptation emerges from sequential updates of an internal belief state. In addition, we show that artificial neural networks trained via reinforcement learning achieve near-optimal performance, implementing Bayesian inference-like mechanisms within their recurrent dynamics. These networks develop flexible internal representations that enable adaptive inference in real-time. Our findings establish flexible perceptual inference as a core principle of cognitive flexibility, offering computational and neural-mechanistic insights into adaptive behavior in uncertain environments.

How do animals adapt their behavior when the rules of their environment change unexpectedly? In this study, we address this fundamental question by training mice to perform a task in which they must infer not only the current state of their environment, but also the context that determines how reliable their sensory cues are. We show that mice quickly adjust their decisions to changes in both state and context, often before receiving any explicit feedback. To understand how such flexible inference could be implemented in the brain, we developed a computational framework based on Bayesian inference and trained artificial neural networks to solve the same task. Our analysis reveals that both mice and neural networks learn to infer hidden contexts from sensory inputs, relying on internal computations rather than just trial-and-error learning. We also demonstrate that the internal dynamics of the networks mirror the complex, context-dependent strategies predicted by optimal Bayesian theory. These findings shed light on the computational principles and potential neural mechanisms that allow animals to flexibly adapt their behavior in uncertain and changing environments.

(a) Task structure. Each trial begins in an unsafe state (s_t = 0) that transitions to a safe state (s_t = 1) with probability per time step. The agent observes binary signals: nogo (x_t = 0) or go (x_t = 1). In the safe state, observations are reliable (go with probability 1). In the unsafe state, observations depend on context parameter θ: agents receive misleading go signals with probability θ and veridical nogo signals with probability 1 − θ. Actions terminate trials immediately. (b) Reward structure and speed-accuracy trade-off. Agents earn rewards only for acting in the safe state. Premature actions in the unsafe state yield no reward and incur time penalties. The optimal strategy requires waiting for sufficient consecutive go signals to ensure state transition, with waiting time depending on the reliability context θ. Inter-trial intervals (ITIs) have uniform duration and unsafe-like statistics, preventing anticipation of trial onset. (c) Block design with latent context changes. The context parameter θ remains constant within blocks but changes unpredictably between blocks, creating periods of reliable (low θ) and unreliable (high θ) sensory evidence. Following block switches, optimal performance requires jointly inferring both the current state and the new context from the observation stream alone, without external cues or feedback.

https://doi.org/10.1371/journal.pcbi.1013675.g001

3.1 Trained mice adapt their behavior to a new context within the first trial

We trained adult mice in a physical implementation of our task. Water-deprived, head-fixed mice were placed in front of a water delivery spout (Fig 2A). In the experiments, we represented the go signals (x_t = 1) as a continuous auditory tone, while the nogo signals (x_t = 0) were represented by silence. The sound sequences were generated using the trial structure described in Sect 2. Each signal lasted 0.2 seconds, with no gaps between consecutive signals. The mice were rewarded with a water drop for successful actions, defined as licking the water delivery port during the safe state. Detailed experimental procedures are provided in the Methods section.

To assess how mice adapt their behavior to different latent states of the environment, we trained and tested a group of 17 mice on our task, switching between blocks of varying values of θ. We focused on analyzing the mice’s waiting times, defined as the duration between the last nogo signal and the first lick. Since the underlying state is hidden, the number of consecutive go cues is therefore the only relevant observable. In the next section, we derive the optimal solution in terms of waiting times.

Fig 2B presents the average waiting times for different contexts, demonstrating that, consistent with normative predictions we detail in Sect 3.2, mice wait longer as θ increases (Pearson’s r = 0.83, p < .005). Although behavioral statistics show that mice modulate their waiting time given the latent state, they do not fully reveal what elicited this behavioral change. We were particularly interested in determining whether mice infer the hidden context from input statistics or adapt their behavior based on rewards (or lack thereof) following their actions. We reasoned that if mice relied solely on reward feedback, we would only observe a significant behavioral change from the second trial after the context had changed (i.e., after experiencing the first success or failure in the new context). In contrast, if mice had already adapted their behavior in the first trial, it would suggest that they utilized sensory-based inference to inform their actions.

3.2 Optimal decision-making in a dynamic environment with latent interactions

The environment introduced in the previous section requires the agent to simultaneously track a latent binary state and a latent contextual parameter that governs observation reliability. These dual uncertainties give rise to a challenging inference problem. To illustrate, consider a sequence of k consecutive go signals . In an unsafe state, such a sequence should prompt the agent to update its estimate of the context parameter θ, since repeated go observations in the unsafe state suggest a higher likelihood of observation flips. Conversely, in a safe state, the same stream of signals carries no information about θ. Further complicating matters, a transition from unsafe to safe can occur mid-sequence, confounding the inference of θ and demanding a simultaneous estimation of both the state and the context.

This complexity reflects a form of negative interaction information [44,45] between the latent state and the context: each new observation not only strengthens their statistical coupling, but also amplifies the uncertainty surrounding each individual variable. While the state and the context may be independent a priori, they become statistically dependent once conditioned on observations. Thus, resolving the joint inference of these hidden variables requires a unified treatment of their evolving likelihoods and the agent’s actions.

Identifying an optimal policy in this setting goes beyond mere inference, as it requires the agent to select actions that maximize rewards in the face of dynamic uncertainty. We thus employ the framework of partially observed Markov decision processes (POMDPs) [46] to develop a rigorous Bayesian account of how an ideal observer should update its internal beliefs and choose actions. In the subsequent sections, we derive the belief-update equations for an optimal Bayesian observer and show how these updates drive a policy that achieves near-maximal reward. This analysis reveals how flexible perceptual inference can emerge through sequential belief-state updates that integrate observations with a learned internal model.

3.3 Dynamics of Bayes-optimal perceptual decision-making in dynamic environments

Having derived a Bayes-optimal solution to account for the joint inference of state and context (Sect 3.2), we now apply this theoretical framework to the change-detection task introduced in Sect 2. The Bayesian solution allows us to quantify how an ideal observer integrates ambiguous observations and adapts to implicit context changes. Crucially, it also allows us to study how an ideal drift-diffusion model would implement simultaneous inference of the latent state and context.

Performance of a Bayes-optimal actor. We first assessed the performance of an agent that updates its belief states according to Eq (4) and follows the optimal policy in Eq (11). As illustrated in Fig 3C (top), the Bayes-optimal actor attains near-perfect performance, closely approximating an agent with explicit knowledge of the context. The modest performance gap arises from residual uncertainty in the estimated context, which broadens the distribution of waiting times (bottom of Fig 3C). This efficiency gap reflects the central challenge of flexible perceptual inference: using noisy and partial observations to infer, online, both the binary state and the reliability of incoming signals.

Adapting to a new context within the first trial. One of the hallmarks of flexible perceptual inference, as laid out in the introduction, is rapid adaptation to changing environments without relying on trial-and-error. Our Bayesian framework meets this requirement by continuously updating the context estimate from each incoming observation. Fig 3D demonstrates how the Bayes-optimal actor adapts within the very first trial of a new context block: correlations between waiting times in the current trial and in previous blocks reveal immediate behavioral shifts, even before any explicit feedback. This ability to adapt in real-time matches the behavior observed in mice (Fig 2C and 2D), suggesting that animals may employ similarly sophisticated inference when faced with dynamic uncertainties.

3.4 Neural mechanisms of approximate Bayesian computation

Having derived the Bayes-optimal solution and examined its computational demands, we now investigate how neural networks solve this task. The joint inference of state and context requires maintaining probability distributions over multiple latent variables—a computation that may exceed biological constraints. Can neural circuits discover efficient approximations through learning alone? To address this question, we trained recurrent neural networks using reinforcement learning, providing only reward signals without explicit instruction for probabilistic computation. By analyzing both behavioral performance and internal representations, we examined whether and how these networks implement the principles of flexible perceptual inference.

3.5 Approaching Bayes-optimal performance requires training the recurrent connections

A recent study showed that optimal Bayesian estimators can emerge from neural networks trained for value prediction [60]. In particular, estimators over beliefs can be linearly decoded even from untrained networks, suggesting that the intrinsic dynamics of random networks might be sufficiently rich to support Bayesian computation. In this case, Bayesian inference can be learned through reservoir computing [61], training only the readout weights.

However, there are key differences between value prediction and perceptual decision-making. In particular, we consider an operant task that demands discrete decisions. The key distinction lies in the behavioral output. Whereas value prediction tasks involve estimating a continuous variable, our task requires categorical choice (act or wait) based on a dynamic decision variable. This difference is crucial, as optimal performance in our task depends on accurate single-trial state estimation. The average network output, which might suffice for value prediction, may not capture the temporal precision required for adaptive behavior in dynamic environments.

To examine whether random networks can maintain belief-like representations in this more challenging task setting, we used a "naive" model. This model consists of an RNN with fixed random weights for both recurrent and input connections. As in reservoir computing, only readout weights were trained, which feed into a softmax layer to produce action probabilities (Fig 6D).

We found that retaining the softmax-based probabilistic action selection was essential for the functionality of the naive model. This differs from the fully trained network, which could learn a deterministic policy as prescribed by the theory (using a hard threshold in the final layer). However, the naive model required stochastic action selection to perform acceptably.

The performance of this random RNN model was poor (Fig 6A). Furthermore, it exhibits a different qualitative behavior. While optimal performance requires longer waiting times at higher θ values, the random RNN model showed the opposite pattern, with decreased waiting times at higher θ values (Fig 6B). This maladaptive behavior is a direct result of stimulus statistics. Without proper context inference, the network responds to go cues by simply increasing its action probability, ignoring the broader temporal structure of the task.

Despite the poor behavioral performance, we could decode Bayesian state and context estimators from the network’s activity with relatively high fidelity. Using the same cross-validation methods applied to trained networks, we found that random dynamics had strong predictive power, albeit slightly lower than fully trained networks ( ± 0.001 for the state and R² = 0.91 ± 0.003 for the context estimators; mean ). These results align with the findings of [60], which successfully decoded Bayesian estimators from random networks in a value prediction task. However, they raise the question of why the network’s performance still is so poor.

Untrained networks fail when it matters the most. The apparent contradiction between the successful decoding of Bayesian estimators and poor behavioral performance (Fig 6G left) in random RNN models stems from a specific limitation in their representation of latent variables. While network activity allows good decoding on average, it fails critically at the most behaviorally relevant moment, after several consecutive go signals, when the agent needs to make a decision. Specifically, when the decision variable approaches the action threshold, the random network output Q_t shows no meaningful differentiation between contexts or the number of consecutive go signals (Fig 6H left).

Random recurrent networks have inherent limitations in their temporal integration capabilities. Classical random networks exhibit exponentially decaying memory traces [62], while random LSTM architectures can maintain somewhat longer dependencies [63]. This memory constraint allows linear decoders to differentiate contexts based on recent history, but the discrimination inevitably fails after several consecutive go signals as the memory trace degrades. Assuming a random network, the dynamics of the output Q_t can be captured by a simple linear integration model

where z_t represents a leaky linear integration of the observations x_t and α controls the memory trace. The output Q_t is a linear transformation of the integrator state with gain β and baseline ε. This minimal model accounts for 99.0 ± 0.2% of the single-trial variance while explaining only 74.0 ± 1.6% of the variance in fully trained networks, confirming that random networks implement a fundamentally simpler computational strategy.

Learning the recurrent weights is required even when explicit context is provided. The failure of random networks to maintain context-dependent representations near decision boundaries raises a fundamental question: Does this limitation stem from an inability to infer context or from an inability to implement proper context-dependent integration even with perfect context knowledge? To differentiate between these possibilities, we modified our random RNN model by directly supplying the context as an additional input (Fig 6E) (see Sect 4 for details). This approach isolates the role of context inference from the computational demands of context-dependent evidence accumulation.

With explicit context information, the random network model demonstrated a significant qualitative improvement in its behavioral strategy. In particular, the mean waiting time increases with θ, as expected from optimal behavior and observed in mouse data (Fig 6B). This increase is deterministic for the trained actors (Fig 6C). This qualitative change suggests that access to context information enables the network to implement a basic form of context-dependent integration. Still, quantitative analysis revealed persistent suboptimal performance (Fig 6G middle), particularly pronounced at high θ values where reward rates remained significantly below those achieved by the Bayesian model and fully trained networks (Fig 6A).

The random RNN model shows decreased performance despite achieving appropriate mean waiting times due to excessive response variability (Fig 6C). This limitation comes from the fixed memory capacity of the network, which restricts its ability to integrate go cues beyond a certain duration reliably. To compensate for limited memory and achieve longer mean waiting times, the network implements a suboptimal strategy: maintaining the output Q_t at a context-dependent distance from the stochastic threshold, effectively trading off precision for approximate temporal control. This bias-variance trade-off allows the network to increase average waiting times through reduced action probabilities but at the cost of high trial-to-trial variability.

Similarly to the random network that did not receive an explicit context input, this mode can be explained by the linear integrator model z_t shown in equation Eq (12), where the readout is given by . Here, the baseline depends on the explicit context through the set of parameters . As before, the linear model explains the output trajectory of the random model (with explicit context) very well (explained variance 99.1 ± 0.6%) while failing to explain fully trained networks (74.0 ± 1.6%) (Fig 6H).

This systematic analysis of random networks reveals fundamental limitations in implementing Bayes-optimal behavior, even when provided with perfect context information. The now necessary reliance on stochastic policies starkly contrasts the deterministic Bayesian solution that uses internal belief states. Linear modeling exposed specific computational deficits: Random networks do not modulate their integration rate or decision threshold. When context is provided, random networks improve their behavior by making crude threshold adjustments. In contrast, trained recurrent networks (Fig 6F) develop internal dynamics that independently control both integration rate and baseline in the precise and complementary manner predicted by Bayesian theory. These results demonstrate that optimal inference in our task emerges specifically through the training of recurrent connectivity, supporting a mechanistic link between reinforcement learning and flexible perceptual inference.

Ethics statement

All protocols and procedures were approved by the Institutional Animal Care and Use Committees at The Hebrew University of Jerusalem and were in accordance with the National Institutes of Health Guide for the Care and Use of Laboratory Animals.

Animals

All experiments were conducted using male C57BL/6 mice under water restriction (age: 2-4 months, weight: 20-30 g). Mouse weight was monitored daily and extra water was provided if needed to ensure that mice maintain no less than 80% of their original weight. Mice were housed in a 12-hour light/dark cycle with ad libitum access to food. All procedures were approved by the Institutional Animal Care and Use Committee (IACUC) of Hebrew University.

Surgical procedures

To implant head bars for head fixation, mice were anesthetized using isoflurane (SomnoSuite, Kent Scientific USA) and placed in a stereotaxic apparatus (KOPF 962, David Kopf instruments, USA). Meloxicam (2 mg/kg) was injected subcutaneously for systemic analgesia and 0.1 ml lidocaine (2% solution) was injected subcutaneously before incising the scalp and exposing the skull. A 20 mm long and 3 mm wide head-post was then cemented to the skull horizontally using dental cement (C&B Super-bond dental cement, Sun Medical, Japan) and secured using dental resin (Pi-ku-plast, Bredent, Germany). Post-operative analgesia was administered and mice were allowed to recover for 2 to 3 days before beginning a 1-week water restriction schedule prior to training.

Animal task implementation

Head-fixed mice were trained to perform an auditory change detection task. Each trial was discretized into 0.2-second time bins, during which the trial remained in one of four latent states: unsafe, safe, premature, or inter-trial interval (ITI). In all states, except the safe state, beeps (i.e., go cues) were generated with probability θ (meaning that if , a beep occurred in 30% of bins). In the safe state, beeps occurred with probability 1 (i.e., in every time bin).

Each trial began in the unsafe state. If the mouse withheld licking, the state transitioned from unsafe to safe after a delay of seconds (where and is the floor of x). Licking prematurely during the unsafe state triggered the premature state, which lasted for 7 seconds, before the state transitioned to the ITI. In the safe state (which could last up to 7 seconds), licking was rewarded with a 4–6 drop of water, and caused an immediate transition to ITI. The ITI duration was uniformly distributed between 1 and 5 seconds, and afterwards the next trial began.

Beeps consisted of a harmonic chord composed of pure tones at 2, 4, 6, 8, and 16 kHz, presented via a speaker adjacent to the mouse at 70 dB SPL. The task was controlled using the Bpod behavior measurement and control system (Sanworks, USA). Licking behavior was monitored using a 200 fps camera (Blackfly BFS-U3-16S2M-CS, FLIR, USA) with a Tamron M118FM16 lens (Tamron, Japan). Licking events were detected online using Bonsai [75].

Animal training

Two batches of mice were trained: the first batch (17 mice) was trained on all tested θ-values (0.1-0.9), while the second batch (7 mice) was trained and tested on and 0.7 only. Before varying θ, mice were initially trained on the task with a single θ, and Pavlovian rewards, in order to condition them to lick the water spout. Safe and Unsafe responses were still rewarded and punished, respectively. Initial training with interfered with the later learning of , so θ was initially set to 0.2 or 0.3. These early sessions taught the task fundamentals by exposing mice to reward-predicting safe go’s and unrewarding unsafe go’s. We began varying θ once mouse behavior indicated that they had learned the association between go cue and reward.

The first batch of mice experienced 3–9 sessions in which θ values ranged from 0.1 to 0.7. Once mice demonstrated task acquisition by reliably responding to the state change, sessions with multiple θ values were introduced, presented in blocks of 30–100 trials. Mice were considered experts upon reaching stable performance, which occurred after ∼10 training sessions. Upon achieving expert-level performance, we collected 10 sessions per mouse for analysis (Fig 2B and 2C).

The second batch of mice (Fig 2D and 2E) underwent sessions with or 0.7 changing in blocks of 30-60 trials. The first 5 sessions with both contexts (sessions 6, 7, 8, 9, and 11; session 10 was excluded as it only included ) were classified as ’novice’ data. The last 5 training sessions (sessions 15–19) were classified as ’expert’ data.

Animal behavior analysis

Behavioral data were analyzed in MATLAB (MathWorks, USA). Trials were included in the analysis if the mouse licked at least once and if the trial occurred within the first 300 trials of the session. The first four trials of each session were excluded from adaptation analyses (Fig 2C–2E) because no prior context was available in the first block for mice to adapt their behavior relative to a preceding condition.

For correlation analyses, block averages were computed using all trials in a block except the first and last four trials. To examine the relationship between waiting time and block transitions (Fig 2C), we calculated the Pearson correlation between waiting times of individual trials and the mean waiting time of the preceding or current block.

Waiting time differences (Fig 2D) were computed as the average first-trial waiting time in minus the average first-trial waiting time in . Two-sided Wilcoxon signed-rank tests were used to compare first-trial waiting times (N = 7 mice) (Fig 2E). Pearson’s correlation coefficients and significance values were computed using MATLAB’s corrcoef function. Error bars indicate standard error across mice.

Neural network training

Architecture. We used the PyTorch library to implement an Actor-Critic framework [76]. Both the Actor and the Critic consisted of an LSTM layer (100 hidden units) followed by a linear readout. The readout of the critic was a scalar value, , which was used to compute the Advantage + R_t − where γ is a discount factor, R_t is reward, and is the discounted future return. As per the standard Advantage Actor-Critic (A2C) algorithm, was used solely to train the actor. The actor readout produced 2 pre-activations, which we termed and . The 2 pre-activations were passed through a softmax to yield action with probability:

where − is the network decision variable. For standard policy execution, we sampled from the softmax distribution to facilitate trainability; when comparing to the optimal Bayesian agent, actions were chosen by argmax:

because the Bayesian policy is deterministic as well given a belief state.

Input and initialization. At each time step, both LSTMs received four inputs: (i) a scalar representing the current go/no-go state, (ii) the previous action, (iii) the previous reward, and (iv) optionally, the scalar context (for the untrained + context variant). Weight parameters were initialized using PyTorch’s default linear layer initialization, and LSTM hidden states were zero-initialized before the first episode. For subsequent episodes, hidden states were carried over from the final time step of the preceding episode to avoid signaling the start of a new episode.

Training procedure. Networks were trained for 10,000 episodes, each comprising 20 trials under a single θ (sampled from 20 equally spaced values in [0, 1]). Each trial ended when the network responded. We applied a standard policy-gradient update with an entropy bonus of 0.05 to the actor, while the critic loss was scaled by 0.01 relative to the Actor [55]. Critic weights were updated by minimizing the advantage. The explicit loss of the critic was + R_t − , with a discount factor . The final return G_T was set equal to . The policy-gradient (with entropy bonus) loss applied to the actor was Optimization used RMSprop with an initial learning rate of 0.0005, decaying by 0.99 every 10 episodes, and momentum/weight decay both set to 1e-3. Gradients were computed via backpropagation through time at the end of each episode, then detached.

Task and reward structure. The unsafe state duration was drawn from an exponential distribution with a mean 10 steps (minimum of 1), always beginning with a forced no-go step. The intertrial interval (ITI) was uniformly distributed between 5 and 25 steps. During training, correct actions in the safe state delivered a reward of +1, whereas acting during the stimulus period incurred a negative reward penalty (r = −3) to stabilize learning. At inference time, this negative reward was removed.

Evaluation. We trained 10 networks per model variant (trained RNN, untrained RNN, and untrained RNN + context). Variability across the 10 networks could result from stochasticity in the weight-initialization, state durations, stimuli, and policy. After reaching comparable performance across all networks, the weights were frozen, and 5,000 episodes per network were evaluated and pooled for the final behavioral analysis.

Analysis of network data

Principle component analysis. PCA was performed on the Short-Term Memory stream of the Actor LSTM. Data from all units across all trails (including intertrial intervals) were centrelized and used to compute the principal components using Pythonn’s scikit-learn PCA package.

Fitting network output to the linear integrator model. To model the output of the recurrent network over time Q_t as a linear integrator, we optimize the model parameters under constraints , , and . Optimization was performed using the first 100,000 data time steps, while the entire dataset was used to evaluate goodness of fit (Fig 4G and 4H). The variability between the network realizations (n = 10) was negligible and did not affect the estimated fit parameters.

Drift-diffusion models. To compare our results with conventional evidence-accumulation approaches, we implemented a leaky integrator to estimate θ:

and an integrator with reset inferred state according to

Performance was evaluated for

. 5000 episodes were collected for each agent.

Vector field analysis

In order to measure the history and cue dependent update on each agent’s decision variable and uncertainty, we randomly planted trials of a specific stimulus. The exact stimulus was x_0:20 = [0,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0]. To avoid the confound of different policies on our analysis we prevent the agents from acting during the planted stimulus. This enabled us to collect an equal amount of data (approximately 500 trials per θ) across agents and values of θ without the risk of selection bias. For each step of the planted stimulus we compute the change in each direction by taking the difference along each axis during and after the cue.

Regression from network activity. To decode ground truth and Bayesian estimates, we performed regression on the short-term memory stream of the actor network, as it represents neural activations relevant to decision-making. The critic network, while essential for learning, does not contribute significantly during inference and was therefore excluded [55].

All regressions were performed on trial data, excluding inter-trial intervals (ITI), as they are not behaviorally relevant. To obtain a linear readout, we applied logistic regression to decode the true state s_t and the Bayes optimal state estimate , with the latter regressing to its log-likelihood ratio:

Context decoding was performed using linear regression to predict the true context and the Bayesian estimate . To decode the Bayesian decision variable, we regressed the network activity to the Bayesian state estimator and the belief threshold using the log-likelihood ratio, then calculated the decision variable as . All regression models were trained on half of the trial data and tested on the remaining half.

Bayesian agent

The Bayesian agent was numerically simulated by performing a sequential update of the joint posterior distribution at each time step, following Eq (4). The agent was assumed to be Bayes-optimal, possessing perfect knowledge of the world model, including the observation likelihood and the latent transition probabilities and .

The agent initiated an action when the state estimate

exceeded the optimal decision threshold,

The optimal action time was determined by maximizing the expected reward,

When the state transitions were determined by the RNN we approximated using for simplicity.

S1 Text. Contains supplemental text, methods and figures.

https://doi.org/10.1371/journal.pcbi.1013675.s001

(PDF)