Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data
Authors: Thomas TCK Zhang
Leonardo Felipe Toso
James Anderson
Nikolai Matni
Venue: ICLR 2024 spotlight (2024)
Keywords: Representation learning
meta learning
multi-task learning
Primary area: learning theory
TL;DR: Existing representation learning methods induce biases even in linear settings, which throttle sample-efficiency in learning over many tasks. We propose a fix..

Abstract: A powerful concept behind much of the recent progress in machine learning is the extraction of common features across data from heterogeneous sources or tasks. Intuitively, using all of one's data to learn a common representation function benefits both computational effort and statistical generalization by leaving a smaller number of parameters to fine-tune on a given task. Toward theoretically grounding these merits, we propose a general setting of recovering linear operators $M$
from noisy vector measurements $y = Mx + w$, where the covariates $x$ may be both non-i.i.d. and non-isotropic. We demonstrate that existing isotropy-agnostic meta-learning approaches incur biases on the representation update, which causes the scaling of the noise terms to lose favorable dependence on the number of source tasks. This in turn can cause the sample complexity of representation learning to be bottlenecked by the single-task data size. We introduce an adaptation, $\texttt{De-bias}$ & $\texttt{Feature-Whiten}$ ($\texttt{DFW}$), of the popular alternating minimization-descent (AMD) scheme proposed in Collins et al., (2021), and establish linear convergence to the optimal representation with noise level scaling down with the $\textit{total}$ source data size. This leads to generalization bounds on the same order as an oracle empirical risk minimizer. We verify the vital importance of $\texttt{DFW}$ on various numerical simulations. In particular, we show that vanilla alternating-minimization descent fails catastrophically even for iid, but mildly non-isotropic data.
Our analysis unifies and generalizes prior work, and provides a flexible framework for a wider range of applications, such as in controls and dynamical systems.

Published as a conference paper at ICLR 2024
SAMPLE-EFFICIENT LINEAR REPRESENTATION LEARN- ING FROM NON-IID NON-ISOTROPIC DATA
Thomas T.C.K. Zhang University of Pennsylvania
Leonardo F. Toso Columbia University New York, NY 1t287 Ccol umbii edu
Philadelphia, PA ttz2@seas upenn. edu
James Anderson
Nikolai Matni
Columbia University
University of Pennsylvania
New York, NY
Philadelphia, PA nmatnn @seas. weees..peennn............ wwwpppeeeeeii
james. wondersnnooollmmm edu
ABSTRACT
A powerful concept behind much of the recent progress in machine learning is the extraction of common features across data from heterogencous sources or tasks. Intuitively, using all of one's data learn common representation function benefits both computational effort and statistical generalization by leaving smaller number of parameters to fine-tune on given task. Toward theoretically grounding these merits, we propose general setting of recovering linear operators M from noisy vector measurements y = Mr + w, where the covariates x may both non-i.i.d. and non-isotropic. We demonstrate that existing isotropy-agnostic meta- learning approaches incur biases on the representation update, which causes the scaling of the noise terms to lose favorable dependence the number of source tasks. This in turn can cause the sample complexity of representation learning to be bottlenecked by the single-task data size. We introduce adaptation, De-bias as & Feature-Whiter Whiten (DFW), of the popular alternating minimization-Aeseee (AMD) scheme proposed in Collins al., (2021) (Collins al., 2021), and establish linear convergence to the optimal representation with noise level scaling down with the total source data size. This leads to generalization bounds on the same order as oracle empirical risk minimizer. We verify the vital importance of DFW on various numerical simulations. particular, we show that vanilla alternating- minimization descent fails catastrophically even for iid, but mildly non-isotropic data. Our analysis unifies and generalizes prior work, and provides a flexible framework for wider range of applications, such as in controls and dynamical systems.
INTRODUCTION
A unifying paradigm belying recent exciting progress in machine learning is learning common feature space or representation for downstream tasks from heterogeneous sources. This forms the core of fields such meta-learning, transfer learning, and federated learning. A shared theme across these fields is the scarcity of data for specific task out of many, such that designing individual models for each task is both computationally and statistically inefficient, impractical, or impossible. Under the assumption that these tasks are similar in some way, natural alternative approach is to use data across many tasks to learn common component, such that fine-tuning to given task involves fitting a much smaller model that acts on the common component. Over the last few years, significant attention has been given framing this problem setting theoretically, providing provable benefits of learning over multiple tasks in the context of linear regression (Bullins et al., 2019; Du et al., 2020; Tripuraneni et al., 2021: Collins et al., 2021: Thekumparampil et al., 2021; Saunshi et al., 2021) and in identification/conttol of linear dynamical systems (Modi et al., 2021; Chen et al., 2023; Zhang et al., 2023). These works study the problem of linear representation learning, where the data for each task is generated noisily from unknown shared latent subspace, and the goal is to efficiently recover representation of the latent space & from data across different task distributions.

---

