Continual learning with the neural tangent ensemble

Authors: ari s benjamin
christian  pehle
kyle  daruwalla
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 21:30:22

Abstract: A natural strategy for continual learning is to weigh a Bayesian ensemble of fixed functions. This suggests that if a (single) neural network could be interpreted as an ensemble, one could design effective algorithms that learn without forgetting. To realize this possibility, we observe that a neural network classifier with N parameters can be interpreted as a weighted ensemble of N classifiers, and that in the lazy regime limit these classifiers are fixed throughout learning. We call these classifiers the *neural tangent experts* and show they output valid probability distributions over the labels. We then derive the likelihood and posterior probability of each expert given past data. Surprisingly,  the posterior updates for these experts are equivalent to a scaled and projected form of stochastic gradient descent (SGD) over the network weights. Away from the lazy regime, networks can be seen as ensembles of adaptive experts which improve over time. These results offer a new interpretation of neural networks as Bayesian ensembles of experts, providing a principled framework for understanding and mitigating catastrophic forgetting in continual learning settings.

Main paper: 
 introduction: 
 Neural networks often forget previous knowledge when trained with gradient descent. In contrast, animals learn from sequential experiences, suggesting that true 'lifelong learners' use different strategies for learning [25] .One strategy to learns without forgetting is to update the posterior distribution over a set of fixed probabilistic models [8] . This includes any fully Bayesian model, such as Bayesian linear regression. The fundamental reason why these algorithms do not forget information is because the posterior over models is invariant to data sequence. Given two permutations of the data, the posterior will be the same. This property of posteriors has inspired many strategies to reduce forgetting by approximating the posterior distribution over neural networks [22, 24, 11, 28, 26, 41, 37] . However, these approximations introduce many new parameters and considerable memory overhead. In general, estimating the full posterior distribution over high-dimensional networks is prohibitive.Here, we shift our perspective and instead interpret a single neural network as an ensemble of many experts. This allows tracking a posterior (over experts, instead of networks) without introducing any memory overhead besides the network itself. This motivates our main result, which we note is generally applicable outside of continual learning. More specifically, we show that neural network classifiers perturbed by a small vector in parameter space can be described as a weighted ensemble of valid classifiers outputting a probability distribution over labels. We call this the Neural Tangent Ensemble (NTE). Inspired by the Neural Tangent Kernel, this result depends on a first-order Taylor expansion around a seed point [19] . As a consequence, it operates as an ensemble of fixed classifiers in the NTK limit of infinite width.In this framework, learning is framed as Bayesian posterior updating rather than optimization. These two approaches might be expected to be quite different, as a posterior update is multiplicative whereas gradient-based optimization is additive. Surprisingly, however, we find that the NTE posterior update rule is approximately stochastic gradient descent (SGD) on the network with batch size 1, thus shedding new light on the dynamics of neural network optimization.Our primary contributions are:• We introduce the Neural Tangent Ensemble (NTE), a novel formulation that interprets networks as ensembles of classifiers, with each parameter contributing one classifier. • We derive the posterior update rule for the NTE for networks in the lazy regime, in which experts are fixed, and show that it is equivalent to single-example stochastic gradient descent (SGD) without momentum, projected to the probability simplex. • This justifies the empirical finding that SGD with no momentum forgets much less than standard optimizer settings. • We demonstrate that catastrophic forgetting in neural networks is associated with the transition from the lazy to the rich regime.2 Motivation: Ensembles are natural continual learners Pruning the set of functions f i to those good for task A, followed by further pruning for tasks B and C, will result in a set of f i still good on A.To demonstrate why Bayesian ensembles make good continual learners, consider a function f (x) that is an ensemble of many experts f i (x) (Fig. 1 ). We will consider what it takes to modify this ensemble so that it performs well on two tasks A and B.A simple strategy for continual learning is to prune away experts. Let S A be the set of functions that are good (and equally good) for task A. A good ensemble can be constructed by sampling from S A :f A (x) = 1 N fi∈S A f i (x).Given a subsequent task B, a new ensemble can be constructed on the fly by continuing to prune away the experts in f A (x) that perform poorly on task B. The remaining ensemble is still composed of experts from S A (assuming that the set intersection is not zero).In contrast to many continual learning strategies for neural networks, this does not require replay, task boundaries, or any additional memory dedicated to preserving old task performance. 
