0.0.1 Introduction↩︎

Modern language models exhibit the remarkable ability to learn novel tasks or solve complex problems from the demonstrations provided within their context window [1]–[4]. Such in-context learning (ICL) offers a novel and effective alternative to traditional fine-tuning techniques. It enables successful prediction across a wide range of tasks simply through a forward pass, eliminating the need for task-specific model weight updates. Since its introduction, ICL capability has become an important feature of LLM with its applications spanning retrieval-augmented generation [5], and reasoning via advanced prompting techniques, such as chain-of-thought [6]. While ICL already exhibits considerable benefits with a small number of demonstrations, i.e., few-shot data, there is a growing interest in extending its benefits to the many-shot settings, potentially realizing even more pronounced benefits [7].

ICL ability also presents an important research avenue to develop stronger theoretical and mechanistic understanding of large language models. To this aim, there has been significant recent interest in demystifying ICL through the lens of function approximation [8], Bayesian inference [9]–[11], and learning and optimization theory [12]–[15]. The latter is concerned with understanding the optimization landscape of ICL, which is also crucial for understanding the generalization properties of the model. A notable result in this direction is the observation that linear attention models [12], [16], [17] implement preconditioned gradient descent (PGD) during ICL [12], [18]. While this line of works provide a fresh perspective to ICL, the existing studies do not address many questions arising from real-life applications nor provide guiding principles for various ICL setups motivated by practical considerations.

To this aim, we revisit the theoretical exploration of ICL with linear data model where we feed an in-context prompt containing \(n\) examples \((\boldsymbol{x}_i,y_i = \boldsymbol{x}_i^\top{\boldsymbol{\beta}}+\xi_i)_{i=1}^{n}\subset\mathbb{R}^{d}\times \mathbb{R}\) and a test instance or query \(\boldsymbol{x}_{n+1} \in \mathbb{R}^d\) to the model, with \(d\) being the feature dimension, \({\boldsymbol{\beta}}\in\mathbb{R}^d\) being the task weight vector, and \((\xi_i)_{i=1}^n\) denoting the noise in individual labels. Given the in-context prompt, the model is tasked to predict \(\hat{y}_{n+1}\) – an estimate for \(y_{n+1}=\boldsymbol{x}_{n+1}^\top{\boldsymbol{\beta}}+\xi_{n+1}\). We aim to provide answers to the following questions by exploring the loss landscape of ICL:

1. Is the ability to implement gradient-based ICL unique to (linear) attention? Can alternative sequence models implement richer algorithms beyond PGD?

2. In language modeling, ICL often works well with few-shot samples whereas standard linear estimation typically requires \({\cal{O}}\left(d\right)\) samples. How can we reconcile this discrepancy between classical learning and ICL?

3. To our knowledge, existing works assume linear-attention is fully parameterized, i.e., key and query projections \(\boldsymbol{W}_k,\boldsymbol{W}_q\in\mathbb{R}^{d\times d}\). What happens when they are low-rank? What happens when there is distribution shift between training and test in-context prompts and we use LoRA [19] for adaptation?

In this work, we conduct a careful investigation of these questions. Specifically, we focus on ICL with 1-layer models and make the following contributions:

1. We jointly investigate the landscape of linear attention and H3 [20], a widely popular state-space model (SSM). We prove that under correlated design, both models implement 1-step PGD (c.f. Proposition 1) and the alignments in Fig. 1 (b) verify that where the dotted curve represents the theoretical PGD result derived from Theorem 1. Our analysis reveals that the gating mechanism in H3 imitates attention. We also empirically show that H3 has the advantage of implementing sample-weighting which allows it to outperform linear attention in temporally-heterogeneous problem settings in Section 0.0.4 and Figure 4.

2. Proposition 1 allows for task and features to be correlated to each other as long as odd moments are zero. Through this, we can assess the impact of distributional alignment on the sample complexity of ICL. Specifically, we characterize the performance of Retrieval Augmented Generation (RAG) (c.f. Theorem 2 and Fig. 1 (c)) and Task-Feature Alignment (c.f. Theorem 3), where the in-context examples are \(\alpha\)-correlated with either the query or the task vector. For both settings, we prove that alignment amplifies the effective sample size of ICL by a factor of \(\alpha^2d+1\), highlighting that aligned data are crucial for the success of ICL in few-shot settings.

3. We show that, under low-rank parameterization, optimal attention-weights still implements PGD according to the truncated eigenspectrum of the fused task-feature covariance (see Section 0.0.3.2). We similarly derive risk upper bounds for LoRA adaptation (c.f. Eq. 16 and Fig. 1 (d)), and show that, these bounds accurately predict the empirical performance.

0.0.2 Problem Setup and Preliminaries↩︎

We begin with a short note on notation. Let bold lowercase and uppercase letters (e.g., \(\boldsymbol{x}\) and \({\boldsymbol{X}}\)) represent vectors and matrices, respectively. The symbol \(\odot\) is defined as the element-wise (Hadamard) product, and \(\ast\) denotes the convolution operator. \({\mathbf{1}}_d\) and \({\mathbf{0}}_d\) denote the \(d\)-dimensional all-ones and all-zeros vectors, respectively; and \({\boldsymbol{I}}_d\) denotes the identity matrix of dimension \(d\times d\). Additionally, let \(\texttt{tr}\left(\boldsymbol{W}\right)\) denote the trace of the square matrix \(\boldsymbol{W}\).

As mentioned earlier, we study the optimization landscapes of 1-layer linear attention [16], [21] and H3 [20] models when training with prompts containing in-context data following a linear model. We construct the input in-context prompt similar to  [12]–[14] as follows.

Linear data distribution. Let \((\boldsymbol{x},y)\in\mathbb{R}^{d}\times\mathbb{R}\) be a (feature, label) pair generated by a \(d\)-dimensional linear model parameterized by \({\boldsymbol{\beta}}\in\mathbb{R}^{d}\), i.e., \(y=\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi\), where \(\boldsymbol{x}\) and \({\boldsymbol{\beta}}\) are feature and task vectors, and \(\xi\) is the label noise. Given demonstrations \((\boldsymbol{x}_i,y_i)_{i=1}^{n+1}\) sampled from a single \({\boldsymbol{\beta}}\), define the input in-context prompt \[\begin{align} \boldsymbol{Z}=[{\boldsymbol{z}}_1~\dots~{\boldsymbol{z}}_n~{\boldsymbol{z}}_{n+1}]^\top=\begin{bmatrix} \boldsymbol{x}_1&\dots&\boldsymbol{x}_n&\boldsymbol{x}_{n+1}\\ y_1&\dots&y_n&0 \end{bmatrix}^\top \in \mathbb{R}^{(n+1) \times (d+1)}.\label{def32Z} \end{align}\tag{1}\] Here, we set \({\boldsymbol{z}}_{i}=\begin{bmatrix} \boldsymbol{x}_{i}\\ y_i \end{bmatrix}\) for \(i\leq n\) and the last/query token \({\boldsymbol{z}}_{n+1}=\begin{bmatrix} \boldsymbol{x}_{n+1}\\ 0 \end{bmatrix}\). Then, given \(\boldsymbol{Z}\), the goal of the model is to predict the correct label \(y_{n+1}\) corresponding to \(\boldsymbol{x}_{n+1}\). For cleaner notation, when it is clear from context, we drop the subscript \(n+1\) and set \(\boldsymbol{x}=\boldsymbol{x}_{n+1},~{\boldsymbol{z}}={\boldsymbol{z}}_{n+1}\). Different from the previous work [12]–[14], [18] where \((\boldsymbol{x}_i)_{i=1}^{n+1}\) and \({\boldsymbol{\beta}}\) are assumed to be independent, our analysis focuses on a more general linear setting that captures the dependency between \((\boldsymbol{x}_i)_{i=1}^{n+1}\) and \({\boldsymbol{\beta}}\).

Model architectures. To start with, we first review the architectures of both Transformer and state-space model (SSM). Similar to the previous work [12]–[14], [17] and to simplify the model structure, we focus on single-layer models and omit the nonlinearity, e.g., softmax operation and MLP activation, from the Transformer. Given the input prompt \(\boldsymbol{Z}\in\mathbb{R}^{(n+1)\times (d+1)}\) in 1 , which can be treated as a sequence of \((d+1)\)-dimensional tokens, the single-layer linear attention \({\texttt{ATT}}\) and H3-like single-layer SSM \({\texttt{SSM}}\) are denoted by \[\tag{2} \begin{align} &{\texttt{ATT}}(\boldsymbol{Z})=(\boldsymbol{Z}\boldsymbol{W}_q\boldsymbol{W}_k^\top\boldsymbol{Z}^\top)\boldsymbol{Z}\boldsymbol{W}_v\tag{3}\\ &{\texttt{SSM}}(\boldsymbol{Z})=\left((\boldsymbol{Z}\boldsymbol{W}_q)\odot((\boldsymbol{Z}\boldsymbol{W}_k\odot\boldsymbol{Z}\boldsymbol{W}_v) \ast\boldsymbol{f})\right)\tag{4} \end{align}\] where \(\boldsymbol{W}_k,~\boldsymbol{W}_q,~\boldsymbol{W}_v\in\mathbb{R}^{(d+1)\times(d+1)}\) denote the key, query and value weight matrices, respectively. In 4 , the parameter \(\boldsymbol{f}\in\mathbb{R}^{n+1}\) is a 1-D convolutional filter that mixes tokens. The Hadamard product \(\odot\) is the gating mechanism [22] between key and query channels, which is crucial for attention-like feature creation. Thus, 4 is more generally a gated-convolution layer. For \(\boldsymbol{f}\) only, we use indexing \(\boldsymbol{f}=[f_0~\dots~f_n]^\top\in\mathbb{R}^{n+1}\) and given any vector \(\boldsymbol{a}\), denote convolution output \((\boldsymbol{a}\ast\boldsymbol{f})_i=\sum_{j=1}^{i} f_{i-j}a_j\). Note that our notation slightly differs from the original H3 model [20] in two ways:

1. SSMs provide efficient parameterization of \(\boldsymbol{f}\) which would otherwise grow with sequence length. In essence, H3 utilizes a linear state-space model \(\boldsymbol{s}_i={\boldsymbol{A}}\boldsymbol{s}_{i-1}+{\boldsymbol{B}}u_i\) and \(y_i={\boldsymbol{C}}\boldsymbol{s}_i\) with parameters \(({\boldsymbol{A}}\in\mathbb{R}^{d\times d},{\boldsymbol{B}}\in\mathbb{R}^{d\times 1},{\boldsymbol{C}}\in\mathbb{R}^{1\times d})\) from which the filter \(\boldsymbol{f}\) is obtained via the impulse response \(f_i={\boldsymbol{C}}{\boldsymbol{A}}^i{\boldsymbol{B}}\) for \(i\geq 0\). Here \(d\) is the state dimension and, in practice, \({\boldsymbol{A}}\) is chosen to be diagonal. Observe that, setting \(d=1\) and \({\boldsymbol{A}}=\rho,{\boldsymbol{C}}={\boldsymbol{B}}=1\), SSM reduces to the exponential smoothing \(f_i=\rho^i\) for \(i\geq 0\). Thus, H3 also captures the all-ones filter as a special instance. As we show in Proposition 1, this simple filter is optimal under independent data model and exactly imitates linear attention. Note that, utilizing a filter \(\boldsymbol{f}\) as in 4 is strictly more expressive than the SSM as it captures all possible impulse responses.

2. H3 also applies a shift SSM to the key embeddings to enable the retrieval of the local context around associative recall hits. We opted not to incorporate this shift operator in our model. This is because unless the features of the neighboring tokens are correlated (which is not the case for the typical independent data model), the entry-wise products between values and shifted keys will have zero mean and be redundant for the final prediction.

We note that we conduct all empirical evaluations with the original H3 model, which displays exact agreement with our theory formalized for 12 , further validating our modeling choice.

0.0.2.1 In-context Linear Estimation

We will next study the algorithms that can be implemented by the single-layer attention and state-space models. Through this, we will show that training \({\texttt{ATT}}\) and \({\texttt{SSM}}\) with linear ICL data is equivalent to the prediction obtained from one step of optimally-preconditioned gradient descent (PGD) and sample-weighted preconditioned gradient descent (WPGD), respectively. We will further show that under mild assumption, the optimal sample weighting for \({\texttt{SSM}}\) (e.g., \(\boldsymbol{f}\)) is an all-ones vector and therefore, establishing the equivalence among PGD, \({\texttt{ATT}}\), and \({\texttt{SSM}}\).

Background: 1-step gradient descent. Consider minimizing squared loss and solving linear regression using one step of PGD and WPGD. Given \(n\) samples \((\boldsymbol{x}_i,y_i)_{i=1}^n\), define \[{\boldsymbol{X}}=[\boldsymbol{x}_1~\cdots~\boldsymbol{x}_n]^\top\in\mathbb{R}^{n\times d}\quad\text{and}\quad \boldsymbol{y}=[y_1~\cdots~y_n]^\top\in\mathbb{R}^n.\] Starting from \({\boldsymbol{\beta}}_0={\mathbf{0}}_d\) and letting \(\eta=1/2\) be the step size, a single-step GD preconditioned with weights \(\boldsymbol{W}\) returns prediction \[\begin{align} \hat{y}=\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{y}:=g_{{\texttt{PGD}}}(\boldsymbol{Z})\label{gd}, \end{align}\tag{5}\] and a single-step sample-weighted GD given weights \(\boldsymbol{\omega}\in\mathbb{R}^{n}\) and \(\boldsymbol{W}\in\mathbb{R}^{d\times d}\) returns prediction \[\begin{align} \hat{y}=\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top(\boldsymbol{\omega}\odot\boldsymbol{y}):=g_{{\texttt{WPGD}}}(\boldsymbol{Z}), \label{wpgd} \end{align}\tag{6}\] where \(\boldsymbol{Z}\) is defined in 1 consisting of \({\boldsymbol{X}},\boldsymbol{y}\) and \(\boldsymbol{x}\). Our goal is to find the optimal \(\boldsymbol{W}\), as well as \(\boldsymbol{\omega}\) in 6 that minimize the population risks defined as follows. \[\tag{7} \begin{align} \min_{\boldsymbol{W}}{\cal{L}}_{{\texttt{PGD}}}({\cal{W}})\quad&\text{where}\quad{\cal{L}}_{{\texttt{PGD}}}({\cal{W}})= \operatorname{\mathbb{E}}\left[(y-g_{{\texttt{PGD}}}(\boldsymbol{Z}))^2\right],\tag{8}\\\min_{\boldsymbol{W},\boldsymbol{\omega}}{\cal{L}}_{{\texttt{WPGD}}}({\cal{W}})\quad&\text{where}\quad{\cal{L}}_{{\texttt{WPGD}}}({\cal{W}})= \operatorname{\mathbb{E}}\left[(y-g_{{\texttt{WPGD}}}(\boldsymbol{Z}))^2\right].\tag{9} \end{align}\] Here, the expectation is over the randomness in \((\boldsymbol{x}_i,\xi_i)_{i=1}^{n+1}\) and \({\boldsymbol{\beta}}\), and we use \({\cal{W}}\) to represent the set of corresponding trainable parameters. The search spaces for \(\boldsymbol{\omega}\) and \(\boldsymbol{W}\) are \(\mathbb{R}^n\) and \(\mathbb{R}^{d\times d}\), respectively.

As per 2 , given input prompt \(\boldsymbol{Z}\in\mathbb{R}^{(n+1)\times (d+1)}\), either of the underlying models outputs a \((n+1)\)-length sequence. Note that the label for the query \(\boldsymbol{x}=\boldsymbol{x}_{n+1}\) is excluded from the prompt \(\boldsymbol{Z}\). Similar to [12], [13], we consider a training objective with a causal mask to ensure inputs cannot attend to their own labels and training can be parallelized. Let \(\boldsymbol{Z}_0=[{\boldsymbol{z}}_1~\dots~{\boldsymbol{z}}_n~0]^\top\) be the features post-causal masking at time/index \(n+1\). Given weights \(\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v\) and the filter \(\boldsymbol{f}\) for \({\texttt{SSM}}\), predictions at the query token \({\boldsymbol{z}}=\begin{bmatrix} \boldsymbol{x}\\0 \end{bmatrix}\) take the following forms following sequence-to-sequence mappings in 2 : \[\begin{align} &g_{{\texttt{ATT}}}(\boldsymbol{Z})=({\boldsymbol{z}}^\top\boldsymbol{W}_q\boldsymbol{W}_k^\top\boldsymbol{Z}_0^\top)\boldsymbol{Z}_0\boldsymbol{W}_v\boldsymbol{v},\\ &g_{{\texttt{SSM}}}(\boldsymbol{Z})=\left(({\boldsymbol{z}}^\top\boldsymbol{W}_q)^\top\odot ((\boldsymbol{Z}_0\boldsymbol{W}_k\odot\boldsymbol{Z}_0\boldsymbol{W}_v) \ast\boldsymbol{f})_{n+1}\right)\boldsymbol{v}, \end{align}\] where \(\boldsymbol{v}\in\mathbb{R}^{d+1}\) is the linear prediction head and \(((\boldsymbol{Z}_0\boldsymbol{W}_k\odot\boldsymbol{Z}_0\boldsymbol{W}_v) \ast\boldsymbol{f})_{n+1}\) returns the last row of the convolution output. Note that SSM can implement the mask by setting \(f_0=0\). Now consider the meta learning setting and select loss function to be the squared loss, same as in 7 . Thus, the objectives for both models take the following forms. \[\tag{10} \begin{align} \min_{\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v,\boldsymbol{v}}{\cal{L}}_{{\texttt{ATT}}}({\cal{W}})\quad&\text{where}\quad{\cal{L}}_{{\texttt{ATT}}}({\cal{W}})= \operatorname{\mathbb{E}}\left[(y-g_{{\texttt{ATT}}}(\boldsymbol{Z}))^2\right]\tag{11},\\ \min_{\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v,\boldsymbol{v},\boldsymbol{f}}{\cal{L}}_{{\texttt{SSM}}}({\cal{W}})\quad&\text{where}\quad{\cal{L}}_{{\texttt{SSM}}}({\cal{W}})= \operatorname{\mathbb{E}}\left[(y-g_{{\texttt{SSM}}}(\boldsymbol{Z}))^2\right].\tag{12} \end{align}\] Here, similarly, the expectation subsumes the randomness of \((\boldsymbol{x}_i,\xi_i)_{i=1}^{n+1}\) and \({\boldsymbol{\beta}}\) and \({\cal{W}}\) represents the set of trainable parameters. The search space for matrices \(\boldsymbol{W}_k\), \(\boldsymbol{W}_q\), \(\boldsymbol{W}_v\) is \(\mathbb{R}^{(d+1)\times(d+1)}\), for head \(\boldsymbol{v}\) is \(\mathbb{R}^{d+1}\), and for \(\boldsymbol{f}\) is \(\mathbb{R}^{n+1}\).

Note that for all the optimization methods (c.f. 7 , 10 ), to simplify the analysis, we train the models without capturing additional bias terms. Therefore, in the following, we introduce the centralized data assumptions such that the models are trained to make unbiased predictions.

To begin with, a cross moment of random variables is defined as the expectation of a monomial of these variables, with the order of the cross moment being the same as order of the monomial. For example, \(\operatorname{\mathbb{E}}[\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{\beta}}]\) is a sum of cross-moments of order 2. Then, it motivates the following data assumptions.

