Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes an infinite-width parameterization ($\pi$-parameterization) that is similar to $\mu$-parameterization of Yang and Hu, with two important tweaks: the initialization involves a random projection of a certain form and the gradient update invokes the same projection. This results in an infinite-with description that is more amenable to computation and analysis. The paper demonstrates experimentally that this method achieves feature learning and better performance than NTK.
main_review: Thanks a lot for the replies to many questions I asked!

Through the rebuttal, some central claims and contributions are clearer to me. But truthfully it turns out the paper now appears less promising and less exciting than I originally thought. Two major points that I learned from the rebuttal:

- The $\pi$-limit appears to give a qualitative view of the $\pi$-net, but not the more precise details of the $\pi$-net unless the width is 10 or 100 times larger than $r$. For example, the second figure in figure 12 shows that at width 2048 or even 32786, for $r=400$, the training accuracy could deviate from the $\pi$-limit as much as 10% at various points of time. Figures 9 and 10 show better matches, but it is for $r=2$ which is atypical and very low.

- The $\pi$-limit can take a very long time to compute; its runtime is multiple times longer than training a $\pi$-net (Table 8). Rather than “efficient computation”, what is demonstrated here is basically computability, against the fact the $\mu$-limit is practically uncomputable.

Without further demonstration of what use the $\pi$-limit can offer, this is quite hard to appreciate the $\pi$-limit, against the empirical performances which not all reviewers here are excited about. Note, throughout the infinite-width literature, be it NTK or mean field, it is quite standard that the identification of the infinite-width limit comes with an additional analysis or a demonstration of use, such as proving theoretical global convergences.

I’m also empathetic with reviewer MBzR’s concern, which is that the $\pi$-parameterization simply goes further away from the usual initialization practice. The fact that the $\mu$-parameterization bears similarity to the usual initialization practice (even though strictly speaking, it actually differs from the practice in a very important way) is probably one major reason why people pay much attention to it. By differing even more, the $\pi$-limit disregards entirely this premise. As such, I would expect to see from the $\pi$-limit something more interesting than a “fix” to the $\mu$-limit.

In various ways, I currently find less reasons to be convinced it’s a promising direction.

I will give the benefit of the doubt and raise my score, not because I want to align with other reviewers, but mainly because I simply like the idea. But I wish the paper were more carefully fleshed out than just this.

=================================================

The proposal is interesting and novel to my knowledge. The fact that it allows for computation of the infinite-width limit, unlike $\mu$-parameterization, while still achieving some nontrivial performance, is valuable in my opinion. I find that very encouraging, despite that the performance is of course nowhere close to strong modern neural nets yet. I also like that the paper provides a short description of the difficulty with $\mu$-limit and takes a detour to the 1-hidden-layer example.

I however find that the paper incomplete in several crucial aspects. The following should be considered in the finalized version:

- The paper does not demonstrate convincingly if a finite-width net under $\pi$-parameterization approaches the $\pi$-limit as widths go to infinity; that is, a verification of Theorem 3.5. Now that the limit can be computed (as compared to $\mu$-limit), this should have been done. In particular, it should demonstrate the closing gap to the $\pi$-limit at every training steps, as indicated by Theorem 3.5. How large should the width be to be visually close to the limit? Is it close in just the classification accuracy metric? These questions should have been the first items to be answered when an infinite-width limit is claimed to appear more practical. 

    Figures 4 and 5 are insufficient in this aspect. Figure 5 looks concerning: as it shows, a width of 40,000 (!) still has a big gap to the limit on CIFAR10. I think this warrants deeper investigation.

- The hyperparameter $r$ appears very important, but there is little discussion. From the experiments, $r$ should be large. But the need for large $r$ complicates the studies.

    - Firstly this questions the relevance of the $\pi$-limit itself, if one takes it as a theoretical limit of $\pi$-net. In particular, $\pi$-limit is derived with $\text{width}\to\infty$ while keeping $r$ constant. Once $r$ is a nontrivial fraction of the width, one should expect some nontrivial random matrix effect to kick in and the $\pi$-nets to have nontrivial deviation from the $\pi$-limit, casting doubt on the validity of $\pi$-limit. This is indeed the case in the experiments.

    - Secondly if the proposal is to consider $\pi$-limit as a new neural architecture with a specific training algorithm rather than a theoretical infinite-width limit of $\pi$-nets, then the comparison should have been done differently. That is, it should not be compared against the simple NTK of feedforward nets, its expressiveness property should be investigated, its runtime should be studied, etc.

- The claim (in the paper title) is that $\pi$-limit can be efficiently computed, but again there is little discussion. In particular, can the paper provide a snapshot of the runtime for a reasonable, easily reproducible setup? There is a very short remark in page 4 that the complexity scales as $O(T^2)$, but this is insufficient to claim efficient computation. Is it running faster than or comparably to SGD on $\pi$-nets, in terms of compute hours? Is it better / faster to just run SGD on $\pi$-nets, and for what range of $r$?

Other issues:

- The paper claims outperformance of $\pi$-limit and $\pi$-nets over others in Table 1, but some of these numbers look too close to make a call. An increase of 2% for CIFAR10 could be borderline (given the nontrivial randomness in the generation of the matrices $A_l$ and $B_l$).

- Footnote 12 is quite concerning: it is well-known that standard-parameterized MLPs should behave differently from $\pi$-limit (and are close to NTK at best) as widths tend to infinity, but this footnote seems to suggest that the experimented $\pi$-nets may operate in a regime far from the $\pi$-limit?!

- Footnote 13 seems unfair. Firstly the design of the learning rate schedule should be considered as a knob in hyperparameter tuning. Secondly the outcome of convex optimization would differ if one switches from a constant learning rate schedule to one with learning rate drops.

- The hyperparameter tuning is rather intense for theoretically inclined audiences. It would be of great service if the paper can report the experimental results of a simple (not the best) hyperparameter setup for easy repetition/reproduction, in conjunction with these more tuned results. The paper should also specify the data processing step.

Questions:

- Does the $\pi$-parameterization inherit most properties of the $\mu$-parameterization, other than being different in its limit’s analytical form and training procedure? 

    For example, the $\mu$-parameterization has a somewhat paradoxical property (despite its name "maximal") that $\Delta w(t) / w(0) \to 0$ as width tends to infinity in most layers, for $\Delta w(t)$ the gradient update at time $t$ and $w(0)$ the initial weight. That is, each weight coordinate moves very little from its initialization. Does the $\pi$-parameterization share the same property?

- How would one initialize the weights in $\pi$-nets if the widths are not all equal? Here the same projection $\Omega$ is used, restricting the widths to be all equal.

- Footnote 10 seems to say that doing gradient accumulation on $\pi$-nets could be difficult due to smoothness of $\mathcal{V}_\phi$. Is there such difficulty with the computation of the $\pi$-limit?

summary_of_the_review: The proposal in the paper is interesting and novel, but the paper has deficiencies in its research methodology. It would be in a much better shape if several more fundamental issues are addressed.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: In this work, the authors propose a new feature learning limit for infinite-width neural networks that resolves some of the computational issues with the previous $\mu$-limit.
This new limit, called the $\pi$-limit, is based on projected gradient descent, and crucially allows the authors to evaluate the model on deep networks, which could not be done before.
The authors show that it outperforms both finite-width and previous infinite-width networks on CIFAR10 and Omniglot.
main_review: Overall, this paper reads well and is clear despite the dense mathematical content, and solves a relevant problem.
The solution is technically sound, and works well empirically.
I would be very happy to see this paper accepted.

I still have some minor comments and concerns, however, that I hope the authors would consider addressing:

- The proposed model ends up being quite different to a standard MLP, both in the way the weights are computed and the way training is performed.
This is unlike the previously proposed $\mu$-net, which is much closer to the way standard MLPs are trained.
Clearly, the choices made are necessary for the computational speedups that this paper is about.
However, I was hoping that the authors could comment more on the intuitive differences between these models?
The authors briefly discuss a similarity to the hypernetworks from Ha et al. (2016), but it would be great if this discussion could be expanded, particularly with respect to $r$.
For instance, would the authors expect the inductive biases to be much different between the two models?
This is likely a difficult question to answer, so I would simply ask the authors for some additional discussion along these lines.
However, this question is quite relevant for determining what we can learn from these infinite width limits, which is one of the main motivations of studying infinite width limits.

- I think that the results from Table 7 in the Appendix should be promoted to the main text.
While not explicitly stated, I believe I am correct that the training loss in Table 1 is the standard softmax for the feature learning methods?
In this case, to ensure a fair comparison between methods, it is important that the same loss be used, as in my experience FKR is very different to training with a proper classification loss.
Along these lines, do the authors have any intuition as to why the $\mu$-net suffers much more than the $\pi$-net when used for FKR?

- Would the authors be able to provide a comparison to $\mu$-net and $\pi$-net under transfer learning?
It would be understandable if not due to the computational requirements of doing so, but I think it would strengthen the paper.
summary_of_the_review: In summary, I think this is a very good paper that I would personally like to see at the conference.
I have some minor reservations which I hope the authors might address, but unless the other reviewers find something critical that I have missed I am unlikely to change my score.
Nevertheless, I look forward to reading the other reviewers' thoughts and the authors' responses.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The authors study a certain variant of an MLP trained using a projected gradient descent-inspired update rule in the infinite width limit. This infinite width model admits closed-form rules for computing the predictor at any training iteration (as long as the V-transform or neural network kernel being known, which is the case for e.g. the arc-cosine kernel). The authors evaluate their model on Omiglot and CIFAR10 against NTK and finite width baselines.
main_review: Strengths:
Infinite-width models in their various flavours (NNGP, NTK, ...) form a useful set of tools for theoretically understanding the behaviour of neural networks in terms of classical kernel methods or GPs. Most (all?) of these types of limits usually mean that the limiting model keeps some notion of feature space constant during training. Motivated by the fact that practical neural networks actually update the feature space during training, the authors introduce a new limiting model (the pi-limit) that allows the features of the kernel to be updated during training. This is an important direction for the progression of infinite-width models.

