Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This submission analyzes the dynamics of adaptive optimization of arbitrary neural network architectures in the infinite-width limit, and as a by-product,  a new tensor program framework is proposed.
strength_and_weaknesses: #### Strength
- Close the gap between adaptive optimizer and NTK analysis.
- Extend the guarantee in NTK and feature learning to the non-linear cases in the sense of the Tensor program language.


#### Weaknesses
- The adaptive optimizer considered in the analysis is memoryless and one-dimensional (Thm.4.1-4.2) given batch size is one, which is far from the practical case.
- The theoretical results seem to be a generalized version of SGD, and we cannot see the difference between AGO and SGD from the theory. In fact, AGO and SGD behaviors differently during training, and we cannot treat SGD as a particular case.
clarity,_quality,_novelty_and_reproducibility: The submission extended the Tensor Programs framework to allow the expression of the computation graph involving adaptive optimizers. From this perspective, the novelty is limited.
summary_of_the_review: The submission adds some new tools in the previously developed framework (Tensor programs) for helping analyze the dynamic of $\infty$-width neural network trained by adaptive optimizers. The only disadvantage is that the theory does not give more insight into the difference between SGD and adaptive optimizers (e.g., Adam).

Overall, I still hold a positive score for this work since it is a necessary expansion of Tensor programs.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper studies two infinite-width limits (kernel and feature learning limits) of fully-connected neural networks (MLPs) under adaptive gradient-based optimization. Both results generalize that of MLPs trained by a non-adaptive way. In addition, the authors modify a framework (called Tensor Program) that allows to express the adaptive gradient processing as well as the convergence guarantee.

strength_and_weaknesses: Strength:

- This paper extends two infinite-width limits of MLPs for adaptive gradient-based optimization, which have not been studied before. The proposed results can generalize both kernel and feature learning limits of non-adaptive settings. It is interesting that the adaptive NTK regime is still data-independent like the non-adaptive NTK.  

- They provide a general Tensor Program framework for the adaptive optimization setting with O(1/sqrt(n)) convergence rate guarantees. 

Weakness:

- Adaptive optimizations have various benefits over the SGD (e.g., fast convergence, better generalization, etc). Can this point of view be captured by the proposed adaptive infinite-width limits? It would be great if more clear motivations and meaningful results of adaptive infinite-width limits are provided. 

- The paper focuses on simple settings (e.g., memoryless adaptive gradient descents, update of parameters in a single layer). But, it is unclear whether the proposed results hold for more general settings. 

clarity,_quality,_novelty_and_reproducibility: The writing quality of this paper is quite poor and very informal. It should be improved by providing more comprehensive and rigorous flows. 

summary_of_the_review: This paper extends infinite-width limits to the adaptive optimization setting, which has never been discovered before. But, it seems that a more in-depth comparison/analysis of the non-adaptive results would make the contributions much stronger. The writing quality needs to be improved for better readability.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: In this work the authors tackle the question of the infinite-width limit of neural networks trained using adaptive optimizers (like ADAM). They show that if the step taken in the optimizer depends on a non-linear, scale-invariant function of the previous step, then the NTK can be computed by applying said function to $\frac{\partial f_{train}}{\partial\theta}\nabla_{\theta}\mathcal{L}_{train} $, and then multiplying by $\frac{\partial f_{out}}{\partial\theta}$. An analogous result holds for the $\mu$-parameterization (the unique family of infinite-width feature learning initializations).
strength_and_weaknesses: The arguments for the memoryless case are very clean. In addition, these additions to the TP framework make it more flexible and can give insight in many of the more complex practical cases which arise when training large models. The authors make a good mix of intuition-building arguments and links to the more rigorous ones.

One point of clarification: in the mean-field limit, does $\epsilon$ have to go to $0$ to get all the scaling arguments right? Is there an intuitive reason that this doesn't mess up the Gaussianity arguments via divisors which are random variables which can have significant probability mass around $0$? The final limits seem to be $\epsilon$-independent.

