PADDL model captures both choice and RT behavior in the RLWM paradigm

We first evaluated the ability of the PADDL model to capture behavior of the participants in the RLWM task. Across both groups, posterior predictive checks revealed that the model shows good alignment with empirical data for both choices and RTs (Figs 1D and 5A). Accuracy increases, and RT decreases, as learning progresses (in terms of the encountered stimulus iterations), and both effects are attenuated with higher set sizes as the WM load increases. The model was particularly effective at capturing the set size-related RT variability during the initial stimulus iterations.

Joint modeling of choice and RT more accurately disentangles RL from WM components

We next evaluated whether including RTs in the PADDL model imposes informative constraints to improve parameter estimation even for parameters of the model that are typically used to explain choice distributions alone. We hypothesized that fitting a choice-only model to the data generated from a joint choice-RT generative process could introduce systematic biases in parameter estimation. We evaluate this hypothesis by fitting the choice-only (softmax) as well as the PADDL model to empirical and synthetic datasets.

We first compared parameter estimates of the choice-only (softmax) and the PADDL model (with the evidence accumulation decision process), when fitted to the empirical data. Fig 2A shows that the softmax (choice-only) model overestimates RL learning rate, WM reliance and WM capacity parameters relative to the PADDL model. Although one cannot yet determine from this exercise alone which parameters are more accurate, we reasoned that maximization of softmax choice likelihood alone provides limited information to disentangle RL vs. WM processes. Conversely, RTs induce a powerful constraint on the joint probability , fundamentally changing the estimation landscape. By requiring the model to simultaneously account for both what was chosen and how quickly it was chosen, the model is forced to estimate parameter values that are consistent with both aspects of the behavioral data. To validate this account, we generated synthetic choice and RT data using the PADDL model, and then fit that data with both the PADDL model and the choice-only model. Fig 2B shows a similar tendency of the softmax model to overestimate parameters such as RL learning rate, WM reliance, and WM capacity. In addition, the noise and perseveration parameters were consistently underestimated in the softmax model. This exercise confirms that when ground truth parameters from the PADDL model are known, the softmax model is less accurate at recovering them.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 2. Comparison of the recovered parameters from the choice-only model and the joint PADDL model.

The plots denote correlations between the recovered parameters on (A) empirical dataset and (B) simulated (synthetic) dataset. The simulated dataset was generated using the PADDL model. Error bars denote 94% HDI ranges for both axes.

https://doi.org/10.1371/journal.pcbi.1014796.g002

How do RTs help inform parameter estimates? Note that although the RL module is more incremental than the WM module, the two will often agree as to which action has highest value in terms of the rank ordering of actions. In contrast with the choice distributions, RTs provide more continuous outputs, and in particular, RTs should decline more incrementally over time when the RL process is dominant (i.e., when capacity is exceeded). RTs are also informative about proactive control strategies, whereby participants may be more cautious with increasing WM demands, even before learning accumulates. To systematically explore the contributions of each mechanism to RTs, we performed two ablation experiments (Fig 3). Each removes or replaces a single mechanism while holding every other learning and decision parameter fixed, so that any resulting misfit to the RT data isolates the contribution of that one mechanism. The first ablation, presented here, targets the RL learning rate; the second, presented in the proactive control section below, targets the set size dependent decision threshold.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 3. Each ablation removes or replaces a single mechanism of the joint model while holding all others fixed, isolating that mechanism’s contribution.

(A-C) Ablation experiments show that PADDL provides better fits and identifies key diagnostic RT patterns. (A) Accuracy (left) and RT trends (right) for the model which uses the estimated learning rate from the choice-only softmax model and the model (model with static decision bounds, i.e., without proactive control). The PADDL trends are plotted in dashed lines for comparison and those from the ablation/alternate models are plotted in solid lines. Static bounds fail to capture the large differences in RTs by set size, whereas softmax produces progressively faster RTs over learning particularly for higher set sizes. (B) Summary statistic showing RT discrepancy by set size in late trials for PADDL compared to the softmax-recovered, overestimated learning rate. (C) Summary statistic showing RT discrepancy by set size in early trials for PADDL compared to the static bounds model. (D-E) Posterior predictive checks show that PADDL captures empirical RT patterns diagnostic of RL learning rate and proactive decision threshold adjustment, as identified in panels B and C. (D) RT differences in empirical and simulated data between low- and high- participant subgroups (in HC) during late/asymptotic learning phase. (E) RT differences in empirical and simulated data between high- and low-m participant subgroups (in HC) during early learning phase. The difference in RT is calculated between the median-split, high and low parameter participant groups as determined from the PADDL model fits. The error bars on simulated trends denote 94% HDI ranges.

https://doi.org/10.1371/journal.pcbi.1014796.g003

The first ablation comparison is diagnostic of how the RL learning rate contributes to RTs. We focus on later trials and higher set sizes because that is where RL, and hence sensitivity to , dominates. In particular, later trials allow RL sufficient experience to accumulate, and as set size increases, the WM contribution is weakened due to capacity limitations [12]; hence RTs in that regime are the most informative read-out of the learning rate. An overestimated implies learned values that separate too quickly, which manifests as RTs that decline too fast over learning; a correctly estimated does not. To test this notion, we simulated RTs from the joint PADDL model twice: once using the s estimated by the PADDL model, and once using those from the softmax model. Posterior predictive checks in Fig 3A and 3B confirmed this interpretation about differential RT predictions from the RL and WM modules. We observed that when learning rates are estimated from the choice-only model, the simulated RTs decline faster than observed empirically, and moreover that this difference became increasingly apparent in the higher set sizes and in later trials. This is because RL is relatively more dominant under high WM load, and the (right) tail of the RT distribution constrains the RL parameter estimates. As predicted, the joint PADDL model more appropriately captures the RT dynamics, specifically at higher set sizes when RL is more dominant. This analysis revealed a novel signature of variations within RL that are manifest in terms of progressive changes in RT over learning in higher set sizes.

