An Adaptive Tangent Feature Perspective of Neural Networks
In order to better understand feature learning in neural networks, we propose
a framework for understanding linear models in tangent feature space where the
features are allowed to be transformed during training. We consider linear
transformations of features, resulting in a joint optimization over parameters
and transformations with a bilinear interpolation constraint. We show that this
optimization problem has an equivalent linearly constrained optimization with
structured regularization that encourages approximately low rank solutions.
Specializing to neural network structure, we gain insights into how the
features and thus the kernel function change, providing additional nuance to
the phenomenon of kernel alignment when the target function is poorly
represented using tangent features. We verify our theoretical observations in
the kernel alignment of real neural networks.
Daniel LeJeune, Sina Alemohammad
# An Adaptive Tangent Feature Perspective

###### Abstract

In order to better understand feature learning in neural networks, we propose a framework for understanding linear models in tangent feature space where the features are allowed to be transformed during training. We consider linear transformations of features, resulting in a joint optimization over parameters and transformations with a bilinear interpolation constraint. We show that this optimization problem has an equivalent linearly constrained optimization with structured regularization that encourages approximately low rank solutions. Specializing to neural network structure, we gain insights into how the features and thus the kernel function change, providing additional nuance to the phenomenon of kernel alignment when the target function is poorly represented using tangent features. In addition to verifying our theoretical observations in real neural networks on a simple regression problem, we empirically show that an adaptive feature implementation of tangent feature classification has an order of magnitude lower sample complexity than the fixed tangent feature model on MNIST and CIFAR-10.

## 1 Introduction

Tremendous research effort has been expended in developing and understanding neural networks [1, 2, 3, 4]. In terms of development, this effort has been met with commensurate tremendous practical success, dominating the state of the art [5, 6, 7] and establishing a new normal of replacing intricately engineered solutions with the conceptually simpler approach of learning from data [8].

Paradoxically, this simple principle of learning has not made the theoretical understanding of the success of neural networks easier [3, 9]. Neural networks stand in stark contrast to the engineered solutions that they have replaced--instead of leveraging vast amounts of human expertise about a particular problem, for which theoretical guarantees can be specifically tailored, neural networks appear to be universal learning machines, adapting easily to a wide range of tasks given enough data. The tall task of the theoretician is to prove that neural networks efficiently learn essentially any function of interest, while being trained through an opaque non-convex learning process on data with often unknown and mathematically uncharacterizable structure.

One promising theoretical direction for understanding neural networks has been through linearizing networks using the neural tangent kernel (NTK) framework [10, 11]. Given an appropriate initialization, infinitely wide neural networks can be shown to have constant gradients throughout training, such that the function agrees with its first-order Taylor approximation and is therefore a linear model, where the (tangent) features are the gradient of the network at initialization. The NTK framework reduces the complexity of neural networks down to linear (kernel) regression, which is significantly better theoretically understood [9, 12, 13, 14]. However, real neural networks still outperform their NTK approximants [15], and the fundamental assumption of the NTK--that the gradients do not change during training--is typically not satisfied in theory or practice [16, 17].

In this work, we take a step towards understanding the effects of the change of the gradients and how these changes allow the features to adapt to data. Rather than considering neural networks directly, we consider a relaxation of the problem in which the tangent features are allowed to adapt to the
problem alongside the regression coefficients. By studying this underlying fundamental problem, we shed light on effects observed in real neural networks. Our specific contributions are as follows.

* We introduce a framework of linear feature adaptivity enabling two complementary views of the same optimization problem: as regression using adaptive features, and equivalently as structured regression using fixed features.
* We show how restricting the adaptivity imposes specific regularizing structure on the solution, resulting in a group approximate low rank penalty on a neural network based model.
* We consider the resulting adapted kernel and provide new insights on the phenomenon of NTK alignment [18, 19, 20, 21], specifically when the target function is poorly represented using the initial tangent features.
* We empirically evaluate our adaptive feature model in neural networks on MNIST and CIFAR-10, which provides an order of magnitude better sample complexity compared to fixed tangent features.

Our work is quite far from an exact characterization of real neural networks, and our empirical results indicate that adaptivity as we propose cannot fully explain the performance of real neural networks. Nevertheless, our framework extends the class of phenomena can be explained by linear models built on tangent features, and so we believe it to be a valuable contribution towards understanding neural networks.

**Related work.** The neural tangent kernel has been proposed and studied extensively in the fixed tangent feature regime [10, 11, 15, 22]. Recent works have studied the alignment of the tangent kernel with the target function: Baratin et al. [18] empirically demonstrate tangent kernel alignment, and Atanasov et al. [19] show that linear network tangent kernels align with the target function early during training. Seleznova and Kutyniok [20] demonstrate alignment under certain initialization scalings, and characterize kernel alignment and neural collapse under a block structure assumption on the NTK [21]. In contrast, by characterizing the adaptive feature learning problem instead of neural networks specifically, we are able to gain more nuanced insights about kernel alignment.

The idea of simultaneously learning features and fitting regression models has appeared in the literature in tuning kernel parameters [23], multiple kernel learning [24], and automatic relevance determination (ARD) [25], which has been shown to correspond to a sparsifying iteratively reweighted \(\ell_{1}\) optimization [26]. Other areas in which joint factorized optimization results in structured models include matrix factorization [27] and adaptive dropout [28], which are equivalent to iteratively reweighted \(\ell_{2}\) optimization. Our work provides generic results on optimization of matrix products with rotationally invariant penalties, which complements the existing literature.

All proofs can be found in Appendix A.

## 2 An adaptive feature framework

We first formalize our adaptive feature framework, which enables us to jointly consider feature learning and regression. Our formulation is motivated by highly complex overparameterized models such as neural networks and specializes to the interpolation regime in which the training dataset can be perfectly fit by the model.