Published as a conference paper at ICLR 2024
For example, in the linear regression setting, one may have data of the form
y: g(t) T + noise, u E RR x RRZz RRxd,
with I,..., N iid data points from t l...., task distributions. Since the representation & is shared across all tasks, one may expect the generalization error of an approximate representation & fit TN data points to scale as d.r. the representation. This is indeed the flavor of statistical guarantees from prior work (Du et al., 2020; Tripuraneni et al., 2021; Thekumparampil et al., 2021; Zhang al., 2023), which concretely demonstrates the benefit of using data across different tasks.
TN- where dgr is the number of parameters determining
However, existing work, especially beyond the scalar measurement setting, is limited in more important components of their analysis. For example, it is common to assume that the covariates z; isotropic across all tasks. Furthermore, statistical analyses often assume access to an empirical risk minimizer, even though the linear representation learning problem is non-convex and ill-posed (Maurer et al., 2016; Du al., 2020; Zhang al., 2023). Our paper addresses these problems under a unified framework of linear operator recovery. i.c. recovering linear operators M E R2, xd, from (noisy) vector measurements y Mr w. where the covariates z may not be independent or isotropic. This setting subsumes the scalar measurement setting, and encompasses many fundamental control and dynamical systems problems, such as linear system identification and imitation learning. In particular, the data in these settings are incompatible with the common distributional assumptions (e.g., independence, isotropy) made prior work.
Contributions: Toward this end, our main contributions are as follows:
We demonstrate that naive implementation of local methods for linear representation learning fail catastrophically even when the data is iid but mildly non-isotropic. We identify the source of the failure as interaction between terms incurring biases in the representation gradient, which do not scale down with the number of tasks.
We address these issues by introducing two practical algorithmic adjustments, be-bias (DFW), which provably mitigate the identified issues. We then show that DFW is necessary for gradient-bart based methods to benefit from the total size of the source dataset.
Feature-Whis
We numerically show our theoretical guarantees are predictive of the efficacy of our proposed algorithm, and of the key importance of individual aspects of our algorithmic framework.
Our main result can be summarized by the following descent guarantee for our proposed algorithm.
Theorem 1.1 (main result, informal) Let & be the current estimate of the representation, and &. the optimal representation. Running one iteration of DEW yields the following improvement
C
dist dist($,,..)) S dist( 2.)+ 2.)
P E (0, C> 0.
V# tasks x # data per task'
Critically, the second term of the right hand side scales jointly in the number tasks and datapoints per task, whereas naively implementing other methods may be bottlenecked by term that scales solely with the amount of data for a single task, which leads to suboptimal sample-effiiencyy
1.1 RELATED WORK
Multi-task linear regression: Directly related to our work are results demonstrating the benefits of multi-task learning for linear regression (Maurer et al., 2016; Bullins et al., 2019; Du al., 2020; Tripuraneni et al., 2021; Collins et al., 2021; Thekumparampil et al., 2021). under the assumption of particular, our proposed algorithm adapted
a shared but unknown linear feature representation. from the alternating optimization scheme independently in Nayer and Vaswani (2022) and Collins et al. (2021), and extends these results to the vector measurement setting and introduces algorithmic modifications to extend its applicability to non-iid and non-isotropic covariates. We also highlight that in the isotropic linear regression setting, Thekumparampil et al. (2021) provide an alternating minimization scheme that results in near minimax optimal representation learning. However, the representation update step simultancously accesses data across tasks, which we avoid in this work due to motivating applications, e.g. distributed learning, that impose locality or data privacy constraints.
2

---

Published as a conference paper at ICLR 2024
Meta/mulli-task RL: There is wealth of literature in reinforcement learning that seeks empirically to solve different tasks with shared parameters (Teh et al., 2017: Hessel et al., 2018; Singh et al., 2020; Deisenroth al., 2014). In parallel, there is body of theoretical work which studies the sample efficiency of representation learning for RI (Lu et al., 2021; Cheng et al., 2022; Maurer et al., 2015). This line of work considers MDP settings, and thus the specific results are often stated with incompatible assumptions (such as bounded states/cost functions and discrete action spaces), and are suboptimal when instantiated in our setting.
Multi-task learning has gained recent attention in controls, e.g. for adaptive control over similar dynamics (Harrison et al., 2018; Richards et al., 2021; Shi et al., 2021: Muthirayan et al., 2022), imitation learning for linear systems (Zhang et al., 2023; Guo et al., 2023), and notably linear system identification (I.i et al., 2022: Wang al., 2022; Xin al., 2023; Modi et al., 2021; Faradonbeh and Modi, 2022; Chen et al., 2023). In many of these works (Li et al., 2022; Wang et al., 2022; Xin et al., 2023), task similarity is quantified by generic norm closeness of the dynamics matrices, and thus the benefit of multiple tasks extends only to radius around optimality, Under the existence of a shared representation, our work provides an efficient algorithm and statistical analysis to establish convergence per-task optimality.
System identification and control:
2 PROBLEM FORMULATION
Notation: the Euclidean norm of a vector x is denoted ||=||- The spectral and Frobenius norms of matrix A are denoted ||A||| and ||l respectively. For symmetric matrices A, B, < denotes B - A is positive semidefinite. The largest/smallest singular and eigenvalues of matrix A are denoted (A), min (A), and (A), Amin (A), respectively. The condition number of a matrix A is denoted (<A) == maxe
(A). Define the indexing shorthand [m] = (1...., ......,90 n). We use big-O notation Oc), (-). S?() to omit universal numerical factors, and Ó(), 6()) 2((+) to additionally omit polylog factors in the argument.
Regression Model. Let a covariate sequence (also denoted a trajectory) be an indexed set {2:):21 C Rdz We denote distribution P. over covariate sequences, which we assume to have bounded second moments for all i 2 1, i.e. E [=,a,]] is finite for all i 2 1. Defining the filtration {Fi) where F, == a af 1m1 (wx))==1) is the o-algebra generated by the covariates up to i + 1 and noise up to i, we assume that {w; h121 is o -uugaussiinn martingale difference sequence (MDS):
E 1
we
Fi-1] 0, E [exp Jw:ww) [F-1]] S exp (1*1Iu)2 d a.s. AA R, E
Assuming a ground truth operator M. E Rd Xdz, our observation model is given by
y: M.x. Wi i 1,
for Yi the labels, and w; the label noise. We further define Ex,N = E E[z,o]]] J. When the marginal distributions of Fú, i 2 are identical, we denote E, = Eg,,N-
Multi-Task Operator Recovery. We consider the following instantiation of the above lin- ear operator regression model over multiple tasks. In particular, we consider heterogeneous data ((x{") v ( t t= 1,..., T' task distributions. For notational convenience, we assume that the length of trajectories N is the same across training tasks. For each task I, the observation model is
Ji=L.==: (
consisting of trajectories of length N, generated independently across
M x + !
where 0 F: 4. admits decomposition into a ground-truth representation 4. grxd common across all tasks t E [T] and a task-specific weight matrix F(!)) E gd, r & dz. We denote the joint distribution over covariates and observations (a("), -/ he2 by p!!) We assume that the representation Φ. is normalized to have orthonormal rows to prevent boundedness issues, since F(()) E(() Q-1, -1 Φ QØ, are valid decompositions for any invertible Q E grxr, To measure closeness of approximate representation & to optimality, we define subspace metric.
3

---

Published as a conference paper at ICLR 2024
Definition 2.1 (Subspace Distance (Stewart and Sun, 1990; Collins et al., 2021)) Let D, D. Rrxds be matrices whose rows are orthonormal. Furthermore, let D. Rd (xx=-)xxx r)xda be a matrix such that [o] 1 D. is an orthogonal matrix. Define the distance between the subspaces spanned by the rows of and v, by
dist( 2.) = 1 1
In particular, the subspace distance quantitatively captures the alignment between two subspaces. interpolating smoothly between (occurring iff span($.) (pan($)) and 1 (occurring iff sp, 1 span(()). We define the task-specific stacked vector notation by capital letters, e.g.,
x
(0) 1
IN E R.wxd=
The goal of multi-task operator recovery is to estimate {F( multiple tasks (xx!"),,"))) ) t = l..... T. Some prior works (Maurer et al., 2016; Du et al., 2020; Zhang et al., 2023) assume access to an empirical risk minimization oracle, i.e. access to
Jt=1 and Φ from data collected across
[ JE=1, & argmin ΣΣ (0 - p Dx. (FO))
focusing on the statistical generalization properties of an ERM solution. However, the above optimization is non-convex even in the linear setting, and thus it is imperative to design and analyze efficient algorithms for recovering optimal matrices (Z=!"))" Jt=1 and 4. To address this problem in the linear regression setting, various works, e.g. FedRep (Collins et al., 2021), ALtGD-Min (Nayer scheme, where on fresh data batch, the weights (F(t) are computed on local data via least- ceast-suures, and subsequently an estimate of the representation gradient is then computed with respect to local data and aggregated across tasks to perform gradient descent. This algorithmic framework is intuitive, and thus forms a reasonable starting point toward provably sample-efficiens algorithm in our setting.
and Vaswani, 2022), propose an alternating minimization-dessee
3
LINEAR REPRESENTATION LEARNING
We begin by describing the vanilla alternating minimization-desee (2021). We show that in our setting with label noise and non -oo-iiotropy. interaction terms arise in the representation gradient, which cause biases to form that do not scale down with the number of tasks T. In $3.2, we propose alterations to the scheme to remove these biases, which we then show in $3.3 lead to fast convergence rates that allow us recover near-oracle ERM generalization bounds.
scheme proposed in Collins et al.
3.1 PERILS OF (VANILLA) GRADIENT DESCENT ON THE REPRESENTATION
We begin with summary of the main components of an alternating minimization-desean method analogous to FedRep (Collins et al., 2021) and Al G.D-Min (Nayer and Vaswani, 2022). During each optimization round, a new data batch is sampled for each task: {(x}") (t) >) E [T]. We then compute task-specific weights f²(t) on the corresponding dataset, keeping the current representation estimate & fixed. For example, ((t) may be the least-squares weights conditioned on 4 (Collins et al., 2021). Define z = dz and the empirical covariance matrices 5-()) := 1 (1), $:! 1/2 ( Z(t, The least squares solution (t)) is given by the convex quadratic minimization (r)) argmin Σ llu." F 122 For each task, we then fix the weight matrix f(0) and perform a descent step with respect the representation conditioned on the local data. The resulting representations are averaged across tasks to form the new representation. When the descent direction is the gradient, the update rule is given by
N
T
17
V
yi
M
p
(1)
T
2N
T
i=1
t=1
4

---

Published as a conference paper at ICLR 2024
where 17 0 is a given step size. We normalize o+ to have orthonormal rows, e.g. by (thin/reduced) QR decomposition (Trefethen and Bau, 2022), to produce the final output Φ+, i... ++ Rd+. R E R'xr, leading to
T T T (f() - p.D 50
T T
R$ 6 T
(4)
NT
As in Collins al. (2021), right-multiply both sides of (4) by $ $ recalling | 1...._11
dist @ P.). Crucially, Collins al. (2021) assume x;!") has mean and identity covariance, and w = across i, t. Therefore, the label noise terms ((() (()) xx()) disappear, and the sample covariance for each task ) concentrates to identity, such that 4, E T 0. Under appropriate choice of 7) and bounding the effect of the orthonormalization factor R, linear convergence the optimal representation can be established. However, two issues arise when label noise wit) is introduced and when x;
(t) has non-identity covariance.
1. When label noise (1) is present, since F(() is computed on y ((),, the gradient noise term is generally biased: NT E (() [ U Even in the simple case that all task distributions Plyy are identical, NT Ei=, P(O W(OT ( concentrates to its bias, and thus for large T the size of noise term is bottlenecked at N F xx)))())))))) This critically causes the noise term to lose scaling in the number of tasks T, even when the tasks are identical.
(t) has non-identity covariance, the decomposition into contraction and covariance 2. When ii concentration term no longer holds, since generally 4 EE I * f O. This causes term whose
norm otherwise concentrates around 0 in the isotropic case to scale with Amaa (5()) Aminn (5())
in the worst case.
This motivates modifying the representation update beyond following the vanilla stochastic gradient.
3.2 TASK-EFFICIENT ALGORITHM: DE-BIAS FEATURE-WHITEH
In the previous section, we identified two fundamental issues: 1. the bias introduced by computing the least squares weights and representation update on the same data batch, and 2. the nuisance term introduced by non-identity second moments of the covariates x;; Toward addressing the first issue, we introduce "de-biasing" step. where each agent computes the least squares weights f()) and the representation update on independent batches of data, e.g. disjoint subsets of trajectories. To adaptation (LeCCun et al., 2002), where
address the second issue, we introduce the gradient estimate sent by each agent is pre-conditioned by its inverse sample covariance matrix. Combining these adaptations, the representation update becomes
T M FOTTOW WOT ( ($!0) ($
R$b $ FOT (f10g- -
7 Σ
7
(5)
where assume {1)) are computed independent data using the aforementioned batching I....,T, T, are all mutually independent, then the first two terms of the
strategy. When x; w update form the contraction, and the last term is an average of zero-mean least-squares-error over tasks, which can be studied using standard tools (Abbasi- Yadkori et al., 2011; Abbasi- Yadkori, 2013). This culminates in convergence rates that scale favorably with the number of tasks ($3.3). To Jiz t [2']. be a dataset available to each agent. For the weights de-biasing step, we sub-sample trajectories NG C [N]. For each agent, we compute least-squares weights from M₁. We then sub- sub-sample trajectories A2 C [N] WM, and compute the cask-conditioned representation gradients from NV.
terms
operationalize our proposed adaptations, let ( = {(x y; )
(t) = argmin
F2
600 = Vo
1
Σ iEN1
Mi
F())
(6)
5

---

11: end for Published as a conference paper at ICLR 2024
Algorithm 1 De -ee-bissed & Feature-white (DFW) Alt. Mininization-Descentt
1: Input: step sizes (Tkr batch sizes (Ne)u21 initial estimate 6g.
2: for == 1...... K do
3:
for / E [2'] (in parallel) do
4:
Obtain samples {(a:]] D/ )) Partition trajectories [Ava] N..1 L.
5:
Compute f(() e.g. via least squares on Na.1 (6).
6:
7:
Compute cask-conditioned representation gradient XX2, on Na.2 (6).
Compute task condiiondi representation update (1) (7).
8:
9: end for
@ thin _R E ())
12: return Representation estimate ox
Lastly, each agent updates its local representation via feature-whitened gradient step yield T The global representation update is computed by averaging the updated task-conditioned representations (0) and performing orthonormalization:
T
(t) = 1- (2NNN))
R$
T
We summarize the full algorithm in Algorithm The above de-biasing and feature whitening steps ensure that the expectation of the representation update is contraction (with high probability):
|[ dist( P.),
E dist dist((.. v.) =
R
la.
" T L POS $(0)
where the task and trajectory-wise independence ensures that the variance of the gradient scales inversely in NT.
3.3 ALGORITHM GUARANTEES
We present main result in the form of convergence guarantees for Algorithm 1. We begin by defining standard measure of dependency along covariate sequences via 3-mixing.
Definition 3.1 (8-mixing) Let {zibiz1 be a B. -alued discrete-tine time stochastic process adapted to
filtration (F We denote the stationary distribution Voc. We define the B-mixing coefficient
B(k) sup (er
EupE i21
[l|Pa (-) Moo Mool
where Ilyllcy denotes the total variation distance between probability measures.
Intuitively, the B-mixing coefficient measures how quickly average process mixes its stationary distribution along any sample path. To instantiate our bounds, we make the following assumptions on the covariates.
the marginal distribu- tions of x to be identical with zero-mean, covariance s((), and to -i-subgaussian across all i, t:
Assumption 3.1 (Subgaussian covariates, geometric mixing) We
E|x =0, 0, E exp Xu x (0))] < exp (X (101)0)12 a.s. V.E
Furthermore, we assume the process {x) Jiz1 is geometrically 3-mixing with ((t) S ((0) r()) > O, (0)) E [0, 1), for each task E [T]. Lastly, we define T := log(r(() N/6)
(og(1/)(9)) V v1).

---

Published as a conference paper at ICLR 2024
Notably, these assumptions subsume fundamental problems in learning over (stable) dynamical systems, in particular linear system identification and imitation learning, where non-iid and non- isotropy across tasks are unavoidable. We discuss these applications in depth in Appendix B. As in prior work, our final convergence rates depend on a notion of task diversity.
Definition 3.2 (Task diversity) We define the quantities
T
min :- Amin
1
Σ
A
Amax
Σ
p(o)
(10)
T
T
[ml
We now present our main result.
Theorem 3.1 (Main result) Let Assumption 3.1 hold. Assume the current representation iterate + satisfies dist P.)
Amin
min, K k(El)))
and
AF
X (A )2 | 1| PAmin (2:1
(t)2
[i] 2 max mix
max
? (r + log(T())
Ow
(dy r (00g(1(000)))
(t)2
(dz + log log(T })
W2/T 2 max mix
(1)
Ω
max
T-² (dz log((/6))
(AF )
mmx
| Ami (E(1)
Given
S
with probability greater than 1 - 5, the next iterate outputted by DFW (Algorithm 1)
32 maxx
satisfies
dist (04, 1.)
1/2
( (.. p.) Ö davg
dr + log (7)
Mmn 2
dist
NT
where Gnvg
T =1=1
o 1 is the task-averaged noise-level.
Amiin (
Remark 3.1 (Initialization) We note that Theorem 3. relies the representation & being suffi- ciently close to 4.. We do not address this issue in this paper, and refer to Collins et al. (2021); Tripuraneni et al. (2021); Thekumparampil et al. (2021) for initialization schemes the iid linear regression setting. Our experiments suggest that initialization is often unnecessary, which mirrors the experimental findings Thekumparampil et al. (2021, Sec. 6). We leave constructing an initialization scheme for our general setting, or showing it is unnecessary, to future work.
key benefit of having the variance term in Theorem 3.1 scale properly in N. T is that we may construct representations on fixed datasets whose error scales on the same order as the oracle empirical risk minimizer by running Algorithm 1 appropriately partitioned subsets of given dataset.
Corollary 3.1 (Approximate ERM) Let the assumptions of Theorem 3.1 hold. Let Py be an initial representation satisfying dist ( x.) v, and define p = 64
XP Let D
XF
H(()) -Vi ( be given dataset. There exists partition of D into independent batches
Bi...., BR, such that iterating DFW on Br, E [K] yields with probability greater than 1 ö:
max, (t) ( + log(L/o))
dist( K, 4.)2500 ((O)
(11)
NT
where C(p) a constant depending on the contraction rate p.
We note that the RHS of (11) has the "correct" scaling: noise level o multiplied by # parameters of the representation, divided by the total amount of data NT. In particular, given a fine-tuning dataset of size N' sampled from task T + 1 that shares the representation t_, computing the least
(0)
7

---

Published as a conference paper at ICLR 2024
squares weights F(T+1) conditioned on Φx yields high-probabille bound (cf. Appendix A.4.1) on the parameter error
| M S dist(0k, $2)²+ (T+1)2 2 -o
log(1/5)
N'
C(p) max (1)2 (dgr r log (og(1/6))
(T+1)2 (dyr (dur log(1 (1/))
ow
NT
N
where we omit task task-related quantities for clarity. We note that the above parameter recovery bound mirrors that of ERM estimates (Du al., 2020; Zhang et al., 2023), where we note the latter fine- tuning term scales with dyr (the number of parameters in F(T+1)) as opposed to in the linear regression setting (dy 1).
NUMERICAL VALIDATION
We present numerical experiments to demonstrate the key importance of aspects of our proposed algorithm. We consider two scenarios: 1. linear regression with non-isotropic iid data, and 2. lin- system identification. The linear regression experiments highlight the breakdown of standard approaches when approached with mildly non-isotropic data, thus highlighting the necessity of our proposed be-biasing Feature wnare-whitening steps, and our system identification experi- ments demonstrate the applicability of our algorithm to sequentially dependent (non-isotropic) data. Additional experiments and experiment details can be found in Appendix C.
4.1 LINEAR REGRESSION WITH IID AND NON-ISOTROPIC DATA
We consider the observation model from (1), where set the operator dimensions and rank as d = d, = 50 and r = 7. We generate the T operators using the following steps: 1. a ground truth representation Φ. E Brxds is randomly generated through applying thin_svd to a random matrix with values i.i.d. (0, 2. a nominal task weight matrix Fu E Rd,xr is generated with elements f..dd. NO0 1). task-specific weights F:!") E Rd,xx, e [7] are generated by applying random rotations to Fo. A non-isotropic covariance matrix Σn shared across all tasks is generated as S.,.,, Mi Tr.).)' where Σx = + = 5- la, + V, with Vi, i.d. Unif(0, 1). We note that by design of U, Σg is only mildly non-isotropic. Figure la compares the performance of Algorithm 1, for a single-task (T = 1) and multiple tasks (T = 25), as well as standard alternating minimization- descent like FedRep (Collins et al., 2021) with T = 25. For each optimization iteration, we sample N = 100 fresh data per task. The figure highlights that DFW is able to make use of all source data to decrease variance and learn representation to near-optimality. As predicted in $3.1, Figure la also shows that vanilla alternating mininization-descen is not able to improve beyond a highly suboptimal bias, despite all tasks sharing the same, rather mildly non-isotropic covariate distribution.
For second experiment, we consider instead of applying vanilla alternating minimization -esseent on linear representation, we parameterize the representation by a neural network. In particular, we consider ReL. U-activvae network with one hidden layer of dimension 64. We keep the same getta-genrratin, with N 100 fresh samples per optimization iteration per task, and compare DFW F'W (on a linear representation) to AMD on the neural net representation. Since the subspace distance is no longer defined for the neural net representation, we measure optimality the learned representations by computing the average least-squares loss with respect to a validation set generated through the nominal weight matrix Fo. We plot the "optimal" loss as that attained by the ground truth operator F. We see in Figure 1b that, despite the much greater feature representation power of the neural net, alternating minimization-desee finding& non-linear parameter representation that allows descent to optimality, taking almost two orders of magnitude more iterations than DFW to reach optimality. This feature learning phase is only exacerbated when there are fewer tasks. We note that in such streaming data setting, the high iteration complexity of AMD translates to a greatly hampered sample-efficiency compared to DFW.
plateaus for many iterations just like the linear case before
4.2 LINEAR SYSTEM IDENTIFICATION
We consider a discrete-time linear system identification (sysID) problem, with dynamics
Fi+1 Ax;
Bu: + Wi, i 0, 0,..., N
8

---

Published as a conference paper at ICLR 2024
0.8
OFW (T-13)
pie (Ta2a
Validte toss
AMD IT-14
IT-11
disico 0..)
Siglle Leak (Tm2)
Vulggr ual (1-23)
0.4
Decovet 25:
02-
0
20
80 Number of Iterations k
40
100
16
10'
0.0
50
100
150
200
250
Number iterations
Number of iterations k
(a) ITD linear regression
(b) Neural net rep. + AMD
(c) Linear system identification
Figure 1: We plot the suboptimality the current and ground truth representation with respect the numher of iterations, comparing between the single and multiple-task settings Algorithm and multi-task alternating iimimization-descen We observe performance improvement and variance reduction for multi-task DPW as predicted. All curves are are plotted as the mean with 95% confidence regions shaded
where z; is the state of the system and u, is the control input. In contrast to the previous example, the covariates are now additionally non-iid due to correlation over time. In particular, we can instantiate multi-task linear sysID in the form of (1),
(1) x ait1
(1) (t) M 24 + wi
(0)
i o......,N
where ( = (A B F( Rd xdz, The state-action pair at time instant for all tasks t [T] is embedded as z; = [-; (t) u. T The process noise w and control action ":
(f) are assumed to be drawn from Gaussian distributions (0, E.,) and N TT, Ta,), respectively, where d, represents the dimension of the control action. We set the state dimension dN = 25, control dimension d,, = 2, latent dimension r = 6, horizon N 100, and input variance o 1.
The generation process of the ground truth system matrices M follows similar approach as described in the linear regression problem, with the addition a normalization step of the nominal weight matrix Fo ensure system stability for all tasks & E [7]. Furthermore, the process noise covariance Dn is parameterized in a similar manner as in the linear regression example, with U 5 Ia, + 2. V. The initial state To distribution A (0, ) which is determined by the solution to the discrete Lyapunov equation T + (()
(t) is drawn iid across tasks from the system's 's stationary
En. We note this implies the covariates (t) are inherently non-isotropic and non-identically distributed across tasks. Figure Ic again demonstrates the advantage of leveraging multi-task data to reduce the error in computing shared representation across the system matrices M line with our theoretical findings, DFW continues to benefit from multiple tasks, even when the data is sequentially dependent. We see that FedRep remains
sx() ( (1) (A
suboptimal in this non-iid, non-isotropic setting.
5 DISCUSSION AND FUTURE WORK
We propose an efficient algorithm to provably recover linear operators across multiple tasks to optimality from non-iid non-isotropic data, recovering near oracle empirical risk minimization rates. We show that the benefit of learning over multiple tasks manifests in a lower noise level in the optimization and smaller sample requirements for individual tasks. These results contribute toward general understanding of representation learning from an algorithmic and statistical perspective. Some immediate open questions are: whether good initialization of the representation is necessary, and whether the convergence rate of DFW can be optimized e.g., through rerreuulai weights f()). Resolving these questions has important implications for the natural extension of our framework: as emphasized in (Collins et al., 2021), the alternating empirical risk minimization (holding representa- tion fixed) and gradient descent (holding task-specific weights fixed) framework naturally extends to the nonlinear setting. Providing guarantees for nonlinear function classes is an exciting and impactful avenue for future work, which concurrent work is moving toward, e.g. for 2-layer ReLU networks (Collins et al., 2023) and kernel ridge regression (Meunier et al., 2023). It remains to be seen whether a computationalll--fiiee -ffficiceet algorithm can be established for nonlinear learning in the non-iid and non-isotropic data regime, while preserving joint scaling in number of tasks and data.

---

Published as a conference paper at ICLR 2024
REFERENCES
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared represen- tations for personalized federated learning. In International Conference on Machine Learning. pages 2089-2099 PMLR, 2021.
Brian Bullins, Elad Hazan, Adam Kalai, and Roi Livni. Generalize across tasks: Efficient algorithms for linear representation learning. In algorithmic learning theory, pages 235-246. PMLR, 2019.
Simon S Du, Wei Hu, Sham M Kakade, Jason D Lee, and Qi Lei. Few-shot learning via learning the representation, provably. arXiv preprint arXiv:2002.0943: 2020.
Nilesh Tripuraneni, Chi Jin, and Michael Jordan. Provable meta-learning of linear representations. In International Conference Machine Learning, pages 10434-10443, PMLR, 2021.
Kiran Koshy Thekumparampil, Prateek Jain, Praneeth Netrapalli, and Sewoong Oh. Sample efficient linear meta-learning by alternating minimization, 2021.
Nikunj Saunshi, Arushi Gupta, and Wei Hu. A representation learning perspective on the importance of vrain-validation splitting in meta-learning, 2021.
Aditya Modi, Mohamad Kazem Shirani Faradonbeh, Ambuj Tewari, and George Michailidis. Joint learning of linear time-invariant dynamical systems. arXiv preprint arXiv:2112.10955 2021.
Yiting Chen, Ana M Ospina, Fabio Pasqualetti, and Emiliano Dall' Anese. Multi-task system identification of similar linear time-invariant dynamical systems. arXiv preprint arXiv:3301.01430, 2023.
Thomas T Zhang. Katie Kang, Bruce D) Lee, Claire Tomlin, Sergey Levine, Stephen Tu, and Nikolai Matni. Multi-task imitation learning for linear dynamical systems. In Learning for Dynamics and Control Conference. pages 586-599. PMLR, 2023.
Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes. The benefit of multitask representation learming. Journal of Machine Learning Research, 17(81):1-32, 2016.
Seyedehsara Nayer and Namrata Vaswani. Fast and sample-efficient federated low rank matrix recovery from column wise linear and quadratic projections. IEEE Transactions on Information Theory, 69(2); 117-112177 2022.
Yee Whye Teh, Victor Bapst, Wojciech Marian Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning, 2017.
Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojeiech Czarnecki, Simon Schmitt, and Hado van Hasselt. Multi-task deep reinforcement learning with popart, 2018.
Avi Singh, Eric Jang, Alexander Irpan, Daniel Kappler, Murtaza Dalal, Sergey Levine, Mohi Khansari, and Chelsea Finn. Scalable multi task imitation learning with autonomous improvement. 2020.
Marc Peter Deisenroth, Peter Englert, Jan Peters, and Dieter Fox. Multi-task policy search for robotics. In 2014 IEEE International Conference on Robotics and Automation (ICRA), pages 3876-3881, 2014. doi: 10.1109/ICRA22211..9991
Rui Lu, Gao Huang, and Simon S. Du. On the power of multitask representation learning in linear mdp, 2021.
Yuan Cheng, Songtao Feng, Jing Yang, Hong Zhang, and Yingbin Liang. Provable benefit of multitask representation learning in reinforcement learning, 2022.
Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes The benefit of multitask representation learning, 2015.
James Harrison, Apoorva Sharma, Roberto Calandra, and Marco Pavone. Control adaptation via meta-learning dynamics. In Workshop Meta-Learning NeurIPS, volume 2018, 2018.
10

---

Published as a conference paper at ICLR 2024
Spencer M Richards, Navid Azizan, Jean-Jacques Slotine, and Marco Pavone. Adaptive-control- oriented meta-learming for nonlinear systems. arXiv preprint arXi:21003.0499, 2021.
Guanya Shi, Kamyar Azizzadenesheli, Michael O' O'Connell, Soon-Jo Chung, and Yisong Yue. Meta-
adaptive nonlinear control: Theory and algorithms. Advances in Neural Information Processing Systems, 34:10013-10025, 2021.
Deepan Muthirayan, Dilcep Kalathil, and Pramod Khargonckar, Meta-learning online control for linear dynamical systems. arXiv preprint arXiv:2208. 10259, 2022.
Taosha Guo, Abed AlRahman Al Makdah, Vishaal Krishnan, and Fabio Pasqualetti. Imitation and transfer learning for lqg control. arXiv preprint arXiv:2303.00000 2023.
Lidong Li, Claudio De Persis, Pietro Tesi, and Nima Monshizadeh. Data-based transfer stabilization in linear systems. arXiv preprint arXiv:2211.00 200:2221.00536, 2022.
Han Wang, Leonardo F Toso, and James Anderson. Fedsysid: A federated approach to sample- efficient system identification. arXiv preprint arXiv:2211.14393, 2022.
Lei Xin, Lintao Ye, George Chiu, and Shreyas Sundaram. Learning dynamical systems by leveraging data from similar systems. arXiv preprint arXiv:2302.04344 2023.
Mohamad Kazem Shirani Faradonbeh and Aditya Modi. Joint learning based stabilization of multiple unknown linear systems. PTPA--,upesOOnLine, 55(12):723-728, 2022.
Gilbert W Stewart and Ji-guang Sun. Matrix perturbation theory. Academic press, 1990.
Lloyd N Trefethen and David Bau. Numerical linear algebra, volume 181. Siam, 2022.
Yann LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. Neural networks: Tricks of the trade, pages 9-50. Springer, 2002.
Yasin Abbasi-Yadeee Dávid Pál, and Csaba Szepesvári. Online least squares estimation with self-normalized processes: An application to bandit problems. arXiv preprint arXiv:1102.2670.
2011.
Yasin Abbasi- Yadkori. Online learning for linearly parametrized control problems. 2013.
Liam Collins, Hamed Hassani, Mahdi Soltanolkotabi, Aryan Mokhtari, and Sanjay Shakkottai. Provable multi-task representation learning by two-layer relu neural networks. arXiv preprint arXiv:2307.06887 2023.
Dimitri Meunier, Zhu Li. Arthur Gretton, and Samory Kpotufe. Nonlinear meta-learning can guarantee faster rates. arXiv preprint arXiv:2307.10870, 2023.
Tuhin Sarkar, Alexander Rakhlin, and Munther A Dahleh. Finite time Iti system identification. The Journal of Machine Learning Research, 22(1): 186-11266,,, 2021.
Vitaly Kuznetsov and Mehryar Mohri. Generalization bounds for non-stationary mixing processes. Machine Learning, 106(1):93-117, 2017.
Stephen Tu and Benjamin Recht. Least-squares temporal difference learning for the linear quadratic
regulator. In International Conference on Machine Learning, pages 5005-5014. PMLR, 2018.
Joel A. Tropp. User-friendly tail bounds for sums of random matrices. Foundations of Computational Mathematics, 2(4)):39-433, aug 2011. doi: 00.1007/11208--01-99900 URL https: htttps://dii... /doi. org/10. 7827 0200 01 9099-2. z.
Roger A Horn and Charles R Johnson. Matrix analysis. Cambridge university press, 2012.
Bin Yu. Rates convergence for empirical processes of stationary mixing sequences. The Annals of Probability, pages 94-1 1994.
Bin Hu, Kaiqing Zhang, Na Li, Mehran Mesbahi, Maryam Fazel, and Tamer Bayar. Toward theoretical foundation of policy optimization for learning control policies. Annual Review of Control, Robotics, and Autonomous Systems, 6: 1:112155,, 2023.
11

---

Published as a conference paper at ICLR 2024
Karl Krauth, Stephen Tu, and Benjamin Recht. Finite-time analysis of approximate policy iteration for the linear quadratic regulator. Advances in Neural Information Processing Systems, 32, 2019.
Stephen Tu and Benjamin Recht. The gap between model-based and model-free methods on the linear quadratic regulator: An asymptotic viewpoint. In Conference on Learning Theory, pages 3036-3083. PMLR, 2019.
Benjamin Recht. A tour of reinforcement learning: The view from continuous control. Annual Review of Control, Robotics, and Autonomous Systems, 2:253 2019.
Maryam Fazel, Rong Ge, Sham Kakade, and Mehran Mesbahi. Global convergence policy gradient methods for the linear quadratic regulator. In International Conference on Machine Learning, pages 1467-1476. PMLR, 2018.
Max Simchowitz and Dylan Foster. Naive exploration is optimal for online lqr. In International Conference on Machine Learning, pages 8937-8948. PMLR, 2020.
Sarah Dean, Horia Mania, Nikolai Matni, Benjamin Recht, and Stephen Tu. Regret bounds for robust adaptive control of the linear quadratic regulator. Advances in Neural Information Processing Systems, 31, 2018
Sarah Dean, Horia Mania, Nikolai Matni, Benjamin Recht, and Stephen Tu. On the sample complexity of the linear quadratic regulator. Foundations of Computational Mathematics, 20(4):633-67) 2020.
Naman Agarwal, Elad Hazan, and Karan Singh. Logarithmic regret for online control. Advances in Neural Information Processing Systems, 32, 2019a.
Naman Agarwal, Brian Bullins, Elad Hazan, Sham Kakade, and Karan Singh. Online control with adversarial disturbances. In International Conference Machine Learning, pages 111-119. PMLR, 2019b.
Ingvar Ziemann, Anastasios Tsiamis, Henrik Sandberg, and Nikolai Matni. How are policy gradient methods affected by the limits of control? In 2022 IEEE 61 Conference Decision and Control (CDC), pages 5992-5999. IEBE, 2022.
Ingvar Ziemann and Henrik Sandberg, Regret lower bounds for learning linear quadratic gaussian systems. arXiv preprint arXiv:2201.016R) 2022.
Bruce Lee, Ingvar Ziemann, Anastasios Tsiamis, Henrik Sandberg, and Nikolai Matni. The fundamental limitations of learning lincar-quadratic regulators. arXiv preprint a7X5v:2303.15637, 2023.
Max Simchowitz, Horia Mania, Stephen Tu, Michael Jordan, and Benjamin Recht. Learning without mixing: Towards sharp analysis of linear system identification. In Conference On Learning Theory, pages 439-473. PMLR, 2018.
Anastasios Tsiamis, Ingvar M Ziemann, Manfred Morari, Nikolai Matni, and George J Pappas. Learning to control linear systems can be hard. In Conference on Learning Theory, pages 3820- 3857. PMLR, 2022.
Yassir Jedra and Alexandre Proutiere. Finite-time identification of stable linear systems optimality of the least-squares estimator. In 2020 59th IEEE Conference on Decision and Control (CDC), pages 996-1001, 2020. doi: 10. DD9CCDCC42440.2000999900
Stephen Tu, Roy Frostig, and Mahdi Soltanolkotabi. Learning from many trajectories. arXiv preprint arXiv:2203.17113, 2022.
Alexander Goldenshluger and Assaf Zeevi. Nonasymptotic bounds for autoregressive time series modeling. The Annals of Statistics, 29(2)4417444, 2001.
12


Reviews:
[{"review_id": "TXmhWmUB2G", "reviewer": "Reviewer_4V27", "summary": "In \"Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data\" proposes a scheme and statistical guarantees to problems stemming from multi-task learning. In this setting, prior works focused on i.i.d. and isotropic data while the proposed work allows for non-i.i.d. and non-isotropic data. In order to design the scheme and provide with statistical guarantees, in which learning all tasks jointly implies a statistical gain, the authors modify a proposed scheme for the i.i.d. and isotropic data by including mini-batches and whitening.\n\nThe obtained results are what is expected in terms of statistical precision, given the total number of samples, tasks and problem dimension.", "questions": "- While both are valid, the title in the pdf file and the title given within the open review system do not match.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is overall well written, and I think the result is good enough. While a criticism can be put forth in that it combines existing known techniques to establish the final result, it is not necessarily obvious that the combination yields the desired statistical result.", "weaknesses": "I am overall happy with the paper. I think the authors did a good job at presenting their work. I mainly have two questions/weaknesess. The result provided by the authors requires a minimum number of samples under which contraction upto a ball of the alignment of the estimated and optimal subspace are. Is there any sense to how tight this bound is from an information theoretical sense, i.e. the scaling with gamma, mu, etc? \n\nSecond, in corollary 3.1. the authors establish the existence of a partition of independent batches that guarantees this result. Can it be guaranteed that such partition is found in practice?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4SVUIo0aCi", "reviewer": "Reviewer_6Cex", "summary": "This paper points out a failed example of traditional algorithms in handling non-isotropic data. To overcome this issue, it proposes an algorithm called De-bias & Feature-Whiten (DFW) for multi-task linear representation learning from non-iid and non-isotropic data. DFW provably recovers the optimal shared representation at a rate that scales favorably with the number of tasks and data samples per task. Numerical verification is also provided to validate the proposed algorithms.", "questions": "1. 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?\n\n2. 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. \n\n3. The proof for the non-iid case simply says after taking the \"blocking technique on each trajectory\", everything is same as the iid case. First, what is the \"blocking technique on each trajectory\"? Has this technique been introduced before? \n\n4. Then 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 supplimentary 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}]$$\nThis 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Regarding the originality, few meta- federated- learning papers are working on non-iid settings. So this paper has its own novelty.  \n\nThe paper is also well-structured and clearly states the necessary backgrounds, though some technical details should be further extended.\n\nThe example on the non-IID non-isotropic data provides a clear motivation of proposing a new algorithm to overcome this issue. It indicates parts of significance of this work.", "weaknesses": "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$.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HwvWf8SJtg", "reviewer": "Reviewer_ZRaH", "summary": "This paper proposes an algorithm about learning the representation is a linear connection between feature and labels. The algorithm is based on gradient descent and QR decomposition on the iterates. The paper further proves a bound about the sample complexity and error of the algorithm, which is optimal in terms of problem parameters (degree of freedom). Numerical experiments validates the performance of the algorithm.", "questions": "No more questions.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper proposes A practical and simple algorithm, and the theories as well as the math proof of the sample complexity (per batch and in total) and error are solid in terms of the degree of the freedom. The logic and the writing is clear. \nEspecially, Remark 3.2 is great where we can see that the lower bound of $N$ makes sense. Some other papers, although claiming optimality with respect to total samples $NT$, there is a strong assumption on lower bound $N$ that makes them trival, e.g., Du et al.", "weaknesses": "On the other hand, does Tripuraneni et al. work when $N = O(1)$? This paper assumes $N = \\Omega(r)$ so there is still a gap from the optimum. The result in this paper is still good because $r$ it's a small number in low rank setting which we are interested in, and it is already better than the papers listed in Remark 3.2. But it would be great to propose why this paper cannot achieve $N = O(1)$. \n\nSince this paper discusses general feature covariances, it would be great to talk more about the impact of the spectrum of the covariance matrix. There are a few papers about how the feature and operator covariances’ spectrums show up in the bound, and how the \"aligned\" covariances help learning, for example,\n\nWu and Xu, On the Optimal Weighted $\\ell_2$ Regularization in Overparameterized Linear Regression\n\nAnd a few relevant ones. \n\nIt would be great to have a notation table, either in main text or appendix, because there are many different notations/definitions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZAcFkuirYH", "reviewer": "Reviewer_St8o", "summary": "The authors study the problem of linear representation in a multi-task regression setting. As a starting point, they use an alternating minimization procedure (AMD) developed in prior works on the same problem. They showed empirically that this procedure can fail to learn the correct representation when there is noise in the observations or non-isotropic covariates, even when the different tasks are identical, and gave a theoretical explanation for the sources of error. Based on their analysis, they propose a modification to the alternating minimization procedure (dubbed DFW) which can handle noisy observations and non-isotropic covariates, and experiments confirm the efficacy of their modification.", "questions": "1. 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?\n\n2. 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$?\n\n3. In Definition 3.1, is there an implicit assumption that the stationary distribution $\\nu_\\infty$ exists, or are there some conditions imposed on the covariate trajectory distributions which guarantee that a stationary distribution will exist as a consequence?\n\n4. Do the results still hold if the equality for $\\beta(k)$ in Assumption 3.1 is replaced with an inequality?\n\n5. 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "**Clarity of exposition.** The paper is very well written and easy to follow. The authors give extensive interpretation of their results which greatly contributed to my understanding of the paper. The precise relationship to previous work is made explicit, so readers unfamiliar with this sub-field can still parse the paper and understand its contribution easily.\n\n**Intuitive and well-motivated algorithm.** The shortcomings of the base algorithm (AMD) are explained clearly, as are the modifications the authors proposed in DFW, making for an intuitive algorithm. The modifications are simple, easy to implement, and obtain near optimal sample complexity rates.\n\n**Technical contribution.** The authors remove strong technical assumptions found in previous work. They show both theoretically and empirically that these strong assumptions are necessary for AMD to succeed, and are not merely artifacts of previous proofs. Their results are strong both statistically (obtaining optimal sample complexity) and algorithmically (not requiring access to optimization oracles for non-convex problems, which were assumed in some previous works). Their algorithm is also constructed in such a way that data does not need to be shared in its explicit form across tasks, making it attractive when data privacy is a concern. (Remark: It is unclear if the representation _updates_ from each task will still leak private information, but anyway this is not the main focus of the paper.)", "weaknesses": "**Theory.** While the assumptions are much weaker than those in related works, some of the assumptions are still very strong. Two in particular stand out.\n1. 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.\n2. 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$.\n\n**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.\n\nA minor point: the title of the OpenReview submission does not match the title on the paper. This should be fixed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)