Posterior predictive checks further confirmed that this signature is apparent in those participants with high vs. low estimated PADDL learning rates: those with lower learning rates showed progressively slower RTs in high set sizes during the asymptotic learning phase, and these discrepancies were well captured by the PADDL model (Fig 3D). Finally, while this ablation experiment shows the consequences of omitting the decision dynamics during learning, leading to RL overestimation, we also confirmed that the small RL learning rates estimated from PADDL still contribute appreciably to observable behavior. To do so, we ran PADDL simulations setting RL , effectively removing the RL process altogether; these simulations confirmed that the choices and RTs failed to match empirical data without RL contributions (Fig B in S1 Text).

PADDL model improves parameter recovery despite added complexity

These conclusions about RL learning rate overestimation in choice-only models depended on the assumption that choices and RTs are generated by a common process (i.e., the synthetic data were generated from the LBA), and it is in theory possible that the altered estimates of RL learning rates resulted from model misspecification. We next turned to a stronger test: we compared parameter recovery for the choice-only model when it was used to generate the data (using a softmax decision policy) against recovery from the PADDL model. An improvement in recovery would be striking here because the joint model is more complex, having all the same parameters plus the decision process parameters. But as noted above, because RL and WM systems both converge on the same optimal discrete policy, they can partially mimic each other [22] and therefore hinder recovery when using choice-only models.

We simulated a synthetic dataset based on the parameters for each model that fit best to our empirical dataset, and attempted to recover the true data-generating parameters. We simulated the dataset using the mean of all participant-level parameter posterior distributions. Fig F (panel A) in S1 Text shows that the PADDL model performed well and was able to recover all parameters with reasonable accuracy. Notably, even for the parameters common between models, it recovered all parameters as well as or better than the choice-only model (Fig F (panel B) in S1 Text). The learning rate was significantly better recovered with the PADDL model (). These improvements in parameter recovery are notable given that they are applied to parameter ranges that best explain the empirical data across both clinical and non-clinical populations. (We note here that parameter recovery quality is expected to vary across parameters, reflecting differences in identifiability, interdependence, and noise levels. Non-identifiability can be localized to specific parameter subspaces, and critically, recovery is often less robust in parameter regions associated with poorer overall task performance, a characteristic typical of clinical datasets. Therefore, beyond the partial parameter trade-offs inherent in the model structure, the recovery statistics must be evaluated relative to the parameter distributions observed in the empirical (and clinical) data and cannot be numerically compared to those from studies of younger or healthier populations, where identifiability (even for choice-only models) may be higher.) An additional point to note here is that, in spite of the fewer free parameters, the choice-only softmax model was flexible enough to account for the group-level accuracy learning curves across set sizes (Fig A in S1 Text).

Taken together with the previous section, these results suggest that the PADDL model is a better generative model for the empirical data and reinforce the notion that modeling of choice and RT data can improve parameter identification in cognitive models [45]. The latter point is specifically important because the RTs can provide more refined continuous data that inform the RL learning rate over trials.

Proactive control over decision boundaries under varying WM load

Beyond improving parameter recovery, modeling RTs also allows us to assess other neurocognitive processes not typically studied in RL (including the RLWM paradigm), such as speed-accuracy trade-offs. In sequential sampling models (including the linear ballistic accumulator used here), speed-accuracy trade-offs are well captured by adjustment of the decision boundary (sometimes referred to as the decision threshold). With a higher threshold/boundary, more evidence is accumulated before committing to a choice, overcoming the impact of accumulator variability and thus improving accuracy at the expense of speed [41,46]. Moreover, it is a well-established finding that when faced with tasks that require cognitive control to maintain high accuracy, participants adjust their decision boundaries as a way to impose response caution, and that this relates to mechanisms within prefrontal cortex and basal ganglia [27,47,48]. In the RLWM paradigm, we hypothesized that participants might proactively adjust the decision boundaries as a function of WM load. At the start of each block, participants were shown all the stimuli they would encounter in that block. This information about set size could potentially allow participants to proactively increase the decision thresholds – imposing more response caution when they know choices might be harder to remember.

To test this idea, the PADDL model allowed the decision boundary a to vary linearly with set size such that , with k reflecting the baseline threshold and m the set size dependent adjustment or slope. Note that the base model without set size-dependent decision bounds is a strictly nested model under the PADDL model because setting m to 0 recovers the base model with static decision bounds. This nesting is what makes the static bounds model the correct comparison for the proactive control claim: it shares every learning and decision parameter with PADDL and differs only in whether the threshold scales with set size, so a difference in fit cannot be attributed to anything else. The diagnostic signature is the difference in RTs across set sizes during the very first trials of each block, before any feedback in that block could have produced differential learning. Because set size is revealed to participants at block onset before any learning, a set size difference in RT at that point cannot reflect differences in learning or WM and must instead reflect a threshold set in anticipation of load. A static bound cannot produce this pattern, whereas the proactive mechanism can. Across all participants (Fig E in S1 Text and Table A in S1 Text), we found that the 94% highest density interval (HDI) of the threshold slope parameter m was recovered as a positive value and did not contain zero. Moreover, posterior predictive checks revealed that this adaptive threshold mechanism was necessary to capture the large differences in RTs across set sizes even from the very early learning phase, and that the corresponding static bounds model failed to do so (Fig 3A and 3C). Summary statistics showing set size-induced RT effects (as a result of proactive decision bound modulation) confirmed that proactive adaptation parameter m accurately characterizes the observed differences in RTs across set sizes (Fig 3E). We also confirmed that this proactive control model outperformed alternative models in which RTs were slowed due to policy uncertainty. (The model proposed in [18] did not improve the quality of fits over the PADDL model as assessed by posterior predictive checks. Moreover, when adding the proactive control mechanism within the model we also observed similar proactive adjustment of decision bounds; hence this effect is additional to any effects of policy uncertainty.) This suggests that the proactive control mechanism introduced in the PADDL model is critical to account for the empirical data.