One of the big weaknesses is that the promise of extending the analysis to the case of an adaptive optimizer with memory (depending on gradients $g$ further back in time) is not clearly made. I believe without that, the result in the memoryless case is interesting but not sufficiently significant for an acceptance rating. It would also be good to have a discussion on how the total time $t$ must scale with $n$ for the arguments to hold.

I think the paper would be significantly strengthened if there was an explicit formulation and proof of at least the NTK for the ADAM algorithm. From the comments in the paper these seem like results that are already completed, or easily derivable for the framework. This paper is positioned to become a standard reference for NTK/$\mu$-parameterization for adaptive algorithms, and it would be strange for the most popular adaptive algorithm to not be explicitly covered.

Due to the limits of my technical ability and the review timescale, I have not been able to rigorously check the main proofs in the TP framework. I followed along with the arguments in the main text and they seem sound. I believe that there needs to be numerical evidence that the proposed scaling limits hold - not to replace the proofs, but to supplement them and decrease the possibility that some key point in the complex machinery did not go awry.

As a small note, on page 3 it is unclear why the labeling of the gradients $g$ has even time indices only.

Update: most of the concerns above have been addressed; my updated review score reflects those changes.
clarity,_quality,_novelty_and_reproducibility: Sections 1-4 are very clear and easy to read. Sections 5 and 6 are more difficult; however, I believe that this is due to the technical nature of the content, and not any fault of the authors who make an effort to build intuition for the proofs. The arguments are, to my knowledge, novel.
summary_of_the_review: The paper does an excellent job of building intuitive and rigorous arguments for memoryless adaptive optimizers. However, I feel that this is not a significant enough contribution on its own, and the case of non-memoryless optimizers like ADAM should be explicitly covered if they are in fact trivial extensions of the framework as suggested by the authors. I did not have the time nor expertise to rigorously check the proofs. I also suggest that numerical evidence of convergence to the theoretical limits would be very helpful to readers, especially those without strong theoretical backgrounds.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper studies the infinite width limit of Multi-layer perceptrons, in particular the existence of feature learning. 
In contrast to previous work, it studies adaptive optimizers instead of Gradient Descent, but results are roughly the same: NTK limit still results in lazy training (and no feature learning), while mu-parameterization supports feature learning.
Unfortunately, no experiments on real (or even synthetic) data are provided in support of theoretical claims (It has been addressed in the rebuttal).

strength_and_weaknesses: Strengths:

- Infinite width limit of MLP studied for the first time with adaptive optimizers.


Weaknesses:

- No empirical data is provided in support of theoretical findings.

- NTK limit is provided (theorem 4.1) in memoryless case only, which is of very limited interest.
The authors state: "we maintain that our analysis is applicable in the more general setting without introducing any additional arguments", but no proof is provided as far as I understand.

(It has been addressed in the rebuttal)

clarity,_quality,_novelty_and_reproducibility: The paper is clear and contributions are novel.

summary_of_the_review: I would be much more convinced of the significance of this work if empirical validation of the theoretical results was provided (It has been addressed in the rebuttal).
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The paper investigates adaptive optimization in the infinite width limit of neural networks. 

* Strengths are a novel and relevant extension of tensor program theory along with analysis of adaptive optimizers advancing the connection of theory and practical settings. 
* Weaknesses are certain limitations of the method to discern between training methods. 

Reviewers mostly agree on the novelty and technical significance. In the rebuttal the authors addressed some of the reviewers' initial comments. Particularly, generalizations of the results to optimizers with memory and numerical evaluations. The final recommendations were mostly on the favorable side. I find that the topic and methods are of interest and conclude that the paper should be accepted. The authors should still carefully check updates conducted during the rebuttal period, particularly on memory and dimension.  
 
justification_for_why_not_higher_score: Although the overall assessment is favorable, there were no advocates for a higher score. 
justification_for_why_not_lower_score: The reviewers agree that the paper should be accepted.
Decision: Accept: poster