**Notation.** Given a vector \(\mathbf{v}_{\mathbf{x}}\in\mathbb{R}^{P}\) parameterized by another vector \(\mathbf{x}\in\mathbb{R}^{Q}\), we use the "denominator" layout of the derivative such that \(\nabla_{\mathbf{x}}\mathbf{v}_{\mathbf{x}}=\partial\mathbf{v}_{\mathbf{x}}/ \partial\mathbf{x}\in\mathbb{R}^{P\times Q}\), and given a scalar \(v_{\mathbf{X}}\in\mathbb{R}\) parameterized by a matrix \(\mathbf{X}\in\mathbb{R}^{P\times Q}\), we orient \(\nabla_{\mathbf{X}}v_{\mathbf{X}}\in\mathbb{R}^{Q\times P}\). The vectorization of a matrix \(\mathbf{X}=[\mathbf{x}_{1}\quad\ldots\quad\mathbf{x}_{Q}]\in\mathbb{R}^{P \times Q}\) is the stacking of columns such that \(\mathrm{vec}(\mathbf{X})^{\mathsf{T}}=\begin{bmatrix}\mathbf{x}_{1}^{\mathsf{ T}}&\ldots&\mathbf{x}_{Q}^{\mathsf{T}}\end{bmatrix}\in\mathbb{R}^{PQ}\). For \(\mathbf{x}\in\mathbb{R}^{\min\{P,Q\}}\), we denote by \(\mathrm{diag}_{P\times Q}(\mathbf{x})\in\mathbb{R}^{P\times Q}\) the (possibly non-square) matrix with \(\mathbf{x}\) along the main diagonal, and we omit the subscript \(P\times Q\) when \(P=Q\). We denote the set \(\{1,\ldots,N\}\) by \([N]\). Given a vector \(\mathbf{x}\in\mathbb{R}^{P}\), we denote by \([\mathbf{x}]_{j}\) for \(j\in[P]\) the \(j\)-th coordinate of \(\mathbf{x}\). Given a square
matrix \(\mathbf{X}\in\mathbb{R}^{P\times P}\), we let \(\lambda_{i}(\mathbf{X})\) denote its \(i\)-th largest eigenvalue. The characteristic function \(\chi_{\mathcal{A}}\) of a set \(\mathcal{A}\) satisfies \(\chi_{\mathcal{A}}(\mathbf{x})=0\) for \(\mathbf{x}\in\mathcal{A}\) and \(\chi_{\mathcal{A}}(\mathbf{x})=\infty\) for \(\mathbf{x}\notin\mathcal{A}\).

### First-order expansion with average gradients

Consider a differentiably parameterized function \(f_{\mathbf{\theta}}\colon\mathbb{R}^{D}\to\mathbb{R}^{C}\) with parameters \(\mathbf{\theta}\in\mathbb{R}^{P}\), such as a neural network. Our goal is to fit this function to data \((\mathbf{x}_{1},\mathbf{y}_{1}),\ldots,(\mathbf{x}_{N},\mathbf{y}_{N})\in \mathbb{R}^{D\times C}\) by solving

\[\operatorname*{minimize}_{\mathbf{\theta}\in\mathbb{R}^{P}}\ \sum_{i=1}^{N}\ell( \mathbf{y}_{i},f_{\mathbf{\theta}}(\mathbf{x}_{i})),\]

where \(\ell\colon\mathbb{R}^{C}\times\mathbb{R}^{C}\to\mathbb{R}\) is some loss function. In order characterize the solution, we need to understand how \(f_{\mathbf{\theta}}\) changes with \(\mathbf{\theta}\). One way that we can understand \(f_{\mathbf{\theta}}\) is through the fundamental theorem of calculus for line integrals. Letting \(\mathbf{\theta}_{0}\) be some reference parameters, such as random initialization or pretrained parameters,

\[f_{\mathbf{\theta}}(\mathbf{x})-f_{\mathbf{\theta}_{0}}(\mathbf{x})=\int_{\mathbf{\theta} _{0}}^{\mathbf{\theta}}\nabla_{\mathbf{\theta}^{\prime}}f_{\mathbf{\theta}^{\prime}}( \mathbf{x})d\mathbf{\theta}^{\prime}=\underbrace{\Big{(}\int_{0}^{1}\nabla_{\mathbf{ \theta}^{\prime}}f_{\mathbf{\theta}^{\prime}}(\mathbf{x})\big{|}_{\mathbf{\theta}^{ \prime}=(1-t)\mathbf{\theta}_{0}+t\mathbf{\theta}}dt\Big{)}}_{\triangleq\overline{ \nabla f_{\mathbf{\theta}}(\mathbf{x})}}(\mathbf{\theta}-\mathbf{\theta}_{0}).\]

That is, our model is a linear predictor using the average **tangent features**\(\overline{\nabla f_{\mathbf{\theta}}}(\mathbf{x})\) and coefficients \(\mathbf{\theta}-\mathbf{\theta}_{0}\). When \(\mathbf{\theta}\approx\mathbf{\theta}_{0}\), we should expect that \(\overline{\nabla f_{\mathbf{\theta}}}(\mathbf{x})\approx\nabla_{\mathbf{\theta}_{0}}f _{\mathbf{\theta}_{0}}(\mathbf{x})\), which are the tangent features at initialization. In fact, this has been shown to hold for very wide neural networks in the "lazy training" regime even for \(\mathbf{\theta}\) at the end of training, in which case the problem can be understood as kernel regression using the neural tangent kernel [10, 11]. However, outside of those special circumstances, it is likely that \(\overline{\nabla f_{\mathbf{\theta}}}(\mathbf{x})\) will change with \(\mathbf{\theta}\). In general, it is difficult to say anything further about how \(f_{\mathbf{\theta}}\) should change with \(\mathbf{\theta}\), or what properties the optimal \(\mathbf{\theta}\) and \(\overline{\nabla f_{\mathbf{\theta}}}(\mathbf{x})\) for a prediction problem should have.

### Adaptive feature model assumptions

In order to proceed with understanding the fitting of such a function, we must make additional assumptions. At this point, we do not wish to specialize to neural networks; instead, we make a few heuristic assumptions based on general properties of overparameterized models.

**Interpolation.** We assume that for any \(\widetilde{\mathbf{y}}_{1},\ldots,\widetilde{\mathbf{y}}_{N}\) of interest, there exists \(\widetilde{\mathbf{\theta}}\in\mathbb{R}^{P}\) such that \(\widetilde{\mathbf{y}}_{i}=f_{\widetilde{\mathbf{\theta}}}(\mathbf{x}_{i})\) for all \(i\in[N]\). This assumption is in line with the growing literature on interpolating predictors [9, 29, 30, 31, 32] and typically holds for neural networks, which are universal function approximators given enough parameters [33]. In particular, we require that this assumption holds for

\[\widehat{\mathbf{y}}_{i}\triangleq\operatorname*{arg\,min}_{\widetilde{\mathbf{ y}}\in\mathbb{R}^{C}}\ell(\mathbf{y}_{i},\widetilde{\mathbf{y}}),\]