belief updating generalizes set intersections: 
 In real ensembles, experts do not perform equally well. This justifies weighing each expert in the ensemble with weights p i which are chosen such that N i p i = 1:EQUATIONThis is particularly convenient if the experts encode the probability or belief about an event, f i (x) = p(y|x, f i ). In this case, one can weigh each function by its posterior probability given previous data:EQUATIONThis is the optimal weighing strategy when the experts can be assumed to be independent [45] .It is useful to contrast the ensemble in Eq. 2 with linear regression using a feature map, f (x) = i w i ϕ i (x), as might be observed in kernel regression. In an ensemble the weights w i are strictly positive, whereas weights in regression may switch sign arbitrarily.The property of Bayesian ensembles that motivates this paper is that the posterior probability of each expert is invariant to the order in which data in seen. This is because, like set intersections, single-task posteriors multiply to form the multi-task posterior:EQUATION)This property is restated more formally in the following Lemma:Lemma 1. Invariance to data ordering in Bayesian Ensembles. Let F = f 1 , ..., f N be a set of fixed experts, W = w 1 , ..., w N be their weights, and D = D 1 , ..., D T be a sequence of datasets from T tasks. Then, for any permutation π of the indices 1, ..., T, p(f i |D) = p(f i |D 1 , ..., D T ) = p(f i |D π(1) , ..., D π(T ) ) Proof. By Bayes' rule, p(f i |D) ∝ p(f i ) T t=1 p(D t |f i ).The right-hand side is a product of terms, one for each dataset. Since multiplication is commutative,T t=1 p(D t |f i ) = T t=1 p(D π(t) |f i ) for any permutation π. Therefore, p(f i |D 1 , ..., D T ) = p(f i |D π(1) , ..., D π(T ) ).Thus, there is no catastrophic forgetting problem for models which are ensembles of fixed, independent probabilistic classifiers. This motivates assessing under what conditions neural networks approach this setting.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: The main limitations of this work are 1) Theorem 1 is valid only in the limit of small parameter changes, and with the assumption of low network curvature. This assumption is stated clearly in the Theorem statement. Then, in the discussion, we mention that weighing experts by their posterior probability is only optimal when the experts are independent, which is not the case. We do not claim this is the optimal weighing strategy. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
