Paper: Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data (ICLR-2024)

Reviewer comment (Theory):
The representation dimension $r$ is required to be at most $\min(d_x, d_y)$, where $d_x$ is the dimension of the covariates and $d_y$ is the dimension of the observations (Section 2, just after equation (1)). In the linear regression setting, this would mean that there must be a one-dimensional representation. This is a very strong assumption. It seems like we should still be able to obtain some benefit if the representation only has a lower dimension than the _covariates_. This would more closely mirror practical settings such as e.g. computer vision, where the data are assumed to belong to a lower-dimensional manifold.

Author response:
The $r \leq \min\{d_y, d_x\}$ is a typo; as one should expect, we only require $r \leq d_x$ for our analysis and can accommodate arbitrary $d_y$, e.g. $d_y = 1 \ll d_x$ for scalar-output linear regression, or $d_y \approx d_x$ in linear system identification (Appendix B.1). We thank the reviewer for identifying this important oversight.

Reviewer comment (Theory):
Assumption 3.1: It is assumed that the $\beta$-mixing coefficient follows an _exact_ geometric decay, i.e., $\beta^{(t)}(k) = \Gamma^{(t)} \mu^{(t)k}$ for each task $t$. This should place strong restrictions on the possible types of covariate trajectory distributions. It seems like we should expect the results if the decay is _at least_ geometric in nature, i.e., $\beta(k) \leq \Gamma \mu^k$ for some $\mu < 1$.

Author response:
The beta-mixing coefficient only needs to be upper bounded by geometric decay. This will be fixed in the revision.

Reviewer comment (Experiments):
The empirical results would be more convincing at showing a fundamental limit on the accuracy for AMD if final accuracy vs. number of tasks was shown at a fixed sample size per task, and showing that this accuracy does not approach 0 as the number of tasks increases. At present, it is just shown for T=25. While DFW does converge in this scenario, in principle, it could just be that DFW has a better sample complexity, but AMD will still eventually converge given enough tasks, albeit at a slower rate. Adding this experiment would strengthen the paper.