Main criticism:
As pointed out by the authors, there is a gap between these models and modern neural networks in terms of the architectures, training schemes, and training regime (i.e. is the network "wide enough" and trained with a "small enough" learning rate?). This means that when performing this style of analysis, one usually has to make some modification to what is typically done in practice. One runs the risk of making a modification so severe that the analysed model does not reflect a neural network that people usually compute with.

In this work, in order for the infinite width model to be tractable, several modifications are required for 1 layer case:
-Initialise readout parameters at 0
-Projected gradients (but only for hidden layers)
-Minibatch size 1 (?)
-Gradients are appended to the network instead of added (!)

In the multi-layer case, an additional modification is required:
-pi-Initialisation. This initialisation appears to be substantially different to the typical scaling by sqrt n that is typical in infinite width works and in vanilla finite-width initialisation schemes (e.g. Glorot init).

These modifications, especially the fact that gradients are appended instead of added/subtracted in the usual way gradient updates are applied, mean that the model studied is quite different to one that would typically be used. This makes me wonder whether the studied model bears any resemblance to "real" neural networks. The model also does not achieve state of the art (it is not capable of expressing convolutional or transformer architectures). Since the model is not reflective of any simple practical neural network, nor does it achieve good performance, I am questioning the interest of this model in the broader ML community.

Questions for the authors:
1. Can you say a bit more about this? "For simplicity, we only consider batch size 1; it’s straightforward to generalize to larger batch sizes". There is a comment in appendix A (A.4), but I do not see any statement or how to modify the result toa accommodate larger batch sizes.

2. In equation (6), we use the V-transform, which is essentially the neural network kernel. The neural network kernel is symmetric in its second and third argument (i.e. swap X and Y and you get the same thing). In equation (6), the second and third argument don't look similar at all. Can you explain this asymmetry?

3. There are number of fiddly modifications performed in the experimental evaluation. For example, the learning rate drop mentioned in footnote 13, gradient clipping, ANIL, cosine annealing, ... . This seems slightly counter to the overall objective of the paper, which was (I think) primarily to study an infinite width network that is able to perform feature learning for its theoretical niceness. When moving to the empirical evaluation, I am not convinced that the same amount of hyperparameter and model tuning was afforded to the NTK/NNGP/finite width models as was given to the pi-limit. I think that the fairest comparison would be to not include any "tricks" in any models.
summary_of_the_review: After the reviewer response, I have upgraded my score to 6. My main concern remains that the model analysed is different to what is typically used in practice.

The paper appears to be technically sound, novel and written with the intention of solving an important problem. But it studies a model so far removed from a typical neural network (so it's theory is not widely applicable), and the model does not achieve state of the art performance (so it is empirically not widely applicable).
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper introduces an approach (named pi-limit) to compute explicitly the infinite-width limit performance of a multilayer perceptron (MLP). The approach is based on projecting the gradient of the loss at each epoch in a fixed direction, chosen upon initialization. If one does so,  the expectation value of the function learned by the network at infinite training time can be computed explicitly by the iterative procedure described in Theorem 3.5. The asymptotic accuracy on a network trained on CIFAR10 is 61.5 %, slightly (0.2 %) better than the accuracy obtained by the nu-net procedure, which does not allow estimating explicitly the limit, and almost 2 % better than the the accuracy of a Neural Tangent Kernel, which, as well known, is feature agnostic. Based on these tests, it is concluded that the pi-limit is the first procedure which allows computing exact expectation values in a MLP and that, at the same time, allows learning features.
main_review: The formal development presented in the paper is interesting and insightful. The idea is first presented in a simple example, in which the input and the output are a single real number, and this allows digesting the pretty complicated math. The idea itself, restricting the parameter dynamics in a few directions, is not new, as also stated by the authors  (it is also exploited in  Direct Feedback Alignment https://arxiv.org/abs/1609.01596, which in my opinion should be cited), but the explicit expression of the limit in the two theorems is, to the best of my knowledge, novel. 

My main concern on the manuscript is on the meaningfulness of the numerical tests. The performance gap with respect to NTK on the CIFAR10 network is only 2 %, making the improvement brought by feature learning and by taking the explicit infinite-width limit difficult to appreciate. I suggest to repeat the test on a fully-connected architecture designed for image recognition, which achieves a test accuracy of 78 %, much closer to the state of the art on this data set: https://arxiv.org/abs/1511.02580

In this more realistic condition it should be possible appreciating better the improvement with respect to NTK and the nu-net.


summary_of_the_review: The theoretical developments presented in the paper are (to my knowledge) novel, and can trigger other ideas for further improvements. The numerical tests presented to support the usefulness of the theoretical framework are not fully convincing, but can be rather easily improved. My score would be a 7, if 7 was available.
review_title: null

Meta-review: This paper studies deep non-linear infinite-width neural networks that go beyond the NTK and learn features. This paper extends the prior result on shallow neural networks to deep neural networks and empirically evaluates the deep inf-wide nn. The reviewers find the contributions in the paper valuable. The meta reviewer agrees and thus recommends acceptance.
Decision: Accept (Poster)