the neural tangent ensemble: 
 How might a neural network be described as an ensemble? One simple strategy would be to take the last network layer as a set of functions, and then to choose the output weights according to their relative performance. However, this is is an expensive strategy to construct a relatively small set of classifiers, and it does not specify how earlier weights might change.Here, we employ a first-order Taylor expansion to show that neural networks are (approximately) large ensembles over N component functions, one for each edge in the network. We will examine a neural network p(y|x, W (t) ) with parameters W (t) whose output represents the probability or confidence of a label y given input x. We can describe this output with a linearization around a very nearby seed point W (0) . Note that we use the notation W (0) and W (t) for consistency and in general W (0) need not be the initialization or on the optimization trajectory at all.EQUATIONAt first glance it does not appear that this Taylor expansion is an ensemble. There seem to be no true classifiers: the gradients are not probabilities over classes y, being neither nonnegative, bounded, nor normalized to 1 across the output labels. Nor are there true weights, as ∆w i is also not nonnegative. However, both of these criteria can be met with some rearrangements and with the assumption that the loss is sufficiently smooth with respect to its parameters. This leads to our main result:Theorem 2. Suppose p(y|x, W (0) ) is a neural network for which the log-likelihood is L-Lipschitz continuous in its parameters, i.e. all gradients of the loss are bounded by a constant L. Let W (0) then be perturbed by a ∆W with ∥∆W ∥ 1 = z. If the perturbation is sufficiently small (with zL < 1), then the network can be described as an ensemble of a set of N classifiers {p(y|x, f i )} N i , each with weight |∆wi| z , plus higher-order contributions which vanish for small z:p(y|x, W (t) ) = N weights i |∆w i | z p(y|x, f i ) + O(∥∆W ∥ 2 2 )Each classifier p(y|x, f i ), which we call the neural tangent expert, outputs a probability distribution over labels y:p(y|x, f i ) = p(y|x, W (0) ) 1 + z sign(∆w i ) ∂ ∂w (0) i log p(y|x, W (0) )The proof is postponed to Appendix 8.1. Informally, it relies two simple rearrangements: splitting the weights into sign and magnitude ∆w i = |∆w i |sign(∆w i ), and bringing the zeroth order term inside the first-order sum. This results in a sum over a term which, surprisingly, sums to 1 over the output labels and is weighted by a term that sums to 1 over experts, meeting the criteria of an ensemble.This simple reformulation invites a change in perspective about the role of each parameter in a deep neural network. Each parameter contributes a separate classifier. The distributed architecture and connected paths of the network matter, but they explicitly contribute through the gradients alone.In the literature on ensembles, a common focus is to examine the quality and diversity of the experts separately. By the bias/variance decomposition, both aspects enter in the generalization error [38, 47] .Here, it is clear that all experts share a factor that is the overall quality of the center of the Taylor expansion, p(y|x, W (0) ). What distinguishes experts from one another is the diversity of network gradients.This paper is motivated by the fact that Bayesian ensembles of fixed experts do not forget past data when learning by posterior updating. Under what conditions is the Neural Tangent Ensemble composed of fixed functions?The answer to this question follows directly from the literature on the Neural Tangent Kernel (NTK) and lazy regime networks [19, 7] . If the network is in the 'lazy' regime, then the Jacobian of the network does not change during gradient descent learning and the linearization remains valid. This occurs in the limit of infinite width for MLPs for certain initializations [19] . (Output scaling also controls laziness [7] , and is a necessary when using softmax nonlinearities even in the infinite width [29] .) As a consequence, the experts in the NTE interpretation are fixed functions in the lazy regime.If a network is secretly an ensemble, how should it learn from new data? The natural next step is to convert the NTE into a Bayesian ensemble. In a Bayesian ensemble, the weight of each function is its posterior probability given past data:EQUATION)This can be seen as the E step in a generalized EM algorithm [45] . In the following section, we will describe how to calculate this posterior probability with an online learning algorithm. For the moment, we assume the experts are fixed functions (i.e. the network is lazy).Lemma 3. For IID data D = {x k , y k } P k=1 , the likelihood of the data under an expert can be written in terms of a log-likelihood loss function ℓ(0) k = -log p(y k |x k , W (0)) of the network at initialization:EQUATIONProof. Starting with the data likelihood,EQUATIONPlugging in the definition of ℓ (0) k yields the above expression.After the data likelihoods are computed for each neural tangent expert, they must be renormalized to obtain the posterior probabilities. In our case, we naturally have access to a very large number of tangent experts and their likelihoods. Indeed, if the width is indeed taken to infinity, this there are infinitely many neural tangent experts in a single network. We propose to use the distribution of likelihoods in the current network as a Monte Carlo estimate of the normalizing constant.EQUATIONThis can be simplified by noting that each e -ℓ (0) k term will cancel; the product k e -ℓ (0) k appears in every additive term in the denominator. Assuming a uniform prior p(f i ), we then have:p(f i |D) = k 1 -z sign(∆w i ) ∂ ∂w (0) i ℓ (0) k i k 1 -z sign(∆w i ) ∂ ∂w (0) i ℓ (0) k (11)We will now link this posterior expression with a neural network update rule. Recall that in Theorem 2, the normalized magnitude of each perturbation is interpreted as the posterior probability of the corresponding neural tangent expert.|∆w i | z ← p(f i |D)This means the parameters can act as a running cache of the posterior as new data is encountered. As in standard belief updating, this involves a likelihood update followed by renormalization. Surprisingly, this multiplicative belief update rule yields an update which is very close to SGD. Lemma 4. For any network that is well-described as a first-order Taylor expansion around around W (0) with perturbation ∥∆W ∥ 1 = z, the posterior belief update given a new example is equivalent to single-example stochastic gradient descent under a cross-entropy loss objective, subject to the constraint that ∥∆W ∥ 1 = z, and using a per-parameter learning rate of z|∆w i |.Proof. The proof is a matter of writing out how the posterior changes with a single example. Multiplying by the likelihood of a new example, the unnormalized posterior updates as:EQUATIONThis multiplicative update for the unnormalized weights can also be written an additive rule. Multiplying by z and by sign(∆w i ), ∆wEQUATIONOne can add the initial parameters to either side to yield a rule in the space of network parameters:EQUATIONThis is true (single-example) stochastic gradient descent projected in the L1 diamond with a learning rate z|∆w i |. Note that this does not allow averaging gradients across examples (a "batch size of 1" update) and that it uses the gradients at initialization (though see section 4.1).To complete the update, the parameters should then be renormalized such that i |∆w(t) i | = z.An alternative normalization scheme is to use a gradient projection algorithm. Adding a Langrage multiplier γ to Eq. 13 and solving for the γ that ensures i |∆w i | = z yields a update which keeps ∥∆W ∥ 1 = z even without renormalization:EQUATION)Not only is the posterior update tractable, then, but it is sufficiently close to gradient descent that it can be interpreted in a standard optimization framework.Although it may seem that our result would depend on the idiosyncratic likelihood function of the NTE, this result is nevertheless similar to previous algorithms that have been proposed as ways to weigh many experts. At high level, our result appears similar to the Multiplicative Weights algorithm described in [1] . Another interpretation of this algorithm is as the approximated exponential gradient descent with positive and negative weights algorithm from [23] but applied to the change in weights ∆W . There, it is derived by minimizing an arbitrary loss function under a constrained change in the relative entropy over ensemble weights to obtain the exponentiated gradient descent algorithm, which is then linearized with a Taylor expansion in the approximated version.The Neural Tangent Ensemble is an interpretation of networks as ensembles of neural tangent experts.Updating the NTE of lazy networks as a Bayesian ensemble creates a perfect continual learner, in the sense that the multitask solution is guaranteed to be the same as the sequential task solution.The posterior probability of each expert in the NTE is surprisingly tractable. Given a new example, the update rule is a simple additive rule in the space of network parameters which can be interpreted as projected gradient descent scaled by the change in parameters since initialization.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: All assumptions are listed upfront, and the proofs are complete. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