It is conceivable that rather than proactive control, participants could instead be reactive [49] and instead of adjusting the initial height of the boundary, they might wait for evidence accumulation to detect conflict (e.g., among multiple items in WM), and then reactively adjust the boundary (e.g.[48],). While such models are typically difficult to estimate, our differentiable likelihood neural network methods rendered their estimation straightforward [38]. We tested this idea by allowing a dynamic decision boundary which collapsed over time (increasing the probability of responding), but where this collapse could be delayed or reduced when conflict (from the higher WM load in higher set sizes) is detected [48]. We thus augmented the LBA decision process model to allow for a collapsing decision threshold with an angle . Posterior distributions (see Table A in S1 Text) however suggested that this angle was close to zero in both groups, and models that allowed this collapse to change with set size also worsened fit, further supporting a proactive rather than reactive control account. This result is sensible here because WM load was blocked and hence participants could adopt a proactive strategy to accommodate the need for more evidence accumulation across the entire block; in contrast, previous findings of conflict-induced changes in collapsing bounds applied in perceptual tasks whereby task condition was unpredictable over trials [48].

PADDL estimated RL predicts out-of-sample test phase choices

The variant of the RLWM task used here [23,42] includes a test phase administered after the learning blocks/train phase, in which participants choose the more rewarding option between pairs of previously learned stimuli without feedback. Because the paired stimuli are drawn across different learning blocks and are encountered after long delays, working memory makes no appreciable contribution at test, and choices are instead governed by the slowly accrued reinforcement learning values, thought to reflect striatal learning [15,50]. The test phase therefore provides a WM-free read-out of retained stimulus-reward values, and a stringent out-of-sample criterion: the learning rate was estimated entirely from the train phase, so test phase performance was never part of the fitting objective. Moreover, performance was never evaluated on the test phase until final model comparison and selection was performed on the learning data.

Because the test phase is thought to reflect RL contributions, we evaluated whether the RL learning rates estimated from PADDL and choice-only models were more predictive of test phase accuracy. For each participant, we extracted the final Q-values from each model at the end of learning, which are fully determined by the history of choices and outcomes and the RL learning rates (, ). Because participants were asked which stimulus was more rewarded (not which response was rewarded), Q-values were weighted according to the frequency of the chosen action, such that each stimulus had an effective action-independent reward value. (Equivalently, these values can be obtained by updating a V-value for each stimulus based on its (stimulus-specific) reward history during learning, ignoring the response selected, but using the same RL learning rate, akin to a “critic” in actor-critic models). We then predicted test phase choices from these values, binned by the within-participant value difference between the two offered stimuli, and compared predicted to empirical accuracy. No parameters were estimated on the test phase data, and an identical softmax choice rule was applied for simulating test phase choices in both models (it is no longer sensible to use the LBA for the test phase; see Methods), so the only quantity differing between the two sets of predictions was the learned values each model implied.

Strikingly, only the PADDL model accorded well with test phase choices, and the two models diverged in a diagnostic way. Supporting the conclusion that it overestimated RL learning rates, the choice-only softmax model systematically overpredicted test accuracy, with predictions exceeding the 94% HDI of the empirical estimates at the medium and high value-difference bins and the discrepancy growing with value difference (Fig 4A). Its psychometric function was correspondingly too steep through the point of indifference (Fig 4B), the signature of over-separated learned values. The PADDL-derived values, in contrast, tracked the empirical curve closely, with only a slight underprediction at the largest value difference, substantially smaller than the softmax overprediction. Quantifying this as the signed deviation of predicted from empirical accuracy, the softmax model was substantially biased upward (mean deviation = 0.0488, CI [0.0313, 0.0657]) whereas PADDL was approximately unbiased (mean deviation = −0.0190, CI [−0.0384, −0.0001]).

Download:

• PNG
  larger image
• TIFF
  original image

Fig 4. Out-of-sample validation of estimated parameters on the held-out test phase (HC group).

Test phase choices were predicted from end-of-training RL values obtained from each participant’s learning phase, with no parameters fit to test data and an identical softmax rule applied to both models, so predictions differed only in the learned values each implied (see Methods). (A) Probability of choosing the higher-valued stimulus by within-participant value-difference tertile, for empirical data (black), PADDL (blue), and the choice-only softmax model (green). The choice-only model overpredicts accuracy, exceeding the empirical estimates at the medium and high bins, whereas PADDL tracks the empirical curve closely. (B) Psychometric function of choosing the right-hand stimulus against the signed value difference (). The choice-only curve is too steep through indifference, the signature of over-separated learned values, while PADDL matches the empirical slope. Shaded bands denote 94% HDI.

https://doi.org/10.1371/journal.pcbi.1014796.g004

