Meta-Learning Operators to Optimality from Multi-Task Non-IID Data
Thomas T. C. K. Zhang Leonardo F. Toso James Anderson Nikolai Matni
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 &
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.