SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} f _ {c} \left(\alpha \mathbf {I} _ {N} + \mathbf {W} \mathrm {d i a g} (\mathbf {g}) \mathbf {W} ^ {\top}\right) = \mathbb {E} _ {\mathbf {s} \sim p (\mathbf {s} | c)} \left[ \max  _ {\mathbf {r} \in \mathbb {R} ^ {N}} \mathrm {T r} \left(2 \mathbf {r s} ^ {\top} - \alpha \mathbf {r r} ^ {\top} - \mathbf {W} \mathrm {d i a g} (\mathbf {g}) \mathbf {W} ^ {\top} \mathbf {r r} ^ {\top}\right) \right] \\ + \alpha N + \operatorname {T r} (\mathbf {W} \operatorname {d i a g} (\mathbf {g}) \mathbf {W} ^ {\top}) \\ = \mathbb {E} _ {\mathbf {s} \sim p (\mathbf {s} | c)} \left[ \max  _ {\mathbf {r} \in \mathbb {R} ^ {N}} \ell (\mathbf {W}, \mathbf {g}, \mathbf {r}, \mathbf {s}) \right] + \alpha N. \\ \end{array}
$$
Substituting into equation 3 and dropping the constant  $\alpha N$  term results in equation 5.
C Offline multi-timescale adaptive whitening algorithm
We consider an algorithm where we directly optimize the objective in equation 3. In this case, we assume that the input to the algorithm is a sequence of covariance matrices  $\mathbf{C}_{ss}(1),\mathbf{C}_{ss}(2),\ldots$  Within each context  $c = 1,2,\dots$  , we take  $J\geq 1$  concurrent gradient descent steps with respect to  $\mathbf{g}$  and  $\mathbf{W}$  ..
$$
\Delta \mathbf {g} = - \eta_ {g} \operatorname {d i a g} \left[ \mathbf {W} ^ {\top} \nabla_ {\mathbf {M}} f _ {c} (\alpha \mathbf {I} _ {N} + \mathbf {W} \operatorname {d i a g} (\mathbf {g}) \mathbf {W} ^ {\top}) \mathbf {W} \right],
$$
$$
\Delta \mathbf {W} = - \eta_ {w} \nabla_ {\mathbf {M}} f _ {c} (\alpha \mathbf {I} _ {N} + \mathbf {W} \operatorname {d i a g} (\mathbf {g}) \mathbf {W} ^ {\top}) \mathbf {W} \operatorname {d i a g} (\mathbf {g}),
$$
where we assume  $\eta_g \gg \eta_w > 0$  as in Algorithm 1 and the gradient of  $f_c(\mathbf{M})$  with respect to  $\mathbf{M}$  is given in equation 10. These updates for  $\mathbf{g}$  and  $\mathbf{W}$  can also be obtained by averaging the corresponding updates in equations 7 and 8 over the conditional distribution  $p(\mathbf{s}|c)$ . This results in Algorithm 2.
Algorithm 2: Offline multi-timescale adaptive whitening
1: Input: Covariance matrices  $\mathbf{C}_{ss}(1),\mathbf{C}_{ss}(2),\ldots$   
2: Initialize:  $\mathbf{W}\in \mathbb{R}^{N\times K};\mathbf{g}\in \mathbb{R}^{K};\alpha \geq 0;J\geq 1;\eta_{g}\gg \eta_{w} > 0$   
3: for  $c = 1,2,\dots$  do  
4: for  $j = 1,\dots ,J$  do  
5:  $\mathbf{M}\gets \alpha \mathbf{I}_N + \mathbf{W}\mathrm{diag}(\mathbf{g})\mathbf{W}^\top$   
6:  $\mathbf{g}\gets \mathbf{g} - \eta_g\mathrm{diag}[\mathbf{W}^\top \nabla_\mathbf{M}f_c(\mathbf{M})\mathbf{W}]$   
7:  $\mathbf{W}\gets \mathbf{W} - \eta_w\nabla_\mathbf{M}f_c(\mathbf{M})\mathbf{W}\mathrm{diag}(\mathbf{g})$   
8: end for  
9: end for
D Adaptive whitening of natural images
In this section, we elaborate on the converged structure of  $\mathbf{W}_T$  using natural image patches. To better visualize the relationship between the learned columns of  $\mathbf{W}$  and sinusoidal basis functions (e.g.
B)
C)
Figure 5: Control experiment accompanying Sec. 5.2. A)  $\mathbf{W}_T$  learned from natural image patches. B) Basis vectors from A displayed as line plots, compared to the 1D DCT, and principal components of  $\mathbb{E}_{c\sim p(c)}[\mathbf{C}_{ss}(c)]$ . C) Control condition.  $\mathbf{W}_T$  learned from spectrally-matched image patches with random eigenvectors.
DCT), we focus on 1-dimensional image patches (rows of pixels). The results are similar with 2D image patches.
It is well known that eigenvectors of natural images are well-approximated by sinusoidal basis functions [e.g. the DCT; 35; 36]. Using the same images from the main text [34], we generated 56 contexts by sampling  $16 \times 1$  pixel patches from separate images, with 2E4 samples each. We train Algorithm 2 with  $K = N = 16$ ,  $\eta_w = 5\mathrm{E} - 2$ , and random  $\mathbf{W}_0 \in O(16)$  on a training set of  $\mathbf{X}$  of the images, presented uniformly at random  $T = 1\mathrm{E}5$  times. Fig 5A,B shows that  $\mathbf{W}_T$  approximates the principal components of the aggregated context-dependent covariance,  $\mathbb{E}_{c \sim p(c)}[\mathbf{C}_{ss}(c)]$ , which are closely aligned with the DCT. To show that this structure is inherent in the spatial statistics of natural images, we generated control contexts,  $\mathbf{C}_{ss}(c)$ , by forming covariance matrices with matching eigenspectra, but each with random and distinct eigenvectors. This destroys the structure induced by natural image statistics. Consequently, the learned vectors in  $\mathbf{W}_T$  are no longer sinusoidal (Fig 5C). As a result, whitening error with  $\mathbf{W}_T$  is much higher on the training set, with  $0.3 \pm 0.02$  error (mean  $\pm$  standard error over 10 random initializations; Eq. 9) on natural image contexts and  $2.7 \pm 0.1$  on the control contexts. While for the natural images, a basis approximating the DCT was sufficient to adaptively whiten all contexts in the ensemble, this is not the case for the generated control contexts.
Finally, we find that as  $K$  increases from  $K = 1$  to  $K = 16$ , the basis vectors in  $\mathbf{W}_T$  progressively learn higher frequency components of the DCT (Fig. 6). This is a sensible solution, due to the  $\ell_2$  reconstruction error of our objective, and the  $1/f$  spectral content of natural image statistics. With more flexibility, as  $K$  increases past  $N$  (i.e. the overcomplete regime), the network continues to improve its whitening error (Fig. 7A) by learning a basis,  $\mathbf{W}_T$ , that can account for within-context information that is insufficiently captured by the DCT (Fig. 7B). Taken together, our model successfully learns a basis  $\mathbf{W}_T$  that exploits the spatial structure present in natural images.
E Modifications for increased biological realism
In this section, we modify Algorithm 1 to be more biologically realistic.
E.1 Enforcing unit norm basis vectors
In our algorithm, there is no constraint on the magnitude of the column vectors of  $\mathbf{W}$ . We can enforce a unit norm (here measured using the Euclidean norm) constraint by adding Lagrange multipliers to the objective in equation 3:
$$
\min  _ {\mathbf {W} \in \mathbb {R} ^ {N \times K}} \max  _ {\mathbf {m} \in \mathbb {R} ^ {K}} \mathbb {E} _ {c \sim p (c)} \left[ \min  _ {\mathbf {g} \in \mathbb {R} ^ {K}} \mathbb {E} _ {\mathbf {s} \sim p (\mathbf {s} | c)} [ g (\mathbf {W}, \mathbf {g}, \mathbf {r}, \mathbf {s}) ] \right], \tag {11}
$$
Figure 6: As  $K$  increases, columns of  $\mathbf{W}$  progressively learn higher frequency components of the DCT.
Figure 7: A) Error on training and test set as a function of  $K$ . B) In the overcomplete regime, the network learns basis vectors  $\mathbf{W}_T$  that further improve the error compared with the  $K \leq N$  regime.
$$
g (\mathbf {W}, \mathbf {g}, \mathbf {r}, \mathbf {s}) = \ell (\mathbf {W}, \mathbf {g}, \mathbf {r}, \mathbf {s}) + \sum_ {i = 1} ^ {K} m _ {i} \left(\left\| \mathbf {w} _ {i} \right\| ^ {2} - 1\right).
$$
Taking partial derivatives with respect to  $\mathbf{w}_i$  and  $\mathbf{m}_i$  results in the updates:
$$
\Delta \mathbf {w} _ {i} = \eta_ {w} \left(n _ {i} \mathbf {r} - \left(g _ {i} + m _ {i}\right) \mathbf {w} _ {i}\right)
$$
$$
\Delta m _ {i} = \left\| \mathbf {w} _ {i} \right\| ^ {2} - 1.
$$
Furthermore, since the weights are constrained to have unit norm, we can replace  $\| \mathbf{w}_i\|^2$  with 1 in the gain update:
$$
\Delta g _ {i} = \eta_ {g} \left(z _ {i} ^ {2} - 1\right).
$$
E.2 Decoupling the feedforward and feedback weights
We replace the primary neuron-to-interneuron weight matrix  $\mathbf{W}^{\top}$  (resp. interneuron-to-primary neuron weight matrix  $-\mathbf{W}$ ) with  $\mathbf{W}_{rn}$  (resp.  $-\mathbf{W}_{nr}$ ). In this case, the update rules are
$$
\mathbf {W} _ {r n} \leftarrow \mathbf {W} _ {r n} + \eta_ {w} \left(\mathbf {n} _ {t} \mathbf {r} _ {t} ^ {\top} - \operatorname {d i a g} (\mathbf {g} + \mathbf {m}) \mathbf {W} _ {r n}\right)
$$
$$
\mathbf {W} _ {n r} \leftarrow \mathbf {W} _ {n r} + \eta_ {w} \left(\mathbf {r} _ {t} \mathbf {n} _ {t} ^ {\top} - \mathbf {W} _ {n r} \mathrm {d i a g} (\mathbf {g} + \mathbf {m})\right).
$$
Figure 8: With stricter constraints on biological realism (Algorithm 3), the model succeeds to whiten synthetic data setup from Section 5.1 when the column vectors of  $\mathbf{V}$  are chosen to be nonnegative. A) Error decreases as training progresses. B) Online algorithm performance before and after training. C) The learned basis (row vectors of  $\mathbf{W}_{nr}$  after learning, red) is well-aligned with the generative basis (dashed black) compared to initialization (row vectors of  $\mathbf{W}_{nr}$  at initialization, blue). D) Asymmetric feedforward weights,  $\mathbf{W}_{rn}$ , and feedback weights,  $-\mathbf{W}_{nr}$ , converge to being symmetric.
Let  $\mathbf{W}_{rn,t}$  and  $\mathbf{W}_{nr,t}$  denote the values of the weights  $\mathbf{W}_{rn}$  and  $\mathbf{W}_{nr}$ , respectively, after  $t = 0,1,\ldots$  iterates. Then for all  $t = 0,1,\ldots$ ,
$$
\mathbf {W} _ {r n, t} ^ {\top} - \mathbf {W} _ {n r, t} = \left(\mathbf {W} _ {r n, 0} ^ {\top} - \mathbf {W} _ {n r, 0}\right) \left(\mathbf {I} _ {N} - \eta_ {w} \mathrm {d i a g} (\mathbf {g} + \mathbf {m})\right) ^ {t}.
$$
Thus, if  $g_{i} + m_{i} \in (0, 2\eta_{w}^{-1})$  for all  $i$  (e.g., by enforcing non-negative  $g_{i}$ ,  $m_{i}$  and choosing  $\eta_{w} > 0$  sufficiently small), then the difference decays exponentially in  $t$  and the feedforward and feedback weights are asymptotically symmetric. This result is similar to that of Kolen and Pollack [63], where the authors trained a network with weight decay to show that weight symmetry naturally arises to induce bidirectional function synchronization during learning.
E.3 Sign-constraining the synaptic weights and gains
The synaptic weight matrix  $\mathbf{W}$  and gains vector  $\mathbf{g}$  are not sign-constrained in Algorithm 1, which is not consistent with biological evidence. We can modify the algorithm to enforce the sign constraints by rectifying the weights and gains at each step. Here  $[\cdot]_{+}$  denote the elementwise rectification operation. This results in the updates
$$
\mathbf {g} \leftarrow \left[ \mathbf {g} + \eta_ {g} (\mathbf {z} \circ \mathbf {z} - \mathbf {1}) \right] _ {+}
$$
$$
\mathbf {W} _ {r n} \leftarrow \left[ \mathbf {W} _ {r n} + \eta_ {w} \left(\mathbf {n} _ {t} \mathbf {r} _ {t} ^ {\top} - \operatorname {d i a g} (\mathbf {g} + \mathbf {m}) \mathbf {W} _ {r n}\right) \right] _ {+}
$$
$$
\left. \right. \mathbf {W} _ {n r} \leftarrow \left[ \mathbf {W} _ {n r} + \eta_ {w} \left(\mathbf {r} _ {t} \mathbf {n} _ {t} ^ {\top} - \mathbf {W} _ {n r} \operatorname {d i a g} (\mathbf {g} + \mathbf {m})\right)\right] _ {+}.
$$
E.4 Online algorithm with improved biological realism
Combining these modifications yields our more biologically realistic multi-timescale online algorithm, Algorithm 3.
Algorithm 3: Biologically realistic multi-timescale adaptive whitening
1: Input:  $\mathbf{s}_1, \mathbf{s}_2, \dots \in \mathbb{R}^N$   
2: Initialize:  $\mathbf{W}_{nr} \in \mathbb{R}_+^{N \times K}$ ;  $\mathbf{W}_{rn} \in \mathbb{R}_+^{K \times N}$ ;  $\mathbf{m}, \mathbf{g} \in \mathbb{R}_+^K$ ;  $\eta_r, \eta_m > 0$ ;  $\eta_g \gg \eta_w > 0$   
3: for  $t = 1, 2, \ldots$  do  
4:  $\mathbf{r}_t \gets \mathbf{0}$   
5: while not converged do  
6:  $\mathbf{z}_t \gets \mathbf{W}_{rn} \mathbf{r}_t$ ; // interneuron inputs  
7:  $\mathbf{n}_t \gets \mathbf{g} \circ \mathbf{z}_t$ ; // gain-modulated interneuron outputs  
8:  $\mathbf{r}_t \gets \mathbf{r}_t + \eta_r (\mathbf{s}_t - \mathbf{W}_{nr} \mathbf{n}_t - \alpha \mathbf{r}_t)$ ; // recurrent neural dynamics  
9: end while  
10:  $\mathbf{m} \gets [\mathbf{m} + \eta_m (\mathrm{diag}(\mathbf{W}_{rn} \mathbf{W}_{nr}) - 1)]_+$ ; // weight normalization update  
11:  $\mathbf{g} \gets [\mathbf{g} + \eta_g (\mathbf{z}_t \circ \mathbf{z}_t - 1)]_+$ ; // gains update  
12:  $\mathbf{W}_{rn} \gets [\mathbf{W}_{rn} + \eta_w (\mathbf{n}_t \mathbf{r}_t^\top - \mathrm{diag}(\mathbf{g} + \mathbf{m}) \mathbf{W}_{rn})]_+$ ; // synaptic weights update  
13:  $\mathbf{W}_{nr} \gets [\mathbf{W}_{nr} + \eta_w (\mathbf{r}_t \mathbf{n}_t^\top - \mathbf{W}_{nr} \mathrm{diag}(\mathbf{g} + \mathbf{m}))]_+$   
14: end for
We test Algorithm 3 on a similar synthetic data setup to what was used in section 5.1, except that we sample the column vectors of  $\mathbf{V}$  from the intersection of the unit circle with the nonnegative quadrant, Figure 8.