To statistically test whether PADDL and the choice-only softmax model translate their internal predictions into empirically accurate choice probabilities during the test phase, we conducted a calibration analysis. This paired prediction analysis places both models in one fully interacted logistic calibration regression (see details in Methods). PADDL’s calibration slope was 1.1179 (95% CI [0.8969, 1.3390]), consistent with ideal calibration (slope = 1). The choice-only softmax slope was 0.6625 (95% CI [0.5527, 0.7722]), indicating overconfident probabilities. The direct choice-only-minus-PADDL slope difference was −0.4554 (95% CI [−0.5895, −0.3212], cluster-robust ). Calibration intercepts did not differ. Thus, PADDL’s predicted choice probabilities closely matched participants’ observed behavior, whereas the softmax model was systematically overconfident, assigning probabilities that were more extreme than the data supported. Thus, PADDL more accurately captured the strength of participants’ choice preferences. (When averaged across all test trials, the participant-level log loss favored PADDL only numerically and did not reach significance (difference 0.0048 0.0067 SEM), but this analysis includes many trials in which RL value differences are small and the two models agree). The better calibration of the PADDL model is the direct out-of-sample consequence of the overestimation documented earlier (Figs 2 and 3): an inflated learning rate produces overly separated learned values that overstate discriminability at test. That PADDL predicts the held-out test phase accurately using parameters estimated solely from the learning phase indicates that its lower learning rate estimates reflect the true incremental RL signal rather than under-identification, and provides independent support for the learning rate estimates within-group (see section Test phase results for SCZ) (Fig C in S1 Text).

Proactive Adaptation of Decision Dynamics in Learning (PADDL) model reveals differences in WM, RL and proactive control across clinical groups

Having established that the proactive control model improves parameter recovery and identifies novel neurocognitive processes, we next set out to determine whether it affords new conclusions about clinical differences between HC and SCZ groups. Previous RLWM studies suggested that SCZ patients show impaired WM but intact RL [19,23,44].

To understand the behavioral differences across clinical groups in terms of the cognitive parameters of the computational model, we compared the means of the group-level posterior distributions obtained from the parameter fitting procedure. We leverage HDI as the Bayesian test of credible difference between the two distributions (see Methods section Clinical group contrasts for details). Fig 5B shows three key findings from parameter estimation. First, replicating prior findings, the SCZ group showed higher rates of WM forgetting/interference as estimated by the WM decay parameter ( = 0.0964, = 0.3658, HDI = [−0.3921, −0.1196]). While the SCZ group also had numerically reduced WM capacity, we did not find credible differences between the two groups ( = 2.6940, = 2.0537, HDI = [−1.1827, 2.5255]). Note, however, that the larger reported above does reduce effective capacity.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 5. (A) Posterior predictive checks for SCZ.

(Left) Learning curves for each set size plotted in terms of accuracy for empirical data and model simulations. The model simulations were generated by sampling the posterior and using the post hoc absolute fit method. (Right) Learning curves for each set size plotted in terms of RT trends for empirical data and model simulations. Shaded regions show 94% HDI ranges of the group mean. (B) Group-level differences between HC and SCZ. Posterior distributions of the group-level means are shown for the three main contrasts – proactive adjustment of decision threshold m, RL learning rate and WM decay . Note that denotes the back-transformed (via exp function) parameter space. (C) (Left) RT differences (or slope) in empirical and simulated data between HC and SCZ groups over the course of learning. The RT slopes across learning are characterized by learning rates – HC group with higher learning rate shows steeper slope than SCZ. The differences in slope increase with set size. The trend in gray shows the same effect on simulated data when SCZ learning rate is adjusted to match to the HC learning rate. Negating the learning rate differences between HC and SCZ ablates these learning rate induced RT slope effects across learning. (Right) RT differences (or slope) in empirical and simulated data between high- and low- participant subgroups (in HC). The error bars on simulated trends denote 94% HDI ranges.

https://doi.org/10.1371/journal.pcbi.1014796.g005

Second, unlike previous studies using choice-only models, PADDL revealed a lower RL learning rate in the SCZ group than in controls ( = 0.0073, = 0.0045, HDI = [0.0009, 0.0046]). In our previous study [44] with the same dataset, we found that fitting a choice-only model revealed differences only in WM decay (), but not in the RL learning rate (). Using summary statistics on empirical and simulated datasets from HC and SCZ, we confirmed that the model was able to account for a relatively slower rate of decrease in RTs over the course of learning (i.e., less steep RT slopes) in the SCZ group than in the HC group (Fig 5C). The difference is more pronounced in higher set sizes because of a more prominent RL contribution. The model is able to account for these group differences in RT slopes owing to accurate characterization of the RL learning rate – removing HC-SCZ learning rate differences cancels out these RT slope effects (Fig 5C). A similar analysis of high- and low- participants in the HC group revealed that the RL learning rate accounts for these differences in RT slopes. Overall, these conclusions were afforded by the better parameter recovery and identifiability of the RL learning rate enabled by our PADDL model, as fitting with the choice-only model again failed to reveal a difference in learning rate between SCZ and HC. As shown above (Fig 3A, 3B and 3D), the appropriate RL learning rate accounts for more incremental speed-ups in RTs over learning. Indeed, Fig 5A and 5C show that SCZ patients showed more gradual declines in RTs over trials compared to HC.

Third, SCZ patients showed reductions in proactive adjustment of the decision boundary as a function of WM load ( = 0.0536, = 0.0408, HDI = [0.0049, 0.0218]). Qualitatively, this effect can be observed by inspecting the RTs during early learning trials: HC participants show more pronounced increases in RT with higher set sizes than do SCZ (Fig 5A). This result reveals a novel computational difference between SCZ and HC in the RLWM paradigm, but is generally consistent with prior literature suggesting cognitive control and performance monitoring deficits in SCZ [51,52].