such that for \(\widehat{\mathbf{\theta}}\in\mathbb{R}^{P}\), we have \(\widehat{\mathbf{y}}_{i}=f_{\widehat{\mathbf{\theta}}}(\mathbf{x}_{i})\) for all \(i\in[N]\), which is typically satisfiable in regression. In classification problems, the minimizers are typically infinite valued and never realized unless there is label noise, in which case we should let \(\widehat{\mathbf{y}}_{i}\triangleq\operatorname*{arg\,min}_{\widetilde{\mathbf{ y}}\in\mathbb{R}^{C}}\sum_{j\colon\mathbf{x}_{j}=\mathbf{x}_{i}}\ell( \mathbf{y}_{j},\widetilde{\mathbf{y}})\).

**Factorized features.** We assume that for each \(\mathbf{\theta}\), there exists a matrix \(\mathbf{M}_{\mathbf{\theta}}\in\mathbb{R}^{P\times P}\) such that for all \(\mathbf{x}\in\mathbb{R}^{D}\), \(\overline{\nabla f_{\mathbf{\theta}}}(\mathbf{x})=\nabla_{\mathbf{\theta}_{0}}f_{\mathbf{ \theta}_{0}}(\mathbf{x})\mathbf{M}_{\mathbf{\theta}}\). That is, the features are a linear transformation of the initial tangent features at the reference parameters \(\mathbf{\theta}_{0}\). This assumption is slightly stronger than the interpolation assumption: it also requires consistency at other values of \(\mathbf{x}\) not in the training data (if the interpolation assumption holds, we should expect \(\mathbf{M}_{\mathbf{\theta}}\) to exist that satisfies this equality for all \(\mathbf{x}_{i}\)). For highly overparameterized models, we believe this assumption roughly approximates reality.

**Independent optimization.** We assume that \(\mathbf{M}_{\mathbf{\theta}}\) can change essentially independently of \(\mathbf{\theta}\), and so we simply write \(\mathbf{M}\) instead of \(\mathbf{M}_{\mathbf{\theta}}\). In this key assumption, we allow ourselves to depart from the rigid parameterization of a particular model (such as a neural network) and ask what a less
restricted model would choose to do. This assumption is motivated by the compositional nature of neural networks and the observation that change in the gradients for a particular parameter is often more heavily affected by changes in parameters in other layers than by changes to the parameter itself. Despite the independence of \(\mathbf{M}\) and \(\boldsymbol{\theta}\), if the change of \(\boldsymbol{\theta}\) is small, we still expect change in \(\mathbf{M}\) to be small as well.

**Local optimization.** We assume that we search for a solution near our reference parameters \(\boldsymbol{\theta}_{0}\). This assumption coincides with the practice of optimization by gradient-based methods. Since we do not want to bias the optimization toward any particular direction, we apply an \(\ell_{2}\) penalty on the deviation of \(\boldsymbol{\theta}\) from \(\boldsymbol{\theta}_{0}\). Since \(\mathbf{M}\) should also not deviate much from its initial value of \(\mathbf{M}=\mathbf{I}_{P}\), we also apply some appropriate regularizer \(\Omega\colon\mathbb{R}^{P\times P}\to\mathbb{R}\), to be chosen later.

## 3 Adaptive feature learning

Combining the above assumptions, we obtain our learning model as the optimization problem

\[\widehat{\mathbf{M}},\,\widehat{\boldsymbol{\theta}}\triangleq\operatorname*{ arg\,min}_{\mathbf{M}\in\mathbb{R}^{P\times P},\,\boldsymbol{\theta}\in \mathbb{R}^{P}}\,\Omega(\mathbf{M})+\|\boldsymbol{\theta}-\boldsymbol{\theta }_{0}\|_{2}^{2}\,\,\,\text{s.t.}\,\,\,\widehat{\mathbf{y}}_{i}-f_{\boldsymbol {\theta}_{0}}(\mathbf{x}_{i})=\nabla_{\boldsymbol{\theta}_{0}}f_{\boldsymbol{ \theta}_{0}}(\mathbf{x}_{i})\mathbf{M}(\boldsymbol{\theta}-\boldsymbol{\theta }_{0})\,\,\forall\,i\in[N]. \tag{1}\]

This learning problem is very similar to regularized linear interpolation, except that it has a bilinear interpolation constraint (linear individually in each \(\mathbf{M}\) and \(\boldsymbol{\theta}\)) rather than a simple linear constraint. This bilinear structure will lead to significant structure in both \(\widehat{\mathbf{M}}\) and \(\widehat{\boldsymbol{\theta}}\). In particular, this structure will vary with choice of \(\Omega\).

Since \(\mathbf{M}_{\boldsymbol{\theta}_{0}}=\mathbf{I}_{P}\), we let \(\Omega\) regularize the deviation of \(\mathbf{M}\) from \(\mathbf{I}_{P}\). In order to capture a wide variety of behaviors, we consider regularizers built from the following class of spectral regularizers for a strictly quasi-convex function \(\omega\colon\mathbb{R}\to\mathbb{R}\) having minimum value \(\omega(1)=0\):

\[\Omega_{\omega}(\mathbf{M})=\sum_{j=1}^{P}\omega(\lambda_{j}(\mathbf{M})).\]

This penalty applies only to the eigenvalues, so the eigenvectors of \(\mathbf{M}\) are free. Simple examples of choices for \(\omega\) are \(\omega(\lambda)=|\lambda-1|^{p}\) for \(p>0\). For any penalty \(\omega\), we also define the **effective penalty**

\[\tilde{\omega}(v)\triangleq\min_{z\geq 1}\omega(z)+\frac{v^{2}}{z^{2}}\quad \text{and}\quad\widetilde{\Omega}_{\omega}(\mathbf{M})\triangleq\Omega_{ \tilde{\omega}}(\mathbf{M}).\]

As we will see, the bilinear constraint results in an equivalent formulation of the optimization problem in eq. (1) as an optimization with a linear constraint with penalty \(\widetilde{\Omega}\).

Without loss of generality, we will consider \(\mathbf{M}\) to be symmetric positive semidefinite, since we can always rotate \(\boldsymbol{\theta}\) accordingly without changing the optimization problem.

When considering solutions, we are also interested in the kernel corresponding to the learned features. Specifically, for \(\mathbf{x},\mathbf{x}^{\prime}\in\mathbb{R}^{D}\) we define the **adapted kernel** as the feature inner products

\[K(\mathbf{x},\mathbf{x}^{\prime})\triangleq\overline{\nabla f_{\widehat{ \boldsymbol{\theta}}}}(\mathbf{x})\overline{\nabla f_{\widehat{\boldsymbol{ \theta}}}}(\mathbf{x}^{\prime})^{\mathsf{T}}=\nabla_{\boldsymbol{\theta}_{0} }f_{\boldsymbol{\theta}_{0}}(\mathbf{x})\widehat{\mathbf{M}}^{2}\nabla_{ \boldsymbol{\theta}_{0}}f_{\boldsymbol{\theta}_{0}}(\mathbf{x}^{\prime})^{ \mathsf{T}}\in\mathbb{R}^{C\times C}.\]