networks away from the lazy regime: 
 In real finite-width networks, gradients change throughout learning. Since each weight's corresponding neural tangent expert changes, there is no guarantee that weights at time t still reflect the cumulative likelihood of past data under that expert.This phenomenon is clearly observed empirically by measuring how much experts change under the NTE update rule as a function of network width. In Fig. 2 , we measure the average change in expert's Jacobian from initialization after training on MNIST as a function of network width with the NTE rule described above. Experts change less in wider networks than in narrow networks. Another way this can be measured is by verifying that, in finite-width networks, the posterior update rule using the gradients around initialization does not lead to effective training. In Figure 3 , we confirm that as the gradients lose correlation with the gradient at initialization, performance begins to rapidly degrade. This echoes the findings of [7] that linearized CNNs do not learn as effectively as their non-lazy counterparts. Thus, the NTE posterior update rule as written above is only effective when the Jacobian is truly static. (informal) Let W (t) be the parameters of a (finite-width) neural network. Choose a nearby seed point W (t) as W (t) + ϵ, with ϵ fixed and ∥ϵ∥ 2 sufficiently small relative to the curvature such that the Jacobians of the log output probabilities of the perturbed and unperturbed networks are identical, J( W (t) ) = J(W (t) ). The network can then be written as an ensemble of adaptive experts:p(y|x, f (t) i ) = p(y|x, W (t) ) 1 + ∥ϵ∥ 1 sign(ϵ i ) ∂ ∂w (t) i log p(y|x, W (t) )If ϵ is set as the uniform vector with values ϵ i = η/N , the learning rate in the posterior update rule reduces to ∥ϵ∥ 1 |ϵ i | = η and we recover stochastic gradient descent with mean-centered gradients and learning rate η:EQUATIONRich-regime learning is thus akin to a particle filter; each expert changes individually, but the prediction is the ensemble vote.A interesting feature of this lemma is the equivalence between the rule that improves each expert (gradient descent on w) and the rule that decides how to weigh the experts in the ensemble (also gradient descent on w). This need not have been the case. As a result, one can perform belief updating assuming a fixed ensemble and end up improving each expert within it.Motivated by this result, we evaluated how well the NTE posterior update rule works when the gradients evaluated at initialization, ∂ ∂w(0) i ℓ (0)K , are replaced with the gradients of the current network∂ ∂w (t) i ℓ (t)K . These converge in the infinite-width limit.To obtain a practical algorithm, we additionally modify the NTE update rule with two hyperparameters that control the learning rate. First, noting that z in Eq. 14 acts as a learning rate, we replace it with a tunable parameter η. Secondly, we introduce a regularization parameter β which keeps the network close to initialization as measured by the relative entropy of the change in parameters (see Appendix 8.2 for derivation). This constrains the amount of information contained in the weights [17] .Pseudocode for the resulting algorithm is in the Appendix 1. We also display the result of sweeps over β and η on the Permuted MNIST task in Fig. 7 .Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: The experiment details are listed in full. Code is also provided that implements all models and optimizers.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