Ethics statement

The experimental procedures were approved by the Institutional Review Board at Washington University in St. Louis. All participants gave informed written consent prior to participation.

Participants

This study was conducted as part of the Cognitive Neuroscience Test Reliability and Computational Applications for Schizophrenia Consortium (CNTRaCS). The participants were recruited from five different US locations – the University of Maryland, the University of California, Davis, the University of Minnesota, the University of Chicago, and Washington University in St. Louis. As part of this larger study, we collected data for a total of 255 participants across the four diagnostic groups (HC = 87, SCZ = 67, MDD = 54, BP = 47). This dataset also included simultaneously recorded EEG. No participants were excluded within any analyzed group.

Diagnostic status across all study groups was determined using the Structured Clinical Interview for DSM-5, as detailed below. All patient participants were clinically stable, with no medication adjustments within the preceding month and none anticipated in the subsequent month. Individuals with MDD satisfied DSM-5 criteria for a minimum of two depressive episodes, with at least one episode occurring within the prior three years. Participants with SCZ were diagnosed with either schizophrenia or schizoaffective disorder, based on the SCID assessment described below. HC participants had no current psychiatric disorder, were not using psychiatric medication, and reported no family history of psychosis. All individuals were between 18 and 63 years, had no history of neurological injury, and met criteria for the absence of substance use disorder during the three months preceding study enrollment.

A detailed analysis of the behavioral and EEG data across clinical groups is reported in our previous work [44]. The main focus of this study is on the behavioral data from HC and SCZ groups because multiple studies in the past have contrasted these groups where SCZ deficits were shown to be limited to WM contributions. However, for completeness, we also report effects in the other populations (BP and MDD) in the supplementary – see section Clinical results for BP and MDD (Fig D and Table B in S1 Text). These groups required EEG-based phenotyping to dissociate specific neurocomputational mechanisms underlying performance deficits [44].

Diagnostic assessment

All participants completed the Structured Clinical Interview for DSM-5 (SCID-5) with a trained researcher under the supervision of a doctoral-level clinician. Raters received training through remote webinars in which scoring procedures and anchor points were reviewed, and they also completed six standardized training videos followed by group discussions. After training, raters were required to reach agreement with “gold standard” ratings provided by experienced clinicians at the Maryland and St. Louis sites across at least six practice interviews. Agreement was defined as having no more than two items differing by more than one point from the reference ratings. To ensure reliability throughout the study, the St. Louis site recorded one interview every 2–4 weeks, which all raters evaluated. These were subsequently discussed in remote consensus meetings to resolve discrepancies and maintain inter-rater consistency.

Reinforcement Learning - Working Memory (RLWM) task

The RLWM task [12,19] was developed to dissociate the respective contributions of reinforcement learning and working memory to stimulus-response acquisition. In this paradigm, participants acquire stimulus-response associations through trial-and-error learning within a three-alternative forced-choice design. Correct responses always yielded positive feedback, with reward magnitude taking a value of 1 or 2 points; incorrect responses yielded 0 points. The overall structure of the task is divided into two distinct phases: a learning phase and a test phase.

During the learning phase, participants complete multiple trial blocks in which WM demands are parametrically manipulated. These manipulations occur through variation of the set size – the number of unique stimuli per block, ranging from two to five as well as through the manipulation of the delay, defined as the number of intervening trials before a stimulus is re-presented. These design features allow for systematic evaluation of WM-related learning by varying WM demands both within and across blocks. In contrast, the RL-related contributions to performance are reflected in progressive improvements in accuracy that accumulate as a function of reward history, reflecting the gradual strengthening of stimulus-response associations. In total, participants completed ten training blocks with varying set sizes, amounting to 360 trials overall. By jointly examining performance under varying WM loads and accuracy improvements driven by reinforcement history, the RLWM task offers a structured framework for parsing the unique and interactive effects of RL and WM on learning processes.

The subsequent test phase is intended to measure longer-term reward retention or value-based learning effects. We do not include the test phase as a constraint during model fitting, since parameter estimation can be conducted accurately using the learning phase data alone (see Computational Modeling). We instead reserve the test phase as an independent, out-of-sample validation of the fitted learning rates (see Results and Methods), for which it is well suited because working memory contributes negligibly to retention across blocks. However, we do note that our prior work has not identified clinically relevant behavioral differences in the test phase [44].

PADDL model: Joint computational model of choice and RT

The RLWM model was utilized to characterize learning as the joint outcome of two concurrent mechanisms: reinforcement learning and working memory. This computational framework was applied to participants’ empirical data, with model fitting performed on both choice behavior and response time distributions.

Conceptually, the RLWM model integrates two systems that operate in parallel, each contributing uniquely to behavior. The WM system supports rapid acquisition of associations but is constrained by limited capacity and temporal decay, whereas the RL system provides slower yet more stable learning that enables long-term retention. Both mechanisms independently encode state-action value representations, which are updated through temporal-difference learning rules. To capture the characteristic one-shot updating of WM, its learning rate is fixed at 1, reflecting immediate encoding but without the iterative adjustments typical of RL. WM capacity is represented by the parameter C, which provides a probabilistic estimate of the number of active “memory” slots or capacity available at a given time. The natural decline of WM representations is modeled by a decay parameter , capturing the process of information loss across trials. The mixture decision policy is determined by a weighted integration of the RL and WM policies, with the parameter indexing the individual’s general propensity to rely on the WM process over the RL process for decision-making.