Note that compared to [12]–[14], Assumption 1 is more general which also subsumes the dependent distribution settings. In this work, we consider the following three linear models (omitting noise) satisfying Assumption 1. Let \(\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}},\boldsymbol{\Sigma}_{\boldsymbol{x}}\in\mathbb{R}^{d\times d}\) represent the task and feature covariance matrices for independent data, and let \(0\leq\alpha\leq1\) be the correlation level when considering data dependency. More specific discussions are deferred to Section 0.0.3.

1. Independent task and data: \({\boldsymbol{\beta}}\sim{\cal{N}}(0,\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}),~\boldsymbol{x}_i\sim{\cal{N}}(0,\boldsymbol{\Sigma}_{\boldsymbol{x}}),~~\text{for all}~~1\leq i\leq n+1\).

2. Retrieval augmented generation: \({\boldsymbol{\beta}},\boldsymbol{x}\sim{\cal{N}}(0,{\boldsymbol{I}}_d),~ \boldsymbol{x}_i{{~\big |~}\boldsymbol{x}}\sim{\cal{N}}(\alpha\boldsymbol{x},(1-\alpha^2){\boldsymbol{I}}_d),~~\text{for all}~~1\leq i\leq n\).

3. Task-feature alignment: \({\boldsymbol{\beta}}\sim{\cal{N}}(0,{\boldsymbol{I}}_d),~\boldsymbol{x}_i{{~\big |~}{\boldsymbol{\beta}}}\sim{\cal{N}}(\alpha{\boldsymbol{\beta}},{\boldsymbol{I}}_d),~~\text{for all}~~1\leq i\leq n+1.\)

Next, we introduce the following result which establishes the equivalence among optimizing 1-layer linear attention (c.f. 11 ), 1-layer H3 (c.f. 12 ), and 1-step gradient descent (c.f. 7 ).

We defer the proof to Appendix 1.0.1.1. Proposition 1 establishes that analyzing the optimization landscape of ICL for both single-layer linear attention and the H3 model can be effectively reduced to examining the behavior of a one-step PGD algorithm. Notably, under the independent, RAG and task-feature alignment data settings discussed above, examples \((\boldsymbol{x}_i,y_i)_{i=1}^n\) are independently sampled given \(\boldsymbol{x}\) and \({\boldsymbol{\beta}}\), and we therefore conclude that \({\cal{L}}_{{\texttt{PGD}}}^\star={\cal{L}}_{{\texttt{ATT}}}^\star={\cal{L}}_{{\texttt{SSM}}}^\star\). Leveraging this result, the subsequent section of the paper concentrate on addressing 8 , taking into account various linear data distributions.

While Proposition 1 demonstrates the equivalence of optimal losses, we also study the uniqueness and equivalence of optimal prediction functions. To this end, we analyze the strong convexity of \({\cal{L}}_{{\texttt{PGD}}}({\cal{W}})\) and derive the subsequent lemmas.

As mentioned above, in this work, we study three specific linear models: with general independent, RAG-related, and task-feature alignment data. Note that for all the three cases, according to Proposition 1, we have \({\cal{L}}^\star_{{\texttt{PGD}}}={\cal{L}}^\star_{{\texttt{ATT}}}={\cal{L}}^\star_{{\texttt{SSM}}}\). Additionally, the second claim in Lemma 2 holds, and \({\cal{L}}_{{\texttt{PGD}}}({\cal{W}})\) is strongly convex. Therefore, following Lemma 1, we have \(g^\star_{{\texttt{PGD}}}=g^\star_{\texttt{ATT}}\). Thanks to the equivalence among PGD, \({\texttt{ATT}}\), and \({\texttt{SSM}}\), in the next section, we focus on the solution of objective 8 under different scenarios, which will reflect the optimization landscapes of \({\texttt{ATT}}\) and \({\texttt{SSM}}\) models.

0.0.3 Main Results↩︎

In light of Proposition 1, optimizing a single layer linear-attention or H3 model is equivalent to solving the objective 8 . Therefore, in this section, we examine the properties of the one-step PGD in 8 . To this end, we consider multiple problem settings, including distinct data distributions and low-rank training. The latter refers to the scenario where the key and query matrices have rank restrictions, e.g., \(\boldsymbol{W}_k,\boldsymbol{W}_q\in\mathbb{R}^{(d+1)\times r}\), as well as LoRA-tuning when adapting the model under distribution shift.

0.0.3.1 Analysis of Linear Data Models

We first consider the standard independent data setting. We will then examine correlated designs.

Independent data model. Let \(\boldsymbol{\Sigma}_{\boldsymbol{x}}\) and \(\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}\) be the covariance matrices of the input feature and task vectors, respectively, and \(\sigma\geq0\) be the noise level. We assume \[\begin{align} {\boldsymbol{\beta}}\sim{\cal{N}}(0,\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}),\quad\boldsymbol{x}_i\sim{\cal{N}}(0,\boldsymbol{\Sigma}_{\boldsymbol{x}}),\quad\xi_i\sim{\cal{N}}(0,\sigma^2),\quad 1\leq i\leq n+1\label{data32ind} \end{align}\tag{13}\] and the label is obtained via \(y_i=\boldsymbol{x}_i^\top{\boldsymbol{\beta}}+\xi_i\). Our following result characterizes the optimal solution of 8 . Note that the data generated from 13 satisfies the conditions in Proposition 1. Therefore, the same results can be applied to both linear-attention and H3 models.

See Appendix 1.0.2.2 for proofs. Note that Theorem 1 is consistent with prior work [12] when specialized to isotropic task covariance, i.e., \(\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}={\boldsymbol{I}}_d\). However, their result is limited as the features and task are assumed to be independent. This prompts us to ask: What is the optimization landscape with correlated in-context samples? Toward this, we consider the following RAG-inspired and task-feature alignment models, where Assumptions 1 and 2 continue to hold and Proposition 1 applies.

Retrieval augmented generation. To provide a statistical model of the practical RAG approaches, given the query vector \(\boldsymbol{x}_{n+1}=\boldsymbol{x}\), we propose to draw ICL demonstrations that are similar to \(\boldsymbol{x}\) with the same shared task vector \({\boldsymbol{\beta}}\). Modeling feature similarity through the cosine angle, RAG should sample the ICL examples \(\boldsymbol{x}_i\), \(i\leq n\), from the original feature distribution conditioned on the event \(\text{cos}(\boldsymbol{x}_i,\boldsymbol{x})\geq \alpha\) where \(\alpha\) is the similarity threshold. As an approximate proxy, under the Gaussian distribution model, we assume that \({\boldsymbol{\beta}}\sim\mathcal{N}(0,{\boldsymbol{I}}_d),~\boldsymbol{x}\sim\mathcal{N}(0,{\boldsymbol{I}}_d)\) and that RAG samples \(\alpha\)-correlated demonstrations \((\boldsymbol{x}_i,y_i)_{i=1}^n\) as follows: \[\begin{align} \boldsymbol{x}_i{~\big |~}\boldsymbol{x}\sim\mathcal{N}(\alpha\boldsymbol{x},(1-\alpha^2){\boldsymbol{I}}_d),\quad\xi_i\sim{\cal{N}}(0,\sigma^2)\quad\text{and}\quad y_i=\boldsymbol{x}_i^\top {\boldsymbol{\beta}}+\xi_i,\quad 1\leq i\leq n.\label{data32rag} \end{align}\tag{14}\] Note that the above normalization ensures that the marginal feature distribution remains \(\mathcal{N}(0,{\boldsymbol{I}}_d)\). The full analysis of RAG is provides in Appendix 1.0.2.3. Specifically, when we carry out the analysis by assuming \(\alpha={\cal{O}}\left(1/\sqrt{d}\right)\) and \(d/n={\cal{O}}\left(1\right)\) where \({\cal{O}}\left(\cdot\right)\) denotes proportionality, our derivation leads to the following result:

Here, ?? is reminiscent of Corollary 1 and has a surprisingly clean message. Observe that, \(\alpha^2d+1\) is the dominant multiplier ahead of \(n\) in both equations. Thus, we deduce that, RAG model follows the same error bound as the independent data model, however, its sample size is amplified by a factor of \(\alpha^2d+1\). \(\alpha=0\) reduces to the result of Corollary 1 whereas we need to set \(\alpha={\cal{O}}\left(1/\sqrt{d}\right)\) for constant amplification. When \(\alpha=1\), RAG achieves the approximate risk \({\cal{L}_\star}\approx 2+\sigma^2\), where the constant bias is due to the higher order moments (e.g., the \(4\)’th and \(6\)’th moments) of the standard Gaussian distribution. As \(d\) increases, the normalized loss \({\cal{L}}_\star/d\to 0\). The full analysis of its optimal solution \(\boldsymbol{W}_\star\) and loss \({\cal{L}}_\star\) are deferred Theorem 4 in Appendix 1.0.2.3.

Task-feature alignment. We also consider another dependent data setting where task and feature vectors are assumed to be correlated. This dataset model has the following motivation: In general, an LLM can generate any token within the vocabulary. However, once we specify the task (e.g. domain of the prompt), the LLM output becomes more deterministic and there are much fewer token candidates. For instance, if the task is “Country”, “France” is a viable output compared to “Helium” and vice versa when the task is “Chemistry”. Formally speaking, this can be formalized as the input \(\boldsymbol{x}\) having a diverse distribution whereas it becomes more predictable conditioned on \({\boldsymbol{\beta}}\). Therefore, it can be captured through a linear model by making the conditional covariance of \(\boldsymbol{x}{~\big |~}{\boldsymbol{\beta}}\) to be approximately low-rank. This formalism can be viewed as a spectral alignment between input and task, which is also well-established in deep learning both empirically and theoretically [23]–[26]. Here, we consider such a setting where the shared task vector is sampled as standard Gaussian distribution \({\boldsymbol{\beta}}~\sim{\cal{N}}(0,{\boldsymbol{I}}_d)\) and letting \(\kappa=\alpha^2d+1\), we sample the \(\alpha\)-correlated ICL demonstrations \((\boldsymbol{x}_i,y_i)_{i=1}^{n+1}\) as follows: \[\begin{align} {\boldsymbol{x}_i{~\big |~}{\boldsymbol{\beta}}} \sim\mathcal{N}(\alpha{\boldsymbol{\beta}},{\boldsymbol{I}}_d),\quad\xi_i\sim{\cal{N}}(0,\sigma^2)\quad\text{and}\quad y_i={\kappa^{-1/2}}\boldsymbol{x}_i^\top {\boldsymbol{\beta}}+\xi_i,\quad 1\leq i\leq n+1.\label{data32tfa} \end{align}\tag{15}\] Above, \({\kappa^{-1/2}}\) is a normalization factor to ensure that label variance remains invariant to \(\alpha\). To keep the exposition cleaner, we defer the full analysis of its optimal solution \(\boldsymbol{W}_\star\) and loss \({\cal{L}}_\star\) to Theorem 5 in Appendix 1.0.2.4. Similar to the RAG setting, by assuming \(\alpha={\cal{O}}\left(1/\sqrt{d}\right)\) and \(d/n={\cal{O}}\left(1\right)\), we obtain the following results for the optimal parameter and risk.

Similar to ?? , ?? contains \(\kappa=\alpha^2+1\) multiplier ahead of \(n\), which reduces the in-context sample complexity and setting \(\alpha=0\) reduces to the results of Corollary 1.

0.0.3.2 Low-rank Parameterization and LoRA

In this section, we investigate training low-rank models, which assume \(\boldsymbol{W}_k,\boldsymbol{W}_q\in\mathbb{R}^{(d+1)\times r}\) where \(r\) is the rank restriction. Equivalently, we consider objective 8 under condition \(\text{rank}\left(\boldsymbol{W}\right)=r\).

Note that \(\texttt{tr}\left(\boldsymbol{\Sigma}\right)=\sum_{i=1}^d\lambda_i\). Removing the rank constraint and considering noiseless data setting, this reduces to the following optimal risk \({\cal{L}_\star}=\sum_{i=1}^d\frac{\lambda_i+M}{n+1+M/\lambda_i}\). See Appendix 1.0.3.1 for more details.

Impact of LoRA: Based on the above lemma, we consider the impact of LoRA for adapting the pretrained model to a new task distribution under jointly-diagonalizable old and new eigenvalues of \(\boldsymbol{\Sigma},~\boldsymbol{\Sigma}^{new}\), \((\lambda_i)_{i=1}^d,~(\lambda_i^{new})_{i=1}^d\). Consider adapting LoRA matrix to the combined key and value weights in attention, which reflects minimizing the population loss \(\tilde{\cal{L}}(\boldsymbol{W}_{lora}):={\cal{L}}(\boldsymbol{W}+\boldsymbol{W}_{lora})\) in 8 with fixed \(\boldsymbol{W}\). Suppose \(\texttt{tr}\left(\boldsymbol{\Sigma}\right)=\texttt{tr}\left(\boldsymbol{\Sigma}^{new}\right)=M\), \(\sigma=0\) and \(\boldsymbol{W}\) is jointly diagonalizable with \(\boldsymbol{\Sigma},~\boldsymbol{\Sigma}^{new}\), then LoRA’s risk is upper-bounded by \[\begin{align} \min_{\text{rank}\left(\boldsymbol{W}_{lora}\right)\leq r}\tilde{\cal{L}}(\boldsymbol{W}_{lora})\leq \min_{|{\mathcal{I}}|\leq r,{\mathcal{I}}\subset[d]}\left(\sum_{i\not\in{\mathcal{I}}}\frac{\lambda_i+M}{n+1+M/\lambda_i}+\sum_{i\in{\mathcal{I}}}\frac{\lambda^{new}_i+M}{n+1+M/\lambda^{new}_i}\right). \label{formula32lora} \end{align}\tag{16}\] Note that, the right hand side is provided assuming the optimal LoRA-updated model \(\boldsymbol{W}_{lora}\) is also jointly diagonalizable with covariances \(\boldsymbol{\Sigma},~\boldsymbol{\Sigma}^{new}\), and \(\boldsymbol{W}\).

0.0.4 Experiments↩︎

We now conduct synthetic experiments to support our theoretical findings and further explore the behavior of different models of interest under different conditions. The experiments are designed to investigate various scenarios, including independent data, retrieval-augmented generation (RAG), task-feature alignment, low-rank parameterization, and LoRA adaption.

Experimental setting. We train 1-layer attention and H3 models for solving the linear regression ICL. As described in Section 0.0.2, we consider meta-learning setting where task parameter \({\boldsymbol{\beta}}\) is randomly generated for each training sequence. In all experiments, we set the dimension \(d=20\). Depending on the in-context length (\(n\)), different models are trained to make in-context predictions. We train each model for \(10000\) iterations with batch size \(128\) and Adam optimizer with learning rate \(10^{-3}\). Since our study focuses on the optimization landscape, and experiments are implemented via gradient descent, we repeat \(20\) model trainings from different initialization and results are presented as the minimal test risk among those \(20\) trails. In all the plots, theoretical predictions are obtained via the corresponding formulae presented in Section 0.0.3 and the test risks are normalized by the dimension \(d\).

\(\bullet\) Equivalence among \({\cal{L}}^\star_{{\texttt{PGD}}}\), \({\cal{L}}_{{\texttt{ATT}}}^\star\) and \({\cal{L}}_{{\texttt{SSM}}}^\star\) (Figure 2). To verify Proposition 1 as well as Theorem 1, we run random linear regression instances where in-context samples are generated obeying 13 . Fig. 2 (a) is identical to Fig. 1 (b) where we set \(\boldsymbol{\Sigma}_{\boldsymbol{x}}=\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}={\boldsymbol{I}}_d\) and \(\sigma=0\). In Fig. 2 (b), set \(\boldsymbol{\Sigma}_{\boldsymbol{x}}=\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}={\boldsymbol{I}}\) and vary noise level \(\sigma^2\) from \(0\) to \(0.3\times d\). In Fig. 2 (c), we consider noiseless labels, \(\sigma=0\), isotropic feature distribution \(\boldsymbol{\Sigma}_{\boldsymbol{x}}={\boldsymbol{I}}_d\) and set task covariance to be \(\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}=\gamma{\mathbf{1}}{\mathbf{1}}^\top+(1-\gamma){\boldsymbol{I}}_d\) by choosing \(\gamma\) in \(\{0, 0.3, 0.6, 0.9\}\). Note that in Fig. 2 (c), we train a sufficient number of models (greater than \(20\)) to ensure the optimal model is obtained. In all the figures, solid and dashed curves correspond to the ICL results from training 1-layer \({\texttt{ATT}}\) and \({\texttt{SSM}}\) models, respectively, and dotted curves are obtained from ?? in Theorem 1. The alignment of solid, dashed and dotted curves validates our Proposition 1 and Theorem 1.

\(\bullet\) Distributional alignment experiments (Figs. 3 (a)&3 (b)). In Figs. 3 (a) and 3 (b), we generate RAG and task-feature alignment data following 14 and 15 , respectively, by setting \(\sigma=0\) and varying \(\alpha\) from \(0\) to \(0.6\). Attention training results are displayed in solid curves, and we generate theory curve (dotted) via the \({\cal{L}_\star}\) formula as described in ?? in Appendix 1.0.2.3 and ?? in Appendix 1.0.2.4. The empirical alignments corroborate Theorems 4 and 5, further confirming that Proposition 1 is applicable to a broader range of real-world distributional alignment data.

\(\bullet\) Low-rank (Fig. 3 (c)) and LoRA (Fig. 3 (d)) experiments. We also run simulations to verify our theoretical findings in Section 0.0.3.2. Consider the independent data setting as described in 13 . In Fig. 3 (c), we set \(\boldsymbol{\Sigma}_{\boldsymbol{x}}={\boldsymbol{I}}_d\), \(\sigma=0\) and task covariance to be diagonal with diagonal entries \(c[1~{2}^{-1}~\cdots~d^{-1}]^\top\) for some normalization constant \(c=d/\sum_{i=1}^di^{-1}\), and parameterize the attention model using matrices \(\boldsymbol{W}_k,\boldsymbol{W}_q\in\mathbb{R}^{(d+1)\times r}\) and vary \(r\) across the set \(\{1,5,10,20\}\). Results show that empirical (solid) and theoretical (dotted, c.f. ?? ) curves overlap. In Fig. 3 (d), we implement two phases of training. Phase 1: Setting \(\boldsymbol{\Sigma}_{\boldsymbol{x}}=\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}={\boldsymbol{I}}_d\) and \(\sigma=0\), we pretrain the model with full rank parameters and obtain weights \(\hat{\boldsymbol{W}}_k,\hat{\boldsymbol{W}}_q,\hat{\boldsymbol{W}}_v\in\mathbb{R}^{(d+1)\times (d+1)}\). Phase 2: We generate new examples with task covariance \(\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}\) being a diagonal matrix with diagonal entries \(c'[{2}^{-1}~2^{-2}~\cdots~2^{-d}]^\top\) for some normalization constant \(c'=d/\sum_{i=1}^d2^{-i}\). Given the rank restriction \(r\), we train additional LoRA parameters \(\boldsymbol{W}_{\text{up}},\boldsymbol{W}_{\text{down}}\in\mathbb{R}^{(d+1)\times r}\) where \(\boldsymbol{W}_{lora}:=\boldsymbol{W}_{\text{up}}\boldsymbol{W}_{\text{down}}^\top\) and 3 becomes \({\texttt{ATT}}(\boldsymbol{Z})=(\boldsymbol{Z}(\hat{\boldsymbol{W}}_q\hat{\boldsymbol{W}}_k^\top+\boldsymbol{W}_{\text{up}}\boldsymbol{W}_{\text{down}}^\top)\boldsymbol{Z}^\top)\boldsymbol{Z}\hat{\boldsymbol{W}}_v\). Fig. 3 (d) presents the results after two phases of training where dotted curves are drawn from the right hand side of 16 directly. Here, note that since \(\boldsymbol{\Sigma},\boldsymbol{\Sigma}^{new}\) are diagonal, the right hand side of 16 returns the exact optimal risk of LoRA and the alignments verify it.

