SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Multi-timescale online algorithm. Combining the neural dynamics, gain modulation and synaptic plasticity yields our online multi-timescale adaptive whitening algorithm, which we express in vector-matrix form with  $\circ$  denoting the Hadamard (elementwise) product of two vectors:
Algorithm 1: Multi-timescale adaptive whitening via synaptic plasticity and gain modulation
1: Input:  $\mathbf{s}_1, \mathbf{s}_2, \dots \in \mathbb{R}^N$   
2: Initialize:  $\mathbf{W} \in \mathbb{R}^{N \times K}$ ;  $\mathbf{g} \in \mathbb{R}^K$ ;  $\eta_r > 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}^\top \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} \mathbf{n}_t - \alpha \mathbf{r}_t)$ ; // recurrent neural dynamics  
9: end while  
10:  $\mathbf{g} \gets \mathbf{g} + \eta_g (\mathbf{z}_t \circ \mathbf{z}_t - \mathrm{diag}(\mathbf{W}^\top \mathbf{W}))$ ; // gains update  
11:  $\mathbf{W} \gets \mathbf{W} + \eta_w (\mathbf{r}_t \mathbf{n}_t^\top - \mathbf{W} \mathrm{diag}(\mathbf{g}))$ ; // synaptic weights update  
12: end for
Alg. 1 is naturally viewed as a unification and generalization of previously proposed neural circuit models for adaptation. When  $\alpha = 0$  and the gains  $\mathbf{g}$  are constant (e.g.,  $\eta_{g} = 0$ ) and identically equal to the vector of ones  $\mathbf{1}$  (so that  $\mathbf{n}_t = \mathbf{z}_t$ ), we recover the synaptic plasticity algorithm from [15]. Similarly, when  $\alpha = 1$  and the synaptic weights  $\mathbf{W}$  are fixed (e.g.,  $\eta_w = 0$ ), we recover the gain modulation algorithm from [27].
Figure 3: Adaptive whitening of a synthetic dataset with  $N = 2$ ,  $\eta_w = 1\mathrm{E} - 5$ ,  $\eta_g = 5\mathrm{E} - 2$ . A) Covariance ellipses (orange) of 4 out of 64 synthesized contexts. Black dashed lines are axes corresponding to the column vectors of V. The unit circle is shown in green. Since the column vectors of V are not orthogonal, these covariance matrices do not share a common set of eigenvectors (orange lines). B) Whitening error with a moving average window spanning the last 10 contexts. C) Error at each stimulus presentation within five different contexts (gray panels), presented with  $\mathbf{W}_0$ , or  $\mathbf{W}_T$ . D) Column vectors of  $\mathbf{W}_0$ ,  $\mathbf{W}_T$ , V (each axis corresponds to the span of one column vector in  $\mathbb{R}^2$ ). E) Smoothed distributions of error (in Frobenius norm) between  $\hat{\mathbf{W}}$  and V across 250 random initializations of  $\mathbf{W}_0$ .
5 Numerical experiments
We test Alg. 1 on stimuli  $\mathbf{s}_1, \mathbf{s}_2, \ldots$  drawn from slowly fluctuating latent contexts  $c_1, c_2, \ldots$ ; that is,  $\mathbf{s}_t \sim p(\mathbf{s} | c_t)$  and  $c_t = c_{t-1}$  with high probability. To measure performance, we evaluate the operator norm on the difference between the expected response covariance and the identity matrix:
$$
\operatorname {E r r o r} (t) = \left\| \mathbf {M} _ {t} ^ {- 1} \mathbf {C} _ {s s} \left(c _ {t}\right) \mathbf {M} _ {t} ^ {- 1} - \mathbf {I} _ {N} \right\| _ {\mathrm {o p}}, \quad \mathbf {M} _ {t} := \alpha \mathbf {I} _ {N} + \mathbf {W} _ {t} \operatorname {d i a g} (\mathbf {g}) \mathbf {W} _ {t} ^ {\top}. \tag {9}
$$
Geometrically, this "worst-case" error measures the maximal Euclidean distance between the ellipsoid corresponding to  $\mathbf{M}_t^{-1}\mathbf{C}_{ss}(c_t)\mathbf{M}_t^{-1}$  and the  $(N - 1)$ -sphere along all possible axes. To compare two synaptic weight matrices  $\mathbf{A},\mathbf{B}\in \mathbb{R}^{N\times K}$ , we evaluate  $\| \hat{\mathbf{A}} -\mathbf{B}\| _F$ , where  $\hat{\mathbf{A}} = \mathbf{AP}$  and  $\mathbf{P}$  is the permutation matrix (with possible sign flips) that minimizes the error.
5.1 Synthetic dataset
To validate our model, we first consider a 2-dimensional synthetic dataset in which an optimal solution is known. Suppose that each context-dependent inverse whitening matrix is of the form  $\mathbf{M}_c = \mathbf{I}_N + \mathbf{V}\boldsymbol{\Lambda}(c)\mathbf{V}^\top$ , where  $\mathbf{V}$  is a fixed  $2 \times 2$  matrix and  $\boldsymbol{\Lambda}(c) = \mathrm{diag}(\lambda_1(c), \lambda_2(c))$  is a context-dependent diagonal matrix. Then, in the case  $\alpha = 1$  and  $K = 2$ , an optimal solution of the objective in Eq. 3 is when the column vectors of  $\mathbf{W}$  align with the column vectors of  $\mathbf{V}$ .
To generate this dataset, we chose the column vectors of  $\mathbf{V}$  uniformly from the unit circle, so they are not generally orthogonal. For each context  $c = 1,\dots ,64$  we assume the diagonal entries of  $\Lambda (c)$  are sparse and i.i.d.: with probability  $1 / 2$ ,  $\lambda_{i}(c)$  is set to zero and with probability  $1 / 2$ ,  $\lambda_{i}(c)$  is chosen uniformly from the interval [0, 4]. Example covariance matrices from different contexts are shown in Fig. 3A (note that they do not share a common eigen-decomposition). Finally, for each context, we generate 1E3 i.i.d. samples with context-dependent distribution  $\mathbf{s} \sim \mathcal{N}(\mathbf{0},\mathbf{M}_c^2)$ .
We test Alg. 1 with  $\alpha = 1$ ,  $K = 2$ ,  $\eta_w = 1$ E-5, and  $\eta_g = 5$ E-2 on these sequences of synthetic inputs with the column vectors of  $\mathbf{W}_0$  chosen uniformly from the unit circle. The model successfully learns to whiten the different contexts, as indicated by the decreasing whitening error with the number of contexts presented (Fig. 3B). At the end of training, the synaptic weight matrix  $\mathbf{W}_T$  is optimized such that the circuit can adapt to changing contexts exclusively by adjusting its gains. This is evidenced by the fact that when the context changes, there is a brief spike in error as the gains adapt to the new context (Fig. 3C, red line). By contrast, the error remains high in many of the contexts when using
$^{2}$ Python code accompanying this study can be found at https://github.com/lyndond/multi-timescale-whitening.
Figure 4: Adaptive whitening of natural images. A) Examples of 2 out of 56 high-resolution images (top) with each image corresponding to a separate context. For each image,  $5 \times 5$  pixel patches are randomly sampled to generate context-dependent stimuli with covariance matrix  $\mathbf{C}_{ss}(c) \in \mathbb{S}_{++}^{25}$  (bottom). B) Mean error during training (Eq. 9) with  $K = N = 25$ . Shaded region is standard deviation over 2E3 random initializations  $\mathbf{W}_0 \in O(25)$ . C) Smoothed distributions of average adaptive whitening error over all 2E3 initializations. The red distribution corresponds to the error on the held-out images with fixed learned synapses  $\mathbf{W}_T$  and modulated gains  $\mathbf{g}$ . The blue (resp. green, purple) distribution corresponds to the same error, but tested on the training images (resp. with fixed gains equal to the average gains over the final 100 iterations, with fixed random synapses  $\mathbf{W}_0$ ). D) The learned weights (re-shaped columns of  $\mathbf{W}_T$ ) approximate orthogonal 2D sinusoids. E) Final error (after  $T = 5\mathrm{E4}$  iterations) as a function of number of interneurons  $K$ . Bars are standard deviations centered on the mean error at each  $K$ . The red horizontal line denotes average error when  $K = 0$  (in which case  $\mathbf{r} = \mathbf{s}$ ). F) Frobenius norm between the eigenbasis of  $\mathbb{E}_{c \sim p(c)}[\mathbf{C}_{ss}(c)]$  (i.e. across all contexts),  $\mathbf{Q}_{*}$ , with  $\mathbf{W}_T$ ,  $\mathbf{W}_0$ , and the eigenbasis of each individual context covariance,  $\mathbf{Q}(c)$ , when  $K = N = 25$ . See Appx. D for additional experiments.
the initial random synaptic weight matrix  $\mathbf{W}_0$  (Fig. 3C, blue line). In particular, the synapses learn (across contexts) an optimal configuration in the sense that the column vectors of  $\mathbf{W}$  learn to align with the column vectors of  $\mathbf{V}$  over the course of training (Fig. 3DE).
5.2 Natural images dataset
By hand-crafting a particular set of synaptic weights, Duong et al. [27] showed that their adaptive whitening network can approximately whiten a dataset of natural image patches with  $\mathcal{O}(N)$  gainmodulating interneurons instead of  $\mathcal{O}(N^2)$ . Here, we show that our model can exploit spatial structure across natural scenes to learn an optimal set of synaptic weights by testing our algorithm on 56 high-resolution natural images [34] (Fig. 4A, top). For each image, which corresponds to a separate context  $c$ ,  $5 \times 5$  pixel image patches are randomly sampled and vectorized to generate context-dependent samples  $\mathbf{s} \in \mathbb{R}^{25}$  with covariance matrix  $\mathbf{C}_{ss}(c) \in \mathbb{S}_{++}^{25}$  (Fig. 4A, bottom). We train our algorithm in the offline setting where we have direct access to the context-dependent covariance matrices (Appx. C, Alg. 2,  $\alpha = 1$ ,  $J = 50$ ,  $\eta_g = 5\mathrm{E - }1$ ,  $\eta_w = 5\mathrm{E - }2$ ) with  $K = N = 25$  and random  $\mathbf{W}_0 \in O(25)$  on a training set of 50 of the images, presented uniformly at random 1E3 total times. We find that the model successfully learns a basis that enables adaptive whitening across different visual contexts via gain modulation, as shown by the decreasing training error (Eq. 9) in Fig. 4B.
How does the network learn to leverage statistical structure that is consistent across contexts? We test the circuit with fixed synaptic weights  $\mathbf{W}_T$  and modulated (adaptive) gains  $\mathbf{g}$  on stimuli from the held-out images (Fig. 4C, red distribution shows the smoothed error over 2E3 random initializations  $\mathbf{W}_0$ ). The circuit performs as well on the held-out images as on the training images (Fig. 4C, red versus blue distributions). In addition, the circuit with learned synaptic weights  $\mathbf{W}_T$  and modulated gains  $\mathbf{g}$  outperforms the circuit with learned synaptic weights  $\mathbf{W}_T$  and fixed gains (Fig. 4C, green distribution), and significantly outperforms the circuit with random synaptic weights  $\mathbf{W}_0$  and modulated gains (Fig. 4C, purple distribution). Together, these results suggest that the circuit
learns features  $\mathbf{W}_T$  that enable the circuit to adaptively whiten across statistical contexts exclusively using gain modulation, and that gain modulation is crucial to the circuit's ability to adaptively whiten. In Fig. 4D, we visualize the learned filters (columns of  $\mathbf{W}_T$ ), and find that they are approximately equal to the 2D discrete cosine transform (DCT, Appx. D), an orthogonal basis that is known to approximate the eigenvectors of natural image patch covariances [35; 36].
To examine the dependence of performance on the number of interneurons  $K$ , we train the algorithm with  $K = 1, \ldots, 2N$  and report the final error in Fig. 4E. There is a steady drop in error as  $K$  ranges from 1 to  $N$ , at which point there is a (discontinuous) drop in error followed by a continued, but more gradual decay in both training and test images error as  $K$  ranges from  $N$  to  $2N$  (the overcomplete regime). To understand this behavior, note that the covariance matrices of image patches approximately share an eigen-decomposition [36]. To see this, let  $\mathbf{Q}(c)$  denote the orthogonal matrix of eigenvectors corresponding to the context-dependent covariance matrix  $\mathbf{C}_{ss}(c)$ . As shown in Fig. 4F (green histogram), there is a small, but non-negligible, difference between the eigenvectors  $\mathbf{Q}(c)$  and the eigenvectors  $\mathbf{Q}_*$  of the average covariance matrix  $\mathbb{E}_{c \sim p(c)}[\mathbf{C}_{ss}(c)]$ . When  $K = N$ , the column vectors of  $\mathbf{W}_T$  learn to align with  $\mathbf{Q}_*$  (as shown in Fig. 4F, blue histogram), and the circuit approximately adaptively whitens the context-dependent stimulus inputs via gain modulation. As  $K$  ranges from 1 to  $N$ ,  $\mathbf{W}_T$  progressively learns the eigenvectors of  $\mathbf{Q}_*$  (Appx. D). Since  $\mathbf{W}_T$  achieves a full set of eigenvectors at  $K = N$ , this results in a large drop in error when measured using the operator norm. Finally, as mentioned, there is a non-negligible difference between the eigenvectors  $\mathbf{Q}(c)$  and the eigenvectors  $\mathbf{Q}_*$ . Therefore, increasing the number of interneurons from  $N$  to  $2N$  allows the circuit to learn an overcomplete set of basis vectors  $\mathbf{W}_T$  to account for the small deviations between  $\mathbf{Q}(c)$  and  $\mathbf{Q}_*$ , resulting in improved whitening error (Appx. D).
We've derived an adaptive statistical whitening circuit from a novel objective (Eq. 3) in which the (inverse) whitening matrix is factorized into components that are optimized at different timescales. This model draws inspiration from the extensive neuroscience literature on rapid gain modulation [25] and long-term synaptic plasticity [17], and concretely proposes complementary roles for these computations: synaptic plasticity facilitates learning features that are invariant across statistical contexts while gain modulation facilitates adaptation within a statistical context. Experimental support for this will come from detailed understanding of natural sensory statistics across statistical contexts and estimates of (changes in) synaptic connectivity from wiring diagrams (e.g., [37]) or neural activities (e.g., [38]).
Our circuit uses local learning rules for the gain and synaptic weight updates that could potentially be implemented in low-power neuromorphic hardware [39] and incorporated into existing mechanistic models of neural circuits with whitened or decorrelated responses [40-45]. However, there are aspects of our circuit that are not biologically realistic. For example, we do not sign-constrain the gains or synaptic weight matrices, so our circuit can violate Dale's law. In addition, the feedforward synaptic weights  $\mathbf{W}^{\top}$  and feedback weights  $-\mathbf{W}$  are constrained to be symmetric. In Appx. E, we describe modifications of our model that can make it more biologically realistic. Additionally, while we focus on the potential joint function of gain modulation and synaptic plasticity in adaptation, short-term synaptic plasticity, which operates on similar timescales as gain modulation, has also been reported [18]. Theoretical studies suggest that short-term synaptic plasticity is useful in multi-timescale learning tasks [46-48] and it may also contribute to multi-timescale adaptive whitening. Ultimately, support for different adaptation mechanisms will be adjudicated by experimental observations.
Our work may also be relevant beyond the biological setting. Decorrelation and whitening transformations are common preprocessing steps in statistical and machine learning methods [49-53], and are useful for preventing representational collapse in recent self-supervised learning methods [54-57]. Therefore, our online multi-timescale algorithm may be useful for developing adaptive self-supervised learning algorithms. In addition, our work is related to the general problem of online meta-learning [58; 59]; that is, learning methods that can rapidly adapt to new tasks. Our solution—which is closely related to mechanisms of test-time feature gain modulation developed for machine learning models for denoising [28], compression [60; 61], and classification [62]—suggests a general approach to meta-learning inspired by neuroscience: structural properties of the tasks (contexts) are encoded in synaptic weights and adaptation to the current task (context) is achieved by adjusting the gains of individual neurons.
[21] Odelia Schwartz and Eero P Simoncelli. Natural signal statistics and sensory gain control. Nature Neuroscience, 4(8):819-825, 2001.
[22] Frances S Chance, Larry F Abbott, and Alex D Reyes. Gain modulation from background synaptic input. Neuron, 35(4):773-782, 2002.
[23] Yousheng Shu, Andrea Hasenstaub, Mathilde Badoual, Thierry Bal, and David A McCormick. Barrages of synaptic activity control the gain and sensitivity of cortical neurons. Journal of Neuroscience, 23(32):10388-10401, 2003.
[24] Pierre-Olivier Polack, Jonathan Friedman, and Peyman Golshani. Cellular mechanisms of brain state-dependent gain modulation in visual cortex. Nature Neuroscience, 16(9):1331-1339, 2013.
[25] Katie A Ferguson and Jessica A Cardin. Mechanisms underlying gain modulation in the cortex. Nature Reviews Neuroscience, 21(2):80–92, 2020.
[26] David Wyrick and Luca Mazzucato. State-dependent regulation of cortical processing speed via gain modulation. Journal of Neuroscience, 41(18):3988-4005, 2021.
[27] Lyndon R Duong, David Lipshutz, David J Heeger, Dmitri B Chklovskii, and Eero P Simoncelli. Adaptive whitening in neural populations with gain-modulating interneurons. Proceedings of the 40th International Conference on Machine Learning, PMLR, 202:8902-8921, 2023.
[28] Sreyas Mohan, Joshua L Vincent, Ramon Manzorro, Peter Crozier, Carlos Fernandez-Granda, and Eero Simoncelli. Adaptive denoising via gaintuning. Advances in Neural Information Processing Systems, 34:23727–23740, 2021.
[29] Yonina C Eldar and Alan V Oppenheim. MMSE whitening and subspace whitening. IEEE Transactions on Information Theory, 49(7):1846-1851, 2003.
[30] Eero P Simoncelli and Bruno A Olshausen. Natural image statistics and neural representation. Annual Review of Neuroscience, 24(1):1193-1216, 2001.
[31] Deep Ganguli and Eero P Simoncelli. Efficient sensory encoding and Bayesian inference with heterogeneous neural populations. Neural computation, 26(10):2103-2134, 2014.
[32] Wiktor F Młynarski and Ann M. Hermundstad. Efficient and adaptive sensory codes. Nature Neuroscience, 24(7):998-1009, 2021.
[33] Xiao-Jing Wang, Yinghui Liu, Maria V Sanchez-Vives, and David A McCormick. Adaptation and temporal decorrelation by single neurons in the primary visual cortex. Journal of Neurophysiology, 89(6):3279-3293, 2003.
[34] JH van Hateren and A van der Schaaf. Independent component filters of natural images compared with simple cells in primary visual cortex. Proceedings: Biological Sciences, 265 (1394):359-366, 1998.
[35] Nasir Ahmed, T Natarajan, and Kamisetty R Rao. Discrete cosine transform. IEEE Transactions on Computers, 100(1):90-93, 1974.
[36] David Bull and Fan Zhang. Intelligent Image and Video Compression: Communicating Pictures. Academic Press, London, 2021.
[37] Adrian A Wanner and Rainer W Friedrich. Whitening of odor representations by the wiring diagram of the olfactory bulb. Nature Neuroscience, 23(3):433-442, 2020.
[38] Scott Linderman, Christopher H Stock, and Ryan P Adams. A framework for studying synaptic plasticity with neural spike train data. Advances in Neural Information Processing Systems, 27, 2014.
[56] Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, and Hang Zhao. On feature decorrelation in self-supervised learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9598-9608, 2021.
[57] Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg: Variance-invariance-covariance regularization for self-supervised learning. International Conference on Learning Representations, 2022.
[58] Sebastian Thrun and Lorien Pratt. Learning to Learn. Springer Science & Business Media, 2012.
[59] Chelsea Finn, Aravind Rajeswaran, Sham Kakade, and Sergey Levine. Online meta-learning. In International Conference on Machine Learning, pages 1920-1930. PMLR, 2019.
[60] Johannes Ballé, Philip A Chou, David Minnen, Saurabh Singh, Nick Johnston, Eirikur Agustsson, Sung Jin Hwang, and George Toderici. Nonlinear transform coding. IEEE Journal of Selected Topics in Signal Processing, 15(2):339-353, 2020.
[61] Lyndon R Duong, Bohan Li, Cheng Chen, and Jingning Han. Multi-rate adaptive transform coding for video compression. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 1-5. IEEE, 2023.
[62] Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell. Tent: Fully test-time adaptation by entropy minimization. arXiv preprint arXiv:2006.10726, 2020.
[63] John F Kolen and Jordan B Pollack. Backpropagation without weight transport. In Proceedings of 1994 IEEE International Conference on Neural Networks (ICNN'94), volume 3, pages 1375-1380. IEEE, 1994.
A Notation
For  $N, K \geq 2$ , let  $\mathbb{R}^N$  denote  $N$ -dimensional Euclidean space equipped with the usual Euclidean norm  $\|\cdot\|$  and let  $\mathbb{R}_+^N$  denote the non-negative orthant. Let  $\mathbf{0} = [0, \dots, 0]^\top$  and  $\mathbf{1} = [1, \dots, 1]^\top$  respectively denote the vectors of zeros and ones, whose dimensions should be clear from context.
Let  $\mathbb{R}^{N\times K}$  denote the set of  $N\times K$  real-valued matrices. Let  $\| \cdot \| _F$  denote the Frobenius norm and  $\| \cdot \|_{\mathrm{op}}$  denote the operator norm. Let  $O(N)$  denote the set of  $N\times N$  orthogonal matrices. Let  $\mathbb{S}^N$  (resp.  $\mathbb{S}_{++}^N$ ) denote the set of  $N\times N$  symmetric (resp. positive definite) matrices. Let  $\mathbf{I}_N$  denote the  $N\times N$  identity matrix.
Given vectors  $\mathbf{u},\mathbf{v}\in \mathbb{R}^N$  , let  $\mathbf{u}\circ \mathbf{v} = [u_1v_1,\dots ,u_Nv_N]^\top \in \mathbb{R}^N$  denote the Hadamard (elementwise) product of  $\mathbf{u}$  and  $\mathbf{v}$ . Let  $\mathrm{diag}(\mathbf{u})$  denote the  $N\times N$  diagonal matrix whose  $(i,i)^{\mathrm{th}}$  entry is  $u_{i}$ . Given a matrix  $\mathbf{M}\in \mathbb{R}^{N\times N}$  , we also let  $\mathrm{diag}(\mathbf{M})$  denote the  $N$  -dimensional vector whose  $i^{\mathrm{th}}$  entry is  $M_{ii}$
B Calculation details
B.1 Minimum of the objective
Here we show that the minimum of  $f_{c}$ , defined in equation 2, is achieved at  $\mathbf{M}_c = \mathbf{C}_{ss}^{1/2}(c)$ . Differentiating  $f_{c}(\mathbf{M})$  with respect to  $\mathbf{M}$  yields
$$
\nabla_ {\mathbf {M}} f _ {c} (\mathbf {M}) = - \mathbf {M} ^ {- 1} \mathbf {C} _ {s s} (c) \mathbf {M} ^ {- 1} + \mathbf {I} _ {N}. \tag {10}
$$
Setting the gradient to zero and solving for  $\mathbf{M}$  yields  $\mathbf{M} = \mathbf{C}_{ss}^{1/2}(c)$ . Substituting into  $f_c$  yields  $f_c(\mathbf{M}_c) = \mathrm{Tr}(\mathbf{M}_c^{-1/2}\mathbf{C}_{ss}(c) + \mathbf{M}_c) = 2\mathrm{Tr}(\mathbf{M}_c)$ .
B.2 Adding neural responses
Here we show that equation 4 holds. First, note that the trace term in equation 4 is strictly concave with respect to  $\mathbf{r}$  (assuming  $\mathbf{M}$  is positive definite) and setting the derivative equal to zero yields
$$
2 \mathbf {s} - 2 \mathbf {M r} = 0.
$$
Therefore, the maximum in equation 4 is achieved at  $\mathbf{r} = \mathbf{M}^{-1}\mathbf{s}$ . Substituting into equation 4 with this form for  $\mathbf{r}$ , we get
$$
\begin{array}{l} \mathbb {E} _ {\mathbf {s} \sim p (\mathbf {s} | c)} \left[ \max  _ {\mathbf {r} \in \mathbb {R} ^ {N}} \operatorname {T r} \left(2 \mathbf {r s} ^ {\top} - \mathbf {M} \mathbf {r r} ^ {\top} + \mathbf {M}\right) \right] = \mathbb {E} _ {\mathbf {s} \sim p (\mathbf {s} | c)} \left[ \operatorname {T r} \left(\mathbf {M} ^ {- 1} \mathbf {s s} ^ {\top} + \mathbf {M}\right) \right] \\ = \operatorname {T r} \left(\mathbf {M} ^ {- 1} \mathbf {C} _ {s s} (c) + \mathbf {M}\right) \\ = f _ {c} (\mathbf {M}), \\ \end{array}
$$
where the second equality uses the linearity of the expectation and trace operators as well as the formula  $\mathbf{C}_{ss}(c) \coloneqq \mathbb{E}_{\mathbf{s} \sim p(\mathbf{s}|c)}[\mathbf{ss}^\top]$ . This completes the proof that equation 4 holds. Next, using this expression for  $f_c$ , we have