Temporal-difference updates in RL and WM processes. Eqs. 1 and 2 define the trial-by-trial updating rules for the RL and WM components, respectively. In these formulations, s refers to the current trial’s stimulus or state, a denotes the action chosen, and indicates the reward outcome obtained on that trial. The RL component is further parameterized by , the learning rate that governs the rate of incremental updating, and by , a perseveration parameter that captures the tendency to ignore negative feedback. For the purposes of model implementation, reward outcomes were binary coded, with 0 corresponding to an incorrect response and 1 to a correct response. Our prior analysis [17,23] has demonstrated that incorporating a three-level reward coding scheme (0–1–2 points) within the Q-learning updates did not improve model fit relative to the simpler binary coding. (The reason is probably that the participants know there is only one correct response for each stimulus, so the best they can do is find that rewarding response.)

(1)

(2)

WM decay. The model accounts for WM decay by progressively reverting the WM Q-values toward their baseline initialization of , where denotes the total number of available actions. This decay process is parameterized by (Eq. 3), which quantifies the extent to which intervening trials disrupt or interfere with WM representations before the same stimulus is presented again.

(3)

Behavioral policy formulation. On each trial, the Q-values generated by the RL and WM components are converted into action choice probabilities through a softmax transformation, as specified in Eq. 4. The expression p(a|s) gives the individual probability mass assigned to each action. The tuple is individually derived to get RL and WM policies. This function incorporates an inverse temperature parameter, , which was fixed at 100 for all participants in order to minimize potential confounds and avoid undesirable trade-offs with other free parameters of the model.

The resulting policies, and , are then combined into a weighted mixture policy, . The relative weight of this mixture depends on the degree of WM engagement in a given block, represented by , whose computation is described in Eq. 5. The parameter reflects a participant’s baseline tendency to rely on WM across different task contexts, while C denotes the individual’s WM capacity. The variable corresponds to the number of unique, intervening stimuli encountered since the last correct response for any given stimulus k. It is important to note here that delay, measured since the last correct response, scales with set size. Finally, Eq. 6 specifies how the combined policy is derived as a weighted integration of and .

(4)

(5)

(6)

Incorporation of undirected noise. The final evaluation policy for the train phase, , is obtained by combining the mixture policy with a uniform policy that assigns equal probability to all possible actions. This addition accounts for random noise or attentional lapses that may be reflected in participants’ behavioral data. The parameter specifies the extent to which such noise or lapses influence decision-making. Eq. 7 formalizes the computation of the final behavioral policy following the integration of these noise-related probabilities.

(7)

Evidence accumulation decision process. The decision-making process is formalized through an evidence accumulation framework. Building on prior work [18], our approach replaces the standard softmax choice rule with a linear ballistic accumulator, allowing simultaneous modeling of choices and RTs. In the LBA, each of the three possible actions is represented by an independent accumulator, with accumulation rates proportional to the action values derived from the RL and WM processes. A decision is registered once one of the accumulators reaches its boundary. Because both RL and WM values evolve dynamically across trials, the associated RTs also shift over time. Additional trial-to-trial RT variability is captured through random variation in the starting points of accumulation.

While the basic LBA provides an analytical likelihood function, for the purposes of both methodological generalization and hypothesis testing, we employ a variant that incorporates linearly collapsing decision bounds (the LBA Angle). This extension allows us to test whether decision thresholds decrease as time progresses and also whether reactive control could be deployed to delay or reduce this collapse [48]. However, because an analytical likelihood function is not yet available for this variant, we employ surrogate likelihoods approximated through LANs [37]. The LAN approach accurately approximates the surrogate likelihood, allowing us to apply a larger class of candidate models without the need for elaborate mathematical derivations for such non-canonical variants of traditional evidence accumulation models. It is important to note that while we employ a model without within-trial accumulation noise (the ballistic LBA), where variability stems from across-trial noise in drift and starting point (as well as learning), our framework and method could be generally applied to a broader class of evidence accumulation models including sequential sampling models such as racing diffusion models.

Within this framework, each action’s accumulator is assigned a drift rate () proportional to the final evaluation policy (). At the start of each trial, accumulation starting points are sampled from a uniform distribution with an upper bound set by the variability parameter (z). The first accumulator to cross the decision threshold (a) determines both the selected action and the RT. Importantly, the threshold a is adaptively modulated by set size (ns) on a trial-by-trial basis. This adjustment follows a linear function of threshold intercept (k) and threshold slope (m), as specified in Eq. 8.

(8)

It is worth noting that we initially evaluated a generalized logistic regression function of set size to model the dependence of decision threshold on ns. However, we ultimately adopted only the a-regression in its linear form, as the parameters recovered from the logistic model indicated that the threshold increased in a manner consistent with a straightforward linear relationship.

To make the inference tractable, we did not estimate the non-decision time in the decision process because it shows severe trade-offs with other learning and decision process parameters [18]. Different parameterizations of the model could be tested to alleviate the issue, or the non-decision time could be estimated by identifying the best-fit parameters over a sparse grid of values.

In summary, the PADDL model is parameterized by the following parameters – .

Model comparison

