Rescorla–Wagner Models with Sparse Dynamic Attention

Abstract

The Rescorla–Wagner (R–W) model describes human associative learning by proposing that an agent updates associations between stimuli, such as events in their environment or predictive cues, proportionally to a prediction error. While this model has proven informative in experiments, it has been posited that humans selectively attend to certain cues to overcome a problem with the R–W model scaling to large cue dimensions. We formally characterize this scaling problem and provide a solution that involves limiting attention in a R–W model to a sparse set of cues. Given the universal difficulty in selecting features for prediction, sparse attention faces challenges beyond those faced by the R–W model. We demonstrate several ways in which a naive attention model can fail explain those failures and leverage that understanding to produce a Sparse Attention R–W with Inference framework (SAR-WI). The SAR-WI framework not only satisfies a constraint on the number of attended cues, it also performs as well as the R–W model on a number of natural learning tasks, can correctly infer associative strengths, and focuses attention on predictive cues while ignoring uninformative cues. Given the simplicity of proposed alterations, we hope this work informs future development and empirical validation of associative learning models that seek to incorporate sparse attention.

Appendices

A Limits

Here, we provide details on the expression for the expectation and covariance matrix of associative strength on trial t: $\theta _t$ and $\varSigma _t$. Throughout, we assume only that observations $(R_t,C_t)$ are mutually independent between trials (Assumption 2). Taking expectation of the R–W update in (1) gives a recursive expression for $\theta _t$:

$$\begin{aligned} \theta _{t+1}= & {} \theta _t + \mathbb {E}\left[ \alpha C_t\left( R_t - C_t' V_t\right) \right] \\= & {} \theta _t + \alpha \mathbb {E}\left[ R_tC_t\right] - \alpha \mathbb {E}\left[ C_tC_t' \right] \theta _t\\= & {} \theta _t - \alpha \mathbb {E}\left[ C_tC_t' \right] (\theta _t-v^*) \\= & {} v^* + \left( I-\alpha \mathbb {E}\left[ C_tC_t' \right] \right) (\theta _t-v^*), \end{aligned}$$

where Eq. 2 allows for the substitution introducing $v^*$. Applying this recursive formula repeatedly over trials t, a direct formula can be recovered:

$$\begin{aligned} \theta _{t+1} = v^* + \left( I-\alpha \mathbb {E}\left[ C_tC_t' \right] \right) ^t (\theta _1-v^*). \end{aligned}$$

Thus, how quickly expected associative strength $\theta _{t+1}$ goes to $v^*$ depends on $I-\alpha \mathbb {E}\left[ C_tC_t' \right]$. Repeated application of the sub-multiplicative property yields the following bound:

$$\begin{aligned} \begin{Vmatrix} \theta _{t+1}-v^* \end{Vmatrix}= & {} \begin{Vmatrix} \left( I-\alpha \mathbb {E}\left[ C_tC_t' \right] \right) ^t (v^*-\theta _1) \end{Vmatrix} \\\le & {} \begin{Vmatrix} I-\alpha \mathbb {E}\left[ C_tC_t' \right] \end{Vmatrix}^t \begin{Vmatrix}v^*-\theta _1\end{Vmatrix}, \end{aligned}$$

where we are using the 2-norm for vectors and matrices. Furthermore, the 2-norm of a symmetric matrix such as $I-\alpha \mathbb {E}\left[ C_tC_t' \right]$ is the absolute value of its largest eigenvalue. Additionally, all eigenvalues $\rho _{C,1},\ldots ,\rho _{C_,N}$ of $\mathbb {E}\left[ C_tC_t' \right]$ are real and positive, since $\mathbb {E}\left[ C_tC_t' \right]$ is positive semidefinite, and the eigenvalues of $I-\alpha \mathbb {E}\left[ C_tC_t' \right]$ are $1-\alpha \rho _{C,1},\ldots ,1-\alpha \rho _{C_,N}$, since $I-\alpha \mathbb {E}\left[ C_tC_t' \right]$ shares eigenvectors with $\mathbb {E}\left[ C_tC_t' \right]$. Thus,

$$\begin{aligned} \begin{Vmatrix} I-\alpha \mathbb {E}\left[ C_tC_t' \right] \end{Vmatrix} = \max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}. \end{aligned}$$

The condition $\max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix} < 1$ ensures mean associative strength $\theta _t$ converges exponentially to $v^*$. Consequently,

$$\begin{aligned} \theta _{t+1} = v^* + \mathcal {O}\left( \max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^t\right) , \end{aligned}$$

as noted in the main text.

For the covariance matrix of $V_{t+1}$, denoted by $\varSigma _{t+1}$, we first note that

$$\begin{aligned} \varSigma _{t+1}= & {} \text {Var}\left[ V_t + \alpha C_t R_t - \alpha C_tC_t'V_t \right] \\= & {} \text {Var}\left[ V_t + \alpha \left( R_t C_t - C_tC_t'\theta _t\right) - \alpha C_tC_t'\left( V_t-\theta _t\right) \right] . \end{aligned}$$

Upon expanding the expression above using properties of variance, we have that

$$\begin{aligned} \varSigma _{t+1}= & {} \varSigma _t + \alpha ^2 \text {Var}\left[ R_t C_t - C_tC_t'\theta _t\right] + \alpha ^2 \text {Var}\left[ C_tC_t'\left( V_t-\theta _t\right) \right] \\&+\, \alpha \mathrm {Cov}\left[ V_t,R_t C_t - C_t C_t' \theta _t\right] + \alpha \mathrm {Cov}\left[ R_t C_t - C_t C_t' \theta _t,V_t\right] \\&-\, \alpha \mathrm {Cov}\left[ V_t,C_t C_t' (V_t-\theta _t)\right] - \alpha \mathrm {Cov}\left[ C_t C_t' (V_t-\theta _t),V_t\right] \\&-\, \alpha ^2 \mathrm {Cov}\left[ R_t C_t - C_t C_t' \theta _t, C_t C_t' (V_t - \theta _t) \right] \\&-\, \alpha ^2 \mathrm {Cov}\left[ C_t C_t' (V_t - \theta _t), R_t C_t - C_t C_t' \theta _t \right] \end{aligned}$$

We can use the fact that $V_t$ is independent from $R_t$ and $C_t$ to show that several terms above are zero, since

$$\begin{aligned}&\mathrm {Cov}\left[ V_t,R_t C_t - C_t C_t' \theta _t\right] \\&\quad = \mathbb {E}\left[ \left( V_t-\theta _t\right) \left( R_t C_t - C_t C_t' \theta _t - \mathbb {E}[R_t C_t - C_t C_t'\theta _t]\right) '\right]&= 0, \\&\mathrm {Cov}\left[ R_t C_t - C_t C_t' \theta _t,V_t\right] \\&\quad = \mathbb {E}\left[ \left( R_t C_t - C_t C_t' \theta _t - \mathbb {E}[R_t C_t - C_t C_t'\theta _t]\right) \left( V_t-\theta _t\right) '\right]&= 0, \\&\mathrm {Cov}\left[ R_t C_t - C_t C_t' \theta _t, C_t C_t' (V_t - \theta _t) \right] \\&\quad = \mathbb {E}\left[ \left( R_t C_t - C_t C_t' \theta _t - \mathbb {E}[R_t C_t - C_t C_t'\theta _t]\right) (V_t - \theta _t)' C_t C_t'\right]&= 0, \\&\mathrm {Cov}\left[ C_t C_t' (V_t - \theta _t), R_t C_t - C_t C_t' \theta _t \right] \\&\quad = \mathbb {E}\left[ C_t C_t' (V_t - \theta _t)\left( R_t C_t - C_t C_t' \theta _t- \mathbb {E}[R_t C_t - C_t C_t'\theta _t]\right) '\right]&= 0. \end{aligned}$$

Setting these terms to zero yields

$$\begin{aligned} \varSigma _{t+1}= & {} \varSigma _t + \alpha ^2 \text {Var}\left[ R_t C_t - C_tC_t'\theta _t\right] + \alpha ^2 \text {Var}\left[ C_tC_t'\left( V_t-\theta _t\right) \right] \\&- \,\alpha \mathrm {Cov}\left[ V_t,C_t C_t' (V_t-\theta _t)\right] - \alpha \mathrm {Cov}\left[ C_t C_t' (V_t-\theta _t),V_t\right] . \end{aligned}$$

We can use $\varSigma _t$ to rewrite some remaining terms in the above expression:

$$\begin{aligned} \text {Var}\left[ C_tC_t'(V_t-\theta _t)\right]= & {} \mathbb {E}\left[ C_tC_t'\varSigma _t C_tC_t'\right] , \\ \mathrm {Cov}\left[ V_t,C_tC_t'(V_t-\theta _t)\right]= & {} \mathrm {Cov}\left[ V_t-\theta _t,C_tC_t'(V_t-\theta _t)\right] = \varSigma _t\mathbb {E}\left[ C_tC_t'\right] ; \\ \mathrm {Cov}\left[ C_tC_t'(V_t-\theta _t),V_t\right]= & {} \mathrm {Cov}\left[ C_tC_t'(V_t-\theta _t),V_t-\theta _t\right] =\mathbb {E}\left[ C_tC_t'\right] \varSigma _t. \end{aligned}$$

Plugging these expressions in the prior expression leads to

$$\begin{aligned} \varSigma _{t+1}= & {} \varSigma _t - \alpha \varSigma _t\mathbb {E}\left[ C_tC_t'\right] - \alpha \mathbb {E}\left[ C_tC_t'\right] \varSigma _t + \alpha ^2 \mathbb {E}\left[ C_tC_t'\varSigma _t C_tC_t'\right] \\&+ \, \alpha ^2 \text {Var}\left[ R_t C_t - C_tC_t'\theta _t\right] . \end{aligned}$$

Because $\varSigma _t$ is involved in left and right matrix multiplication, we can use vectorization to simplify this expression. Namely,

$$\begin{aligned}&\mathtt {vec}\left( \varSigma _t - \alpha \varSigma _t\mathbb {E}\left[ C_tC_t'\right] - \alpha \mathbb {E}\left[ C_tC_t'\right] \varSigma _t + \alpha ^2 \mathbb {E}\left[ C_tC_t'\varSigma _t C_tC_t'\right] \right) \\&\quad = \mathtt {vec}(\varSigma _t) - \alpha \, \mathtt {vec}\left( \varSigma _t\mathbb {E}\left[ C_tC_t'\right] \right) - \alpha \, \mathtt {vec}\left( \mathbb {E}\left[ C_tC_t'\right] \varSigma _t\right) \\&\qquad + \alpha ^2 \mathtt {vec}\left( \mathbb {E}\left[ C_tC_t'\varSigma _t C_tC_t'\right] \right) \\&\quad = \mathbb {E}\left[ I - \alpha (I\otimes C_t C_t')-\alpha (C_t C_t'\otimes I) +\alpha ^2 (C_t C_t'\otimes C_t C_t') \right] \mathtt {vec}(\varSigma _t) \\&\quad = (I-\alpha K_{\alpha }) \mathtt {vec}(\varSigma _t), \end{aligned}$$

where $\otimes$ is the Kronecker product, $\mathtt {vec}$ reshapes a matrix into a vector, and

$$\begin{aligned} K_{\alpha } = \mathbb {E}\left[ (I\otimes C_t C_t')+(C_t C_t'\otimes I) -\alpha (C_t C_t'\otimes C_t C_t') \right] . \end{aligned}$$

Thus, vectorizing the equation for $\varSigma _{t+1}$ gives

$$\begin{aligned} \mathtt {vec}(\varSigma _{t+1})= & {} (I-\alpha K_{\alpha })\mathtt {vec}(\varSigma _t) + \alpha ^2 \mathtt {vec}\left( \text {Var}\left[ R_tC_t-C_tC_t'\theta _t\right] \right) . \end{aligned}$$

This recursive formula can be used to recover a direct formula:

$$\begin{aligned} \mathtt {vec}(\varSigma _{t+1})&= (I-\alpha K_{\alpha })^t\mathtt {vec}(\varSigma _1)\\&\quad + \alpha ^2 \sum _{s=1}^{t} (I-\alpha K_{\alpha })^{t-s} \mathtt {vec}(\text {Var}\left[ R_s C_s-C_s C_s' \theta _s\right] ) \\&=\alpha ^2 \sum _{s=1}^{t} (I-\alpha K_{\alpha })^{t-s} \mathtt {vec}(\text {Var}\left[ R_s C_s-C_s C_s' \theta _s\right] ), \end{aligned}$$

where $\varSigma _1=0$ since $V_1$ is set to a fixed value.

To arrive at the final expression provided in the main next, we need the following three results:

• (1) If $K_{\alpha }^{-1}$ exists, then $\alpha ^2 \sum _{s=1}^{t} (I-\alpha K_{\alpha })^{t-s} = \alpha K_{\alpha }^{-1}\left( I- (I-\alpha K_{\alpha })^t \right) .$

• (2) If $\max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix} < 1$, then $\text {Var}\left[ R_t C_t-C_t C_t' \theta _t\right] = \text {Var}\left[ R_t C_t-C_t C_t' v^*\right] + \mathcal {O}\left( \max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^t\right)$.

• (3) $I - \alpha K_{\alpha } = \mathcal {O}\left( 1-\alpha \rho _{K,\alpha }\right)$ where $\rho _{K,\alpha }$ is the smallest eigenvalue of $K_{\alpha }$.

Result 1 follows the existence of $K_{\alpha }^{-1}$, since then

$$\begin{aligned} \sum _{s=0}^{\infty } (I-\alpha K_{\alpha })^s = (I-(I-\alpha K_{\alpha }))^{-1} = \frac{1}{\alpha } K_{\alpha }^{-1}, \end{aligned}$$

and

$$\begin{aligned} \alpha ^2 \sum _{s=1}^{t} (I-\alpha K_{\alpha })^{t-s}&= \alpha ^2 \sum _{s=0}^{t-1} (I-\alpha K_{\alpha })^{s} \\&= \alpha ^2 \sum _{s=0}^{\infty } (I-\alpha K_{\alpha })^s - \alpha ^2 \sum _{s=t}^{\infty } (I-\alpha K_{\alpha })^s \\&= \alpha ^2 \sum _{s=0}^{\infty } (I-\alpha K_{\alpha })^s - \alpha ^2 (I-\alpha K_{\alpha })^t \sum _{s=0}^{\infty } (I-\alpha K_{\alpha })^s\\&= \alpha ^2 \sum _{s=0}^{\infty } (I-\alpha K_{\alpha })^s \left( I- (I-\alpha K_{\alpha })^t\right) \\&= \alpha K_{\alpha }^{-1}\left( I- (I-\alpha K_{\alpha })^t \right) . \end{aligned}$$

Remark 4

Assuming $K_{\alpha }$ is invertible is amenable when exploring convergence of the covariance matrix $\varSigma _t$, since to guarantee convergence, we later need to impose an stronger assumption that all eigenvalues of $K_{\alpha }$ are strictly positive. Strictly positive eigenvalues ensure that $K_{\alpha }$ is invertible. Thus $K_{\alpha }$ is invertible in the situation we care about: when $\varSigma _t$ converges. Of course, the covariance matrix does not always converge as we saw in the haystack problem described in the main text, and in these situations, $K_{\alpha }$ may or may not be invertible.

Result 2 uses $\mathbb {E}[R_t C_t-C_t C_t'v^*]=0$ from the definition of $v^*$ and $(\theta _t-v^*)=\mathcal {O}\left( \max _i \begin{vmatrix} 1 - \alpha \rho _{C,i}\end{vmatrix}^t\right)$ from our analysis of the mean of $V_t$ to get that

$$\begin{aligned} \text {Var}\left[ R_t C_t-C_t C_t' \theta _t\right] =&\text {Var}\left[ R_t C_t-C_t C_t' v^* - C_t C_t'(\theta _t-v^*)\right] \\ =&\text {Var}\left[ R_t C_t-C_t C_t' v^*\right] +\,\text {Var}\left( C_t C_t'(\theta _t-v^*)\right) \\&-\,\mathrm {Cov}\left( R_t C_t-C_t C_t' v^*, C_t C_t'(\theta _t-v^*)\right) \\&-\,\mathrm {Cov}\left( C_t C_t'(\theta _t-v^*),R_t C_t-C_t C_t' v^*\right) \\ =&\text {Var}\left[ R_t C_t-C_t C_t' v^*\right] \\&+\,\mathbb {E}\left[ (C_t C_t'-\mathbb {E}[C_t C_t'])(\theta _t-v^*)(\theta _t-v^*)'(C_t C_t'-\mathbb {E}[C_t C_t'])\right] \\&-\,\mathbb {E}\left[ \left( R_t C_t-C_t C_t' v^*\right) (\theta _t-v^*)'\left( C_t C_t'-\mathbb {E}[C_t C_t']\right) \right] \\&-\,\mathbb {E}\left[ \left( C_t C_t'-\mathbb {E}[C_t C_t']\right) (\theta _t-v^*)\left( R_t C_t-C_t C_t' v^*\right) '\right] \\ =&\text {Var}\left[ R_t C_t-C_t C_t' v^*\right] \\&+ \mathcal {O}\left( \max \{\max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^t, \max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^{2t} \}\right) . \end{aligned}$$

If $\max _{i}\begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix} < 1$ (an assumption we later need to ensure convergence of the covariance matrix), then we can drop $\max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^{2t}$ from the last term since then $\begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^2< \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix} < 1$ for all i. Hence, we arrive at the desired result:

$$\begin{aligned} \text {Var}\left[ R_t C_t-C_t C_t' \theta _t\right] =&\text {Var}\left[ R_t C_t-C_t C_t' v^*\right] + \mathcal {O}\left( \max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^t \right) \end{aligned}$$

Result 3 follows from noting that symmetric matrix $I - \alpha K_{\alpha }$ shares eigenvectors with symmetric matrix $K_{\alpha }$ and is also positive semi-definite, since

$$\begin{aligned} I - \alpha K_{\alpha }= & {} \mathbb {E}\left[ (I-\alpha C_t C_t')\otimes (I-\alpha C_t C_t') \right] . \end{aligned}$$

Thus, eigenvalues of $I - \alpha K_{\alpha }$ are nonnegative and are all of the form $1-\alpha \rho$ for each eigenvalue $\rho$ of $K_{\alpha }.$ Consequently,

$$\begin{aligned} I - \alpha K_{\alpha } = \mathcal {O}\left( 1-\alpha \rho _{K,\alpha }\right) \end{aligned}$$

where $\rho _{K,\alpha }$ is the smallest eigenvalue of $K_{\alpha }$.

Putting these last three results together when $K_{\alpha }^{-1}$ exists and $\max _i\begin{vmatrix}1-\alpha \rho _{C,i}\end{vmatrix} < 1$ leads to

$$\begin{aligned} \mathtt {vec}(\varSigma _{t+1})&= \alpha ^2 \sum _{s=1}^{t} (I-\alpha K_{\alpha })^{t-s} \mathtt {vec}\left( \text {Var}\left[ R_s C_s-C_s C_s' v^*\right] \right. \\&\quad \left. + \mathcal {O}\left( \max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^s\right) \right) \\&= \alpha K_{\alpha }^{-1} (I-(I-\alpha K_{\alpha })^t)\,\mathtt {vec}\left( \text {Var}\left[ R_t C_t-C_t C_t' v^*\right] \right) \\&\quad +\, \alpha ^2 \sum _{s=1}^{t} (I-\alpha K_{\alpha })^{t-s} \mathcal {O}\left( \max _{i} \begin{vmatrix} 1-\alpha \rho _{C,i} \end{vmatrix}^s\right) \\&= \alpha K_{\alpha }^{-1}\left( I-\mathcal {O}\left( \left( 1-\alpha \rho _{K,\alpha }\right) ^t\right) \right) \,\mathtt {vec}\left( \text {Var}\left[ R_t C_t-C_t C_t' v^*\right] \right) \\&\quad +\, \alpha ^2 \sum _{s=1}^{t} \mathcal {O}\left( \max \{\max _i \begin{vmatrix} 1-\alpha \rho _{C,i}\end{vmatrix}^t,\,\left( 1-\alpha \rho _{K,\alpha }\right) ^t\}\right) \\&=\alpha K_{\alpha }^{-1} \mathtt {vec}\left( \text {Var}\left[ R_tC_t-C_tC_t'v^*\right] \right) \\&\quad +\, \mathcal {O}\left( t\max \{\max _i \begin{vmatrix} 1-\alpha \rho _{C,i}\end{vmatrix}^t,\,\left( 1-\alpha \rho _{K,\alpha }\right) ^t\}\right) . \end{aligned}$$

In conclusion, if $\rho _{K,\alpha } > 0$ (which implies $K_{\alpha }^{-1}$ exists) and $\max _i\begin{vmatrix}1-\alpha \rho _{C,i}\end{vmatrix} < 1$, then the covariance matrix $\varSigma _t$ converges to a long-term covariance matrix that scales linearly with $\alpha$.

B Adding a Decay Term

Associative models of learning often incorporate a decay term into model updates which shrinks associative strengths to zero. While usually only absent cues have their associative strengths decay (c.f., Niv et al. 2015), to simplify subsequent analysis, we consider an R–W model modified to decay all associative strengths to zero using the following update:

$$\begin{aligned} V_{t+1} = V_t + \alpha C_t \left( R_t - C_t'V_t\right) - \alpha \eta V_t, \end{aligned}$$

for some decay parameter $\eta$ with $1>\alpha \eta \ge 0$. Under Assumptions 1– 2, this version of the R–W model is stochastic gradient descent applied to the following objective function:

$$\begin{aligned} \mathbb {E}\left[ \left( R_t - C_t'V_t\right) ^2 +\eta V_t'V_t\right] , \end{aligned}$$

which tries to find associative strengths that minimize a combination of square prediction error and its square norm. Thus, this R–W model with decay addresses the L2 regularized version of the original prediction problem addressed by the R–W model. A solution $v^*_{\eta }$ to this regularized objective function satisfies:

$$\begin{aligned} \mathbb {E}\left[ \eta I + C_t C_t'\right] v^*_{\eta } = \mathbb {E}\left[ R_t C_t \right] . \end{aligned}$$

Following similar arguments in the R–W model, we can recover an expression for the expectation $\theta _t$ and covariance matrix $\varSigma _t$ of associative strength on trial t for this R–W model with decay. Expectation of the update rule gives:

$$\begin{aligned} \theta _{t+1}= & {} \theta _t + \mathbb {E}\left[ \alpha C_t\left( R_t - C_t' V_t\right) \right] - \alpha \eta \theta _t\\= & {} v^*_{\eta } + \left( I - \alpha (\eta I + \mathbb {E}\left[ C_tC_t' \right] )\right) (\theta _t-v^*_{\eta }), \end{aligned}$$

and thus:

$$\begin{aligned} \theta _{t+1} = v^*_{\eta } + \left( (I-\alpha (\eta I + \mathbb {E}\left[ C_tC_t' \right] )\right) ^t (\theta _1-v^*_{\eta }). \end{aligned}$$

In other words, $\theta _{t+1}$ approaches $v^*_{\eta }$ at a rate that depends on

$$\begin{aligned} \begin{Vmatrix}I-\alpha \left( \eta I + \mathbb {E}\left[ C_tC_t' \right] \right) \end{Vmatrix} = \max _{i} \begin{vmatrix} 1-\alpha (\eta + \rho _{C,i}) \end{vmatrix} \end{aligned}$$

provided $\max _{i} \begin{vmatrix} 1-\alpha (\eta + \rho _{C,i}) \end{vmatrix} < 1$. As before, $\rho _{C,i}$ denotes an eigenvalue of $\mathbb {E}[C_t C_t']$. For sufficiently small $\alpha$ and $\eta$, the decay term can speed up convergence in the mean.

For the covariance matrix $\varSigma _{t+1}$, we have

$$\begin{aligned} \varSigma _{t+1}= & {} \text {Var}\left[ V_t + \alpha C_t R_t - \alpha C_tC_t'V_t - \eta V_t \right] \\= & {} \text {Var}\left[ V_t + \alpha \left( R_t C_t - (\eta I+ C_tC_t')\theta _t\right) - \alpha (\eta I + C_tC_t')\left( V_t-\theta _t\right) \right] , \end{aligned}$$

which expands to

$$\begin{aligned} \varSigma _{t+1}= & {} \varSigma _t - \alpha \varSigma _t\mathbb {E}\left[ \eta I + C_tC_t'\right] - \alpha \mathbb {E}\left[ \eta I + C_tC_t'\right] \varSigma _t\\&+ \alpha ^2 \mathbb {E}\left[ (\eta I + C_tC_t')\varSigma _t (\eta I + C_tC_t')\right] \\&+ \, \alpha ^2 \text {Var}\left[ R_t C_t - (\eta I +C_tC_t')\theta _t\right] . \end{aligned}$$

Letting

$$\begin{aligned} K_{\alpha ,\eta } = \mathbb {E}\left[ I\otimes (\eta I + C_t C_t')+(\eta I + C_t C_t')\otimes I -\alpha (\eta I + C_t)\otimes (\eta I + C_t C_t') \right] \end{aligned}$$

with minimum eigenvalue $\rho _{K,\alpha ,\eta }$, then vectorization yields

$$\begin{aligned} \mathtt {vec}(\varSigma _{t+1})&= (I-\alpha K_{\alpha ,\eta })\mathtt {vec}(\varSigma _t) + \alpha ^2 \mathtt {vec}(\text {Var}\left[ R_t C_t-(\eta I + C_t C_t') \theta _t\right] ) \\&=\alpha ^2 \sum _{s=1}^{t} (I-\alpha K_{\alpha ,\eta })^{t-s} \mathtt {vec}(\text {Var}\left[ R_s C_s-(\eta I + C_s C_s') \theta _s\right] ). \end{aligned}$$

Using similar arguments before, we have that when $K_{\alpha ,\eta }^{-1}$ exists:

$$\begin{aligned} \mathtt {vec}(\varSigma _{t+1})&= \alpha K_{\alpha ,\eta }^{-1} \mathtt {vec}\left( \text {Var}\left[ R_tC_t-(\eta I + C_tC_t')v^*_{\eta }\right] \right) \\&\quad +\, \mathcal {O}\left( t\max \{\max _i \begin{vmatrix} 1-\alpha (\eta + \rho _{C,i})\end{vmatrix}^t,\,\left( 1-\alpha \rho _{K,\alpha ,\eta }\right) ^t\}\right) \end{aligned}$$

Note when $\eta =0$, then $v^*_0 = v^*$, $K_{\alpha ,0}=K_{\alpha }$, and $\rho _{K,\alpha ,0}=\rho _{K,\alpha }$, and hence, we recover our expressions for a R–W model without decay.

Based on these derived expressions, we reproduced figures from the main text for the R–W model with decay. Comparing Figs. 11 to 3 from the main text, the R–W model with decay leads to faster convergence of the mean and variance of associative strengths over the R–W model without decay for sufficiently small $\alpha$ and N, but does not expand the region of $\alpha ,N$ values with convergent means and variances. Decay has little effect on long-term variance.

Fig. 11

Statistical properties of associative strength $V_t$ for the R–W model with decay in the haystack problem. Dashed line marks the boundary between regions where long-term variance exists and does not exist. Decay parameter $\eta$ was set to 0.5 (Color figure online)

Fig. 12

Statistical properties of associative strength $V_t$ for the Sparse R–W model with decay in the haystack problem. The dashed line is $k=N$, which represents the R–W model with decay. Decay parameter $\eta$ was set to 0.5, and learning rate $\alpha$ was 0.2 (Color figure online)

Fig. 13

Estimated error in prediction for R–W, Leong, SAR-WI, and SAR-WI with decay on four learning tasks. All models have $\alpha = 0.1$, $N=20$, the Leong model was run with $\beta =10$, the two SAR-WI models was run with an attention bandwidth k of 5, and the SAR-WI model with decay was run with $\eta =0.5$. Expectations were estimated by averaging $4 \times 10^4$ simulations of the task (Color figure online)

We can extend the R–W model with decay to the sparse attention setting:

$$\begin{aligned} V_{t+1} = V_t + \alpha X_t C_t \left( R_t - (X_t C_t)'V_t\right) - \alpha \eta V_t, \end{aligned}$$

Notice that this version of association decay applies to all associations equally, and thus represents a generalized decay or forgetfulness likely due to some general biological or cognitive process which decays or rewrites on top of all associations. Since attended variables are also updated, computations are not sparse. The net effect of this model would be that consecutively unattended associations would decay, while attended ones may sustain associations. As before, the derivations for mean and variance hold with $C_t$ replaced with $X_t C_t$. Comparing Figs. 12 to 4 from the main text, the Sparse R–W model with decay leads to faster convergence of the mean and variance of associative strengths over the Sparse R–W model without decay, but leads to worse long-term variance.

We can also add a general decay term to the SAR-WI framework:

$$\begin{aligned} V_{t+1}&= V_t + \alpha X_t C^0_t(R^0_t -(X_tC^0_t)'V_t) - \eta \alpha V_t, \\ \bar{C}_{t+1}&= \bar{C}_t + \alpha X_t(C_t-\bar{C}_t) - \eta \alpha \bar{C}_t \\ \bar{R}_{t+1}&= \bar{R}_t + \alpha (R_t-\bar{R}_t) - \eta \alpha \bar{R}_t \end{aligned}$$

where as before

$$\begin{aligned} C^0_t&= C_t - \bar{C_t}\\ R^0_t&= R_t - \bar{R_t}. \end{aligned}$$

Figure 13 adds the SAR-WI model with decay to the comparison of models on four tasks. On these tasks and for these sets of parameters, the SAR-WI model with decay performs worse than the SAR-WI model without decay.

C Selecting Cues Without a Fixed Attention Bandwidth

When satisfying the fixed attention bandwidth constraint is less important, or when it is natural to consider the number of attended cues changing in response to their associations, another approach is to sample cues independently from each other:

$$\begin{aligned} P(X_{t,i,i}) = \frac{ V_{t,i}^2 \bar{C}_{t,i}(1-\bar{C}_{t,i}) + \epsilon }{D + \epsilon }, \end{aligned}$$

(12)

for some $\epsilon >0$ and some normalization constant D. If rewards are Bernoulli, then one could use $D = \bar{R}_t(1-\bar{R}_t)$, since independent Bernoulli cue between trials ensures

$$\begin{aligned} \sum _{i} (v_{i}^*)^2 \mathbb {E}[C_{t,i}] (1-\mathbb {E}[C_{t,i}]) \le \mathbb {E}[R_t](1-\mathbb {E}[R_t]). \end{aligned}$$