We correspondingly define the initial kernel \(K_{0}(\mathbf{x},\mathbf{x}^{\prime})=\nabla_{\boldsymbol{\theta}_{0}}f_{ \boldsymbol{\theta}_{0}}(\mathbf{x})\nabla_{\boldsymbol{\theta}_{0}}f_{ \boldsymbol{\theta}_{0}}(\mathbf{x}^{\prime})^{\mathsf{T}}\) for \(\mathbf{M}=\mathbf{I}_{P}\), which is equal to the standard neural tangent kernel in neural networks.

### Structureless feature learning

It is instructive to first consider the effect of unrestricted optimization of \(\mathbf{M}\); that is, if the features were allowed to change without any structural constraints on how features and parameters must interact. In this case, we simply solve eq. (1) directly, using \(\Omega=\Omega_{\omega}\) applied to the full \(\mathbf{M}\).

**Theorem 1**.: _There is a solution to eq. (1) with \(\Omega=\Omega_{\omega}\) satisfying_

\[\widehat{\mathbf{M}}=\mathbf{I}_{P}+(s-1)\|\widehat{\boldsymbol{\beta}}\|_{2} ^{-2}\widehat{\boldsymbol{\beta}}\widehat{\boldsymbol{\beta}}^{\mathsf{T}} \quad\text{and}\quad\widehat{\boldsymbol{\theta}}=\boldsymbol{\theta}_{0}+s^{ -1}\widehat{\boldsymbol{\beta}}.\]
_where \(s=\arg\min_{z\geq 1}\omega(z)+\frac{\|\widehat{\mathbf{\beta}}\|_{2}^{2}}{z^{2}}\) and_

\[\widehat{\mathbf{\beta}}=\operatorname*{arg\,min}_{\mathbf{\beta}\in\mathbb{R}^{P}}\| \mathbf{\beta}\|_{2}\ \ \text{s.t.}\ \ \widehat{\mathbf{y}}_{i}-f_{\mathbf{\theta}_{0}}(\mathbf{x}_{i})=\nabla_{\mathbf{\theta}_{0} }f_{\mathbf{\theta}_{0}}(\mathbf{x}_{i})\mathbf{\beta}\ \forall\ i\in[N].\]

_Furthermore, the adapted kernel for this solution is given by_

\[K(\mathbf{x},\mathbf{x}^{\prime})=K_{0}(\mathbf{x},\mathbf{x}^{\prime})+(s^{2 }-1)\|\widehat{\mathbf{\beta}}\|_{2}^{-2}\underbrace{(f_{\widehat{\mathbf{\theta}}}( \mathbf{x})-f_{\mathbf{\theta}_{0}}(\mathbf{x}))(f_{\widehat{\mathbf{\theta}}}( \mathbf{x}^{\prime})-f_{\mathbf{\theta}_{0}}(\mathbf{x}^{\prime}))^{\mathsf{T}}}_ {\triangleq K_{\mathbf{\beta}(\mathbf{x},\mathbf{x}^{\prime})}}.\]

Note that the equivalent formulation in terms of \(\mathbf{\beta}\) corresponds exactly to standard ridgeless regression [30] using the initial tangent kernel features \(\nabla_{\mathbf{\theta}_{0}}f_{\mathbf{\theta}_{0}}(\mathbf{x})\)--therefore, when no structure is imposed on the adaptive features, the resulting predictions are exactly the same as in NTK regression. However, even though the predictions are no different, we already can see a qualitative difference in the description of the system from NTK analysis. Specifically, this adaptive feature perspective reveals how the adapted kernel itself changes: it is a low rank perturbation to the original kernel that directly captures the model output as the label kernel \(K_{\widehat{\mathbf{y}}}\). This kernel alignment effect has been empirically observed in real neural networks that depart from the lazy-training regime [18, 19, 20].

Moreover, the strength of the kernel alignment is directly related to the difficulty of the regression task as measured by the norm of \(\widehat{\mathbf{\beta}}\). Note that \(s\) takes minimum value at \(\|\widehat{\mathbf{\beta}}\|_{2}=0\) and is increasing otherwise, which means that if the initial tangent features \(\nabla_{\mathbf{\theta}_{0}}f_{\mathbf{\theta}_{0}}(\mathbf{x})\) are sufficient to fit \(\widehat{\mathbf{y}}_{i}\) with a small \(\widehat{\mathbf{\beta}}\), then \(s\approx 1\) and \(K(\mathbf{x},\mathbf{x}^{\prime})\approx K_{0}(\mathbf{x},\mathbf{x}^{\prime})\).2 It is only when the task is difficult and a larger \(\widehat{\mathbf{\beta}}\) is required that \(s^{2}-1\gg 0\) and we observe kernel alignment with the label kernel \(K_{\widehat{\mathbf{y}}}\). We illustrate this in an experiment with real neural networks on an MNIST regression task in Figure 1.

Footnote 2: While there is a \(\|\widehat{\mathbf{\beta}}\|_{2}^{-2}\) factor as well, this is always canceled by the implicit \(\|\widehat{\mathbf{\beta}}\|_{2}^{2}\) in \(K_{\widehat{\mathbf{y}}}\), such that the overall scale of \(K_{\widehat{\mathbf{y}}}\) is determined entirely by the factor \((s^{2}-1)\).

### Neural network structure

We now wish to specialize to structures that arise in neural networks. The developments in this section will be built around parameterization of the network by matrices, but these results could be extended straightforwardly to higher order tensors (such as filter kernels in convolutional neural networks). Generically, the parameters of a neural network are \(\mathbf{\theta}=\operatorname{concat}(\operatorname{vec}(\mathbf{W}_{1}),\dots, \operatorname{vec}(\mathbf{W}_{L}))\), where \(\mathbf{W}_{\ell}\in\mathbb{R}^{P_{\ell}\times Q_{\ell}}\) are parameter matrices such that \(P=\sum_{\ell=1}^{L}P_{\ell}Q_{\ell}\). These matrices comprise operations in a directed acyclic graph along with other operations such as nonlinearities and batch or layer normalizations, which do not have parameters and thus do not contribute tangent features. We first state our neural network model, and then we explain the motivation of each component.