To perform model comparison, we employed the Savage-Dickey ratio test [71] since the candidate models were nested under the PADDL model. The Savage-Dickey method computes a Bayes factor for a nested (point) hypothesis, such as a parameter being zero, as the ratio of the posterior to the prior density evaluated at that point; a ratio below one favors the parameter being non-zero, and a ratio above one favors it being zero. We examined the two possible nested model alternatives that can account for speed-accuracy trade-offs on a within-trials and across-trials basis. First, we investigated the across-trial trade-offs accounted for by the proactive set size-dependent threshold adjustment. We compare the PADDL model with a Base model (model without the threshold adjustment) by examining the proactive adjustment slope parameter m and threshold intercept parameter k estimates to check if the introduction of the additional proactive mechanism was better able to account for the choice and RT distributions. Setting m = 0 trivially reduces the PADDL model to the Base model. With the PADDL model, the (HC) group-level m parameter was recovered with mean = 0.0536 and 94% HDI = [0.0482, 0.0590]. The 94% HDI does not contain 0, suggesting a credible effect of set size-dependency on the decision threshold. We then compared the null hypothesis H₀: m = 0 against the alternative using the Savage-Dickey density ratio test (). The resulting Bayes factor B₀₁ < 0.001 confirmed the result indicating strong evidence for the alternative. Second, to examine the trade-offs on a within-trial basis, we compare the PADDL model with an alternative model without a collapsing bound decision process by examining the collapse rate parameter . With the PADDL model, the group-level parameter was recovered with mean = 0.0007 and 94% HDI = [0.0000, 0.0023]. The Savage-Dickey density ratio test (, ) indicated strong evidence B₀₁ = 688.5403 for the null, suggesting that the model without a collapsing bound better describes the data.

Model fitting

For Bayesian parameter inference throughout this work we relied on Variational Inference (VI) [72]. VI approximates the posterior distribution by optimizing over a parameterized family of candidate distributions so as to minimize the KL divergence (or another divergence measure) from the true posterior. In our framework, the availability of differentiable surrogate likelihoods through LANs, in conjunction with a differentiable RL implementation, makes it possible to compute gradients of the evidence lower bound (ELBO) [73] with respect to the variational parameters. These gradients allow efficient parameter updates using gradient-based VI methods. We elected to use VI rather than MCMC for several reasons: (i) VI generally converges more rapidly than MCMC and often demonstrates more stable optimization behavior, and (ii) VI permits the exploitation of LANs in high-dimensional parameter spaces, thereby enabling a novel and tractable Bayesian inference strategy.

For all model-fitting procedures, we employed Full-rank Automatic Differentiation Variational Inference (ADVI; [74]) as implemented in PyMC [75]. Full-rank ADVI accounts for correlations within the posterior distribution, thereby yielding more reliable and accurate estimates of marginal variances. Parameters were estimated hierarchically with non-informative priors, and the optimization was run for 500,000 iterations with a learning rate of 0.001 using the Adagrad (window) optimizer. Following model fitting, we sampled 2,000 draws from the variational posterior for each parameter. To reduce the risk of converging to suboptimal local minima, each optimization was repeated three times, and the run yielding the highest log-probability was retained for subsequent analyses. Each run was initialized with a randomized starting point in the valid range of parameter space. All runs were similar and yielded similar results in terms of the estimated parameters and resulting posterior predictive checks (Fig H in S1 Text). The ELBO loss was continuously monitored to evaluate both convergence and stability. The same VI settings were applied consistently across all experiments reported in this study.

The likelihood functions were implemented in JAX [76] and wrapped in PyTensor [77] to interface with PyMC [75] for automatic differentiation and numerical computations. To approximate joint choice and RT likelihoods, we trained a multilayer perceptron with five hidden layers (each of size 120), following the procedures introduced in the LAN framework [37]. While the standard LBA has an analytical likelihood, here we relied on a variant with linearly collapsing decision bounds, parameterized by angle , which exemplifies a case where LAN-based surrogate likelihoods provide a clear advantage. The resulting neural network takes trial-level data and model parameters as inputs and produces approximate log-likelihood values as outputs. Training was carried out on parameter combinations, with simulations each using an efficient implementation of the LBA Angle model [78], using PyTorch [79].

The non-decision time was set to 150 ms [18]. This assumption helps constrain other parameter estimates and seems reasonable as it likely did not differ across groups because – first, the main clinical differences from our study are about how RTs change over learning and with set size whereas the non-decision time just shifts the entire RT distribution by some offset; and second, the participants from clinical groups did not show slower RTs in the first trial but rather just reduced learning. To avoid undue influence of trials with very fast or slow RTs on the parameter estimates, we masked out trial-wise likelihoods for trials with RTs < 300 ms and those with RTs > 2000 ms. The choice data from these trials were still used to update the state-action values though. Across groups, the proportion of masked trials was 2% (HC: 1.7%; SCZ: 2.7%; BP: 1.6%; MDD: 1.7%).

Examining the posterior pair plots suggested that no severe parameter trade-offs exist at the group level (Fig I in S1 Text).

Posterior predictive checks

After fitting, the posterior distribution from the best-fit run was sampled to generate 40 simulated datasets. The trial and feedback sequences were fixed to match the empirically observed sequences to control for variance in behavior attributed to different learning trajectories [80]. The behavior on the first stimulus encounter was also fixed to match the empirically observed choice and RT, which cannot be predicted before any learning. The mean and HDI of the group-averaged measures are computed across the simulation runs to denote the confidence in estimating the group-level means.

Test phase analysis and modeling (Out-of-sample validation)

To validate the recovered parameter estimates from the train phase data against out-of-sample data, we analyzed choices in the test phase, which asked participants to select which of two stimuli had been more rewarding during learning (without feedback). Note that this is a response-independent test; the question is how often this stimulus was rewarded, which is a property of the history of rewards experienced, regardless of whether either stimulus was rewarded for any of the three responses. Because the paired stimuli were drawn from different learning blocks and presented after long delays, the WM system contributes none or negligibly at test, isolating the contribution of the slowly accrued RL values. The test phase was not used during model fitting; all parameters were estimated from the learning phase alone.

