Modify Training Direction in Function Space to Reduce Generalization Error

• The Gaussian likelihood assumption 1 seems to be restrictive in the sense that it only holds for Gaussian data, is it possible to relax this assumption?
• The NTK regime assumption 2 also restricts the model to be a random feature model, which can simplify the analysis but overparameterization requirement is also restrictive.
• The optimization result (Theorem 1) is based on the convergence result in [32], which limits the novelty.
• It would be more convincing if more numerical experiments are conducted to show the effectiveness of the modified NGD method. For example, the image classification task on the CIFAR10/CIFAR100 dataset.

Some minor issues:

• the last step in equation (26) should be $1 > \frac{1}{2}$.
• the left hand side of equation (41) should be $f_{\theta_{\infty}}(x)$.

[32] Tim GJ Rudner, Florian Wenzel, Yee Whye Teh, and Yarin Gal. The natural neural tangent kernel: Neural network training dynamics under natural gradient descent. 4th workshop on Bayesian Deep Learning (NeurIPS 2019), 2019.

Overall, the paper feels rushed, with several typos in different places. This is not critical, but there are more conceptual concerns.

The role of NGD. When focusing on the main goal of the paper, as given by the title, the authors analyze the generalization performance of the final predictor $f_{\theta_\infty}$ given by (21) with modification $\varphi$ given within the class described by eq. (7) and (12). However, such a predictor simply corresponds to a complete fit of training residual $\alpha(\mathcal{X},\mathcal{Y})$ projected on a set of eigenspaces selected according to criterion $c$. From this perspective, Modified NGD seems to be only one of many possibilities (and a fairly complicated one) for obtaining such a predictor. For example, first projecting the residual on the selected eigenspaces and then training basic Gradient descent till convergence would produce the same predictor. This raises a question of the actual role of NGD in the analysis and results provided in the paper.

Implication of final training dataset size. The authors treat quite lightly the difference between quantities related to finite training size and those related to (presumably) infinite population distribution. For example, the equality (24) seems to be impossible to hold. As can be seen from its other form in eq. (50), the left-hand side can have only $N$ non-zero entries, while all entries of the right-hand side are presumably non-zero. This is just one implication of the broader issue of learning population quantities from finite number of observations, which is one of the main technical difficulties in references [8,10,17,26] provided in the manuscript and in some other works studying generalization error of kernel methods.

The upper bound. Theorem 2 presents the modification $\varphi$ that improves not the true error but its upper bound (54) derived in the appendix. In this regard, it would be reasonable to discuss this upper bound in the main text so there is more understanding of what is actually minimized. For example, I am not able to tell whether this bound based on the equivalence of norms in finite-dimensional spaces is typical in analysis of generalization error, and the authors do not give any references or comments about it. Also, I did not understand the transition from the first to the second line in eq. (54) - I would expect an additional term with the contribution from components of $\xi^\star$ orthogonal to directions with non-zero singular values of $\mathbf{J}$.

Differences between theoretical and experimental settings. The experimental and theoretical settings are quite different. Most importantly, in theoretical settings, the network is in NTK regime with constant Jacobian, while the experiments are performed with realistic finite networks with Jacobian evolving through training. Because of this, the authors regularly update Fisher during training. This aspect, often termed Feature Learning, is one of the main challenges to NTK theory, and it would be good to address it more clearly in the experiment design to make the relation to the theory more transparent.

Smaller remarks and typos which do not significantly affect the quality of the paper

• $1/N$ factor is missing in eq. (35). This is the reason for the not-very natural $N$ factor in the exponent $e^{-\eta N \sigma_0^2 t}$ in eq. (20).
• It seems strange that predictor $f$ (not the conditional probability $\widetilde{p}(y|x)$) depends on output $y$ in the definition of the Fisher matrix (just below eq. (5)).
• Theorem 1 gives the solution for (continuous) gradient flow algorithm, but formula (8) describes (discrete) gradient descent. It would be better to explicitly mention that the theorem is for gradient flow.

Although I feel the connection is less well-known, overall, it feels like the entire idea is to apply something already well-established in kernel method to neural tangent kernel, which is a special kind of kernel. Another point I want to point out is that the neural network used in practice has been shown in various scenarios that it doesn't behave like a kernel. Thus, it is not clear whether the improved generalization from modified NGD can still be applied to neural networks not in NTK regime. With that, I feel the contribution of this work is somewhat limited. In the experiment section, the authors compare modified NGD to NGD. As a reader, I am also curious about the comparison of modified NGD and NGD to GD.

My main concern is that the results are limited in scope (to NTK, Gaussian data-likelihood) and hence not sufficiently novel and broadly applicable.

1. Under the NTK assumption, the dynamics are linear and hence it is analytically easy to derive the results in Theorem 2. None of the proof outlines are given.

2. Any modification to the gradient does change the functional gradient of the loss function. This is not a unique novelty of this paper. Providing more intuition on the modifications for out-of-distribution generalization can enhance the results.

3. Figure captions could be more informative of what the inferences are.