predictions and results: 
 Our findings suggest several ways to reduce forgetting in finite networks. First, networks closer to the lazy regime will better remember old tasks as long as the update rule is sufficiently similar to the NTE posterior update rule. Second, one should be able to reduce forgetting by ablating standard optimization methods like momentum and moving towards the NTE posterior update rule.Below, we verify these predictions on the Permuted MNIST task with MLPs and on the taskincremental CIFAR100 with modern CNN architectures. In the Permuted MNIST task, an MLP with 10 output units is tasked with repeatedly classifying MNIST, but in each task the pixels are shuffled with a new static permutation. In task-incremental CIFAR100, a convolutional net with 100 output units sees only 10 classes each task. In the terminology of van de Ven et al, this is a task-incremental task, whereas Permuted MNIST is a domain-incremental task [48] .Momentum is not appropriate in a posterior update framework because it over-counts the likelihood of past data. Furthermore, it is a history-dependent factor. By contrast, posterior update rules are multiplicative and give identical results regardless of the order of data presentation.Here, we report that any amount of momentum with SGD is harmful for remembering past tasks. To our knowledge, this has not been noted by previous empirical studies on catastrophic forgetting [13, 36, 35, 2] . As can been seen in Fig. 4 , increasing momentum monotonically increases forgetting a first task on Permuted MNIST. Similar trends exist for ResNet18 and ConvNeXtTiny on the CIFAR100 task (see Fig. 8 ) [30] . Note that the momentum buffers were not reset between tasks; when they are reset, the momentum curve is nonmonotonic (see Fig. 9 ). Although momentum assists single-task performance, any amount of momentum will lead to forgetting previous knowledge. As networks grow wider and (slowly) approach their infinite-width limit, they should remember better if one uses the appropriate posterior update rule over the Neural Tangent Ensemble.Previous literature confirms that this is generally the case. In [40] , the authors report the benefits of scale are robust across architectures, tasks, and pretraining strategies, although they largely use SGD with momentum β = 0.9. In [35] , the authors report similar results and investigate other continual learning benchmark algorithms such as EWC ( [22] ). Forgetting seems to be largely solved by scale.The reason for this in our framework differs from the reason cited by both [35, 40] , which state that the gradients on different tasks will be more orthogonal in high dimensions, which reduces interference. Our interpretation is somewhat different and instead depends on the Jacobian of the network changes. We place no condition on gradient orthogonality between tasks. If the neural tangent experts are indeed fixed, the NTE update rule will find the multitask solution.If this is the case, then wide networks should better remember only if the optimizer can be interpreted as a posterior update. In Fig. 5 , we report that Adam's amnesia is not helped with increasing scale for the Permuted MNIST task. Although this could be for multiple reasons, we argue it stems from a divergence from a valid interpretation as a posterior update.In Section 4.2, we introduced a modified version of NTE posterior update rule in which the Jacobian at initialization replaced with the current Jacobian. As networks get wider, this algorithm will converges to the proper update rule due to the fact that the network Jacobian does not change in the lazy regime. This predicts that this rule will improve with scale. To test this, we trained an ML on Permuted MNIST and ConvNeXtTiny on task-incremental CIFAR100 with this approximate rule. We find that both single-task and multitask learning are greatly improved with width (Fig. 5 and Fig. 10 ). We take this as empirical evidence that the proper NTE posterior update (with a static Jacobian) would work well in the infinite-width limit. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: Code is provided in supplementary information.Guidelines:• The answer NA means that paper does not include experiments requiring code. Justification: All experimental details are provided.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