Figure 1: **A more difficult task yield higher label kernel alignment.** We perform regression using a multi-layer perceptron on 500 MNIST digits from classes 2 and 3. We construct target labels \(y_{i}\) as the best linear fit of binary \(\pm 1\) labels using random neural network features. Then we train two networks, one (**left**) trained to predict \(y_{i}\), and one (**right**) trained to predict \(\operatorname{sign}(y_{i})\). We present the adapted kernel and label kernel matrices for data points ordered according to \(y_{i}\) and report the cosine similarity of the adapted kernel and the label kernel. The harder task of regression with binarized labels has a higher label kernel alignment. Further details are given in Appendix B.1.
**Model A**.: _The neural network model has the following properties:_

1. _The parameter vector_ \(\mathbf{\theta}-\mathbf{\theta}_{0}\in\mathbb{R}^{P}\) _consists of_ \(L\) _matrices_ \(\mathbf{W}_{1}\in\mathbb{R}^{P_{1}\times Q_{1}}\)__,...,_ \(\mathbf{W}_{L}\in\mathbb{R}^{P_{L}\times Q_{L}}\)_._
2. _The feature transformation operator_ \(\mathbf{M}\in\mathbb{R}^{P\times P}\) _is parameterized by_ \(\mathbf{M}_{\ell}^{(1)}\in\mathbb{R}^{P_{\ell}\times P_{\ell}}\) _and_ \(\mathbf{M}_{\ell}^{(2)}\in\mathbb{R}^{Q_{\ell}\times Q_{\ell}}\) _for each_ \(\ell\in[L]\) _such that application of_ \(\mathbf{M}\) _to_ \(\mathbf{\theta}\) _results in the mapping_ \((\mathbf{W}_{\ell})_{\ell=1}^{L}\mapsto(\mathbf{M}_{\ell}^{(1)}\mathbf{W}_{ \ell}\mathbf{M}_{\ell}^{(2)})_{\ell=1}^{L}\)_._
3. _For strictly quasi-convex functions_ \(\omega_{\ell}^{(1)}\) _and_ \(\omega_{\ell}^{(2)}\) _for each_ \(\ell\in[L]\) _minimized at_ \(\omega_{\ell}^{(1)}(1)=0\) _and_ \(\omega_{\ell}^{(2)}(1)=0\)_, the regularizer is given by_ \(\Omega(\mathbf{M})=\sum_{\ell=1}^{L}\Omega_{\omega_{\ell}^{(1)}}(\mathbf{M}_{ \ell}^{(1)})+\Omega_{\omega_{\ell}^{(2)}}(\mathbf{M}_{\ell}^{(2)})\)_._
4. _The final matrix_ \(\mathbf{W}_{L}\) _has_ \(Q_{L}=C\) _and a fixed transformation of_ \(\mathbf{M}_{L}^{(2)}=\mathbf{I}_{C}\) _(corresponding to_ \(\omega_{L}^{(2)}=\chi_{\{1\}}\)_), and there is a mapping_ \(\mathbf{z}\colon\mathbb{R}^{D}\to\mathbb{R}^{P_{\ell}}\) _such that_ \(\nabla_{\mathrm{vec}(\mathbf{W}_{L})}f_{\mathbf{\theta}_{0}}(\mathbf{x})=\mathbf{I }_{C}\otimes\mathbf{z}(\mathbf{x})^{\mathrm{T}}\)_._

The first component is simply a reparameterization of \(\mathbf{W}_{\ell}\) as the difference from initialization, to simplify notation. The other components are motivated as follows.

**Feature transformations.** Consider a single weight matrix \(\mathbf{W}_{\ell}\) and the gradient of the \(k\)-th output \(f_{\mathbf{\theta}}^{(k)}(\mathbf{x})\). The matrix structure of \(\mathbf{W}_{\ell}\) limits the way in which the gradient tends to change. As such, rather than considering a \(P_{\ell}Q_{\ell}\times P_{\ell}Q_{\ell}\) matrix \(\mathbf{M}_{\ell}\) such that \(\overline{\nabla_{\mathrm{vec}(\mathbf{W}_{\ell})_{\ell}}f_{\mathbf{\theta}_{0}} }(\mathbf{x})=\nabla_{\mathrm{vec}(\mathbf{W}_{\ell})}f_{\mathbf{\theta}_{0}}( \mathbf{x})\mathbf{M}_{\ell}\), it is more natural to consider a Kronecker factorization \(\mathbf{M}_{\ell}=\mathbf{M}_{\ell}^{(2)\mathrm{T}}\otimes\mathbf{M}_{\ell}^{ (1)}\) such that

\[\overline{\nabla_{\mathbf{W}_{\ell}}f_{\mathbf{\theta}}^{(k)}}(\mathbf{x})=\mathbf{ M}_{\ell}^{(2)}\nabla_{\mathbf{W}_{\ell}}f_{\mathbf{\theta}_{0}}^{(k)}(\mathbf{x}) \mathbf{M}_{\ell}^{(1)},\]

which corresponds to the mapping \(\mathbf{W}_{\ell}\mapsto\mathbf{M}_{\ell}^{(1)}\mathbf{W}_{\ell}\mathbf{M}_{ \ell}^{(2)}\).

**Independent optimization.** We assume that the feature transformations \(\mathbf{M}_{\ell}^{(1)}\) and \(\mathbf{M}_{\ell}^{(2)}\) are independently optimized from each other and from the feature transformations corresponding to other weights \(\ell^{\prime}\neq\ell\). The motivation for this assumption comes from the fact that at initialization in wide fully connected neural networks, the gradients at each layer are known to be uncorrelated [34]. Because of independence within each layer, we also need to define the **joint penalty** for \(v\geq 1\) as

\[(\omega_{1}\oplus\omega_{2})(v)\triangleq\inf_{1\leq z\leq v}\omega_{1}(z)+ \omega_{2}(\tfrac{v}{z}).\]

It is straightforward to verify that if \(\omega_{1}\) and \(\omega_{2}\) are strictly quasi-convex such that \(\omega_{1}(1)=\omega_{2}(1)=0\), then \(\omega_{1}\oplus\omega_{2}\) has the same properties; see Appendix A.1.

