Continual learning with the neural tangent ensemble

Summary

This is a theoretical paper that seeks to find an algorithm to train neural networks without forgetting in continual learning settings. The authors show that in the infinite width limit, a single classification network can be reformulated as a weighted ensemble of fixed classifiers (fixed experts). They provide a learning rule, which computes the posterior distribution of each expert. This learning rule turns out to be similar to SGD, where there is a per-parameter learning rate and the different of the parameters to their initial state is constrained to be L1 normalized (unsure of the exact correct phrasing of the L1 projection step). In the finite case, the experts can no longer be viewed as static. Experiments show that any amount of momentum hurts continual learning performance, but momentum does help when learning a single task from scratch.

Strengths

Theoretical understanding of how to prevent forgetting in networks as training continues is important and helps move the field forward. The connection of SGD to posterior updates is interesting and (as far as I'm aware) novel. The experiments exploring the effect of momentum on a networks ability to continually learn are interesting and useful from a practitioner's perspective. Experiments demonstrate that the NTE update rule becomes more appropriate as width increases. The work opens up future directions that may be able to estimate uncertainty based on the variance of experts under the NTE formulation.

Weaknesses

The claim that Bayesian ensembles of fixed functions are natural continual learners does not seem obvious to me. In the infinite-width I buy it, but for finite networks, it seems like it would be possible to prune all of the experts away leaving you with a degenerate model. Only one continual learning benchmark dataset (permuted MNIST) is used in the experiments. The details of this task are not well explained. The introduction highlights the value of Bayesian posterior updates as being order-independent. However, there is no experiment exploring the sensitivity of the NTE update rule to the ordering of the data. Section 5.4 seems out of place. It isn't a result, and it seems a stretch to call it a prediction. It may be useful reducing it to a few sentences in the related work section so other sections can include more clarifying text. Overall, the proofs seem correct, but that required a significant amount of time and further reading for me to come to that conclusion. Such theoretical work would be much stronger if it codified its proofs in the form of a proof assistant like Lean4 or Coq. This would require the reviewer to only check that the claim is correctly formulated. The proof itself would be verified automatically, and there would be no question about its correctness. Codified proofs would allow readers learn more about any unfamiliar identities used with full confidence and remove all ambiguity stemming from unfamiliar notation. Typo: Line 205. "This is constrains" should be "This constrains"

Questions

What happens if there is label noise? This analysis seems to assume that the dataset is perfect, which raises the question: if you have an idealized network that cannot forget, what is the impact of a mislabeled example? The paper claims that in the infinite-width network each parameter/weight/edge contributes a classifier to the ensemble. This seems to imply that weights in earlier layers are included here, is that true? I'm confused about how to picture this. In the first layer, there is a weight connecting one of the finite input neurons to one of the infinite hidden neurons. That particular edge isn't able to see much information, but its output does connect to an infinite number of neurons in the second hidden layer. Is its interpretation as a weight on the neural-tangent-experts somehow implicit wrt to the rest of the network depth? Is it that the edge allows energy to accumulate in the "right spots"? Similarly, in the context of finite networks, is there a way to extract the individual experts? Are the number of experts still equal to the number of model parameters in this case? Figure 3: I'm a confused about the continual learning task setup for permuted MNIST. It would be useful to remind readers what the mechanism for predicting on task 1 after 5 tasks is. Is there a new head? Do the number of output units stay fixed throughout all tasks? In theorem 1: why is the perturbation important? Is ΔW the linear update that brings an untrained NTK weights (W⁰) to the trained NKT weights (Wᵗ)? If this is true, it might help to the statement of theorem 1 to call that out. Can anything be said theoretically about how fast increasing the width of a network converges towards the NTE limit? In Figure 4 are the curves over just the first task? In other words, is this just training on regular MNIST to show how the NTE update rule begins to fail as the weights drift from their initial state?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors are clear about the technical limitations of the work, and do not make any unreasonable claims. Code to reproduce looks like it exists.

Summary

This is a theoretical research on preventing forgetting by considering a single network trained with a lazy-regime as an ensemble model of multiple functions and adjusting their weights.

Strengths

The discussion is based on solid theory. This paper obtained an insight that the posterior update rule for the NTE is equivalent to a scaled and projected form of stochastic gradient descent without momentum. Because of this relationship, this paper successfully demonstrates the disadvantages of using momentum and the advantages of increasing width which may induce lazy training.

Weaknesses

I feel that the contribution is valuable, but some settings might be difficult to realize in the context of NTK, even with an infinite limit. I might be able to deepen my understanding through the answers to the questions, so please refer to the Questions section.

Questions

**(1) Non-linear output** Many discussions are based on the formulation of Equation (5), and I assume this setting is intended to be used with non-linear transformation such as softmax. However, in my understanding, the NTK does not remain constant in such cases [1]. What are your thoughts on this perspective? If I have misunderstood, please correct me. **(2) NTK regime** There is a lack of discussion regarding the settings of the model used in this study. In order for the NTK to remain constant during training, parameter scaling or output scaling must be set appropriately. In some settings, lazy training would not be realized, even if the width is increased [2]. In the experiments, does the NTK indeed stop changing as the width increases? **(3) Suboptimality** As mentioned in the Discussion section, experts are not actually independent, so the weighting strategy discussed is not actually optimal. In the Discussion section, various ideas are presented for training experts to maintain independence, but these approaches have not been evaluated in this paper. Therefore, it seems important to show that the impact of suboptimality is not significant. Is that possible? If the impact is too large, I feel that the premise of the discussion might be a bit fragile. [1] Liu et al., On the linearity of large non-linear models: when and why the tangent kernel is constant, NeurIPS 2020. [2] Yang and Hu, Feature Learning in Infinite-Width Neural Networks, ICML2021

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

Discussion is only effective within a lazy regime. Additionally, experts are required to be independent.

Summary

The paper introduces a novel approach to mitigating catastrophic forgetting in neural networks by introducing the concept of Neural Tangent Ensemble (NTE), a formulation interpreting a single neural network as an ensemble of fixed classifiers, leveraging the Neural Tangent Kernel (NTK) framework. This interpretation allows for the derivation of a Bayesian posterior updating rule, which is shown to be equivalent to a scaled and projected form of stochastic gradient descent (SGD). The paper presents theoretical insights into how neural networks can be viewed as Bayesian ensembles, offering an optimization-guided framework for understanding and addressing forgetting in continual learning scenarios.

Strengths

1. Theoretical rigor: The authors provide a detailed theoretical analysis, including proofs and derivations. 2. Connection to SGD: The finding that the NTE posterior update rule is connected to SGD is insightful and could have implications for optimization strategies in neural networks 3. Clear presentation: The paper is clearly presented with fluent logic.

Weaknesses

To me, the largest limitation of this paper is the empirical validation. Although the paper includes empirical results that support the theoretical claims, the scope of experiments could be expanded. The experiment is implemented on a relatively simple dataset (Permunated-MNIST). Additional experiments on more complex datasets would be appreciated. Beyond that, the comparison is relatively insufficient. Could the authors provide comparisons with other optimization-based continual learning methods, like OGD or others, to validate the effectiveness of the proposed method? Or are there any other reasons that they can not be directly compared?

Questions

1. Are there any additional experiments that could be conducted to further validate the effectiveness of the NTE approach in more complex or real-world continual learning tasks? 2. How does the performance of NFE compare with other optimization-based methods? 3. Could the author elaborate more on Equation(3) and line 78, "It is easy to create multitask Bayesian ensembles even when tasks are seen sequentially," which is not very direct to me?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have outlined the limitations and the potential societal impact of the work.

Summary

The paper suggested that linearized networks (under lazy learning regime) at the initialization can be understood as an ensemble model where each ensemble component is a function parameterized by a single parameter in the network (i.e. if the network has N parameters then the prediction is an ensemble over N functions). The weight of each ensemble is controlled by the posterior weights of each ensemble/expert conditioned on the data in a particular task. Under this regime, when a new task comes in, experts good at previous tasks but bad at this task would be down-weight, as such the ensemble would continuously be adapted to new tasks without losing ability in previous tasks. Interestingly, the paper shows that the posterior update formula is almost identical to SGD update (without momentum). Additionally, the authors proposed a more practical version of the framework under the rich regime, where each ensemble/expert is evolving throughout time, under the assumption that the parameter does not move too far away from the initialization. Lastly, the authors draw three useful insights from the ensemble interpretation, and empirically verifies the insights on 2 layer NN on permuted MNIST. Overall I find the connection between linearized network, ensemble and continual learning very interesting and novel. The theory is overall sound to me and the empirical results support the theory well.

Strengths

- The interpretation proposed is novel and interesting. - The posterior update perspective guarantees that under the proposed NTE theory, the model after seeing multiple tasks sequentially is identical to the one that sees all tasks jointly. - The problem studied (catastrophic forgetting) is of great importance. - The theory is mostly sound and the assumptions are reasonable. - The application of the interpretation gives useful insights, such as momentum causes forgetting.

Weaknesses

- It is unclear to me, under the Neural Tangent Ensemble theory, how far can the model change from the initialization, if the model is not changing too much, then it is not surprising that the model is not forgetting, since it is not learning a lot information from the training data. - Similarly, the theory suggest that the model forgets less when getting wider, but the representation learning flexibility would also drop. So the model is learning less while forgetting less. However, good representation learning is crucial for the performance of many deep learning tasks, does it mean there would be inevitable performance v.s. forgetting tradeoff? - The experiments seem to be of rather small scale: Although sufficient for supporting the theory, it would still be nice to see if the results can be transferred to real world applications.

Questions

- What does the subscript k (y_k and x_k) mean in the equation at the bottom of page 3.

Rating

6

Confidence

2

Soundness

4

Presentation

4

Contribution

4

Limitations

- It is unclear how the proposed interpretation can be generalized to more complicated networks or settings, e.g. deeper network or when the parameters move far away from the initialization.

Further Question

It seems that in more complex scenarios in CIFAR-100, the proposed NFE behaves much poorer than ADAM and SGD in small parameter sizes. Could the authors kindly elaborate on the possible reasons, especially the comparison between SGD and NEF, as long as neither introduces momentum?