related work: 
 There is a long history of interpreting networks as ensembles. Networks with dropout, for example, allow this interpretation [12, 14] . This is also closely related to Mixture of Experts models in classic [18, 20] and recent [43, 54] work. These approaches explicitly encode the experts within the network, and unlike our work do not use a Taylor expansion to establish the ensemble experts.The idea of a Bayesian ensemble over networks is also well-studied. Such ensembles can either be assembled empirically through sampling [51, 50] , built via a Laplace approximation [32] , or optimized [3] . Bayesian posteriors are also common players in theoretical works using methods from statistical physics and PAC-Bayes [44, 27] . Some treatments of infinite-width limits study the ensemble of lazy learners [16] . While similar in spirit, these methods study groups of many networks rather than view a single network as an ensemble.Finally, there is related work that uses ideas from ensembles for continual learning. Many of these are in the category of methods that continually learn by training new modules for each task [49, 5, 42, 52, 39, 21] . Most directly related to this current work are papers that take a Bayesian approach and track statistics about the approximated posterior over networks [22, 10, 11, 28, 41, 37] . Many of these works in both categories require task boundaries. Furthermore, by introducing new modules or tracking statistics, these methods require additional memory to prevent forgetting.Our framework justifies the strategy of encouraging parameters to change mostly in directions of low curvature. Such regularization methods are already well-established and proven to reduce forgetting [24, 31, 41] . Although not directly equivalent, this is also similar to Elastic Weight Consolidation, which penalizes by the Fisher Information matrix (an expected second-order derivative of the loglikelihood, rather than the likelihood) [22] . Another proximal method is Synaptic Intelligence, which penalizes parameter changes proportional to their integrated gradients along the path, which in the special case of diagonal, quadratic loss functions, is equivalent the Hessian [53] . Thus, a second interpretation of why these methods work well (and improve with scale [35] ) is that they ensure the tangent experts in the NTE do not change much while learning. 
discussion: 
 Here, we described how networks in the lazy regime can be seen as ensembles of fixed classifiers. With this perspective, we proposed weighing each expert by its posterior probability to form a Bayesian ensemble, and derived the update rule. This strategy of learning by posterior updating has the benefit that the order of data presentation does not matter -sequential experience and interleaved experience lead to identical ensembles.The posterior update rule to the Tangent Ensemble is surprisingly similar to SGD on the model weights. However, it is interesting to note that this update rule is suboptimal. Posterior probabilities are the optimal ensemble weights only when the experts are independent [47, 34, 38] and wellspecified [33, 6] . This assumption is violated by the use of shared data, as well as the fact that neural network architectures introduce dependencies between gradients. Although this does not affect the equivalence between the interleaved and sequential task performance (i.e. forgetting), this will reduce the performance of networks trained with the NTE posterior update. This suggests avenues for improving SGD. This suboptimality could be addressed in multiple ways. In the ensemble literature, there are many strategies to diversify the expert pool [4] such as repulsion [9] . Different experts might also be trained on different data [46] , and one might even take a boosting approach [15] . It is very likely that these approaches would yield neural networks that outperform standard networks trained by updating the posterior distribution over tangent experts.The ability in interpret single networks as ensembles opens many avenues for future research. These extend beyond continual learning; for example, one might be able to obtain a measure of uncertainty of the network output via the variance of the experts [12] . We are hopeful that this insight will lead to deep learning systems that are better understood as their use expands within society.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Error bars where present are clearly described as representing the standard deviation over seeds.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: Although this theoretical work has little direct impact, we mention in the discussion the impacts of continual learning devices upon society. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 