\(\bullet\) H3 outperforms linear attention (Figure 4). Until now, our analysis has established the equivalence between linear attention and H3 models in solving linear ICL problem. Furthermore, we also investigate settings where H3 could outperform linear attention due to its sample weighting ability. In Figs. 4 (a) and 4 (b), instead of training separate models to fit the different context lengths, we train a single model with fixed max-length \(n_{\max}\) and loss is evaluated as the average loss given samples from \(1\) to \(n_{\max}\). Such setting has been wildly studied in the previous ICL work [27]–[29]. We generate data according to 13 with \(\boldsymbol{\Sigma}_{\boldsymbol{x}}=\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}={\boldsymbol{I}}_d\) and \(\sigma=0\), and train 1-layer linear attention (Fig. 4 (a)) and H3 (Fig. 4 (b)) models with different max-lengths \(n_{\max}=30,50,80\). Comparison between Fig. 4 (a) and 4 (b) shows that 1-layer attention and H3 implement different algorithms in solving the averaged linear regression problem and H3 is more consistent in generalizing to longer context lengths. In Fig. 4 (c), we plot the averaged risks for each model and H3 outperforms linear attention. Furthermore, in Fig. 4 (d), we focus on the setting where in-context examples are generated using evolving task vector \({\boldsymbol{\beta}}\). Specifically, consider that each sequence corresponds to two individual task parameters \({\boldsymbol{\beta}}_1\sim{\cal{N}}(0,{\boldsymbol{I}}_d)\) and \({\boldsymbol{\beta}}_2\sim{\cal{N}}(0,{\boldsymbol{I}}_d)\). Then the \(i\)’th sample is generated via \(\boldsymbol{x}_i\sim{\cal{N}}(0,{\boldsymbol{I}}_d)\) and \(y_i={\boldsymbol{\beta}}_i^\top\boldsymbol{x}_i\) where \({\boldsymbol{\beta}}_i=\lambda_i{\boldsymbol{\beta}}_1+(1-\lambda_i){\boldsymbol{\beta}}_2\) and \(\lambda_i=i/n\). The results are reported in Fig. 4 (d) which again shows that H3 achieves better performance compared to linear attention, as H3 may benefit from the additional convolutional filter (c.f. \(\boldsymbol{f}\) in 4 ). Here, dotted curve represent the theoretical results under i.i.d. and noiseless setting, derived from Corollary 1.

0.0.5 Related Work↩︎

There is growing interest in understanding the mechanisms behind ICL [1], [30], [31] in large language models (LLMs) due to its success in continuously enabling novel applications for LLMs [2]–[4]. Towards this, [10] explain ICL by language model’s ability to perform implicit Bayesian inference where, under specific assumptions on the pre-training data distribution, the model infers a shared latent concept among the in-context examples and leverages the concept to make a prediction. [9], [32], [33] introduce prior-data fitted network (PFN) to approximate Bayesian inference on synthetic datasets and use it to perform downstream tasks such as tabular dataset classification. On the other hand, [34] posit induction heads as the key mechanism enabling ICL in Transformers. [35] study how various distributional properties of training data aid in the emergence of ICL in Transformers.

In the previous work, [27] explored ICL ability of Transformers. In particular, they considered in-context prompts where each in-context example is labeled by a target function from a given function class, including linear models. A number of works have studied this and related settings to develop a theoretical understanding of ICL [28], [36]–[43]. [28] focus on linear regression and provide a construction of Transformer weights that can enable a single step of GD based on in-context examples. They further show that Transformers trained on in-context prompts exhibit behaviors similar to the models recovered via explicit learning algorithm on the in-context examples in a prompt. Along the similar line, [17] provide a construction of weights in linear attention-only Transformers that can emulate GD steps on in-context examples for a linear regression task. Interestingly, they find similarity between their constructed networks and the networks resulting from training on in-context prompts corresponding to linear regression tasks. Similar to this line of work, [44] argue that pre-trained language models act as meta-optimizer which utilize attention to apply meta-gradients to the original language model based on the in-context examples. Focusing on various NLP tasks, they further connect it to a specific form of explicit fine-tuning that performs gradient updates to the attention-related parameters. Inspired by the connection between linear attention and GD, they developed a novel attention mechanism that mirrors the behavior of GD with momentum. Beyond Transformers, existing work [45]–[47] demonstrate that other model architectures, such as SSM and RNNs, are also capable of in-context learning (ICL).

Building on these primarily empirical studies, [12]–[15] focus on developing a theoretical understanding of Transformers trained to perform ICL. For single-layer linear attention model trained on in-context prompts for random linear regression tasks with isotropic Gaussian features and isotropic Gaussian weight vectors, [12], [13] show that the resulting model implements a single step of GD on in-context examples in a test prompt, thereby corroborating the findings of [17]. They also show that the learned model implements a PGD step, when faced with anisotropic Gaussian features, with [13] also considering anisotropic Gaussian weight vectors. [12] further study multi-layer model and show that the trained model can implement a generalization of GD++ algorithm, supporting an empirical observation in [17]. On the other hand, [13] extend their single-layer setup to consider suitable non-linear target functions, showing that learned Transformer again implements a single step of GD on lineare regression objective.

For a single-layer linear attention model, [14] study the optimization dynamics of gradient flow while training such a model on in-context prompts for random linear regression tasks. Despite the non-convexity of the underlying problem, they show the convergence to the global minimum of the population objective. Similar to [12], [13], they show that the trained model implements a single step of GD and PGD for isotropic and anisotropic Gaussian features, respectively. In addition, they also characterize the test-time prediction error for the trained model while highlighting its dependence on train and test prompt lengths. Interestingly, [14] further explore the effect of various distributional shifts, including the shift in task weight vector distributions between train and test time as well as the covariate shifts among train and test in-context prompts. Interestingly, they find that while linear-attention models are robust to most shifts, they exhibit brittleness to the covariate shifts. While our work shares similarities with this line of works, as discussed in our contributions in the introduction, we expand the theoretical understanding of ICL along multiple novel dimensions, which includes the first study of LoRA adaptation for ICL in the presence of a distributional shift. Furthermore, we strive to capture the effect of retrieval augmentation [5], [48] on ICL through our analysis. Retrieval augmentation allows for selecting most relevant demonstration out of a large collection for a test instance, e.g., via a dense retrieval model [49], which can significantly outperform the typical ICL setup where fixed task-specific demonstrations are provided as in-context examples [50], [51]. Through a careful modeling of retrieval augmentation via correlated design, we show that it indeed has a desirable amplification effect where the effective number in-context examples becomes larger with higher correlation which corresponds to preforming a successful retrieval of query-relevant demonstrations in a practical retrieval augmented setup.

Recently, state space models (SSMs) [20], [52]–[54] have appeared as potential alternatives to Transformer architecture, with more efficient scaling to input sequence length. Recent studies demonstrate that such SSMs can also perform ICL for simple non-language tasks [35], [47] as well as complex NLP tasks [47]. That said, a rigorous theoretical understanding of ICL for SSMs akin to [12]–[14] is missing from the literature. In this work, we provide the first such theoretical treatment for ICL with SSMs. Focusing on H3 architecture [20], we highlight its advantages over linear attention in specific ICL settings.

0.0.6 Discussion↩︎

In this work, we revisited the loss landscape of in-context learning with 1-layer sequence models. We have established a general connection between ICL and gradient methods that accounts for correlated data, non-attention architectures (specifically SSMs), and the impact of low-rank parameterization including LoRA adaptation. Our results elucidate two central findings: (1) The functions learned by different sequence model architectures exhibit a strong degree of universality and (ii) Dataset and prompt design, such as RAG, can substantially benefit ICL performance.

Future directions and limitations. The results of this work fall short of being a comprehensive theory for ICL in LLMs and can be augmented in multiple directions. First, while the exact equivalence between H3 and linear attention is remarkable, we should examine whether it extends to other SSMs. Secondly, while empirically predictive, our RAG and LoRA analyses are not precise and fully formal. Thirdly, it is desirable to develop a deeper understanding of multilayer architectures and connect to iterative GD methods as in [12], [17]. Finally, we have studied the population risk of ICL training whereas one can also explore the sample complexity of pretraining [55], [56]. Moving beyond the theoretically tractable setup of this work, our simplified models are trained on in-context prompts from random initialization. Therefore, this theoretical study doesn’t address more challenging in-context learning tasks, such as question answering, where both in-context demonstration and general knowledge from pretraining are required. Future work in this area could also shed light on how certain contexts might elicit undesirable behaviors acquired by an LLM during pretraining, an aspect not covered in our current analysis. This work also studies a theoretical model for retrieval augmentation-based ICL. In a real-life retrieval augmentation-based ICL, one needs to account for the quality of the collection of the retrievable demonstrations and its (negative) impacts on the final predictions.

Acknowledgements

This work was supported in part by the National Science Foundation grants CCF-2046816, CCF-2403075, the Office of Naval Research award N000142412289, an Adobe Data Science Research award, and a gift by Google Research.

1 Appendix↩︎

1.0.1 Equivalence among Gradient Descent, Attention, and State-Space Models↩︎

In this section, we present the proofs related to Section 0.0.2. Recap that given data \[\begin{align} &{\boldsymbol{X}}=[\boldsymbol{x}_1~\cdots~\boldsymbol{x}_n]^\top\in\mathbb{R}^{n\times d},\\ &\boldsymbol{\xi}=[\xi_1~\cdots~\xi_n]^\top\in\mathbb{R}^n,\\ &\boldsymbol{y}=[y_1~\cdots~y_n]^\top={\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi}\in\mathbb{R}^n,\\ & \boldsymbol{Z}_0=[{\boldsymbol{z}}_1~\dots~{\boldsymbol{z}}_n~{\mathbf{0}}_{d+1}]^\top=\begin{bmatrix} \boldsymbol{x}_1&\dots&\boldsymbol{x}_n&{\mathbf{0}}_d\\ y_1&\dots&y_n&0 \end{bmatrix}^\top \in \mathbb{R}^{(n+1) \times (d+1)}, \end{align}\] and corresponding prediction functions \[\tag{17} \begin{align} &g_{{\texttt{PGD}}}(\boldsymbol{Z})=\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{y},\tag{18}\\ &g_{{\texttt{WPGD}}}(\boldsymbol{Z})=\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top(\boldsymbol{\omega}\odot\boldsymbol{y}),\tag{19}\\ &g_{{\texttt{ATT}}}(\boldsymbol{Z})=({\boldsymbol{z}}^\top\boldsymbol{W}_q\boldsymbol{W}_k^\top\boldsymbol{Z}_0^\top)\boldsymbol{Z}_0\boldsymbol{W}_v\boldsymbol{v},\tag{20}\\ &g_{{\texttt{SSM}}}(\boldsymbol{Z})=\left(({\boldsymbol{z}}^\top\boldsymbol{W}_q)^\top\odot((\boldsymbol{Z}_0\boldsymbol{W}_k\odot\boldsymbol{Z}_0\boldsymbol{W}_v) \ast\boldsymbol{f})_{n+1}\right)\boldsymbol{v},\tag{21} \end{align}\] we have objectives \[\tag{22} \begin{align} \min_{\boldsymbol{W}}{\cal{L}}_{{\texttt{PGD}}}({\cal{W}})\quad&\text{where}\quad{\cal{L}}_{{\texttt{PGD}}}({\cal{W}})= \operatorname{\mathbb{E}}\left[(y-g_{{\texttt{PGD}}}(\boldsymbol{Z}))^2\right]\tag{23},\\ \min_{\boldsymbol{W},\boldsymbol{\omega}}{\cal{L}}_{{\texttt{WPGD}}}({\cal{W}})\quad&\text{where}\quad{\cal{L}}_{{\texttt{WPGD}}}({\cal{W}})= \operatorname{\mathbb{E}}\left[(y-g_{{\texttt{WPGD}}}(\boldsymbol{Z}))^2\right]\tag{24},\\ \min_{\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v,\boldsymbol{v}}{\cal{L}}_{{\texttt{ATT}}}({\cal{W}})\quad&\text{where}\quad{\cal{L}}_{{\texttt{ATT}}}({\cal{W}})= \operatorname{\mathbb{E}}\left[(y-g_{{\texttt{ATT}}}(\boldsymbol{Z}))^2\right]\tag{25},\\ \min_{\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v,\boldsymbol{v},\boldsymbol{f}}{\cal{L}}_{{\texttt{SSM}}}({\cal{W}})\quad&\text{where}\quad{\cal{L}}_{{\texttt{SSM}}}({\cal{W}})= \operatorname{\mathbb{E}}\left[(y-g_{{\texttt{SSM}}}(\boldsymbol{Z}))^2\right].\tag{26} \end{align}\] Here, the expectation is over the randomness in \((\boldsymbol{x}_i,\xi_i)_{i=1}^{n}\) and \({\boldsymbol{\beta}}\), and the search space for \(\boldsymbol{W}\) is \(\mathbb{R}^{d\times d}\), for \(\boldsymbol{\omega}\) is \(\mathbb{R}^{n}\), for \(\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v\) is \(\mathbb{R}^{(d+1)\times(d+1)}\), for \(\boldsymbol{v}\) is \(\mathbb{R}^{d+1}\), and for \(\boldsymbol{f}\) is \(\mathbb{R}^{n+1}\).

1.0.1.1 Proof of Proposition 1

Consider the problem setting as discussed in Section 0.0.2, Proposition 1 can be proven by the following two lemmas.

Recap the linear attention estimator from 20 and denote \[\boldsymbol{W}_q\boldsymbol{W}_k^\top=\begin{bmatrix} \boldsymbol{\bar{W}}&\boldsymbol{w}_1\\ \boldsymbol{w}_2^\top&w \end{bmatrix}\qquad\text{and}\qquad\boldsymbol{W}_v\boldsymbol{v}=\begin{bmatrix} \boldsymbol{v}_1\\ v \end{bmatrix},\] where \(\boldsymbol{\bar{W}}\in\mathbb{R}^{d\times d}\), \(\boldsymbol{w}_1,\boldsymbol{w}_2,\boldsymbol{v}_1\in\mathbb{R}^d\), and \(w,v\in\mathbb{R}\). Then we have \[\begin{align} g_{{\texttt{ATT}}}(\boldsymbol{Z})&=({\boldsymbol{z}}^\top\boldsymbol{W}_q\boldsymbol{W}_k^\top\boldsymbol{Z}_0^\top)\boldsymbol{Z}_0\boldsymbol{W}_v\boldsymbol{v}\nonumber\\ &=[\boldsymbol{x}^\top~0]\begin{bmatrix} \boldsymbol{\bar{W}}&\boldsymbol{w}_1\\ \boldsymbol{w}_2^\top&w \end{bmatrix}\begin{bmatrix} {\boldsymbol{X}}^\top&{\mathbf{0}}_d\\ \boldsymbol{y}^\top&0 \end{bmatrix}\begin{bmatrix} {\boldsymbol{X}}&\boldsymbol{y}\\ {\mathbf{0}}_d^\top&0 \end{bmatrix}\begin{bmatrix} \boldsymbol{v}_1\\ v \end{bmatrix}\nonumber\\ &=(\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top+\boldsymbol{x}^\top\boldsymbol{w}_1\boldsymbol{y}^\top)({\boldsymbol{X}}\boldsymbol{v}_1+\boldsymbol{y}v)\nonumber\\ &=\boldsymbol{x}^\top(v\boldsymbol{\bar{W}}){\boldsymbol{X}}^\top\boldsymbol{y}+\boldsymbol{x}^\top\boldsymbol{w}_1\boldsymbol{y}^\top{\boldsymbol{X}}\boldsymbol{v}_1+\boldsymbol{x}^\top\left(\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1+v\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{w}_1\right)\nonumber\\ &=\boldsymbol{x}^\top(v\boldsymbol{\bar{W}}+\boldsymbol{w}_1\boldsymbol{v}_1^\top){\boldsymbol{X}}^\top\boldsymbol{y}+\boldsymbol{x}^\top\left(\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1+v\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{w}_1\right)\nonumber\\ &=\underset{\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})}{\underbrace{\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}^\top\boldsymbol{y}}}+\underset{\varepsilon}{\underbrace{\boldsymbol{x}^\top\left(\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1+v\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{w}_1\right)}},\label{gd61att32eq321} \end{align}\tag{27}\] where \(\tilde{\boldsymbol{W}}:=v\boldsymbol{\bar{W}}+\boldsymbol{w}_1\boldsymbol{v}_1^\top\).

We first show that for any given parameters \(\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v,\boldsymbol{v}\), \[\begin{align} \operatorname{\mathbb{E}}\left[(g_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right]\geq\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right].\label{lower32bound} \end{align}\tag{28}\] To this goal, we have \[\begin{align} \operatorname{\mathbb{E}}\left[(g_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right]-\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right] &=\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})+\varepsilon-y)^2\right]-\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}[\varepsilon^2]+2\operatorname{\mathbb{E}}[(\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})-y)\varepsilon]\label{risk32gap} \end{align}\tag{29}\] where we have decomposition \[\begin{align} (\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})-y)\varepsilon&=(\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}^\top\boldsymbol{y}-y)\boldsymbol{x}^\top\left(\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1+v\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{w}_1\right)\\ &=\boldsymbol{y}^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\left(\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1+v\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{w}_1\right)-y\boldsymbol{x}^\top\left(\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1+v\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{w}_1\right)\\ &=\underset{(a)}{\underbrace{\boldsymbol{y}^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1}}+\underset{(b)}{\underbrace{v\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{y}^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{w}_1}}-\underset{(c)}{\underbrace{y\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1}}-\underset{(d)}{\underbrace{ vy\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{x}^\top\boldsymbol{w}_1}}. \end{align}\] In the following, we consider the expectations of \((a),(b),(c),(d)\) sequentially, which return zeros under Assumptions 1 and 2. Note that since Assumption 1 holds, expectation of any odd order of monomial of the entries of \({\boldsymbol{X}},\boldsymbol{x},{\boldsymbol{\beta}}\) returns zero, i.e., order of \(\boldsymbol{x}^\top{\boldsymbol{\beta}}\boldsymbol{x}\) is 3 and therefore \(\operatorname{\mathbb{E}}[\boldsymbol{x}^\top{\boldsymbol{\beta}}\boldsymbol{x}]={\mathbf{0}}_d\). \[\begin{align} (a):\quad&\operatorname{\mathbb{E}}\left[\boldsymbol{y}^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1\right]\\ &=\operatorname{\mathbb{E}}\left[({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1\right]\\ &=\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{X}}^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1\right]+\operatorname{\mathbb{E}}\left[\boldsymbol{\xi}^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1\right]\\ &=0. \end{align}\] \[\begin{align} (b):\quad&\operatorname{\mathbb{E}}\left[v\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{y}^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{w}_1\right]\\ &=\operatorname{\mathbb{E}}\left[v({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})^\top({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{w}_1\right]\\ &=\operatorname{\mathbb{E}}\left[v\left\|{\boldsymbol{\xi}}\right\|_{\ell_2}^2\boldsymbol{\xi}^\top{\boldsymbol{X}}\tilde{\boldsymbol{W}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{w}_1\right]\\ &=0. \end{align}\] \[\begin{align} (c):\quad&\operatorname{\mathbb{E}}\left[y\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1\right]\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi)\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1\right]\\ &=\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1\right]+\operatorname{\mathbb{E}}\left[\xi\boldsymbol{x}^\top\boldsymbol{\bar{W}}{\boldsymbol{X}}^\top{\boldsymbol{X}}\boldsymbol{v}_1\right]\\ &=0. \end{align}\] \[\begin{align} (d):\quad&\operatorname{\mathbb{E}}\left[vy\left\|{\boldsymbol{y}}\right\|_{\ell_2}^2\boldsymbol{x}^\top\boldsymbol{w}_1\right]\\ &=v\operatorname{\mathbb{E}}\left[({\boldsymbol{\beta}}^\top\boldsymbol{x}+\xi)({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})^\top({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})\boldsymbol{x}^\top\boldsymbol{w}_1\right]\\ &=v\operatorname{\mathbb{E}}\left[\xi\left\|{\boldsymbol{\xi}}\right\|_{\ell_2}^2\boldsymbol{x}^\top\boldsymbol{w}_1\right]\\ &=0. \end{align}\] Combining the results with 29 returns that \[\begin{align} \operatorname{\mathbb{E}}\left[(g_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right]-\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right]=\operatorname{\mathbb{E}}[\varepsilon^2]\geq0\label{gd61att32eq322} \end{align}\tag{30}\] which completes the proof of 28 . Therefore, we obtain \[\min_{\boldsymbol{W}_q,\boldsymbol{W}_k,\boldsymbol{W}_v,\boldsymbol{v}}\operatorname{\mathbb{E}}\left[(g_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right]\geq\min_{\tilde{\boldsymbol{W}}}\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{ATT}}}(\boldsymbol{Z})-y)^2\right]=\min_{\boldsymbol{W}}\operatorname{\mathbb{E}}\left[(g_{{\texttt{PGD}}}(\boldsymbol{Z})-y)^2\right].\] We conclude the proof of this lemma by showing that for any \(\boldsymbol{W}\in\mathbb{R}^{d\times d}\) in \(g_{{\texttt{PGD}}}\), there exist \(\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v,\boldsymbol{v}\) such that \(g_{\texttt{ATT}}(\boldsymbol{Z})=g_{{\texttt{PGD}}}(\boldsymbol{Z})\). Let \[\boldsymbol{W}_k=\boldsymbol{W}_v={\boldsymbol{I}}_{d+1},\qquad\boldsymbol{W}_q=\begin{bmatrix} \boldsymbol{W}&{\mathbf{0}}_d\\ {\mathbf{0}}_d^\top&0 \end{bmatrix},\quad\text{and}\quad\boldsymbol{v}=\begin{bmatrix} {\mathbf{0}}_d\\1 \end{bmatrix}.\] Then we obtain \[\begin{align} g_{{\texttt{ATT}}}(\boldsymbol{Z})=\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{y}=g_{{\texttt{PGD}}}(\boldsymbol{Z}), \label{gd61att32eq323} \end{align}\tag{31}\] which completes the proof. height6pt width6pt depth0pt

