Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data

We thank the reviewer for their detailed comments. To address the remaining questions/concerns:

• 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.

• 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.

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

• 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). 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.

• 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)

We thank the reviewer for their evaluation. To address the remaining questions:

• We are introducing a notation table to our revised appendix.

• Regarding the burn-in requirement $\Omega(r)$: the requirement in Tripuraneni et al. (and similar papers) is indeed $\Omega(1)$. However, we strongly emphasize that the proposed algorithm in that paper critically relies on the fact that the covariates across all tasks are iid and isotropic. From an algorithmic perspective, their proposed method-of-moments estimator is designed with the implicit foreknowledge that $\mathbb E[x x^\top] = I$, and is only guaranteed to be consistent then. On the other hand, their proposed lower bound on the subspace distance, by assuming all covariates are iid $\mathcal N(0,I)$, in fact does not even require $N > 0$ for any given task, as long as the total samples $NT$ is large enough. However, this quirk is inextricably tied to the iid covariate assumption, which is reflected in the proof of the result (Lemma 22 of their appendix). Intuitively, we should not expect $N = \Theta(1)$ to be possible, in the sense that allowing the covariances across each task to differ permits the design of weights $F^{(t)}$ and covariances $\Sigma_x^{(t)}$ that are arbitrarily ill-conditioned such that distinguishing $\Phi_\star$ and a perturbed $\Phi’_\star$ from sampled-data is provably hard unless a sufficient per-task burn-in on $N$ (e.g. proportional to the row-rank of $\Phi$) is satisfied. Formalizing this lower bound, as well as extracting the task-diversity/task-coverage assumptions of the task-wise covariances $\Sigma_x^{(t)}$ that do not manifest in the task-wise iid setting is ongoing work.

• We agree that random features in addition to random covariate design is an interesting consideration. We note that when we set the weights $F^{(t)}$ to be random (e.g. Gaussian random matrix), keeping $\Phi_\star$ as deterministic (otherwise the low-rank representation learning problem changes), our analysis actually goes through mutatis mutandis: in particular, instead of deterministic task diversity parameters $\lambda_{\min}^{\mathbf F}, \lambda_{\max}^{\mathbf F}$ (Definition 3.2), we can replace them with high-probability variants, carrying out the rest of the analysis as written and simply accruing the additional failure probability in the final bound. Regarding the alignment of covariances, we note some subtlety in Theorem 3.1 that was washed out for simplicity: the dependence on task-specific quantities, such as the smallest covariance eigenvalue $\lambda_{\min}(\Sigma^{(t)}_x)$, are actually averaged across tasks in a certain way, where they currently appear as a max over tasks $t$ for simplicity. After inverting for sample complexity bounds as in Corollary 3.1, these task "coverage" / "overlap" quantities naturally manifest in the bound, which may be seen as notions of how the alignment and/or coverage of source task covariances affect learning. Lastly, a theory for an overparameterized setting is quite interesting, and seems to be getting concurrent attention; however, the data / feature assumptions are necessarily quite distinct from ours, and thus we cannot claim to easily extrapolate the trends in our results there.

We thank the reviewer for their detailed comments. To address their questions:

• 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.

• 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:

  • The discussion of the “blocking” technique is located around Definition A.1 ($\beta$-mixing) and Lemma A.3. We will make references to blocking clearer in our revision.
  • 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.

• 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.

In light of this discussion, we hope that this answers the reviewer’s concerns about whether our proposed algorithm and analyses fully overcome the identified non-isotropy and non-iid issues, as far as the $\beta$-mixing assumption holds.

We thank the reviewer for their comments. To address the reviewer’s questions:

• 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.

• Regarding Corollary 3.1, the requisite partitions are actually very simple to construct in practice. In short, given the assumptions of Theorem 3.1, we are guaranteed the distance-to-optimality of the next iterate is decomposed into a contraction of the previous iterate’s distance and an additive noise term. Therefore, the idea that yields Corollary 3.1 is to simply partition a given offline dataset into exponentially growing chunks, such that the noise term is sufficiently small with respect to the contraction. Namely, if $d_{t+1} \leq \rho \cdot d_t + \frac{\sigma}{\sqrt{n}}$, then any schedule for $n$ that ensures $d_{t+1} \leq \rho’ \cdot d_t$, $\rho’ < 1$, $t =0,1,\dots$ suffices. For an offline dataset of size $N$, an exponential growing partition can be accommodated $K = O(\log N)$ times, which yields $d_K \lesssim C(\rho’) \frac{\sigma}{\sqrt{N}}$, hence approximating “ERM-like” rates of $\approx \frac{\sigma}{\sqrt{N}}$. We note that this is a straightforward adaptation of the standard “doubling trick” in online learning, and discussed in further detail in Lemma A.11. Any exponential schedule can be used in principle, as long as contraction is enforced (visualized in practice, for example, via a validation loss), without hurting the rate of the final bound.

• Regarding tightness of dependence on various problem quantities. Starting with our per-task sample requirement $N \geq \Omega(r + \sigma_w^2 \max\{d_y, r\})$, we do not provide a formal minimax lower bound in this paper, but we have reasons to believe it cannot be significantly improved, e.g. $N = \mathcal O(1)$, in general. For $d_y = 1$, the lower bound of $N = \mathcal O(1)$ provided by Tripuraneni et al. critically depends on the covariates being iid $\mathcal N(0,I)$ across all tasks. Allowing covariances to change across tasks opens the door to arbitrarily ill-conditioned interactions between weights $F^{(t)}$ and covariances $\Sigma_x^{(t)}$ which cannot happen when $\Sigma_x^{(t)} = I$. Formalizing the resulting local-minimax lower bound is ongoing work. Whether the mixing coefficients $\Gamma, \mu$ in our final bounds is optimal is a challenging question. Indeed, recent work [1] establishes in the general case that a final error bound on least squares regression scaling as $\frac{\tau_{\mathrm{mix}}}{N}$ (for $N$ total data points) is minimax-optimal, where $\tau_{\mathrm{mix}}$ is proportional to the dependence on $\log(N), \Gamma, \mu$ in our bounds, at least in the Markov chain setting. However, in parallel, even more recent work [2] demonstrates that under a related notion of mixing, the dependence on the mixing time becomes lower order after sufficiently large burn-in on $N$, after which the iid regression rate $\mathcal O(1/N)$ is recovered. Regardless, these works are specific to estimators such as the empirical risk minimizer which are distinct from ours, not to mention the additional bilinear structure on our multi-task operators. In short, fully determining whether the dependence on mixing coefficients in our bounds is optimal is non-trivial, though we hypothesize that it is not, thanks to the certain amenable features of our model, such as realizability, subgaussian covariates etc.

[1] Guy Bresler, Prateek Jain, Dheeraj Nagaraj, Praneeth Netrapalli, Xian Wu. “Least Squares Regression with Markovian Data: Fundamental Limits and Algorithms”

[2] Ingvar Ziemann and Stephen Tu. “Learning with little mixing”