proof of theorem 2: 
 Proof. We begin by noting that the change in weights we can be split up the sign and magnitude,∆w i = |∆w i |sign(∆w i ).We will then interpret |∆w i | as the unnormalized component weight. The remaining terms must be the component functions.To identify these functions, and show that they satisfy the properties of a probability distribution, we will rearrange terms. First, noting that i |∆w i | = z for some constant z (potentially z = 1),EQUATIONEQUATIONEQUATIONWe call the term p(y|x, f i ) the neural tangent expert.The neural tangent expert provides a valid probability distribution for small ∆W . First, see that it satisfies j p(y j |x, f i ) = 1. This can be seen from the fact that the right term inside p(y|x, f i ) (the parentheses in the middle line) sums to 0 over the output label:j p(y j |x, f i ) = j z sign(∆w i ) ∂p(y j |x, W (0) ) ∂w (0) i = z sign(∆w i ) j p(y j |x, W (0) ) ∂ log p(y j |x, W (0) ) ∂w (0) i = 0Here we have used the identity that the expectation of a score function is zero.Next, we will show that 1 ≥ p(y j |x, f i ) ≥ 0. First, since each p(y j |x, f i ) sum to 1 over j, no component can be greater than 1 if all components are nonnegative. Thus, it only needs to be shown that p(y|x, f i ) ≥ 0. While this cannot be guaranteed in general, by construction we have assumed that zL < 1. This Lipschitz continuity bounds the L2 norm of the gradients of the log likelihood, which in turn bounds the L1 norm and implies that any individual gradient has magnitude less than 1 z :EQUATIONThus, whether sign(∆w i ) = 1 or sign(∆w i ) = -1, the term in parenthesis is nonnegative.1 + z sign(∆w i ) ∂ ∂w (0) i log p(y|x, W (0) ) > 0.The continual learning ability of a Bayesian ensemble derives from learning to weight a set of fixed functions. If these functions change over time, then there is no guarantee that the likelihood of each function at time t still reflects the cumulative likelihood of past data under the current function.One good way to ensure this is does not occur is to ensure that the parameters change as little as possible from initialization. Although it is typical to measure this distance with ∥∆W ∥ 2 , we instead measure distance as the relative entropy of the change in parameters from the uniform perturbation, due to the simplicity of its result. These have the same minimum; remembering that ∥∆W ∥ 1 = 1, by normalization, the smallest Euclidean distance ∥∆W ∥ 2 will occur when all ∆w i are equal.To derive the maximum-entropy vector |∆W | that is as similar as possible to p(f i |D), we will follow the steps of [23] . A first step is to set the notion of similarity L between |∆W | and p(f i |D). We will then find the value that minimizes:U (|∆W |) = -H[|∆W |] + βL(|∆W |, {p(f i |D)}) + γ(∥∆W ∥ 1 -1)Here β is a parameter that trades off between entropy and matching p(f i |D), and γ is a Langrange multiplier that ensures the parameters remain normalized.If one chooses to maximize the dot product |∆W | T p(f i |D), one obtains the following relation:w i = e β p(fi|D) i e β p(fi|D)Alternatively, if one chooses to minimize the relative entropy KL(|∆W |, p(f i |D)), then one obtainsw i = p(f i |D) β i p(f i |D) βWe implement this second term. If the posterior likelihoods are maintained in log space, β acts as a multiplicative scale upon the log data likelihood. k |x k , f i ) = 1 -η sign(∆w i ) ∂ ∂w (t) i ℓ (t) kUpdate perturbation multiplicatively ∆w i ← ∆w i p(y k |x k , f i ) β end for Renormalize perturbation such that i |∆w i | = z Optionally clip the change in parameters to prevent large changes, such that |∆w i | = 1 end forAll MNIST experiments were completed on two NVIDIA RTX 6000 cards, and all CIFAR100 experiments were conducted on NVIDIA H100 cards. Over 1,500 individual models were trained across all seeds and conditions, amounting to roughly 440 GPU-hours of compute time. 2 A single MLP was trained with 1,000 hidden units per layer and 2 hidden layers using ReLU nonlinearities. The model perturbed from initialization with a random normal vector with scale 0.001, and then was trained with the NTE update rule (Algorithm 1) but using the Jacobian of the model at initialization. The batch size was 24 and the parameters of the NTE algorithm were η = 0.01 and β = 1.We first created a Permuted MNIST task and code to measure the test accuracy on all tasks after training on each task sequentially. All reported results use 5 tasks.We trained an MLP on this task with ReLU nonlinearities and 1,000 hidden units in 2 hidden layers. We used SGD with batch size 128, learning rate 0.01, and momentum swept from 0 to 1. The momentum buffer was not reset between tasks. We report the standard deviation of 10 random seeds. Here, we used the same continual learning task as Figure 3 , but swept the width of the two hidden layers from 10 to 10,000. All batch sizes were 128. 4 , we trained a 2L MLP with 1,000 hidden units on the Permuted MNIST task and varied the momentum of SGD. This time, we reset the momemtum buffer between tasks. Interestingly, this introduces a nonmonotonic behavior and one can attain good performance with momentum near 0.99.Figure 10 : Performance with width for CIFAR-100. We scaled the number of convolutional filters in all layers of a ConvNeXtTiny by a constant factor, and then trained on the CIFAR-100 taskincremental task for each network. Subpanels represent identical information as Fig. 8 . • The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
 
 Appendix:  