Recap the SSM estimator from 21 and let \[\begin{align} &\boldsymbol{W}_q=\begin{bmatrix} \boldsymbol{w}_{q1}& \boldsymbol{w}_{q2}& \cdots& \boldsymbol{w}_{q,{d+1}} \end{bmatrix},\\ &\boldsymbol{W}_k=\begin{bmatrix} \boldsymbol{w}_{k1}& \boldsymbol{w}_{k2}& \cdots& \boldsymbol{w}_{k,{d+1}} \end{bmatrix},\\ &\boldsymbol{W}_v=\begin{bmatrix} \boldsymbol{w}_{v1}& \boldsymbol{w}_{v2}& \cdots& \boldsymbol{w}_{v,{d+1}} \end{bmatrix}, \end{align}\] where \(\boldsymbol{w}_{qj},\boldsymbol{w}_{kj},\boldsymbol{w}_{vj}\in\mathbb{R}^{d+1}\) for \(j\leq d+1\), and let \[\boldsymbol{v}=\begin{bmatrix} v_1\\ v_2\\ \cdots\\ v_{d+1} \end{bmatrix},\quad\text{and}\quad\boldsymbol{f}=\begin{bmatrix} f_{0}\\ f_{1}\\ \cdots\\ f_{n} \end{bmatrix}.\] Then we have \[\begin{align} g_{\texttt{SSM}}(\boldsymbol{Z})&=\left(({\boldsymbol{z}}^\top\boldsymbol{W}_q)^\top\odot((\boldsymbol{Z}_0\boldsymbol{W}_k\odot\boldsymbol{Z}_0\boldsymbol{W}_v) \ast\boldsymbol{f})_{n+1}\right)\boldsymbol{v}\\ &=\sum_{i=1}^nf_{n+1-i}\cdot\boldsymbol{v}^\top\left(\begin{bmatrix} \boldsymbol{w}_{q1}^\top{\boldsymbol{z}}\\ \cdots\\ \boldsymbol{w}_{q,{d+1}}^\top{\boldsymbol{z}} \end{bmatrix}\odot\begin{bmatrix} \boldsymbol{w}_{k1}^\top{\boldsymbol{z}}_i\boldsymbol{w}_{v1}^\top{\boldsymbol{z}}_i\\ \cdots\\ \boldsymbol{w}_{k,{d+1}}^\top{\boldsymbol{z}}_i\boldsymbol{w}_{v,{d+1}}^\top{\boldsymbol{z}}_i \end{bmatrix}\right)\\ &=\sum_{i=1}^nf_{n+1-i}\cdot\boldsymbol{v}^\top\begin{bmatrix} \boldsymbol{w}_{q1}^\top{\boldsymbol{z}}\boldsymbol{w}_{k1}^\top{\boldsymbol{z}}_i\boldsymbol{w}_{v1}^\top{\boldsymbol{z}}_i\\ \cdots\\ \boldsymbol{w}_{q,{d+1}}^\top{\boldsymbol{z}}\boldsymbol{w}_{k,{d+1}}^\top{\boldsymbol{z}}_i\boldsymbol{w}_{v,{d+1}}^\top{\boldsymbol{z}}_i \end{bmatrix}. \end{align}\] Next for all \(j\leq d+1\), let \[\boldsymbol{w}_{qj}=\begin{bmatrix} \bar{\boldsymbol{w}}_{qj}\\ w_{qj} \end{bmatrix},\quad\boldsymbol{w}_{kj}=\begin{bmatrix} \bar{\boldsymbol{w}}_{kj}\\ w_{kj} \end{bmatrix},\quad\boldsymbol{w}_{vj}=\begin{bmatrix} \bar{\boldsymbol{w}}_{vj}\\ w_{vj} \end{bmatrix}\] where \(\bar{\boldsymbol{w}}_{qj},\bar{\boldsymbol{w}}_{kj},\bar{\boldsymbol{w}}_{vj}\in\mathbb{R}^d\) and \(w_{qj},w_{kj},w_{vj}\in\mathbb{R}\). Then we have \[\begin{align} \boldsymbol{w}_{qj}^\top{\boldsymbol{z}}\boldsymbol{w}_{kj}^\top{\boldsymbol{z}}_i\boldsymbol{w}_{vj}^\top{\boldsymbol{z}}_i&=\left(\bar{\boldsymbol{w}}_{qj}^\top\boldsymbol{x}\right)\left(\bar{\boldsymbol{w}}_{kj}^\top\boldsymbol{x}_i+w_{kj}y_i\right)\left(\bar{\boldsymbol{w}}_{vj}^\top\boldsymbol{x}_i+w_{vj}y_i\right)\\ &=\boldsymbol{x}^\top\bar{\boldsymbol{w}}_{qj}\left(w_{vj}\bar{\boldsymbol{w}}_{kj}^\top+w_{kj}\bar{\boldsymbol{w}}_{vj}^\top\right)\boldsymbol{x}_iy_i+\left(\bar{\boldsymbol{w}}_{qj}^\top\boldsymbol{x}\right)\left(\bar{\boldsymbol{w}}_{kj}^\top\boldsymbol{x}_i\right)\left(\bar{\boldsymbol{w}}_{vj}^\top\boldsymbol{x}_i\right)+\left(w_{kj}w_{vj}\bar{\boldsymbol{w}}_{qj}^\top\boldsymbol{x}y_i^2\right)\\ &=\boldsymbol{x}^\top\boldsymbol{W}'_j\boldsymbol{x}_i y_i+\delta_j(\boldsymbol{x},\boldsymbol{x}_i,\boldsymbol{x}_i)+{\boldsymbol{w}_j'}^\top \boldsymbol{x}y_i^2 \end{align}\] where \[\begin{align} &\boldsymbol{W}'_j:=\bar{\boldsymbol{w}}_{qj}\left(w_{vj}\bar{\boldsymbol{w}}_{kj}^\top+w_{kj}\bar{\boldsymbol{w}}_{vj}^\top\right)\in\mathbb{R}^{d\times d},\\ &\boldsymbol{w}_j':=w_{kj}w_{vj}\bar{\boldsymbol{w}}_{qj}\in\mathbb{R}^d,\\ &\delta_j(\boldsymbol{x},\boldsymbol{x}_i,\boldsymbol{x}_i):=\left(\bar{\boldsymbol{w}}_{qj}^\top\boldsymbol{x}\right)\left(\bar{\boldsymbol{w}}_{kj}^\top\boldsymbol{x}_i\right)\left(\bar{\boldsymbol{w}}_{vj}^\top\boldsymbol{x}_i\right)\in\mathbb{R}. \end{align}\] Then \[\begin{align} g_{{\texttt{SSM}}}(\boldsymbol{Z})&=\sum_{i=1}^nf_{n+1-i}\cdot\sum_{j=1}^{d+1}v_j\left(\boldsymbol{x}^\top\boldsymbol{W}'_j\boldsymbol{x}_i y_i+\delta_j(\boldsymbol{x},\boldsymbol{x}_i,\boldsymbol{x}_i)+{\boldsymbol{w}_j'}^\top \boldsymbol{x}y_i^2\right)\\ &=\boldsymbol{x}^\top\left(\sum_{j=1}^{d+1}v_j\boldsymbol{W}'_j\right){\boldsymbol{X}}(\boldsymbol{y}\odot\tilde{\boldsymbol{f}})+\sum_{i=1}^nf_{n+1-i}\cdot\sum_{j=1}^{d+1}v_j\cdot\delta_j(\boldsymbol{x},\boldsymbol{x}_i,\boldsymbol{x}_i)+\left(\sum_{j=1}^{d+1}v_j{\boldsymbol{w}_j'}^\top\right)\boldsymbol{x}\boldsymbol{y}^\top(\boldsymbol{y}\odot\tilde{\boldsymbol{f}})\\ &=\underset{\tilde{g}_{{\texttt{SSM}}}(\boldsymbol{Z})}{\underbrace{\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}\tilde{\boldsymbol{y}}}}+\underset{\varepsilon_1}{\underbrace{\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})}}+\underset{\varepsilon_2}{\underbrace{\tilde{\boldsymbol{w}}^\top\boldsymbol{x}\boldsymbol{y}^\top\tilde{\boldsymbol{y}}}}. \end{align}\] where \[\begin{align} &\tilde{\boldsymbol{f}}:=[f_n~\cdots~f_1]^\top\in\mathbb{R}^n,\\ &\tilde{\boldsymbol{y}}:=\boldsymbol{y}\odot\tilde{\boldsymbol{f}}\in\mathbb{R}^n,\\ &\tilde{\boldsymbol{W}}:=\sum_{j=1}^{d+1}v_j\boldsymbol{W}'_j\in\mathbb{R}^{d\times d},\\ &\tilde{\boldsymbol{w}}:=\sum_{j=1}^{d+1}v_j{\boldsymbol{w}_j'}\in\mathbb{R}^d,\\ &\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}}):=\sum_{i=1}^nf_{n+1-i}\cdot\sum_{j=1}^{d+1}v_j\cdot\delta_j(\boldsymbol{x},\boldsymbol{x}_i,\boldsymbol{x}_i)\in\mathbb{R}. \end{align}\] Next we will show that for any \(\boldsymbol{W}_k,\boldsymbol{W}_q,\boldsymbol{W}_v,\boldsymbol{v}\), \[\operatorname{\mathbb{E}}\left[(g_{{\texttt{SSM}}}(\boldsymbol{Z})-y)^2\right]\geq\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{SSM}}}(\boldsymbol{Z})-y)^2\right].\] To start with, we obtain \[\begin{align} \operatorname{\mathbb{E}}\left[(g_{\texttt{SSM}}(\boldsymbol{Z})-y)^2\right]&=\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{SSM}}}(\boldsymbol{Z})+\varepsilon_1+\varepsilon_2-y)^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{SSM}}}(\boldsymbol{Z})-y)^2\right]+\operatorname{\mathbb{E}}\left[(\varepsilon_1+\varepsilon_2)^2\right]+2\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{SSM}}}(\boldsymbol{Z})-y)(\varepsilon_1+\varepsilon_2)\right]\label{risk32gap32ssm} \end{align}\tag{32}\] where there is decomposition \[\begin{align} (\tilde{g}_{{\texttt{SSM}}}(\boldsymbol{Z})-y)(\varepsilon_1+\varepsilon_2)=\underset{(a)}{\underbrace{\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})\cdot\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}\tilde{\boldsymbol{y}}}}-\underset{(b)}{\underbrace{\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})y}}+\underset{(c)}{\underbrace{\tilde{\boldsymbol{w}}^\top\boldsymbol{x}\boldsymbol{y}^\top\tilde{\boldsymbol{y}}\cdot\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}\tilde{\boldsymbol{y}}}}-\underset{(d)}{\underbrace{y\cdot\tilde{\boldsymbol{w}}^\top\boldsymbol{x}\boldsymbol{y}^\top\tilde{\boldsymbol{y}}}}. \end{align}\] In the following, similar to the proof of Lemma 4, we consider the expectations of \((a),(b),(c),(d)\) sequentially, which return zeros under Assumptions 1 and 2. Note that \(\delta_j(\boldsymbol{x},\boldsymbol{x}_i,\boldsymbol{x}_i)\)’s and \(\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})\) are summation of monomials of entries of \((\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{\beta}})\) with order 3, and entries of \(\boldsymbol{y}\) and \(y\) are summation of monomials of entries of \((\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{\beta}})\) with even orders: e.g., \(y=\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi\) where \(\xi\) is of oder 0 and \(\boldsymbol{x}^\top{\boldsymbol{\beta}}\) is of order 2.

\[\begin{align} (a):\quad&\operatorname{\mathbb{E}}\left[\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})\cdot\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}\tilde{\boldsymbol{y}}\right]\\ &=\operatorname{\mathbb{E}}\left[\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})\cdot\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}({\boldsymbol{X}}{\boldsymbol{\beta}}\odot\tilde{\boldsymbol{f}})\right]+\operatorname{\mathbb{E}}\left[\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})\cdot\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}(\boldsymbol{\xi}\odot\tilde{\boldsymbol{f}})\right]\\ &=\operatorname{\mathbb{E}}\left[\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})\cdot\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}\right]\operatorname{\mathbb{E}}\left[\boldsymbol{\xi}\odot\tilde{\boldsymbol{f}}\right]\\ &=0. \end{align}\] \[\begin{align} (b):\quad&\operatorname{\mathbb{E}}\left[\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})y\right]\\ &=\operatorname{\mathbb{E}}\left[\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})(\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi)\right]\\ &=\operatorname{\mathbb{E}}\left[\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})\boldsymbol{x}^\top{\boldsymbol{\beta}}\right]+\operatorname{\mathbb{E}}\left[\tilde{\delta}(\boldsymbol{x},{\boldsymbol{X}},{\boldsymbol{X}})\xi\right]\\ &=0. \end{align}\] \[\begin{align} (c):\quad&\operatorname{\mathbb{E}}\left[\tilde{\boldsymbol{w}}^\top\boldsymbol{x}\boldsymbol{y}^\top\tilde{\boldsymbol{y}}\cdot\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}\tilde{\boldsymbol{y}}\right]\\ &=\operatorname{\mathbb{E}}\left[\tilde{\boldsymbol{w}}^\top\boldsymbol{x}({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})^\top({\boldsymbol{X}}{\boldsymbol{\beta}}\odot\tilde{\boldsymbol{f}}+\boldsymbol{\xi}\odot\tilde{\boldsymbol{f}})\cdot\boldsymbol{x}^\top\tilde{\boldsymbol{W}}{\boldsymbol{X}}({\boldsymbol{X}}{\boldsymbol{\beta}}\odot\tilde{\boldsymbol{f}}+\boldsymbol{\xi}\odot\tilde{\boldsymbol{f}})\right]\\ &=0. \end{align}\] \[\begin{align} (d):\quad&\operatorname{\mathbb{E}}\left[y\cdot\tilde{\boldsymbol{w}}^\top\boldsymbol{x}\boldsymbol{y}^\top\tilde{\boldsymbol{y}}\right]\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi)\cdot\tilde{\boldsymbol{w}}^\top\boldsymbol{x}({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})^\top({\boldsymbol{X}}{\boldsymbol{\beta}}\odot\tilde{\boldsymbol{f}}+\boldsymbol{\xi}\odot\tilde{\boldsymbol{f}})\right]\\ &=0. \end{align}\] Combining the results with 32 results that \[\begin{align} \operatorname{\mathbb{E}}\left[(g_{\texttt{SSM}}(\boldsymbol{Z})-y)^2\right]-\operatorname{\mathbb{E}}\left[(\tilde{g}_{\texttt{SSM}}(\boldsymbol{Z})-y)^2\right]=\operatorname{\mathbb{E}}\left[(\varepsilon_1+\varepsilon_2)^2\right]\geq0. \end{align}\] Therefore we obtain, \[\min_{\boldsymbol{W}_q,\boldsymbol{W}_k,\boldsymbol{W}_v,\boldsymbol{v},\boldsymbol{f}}\operatorname{\mathbb{E}}\left[(g_{{\texttt{SSM}}}(\boldsymbol{Z})-y)^2\right]\geq\min_{\tilde{\boldsymbol{W}},\tilde{\boldsymbol{f}}}\operatorname{\mathbb{E}}\left[(\tilde{g}_{{\texttt{SSM}}}(\boldsymbol{Z})-y)^2\right]=\min_{\boldsymbol{W},\boldsymbol{\omega}}\operatorname{\mathbb{E}}\left[( g_{{\texttt{WPGD}}}(\boldsymbol{Z})-y)^2\right].\] Next we show that for any choices of \(\boldsymbol{W}\) and \(\boldsymbol{\omega}\) in \(g_{{\texttt{WPGD}}}\), there are \(\boldsymbol{W}_{q,k,v},\boldsymbol{v},\boldsymbol{f}\) such that \(g_{{\texttt{SSM}}}\equiv g_{{\texttt{WPGD}}}\). To this end, given \(\boldsymbol{\omega}=[\omega_1~\dots~\omega_n]^\top\), let \[\boldsymbol{W}_q={\boldsymbol{I}}_{d+1},\quad\boldsymbol{W}_k=\begin{bmatrix} \boldsymbol{W}^\top&{\mathbf{0}}_d\\ {\mathbf{0}}_d^\top&0 \end{bmatrix},\quad\boldsymbol{W}_v=\begin{bmatrix} {\mathbf{0}}_{d\times d}&{\mathbf{0}}_d\\ {\mathbf{1}}_d^\top&0 \end{bmatrix},\quad\boldsymbol{v}=\begin{bmatrix} {\mathbf{1}}_d\\ 0 \end{bmatrix}\quad\text{and}\quad\boldsymbol{f}=\begin{bmatrix} 0\\ \omega_n\\ \cdots\\ \omega_1 \end{bmatrix}.\] Then we get \[\begin{align} ((\boldsymbol{Z}_0\boldsymbol{W}_k\odot\boldsymbol{Z}_0\boldsymbol{W}_v)\ast\boldsymbol{f})_{n+1}&=\left(\left(\begin{bmatrix} {\boldsymbol{X}}\boldsymbol{W}^\top&{\mathbf{0}}_n\\ {\mathbf{0}}_d&0 \end{bmatrix}\odot\begin{bmatrix} \boldsymbol{y}{\mathbf{1}}_d^\top&{\mathbf{0}}_n\\ {\mathbf{0}}_d&0 \end{bmatrix}\right)\ast\boldsymbol{f}\right)_{n+1}\\ &=\begin{bmatrix}\sum_{i=1}^n\omega_i\cdot y_i\boldsymbol{W}\boldsymbol{x}_i\\ 0\end{bmatrix}\\ &=\begin{bmatrix}\boldsymbol{W}{\boldsymbol{X}}^\top(\boldsymbol{y}\odot\boldsymbol{\omega})\\0\end{bmatrix}, \end{align}\] and therefore \[g_{\texttt{SSM}}(\boldsymbol{Z})=\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top(\boldsymbol{y}\odot\boldsymbol{\omega})=g_{{\texttt{WPGD}}}(\boldsymbol{Z}),\] which completes the proof of ?? .

