E FFICIENT C OMPUTATION OF D EEP N ONLINEAR ∞ - W IDTH N EURAL N ETWORKS T HAT L EARN F EATURES
Anonymous authors
Paper under double-blind review
A BSTRACT
While a popular limit of inﬁnite-width neural networks, the Neural Tangent Kernel (NTK) often exhibits performance gaps from ﬁnite-width neural networks on standard datasets, due to lack of feature learning.
Although the feature learning maximal update limit , or µ -limit (Yang and Hu, 2020) of wide networks has closed the gap for 1-hidden-layer linear models, no one has been able to demonstrate this for deep nonlinear multi-layer perceptrons (MLP) because of µ -limit’s computational difﬁculty in this setting.
Here, we solve this problem by proposing a novel feature learning limit, the π -limit , that bypasses the computational issues.
The π -limit, in short, is the limit of a form of projected gradient descent , and the π -limit of an MLP is roughly another MLP where gradients are appended to weights during training.
We prove its almost sure convergence with width using the Tensor Programs technique .
We evaluate it on CIFAR10 and Omniglot against
NTK as well as ﬁnite networks, ﬁnding the π -limit outperform ﬁnite-width models trained normally (without projection) in both settings, closing the performance gap between ﬁnite- and inﬁnite-width neural networks previously left by NTK.
[Figure]
We compare the representations from our best performing ﬁnite-width µ -parametrized network, our proposed ∞ -width π -limit, and NTK.
The representations from former two form neat clusters according to their classes while those from the latter are mixed up together.
See Appendix B.5.3.
I NTRODUCTION
The theory of Neural Tangent Kernel (NTK) (Jacot et al., 2018) led to many theoretical discoveries of neural networks, but they ultimately apply only to a limited, unrealistic regime where the network does not learn features.
Recently, Yang and Hu (2020) discovered the maximal update, or µ , parametrization that induces a feature learning ∞ -width limit, called the µ -limit of deep neural networks, which generalizes the mean ﬁeld limit of shallow networks (Sirignano and Spiliopoulos, 2018; Mei et al., 2018; Chizat and Bach, 2018; Rotskoff and Vanden-Eijnden, 2018).
They explicitly trained such ∞ -width limits on Word2Vec and MAML, tasks relying crucially on feature learning.
These µ -limits outperformed both NTK baselines and ﬁnite-width networks.
Thus, the µ -limit seems to capture key properties of neural networks in practice and is a promising direction to further our theoretical understanding of them.
However, a signiﬁcant roadblock to this end is that the µ -limit for general neural networks is both hard to compute empirically and hard to analyze theoretically.
Indeed, Yang and Hu (2020) restricted main experiments to linear MLPs where exact calculation is fast as a special case.
According to them, the main complications come from the interleaving of forward and backward propagations that leads to the need to evaluate Gaussian integrals involving nested nonlinearities.
Likewise, these complicated integrals also contribute to the difﬁculty of theoretical analysis, especially in deep networks.
Contributions In this work, we bypass these issues by considering a form of projected gradient descent instead.
Roughly speaking, with V φ denoting the V-transform (aka dual) of φ (Deﬁnition 3.1).
[Table]
( (cid:63) )
Here gradient accumulation is just the normal training process where gradient is added to a parameter, whereas in gradient concatenation , gradient is appended to a parameter.
The RHS retains the feature learning capability while making the computation and the mathematics much simpler compared to the µ -limit.
We explicitly evaluate our new limit, which we call the π -limit , on CIFAR10 and Omniglot, and compare against ﬁnite-width neural networks as well as NTK and NNGP.
We ﬁnd the π -limit always outperform wide standard-parametrized or µ -parametrized networks, and they all signiﬁcantly outperform the kernels.
For example, on Omniglot, this is evidenced by Fig.
We also demonstrate transfer learning from ImageNet to CIFAR10.
To our knowledge, this is the ﬁrst time deep nonlinear feature learning ∞ -width limits are evaluated on nontrivial natural datasets.
W HY IS THE µ -L IMIT H ARD TO C OMPUTE ? A N E XAMPLE
While the µ -limit is the “optimal” feature learning limit of a neural network as described in Yang and Hu (2020), it is difﬁcult to calculate analytically.
Let’s demonstrate this difﬁculty with a simple 1-hidden-layer µ -parametrized network (abbreviated µ -net , whose ∞ -width limit is the µ -limit )
f with 1-dimensional input ξ ∈ R and output f ( ξ ) ∈ R :
[Equation]
(1) where u, v ∈ R n are column vectors and φ : R → R is the nonlinearity (e.g. relu).
Note here n is the width of f .
According to µ -parametrization, we initialize u α , v α ∼ N (0 , 1 n ) .
At initialization, one easily sees lim n →∞ f ( ξ ) =
because
[Equation]
is the average of a large number of iid random variables ( √ nv ) α φ ( ξ √ nu α ) with zero mean.
In the language of Tensor Programs (Yang, 2019b;a; 2020a;b), we have an almost sure convergence
[Equation]
at initialization, where Z √ nv and Z √ nu represent the independent standard Gaussian random variables describing the coordinate distributions of √ nv and √ nu .
However, one can see that f ( ξ ) has a limit of the form
[Equation]
after 1 step of SGD
[Equation]
where u and v still represent values at initialization, and · · · represent terms coming from the gradient update.
Even though the · · · appearing inside φ () turn out to be all Gaussians, this expectation is hard to evaluate analytically due to the nesting of φ (cid:48) inside φ , unless φ is polynomial.
Furthermore, as more SGD steps are taken, this nesting quickly becomes more complicated, so even if φ is polynomial, the time needed for evaluation compounds exponentially.
S OLUTION : P ROJECT THE G RADIENT
To alleviate this difﬁculty, we propose to study the limit of training a neural network under a projected form of gradient descent which we call π -SGD , i.e., we update θ ←
θ − η Π ∇ θ L where Π is a linear projection and θ is the vector of all parameters.
We shall ﬁrst describe this projection Π for the simple motivating example above before doing so for the general case.
3. C ONTINUING THE 1-H IDDEN -L AYER E XAMPLE
In the example of Eq.
(1), Π acts on the gradients as follows: 1) Π leaves the output layer gradient ∇ v L untouched, but 2) projects the input layer gradient 2 ∇ u L = ( · · · ) φ (cid:48) ( · · · ) to the linear span of u 0 , the initial value of input weights u :
[Equation]
[Equation] [Equation]
Here Π u 0 denotes the orthogonal projection to the span of u 0 .
In particular, after projection, ∇ u L now has roughly iid Gaussian coordinates (proportional to u 0 ).
Then unlike Eq.
(2), we avoid the nesting of φ (cid:48) inside φ :
[Equation]
for some deterministic scalar ˜ c .
This expectation can then be evaluated routinely using the Vtransform of φ (e.g. Cho and Saul (2009) for relu; see Deﬁnition 3.1 below), and likewise, so can the expectations involved in all later steps.
After formalization inside a Tensor Program, this rigorously gives rise to the ∞ -width limit of f trained under π -SGD, which we call the π -limit of f .
See
Theorem 3.2.
Deﬁnition 3.1.
Let V φ : R 3 → R denote the V-transform of φ , such that V φ ( E XY, E X 2 , E Y 2 ) =
E φ ( X ) φ ( Y ) for any centered jointly Gaussian random variables ( X, Y ) ∈ R 2 .
We will treat V φ like an activation function and automatically vectorize when V φ is applied to 3 vectors of the same length.
3.1. C OMPUTING THE π -L IMIT
Here we sketch how to compute the π -limit of Eq.
Let u t , v t , f t denote corresponding objects at time t , with t = 0 denoting initialization.
Memory and Time Requirements What information do we need to store in order to calculate the π -limit?
Because of the projection, u t =
c t u 0 for some c t ∈ R .
So to calculate the limit of f t , we need to track c t (or rather its ∞ -width limit).
This memory requirement is Θ(1) in training time t .
To see the form of v t , it helps to simplify our setup a bit more by 1) initializing v 0 = 0 and by 2) assuming each π -SGD step involves a minibatch containing a sole input ξ t .
Correspondingly, because the gradient of v t is always proportional to φ ( ξ t √ nu t ) =
φ ( ξ t c t √ nu 0 ) , there exist coefﬁcients { a s ∈ R } t − 1 s =0 , { b s ∈ R } t − 1 s =0 (which are random, but deterministic conditioned on u 0 , v 0 ) such that
[Equation]
Here b s = ξ s c s and a s is formed from the learning rate and loss derivative.
So to calculate the limit of f t , we need { a s ∈ R } t − 1 s =0 , { b s ∈ R } t − 1 s =0 (or rather their n → ∞ limits).
Note that c t , a s ,
b s all have nontrivial ﬂuctuations for ﬁnite n , but become deterministic as n → ∞ .
This implies a memory requirement of Θ( t ) , which is also the total requirement for computing the limit of f t .
As we will see, each forward and backward pass has runtime Θ( t ) as well, so the total runtime for calculating lim n →∞ f t is Θ( t 2 ) , 5 compared to the exponential runtime of general µ -limit.
Forward Pass Using Eq.
(4), we can intuit (and formalize using Tensor Programs)
that f t
[Equation]
=
[Equation]
n
[Equation]
where ˚ a s , ˚ b s , ˚ c t denote the deterministic limits of the corresponding quantities.
The expectation in the
RHS can be evaluated using V-transforms.
The t terms in the summation implies a runtime of Θ( t ) .
Backward Pass
The gradient update for the output weights v t is clearly represented by setting ( a t +1 , b t +1 ) ← ( − η L (cid:48) , ξc t ) , where L (cid:48) denotes the loss derivative ∂
L ( f t ( ξ ) , label) /∂f t ( ξ ) .
However, a priori, it seems unclear how to calculate the limit of the projected gradient Π u 0 ∇ u t L of the input weights.
Fortunately, because we can express the entire unrolled π -SGD training inside a Tensor
Program, the Master Theorem (Yang, 2020b) can automatically give us the exact formulas for the gradient limit; see Appendix F.
But in fact, there is a more intuitive way of obtaining the backward pass limit by recognizing that a projected gradient is just the maximal ascent direction in the projected space.
In our case, Π u 0 ∇ u t L
Assuming minibatch size is constant, and we train for a constant number of epochs, this translates to Θ( N 2 ) to train lim f t , where N is dataset size, and Θ( N ) to do inference on a single input.
This compares favorably with Gaussian Processes, which requires Θ( N 3 ) to “train” (i.e. inverting the kernel), and Θ( N ) for inference.
However, in our experiments here, the constant in our Θ( N 2 ) in practice will make training lim f t slower than the corresponding NTK or NNGP kernel regression.
Summary Below, for brevity, we say training routine to mean the package of learning rate η , training sequence of singleton minibatches { ( ξ t , y t ) } t ≥ 0 (where ξ t is the input and y t is the label),and a loss function L ( f ( ξ ) , y ) that is continuously differentiable in the prediction of the model f ( ξ ) .
Theorem 3.2.
Consider the simple motivating example of 1-hidden-layer network f in Eq.
(1) with 1dimensional input and outputs, initialized by u α ∼ N (0 , 1 /n ) and v α ← 0 .
Suppose its nonlinearity φ has a polynomially bounded 2nd derivative.
For any training routine, f trained by π -SGD for T steps has an ∞ -width limit ˚ f T , in the sense that,
[Equation]
where the limit is almost sure.
˚ f T is given as follows:
Initialize A 0 , B 0 as empty column vectors, and initialize C 0 ∈ R as 1.
Recall V φ is the V-transform of φ .
For each t = 0 , 1 , . .
. , deﬁne the function (where dot product of empty vectors is 0)
[Equation]
where ◦ denotes coordinatewise product and 1 denotes the all-1s vector of appropriate shape, and
A t , B t , C t are inductively given by C t +1 def =
C t − η ∇ C t L (˚ f t
[Equation]
ξ t
[Equation]
, y t
[Equation]
and
[Equation]
where append ( v, p ) appends element p to the end of vector v , increasing the dimension of v by 1.
Here A t , B t correspond to the column vectors formed from the ∞ -width limits of { a s ∈ R } t − 1 s =0 , { b s ∈
R } t − 1 s =0 , and C t corresponds to the limit of c t .
As discussed above, computing ˚ f T requires Θ( T ) memory and Θ( T 2 ) time.
Theorem 3.2 is summarized by Fig.
3. π -P ARAMETRIZATION FOR D EEP N ETWORKS
We can straightforwardly generalize π -SGD and the π -limit theorem (Theorem 3.2) to deep MLPs.
However, due to the n × n Gaussian matrix initialization in the middle of the network, the memory requirement will be Θ( T 2 ) and runtime will be Θ( T 3 ) for training T steps, for the same reason as discussed in Yang and Hu (2020, Sec 8).
This is not scalable to datasets like CIFAR10.
Therefore, we propose a different initialization that brings down the memory requirement to Θ( T ) and runtime to Θ( T 2 ) , just like the 1-hidden-layer case.
Consider an L -hidden-layer µ -net f : R d → R d out : For weight matrices w 1 ∈ R n × d and w 2 , . .
. , w L ∈ R n × n , and nonlinearity φ : R → R , such a neural network on input ξ ∈ R d is given by h 1 ( ξ ) = √ nw 1 ξ ∈ R n , and x l ( ξ ) =
φ ( h l ( ξ ))
∈ R n , h l +1 ( ξ ) =
w l +1 x l ( ξ ) ∈ R n , for l = 1 , . . .
, L − 1 , (7) and the network output is f ( ξ ) =
n − 1 / 2 w L +1 x L ( ξ ) for w L +1 ∈ R d out × n .
In µ -parametrization, we would initialize w lαβ ∼ N (0 , 1 /n ) for any α, β, l .
Next, we describe our alternative proposal.
π -Initialization Choose integers r, M , to be explained shortly; these numbers should all be thought of as constant in width n .
Suppose we are given a collection P of matrices
[Equation]
Then we can initialize weights w l by ﬁrst sampling a standard random Gaussian matrix Ω ∈ R n × r , Ω αi ∼ N (0 , 1) , before setting
[Equation]
As an example, the initialization of the 1-hidden-layer case in Theorem 3.2 corresponds to d = d out = r = 1 , M = 0 and A 1 = 1 .
See Fig. 3(Left) for an illustration.
Let’s digest this initialization scheme: 1) Of course, A l , B l need to be initialized themselves, and in this paper we will just do so with Gaussian random initialization; see Appendix B.1 for the pseudo-algorithm.
A l , B l , l ∈ [2 , L ] , will play the same roles as A, B in the 1-hidden-layer example above, whereas A 1 plays the same role as C there.
3) r is a measure of dimension of the the projection.
More precisely, we will project hidden weight gradients from R n × n to R n × r ; see below.
A l , B l will grow in the M dimension with training time, just like how the sizes of A, B grow in the 1-hidden-layer example.
5) Eq.
(9) can be interpreted as saying w l are generated by 1-hidden-layer MLPs F A l ,B l : R r → ( R r or R d out ) with weights A l , B l , like so:
[Equation]
, F
[Equation]
l ,B l (Ω
[Equation]
) (cid:105) , for all hidden l = 2 , . . .
, L ;
[Equation]
This is reminiscent of hypernetworks (Ha et al., 2016), but the crucial difference is that the “batch dimension” (indexed by α, β ) of the generator F A l ,B l becomes the width dimension of the generated network f , so that the same generators can generate networks of arbitrary width.
π -Projection
We leave the output weights w L +1 alone, but for any w ∈ { w 1 , w 2 , . . .
, w L } , where w has shape n × n or n × d , we project the gradient ∇ w L of the same shape by left multiplying by Π Ω , the projection matrix to the r -dimensional space spanned by the columns of Ω .
This means that preactivations h l are always in this space.
However, the input side of w is not projected, so we are optimizing the hidden weights in a ( r × n ) -dimensional space, which still becomes inﬁnitedimensional as n → ∞ .
We refer to SGD with π -projected gradients as π -SGD .
This brings us to the following
Deﬁnition 3.3 ( π -Parametrization and π -Limit) .
We deﬁne π -parametrization as the package of π -initialization and π -SGD, 9 where r and the initial M are clear from context.
A network in π -parametrization is abbreviated π -network or π -net , and we deﬁne π -limit as its ∞ -width limit.
3.2. C OMPUTING THE π -L IMIT OF D EEP MLP
Here we show the π -limit forward and backward passes can be efﬁciently computed.
π -Limit Forward
Pass Let’s write f P for the (random) neural network π -initialized by P as in
Eqs.
(8) and (9).
Then it’s straightforward to see f P has a deterministic limit ˚ f P expressible as a composition of matrix multiplications and V-transforms:
Theorem 3.4 ( π -Limit Forward Pass) .
Suppose φ has a polynomially bounded 2nd derivative.
Then as n → ∞ , f P ( ξ ) → ˚ f P ( ξ ) , for any ξ ∈ R d , where convergence is almost sure, and ˚ f P ( ξ ) is deﬁned as follows.
Write g 1 def
=
A 1 (cid:62) ξ ∈ R r ,
[Equation]
(cid:26) R r for l = 2 , . . .
, L
R d out for l = L + 1 , and ˚ f P ( ξ ) def =
g L +1 , where B ◦ B yields a size- M column vector of squared norms of B ’s rows.
π -Limit Backward Pass Just like the 1-hidden-layer case, we can interpret the projected gradient as the maximal ascent direction in the column space of Ω .
Formalizing the intuition using Tensor
Programs yields the following theorem that we also empirically verify in Appendix E.
Theorem 3.5 ( π -Limit Backward Pass) .
Let P 0 denote the matrices (c.f. Eq.
(8) ) used in the πinitialization of f as in Eq.
Suppose its nonlinearity φ has a polynomially bounded 2nd derivative.
Then for any training routine, f trained by π -SGD for T steps has an ∞ -width limit which is equal to ˚ f P T for some P T given below, i.e. as n → ∞ , f T ( ξ ) → ˚ f P T ( ξ ) , for any ξ ∈ R d , where the limit is almost sure.
P T is given inductively through its elements A lT , B lT as follows:
A 1 t +1 def
=
A 1 t − η ∇ A 1 L (˚ f P t ( ξ t ) , y t ) , and, for l = 2 , . . .
, L + 1 ,
[Equation]
[Equation]
[Equation] [Equation] l − 1 t ) . Here g lt corresponds to g l in Theorem 3.4 evaluated for input ξ t and function ˚ f P t , and append ( B, g ) means appending g as a new row vector of B , increasing by 1 the dimension represented by M .
Theorems 3.4 and 3.5 are summarized by Fig.
10 As conveyed by Eq.
( (cid:63) ), the π -limit can be thought of as another MLP with activation V φ and trained by gradient concatenation.
11 Because π -parametrization has similar scaling with width as µ -parametrization, it is easy to show that the former admits feature learning in the same way the latter does, i.e. the feature kernel of every layer evolves nontrivially during training.
Using Tensor Programs, it is straightforward to extend the above theorems to more general settings such as biases, large batch size, per layer learning rate, or metalearning (Appendix A).
Appendix D also makes several observations on the π -limit theorems.
π -Limit vs µ -Limit While the projection means optimization is slower in the π -limit, we believe with sufﬁciently large r , the π -limit should perform similarly to the µ -limit.
Indeed, prior works such as Li et al.
(2018) has shown that optimizing a neural network in a much smaller, random parameter subspace can recover most of the performance of training all parameters.
This is also evidenced by our experimental results below, where the π -limit generally only slightly outperform wide µ -networks.
[Table]
E XPERIMENTS
Here we compare the performance of the relu π -limit on CIFAR10 (Krizhevsky, 2009) and Omniglot (Lake et al., 2015) against that of NTK, NNGP, and ﬁnite-width π - and µ -nets.
12 As we will see, 1) π -limits of large enough r beat ﬁnite µ -nets; 2) ﬁnite π -nets underperform the above on CIFAR but, interestingly, outperform them on Omniglot; 3) all of the above beat NTK and NNGP.
4. C LASSIFICATION ON CIFAR
Experimental Setup For π - and µ -networks, inﬁnite or ﬁnite, we train for 50 epochs.
We adopt a step learning rate schedule, with a learning rate drop of 0.15 at a certain milestone, which is a hyperparameter.
We sweep over a variety of hyperparameters such as the learning rate, gradient clipping, weight decay, the LR drop milestone, etc, as well as width, r , and depth.
For NTK and
NNGP, we perform kernel regression following Lee et al.
(2018) using centered labels.
For them, we sweep over the initialization variances and learning rate multipliers, along with ridge coefﬁcient.
See
Appendix B for more details.
Results In the literature, the performance gap between CNN and its NTK (Arora et al., 2019) is often cited for the deﬁciency of the NTK theory for explaining neural networks in practice.
However, in fact, on MLPs we already see a nontrivial gap, as seen in Table 1 (compare µ -Net with NTK ,
NNGP ).
The π -limit closes this gap, having almost a 2-point advantage over NTK.
The ﬁnite-width π -net outperforms NTK and underperforms π -limit both by about 1 point.
Feature Learning in π -Limit and Finite Networks
We show the advantage of π -limit over NTK and NNGP is due to feature learning.
To do so, we track the kernel regression performance of its (last-layer) feature kernel over the course of training.
As seen in Fig. 4(Left), while the feature kernel at initialization underperforms both NNGP and NTK, it improves consistently over time to eventually exceed them.
Similarly, the kernel performance of the best π - and µ -nets improves over time as well, though the accuracy of the feature kernel regression is slightly less than the network itself, likely due to the low rank property of the feature kernel.
On the contrary, the π -limit beneﬁts from feature kernel regression, improving the test accuracy from 61.5% to 61.85%; see Table 7.
Effects of Width, r , and Depth As shown in Fig.
4(Middle), accuracy of π -net increases monotonically with width and with r , approaching the accuracy of π -limit and µ -net, respectively, from below.
This can also be seen in feature kernel regression across training time, Fig.
In contrast, performance is not monotonic in depth for any of µ -net, π -net, or π -limit, beyond a consistent improvement from 1 to 2 hidden layers; see Fig. 6.
[Figure]
4. T RANSFER L EARNING FROM I MAGE N ET TO CIFAR
Pretraining-and-transfer-learning is an important setting where feature learning is crucial.
As pointed out in Yang and Hu (2020), the NTK limit trivializes pretraining, while the µ -limit both theoretical and empirically beneﬁt from it.
Here we investigate pretraining for the π -limit in the image domain by pretraining on ImageNet32 (Chrabaszcz et al., 2017) 15 and transferring to CIFAR10.
Experimental Setup We pretrain the π -limit with r = 400 for 30 epochs on a ﬁxed subset of
ImageNet32 with 250 (out of 1000) randomly subsampled classes.
To evaluate on CIFAR10, we compute the kernel induced by the pretrained ﬁnal-layer-features and perform kernel regression.
Results As shown in Table 1, ImageNet32 pretraining nontrivially raises the downstream CIFAR performance over without pretraining, altogether creating a 5% gap compared to NTK (which has the same test accuracy whether or not it is pretrained on ImageNet32 because of the disparate classes, as shown in Yang and Hu (2020)).
This again demonstrates the feature learning capability of the π -limit.
4. F EW -S HOT L EARNING VIA M ETALEARNING ON O MNIGLOT
Following Yang and Hu (2020), we also evaluate few-shot learning on Omniglot.
Compared to traditional classiﬁcation settings like CIFAR10, doing well on Omniglot requires learning features that can rapidly be adapted to new unseen data (Raghu et al., 2019).
We will adopt a metalearning approach to this, following Finn et al.
Unlike Yang and Hu (2020), we will train our models using ANIL ( Almost No Inner Loop ) (Raghu et al., 2019) rather than ﬁrst-order MAML ( Model Agnostic Meta-Learning ).
Brieﬂy, ANIL is a variant of second-order MAML, where in the inner loop, we only adapt the output layer, while in the outer loop, we only train the network body.
We adopt ANIL because: 1) We would like to train deep MLPs with SGD and without Adam or batchnorm, but empirically this makes optimization of ﬁrst-order MAML difﬁcult for both ﬁnite and ∞ -width networks compared to second-order MAML.
2) While a π -parametrized network trained by second-order MAML has an ∞ -width limit calculable using Tensor Programs, it does not stay in the f P form of Theorem 3.4 because of the second-order gradient, rendering the limit computation inefﬁcient.
Fortunately, π -networks trained by ANIL do not have this issue because the inner and outer loop gradients are on different parameters.
In addition, ANIL performs on par with MAML on standard ﬁnite-width networks (Raghu et al., 2019).
4) Furthermore, ANIL training more clearly delineates the role of the network body for learning reusable features and the role of the head for adapting to new data, which is more ﬁtting for our goals in this paper.
Experimental Setup We focus on the 5-way, 1-shot task 16 , with only 1 step of ANIL adaption.
For π - and µ -networks, inﬁnite or ﬁnite, we train for 50 epochs, 1000 batches per epoch, and 8 tasks per batch.
Following Antoniou et al.
(2019), we use cosine annealing schedule on the meta-learning rate.
For NTK and NNGP, ANIL meta-training has no effect, and meta-testing amounts to just taking kernel gradient descent step on each task.
17 We sweep over a variety of hyperparameters such as the outer and inner learning rates, gradient clipping, etc, as well as width, r , and depth.
See Appendix B for more details.
Results As seen in Table 1, µ -net, π -net, and π -limit all outperform NNGP and NTK by aboutpoints.
Interestingly, while π -limit is slightly better than µ -net, they are both outperformed by π -net.
This is related to the width nonmonotonicity seen on Omniglot, as we describe next.
Effect of Width, r , and Depth While π -net performance with r is still roughly monotonic, in contrast to in CIFAR10, here it can decrease with width, as seen in Fig.
4(Right).
In addition, µ -net seems to underperform π -net of the same width for sufﬁciently large r .
We ﬁnd this is primarily due to optimization and not generalization because Omniglot is hard to overﬁt; see Fig.
Counterintuitively, π -net of large r can optimize faster than µ -net despite the projection.
This is likely a side effect of π -initialization, as we see this advantage persist when r > width , where projection is a no-op.
We do not see width nonmonotonicity on CIFAR10 because overﬁtting is much easier there, so the results depend much more on the implicit bias of each model.
On the other hand, performance seems to monotonically increase along diagonals of ﬁxed r/width ratio, peaking around r/width ≈ 1 / 2 .
This suggests a different limit of r ∝ width → ∞ , which warrants investigation in future works.
Unlike CIFAR10, test accuracy is more monotonic in depth here (Fig. 6).
Again, this is probably because the extra weights help with optimization, which is the bottleneck on Omniglot.
R ELATED W ORKS
The π -limit bears some superﬁcial similarities to hierarchical kernel processes such as deep GPs (Damianou and Lawrence, 2013; Salimbeni and Deisenroth, 2017) or deep kernel processes (Aitchison et al., 2021).
The crucial difference here is that 1) during inference, the π -limit is deterministic and not a hierarchically sampled process, and 2) the π -limit is not trained variationally.
Relatedly, deep GPs or kernel processes can be construed as a kind of randomly initialized inﬁnitewidth neural network with ﬁnite bottleneck layers (Agrawal et al., 2020; Aitchison, 2020).
However, their variational inference procedures do not correspond to the limit of SGD.
In contrast, Littwin et al.
(2021) derived such an SGD limit where the network between consecutive bottlenecks is in
NTK parametrization, and the limit takes the form of a kind of intertwined system of kernel gradient descent.
However, the computation of this limit is not scalable to datasets like CIFAR10.
During training, the π -limit’s A l and B l grow in size, reminiscent of works on growing neural networks (Liu et al., 2019; Wu et al., 2021; Gu et al., 2021;
Gong et al., 2019).
Of course, these works focus on improving training efﬁciency of practical models, which is not our goal here.
Nevertheless, it will be interesting to explore whether insights from π -limit can contribute back to this literature.
Our π -initialization (given ﬁxed P ) can be construed as a special form of the initialization in deep mean ﬁeld limits of MLPs (Araújo et al., 2019; Sirignano and Spiliopoulos, 2020; Fang et al., 2020;
Nguyen, 2019; Nguyen and Pham, 2020).
However, without π -projection, these limits’ analytical computation suffers exponential runtime blowup like the µ -limit.
C ONCLUSION
A good model for studying wide neural networks should 1) capture the desired behavior of neural networks in practice, including feature learning and performance on real datasets, and 2) come with a set of theoretical tools for working researchers.
But an attractive model must also balance these properties with 3) computational simplicity for empirical investigations, and 4) mathematical simplicity for theoretical investigations.
Previously, NTK and NNGP satisfy 2), 3), 4) but not 1), as shown in this and prior works.
The µ -limit satisﬁes 1) and 2) but, arguably, not 3) and 4).
In this work, we presented the π -limit which, we believe, satisﬁes all 4 properties and that can prove fruitful for future researches.
While our work closed the ∞ -width performance gap, it also opens many new questions, e.g., Can we derive the π -limit for a ﬁxed r/width ratio?
When precisely does the width nonmonotonicity of π -net occur?
What about modern architectures, beyond MLP?
We leave these questions to future work.
A E XTENSIONS OF π -L IMIT A. B IASES
Biases are straightforward to add.
Consider an L -hidden-layer µ -net f : R d → R d out with biases: such a neural network on input ξ ∈ R d is given by h 1 ( ξ )
= √ nw 1 ξ + √ nb 1 ∈ R n , and x l ( ξ )
=
φ ( h l ( ξ ))
∈ R n , h l +1 ( ξ ) =
w l +1 x l ( ξ ) + √ nb l ∈ R n , for l = 1 , . . .
, L − 1 , (11) and the network output is f ( ξ ) =
n − 1 / 2 w L +1 x L ( ξ ) + b L +1 for w L +1 ∈ R d out × n .
In µparametrization, the biases are initialized as b lα ∼ N (0 , σ 2 b /n ) for l = 1 , . . .
, L , and b L +1 i ∼
N (0 , σ 2 b ) , for a hyperparameter σ b .
However, in π -parametrization, we will do this a bit differently.
π -Initialization Suppose we are given a collection P of matrices and vectors
[Equation]
Then we can initialize weights and biases by ﬁrst sampling a standard random Gaussian matrix Ω ∈ R n × r , Ω αi ∼ N (0 , 1) , and 1) set w l as in Eq.
(9), and 2) set biases as b L +1 ← β L + and b l ← 1 √ n Ω β l for all l = 1 , . . . , L.
This constitutes the π -initialization of biases.
π -Projection
For π -projection, we also project the bias gradients by Π Ω .
π -Limit Calculation
Theorem A.1 ( π -Limit Forward Pass with Bias) .
Let P be some collection as in Eq.
As n → ∞ , f P
[Equation]
, for any ξ ∈ R d , where convergence is almost sure, and ˚ f P ( ξ ) is deﬁned as follows.
Write g 1 def
=
A 1 (cid:62) ξ + β 1 ∈ R r , g
[Equation]
def =
[Equation] [Equation]
g
[Equation]
+
[Equation]
r for l = 2 , . . .
, L
R d out for l = L + 1 , and ˚ f P ( ξ ) def =
g L +1 , where B ◦ B yields a size- M column vector of squared norms of B ’s rows.
Theorem A.2 ( π -Limit Backward Pass) .
For the same setting as in Theorem 3.5 but with P as in
Eq.
(12) , A l , B l are updated exactly as in Theorem 3.5, and β l is updated by gradient accumulation like A 1 :
[Equation]
for all l = 1 , . . .
, L +
A.
P ARAMETER M ULTIPLIERS
We can insert to Eq.
(11) constant parameter multipliers α w for each parameter w like so
[Equation]
for l
[Equation] [Equation]
These multipliers are tuneable hyperparameters.
They affect both the forward and backward passes of the network.
In the π -limit, the forward pass is the same as in Theorem A.1 except we replace A l with α w l A l and β l with α b l β l .
The backward pass is the same as in Theorem A.2, but we just have to make sure that we backprop through the multipliers (in contrast, if we instead have absorbed the multipliers into the initialization, then we would not backprop through the multilpliers).
In our experiments, we only consider the input weight multiplier, output weight multiplier, and a single multiplier for all biases.
A. L EARNING R ATE M ULTIPLIERS
We may have custom learning rates for speciﬁc weights or biases.
In our experiments, we implement this with learning rate multipliers relative to the global learning rate η , e.g., the learning rate of a parameter w becomes γ w η if the multiplier is denoted γ w .
Then in the limit, we just need to replace η in Theorem 3.5 or Theorem A.2 with γ w l η , i.e.
[Equation]
for all l = 2 ,
[Equation]
, L
[Equation] [Equation]
for all l = 1 , . . .
, L + 1 .
But note we do not modify the B l update.
In our experiments, we will only sometimes use a LR multiplier on the input layer (“Input Layer LR Mult”), one on the output layer (“Output Layer LR
Mult”), and/or a single multiplier for all biases (“Bias LR Mult”).
A. L ARGE B ATCH S IZE
For batch size S > 1 , we still accumulate gradients into A 1 and β l as if they are regular parameters, and for the hidden weights, we just append the S gradient vectors as if they are from S unit-sized batches.
A. G RADIENT C LIPPING
The Frobenius norms of the projected gradients of the input weight w 1 and all biases b l converge to exactly the Frobenius norms of the gradients of A 1 and β l , i.e.
[Equation]
where the convergence is almost sure.
For any hidden weights w l , suppose its gradient over an
S -sized batch in the π -limit is given by
[Equation]
where ˜ A, ˜ B ∈ R S × r , and append ( B, ˜ B ) means appending all S rows of ˜ B into B , increasing the latter’s column length by S .
Then
[Equation]
where the convergence is almost sure.
Then clipping the gradient in the π -limit just means clipping the gradients of A 1 and β l , and, for hidden weights, rescale ˜
A such that
[Equation]
is at most the clipping threshold.
A. W EIGHT D ECAY
The limit of decaying w l ← w l (1 − lr · wd ) is just A l ← A l (1 − lr · wd ) .
A. D ECOUPLING L AYERS
We can actually use different r s and M s for every layer, and a similar theorem can be proved.
We can also use different, independently sampled Ω for each layer.
However, the limit would be exactly the same as before.
Nevertheless, this could make ﬁnite-width networks a less biased estimator of the ∞ -width limit.
B E XPERIMENTAL D ETAILS
All of our experiments are done on V100 GPUs.
All of our networks use relu activation.
The
V-transform of relu is (Cho and Saul, 2009)
[Equation]
[Table]
[Equation]
[Table]
B. N ETWORK I NITIALIZATION AND P ARAMETRIZATION
At initialization, 1) biases are always 0, 2) M is set to r so that A l , B l are square matrices for l = 2 , . . .
, L , and 3) we sample A l , B l as standard Gaussians and then scale them as follows:
A 1 =
A 1 /A 1 .norm ( dim = 0 , keepdim =
True )
B l = B l /B l .norm ( dim = 1 , keepdim =
True ) for all l = 2 , . . .
, L
A l = A l / (cid:112)
A 1 .shape [0] for all l = 2 , . . .
, L , and
A L +1 =
B.
F EATURE K ERNEL
When we talk about the “feature kernel” of a π -limit ˚ f P , we always mean the n → ∞ limit of the feature kernel of f P , and not the kernel induced by g L in ˚ f P .
This feature kernel K on inputs { ξ 1 , . . . , ξ k } is calculated as
[Equation]
where g Li and g Lj are the g L in Theorem 3.4 evaluated for two inputs ξ i and ξ j .
B. CIFAR B.3. µ -N ET , π -N ET , AND π -L IMIT
For π - and µ -networks, inﬁnite or ﬁnite, we train for 50 epochs.
We adopt a step learning rate schedule, with a learning rate drop of 0.15 at a milestone hyperparameter.
We also clip gradients with a hyperparameter threshold, where the clip is triggered individually for each parameter by the parameter’s norm, rather than total parameter norm.
Hyperparameter Optimization We ﬁrst perform random search on hyperparameters listed in
Table 2 for 2-hidden-layer π -net, π -limit, and µ -net: We sample at least 512 hyperparameter combinations 1) for each ( width, r ) ∈ { 128 , 256 , 512 , 1024 , 2048 } × { 50 , 100 , 200 , 400 } for π -net, 2) for each width ∈ { 128 , 256 , 512 , 1024 , 2048 } for µ -net, and 3) for each r ∈ { 50 , 100 , 200 , 400 } for π -limit.
The best accuracies per ( width, r ) are shown in Fig.
4(Middle).
Then we take the best hyperparameter combinations overall for π -net, π -limit, and µ -net, and for depth ∈ { 1 , 2 , 3 , 4 } (where depth denotes number of hidden layers), we resweep only learning rate and weight decay in a grid search, where the grid is ( lr ∗ · { 2 − 2 , 2 − 1 . 5 , . . .
, 2 1 . 5 } ) × ( wd ∗ · { 2 − 2 , 2 − 1 . 5 , . . .
, 2 1 . 5 } ) and lr ∗ and wd ∗ are the optimal hyperparameters from the 2-hidden-layer sweep.
The best test accuracies per depth are shown in Fig.
6, while the overall best test accuracies over all depth are shown in Table 1.
B.3. NNGP AND NTK
For NNGP and NTK, we perform kernel regression following Lee et al.
(2018) using centered labels.
For each depth ∈ { 1 , 2 , 3 , 4 } , we sweep over hyperparameters listed in Table 3 (which change the kernels) using (complete) grid search, with weight initialization variance ﬁxed at 1 in the NTK
[Table] [Table]
parametrization.
For each of NNGP and NTK, the best test accuracy over all depths is listed in
Table 1.
B. I MAGE N ET T RANSFER
We pretrain the π -limit with r = 400 and 2 hidden layers for 30 epochs on a ﬁxed subset of
ImageNet32 with 250 (out of 1000) randomly subsampled classes.
To evaluate on CIFAR10, we compute the kernel induced by the pretrained ﬁnal-layer-features and perform kernel ridge regression.
We optimize the hyperparameters in Table 4 via random search.
B. O MNIGLOT
We focus on the 5-way, 1-shot task, with only 1 step of ANIL adaption.
B.5. µ -
N ET , π -N ET , AND π -L IMIT
For π - and µ -networks, inﬁnite or ﬁnite, we train for 50 epochs, 1000 batches per epoch, and 8 tasks per batch.
In each epoch, we validate on 500 batches from the validation set.
Following Antoniou et al.
(2019), we use cosine annealing learning rate schedule.
We also clip gradients with a hyperparameter threshold, where the clip is triggered by the total parameter norm.
Hyperparameter Optimization We ﬁrst perform random search on hyperparameters listed in
Table 5 for 2-hidden-layer π -net, π -limit, and µ -net: We sample at least 512 hyperparameter combinations 1) for each ( width, r ) ∈ { 128 , 256 , 512 , 1024 , 2048 } × { 50 , 100 , 200 , 400 } for π -net, 2) for each width ∈ { 128 , 256 , 512 , 1024 , 2048 } for µ -net, and 3) for each r ∈ { 50 , 100 , 200 , 400 } for π -limit.
Then we take the best hyperparameter combinations overall (based on validation accuracy) for π -net, π -limit, and µ -net, and for depth ∈ { 1 , 2 , 3 , 4 } (where depth denotes number of hidden layers), we resweep only the meta learning rate (i.e. outer learning rate) and step size (i.e. inner learning rate) in a grid search, where the grid is ( ilr ∗ · { 2 − 2 , 2 − 1 . 5 , . . .
, 2 1 . 5 } ) × ( olr ∗ · { 2 − 2 , 2 − 1 .
, 2 1 . 5 } ) and ilr ∗ and olr ∗ are the optimal inner and outer learning rates from the 2-hidden-layer sweep.
The best validation accuracies per depth are shown in Fig.
Then we take the models with overall best validation accuracies over all depth and evaluate them on the test set using 10000 batches.
These test results are shown in Table 1.
To investigate the width non-monotonicity more thoroughly, we further perform random search on hyperparameters listed in Table 5 for 2-hidden-layer π -net and µ -net.
We sample at least 512 hyperparameter combinations 1) for each ( width, r ) ∈ { 128 , 256 , 512 , . . . , 8192 } × { 50 , 100 , 200 , . . . , 32000 }
[Table] [Table]
for π -net
[Table]
and 2) for each width ∈ { 128 , 256 , 512 , . .
. , 8192 } for µ -net.
The best validation accuracies per ( width, r ) are shown in Fig.
4(Right).
B.5. NNGP AND NTK
As discussed in Section 4.3, for NTK and NNGP, ANIL meta-training has no effect
[Table]
and meta-testing amounts to just taking 1 kernel gradient descent step on each task.
We ﬁx the step size (i
[Table]
e. inner learning rate) at 0.5.
For each depth ∈ { 1 , 2 , 3 , 4 } , we sweep over hyperparameters listed in Tableusing (complete) grid search, with weight initialization variance ﬁxed at 1 in the NTK parametrization.
For each of NNGP and NTK, the best test accuracy over all depths is listed in Table 1.
B.5. V ISUALIZATION OF I MAGE R EPRESENTATIONS
We sample 5 random classes and 10 random images from each class from the Omniglot test set, for a total of 50 images.
We take the best performing NTK, µ -net, and π -limit and evaluate their feature kernels (c.f. Appendix B.2) on the 50 images.
We then do PCA on these kernels to produce Fig.
In
Fig.
8, we also do the same for our best performing π -net.
C D ETAILED C ALCULATIONS OF 1-S TEP SGD
Suppose we present an input ξ 0 to the network and perform a step of gradient descent with learning rate η and loss L .
Then simple calculations show that the updates ∆ u, ∆ v to u, v are
[Equation]
(15) where c = − η L (cid:48) .
Then, via Tensor Programs, f ( ξ ) for any ξ now has a limit of the form
[Equation]
where ˚ c is the deterministic limit of c that is shown to exist by Tensor Programs.
D R EMARKS ON THE π -L IMIT T HEOREMS
π -projection actually ensures that, even for ﬁnite n , f t = f P for some P .
However, this P is random, with ﬂuctuation coming from the sampling of Ω (which is ﬁxed at initialization).
Taking n → ∞ reduces this ﬂuctuation
to 0. 2.
Again, even with the projection, we are optimizing in an inﬁnite-dimensional space, though “linearly inﬁnite” r × ∞ instead of “quadratically inﬁnite” ∞ × ∞ like in the µ -limit.
The forward pass of ˚ f P can be thought of as that of another MLP with nonlinearity V φ , as illustrated in Fig.
In this view, M becomes the width of this MLP.
[Table]
vs Table
[Table]
[Figure]
E N UMERICAL V ERIFICATION OF THE π -L IMIT T HEOREMS
Here we numerically show that wide π -nets have nearly identical loss curves as the π -limit.
See
Figs. 9 and 10.
In addition, Fig.
11 veriﬁes the convergence of the feature kernel to its inﬁnite-width limit.
F P ROOFS
We will just prove Theorems 3.4 and 3.5, as Theorem 3.2 is a special case of them.
At a high level, we need to do two things: 1) Show that f T converges almost surely to something , and 2) this something is ˚ f P T .
Here we assume the reader is familiar with Tensor Programs, in particular the techniques used in Yang and Hu (2020).
F. A LMOST S URE C ONVERGENCE
Showing almost sure convergence is straightforward using the Tensor Programs technique, i.e. express π -initialization and the entire π -SGD training trajectory inside a Tensor Program (a N ETSOR (cid:62) + program in particular) and apply the Master Theorem (c.f. Yang (2020b) in general and Yang and Hu (2020, Sec H.3, H.4) in particular).
Below we sketch the program construction.
Initial Vectors and Matrices Unlike the program for the µ -limit, our program does not have initial matrix variables because in π -parametrization we do not initialize any n × n matrices as iid Gaussians.
This means that we do not use MatMul instructions in our program.
The initial vector variables in our program are just the r columns Ω :1 , . . .
, Ω : r of Ω .
Network Preactivations All preactivations of the network will be of the form Ω C for some C ∈ R r whose entries are scalar variables in the program, so that Ω C can be expressed as a vector variable using Nonlin + . [Figure] [Figure]
Weight Matrices We will sketch the constructions surrounding hidden weight matrices; the input and output weight matrices are similar and easier.
Like in Eq.
[Figure] each hidden weight matrix in deep π -nets can be written mathematically as a sum of vector outer products.
[Equation]
time
[Equation]
[Equation]
where A ls , B ls ∈ R r are the s th rows of A l and B l .
Here φ (Ω B ls ) is the activation going into w l at time s , and Ω A ls is the projected gradient Π Ω ∇ h l L at time s .
Note the sum here is from 1 to M + t , where M comes from the initialization and t is from t steps of training.
All entries of A l and B l will be constructed as scalar variables in the program, so Ω A ls and φ (Ω B ls ) are both vector variables.
In the program, we do not express w l directly, but rather through its matrix-vector product with vectors such as the incoming activation x l − 1 , which would be expressed via a combination of Moment and
Nonlin + instructions like so:
[Equation]
where θ s =
[Equation]
(cid:104)
[Equation] (Ω [Equation] ls
[Equation]
, x [Equation] (cid:105) [Equation]
R
Moment [Equation]
We express ( w l ) (cid:62) indirectly through its matrix-vector products likewise, just with the roles of (Ω A ls ) , φ (Ω B ls ) reversed.
Gradient Projection
In the program, we do not express Π Ω directly, but rather indirectly by expressing the matrix-vector product Π Ω v for vector variables v , such as v = ∇ h l L .
The projection matrix Π Ω is mathematically equal to Ω(Ω (cid:62) Ω) + Ω (cid:62) (for any width), where () + denotes pseudoinverse.
In the program, we would ﬁrst express (the entries of) 1 n Ω (cid:62) Ω ∈ R r × r using many Moment NTK -Net -Net -Limit [Table]
instructions
[Equation]
Moment
Then its pseudo-inverse can be expressed as another Moment instruction (with purely scalar arguments).
[Equation]
Moment
[Equation]
Moment where f ij takes a matrix to the ij th entry of its pseudoinverse.
Note the above expressions depend on Ω only [Table] and not on v .
Finally, we express Π Ω v = Ω γ ∈ R n , γ
[Equation]
n Ω
[Equation]
Ω) + θ ∈ R r , θ = 1 n Ω (cid:62) v ∈ R r in the program like so
[Equation]
Nonlin
[Equation]
where γ i = r
[Equation]
= (
[Equation]
+ ij θ j ∈ R
Moment (19) where θ j = 1 n (cid:104) Ω : j , v
(cid:105) ∈
R
Moment (20)
Wrapping Up Other than what is discussed above, the unrolling of π -SGD follows identical to the unrolling of SGD in Yang and Hu (2020, Sec H.3, H.4).
In particular, f t ( ξ ) for any input ξ and time t is a scalar variable in the program.
Getting Almost Sure
Convergence We apply the Master Theorem (Yang and Hu (2020, Thm 7.4) or Yang (2020b, Thm E.15)) to the program to get almost sure convergence to some limit .
We just need to check the conditions of the theorem.
They are all straightforward except that we need to check the pseudoinverse operation we took is almost surely continuous (in order to satisfy Yang and
Hu (2020, Assm F.4(1))).
However, the only pseudoinverse we took was ( 1 n Ω (cid:62) Ω) + , and Ω has [Table] [Figure] rank r (full rank) almost surely for any n > r .
Therefore, our pseudoinverse operation is almost surely continuous as pseudoinverse is continuous on matrices of constant rank.
F. F ORM OF THE L IMIT
Forward Pass (Theorem 3.4)
In the large- n limit, by the Master Theorem, for hidden weights,
Eq.
(17) becomes
[Equation]
where θ s =
E
[Equation]
Inductively, if g l ∈ R r represents the coefﬁcients of Z h l =
Z w l x l − 1 in terms of Z Ω def =
( Z Ω :1 , . . .
, Z Ω : r ) (which is distributed as a standard isotropic Gaussian vector), then
[Equation]
, which can be rearranged straightforwardly into the equation of Theorem 3.4.
The equations for input and output layers can be derived similarly.
[Figure]
[Figure] t r a i n i n g l o ss width 851232768inf  8 10 12 14 log 2 ( width ) l o g 2 ( m e d l o ss ) legendmed [Figure]
C /
[Figure]
[Figure]

[Figure]
  
[Figure]
Backward Pass (Theorem 3.5) For hidden weight w l , since we maintain w l in the form of Eq.
(16), the gradient update
[Equation]
in the limit corresponds to
[Equation]
where coef ( Z v ) is the coefﬁcient of Z v in terms of Z Ω
= ( Z Ω :1 , . . .
, Z Ω : r ) .
(Note the factor of n in
Z − nη Π Ω ∇ hl L ) is just there so that nη Π Ω ∇ h l L has Θ(1) -sized coordinates).
Of course, coef ( Z h l − 1 ) is just g l − 1 ∈ R r by deﬁnition.
It remains to show that coef ( Z − nη Π Ω ∇ hl L )
=
∇ g l lim L , where lim L = lim n →∞ L ( f ( ξ ) , y ) is the loss at the limit (which is deterministic).
Using the Master Theorem, it is not hard to see that, for each preactivation h l , Z n ∇
hl L is the Frechet derivative ∂ lim L ∂Z hl with respect to the random variable Z h l , where the Frechet derivative is deﬁned with respect to the Hilbert space H of square integrable random variables in the σ -algebra generated by Z Ω def =
( Z Ω :1 , . . .
, Z Ω : r ) .
Furthermore, using the limits of Eqs.
(18) to (20) obtained from the
Master Theorem, it is easy to see that
[Equation]
where Π Z Ω is the orthogonal projection to the linear subspace of H spanned by the random variables
Z Ω :1 , . . .
, Z Ω : r .
Then Lemma F.1 applies and we get that
[Equation]
where the derivative is now an ordinary partial derivative.
Lemma F.1.
Let H be a Hilbert space and V be a k -dimensional subspace of V , where k is ﬁnite.
Let L : H → R be a Frechet differentiable function.
Suppose Γ have for columns an orthonormal set of basis for V , and w = Γ b ∈ V for some b ∈ R k .
Then
[Equation]
where Π V : H → V is the orthogonal projection to V .
Proof.
By the theory of proximal gradients,
[Equation]
Changing coordinates via Γ , we have
[Equation]