**Final weight matrix.** The final operation in most neural networks is linear matrix multiplication by the final weight matrix \(\mathbf{W}_{L}\in\mathbb{R}^{P_{L}\times C}\), such that if \(\mathbf{z}_{\mathbf{\theta}}(\mathbf{x})\in\mathbb{R}^{P_{L}}\) are the penultimate layer features, then the output is a vector \(f_{\mathbf{\theta}}(\mathbf{x})=\mathbf{W}_{L}^{\mathrm{T}}\mathbf{z}_{\mathbf{ \theta}}(\mathbf{x})\in\mathbb{R}^{C}\). Introducing \(\mathbf{M}_{L}\in\mathbb{R}^{P_{L}\times P_{L}}\) such that \(\mathbf{z}_{\mathbf{\theta}}(\mathbf{x})=\mathbf{M}_{L}\mathbf{z}_{\mathbf{\theta}_{0}} (\mathbf{x})\), we have the Kronecker formulation

\[f_{\mathbf{\theta}}(\mathbf{x})=\underbrace{(\mathbf{I}_{C}\otimes(\mathbf{z}_{\bm {\theta}_{0}}(\mathbf{x})^{\mathrm{T}}\mathbf{M}_{L}))}_{\overline{\nabla_{ \mathrm{vec}(\mathbf{W}_{L})}}f_{\mathbf{\theta}}(\mathbf{x})}\mathrm{vec}(\mathbf{ W}_{L}).\]

With the above neural network model in hand, this brings us to our main result on the solution to the adaptive feature learning problem.

**Theorem 2**.: _There is a solution to eq. (1) under Model A such that for each \(\ell\in[L]\),_

\[\widehat{\mathbf{M}}_{\ell}^{(1)}=\mathbf{U}_{\ell}\mathbf{S}_{\ell}^{(1)} \mathbf{U}_{\ell}^{\mathrm{T}},\quad\widehat{\mathbf{W}}_{\ell}=\mathbf{U}_{ \ell}\mathbf{\Sigma}_{\ell}\mathbf{V}_{\ell}^{\mathrm{T}},\quad\widehat{\mathbf{M }}_{\ell}^{(2)}=\mathbf{V}_{\ell}\mathbf{S}_{\ell}^{(2)}\mathbf{V}_{\ell}^{ \mathrm{T}},\]

_where \(\mathbf{U}_{\ell}\in\mathbb{R}^{P_{\ell}\times P_{\ell}}\), \(\mathbf{V}_{\ell}\in\mathbb{R}^{Q_{\ell}\times Q_{\ell}}\) are orthogonal matrices and \(\mathbf{S}_{\ell}^{(1)}=\mathrm{diag}(\mathbf{s}_{\ell}^{(1)})\), \(\mathbf{\Sigma}_{\ell}=\mathrm{diag}_{P_{\ell}\times Q_{\ell}}(\mathbf{\sigma}_{\ell})\), \(\mathbf{S}_{\ell}^{(2)}=\mathrm{diag}(\mathbf{s}_{\ell}^{(2)})\) are given by minimizers_

\[[\mathbf{s}_{\ell}^{(1)}]_{j},[\mathbf{\sigma}_{\ell}]_{j},[\mathbf{s}_{\ell}^{(2) }]_{j}=\operatorname*{arg\,min}_{s_{1},s_{2}\geq 1,\sigma\geq 0}\omega_{1}(s_{1})+ \omega_{2}(s_{2})+\sigma^{2}\text{ \ s.t. \ s}_{1}\sigma s_{2}=[\mathbf{d}_{\ell}]_{j}\text{ \ for }j\leq\min\{P_{\ell},Q_{\ell}\}\]
_and \([\mathbf{s}_{\ell}^{(1)}]_{j}=1\), \([\mathbf{s}_{\ell}^{(2)}]_{j}=1\) for \(j>\min\{P_{\ell},Q_{\ell}\}\), such that \(\widehat{\mathbf{B}}_{\ell}=\mathbf{U}_{\ell}\mathrm{diag}_{P_{\ell}\times Q_{ \ell}}(\mathbf{d}_{\ell})\mathbf{V}_{\ell}^{\mathsf{T}}\) satisfy_

\[(\widehat{\mathbf{B}}_{\ell})_{\ell=1}^{L}\in\operatorname*{arg\,min}_{ \mathbf{B}_{\ell}\in\mathbb{R}^{P_{\ell}\times Q_{\ell}}}\sum_{\ell=1}^{L} \widetilde{\Omega}_{\omega_{\ell}^{(1)}\oplus_{\omega_{\ell}^{(2)}}}(\mathbf{ B}_{\ell})\ \text{ s.t. }\ \widehat{\mathbf{y}}_{i}-f_{\mathbf{\theta}_{0}}(\mathbf{x}_{i})=\sum_{\ell=1}^{L} \nabla_{\mathrm{vec}(\mathbf{W}_{\ell})}f_{\mathbf{\theta}_{0}}(\mathbf{x}_{i}) \mathrm{vec}(\mathbf{B}_{\ell})\ \forall\ i\in[N].\]

In other words, we have reduced the bilinearly constrained optimization over \(\mathbf{M}\) and \(\mathbf{\theta}\) to a linearly constrained optimization over \((\mathbf{B}_{\ell})_{\ell=1}^{L}\), just as we did in the unstructured case. However, this time, due to the matrix structure of \(\mathbf{W}_{\ell}\) and corresponding structure in \(\mathbf{M}_{\ell}^{(1)}\) and \(\mathbf{M}_{\ell}^{(2)}\), the new equivalent optimization is much richer than simple minimum norm regression using tangent features. To better understand the regularization in this new problem, we have the following result.

**Proposition 3**.: _Let \(\omega\) be a continuous strictly quasi-convex function minimized at \(\omega(1)=0\). Then \(v^{2}\mapsto\tilde{w}(v)\) is an increasing concave function, and \(\tilde{\omega}(v)=v^{2}+o(v^{2})\)._

That is, no matter the original penalties \(\omega_{\ell}^{(1)}\) and \(\omega_{\ell}^{(2)}\), the resulting \(\widetilde{\Omega}_{\omega_{\ell}^{(1)}\oplus_{\omega_{\ell}^{(2)}}}\) will always be a spectral penalty with 1) sub-quadratic tail behavior, and 2) quadratic behavior for small values. We illustrate this for a few examples in Figure 2. For one example, when \(\omega(v)=(v-1)^{2}\), the effective penalty \(\tilde{\omega}(v)\) behaves like \(|v|\) for large \(v\), making \(\widetilde{\Omega}_{\omega}(\mathbf{B})\) like the nuclear norm for large singular values and like the Frobenius norm for small singular values. In general, \(\tilde{\omega}\) has slower tails than \(\omega\). This behavior is highly related of the equivalence of the nuclear norm as the sum of Frobenius norms of two factors [27], which coincides with the case \(\omega(v)=v^{2}\); however, since we constrain the \(\mathbf{M}\) to be near \(\mathbf{I}_{P}\), we retain Frobenius norm behavior near 0.