Next, to show ?? , for any \(\boldsymbol{W}\in\mathbb{R}^{d\times d}\), let \({\cal{L}}(\boldsymbol{\omega})=\operatorname{\mathbb{E}}\left[\left(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top(\boldsymbol{y}\odot\boldsymbol{\omega})-y\right)^2\right]\). Then we have \[\begin{align} \frac{\partial{\cal{L}}(\boldsymbol{\omega})}{\partial \omega_i}&=\operatorname{\mathbb{E}}\left[2\left(\boldsymbol{x}^\top\boldsymbol{W}\sum_{j=1}^n\omega_jy_j\boldsymbol{x}_j-y\right)\left(\boldsymbol{x}^\top\boldsymbol{W}y_i\boldsymbol{x}_i\right)\right]\\ &=2\sum_{j=1}^n\omega_j\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}y_j\boldsymbol{x}_j)(\boldsymbol{x}^\top\boldsymbol{W}y_i\boldsymbol{x}_i)\right]-2\operatorname{\mathbb{E}}\left[y\boldsymbol{x}^\top\boldsymbol{W}y_i\boldsymbol{x}_i\right]. \end{align}\] Here since \((\boldsymbol{x}_i,y_i)_{i=1}^n\) follow the same distribution and are conditionally independent given \(\boldsymbol{x}\) and \({\boldsymbol{\beta}}\), for any \(i\neq j\neq j'\), \(\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}y_i\boldsymbol{x}_i)^2\right]=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}y_{j}\boldsymbol{x}_{j})^2\right]\) and \(\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}y_j\boldsymbol{x}_j)(\boldsymbol{x}^\top\boldsymbol{W}y_i\boldsymbol{x}_i)\right]=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}y_{j'}\boldsymbol{x}_{j'})(\boldsymbol{x}^\top\boldsymbol{W}y_i\boldsymbol{x}_i)\right]\). Then let \[\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}y_j\boldsymbol{x}_j)(\boldsymbol{x}^\top\boldsymbol{W}y_i\boldsymbol{x}_i)\right]=\begin{cases} c_1,& i\neq j\\ c_2,& i=j \end{cases}\quad\text{and}\quad\operatorname{\mathbb{E}}\left[y\boldsymbol{x}^\top\boldsymbol{W}y_i\boldsymbol{x}_i\right]=c_3,\] where \((c_1,c_2,c_3):=(c_1(\boldsymbol{W}),c_2(\boldsymbol{W}),c_3(\boldsymbol{W}))\). We get \[\frac{\partial{\cal{L}}(\boldsymbol{\omega})}{\partial \omega_i} =2c_1\boldsymbol{\omega}^\top{\mathbf{1}}_n+2(c_2-c_1)\omega_i-2c_3.\] If \(c_2-c_1=0\), then \(\frac{\partial{\cal{L}}(\boldsymbol{\omega})}{\partial \omega_i}\equiv2c_1\boldsymbol{\omega}^\top{\mathbf{1}}_n-2c_3\) for all \(i\leq n\) and any \(\boldsymbol{\omega}\in\mathbb{R}^n\) achieves the same performance.

If \(c_2-c_1\neq0\), setting \(\frac{\partial{\cal{L}}(\boldsymbol{\omega})}{\partial \omega_i}=0\) returns \[\omega_i=\frac{c_3-c_1\sum_{j=1}^n\omega_j}{c_2-c_1}:=C\quad\text{for all }i\leq n.\] Therefore the optimal loss is achieved via setting \(\boldsymbol{\omega}=C{\mathbf{1}}_n\). Without loss of generality, we can update \(\boldsymbol{W}\to C\boldsymbol{W}\). Then \(\boldsymbol{\omega}={\mathbf{1}}_n\), and we obtain \[\min_{\boldsymbol{W},\boldsymbol{\omega}}\operatorname{\mathbb{E}}\left[\left(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top(\boldsymbol{y}\odot\boldsymbol{\omega})-y\right)^2\right]=\min_{\boldsymbol{W}}\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{y}-y)^2\right]\] which completes the proof of ?? . height6pt width6pt depth0pt

1.0.1.2 Proof of Lemma 1

Recap the loss \({\cal{L}}_{{\texttt{PGD}}}({\cal{W}})\) in 23 and prediction \(g_{{\texttt{PGD}}}(\boldsymbol{Z})\) in 18 , we have \[\begin{align} {\cal{L}}_{{\texttt{PGD}}}({\cal{W}})&=\operatorname{\mathbb{E}}[(y-g_{{\texttt{PGD}}}(\boldsymbol{Z}))^2]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})\right)^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2+2(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})+(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2+(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]+2\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})]\nonumber\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2+(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]\label{cvx32eq321}\\ &=\underset{f_1(\boldsymbol{W})}{\underbrace{\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2+(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]}}\underset{f_2(\boldsymbol{W})}{\underbrace{-2\operatorname{\mathbb{E}}[{\boldsymbol{\beta}}^\top\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}+\xi\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi}]}}+\underset{\text{constant}}{\underbrace{\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top{\boldsymbol{\beta}})^2+\xi^2]}}\nonumber \end{align}\tag{33}\] where 33 follows Assumption 2. Since \(f_2(\boldsymbol{W})\) is convex, \({\cal{L}}_{{\texttt{PGD}}}({\cal{W}})\) is strongly-convex if and only if \(f_1(\boldsymbol{W})\) is strongly-convex, which completes the proof of strong convexity.

Next, 30 and 31 in the proof of Lemma 4 demonstrate that the optimal loss is achievable and is achieved at \(\varepsilon=0\). Subsequently, 27 indicates that \(g^\star_{{\texttt{ATT}}}\) has the same form as \(g^\star_{{\texttt{PGD}}}\). Under the strong convexity assumption, \(g^\star_{{\texttt{PGD}}}\) is unique, which leads to the conclusion that \(g^\star_{{\texttt{PGD}}} = g^\star_{{\texttt{ATT}}}\). height6pt width6pt depth0pt

1.0.1.3 Proof of Lemma 2

According to Lemma 1, \({\cal{L}}_{{\texttt{PGD}}}({\cal{W}})\) is strongly-convex as long as either \(\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]\) or \(\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2]\) is strongly-convex. Therefore, in this lemma, the two claims correspond to the strong convexity of \(\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2]\) and \(\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]\) terms, respectively.

Suppose the decomposition claim holds. Without losing generality, we may assume \((\boldsymbol{x}_1,{\boldsymbol{\beta}}_1,{\boldsymbol{X}}_1)\) are zero-mean because we can allocate the mean component to \((\boldsymbol{x}_2,{\boldsymbol{\beta}}_2,{\boldsymbol{X}}_2)\) without changing the covariance.

\(\bullet\) Claim 1: Let \(\boldsymbol{\bar{\Sigma}}_{\boldsymbol{x}}=\operatorname{\mathbb{E}}[\boldsymbol{x}_1\boldsymbol{x}_1^\top]\), \(\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{\beta}}}=\operatorname{\mathbb{E}}[{\boldsymbol{\beta}}_1{\boldsymbol{\beta}}_1^\top]\), and \(\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{X}}}=\operatorname{\mathbb{E}}[{\boldsymbol{X}}_1^\top{\boldsymbol{X}}_1]\). If the first claim holds, using independence, observe that we can write \[\begin{align} \operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2]&=\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{X}}_1^\top\boldsymbol{\xi})^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{X}}_2^\top\boldsymbol{\xi})^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_2^\top\boldsymbol{W}{\boldsymbol{X}}_1^\top\boldsymbol{\xi})^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_2^\top\boldsymbol{W}{\boldsymbol{X}}_2^\top\boldsymbol{\xi})^2],\nonumber \end{align}\] where the last three terms of the right hand side are convex and the first term obeys \[\begin{align} \operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{X}}_1^\top\boldsymbol{\xi})^2]&=\sigma^2\operatorname{\mathbb{E}}[\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{X}}_1^\top{\boldsymbol{X}}_1\boldsymbol{W}^\top \boldsymbol{x}_1]\\ &=\sigma^2\texttt{tr}\left(\operatorname{\mathbb{E}}[\boldsymbol{x}_1\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{X}}_1^\top{\boldsymbol{X}}_1\boldsymbol{W}^\top]\right)\\ &=\sigma^2\texttt{tr}\left(\boldsymbol{\bar{\Sigma}}_{\boldsymbol{x}}\boldsymbol{W}\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{X}}}\boldsymbol{W}^\top\right)\\ &=\sigma^2\left\|{\sqrt{\boldsymbol{\bar{\Sigma}}_{\boldsymbol{x}}}\boldsymbol{W}\sqrt{\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{X}}}}}\right\|_{F}^2. \end{align}\] Since noise level \(\sigma>0\), using the full-rankness of covariance matrices \(\boldsymbol{\bar{\Sigma}}_{\boldsymbol{x}}\) and \(\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{X}}}\), we conclude with strong convexity of \(\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2]\).

\(\bullet\) Claim 2: Now recall that \(\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{X}}}=\operatorname{\mathbb{E}}[{\boldsymbol{X}}_1^\top{\boldsymbol{X}}_1]\) and set \({\boldsymbol{A}}={\boldsymbol{X}}_1^\top{\boldsymbol{X}}_1-\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{X}}}\) and \({\boldsymbol{B}}={\boldsymbol{X}}_2^\top{\boldsymbol{X}}_2+\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{X}}}\). Observe that \(\operatorname{\mathbb{E}}[{\boldsymbol{A}}]=0\). If the second claim holds, \(\operatorname{\mathbb{E}}[{\boldsymbol{X}}^\top{\boldsymbol{X}}]=\operatorname{\mathbb{E}}[{\boldsymbol{A}}+{\boldsymbol{B}}]\). Note that \(({\boldsymbol{A}},{\boldsymbol{\beta}}_1,\boldsymbol{x}_1)\) are independent of each other and \(({\boldsymbol{B}},{\boldsymbol{\beta}}_2,\boldsymbol{x}_2)\). Using independence and \(\operatorname{\mathbb{E}}[{\boldsymbol{A}}]=0\), similarly write \[\begin{align} \operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]&=\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{A}}{\boldsymbol{\beta}})^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{B}}{\boldsymbol{\beta}})^2]. \end{align}\] Now using \(\operatorname{\mathbb{E}}[{\boldsymbol{\beta}}_1]=\operatorname{\mathbb{E}}[\boldsymbol{x}_1]=0\) and their independence from rest, these terms obeys \[\begin{align} &\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{A}}{\boldsymbol{\beta}})^2]=\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{A}}{\boldsymbol{\beta}}_1)^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{A}}{\boldsymbol{\beta}}_2)^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_2^\top\boldsymbol{W}{\boldsymbol{A}}{\boldsymbol{\beta}}_1)^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_2^\top\boldsymbol{W}{\boldsymbol{A}}{\boldsymbol{\beta}}_2)^2]\\ &\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{B}}{\boldsymbol{\beta}})^2]=\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{B}}{\boldsymbol{\beta}}_1)^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{B}}{\boldsymbol{\beta}}_2)^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_2^\top\boldsymbol{W}{\boldsymbol{B}}{\boldsymbol{\beta}}_1)^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_2^\top\boldsymbol{W}{\boldsymbol{B}}{\boldsymbol{\beta}}_2)^2]. \end{align}\] In both equations, the last three terms of the right hand side are convex. To proceed, we focus on the first terms. Using independence and setting \(\boldsymbol{\Sigma}_{{\boldsymbol{X}} }=\operatorname{\mathbb{E}}[{\boldsymbol{X}}^\top{\boldsymbol{X}}]\succeq \boldsymbol{\bar{\Sigma}}_{{\boldsymbol{X}}}\succ 0\), we note that \[\begin{align} \operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{A}}{\boldsymbol{\beta}}_1)^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{B}}{\boldsymbol{\beta}}_1)^2]=\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}_1)^2] \end{align}\] where \(\boldsymbol{x}_1,{\boldsymbol{\beta}}_1,{\boldsymbol{X}}\) are independent and full-rank covariance. To proceed, note that \[\begin{align} \operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}_1)^2]=\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}\boldsymbol{\Sigma}_{{\boldsymbol{X}}}{\boldsymbol{\beta}}_1)^2]+\operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}({\boldsymbol{X}}^\top{\boldsymbol{X}}-\boldsymbol{\Sigma}_{{\boldsymbol{X}}}){\boldsymbol{\beta}}_1)^2]. \end{align}\] Observing the convexity of the right hand side and focusing on the first term, we get \[\begin{align} \operatorname{\mathbb{E}}[(\boldsymbol{x}_1^\top\boldsymbol{W}\boldsymbol{\Sigma}_{{\boldsymbol{X}}}{\boldsymbol{\beta}}_1)^2]=\texttt{tr}\left(\boldsymbol{\bar{\Sigma}}_{\boldsymbol{x}}\boldsymbol{W}\boldsymbol{\Sigma}_{{\boldsymbol{X}}}\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{\beta}}}\boldsymbol{\Sigma}_{{\boldsymbol{X}}}\boldsymbol{W}^\top\right)=\left\|{\sqrt{\boldsymbol{\bar{\Sigma}}_x}\boldsymbol{W}\boldsymbol{\Sigma}_{{\boldsymbol{X}}}\sqrt{\boldsymbol{\bar{\Sigma}}_{\boldsymbol{\beta}}}}\right\|_{F}^2. \end{align}\] Using the fact that covariance matrices, \(\boldsymbol{\bar{\Sigma}}_{\boldsymbol{x}},\boldsymbol{\Sigma}_{{\boldsymbol{X}}},\boldsymbol{\bar{\Sigma}}_{{\boldsymbol{\beta}}}\), are full rank concludes the strong convexity proof of \(\operatorname{\mathbb{E}}[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]\). height6pt width6pt depth0pt

1.0.2 Analysis of General Data Distribution↩︎

In this section, we provide the proofs in Section 0.0.3, which focuses on solving Objective 8 . For the sake of clean notation, let \({\cal{L}}(\boldsymbol{W}):={\cal{L}}_{{\texttt{PGD}}}({\cal{W}})\) and \(g:=g_{{\texttt{PGD}}}\) in this section.

1.0.2.1 Supporting Results

We begin by deriving the even moments of random variables.

\(\bullet\) \(2n\)’th moment of a normally distributed variable: Let \(u\sim{\cal{N}}(0,\sigma^2)\). Then we have \[\begin{align} \operatorname{\mathbb{E}}[u^{2n}] = \sigma^{2n}(2n-1)!!.\label{moment322n} \end{align}\tag{34}\]

\(\bullet\) \(4\)’th moment: Let \({\boldsymbol{u}}\sim{\cal{N}}(0,{\boldsymbol{I}}_d)\). Then for any \(\boldsymbol{W},\boldsymbol{W}'\in\mathbb{R}^{d\times d}\), we have \[\begin{align} &\operatorname{\mathbb{E}}\left[({\boldsymbol{u}}^\top\boldsymbol{W}{\boldsymbol{u}})({\boldsymbol{u}}^\top\boldsymbol{W}'{\boldsymbol{u}})\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i,j=1}^dW_{ij}u_iu_j\right)\left(\sum_{i,j=1}^dW'_{ij}u_iu_j\right)\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^dW_{ii}u_i^2\right)\left(\sum_{i=1}^dW'_{ii}u_i^2\right)\right]+\operatorname{\mathbb{E}}\left[\left(\sum_{i\neq j}W_{ij}u_iu_j\right)\left(\sum_{i\neq j}W_{ij}'u_iu_j\right)\right]\nonumber\\ &=\sum_{i=1}^dW_{ii}W'_{ii}\operatorname{\mathbb{E}}\left[u_i^4\right]+\sum_{i\neq j}W_{ii}W'_{jj}\operatorname{\mathbb{E}}[u_i^2]\operatorname{\mathbb{E}}[u_j^2]+\sum_{i\neq j}W_{ij}W'_{ij}\operatorname{\mathbb{E}}[u_i^2]\operatorname{\mathbb{E}}[u_j^2]+\sum_{i\neq j}W_{ij}W'_{ji}\operatorname{\mathbb{E}}[u_i^2]\operatorname{\mathbb{E}}[u_j^2]\nonumber\\ &=3\sum_{i=1}^dW_{ii}W'_{ii}+\sum_{i\neq j}W_{ii}W'_{jj}+\sum_{i\neq j}W_{ij}W'_{ij}+\sum_{i\neq j}W_{ij}W'_{ji}\nonumber\\ &=\sum_{i,j=1}^dW_{ii}W'_{jj}+\sum_{i,j=1}^dW_{ij}W'_{ij}+\sum_{i,j=1}^dW_{ij}W'_{ji}\nonumber\\ &=\texttt{tr}\left(\boldsymbol{W}\right)\texttt{tr}\left(\boldsymbol{W}'\right)+\texttt{tr}\left(\boldsymbol{W}'\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}'\right).\label{moment324} \end{align}\tag{35}\]

\(\bullet\) \(4\)’th cross-moment: Let \({\boldsymbol{u}},\boldsymbol{v}\sim{\cal{N}}(0,{\boldsymbol{I}}_d)\) and for any \(\boldsymbol{W}\in\mathbb{R}^{d\times d}\), let \({\boldsymbol{\Lambda}}_{\boldsymbol{W}}=\boldsymbol{W}\odot{\boldsymbol{I}}_d\). Then we have \[\begin{align} &\operatorname{\mathbb{E}}\left[({\boldsymbol{u}}^\top\boldsymbol{W}\boldsymbol{v}\boldsymbol{v}^\top{\boldsymbol{u}})^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i,j=1}^dW_{ij}u_iv_j\right)^2\left(\sum_{i=1}^du_iv_i\right)^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\scalebox{0.9}{\displaystyle \left[\left(\sum_{i,j=1}^dW_{ij}^2u_i^2v_j^2+\sum_{i\neq i'}W_{ij}W_{i'j}u_iu_{i'}v_j^2+\sum_{j\neq j'}W_{ij}W_{ij'}u_i^2v_jv_{j'}+\sum_{i'\neq i,j'\neq j}W_{ij}W_{i'j'}u_iu_{i'}v_jv_{j'}\right)\left(\sum_{i=1}^du_i^2v_i^2+\sum_{i\neq j}u_iu_jv_iv_j\right)\right]}\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i,j=1}^dW_{ij}^2u_i^2v_j^2\right)\left(\sum_{i=1}^du_i^2v_i^2\right)+\left(\sum_{i\neq j}W_{ij}W_{ji}u_i^2u_j^2v_i^2v_j^2\right)\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^dW_{ii}^2u_i^2v_i^2+\sum_{i\neq j}W_{ij}^2u_i^2v_j^2\right)\left(\sum_{i=1}^du_i^2v_i^2\right)\right]+\sum_{i\neq j}W_{ij}W_{ji}\nonumber\\ &=\operatorname{\mathbb{E}}\scalebox{0.9}{\displaystyle \left[\left(\sum_{i=1}^dW_{ii}^2u_i^4v_i^4+\sum_{i\neq j}W_{ii}^2u_i^2v_i^2u_j^2v_j^2\right)\right]+\operatorname{\mathbb{E}}\left[\left(\sum_{i\neq j}W_{ij}^2u_i^4v_j^2v_i^2+\sum_{i\neq j}W_{ij}^2u_i^2v_j^4u_j^2+\sum_{i\neq j\neq k}W_{ij}^2u_i^2v_j^2u_k^2v_k^2\right)\right]+\sum_{i\neq j}W_{ij}W_{ji}}\nonumber\\ &=9\sum_{i=1}^dW_{ii}^2+(d-1)\sum_{i=1}^dW_{ii}^2+6\sum_{i\neq j}W_{ij}^2+(d-2)\sum_{i \neq j}W_{ij}^2+\sum_{i\neq j}W_{ij}W_{ji}\nonumber\\ &=3\sum_{i=1}^dW_{ii}^2+(d+4)\sum_{i,j=1}^dW_{ij}^2+\sum_{i,j=1}^dW_{ij}W_{ji}\nonumber\\ &=3\texttt{tr}\left({\boldsymbol{\Lambda}}_{\boldsymbol{W}}^2\right)+(d+4)\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}^2\right).\label{cross32moment324} \end{align}\tag{36}\]