Decision: The paper has several interesting theoretical contributions, mainly new insights on interpreting a model as ensemble of classifiers and analyzing forgetting. The experimental design is interesting and further supports the theoretical findings. I recommend accept.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: - What does the subscript k (y_k and x_k) mean in the equation at the bottom of page 3.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: The authors have outlined the limitations and the potential societal impact of the work.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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
limitations: limitations: Discussion is only effective within a lazy regime. Additionally, experts are required to be independent.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: 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.

Comments: title: Further Question
comment: I would like to thank the author for the detailed response! They resolved my concerns, I've increased my score.

Overall I think this is a good paper that should be published, and I don't have too much concerns about the empirical evaluation's scale as I think the value of the paper is to provide a new framework for understanding continual learning.
Thank you for your explanation. I appreciate the additional experiments and I will raise my score.
Thank you for the response this clarifies much of what I was unclear on.

> For finite networks, a single column of the Jacobian can be understood as an expert, and the number of experts (length of the column) is equal to the number of parameters.

This clears things up for me and may be worth stating in the paper itself. 

There is an additional question I have with respect to the momentum experiment: when using momentum, was the optimizer state reset between task changes? I think the experiment is interesting regardless of this, but I think this is an important detail to state.

Overall I think this is a strong paper and recommend accepting.
What we intent to show with this plot is that the NTE rule improves with network width, which it does - and note that for very large networks, it actually outperforms SGD. Underperformance is only true for smaller networks. This is to be expected, as the Neural Tangent Ensemble rule is derived from an infinite width limit. We think this is more pronounced for this architecture because for ConvNets the relevant NTK limit is over the number of filters per layer, rather than the overall number of parameters per layer. In the plot above, the number of filters ranges from 12 to 1028. Thus we are displaying more of the regime father away from the NTK limit, yet as we approach it, NTE rapidly improves.
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?
Thank you for your response. My concerns have been addressed by the additional experiments. Therefore, I will update my score from 4 to 5.