The result of this effective regularization is a model that is able to leverage structures through the group approximate low-rank penalty while also being robust to noise and model misspecification through the Frobenius norm penalty for small singular values. From this perspective, we can conceptually consider a decomposition of the matrices \(\mathbf{B}_{\ell}\approx\mathbf{B}_{\ell}^{\mathrm{LO}}+\mathbf{B}_{\ell}^{ \mathrm{NTK}}\), where \(\mathbf{B}_{\ell}^{\mathrm{LO}}\) are low rank and capture the structure learned from the data that is predictive of the training labels, while \(\mathbf{B}_{\ell}^{\mathrm{NTK}}\) form the component that fits the residual (after regression using \(\mathbf{B}_{\ell}^{\mathrm{LO}}\)) via standard NTK interpolation with Frobenius norm minimization. In this way, the adaptiveness of the neural network is able to get the benefits of both strategies.

We also see special structure in the final layer weights which are directly connected to the outputs. For example, for a single output (\(C=1\)) when \(\mathbf{B}_{L}\) is simply a vector \(\mathbf{\beta}_{L}\), the final layer contribution to the adapted kernel (decomposing as \(K(\mathbf{x},\mathbf{x}^{\prime})=\sum_{\ell=1}^{L}K_{\ell}(\mathbf{x}, \mathbf{x}^{\prime})\)) is similar to the structureless feature learning setting, but with only the final layer features:

\[K_{L}(\mathbf{x},\mathbf{x}^{\prime})=\mathbf{z}(\mathbf{x})^{\mathsf{T}}( \widehat{\mathbf{M}}_{L}^{(1)})^{2}\mathbf{z}(\mathbf{x}^{\prime})=\mathbf{z} (\mathbf{x})^{\mathsf{T}}\mathbf{z}(\mathbf{x}^{\prime})+(s_{L}^{2}-1)\| \widehat{\mathbf{\beta}}_{L}\|_{2}^{-2}(\mathbf{z}(\mathbf{x})^{\mathsf{T}} \widehat{\mathbf{\beta}}_{L})(\mathbf{z}(\mathbf{x}^{\prime})^{\mathsf{T}} \widehat{\mathbf{\beta}}_{L}).\]

That is, the contribution of the final layer to the kernel is like the label kernel, but instead of the labels, it consists of outputs representable by the final layer tangent features at initialization \(\mathbf{z}(\mathbf{x})^{\mathsf{T}}\widehat{\mathbf{\beta}}_{L}\). We can see this effect in Figure 1, where even in the case where the neural network is trained on binarized labels (right), the kernel bears some resemblance to the label kernel of the unbinarized labels (left).

Any residual component of the target function that is not representable by the final layer tangent features \(\mathbf{z}(\mathbf{x})\) must be built with the remaining weights \((\widehat{\mathbf{B}}_{\ell})_{\ell=1}^{L-1}\). If this residual is small, then the regularization behaves like the Frobenius norm, and the contribution to the kernel is simply the NTK. It may be the case, however, that for some data points, the residual is quite large, in which case it will take a large value of \(\widehat{\mathbf{B}}_{\ell}\) to close the gap, which will result in a change to the kernel.

To understand this change the kernel, consider the extremely simple case where \(\widehat{\mathbf{B}}_{\ell}\) has only a rank-one dimension in which it escapes the Frobenius norm penalty (call this \(\widehat{\mathbf{B}}_{\ell}^{\mathrm{LO}}\)). Then we observe the same form of the contribution as in the final layer case:

\[K_{\ell}(\mathbf{x},\mathbf{x}^{\prime}) =\nabla_{\mathrm{vec}(\mathbf{W}_{\ell})}f_{\mathbf{\theta}_{0}}( \mathbf{x})\nabla_{\mathrm{vec}(\mathbf{W}_{\ell})}f_{\mathbf{\theta}_{0}}( \mathbf{x}^{\prime})^{\mathsf{T}}\] \[\quad+(s_{\ell}^{2}-1)\|\widehat{\mathbf{B}}_{\ell}^{\mathrm{LO}} \|_{F}^{-2}\nabla_{\mathrm{vec}(\mathbf{W}_{\ell})}f_{\mathbf{\theta}_{0}}(\mathbf{x })\widehat{\mathbf{B}}_{\ell}^{\mathrm{LO}}\widehat{\mathbf{B}}_{\ell}^{ \mathrm{LO}^{\mathsf{T}}}\nabla_{\mathrm{vec}(\mathbf{W}_{\ell})}f_{\mathbf{\theta}_{0}} (\mathbf{x}^{\prime})^{\mathsf{T}}.\]
That is, it also contributes something akin to the label kernel; however, the magnitude of this contribution increases with the "difficulty" of the task in terms of the magnitude of \(\|\widehat{\mathbf{B}}_{F}^{\text{LO}}\|_{F}\) and correspondingly \(s_{\ell}\), and can be much larger than the label kernel. However, in the high dimensional tangent feature space, combined with the low rank regularization, this residual kernel will typically only be aligned with the data points that have a large residual. We can see this in Figure 1 and even more clearly in Figure 3, where for a network trained on binarized targets \(\operatorname{sign}(y_{i})\), the kernel has large values in the middle (corresponding to \(y_{i}\approx 0\), which are the most difficult points, being the place where the residual \(\operatorname{sign}(y)-y\) has a large jump discontinuity).

For higher rank perturbations, the result can be viewed as a sum of such rank-one residual kernels; however, the meaning of these components and the relative magnitude of their weightings becomes more difficult to reason about. As the target function becomes more complex (such as under excessive label noise), the weights are likely to become high rank in order to capture this complexity.

## 4 Discussion

In this work, we have proposed a framework of adaptive feature learning using tangent features that sheds light on some phenomena observed in real neural networks. The most important question one should ask, however, is the degree to which real neural networks actually fit this model, and whether there still remains a gap in our understanding.

**Sample complexity.** In terms of practical performance, we have found that the adaptive feature learning model offers significant sample complexity advantages over fixed neural tangent kernel features. In Figure 4, we compare the two models and real neural networks on MNIST and CIFAR-10. For the same test error, the adaptive feature model requires an order of magnitude fewer samples. However, given enough samples, both the adaptive and fixed feature models appear to converge to the same error floor, while the real neural networks significantly outperform both models and appear to have a faster convergence rate in the case of CNNs on CIFAR-10. This suggests that there are fundamental limitations of the tangent kernel feature space that even adaptive feature learning cannot avoid, although adaptivity does close a non-negligible part of the gap in the low-sample regime. An interesting direction for future work is to analyze the convergence rate of adaptive feature learning under Model A and determine under what conditions there might be improved rates over fixed tangent features.