\(\bullet\) \(6\)’th moment: Let \({\boldsymbol{u}}\sim{\cal{N}}(0,{\boldsymbol{I}}_d)\). Then for any \(\boldsymbol{W},\boldsymbol{W}'\in\mathbb{R}^{d\times d}\), we have \[\begin{align} &\operatorname{\mathbb{E}}\left[({\boldsymbol{u}}^\top\boldsymbol{W}{\boldsymbol{u}})({\boldsymbol{u}}^\top\boldsymbol{W}'{\boldsymbol{u}})\left\|{{\boldsymbol{u}}}\right\|_{\ell_2}^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i,j=1}^dW_{ij}u_iu_j\right)\left(\sum_{i,j=1}^dW'_{ij}u_iu_j\right)\left(\sum_{i=1}^du_i^2\right)\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^dW_{ii}u_i^2\right)\left(\sum_{i=1}^dW'_{ii}u_i^2\right)\left(\sum_{i=1}^du_i^2\right)\right]+\operatorname{\mathbb{E}}\left[\left(\sum_{i\neq j}W_{ij}u_iu_j\right)\left(\sum_{i\neq j}W_{ij}'u_iu_j\right)\left(\sum_{i=1}^du_i^2\right)\right]\nonumber\\ &=\sum_{i=1}^dW_{ii}W'_{ii}\operatorname{\mathbb{E}}\left[u_i^4\left(\sum_{i'=1}^du_{i'}^2\right)\right]+\sum_{i\neq j}W_{ii}W'_{jj}\operatorname{\mathbb{E}}\left[u_i^2u_j^2\left(\sum_{i'=1}^du_{i'}^2\right)\right]\nonumber\\ &\quad+\sum_{i\neq j}W_{ij}W'_{ij}\operatorname{\mathbb{E}}\left[u_i^2u_j^2\left(\sum_{i'=1}^du_{i'}^2\right)\right]+\sum_{i\neq j}W_{ij}W'_{ji}\operatorname{\mathbb{E}}\left[u_i^2u_j^2\left(\sum_{i'=1}^du_{i'}^2\right)\right]\nonumber\\ &=(d+4)\left(3\sum_{i=1}^dW_{ii}W'_{ii}+\sum_{i\neq j}W_{ii}W'_{jj}+\sum_{i\neq j}W_{ij}W'_{ij}+\sum_{i\neq j}W_{ij}W'_{ji}\right)\tag{37}\\ &=(d+4)\left(\sum_{i,j=1}^dW_{ii}W'_{jj}+\sum_{i,j=1}^dW_{ij}W'_{ij}+\sum_{i,j=1}^dW_{ij}W'_{ji}\right)\nonumber\\ &=(d+4)\left(\texttt{tr}\left(\boldsymbol{W}\right)\texttt{tr}\left(\boldsymbol{W}'\right)+\texttt{tr}\left(\boldsymbol{W}'\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}'\right)\right),\tag{38} \end{align}\] where 37 is obtained by following \[\begin{align} &\operatorname{\mathbb{E}}\left[u_i^4\left(\sum_{i'=1}^du_{i'}^2\right)\right]=\operatorname{\mathbb{E}}[u^6]+(d-1)\operatorname{\mathbb{E}}[u^4]\operatorname{\mathbb{E}}[u^2]=3(d+4),\\ &\operatorname{\mathbb{E}}\left[u_i^2u_j^2\left(\sum_{i'=1}^du_{i'}^2\right)\right]=2\operatorname{\mathbb{E}}[u^4]\operatorname{\mathbb{E}}[u^2]+(d-2)\operatorname{\mathbb{E}}[u^2]\operatorname{\mathbb{E}}[u^2]\operatorname{\mathbb{E}}[u^2]=d+4. \end{align}\]

\(\bullet\) \(8\)’th moment: Let \({\boldsymbol{u}}\sim{\cal{N}}(0,{\boldsymbol{I}}_d)\). Then for any \(\boldsymbol{W},\boldsymbol{W}'\in\mathbb{R}^{d\times d}\), we have \[\begin{align} &\operatorname{\mathbb{E}}\left[({\boldsymbol{u}}^\top\boldsymbol{W}{\boldsymbol{u}})({\boldsymbol{u}}^\top\boldsymbol{W}'{\boldsymbol{u}})\left\|{{\boldsymbol{u}}}\right\|_{\ell_2}^4\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i,j=1}^dW_{ij}u_iu_j\right)\left(\sum_{i,j=1}^dW'_{ij}u_iu_j\right)\left(\sum_{i,j=1}^du_i^2u_j^2\right)\right]\nonumber\\ &=\operatorname{\mathbb{E}}\scalebox{0.9}{\displaystyle \left[\left(\sum_{i=1}^dW_{ii}u_i^2\right)\left(\sum_{i=1}^dW'_{ii}u_i^2\right)\left(\sum_{i=1}^du_i^4+\sum_{i\neq j}u_i^2u_j^2\right)\right]+\operatorname{\mathbb{E}}\left[\left(\sum_{i\neq j}W_{ij}u_iu_j\right)\left(\sum_{i\neq j}W_{ij}'u_iu_j\right)\left(\sum_{i=1}^du_i^4+\sum_{i\neq j}u_i^2u_j^2\right)\right]}\nonumber\\ &=\sum_{i=1}^dW_{ii}W'_{ii}\operatorname{\mathbb{E}}\left[u_i^4\left(\sum_{i'=1}^du_{i'}^4+\sum_{i'\neq j'}u_{i'}^2u_{j'}^2\right)\right]+\sum_{i\neq j}W_{ii}W'_{jj}\operatorname{\mathbb{E}}\left[u_i^2u_j^2\left(\sum_{i'=1}^du_{i'}^4+\sum_{i'\neq j'}u_{i'}^2u_{j'}^2\right)\right]\nonumber\\ &\quad+\sum_{i\neq j}W_{ij}W'_{ij}\operatorname{\mathbb{E}}\left[u_i^2u_j^2\left(\sum_{i'=1}^du_{i'}^4+\sum_{i'\neq j'}u_{i'}^2u_{j'}^2\right)\right]+\sum_{i\neq j}W_{ij}W'_{ji}\operatorname{\mathbb{E}}\left[u_i^2u_j^2\left(\sum_{i'=1}^du_{i'}^4+\sum_{i'\neq j'}u_{i'}^2u_{j'}^2\right)\right]\nonumber\\ &=(d+4)(d+6)\left(3\sum_{i=1}^dW_{ii}W'_{ii}+\sum_{i\neq j}W_{ii}W'_{jj}+\sum_{i\neq j}W_{ij}W'_{ij}+\sum_{i\neq j}W_{ij}W'_{ji}\right)\tag{39}\\ &=(d+4)(d+6)\left(\sum_{i,j=1}^dW_{ii}W'_{jj}+\sum_{i,j=1}^dW_{ij}W'_{ij}+\sum_{i,j=1}^dW_{ij}W'_{ji}\right)\nonumber\\ &=(d+4)(d+6)\left(\texttt{tr}\left(\boldsymbol{W}\right)\texttt{tr}\left(\boldsymbol{W}'\right)+\texttt{tr}\left(\boldsymbol{W}'\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}'\right)\right).\tag{40} \end{align}\] where 39 is obtained by following \[\begin{align} &\operatorname{\mathbb{E}}\left[u_i^4\left(\sum_{i'=1}^du_{i'}^4+\sum_{i'\neq j'}u_{i'}^2u_{j'}^2\right)\right]\\ &=\operatorname{\mathbb{E}}[u^8]+(d-1)\operatorname{\mathbb{E}}[u^4]\operatorname{\mathbb{E}}[u^4]+2(d-1)\operatorname{\mathbb{E}}[u^6]\operatorname{\mathbb{E}}[u^2]+(d-1)(d-2)\operatorname{\mathbb{E}}[u^4]\operatorname{\mathbb{E}}[u^2]\operatorname{\mathbb{E}}[u^2]\\ &=105+9(d-1)+30(d-1)+3(d-1)(d-2)\\ &=3(d+4)(d+6),\\ &\operatorname{\mathbb{E}}\left[u_i^2u_j^2\left(\sum_{i'=1}^du_{i'}^4+\sum_{i'\neq j'}u_{i'}^2u_{j'}^2\right)\right]\\ &=2\operatorname{\mathbb{E}}[u^6]\operatorname{\mathbb{E}}[u^2]+(d-2)\operatorname{\mathbb{E}}[u^4](\operatorname{\mathbb{E}}[u^2])^2+2\operatorname{\mathbb{E}}[u^4]\operatorname{\mathbb{E}}[u^4]+4(d-2)\operatorname{\mathbb{E}}[u^4](\operatorname{\mathbb{E}}[u^2])^2+(d-2)(d-3)(\operatorname{\mathbb{E}}[u^2])^4\\ &=30+3(d-2)+18+12(d-2)+(d-2)(d-3)\\ &=(d+4)(d+6). \end{align}\]

1.0.2.2 Independent Data with General Covariance

Consider a general independent linear model as defined in 13 where \(\boldsymbol{\Sigma}_{\boldsymbol{x}}\) and \(\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}\) are full-rank feature and task convariance matrices and \[\boldsymbol{x}\sim{\cal{N}}(0,\boldsymbol{\Sigma}_{\boldsymbol{x}}),\quad{\boldsymbol{\beta}}\sim{\cal{N}}(0,\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}),\quad\xi\sim{\cal{N}}(0,\sigma^2),\quad\text{and}\quad y=\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi.\] Let \[{\boldsymbol{X}}=[\boldsymbol{x}_1~\cdots~\boldsymbol{x}_n]^\top,\quad\boldsymbol{\xi}=[\xi_1~\cdots~\xi_n]^\top,\quad\text{and}\quad \boldsymbol{y}=[y_1~\cdots~y_n]^\top={\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi}.\]

To simplify and without loss of generality, let \(\boldsymbol{\bar{x}}=\boldsymbol{\Sigma}_{\boldsymbol{x}}^{-1/2}\boldsymbol{x}\), \(\boldsymbol{\bar{X}}={\boldsymbol{X}}\boldsymbol{\Sigma}_{\boldsymbol{x}}^{-1/2}\), \(\bar{\boldsymbol{\beta}}=\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}{\boldsymbol{\beta}}\) where we have \[\boldsymbol{\bar{x}}\sim{\cal{N}}(0,{\boldsymbol{I}}),\qquad\bar{\boldsymbol{\beta}}\sim{\cal{N}}(0,\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2})\] and \[y=\boldsymbol{\bar{x}}^\top\bar{\boldsymbol{\beta}}+\xi,\qquad\boldsymbol{y}=\boldsymbol{\bar{X}}\bar{\boldsymbol{\beta}}+\boldsymbol{\xi}.\]

Then recap the loss from 8 , and we obtain \[\begin{align} {\cal{L}}(\boldsymbol{W})&=\operatorname{\mathbb{E}}\left[(y-g(\boldsymbol{Z}))^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})\right)^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2+2(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})+(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]+\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]+\sigma^2,\label{loss32expanded} \end{align}\tag{41}\] where the last equality comes from the independence of label noise \(\xi,\boldsymbol{\xi}\).

We first consider the following term \[\begin{align} \operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]=\operatorname{\mathbb{E}}\left[(\boldsymbol{\bar{x}}^\top(\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\boldsymbol{W}\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2})\boldsymbol{\bar{X}}^\top\boldsymbol{\xi})^2\right]=n\sigma^2\cdot\texttt{tr}\left(\boldsymbol{\bar{W}}\boldsymbol{\bar{W}}^\top\right) \end{align}\] where we define \(\boldsymbol{\bar{W}}=\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\boldsymbol{W}\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\). Next, focus on the following \[\begin{align} \operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]&=\operatorname{\mathbb{E}}\left[(\boldsymbol{\bar{x}}^\top\bar{\boldsymbol{\beta}}-\boldsymbol{\bar{x}}^\top\boldsymbol{\bar{W}}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\bar{\boldsymbol{\beta}})^2\right]\\ &=\operatorname{\mathbb{E}}\left[\left(\boldsymbol{\bar{x}}^\top\left({\boldsymbol{I}}-\boldsymbol{\bar{W}}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\right)\bar{\boldsymbol{\beta}}\right)^2\right]\\ &=\texttt{tr}\left(\operatorname{\mathbb{E}}\left[\left({\boldsymbol{I}}-\boldsymbol{\bar{W}}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\right)\boldsymbol{\Sigma}\left({\boldsymbol{I}}-\boldsymbol{\bar{W}}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\right)^\top\right]\right)\\ &=\texttt{tr}\left(\boldsymbol{\Sigma}\right)-\texttt{tr}\left(\boldsymbol{\Sigma}(\boldsymbol{\bar{W}}+\boldsymbol{\bar{W}}^\top)\operatorname{\mathbb{E}}[\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}]\right)+\texttt{tr}\left(\boldsymbol{\bar{W}}^\top\boldsymbol{\bar{W}}\operatorname{\mathbb{E}}[\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\boldsymbol{\Sigma}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}]\right)\\ &=\texttt{tr}\left(\boldsymbol{\Sigma}\right)-2n\cdot\texttt{tr}\left(\boldsymbol{\Sigma}\boldsymbol{\bar{W}}\right)+\texttt{tr}\left(\boldsymbol{\bar{W}}^\top\boldsymbol{\bar{W}}\operatorname{\mathbb{E}}[\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\boldsymbol{\Sigma}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}]\right), \end{align}\] where \(\boldsymbol{\Sigma}:=\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\).

Let \(\boldsymbol{\bar{x}}_i\in\mathbb{R}^n\) be the \(i\)’th column of \(\boldsymbol{\bar{X}}\) and \(\boldsymbol{\Sigma}_{ij}\) be the \((i,j)\)’th entry of \(\boldsymbol{\Sigma}\). Then the \((i,j)\) entry of matrix \(\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\boldsymbol{\Sigma}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\) is \[(\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\boldsymbol{\Sigma}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}})_{ij}=\sum_{k=1}^d\sum_{p=1}^d\boldsymbol{\Sigma}_{kp}\boldsymbol{\bar{x}}_i^\top\boldsymbol{\bar{x}}_k\boldsymbol{\bar{x}}_p^\top\boldsymbol{\bar{x}}_j.\] Then we get \[\begin{align} i\neq j:\quad\operatorname{\mathbb{E}}\left[\left(\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\boldsymbol{\Sigma}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\right)_{ij}\right]&=\boldsymbol{\Sigma}_{ij}\operatorname{\mathbb{E}}[\boldsymbol{\bar{x}}_i^\top\boldsymbol{\bar{x}}_i\boldsymbol{\bar{x}}_j^\top\boldsymbol{\bar{x}}_j]+\boldsymbol{\Sigma}_{ji}\operatorname{\mathbb{E}}[\boldsymbol{\bar{x}}_i^\top\boldsymbol{\bar{x}}_j\boldsymbol{\bar{x}}_i^\top\boldsymbol{\bar{x}}_j]=n^2\boldsymbol{\Sigma}_{ij}+n\boldsymbol{\Sigma}_{ji}\\ i=j:\quad\operatorname{\mathbb{E}}\left[\left(\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\boldsymbol{\Sigma}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\right)_{ii}\right]&=\boldsymbol{\Sigma}_{ii}\operatorname{\mathbb{E}}\left[\boldsymbol{\bar{x}}_i^\top\boldsymbol{\bar{x}}_i\boldsymbol{\bar{x}}_i^\top\boldsymbol{\bar{x}}_i\right]+\sum_{j\neq i}\boldsymbol{\Sigma}_{jj}\operatorname{\mathbb{E}}\left[\boldsymbol{\bar{x}}_i^\top\boldsymbol{\bar{x}}_j\boldsymbol{\bar{x}}_j^\top\boldsymbol{\bar{x}}_i\right]\\ &=\boldsymbol{\Sigma}_{ii}\operatorname{\mathbb{E}}\left[(x_{i1}^2+\cdots+x_{in}^2)^2\right]+n\sum_{j\neq i}\boldsymbol{\Sigma}_{jj}\\ &=\boldsymbol{\Sigma}_{ii}(3n+n(n-1))+n\sum_{j\neq i}\boldsymbol{\Sigma}_{jj}\\ &=n\left(\boldsymbol{\Sigma}_{ii}(n+1)+\sum_{j=1}^d\boldsymbol{\Sigma}_{jj}\right)\\ &=n\left(\boldsymbol{\Sigma}_{ii}(n+1)+\texttt{tr}\left(\boldsymbol{\Sigma}\right)\right). \end{align}\] Therefore \[\operatorname{\mathbb{E}}[\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}\boldsymbol{\Sigma}\boldsymbol{\bar{X}}^\top\boldsymbol{\bar{X}}]=n(n+1)\boldsymbol{\Sigma}+n\cdot\texttt{tr}\left(\boldsymbol{\Sigma}\right){\boldsymbol{I}}.\] Combining all together results in \[\begin{align} {\cal{L}}(\boldsymbol{W})&=\texttt{tr}\left(\boldsymbol{\Sigma}\right)-2n\texttt{tr}\left(\boldsymbol{\Sigma}\boldsymbol{\bar{W}}\right)+n(n+1)\texttt{tr}\left(\boldsymbol{\Sigma}\boldsymbol{\bar{W}}^\top\boldsymbol{\bar{W}}\right)+n(\texttt{tr}\left(\boldsymbol{\Sigma}\right)+\sigma^2)\texttt{tr}\left(\boldsymbol{\bar{W}}\boldsymbol{\bar{W}}^\top\right)+\sigma^2,\nonumber\\ &=M-2n\texttt{tr}\left(\boldsymbol{\Sigma}\boldsymbol{\bar{W}}\right)+n(n+1)\texttt{tr}\left(\boldsymbol{\Sigma}\boldsymbol{\bar{W}}^\top\boldsymbol{\bar{W}}\right)+nM\texttt{tr}\left(\boldsymbol{\bar{W}}\boldsymbol{\bar{W}}^\top\right),\label{loss} \end{align}\tag{42}\] where \(M:=\texttt{tr}\left(\boldsymbol{\Sigma}\right)+\sigma^2\). Setting \(\nabla_{\boldsymbol{\bar{W}}}{\cal{L}}(\boldsymbol{W})=0\) returns \[-2n\cdot\boldsymbol{\Sigma}+2n(n+1)\cdot\boldsymbol{\Sigma}\boldsymbol{\bar{W}}+2nM\boldsymbol{\bar{W}}=0\Longrightarrow\boldsymbol{\bar{W}}_\star=\left((n+1){\boldsymbol{I}}+M\boldsymbol{\Sigma}^{-1}\right)^{-1}.\] Then we have \[\boldsymbol{W}_\star=\boldsymbol{\Sigma}_{\boldsymbol{x}}^{-1/2}\left((n+1){\boldsymbol{I}}+M\boldsymbol{\Sigma}^{-1}\right)^{-1}\boldsymbol{\Sigma}_{\boldsymbol{x}}^{-1/2}\] and \[{\cal{L}}_\star={\cal{L}}(\boldsymbol{W}_\star)=M-n\texttt{tr}\left(((n+1)\boldsymbol{\Sigma}^{-1}+M\boldsymbol{\Sigma}^{-2})^{-1}\right).\] height6pt width6pt depth0pt