Author response:
**Regarding the requested experiment for alternating minimization-descent (AMD)**: we took the same experimental parameters for linear representations as described in Section 4.1 (corresponding to Fig 1a), except varying the number of total tasks $T \in \{1, 5, 10, 25, 50, 100\}$. For each $T$, we run alternating minimization-descent for a single run of 5000 iterations ($\gg 100$ iterations shown in Fig 1a and more than in the MLP experiment) and record the subspace distances. For comparison, we run DFW with a single task $T=1$. This generates the following figure ([anonymous google drive link](https://drive.google.com/file/d/1ztBpxcl0MTcoupD8L_DZ1gsRENrCQPGr/view?usp=sharing)). The plot demonstrates that: 1. increasing the number of tasks *cannot push AMD past a (large) threshold*, only serving to reduce the variance, precisely as our theory predicts, 2. *AMD does not encounter a feature learning phase* as does an MLP (Figure 1b) that drives the distance to $0$, even after many iterations. In comparison, DFW even for a single task quickly converges to $0$, confirming that the suboptimality of AMD is due to a fundamental bias, rather than the noise level of the problem.

Reviewer comment (Presentation):
A minor point: the title of the OpenReview submission does not match the title on the paper. This should be fixed.

Author response:
The most updated title of our paper is the one on OpenReview: “Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data”. We apologize for the confusion.

Reviewer comment (Theory):
I am curious why the required number of samples $N$ grows (moderately) with the number of tasks $T$. I assume this is to enforce some sort of uniform bound on the random fluctuations across all of the tasks. Can the authors confirm if this intuition is correct?

Author response:
**Regarding the moderate growth of $N$ with respect to tasks $T$**: this is a question with some subtlety. There are two sources to the $\log(T)$ dependence–one that arises in the analysis of the iid setting, and the other from $\beta$-mixing.
   - In the iid setting, as the reviewer suspects, the $\log(T)$ dependence comes from an eventual union bound (a.k.a. uniform control) over the noise processes. In particular, to bound the largest singular value of the average of self-normalized martingale (SNM) noise across tasks $T$, we appeal to a Matrix Hoeffding inequality (Lemma A.4). However, like the traditional Hoeffding inequality, this requires boundedness in the psd order of each summand (i.e. each task’s SNM). Since we have concentration inequalities for individual SNMs (Propositions A.1 and A.2), a simple way to simulate boundedness is to condition on the high-probability boundedness event of each task’s SNM. However, conditioning simultaneously on each task’s boundedness equates to a union bound, which leads to a $\log(T)$ factor when inverted for a desired failure probability $\delta$. The $\log(T)$ factor cannot be avoided when going through the route of boundedness/truncation even though SNMs are independent between different tasks (see e.g. subgaussian maximal inequalities). *However*, we conjecture that this $\log(T)$ factor is only technical, and *may be removable in the iid setting*. This requires going from a (matrix) Hoeffding to a Bernstein-type inequality, which in turn requires control of the higher-order moments of the SNM. As far as we can tell, this would require highly non-trivial, novel analysis of self-normalized martingales, and importantly, **the mild savings will be washed away (order-wise) in the $\beta$-mixing setting**, where they arise for a fundamental reason. (continued in next comment)

Reviewer comment (Theory):
Is there some intuition for why the representation dimension $r$ must be smaller than both the covariate _and_ measurement dimensions? If this is a necessary assumption, can the authors comment on how they would justify this restriction, especially in the linear regression case when $d_y=1$?

Author response:
The $r \leq \min\{d_y, d_x\}$ is a typo; as one should expect, we only require $r \leq d_x$ for our analysis and can accommodate arbitrary $d_y$, e.g. $d_y = 1 \ll d_x$ for scalar-output linear regression, or $d_y \approx d_x$ in linear system identification (Appendix B.1). We thank the reviewer for identifying this important oversight.

Reviewer comment (Theory):
Do the results still hold if the equality for $\beta(k)$ in Assumption 3.1 is replaced with an inequality?

Author response:
The beta-mixing coefficient only needs to be upper bounded by geometric decay. This will be fixed in the revision.

Reviewer comment (Evaluation):
It is very interesting that the use of an MLP allows the original AMD algorithm to overcome the fundamental lower bound on the error present when learning a linear representation (even if the sample complexity is much worse than DFW). Is this just because the quantity being measured (validation loss instead of subspace distance) is different, or would AMD with a linear representation fail to converge to 0 validation loss in this setting? If this is particular to the MLP representation, do the authors have any intuition for why this might be the case?

Author response:
**Regarding the MLP representation experiments**: as suggested by the additional experiment described earlier, a linear parameterization in conjunction with vanilla alternating minimization-descent does not converge no matter how much data is provided. The ability of the MLP to break past this barrier likely extends beyond the measured quantity (validation loss), since the validation set is sampled from a distinct task that is not used in training. In terms of why the MLP parameterization is eventually able to descend to optimality: we do not claim to be experts on the theory of neural network optimization, but a simple hypothesis is that the MLP is expressive enough to find a non-linear parameterization that leads to a path that escapes the local minima, which seems to match the folk understanding that SGD tends to find globally near-optimal solutions for many non-convex objectives in deep learning (e.g. [5]). We generated a figure ([anonymized google drive link](https://drive.google.com/file/d/1h7zpLqaLmCVqPSKBLCWcOJoHHub8ffgB/view?usp=sharing)) that shows the norm of the MLP gradient evolving in parallel to Figure 1b, which demonstrates that SGD does indeed approach a local minima before finding a path out toward the global minima, reflected by the noticeable bump in the gradient norm mirroring the generalization phase in Figure 1b.

Reviewer comment (Presentation):
The title "META-LEARNING OPERATORS TO OPTIMALITY FROM MULTI-TASK NON-IID DATA" is so vague. It is really hard to understand what this paper studies from the title. It should indicate that the goal is to learn the shared parameter $\Phi$.

Author response:
The title in the submitted pdf is in error. The most updated version is the one seen on the OpenReview page “Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data“, which clarifies the goal of learning a shared representation.

Reviewer comment (Evaluation):
The failed example given in Section 3.1 serves as the main motivation of introducing new algorithms. However, these two crucial issues in this example are not really resolved. I am concerned if the de-bias and feature-whiten steps could really resolve these issues. I put more comments in the next section.

Author response:
In general, **regarding the possible correlation due to sampling from the same $\beta$-mixing process**, our main technical workhorse is Lemma A.3, which is a standard tool used in the analysis of mixing processes. In short, the expected value of a bounded measurable function of a given $\beta$-mixing process and the same function on the iid version of the process (where each covariate $x_i$, $i=1,\dots,$ is sampled from its population distribution), can be bounded in terms of the mixing function $\beta(\cdot)$. In particular, if said function is an indicator of an event of interest, e.g. the complement of the descent guarantee Eq (11) in Theorem 3.1, then this implies the probability of the event occurring on the mixing process versus the independent version is the same up to an additive factor that depends on the mixing function. Therefore, the iid results hold for the mixing case, albeit with a smaller effective sample size in order to achieve the same failure rate $\delta$. For geometric mixing processes, this essentially amounts to reducing $N \to N/\log(N)$ samples. This addresses a couple of questions:

Reviewer comment (Theory):
First about clarifying the key idea de-bias and feature-whiten methods. In Section 3.2, it says that $\hat{F}^{(t)}$ is computed on independent data. It is not clear to me why $\hat{F}^{(t)}$ is independent from $X^{(t)}$. To my understanding, for example, the Partition trajectories step (Line 5, Algorithm 1) splits the dataset $N$ to $N_{1}=\{x_1, x_2, \dots, x_n \}$ and $N_{2}=\{x_{n+1}, x_{n+2}, \dots, x_{n+N} \}$. But they come from the same $\beta$-mixing stochastic process, will they become independent?

Author response:
**Regarding independence of $X^{(t)}$ and $\hat F^{(t)}$: in the mixing case**, they are indeed possibly not independent. However, as discussed, their behavior can be bounded by the analysis of the iid case, paying the appropriate costs for dependency in the mixing coefficients. Similarly, the decomposition of $F, W, X\Sigma^{-1}$ need only hold for the iid setting.

Reviewer comment (Writing):
What is "the aforementioned batching strategy" mentioned in Section 3.2 right after Eq.(5)? It seems that there is no batching strategy mentioned before.

Author response:
**The “aforementioned batching strategy” mentioned after Eq (5)** refers to the discussion just preceding Eq (5): “each agent computes the least squares weights $\hat F^{(t)}$ and the representation update on independent batches of data, e.g. disjoint subsets of trajectories”. We will adjust the wording in the revision, as we indeed do not discuss a concrete strategy until the subsequent paragraphs.

Reviewer comment (Theory):
Regarding the proof for the non-iid case, I am mainly concerned if the iid case could be simply immigrated to the non-iid case. For example, on page 16 of the supplementary material, it says "We observe that since $\hat{F}(t)$ is by construction independent of ... " and obtains $$E[FWX\Sigma^{-1}]=E[F]E[W]E[X\Sigma^{-1}]$$ This equation won't hold for the non-iid case. It is because $F$ here is evaluated using a part of the process $\{x\}$ and $\Sigma$ is estimated using another part of the same process.

Author response:
**Regarding independence of $X^{(t)}$ and $\hat F^{(t)}$: in the mixing case**, they are indeed possibly not independent. However, as discussed, their behavior can be bounded by the analysis of the iid case, paying the appropriate costs for dependency in the mixing coefficients. Similarly, the decomposition of $F, W, X\Sigma^{-1}$ need only hold for the iid setting.