For each participant, we reconstructed the end-of-training RL values by re-simulating the trial-by-trial value updates over the observed learning phase sequence using that participant’s fitted learning rates and (which reflects reduced learning from errors). The stimulus values were initialized to chance performance in the three-choice training task. Training trials were then replayed in their observed order. For a trial t containing stimulus and binary feedback , the stimulus values were learned using the standard RL update rule, independent of each response, in Eq. 9.

(9)

In sum, we extracted the final learned stimulus values using the same update equations from the learning phase, but just averaged across responses. The final learned values are then used to simulate test phase choices.

For each test trial, we then used the learned values of the two stimulus options (presented on left and right for each test trial) and predicted the probability of choosing the higher-valued option based on value-differences (Eq. 10; value difference policy ). The inverse temperature was fixed at for both models, identical to its treatment in the learning phase, so that the comparison reflected only the learned values and the shared temperature. This probability was mixed with a participant-specific lapse process estimated using undirected noise parameter from the train phase. The resulting mixture policy gives the final decision policy for the test phase (Eq. 11).

(10)

(11)

We binned test trials into within-participant absolute value difference tertiles (low, medium, and high) and compared model-predicted to empirical accuracy in each bin. To derive the psychometric curve, the signed differences were divided into seven equal-frequency bins across trials.

For the test phase decision rule we used a softmax choice function (Eq. 10) combined with the fitted undirected-noise parameter (Eq. 11), applied identically to both the PADDL and choice-only models. We used softmax rather than the LBA decision process employed during learning for two reasons. First, the analysis concerns whether the learned values predict out-of-sample choice, so the decision rule should be a minimal, assumption-light mapping from value to choice that is common to both models; holding the rule fixed across models ensures that any difference in predicted accuracy is attributable to the learned values rather than to the decision process. Second, it would not be sensible to use the learning phase LBA in the test phase, given that the former required selecting among three instrumental responses for a single stimulus, and the latter required assessing which of two stimuli was more rewarded. Accordingly, the train and test response time distributions differ substantially in their summary statistics (minimum, median, and maximum). (Additionally, the learning phase decision threshold was parameterized as an explicit function of set size () for which the test phase has no analog, and the start-point bias would have to be assumed invariant across phases despite the distributional shift.) Because the analysis targets choice given learned value rather than within-trial decision dynamics, we did not model test phase response times.

We evaluated probabilistic calibration on held-out test choices using logistic regression. For each model, observed right choice was regressed on the logit of its posterior-predictive right-choice probability; ideal calibration corresponds to intercept 0 and slope 1. To compare models directly, trial outcomes were stacked once per model and regressed on predicted log odds, model indicator, and their interaction. The interaction estimated the choice-only-minus-PADDL calibration-slope difference. Sandwich standard errors were clustered by participant, accounting for repeated trials and the paired predictions for each outcome.

Parameter recovery

To validate our modeling approach, we conducted a parameter recovery analysis. We generated a synthetic dataset for 87 simulated participants using the mean of all participant-level parameter posteriors obtained from fitting the empirical data. We then fit both the joint PADDL model and the choice-only model to this synthetic dataset using the same hierarchical structure and priors as in the model fitting step to assess whether the true data-generating parameters could be accurately recovered.

To assess recovery accuracy, we calculated Spearman’s rank correlation between the true data-generating parameters and the recovered estimates for each model. For statistical comparison, each correlation was converted to a Pearson-equivalent value via the sine transform. These values were then Fisher z-transformed. To formally compare the recovery performance of the joint and choice-only models, we contrasted their resulting z-values to obtain the p-values.

Clinical group contrasts

We compared the group-level parameter distributions estimated from the model fit to link behavioral differences in clinical groups to the parameters of the generative cognitive process. To assess the significance of these differences, we estimate the 94% HDI of the contrast where , and N is the number of posterior samples. If zero is contained in the contrast HDI range, we conclude that there is no credible difference between the two groups.

Summary statistics

In our modeling results (see Fig 3B-3E), we examine two critical characteristics of the PADDL model – the proactive control effect and RL effect. The effects account for the subtle but crucial behavioral phenomena (in RT distributions) during early- and late-stage learning. The proactive control effect (characterized by parameter m – via proactively adjusted decision threshold as a function of WM load or set size) describes the model’s ability to capture the RTs during the initial stimulus iterations whereas the RL effect (characterized by parameter – via the contribution of slow, incremental learning in the RL system) describes the ability to capture the late-stage RTs. To assess these effects, we plot the difference in mean RTs across set sizes for across-model comparison (e.g., ablation studies on PADDL vs. PADDL variants) and across-participant comparison (e.g., contrasting high- and low-parameter subgroups within HC). For proactive control effect, the difference in the mean RTs is plotted for stimulus iterations 2 and 3. For the RL effect, the difference in the mean RTs is plotted for the last two stimulus iterations (9 and 10). The mean and HDI of the group-averaged measures are computed across the simulation runs to denote the confidence in estimating the group-level means.

Similarly, to characterize the differences in RT slopes across set sizes in the HC and SCZ groups (see Fig 5C), we plot the slope (or change in RT) over the course of learning as a summary statistic. We compute the difference in RTs between stimulus iterations 1–2 and 9–10. The mean and HDI of the group-averaged measures are computed across the simulation runs to denote the confidence in estimating the group-level means.