1.0.2.3 Retrieval Augmented Generation with \(\alpha\) Correlation

In this section, we consider the retrieval augmented generation (RAG) linear model similar to 14 , where we first draw the query vector \(\boldsymbol{x}\) and task vector \({\boldsymbol{\beta}}\) via \[\boldsymbol{x}\sim{\cal{N}}(0,{\boldsymbol{I}})\quad\text{and}\quad{\boldsymbol{\beta}}\sim{\cal{N}}(0,{\boldsymbol{I}}).\] We then draw data \((\boldsymbol{x}_i)_{i=1}^n\) to be used in-context according to the rule \(\text{corr\_coef}(\boldsymbol{x},\boldsymbol{x}_i)\geq \alpha\geq0\). Hence, for \(i\leq n\) we sample \[\begin{align} \boldsymbol{x}_i{~\big |~}\boldsymbol{x}\sim\mathcal{N}(\alpha\boldsymbol{x},\gamma^2{\boldsymbol{I}}),\quad\xi_i\sim{\cal{N}}(0,\sigma^2)\quad\text{and}\quad y_i=\boldsymbol{x}_i^\top {\boldsymbol{\beta}}+\xi_i,\label{data32rag32app} \end{align}\tag{43}\] which results in 14 by setting \(\gamma^2=1-\alpha^2\).

Here, for clean notation and without loss of generality, we define and rewrite 43 via \[{\boldsymbol{g}}_i\sim{\cal{N}}(0,{\boldsymbol{I}}),\quad\xi_i\sim{\cal{N}}(0,\sigma^2)\quad\text{and}\quad\boldsymbol{x}_i=\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}}_i,\quad y_i=(\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}}_i)^\top{\boldsymbol{\beta}}+\xi_i.\] Then we obtain \[\begin{align} {\cal{L}}(\boldsymbol{W})&=\operatorname{\mathbb{E}}\left[(y-g(\boldsymbol{Z}))^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top({\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})\right)^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2+2(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})+(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]+\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]+\sigma^2\label{ind32loss32eq321}. \end{align}\tag{44}\] To begin with, let \[\begin{align} N_1 = \texttt{tr}\left(\boldsymbol{W}\right)^2+\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}^2\right),\quad N_2=\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right),\quad\text{and}\quad N_3=\texttt{tr}\left(\boldsymbol{W}\right).\nonumber \end{align}\] We first focus on the second term in 44 \[\begin{align} \operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]&=\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\xi_i}\boldsymbol{x}^\top\boldsymbol{W}(\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}}_i)\right)^2\right]\\ &=n\sigma^2\operatorname{\mathbb{E}}\left[\boldsymbol{x}^\top\boldsymbol{W}(\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}})(\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}})^\top\boldsymbol{W}^\top\boldsymbol{x}\right]\\ &=n\sigma^2\left(\alpha^2\operatorname{\mathbb{E}}[\boldsymbol{x}^\top\boldsymbol{W}\boldsymbol{x}\boldsymbol{x}^\top\boldsymbol{W}^\top\boldsymbol{x}]+\gamma^2\operatorname{\mathbb{E}}[\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{g}}{\boldsymbol{g}}^\top\boldsymbol{W}^\top\boldsymbol{x}]\right)\\ &=n\sigma^2\left(\alpha^2N_1+\gamma^2N_2\right). and independence of \boldsymbol{x},{\boldsymbol{g}}.} \end{align}\] Next, the first term in 44 can be decomposed into \[\begin{align} \operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]&=\underset{(a)}{\underbrace{\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}})^2\right]}}+\underset{(b)}{\underbrace{\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]}}-2\underset{(c)}{\underbrace{\operatorname{\mathbb{E}}\left[\boldsymbol{x}^\top{\boldsymbol{\beta}}\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}\right]}}. \end{align}\] In the following, we consider solving \((a)\)-\((c)\) sequentially. \[\begin{align} (a):\quad&\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}})^2\right]=d. \end{align}\] \[\begin{align} (b):\quad&\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]\\ &=\operatorname{\mathbb{E}}\left[\left(\boldsymbol{x}^\top\boldsymbol{W}\sum_{i=1}^n(\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}}_i)(\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}}_i)^\top{\boldsymbol{\beta}}\right)^2\right]\\ &=\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n\boldsymbol{x}^\top\boldsymbol{W}(\alpha^2\boldsymbol{x}\boldsymbol{x}^\top+\gamma^2{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top+\alpha\gamma\boldsymbol{x}{\boldsymbol{g}}_i^\top+\alpha\gamma{\boldsymbol{g}}_i\boldsymbol{x}^\top){\boldsymbol{\beta}}\right)^2\right]\\ &=\scalebox{0.85}{\displaystyle \alpha^4n^2\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}\boldsymbol{x}\boldsymbol{x}^\top{\boldsymbol{\beta}})^2\right]+\gamma^4\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2\gamma^2\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n\boldsymbol{x}^\top\boldsymbol{W}\boldsymbol{x}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2\gamma^2\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{g}}_i\boldsymbol{x}^\top{\boldsymbol{\beta}}\right)^2\right]}\\ &\quad\scalebox{0.9}{\displaystyle +2\alpha^2\gamma^2n^2\operatorname{\mathbb{E}}\left[\boldsymbol{x}^\top\boldsymbol{W}\boldsymbol{x}\boldsymbol{x}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{g}}{\boldsymbol{g}}^\top\boldsymbol{W}^\top\boldsymbol{x}\right]+2\alpha^2\gamma^2 n\operatorname{\mathbb{E}}\left[\boldsymbol{x}^\top\boldsymbol{W}\boldsymbol{x}{\boldsymbol{g}}^\top{\boldsymbol{\beta}}\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{g}}\boldsymbol{x}^\top{\boldsymbol{\beta}}\right]}\\ &=\scalebox{0.9}{\displaystyle \left(\alpha^4n^2(d+4)N_1+\gamma^4n(d+n+1)N_2\right)+\left(\alpha^2\gamma^2ndN_1+\alpha^2\gamma^2n(d+2)N_2\right)+\left(2\alpha^2\gamma^2n^2N_1+2\alpha^2\gamma^2nN_1\right)}\\ &=\scalebox{0.9}{\displaystyle \left(\alpha^4n^2(d+4)+\alpha^2\gamma^2n(2n+d+2)\right)N_1+\left(\alpha^2\gamma^2n(d+2)+\gamma^4n(d+n+1)\right)N_2}\\ &=A_1N_1+A_2N_2. \end{align}\] \[\begin{align} (c):\quad&\operatorname{\mathbb{E}}\left[\boldsymbol{x}^\top{\boldsymbol{\beta}}\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}\right]=\operatorname{\mathbb{E}}\left[\sum_{i=1}^n\boldsymbol{x}^\top{\boldsymbol{\beta}}\boldsymbol{x}^\top\boldsymbol{W}(\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}}_i)(\alpha\boldsymbol{x}+\gamma{\boldsymbol{g}}_i)^\top{\boldsymbol{\beta}}\right]\\ &=\operatorname{\mathbb{E}}\left[\sum_{i=1}^n\boldsymbol{x}^\top{\boldsymbol{\beta}}\boldsymbol{x}^\top\boldsymbol{W}(\alpha^2\boldsymbol{x}\boldsymbol{x}^\top+\gamma^2{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top+\alpha\gamma\boldsymbol{x}{\boldsymbol{g}}_i^\top+\alpha\gamma{\boldsymbol{g}}_i\boldsymbol{x}^\top){\boldsymbol{\beta}}\right]\\ &=\alpha^2n\operatorname{\mathbb{E}}\left[\boldsymbol{x}^\top{\boldsymbol{\beta}}\boldsymbol{x}^\top\boldsymbol{W}\boldsymbol{x}\boldsymbol{x}^\top{\boldsymbol{\beta}}\right]+\gamma^2n\operatorname{\mathbb{E}}\left[\boldsymbol{x}^\top{\boldsymbol{\beta}}\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{g}}{\boldsymbol{g}}^\top{\boldsymbol{\beta}}\right]\\ &=\alpha^2n(d+2)\texttt{tr}\left(\boldsymbol{W}\right)+\gamma^2n\texttt{tr}\left(\boldsymbol{W}\right)\\ &=\left(\alpha^2n(d+2)+\gamma^2n\right)N_3\\ &=A_3N_3. \end{align}\] Here, \((b)\) utilizes the \(4\)’th and \(6\)’th moment results 35 and 38 and we define \[\begin{align} &A_1=\alpha^4n^2(d+4)+\alpha^2\gamma^2n(2n+d+2)\\ &A_2=\alpha^2\gamma^2n(d+2)+\gamma^4n(d+n+1)\\ &A_3=\alpha^2n(d+2)+\gamma^2n. \end{align}\] Then combining all together results in \[\begin{align} {\cal{L}}(\boldsymbol{W})=A_1N_1+A_2N_2-2A_3N_3+n\sigma^2(\alpha^2N_1+\gamma^2N_2)+d+\sigma^2. \end{align}\] To find the optimal solution, set \(\nabla{\cal{L}}(\boldsymbol{W})=0\) and we obtain \[\begin{align} A_1\nabla N_1+A_2\nabla N_2-2A_3\nabla N_3+n\sigma^2(\alpha^2\nabla N_1+\gamma^2\nabla N_2)=0.\label{aug32diff321} \end{align}\tag{45}\] Note that we have \[\begin{align} &\nabla N_1=\nabla\left(\texttt{tr}\left(\boldsymbol{W}\right)^2+\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}^2\right)\right)=2\texttt{tr}\left(\boldsymbol{W}\right){\boldsymbol{I}}+2\boldsymbol{W}+2\boldsymbol{W}^\top\\ &\nabla N_2=\nabla\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)=2\boldsymbol{W}\\ &\nabla N_3=\nabla\texttt{tr}\left(\boldsymbol{W}\right)={\boldsymbol{I}}. \end{align}\] Therefore, 45 returns \[\begin{align} 2A_1\left(\texttt{tr}\left(\boldsymbol{W}\right){\boldsymbol{I}}+\boldsymbol{W}+\boldsymbol{W}^\top\right)+2A_2\boldsymbol{W}-2A_3+2n\sigma^2(\alpha^2(\texttt{tr}\left(\boldsymbol{W}\right){\boldsymbol{I}}+\boldsymbol{W}+\boldsymbol{W}^\top)+\gamma^2\boldsymbol{W}){\boldsymbol{I}}=0,\label{aug32diff322} \end{align}\tag{46}\] which implies that the optimal solution \(\boldsymbol{W}_\star\) has the form of \(c{\boldsymbol{I}}\) for some constant \(c\). Then suppose \(\boldsymbol{W}_\star=c{\boldsymbol{I}}\), we have \(\texttt{tr}\left(\boldsymbol{W}\right)=cd\) and 46 returns \[2A_1(d+2)c{\boldsymbol{I}}+2A_2c{\boldsymbol{I}}-2A_3{\boldsymbol{I}}+2n\sigma^2(\alpha^2(d+2)c{\boldsymbol{I}}+\gamma^2c{\boldsymbol{I}})=0\] \[\begin{align} \Longrightarrow c&=\frac{A_3}{A_1(d+2)+A_2+n\sigma^2(\alpha^2(d+2)+\gamma^2)}\\ &=\frac{\alpha^2(d+2)+\gamma^2}{\alpha^4n(d+2)(d+4)+\alpha^2\gamma^2(d+2)(d+2n+3)+\gamma^4(d+n+1)+\sigma^2(\alpha^2(d+2)+\gamma^2)}. \end{align}\] Then the optimal loss is obtained by setting \(\boldsymbol{W}_\star=c{\boldsymbol{I}}\) and \[\begin{align} {\cal{L}}_\star={\cal{L}}(\boldsymbol{W}_\star)&=A_1c^2d(d+2)+A_2c^2d-2A_3cd+n\sigma^2c^2d(\alpha^2(d+2)+\gamma^2)+d+\sigma^2\\ &=c^2d\left(A_1(d+2)+A_2+n\sigma^2(\alpha^2(d+2)+\gamma^2)\right)-2A_3cd+d+\sigma^2\\ &=d+\sigma^2-A_3cd. \end{align}\] It completes the proof of ?? . Now if assuming \(\alpha={\cal{O}}\left(1/\sqrt{d}\right)\), \(d/n={\cal{O}}\left(1\right)\) and sufficiently large dimension \(d\), we have the approximate \[\begin{align} c&\approx\frac{\alpha^2d+1}{\alpha^4d^2n+\alpha^2d(d+2n)+(d+n)+\sigma^2(\alpha^2d+1)}\\ &=\frac{\alpha^2d+1}{(\alpha^2d+1)^2n+(\alpha^2d+1)d+\sigma^2(\alpha^2d+1)}\\ &=\frac{1}{(\alpha^2d+1)n+d+\sigma^2} \end{align}\] and \[\begin{align} {\cal{L}}_\star&\approx d+\sigma^2-\frac{(\alpha^2d+1)nd}{(\alpha^2d+1)n+d+\sigma^2}. \end{align}\] height6pt width6pt depth0pt

1.0.2.4 Task-feature Alignment with \(\alpha\) Correlation

In this section, we consider the task-feature alignment data model similar to 15 , where we first draw task vector \({\boldsymbol{\beta}}\) via \[{\boldsymbol{\beta}}\sim{\cal{N}}(0,{\boldsymbol{I}}).\] Then we generate examples \((\boldsymbol{x}_i,y_i)_{i=1}^{n+1}\) according to the rule \(\text{corr\_coef}(\boldsymbol{x}_i,{\boldsymbol{\beta}})\geq \alpha\geq0\) via \[\begin{align} \boldsymbol{x}_i{~\big |~}{\boldsymbol{\beta}}\sim\mathcal{N}(\alpha{\boldsymbol{\beta}},{\boldsymbol{I}}),\quad\xi_i\sim{\cal{N}}(0,\sigma^2)\quad\text{and}\quad y_i=\gamma\cdot\boldsymbol{x}_i^\top {\boldsymbol{\beta}}+\xi_i,\label{data32feature32app} \end{align}\tag{47}\] which results in 15 by setting \(\gamma^2=1/(\alpha^2d+1)\).

Here, for clean notation and without loss of generality, we define and rewrite 47 via \[{\boldsymbol{g}}_i\sim{\cal{N}}(0,{\boldsymbol{I}}),\quad\xi_i\sim{\cal{N}}(0,\sigma^2)\quad\text{and}\quad\boldsymbol{x}_i=\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}_i,\quad y_i=\gamma\boldsymbol{x}_i^\top{\boldsymbol{\beta}}+\xi_i=\gamma\cdot(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}_i)^\top{\boldsymbol{\beta}}+\xi_i.\] Recap the loss function from 8 , we obtain \[\begin{align} {\cal{L}}(\boldsymbol{W})&=\operatorname{\mathbb{E}}\left[(y-g(\boldsymbol{Z}))^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\left(\gamma\boldsymbol{x}^\top{\boldsymbol{\beta}}+\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top(\gamma{\boldsymbol{X}}{\boldsymbol{\beta}}+\boldsymbol{\xi})\right)^2\right]\nonumber\\ &=\operatorname{\mathbb{E}}\left[\gamma^2(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2+2\gamma(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})+(\xi-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]\nonumber\\ &=\gamma^2\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]+\operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]+\sigma^2.\label{feature32loss32eq1} \end{align}\tag{48}\] Similar to Appendix 1.0.2.3, to begin with, let \[\begin{align} N_1 = \texttt{tr}\left(\boldsymbol{W}\right)^2+\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}^2\right),\quad N_2=\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right),\quad\text{and}\quad N_3=\texttt{tr}\left(\boldsymbol{W}\right),\nonumber \end{align}\] and additionally, given \({\boldsymbol{\Lambda}}_{\boldsymbol{W}}=\boldsymbol{W}\odot{\boldsymbol{I}}\), let \[\begin{align} N_4 = 3\texttt{tr}\left({\boldsymbol{\Lambda}}_{\boldsymbol{W}}^2\right)+(d+4)\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}^2\right).\nonumber \end{align}\] We first focus on the second term in 48 \[\begin{align} \operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top\boldsymbol{\xi})^2\right]&=\operatorname{\mathbb{E}}\left[\left((\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}})^\top\boldsymbol{W}\sum_{i=1}^n\xi_i(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}_i)\right)^2\right]\\ &=n\sigma^2\operatorname{\mathbb{E}}\left[\left((\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}})^\top\boldsymbol{W}(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}')\right)^2\right]\\ &=n\sigma^2\left(\alpha^4\operatorname{\mathbb{E}}\left[({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}})^2\right]+2\alpha^2\operatorname{\mathbb{E}}\left[({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}')^2\right]+\operatorname{\mathbb{E}}\left[({\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}')^2\right]\right)\\ &=n\sigma^2\left(\alpha^4\left(\texttt{tr}\left(\boldsymbol{W}\right)^2+\texttt{tr}\left(\boldsymbol{W}^2\right)+\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)\right)+(2\alpha^2+1)\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)\right)\\ &=n\sigma^2\left(\alpha^4N_1+(2\alpha^2+1)N_2\right). and independence of {\boldsymbol{\beta}},{\boldsymbol{g}},{\boldsymbol{g}}'.} \end{align}\]