Figure 3: **Adapted kernel reveals difficult structure.** For the neural network from Figure 1 trained on binarized labels \(\operatorname{sign}(y_{i})\) (\(\mathbf{left}\)), the target target function (green, solid) is difficult while the function \(\mathbf{x}\mapsto y\) (black, dotted) is easily predicted using tangent features. The network must learn (**right**) to fit the residual (red, dashed), which results in the kernel (orange, \(\times\)) being highly influenced by difficult training points (near \(y=0\)).
#### Average features vs. final features.

In our analysis we considered the average features over a linear path in parameter space, but this is difficult to work with in practice, compared to for example the features at the end of training. Of course, we can equivalently write the average features as an average transformation by \(\mathbf{M}_{\boldsymbol{\theta}^{\prime}}\) such that \(\nabla_{\boldsymbol{\theta}^{\prime}}f_{\boldsymbol{\theta}^{\prime}}(\mathbf{x })=\nabla_{\boldsymbol{\theta}_{0}}f_{\boldsymbol{\theta}_{0}}(\mathbf{x}) \mathbf{M}_{\boldsymbol{\theta}^{\prime}}\) along the path from \(\boldsymbol{\theta}_{0}\) to \(\boldsymbol{\theta}\):

\[\overline{\nabla f_{\boldsymbol{\theta}}}(\mathbf{x})=\nabla_{\boldsymbol{ \theta}_{0}}f_{\boldsymbol{\theta}_{0}}(\mathbf{x})\mathbf{M}=\nabla_{ \boldsymbol{\theta}_{0}}f_{\boldsymbol{\theta}_{0}}(\mathbf{x})\Big{(}\mathbf{ I}_{P}+\int_{0}^{1}(\mathbf{M}_{\boldsymbol{\theta}^{\prime}}-\mathbf{I}_{P}) \big{|}_{\boldsymbol{\theta}^{\prime}=(1-t)\boldsymbol{\theta}_{0}+t \boldsymbol{\theta}}dt\Big{)}.\]

In general, an integral of matrices, even if individually low rank, should not be a low rank matrix. However, the average \(\mathbf{M}\) is often a low-rank perturbation to \(\mathbf{I}_{P}\), which is a remarkable coincidence unless all \(\mathbf{M}_{\boldsymbol{\theta}^{\prime}}\) along this path have the same principal subspace as \(\mathbf{M}\) but with different eigenvalues. We thus expect the average features and final features to be similar, but not necessarily the same (in general, due to the averaging effect, the final tangent features will be larger for example). This difference can be important in downstream analysis--for example, in transfer learning, we would let \(\boldsymbol{\theta}_{0}\) be the solution to a previous optimization problem, and the initial tangent features would be the final tangent features of that optimization, rather than the average tangent features. Thus another direction for further study is the extent to which average features and final features coincide. In our limited (unpublished) observations, the final feature kernel closely resembles a re-scaled version of the average feature kernel, which coincides with recent results for linear networks [19].

The above limitations notwithstanding, we believe this framework still provides valuable insight into real phenomena in neural networks. We discuss two examples.

#### Benign overfitting.

Neural networks have shown remarkable resilience to noisy labels, sparking the recent theoretical research area of studying models that interpolate noisy labels yet still generalize well [3, 9, 13, 29]. Much research in this area has been concerned with finding fixed feature regimes in which this "benign overfitting" can occur in ridge(less) regression settings, but Donhauser et al. [32] have shown that when the ground truth function is sparsely represented by the features, remarkably, the optimal \(\ell_{p}\) penalty is not \(p\in\{1,2\}\), but rather in between, since some resemblance to sparsity-inducing \(p=1\) encourages learning structure, but something like \(p=2\) is necessary to absorb noise without harming prediction (for more discussion regarding the latter point, see [29]). Our results reflect these desired optimal properties precisely, since the effective penalties always have sub-quadratic tail behavior (promoting structure) and quadratic behavior near zero (absorbing noise). This raises another future research question, regarding the optimality of these effective penalties compared to \(\ell_{p},\,p\in(1,2)\) penalties.

Figure 4: **Adaptive feature learning improves low-sample performance.** We perform 10-class classification on MNIST (**left**) and CIFAR-10 (**right**) using three different models: a linear model using fixed tangent features (blue, solid), the adaptive feature optimization in eq.1 under Model A (orange, dashed), and standard neural networks (green, dash–dot). The adaptive feature achieves the same performance as the non-adaptive tangent feature model with an order of magnitude fewer samples. Further details are given in Appendix B.2.
Low rank optimization.Our framework sheds interesting insight on the success of low rank deviation optimizations in neural networks, which have been used to characterize task difficulty and compress networks [35] and recently to efficiently fine-tune pretrained large language models in the low rank adaptation (LoRA) method [36]. In LoRA, for example, each weight matrix is parameterized as \(\mathbf{W}_{\ell}=\mathbf{U}_{\ell}\mathbf{V}_{\ell}^{\mathsf{T}}\) for \(\mathbf{U}_{\ell}\in\mathbb{R}^{P_{\ell}\times R_{\ell}}\) and \(\mathbf{V}_{\ell}\in\mathbb{R}^{Q_{\ell}\times R_{\ell}}\), where \(R_{\ell}\ll P_{\ell},Q_{\ell}\). Due to the sub-quadratic spectral regularization of the adaptive feature optimization, if the target task is sufficiently related to the source task, such that the target function is well represented by the tangent features at \(\mathbf{\theta}_{0}\) of the pretrained model, then according to Theorem 2, the model is inclined to learn an approximately low-rank deviation from the initial parameters even if the \(\mathbf{W}_{\ell}\) were not explicitly constrained to be low rank. By constraining the weights to be low rank by design, LoRA can essentially recover the same solution, if not an even more aggressively structured one, at a fraction of the computational time and memory cost.

## Acknowledgements

The experiments in this paper were written in PyTorch [37] and JAX [38] with code assistance from GitHub Copilot and ChatGPT. DL was supported by ARO grant 2003514594. SA and computing resources were supported by Richard G. Baraniuk via NSF grants CCF-1911094, IIS-1838177, and IIS-1730574; ONR grants N00014-18-1-2571, N00014-20-1-2534, and MURI N00014-20-1-2787; AFOSR grant FA9550-22-1-0060; and a Vannevar Bush Faculty Fellowship, ONR grant N00014-18-1-2047.