Next, the first term of 48 (omitting \(\gamma^2\)) returns the following decomposition: \[\begin{align} \operatorname{\mathbb{E}}\left[(\boldsymbol{x}^\top{\boldsymbol{\beta}}-\boldsymbol{x}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]&=\operatorname{\mathbb{E}}\left[((\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}})^\top({\boldsymbol{\beta}}-\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}))^2\right]\\ &=\operatorname{\mathbb{E}}\left[\left(\alpha{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}-\alpha{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}+{\boldsymbol{g}}^\top{\boldsymbol{\beta}}-{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}\right)^2\right]\\ &=\alpha^2\operatorname{\mathbb{E}}[({\boldsymbol{\beta}}^\top{\boldsymbol{\beta}})^2]+\alpha^2\operatorname{\mathbb{E}}[({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]+\operatorname{\mathbb{E}}[({\boldsymbol{g}}^\top{\boldsymbol{\beta}})^2]+\operatorname{\mathbb{E}}[({\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]\\ &\quad-2\alpha^2\operatorname{\mathbb{E}}[{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}]-2\operatorname{\mathbb{E}}[{\boldsymbol{\beta}}^\top{\boldsymbol{g}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}]\\ &=\alpha^2d(d+2)+\alpha^2\underset{(a)}{\underbrace{\operatorname{\mathbb{E}}[({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]}}+d+\underset{(b)}{\underbrace{\operatorname{\mathbb{E}}[({\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]}}\\ &\quad-2\alpha^2\underset{(c)}{\underbrace{\operatorname{\mathbb{E}}[{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}]}}-2\underset{(d)}{\underbrace{\operatorname{\mathbb{E}}[{\boldsymbol{\beta}}^\top{\boldsymbol{g}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}]}}. \end{align}\] Consider solving \((a)\)-\((d)\) sequentially as follows:

To begin with, we use the following decomposition for all \((a)\)-\((d)\): \[\begin{align} {\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}&=\sum_{i=1}^n\boldsymbol{x}_i\boldsymbol{x}_i^\top{\boldsymbol{\beta}}\\ &=\sum_{i=1}^n(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}_i)(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}_i)^\top{\boldsymbol{\beta}}\\ &=\sum_{i=1}^n\alpha^2{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}+\alpha{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}+\alpha{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}+{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}. \end{align}\] Then, we have \[\begin{align} (a):\quad&\operatorname{\mathbb{E}}[({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2]\nonumber\\ &=\scalebox{0.85}{\displaystyle \operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n\alpha^2{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}+\alpha{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}+\alpha{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}+{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]}\nonumber\\ &=\scalebox{0.85}{\displaystyle \alpha^4n^2\operatorname{\mathbb{E}}\left[\left({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right)^2\right]+\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]}\nonumber\\ &\scalebox{0.85}{\displaystyle \quad+2\alpha^2n\operatorname{\mathbb{E}}\left[\sum_{i=1}^n{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right]+2\alpha^2\operatorname{\mathbb{E}}\left[\sum_{i=1}^n{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right]}\nonumber\\ &=\scalebox{0.85}{\displaystyle \alpha^4n^2\operatorname{\mathbb{E}}\left[\left({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2n\operatorname{\mathbb{E}}\left[\left({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2n\operatorname{\mathbb{E}}\left[\left({\boldsymbol{\beta}}^\top\boldsymbol{W}{{\boldsymbol{g}}'}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right)^2\right]+\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]}\nonumber\\ &\scalebox{0.85}{\displaystyle \quad+2\alpha^2n^2\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}'{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}\right]+2\alpha^2n\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right]}\nonumber\\ &=\alpha^4n^2(d+4)(d+6)N_1+\alpha^2n(d+4)N_1+\alpha^2n(d+2)(d+4)N_2\tag{49}\\ &\quad+n(n-1)N_1+nN_4\tag{50}\\ &\quad+2\alpha^2n^2(d+4)N_1+2\alpha^2n(d+4)N_1\tag{51}\\ &=\left(\alpha^2n(d+4)(\alpha^2n(d+6)+2n+3)+n(n-1)\right)N_1+\alpha^2n(d+2)(d+4)N_2+nN_4\\ &=B_1N_1+B_2N_2+nN_4,\nonumber \end{align}\] where 49 and 51 utilize 38 and 40 , and 50 is obtained via \[\begin{align} \operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]&=n\operatorname{\mathbb{E}}\left[\left({\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}'{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}\right)^2\right]+n(n-1)\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}'{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{g}}''{{\boldsymbol{g}}''}^\top{\boldsymbol{\beta}}\right]\\ &=nN_4+n(n-1)N_1, \end{align}\] which follows 35 and 36 .

\[\begin{align} (b):\quad&\operatorname{\mathbb{E}}\left[({\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}})^2\right]\nonumber\\ &=\scalebox{0.85}{\displaystyle \operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n\alpha^2{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}+\alpha{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}+\alpha{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}+{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]}\nonumber\\ &=\scalebox{0.85}{\displaystyle \alpha^4n^2\operatorname{\mathbb{E}}\left[\left({\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right)^2\right]+\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]}\nonumber\\ &\scalebox{0.85}{\displaystyle \quad+2\alpha^2n\operatorname{\mathbb{E}}\left[\sum_{i=1}^n{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right]+2\alpha^2\operatorname{\mathbb{E}}\left[\sum_{i=1}^n{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right]}\nonumber\\ &=\scalebox{0.85}{\displaystyle \alpha^4n^2\operatorname{\mathbb{E}}\left[\left({\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2n\operatorname{\mathbb{E}}\left[\left({\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}\right)^2\right]+\alpha^2n\operatorname{\mathbb{E}}\left[\left({\boldsymbol{g}}^\top\boldsymbol{W}{{\boldsymbol{g}}'}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right)^2\right]+\operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]}\nonumber\\ &\quad\scalebox{0.85}{\displaystyle +2\alpha^2n^2\operatorname{\mathbb{E}}\left[{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}'{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}\right]+2\alpha^2n\operatorname{\mathbb{E}}\left[{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right]}\nonumber\\ &=\alpha^4n^2(d+2)(d+4)N_2+\alpha^2n(d+2)N_2+\alpha^2nd(d+2)N_2+n(d+n+1)N_2\tag{52}\\ &\quad+2\alpha^2n^2(d+2)N_2+2\alpha^2n(d+2)N_2\tag{53}\\ &=\left(\alpha^2n(d+2)(\alpha^2n(d+4)+2n+d+3)+n(d+n-1)\right)N_2\nonumber\\ &=B_3N_2,\nonumber \end{align}\] where 52 and 53 are obtained using 35 , 38 and \[\begin{align} \operatorname{\mathbb{E}}\left[\left(\sum_{i=1}^n{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}_i{\boldsymbol{g}}_i^\top{\boldsymbol{\beta}}\right)^2\right]&=n\operatorname{\mathbb{E}}\left[\left({\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}'{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}\right)^2\right]+n(n-1)\operatorname{\mathbb{E}}\left[{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}'{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}''{{\boldsymbol{g}}''}^\top{\boldsymbol{\beta}}\right]\\ &=n(d+2)N_2+n(n-1)N_2=n(n+d+1)N_2. \end{align}\] \[\begin{align} (c):\quad&\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}\right]\\ &=n\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}')(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}')^\top{\boldsymbol{\beta}}\right]\\ &=\alpha^2n\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right]+n\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top\boldsymbol{W}{{\boldsymbol{g}}'}{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}\right]\\ &=\alpha^2n(d+2)(d+4)\texttt{tr}\left(\boldsymbol{W}\right)+n(d+2)\texttt{tr}\left(\boldsymbol{W}\right)\\ &=\left(\alpha^2n(d+2)(d+4)+n(d+2)\right)N_3\\ &=B_4N_3. \end{align}\] \[\begin{align} (d):\quad&\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{g}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{X}}^\top{\boldsymbol{X}}{\boldsymbol{\beta}}\right]\\ &=n\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{g}}{\boldsymbol{g}}^\top\boldsymbol{W}(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}')(\alpha{\boldsymbol{\beta}}+{\boldsymbol{g}}')^\top{\boldsymbol{\beta}}\right]\\ &=\alpha^2n\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{g}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{\beta}}{\boldsymbol{\beta}}^\top{\boldsymbol{\beta}}\right]+n\operatorname{\mathbb{E}}\left[{\boldsymbol{\beta}}^\top{\boldsymbol{g}}{\boldsymbol{g}}^\top\boldsymbol{W}{\boldsymbol{g}}'{{\boldsymbol{g}}'}^\top{\boldsymbol{\beta}}\right]\\ &=\alpha^2n(d+2)\texttt{tr}\left(\boldsymbol{W}\right)+n\texttt{tr}\left(\boldsymbol{W}\right)\\ &=\left(\alpha^2n(d+2)+n\right)N_3\\ &=B_5N_3. \end{align}\] Here we define \[\begin{align} B_1&=\alpha^2n(d+4)(\alpha^2n(d+6)+2n+3)+n(n-1)\\ B_2&=\alpha^2n(d+2)(d+4)\\ B_3&=\alpha^2n(d+2)(\alpha^2n(d+4)+2n+d+3)+n(d+n-1)\\ B_4&=\alpha^2n(d+2)(d+4)+n(d+2)\\ B_5&=\alpha^2n(d+2)+n. \end{align}\] Then combining all together results in \[\begin{align} {\cal{L}}(\boldsymbol{W})&=\scalebox{0.85}{\displaystyle \gamma^2\left(\alpha^2d(d+2)+d+\alpha^2(B_1N_1+B_2N_2+nN_4)+B_3N_2-2\alpha^2 B_4N_3-2 B_5N_3\right)+n\sigma^2(\alpha^4N_1+(2\alpha^2+1)N_2)+\sigma^2}\\ &=\scalebox{0.85}{\displaystyle \gamma^2\left(\alpha^2B_1N_1+(\alpha^2B_2+B_3)N_2-2(\alpha^2B_4+B_5)N_3 +\alpha^2nN_4\right)+n\sigma^2(\alpha^4N_1+(2\alpha^2+1)N_2)+\gamma^2d\left(\alpha^2(d+2)+1\right)+\sigma^2} \end{align}\] and differentiating it results in \[\nabla{\cal{L}}(\boldsymbol{W})=\gamma^2\left(\alpha^2B_1\nabla N_1+(\alpha^2B_2+B_3)\nabla N_2-2(\alpha^2B_4+B_5)\nabla N_3 +\alpha^2n\nabla N_4\right)+n\sigma^2(\alpha^4\nabla N_1+(2\alpha^2+1)\nabla N_2).\] Similar to the proof in Appendix 1.0.2.3, \(\boldsymbol{W}_\star\) has the form of \(\boldsymbol{W}_\star=c{\boldsymbol{I}}\) and we have \[\begin{align} \nabla N_1&=\nabla\left(\texttt{tr}\left(\boldsymbol{W}\right)^2+\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}^2\right)\right)=2\texttt{tr}\left(\boldsymbol{W}\right){\boldsymbol{I}}+2\boldsymbol{W}+2\boldsymbol{W}^\top=2c(d+2){\boldsymbol{I}}\\ \nabla N_2&=\nabla\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)=2\boldsymbol{W}=2c{\boldsymbol{I}}\\ \nabla N_3&=\nabla\texttt{tr}\left(\boldsymbol{W}\right)={\boldsymbol{I}}\\ \nabla N_4&=\nabla\left(3\texttt{tr}\left({\boldsymbol{\Lambda}}_{\boldsymbol{W}}^2\right)+(d+4)\texttt{tr}\left(\boldsymbol{W}\boldsymbol{W}^\top\right)+\texttt{tr}\left(\boldsymbol{W}^2\right)\right)\\ &=6\cdot\text{diag}\left({\boldsymbol{\Lambda}}_{\boldsymbol{W}}\right)+2(d+4)\boldsymbol{W}+2\boldsymbol{W}^\top\\ &=2c(d+8){\boldsymbol{I}}. \end{align}\] Therefore, setting \(\nabla{\cal{L}}(\boldsymbol{W})=0\) returns \[\gamma^2\left(2c(d+2)\alpha^2B_1+2c(\alpha^2B_2+B_3)-2(\alpha^2B_4+B_5) +2c(d+8)\alpha^2n\right)+2cn\sigma^2(\alpha^4(d+2)+2\alpha^2+1)=0\] \[\begin{align} \Longrightarrow c&=\frac{\alpha^2 B_4+B_5}{(d+2)\alpha^2B_1+(\alpha^2B_2+B_3) +(d+8)\alpha^2n+n\sigma^2(\alpha^4(d+2)+2\alpha^2+1)/\gamma^2}\\ &=\scalebox{0.7}{\displaystyle \frac{\alpha^4n(d+2)(d+4)+2\alpha^2n(d+2)+n}{\alpha^6n^2(d+2)(d+4)(d+6)+\alpha^4n(d+2)(d+4)(3n+4)+\alpha^2n((d+2)(3n+d+3)+(d+8))+n(d+n+1)+n\sigma^2(\alpha^4(d+2)+2\alpha^2+1)/\gamma^2}}\\ &=\scalebox{0.7}{\displaystyle \frac{\alpha^4(d+2)(d+4)+2\alpha^2(d+2)+1}{\alpha^6n(d+2)(d+4)(d+6)+\alpha^4(d+2)(d+4)(3n+4)+\alpha^2((d+2)(3n+d+3)+(d+8))+(d+n+1)+\sigma^2(\alpha^4(d+2)+2\alpha^2+1)/\gamma^2}}. \end{align}\] Then the optimal loss is obtained by setting \(\boldsymbol{W}_\star=c{\boldsymbol{I}}\) and \[{\cal{L}}_\star={\cal{L}}(\boldsymbol{W}_\star)=\gamma^2d(\alpha^2(d+2)+1)+\sigma^2-\gamma^2(\alpha^2B_4+B_5) cd.\] It completes the proof of ?? . Now if assuming \(\alpha={\cal{O}}\left(1/\sqrt{d}\right), d/n={\cal{O}}\left(1\right)\), \(\gamma^2=1/(\alpha^2d+1)\) and sufficiently large dimension \(d\), we have the approximate \[\begin{align} c&\approx\frac{\alpha^4d^2+2\alpha^2d+1}{n\alpha^6d^3+3n\alpha^4d^2+(3n+d)\alpha^2d+d+n+\sigma^2(\alpha^4d+2\alpha^2+1)/\gamma^2}\\ &\approx\frac{(\alpha^2d+1)^2}{n(\alpha^2d+1)^3+d(\alpha^2d+1)+\sigma^2(\alpha^2d+1)}\\ &\approx\frac{1}{(\alpha^2d+1)n+(d+\sigma^2)/(\alpha^2d+1)} \end{align}\] and \[\begin{align} {\cal{L}}_\star&\approx \gamma^2d(\alpha^2d+1)+\sigma^2-\frac{\gamma^2(\alpha^2d+1)^2nd}{(\alpha^2d+1)n+(d+\sigma^2)/(\alpha^2d+1)}\\ &=d+\sigma^2-\frac{(\alpha^2d+1)nd}{(\alpha^2d+1)n+(d+\sigma^2)/(\alpha^2d+1)}. \end{align}\] height6pt width6pt depth0pt

1.0.3 Analysis of Low-Rank Parameterization↩︎

1.0.3.1 Proof of Lemma 3

Recall the loss function from 42 \[\begin{align} {\cal{L}}(\boldsymbol{W})=M-2n\texttt{tr}\left(\boldsymbol{\Sigma}\boldsymbol{\bar{W}}\right)+n(n+1)\texttt{tr}\left(\boldsymbol{\Sigma}\boldsymbol{\bar{W}}^\top\boldsymbol{\bar{W}}\right)+nM\texttt{tr}\left(\boldsymbol{\bar{W}}\boldsymbol{\bar{W}}^\top\right) \end{align}\] where \(\boldsymbol{\bar{W}}=\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\boldsymbol{W}\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\), \(\boldsymbol{\Sigma}=\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}\boldsymbol{\Sigma}_{\boldsymbol{x}}^{1/2}\) and \(M=\texttt{tr}\left(\boldsymbol{\Sigma}\right)+\sigma^2\). For any \(\boldsymbol{\bar{W}}\), let us parameterize \(\boldsymbol{\bar{W}}={\boldsymbol{U}}{\boldsymbol{E}}{\boldsymbol{U}}^\top\) where \({\boldsymbol{U}}\in\mathbb{R}^{d\times r}\) denotes the eigenvectors of \(\boldsymbol{\bar{W}}\) and \({\boldsymbol{E}}\in\mathbb{R}^{r\times r}\) is a symmetric square matrix. We will first treat \({\boldsymbol{U}}\) as fixed and optimize \({\boldsymbol{E}}\). We will then optimize \({\boldsymbol{U}}\). Fixing \({\boldsymbol{U}}\), setting \(\boldsymbol{\bar{\Sigma}}={\boldsymbol{U}}^\top\boldsymbol{\Sigma}{\boldsymbol{U}}\), we obtain \[\begin{align} {\cal{L}}({\boldsymbol{E}})&=M-2n\texttt{tr}\left(\boldsymbol{\bar{\Sigma}}{\boldsymbol{E}}\right)+n(n+1)\texttt{tr}\left(\boldsymbol{\bar{\Sigma}}{\boldsymbol{E}}^2\right)+nM\texttt{tr}\left({\boldsymbol{E}}^2\right). \end{align}\] Differentiating, we obtain \[0.5n^{-1}\nabla{\cal{L}}({\boldsymbol{E}})=-\boldsymbol{\bar{\Sigma}}+(n+1)\boldsymbol{\bar{\Sigma}}{\boldsymbol{E}}+M{\boldsymbol{E}}.\] Setting \(\nabla{\cal{L}}({\boldsymbol{E}})=0\) returns \[\begin{align} {\boldsymbol{E}}_\star=(M{\boldsymbol{I}}+(n+1)\boldsymbol{\bar{\Sigma}})^{-1}\boldsymbol{\bar{\Sigma}}.\label{hxbidnkj} \end{align}\tag{54}\] Let \(\bar{\lambda}_i\) denote the \(i\)’th largest eigenvalue of \(\boldsymbol{\bar{\Sigma}}\). Plugging in this value, we obtain the optimal risk as a function of \({\boldsymbol{U}}\) is given by \[\begin{align} {\cal{L}}_\star({\boldsymbol{U}})&=M-n\cdot\texttt{tr}\left(\boldsymbol{\bar{\Sigma}}{\boldsymbol{E}}_\star\right)=M-n\cdot\texttt{tr}\left((M{\boldsymbol{I}}+(n+1)\boldsymbol{\bar{\Sigma}})^{-1}\boldsymbol{\bar{\Sigma}}^2\right)\\ &=M-n\sum_{i=1}^r \frac{\bar{\lambda}_i^2}{(n+1)\bar{\lambda}_i+M}=M-n\sum_{i=1}^r \frac{\bar{\lambda}_i}{n+1+M\bar{\lambda}_i^{-1}}. \end{align}\] Now observe that, the right hand side is strictly decreasing function of the eigenvalues \(\bar{\lambda}_i\) of \(\boldsymbol{\bar{\Sigma}}={\boldsymbol{U}}^\top\boldsymbol{\Sigma}{\boldsymbol{U}}\). Thus, to minimize \({\cal{L}}_\star({\boldsymbol{U}})\), we need to maximize \(\sum_{i=1}^r \frac{\bar{\lambda}_i}{n+1+M\bar{\lambda}_i^{-1}}\). It follows from Cauchy interlacing theorem that \(\bar{\lambda}_j\leq \lambda_i\) where \(\lambda_i\) is the \(i\)’th largest eigenvalue of \(\boldsymbol{\Sigma}\) since \(\boldsymbol{\bar{\Sigma}}\) is an orthogonal projection of \(\boldsymbol{\Sigma}\) on \({\boldsymbol{U}}\). Consequently, we find the desired bound where \[{\cal{L}}_\star=M-n\sum_{i=1}^r \frac{\lambda_{i}}{n+1+M\lambda_{i}^{-1}}.\] The equality holds by setting \({\boldsymbol{U}}\) to be the top-\(r\) eigenvectors of \(\boldsymbol{\Sigma}\) and \({\boldsymbol{E}}={\boldsymbol{E}}_\star({\boldsymbol{U}})\) to be the diagonal matrix according to [ebst32eq]. height6pt width6pt depth0pt

References↩︎