Learning in temporally structured environments
Authors: Matt Jones
Tyler R. Scott
Mengye Ren
Gamaleldin Fathy Elsayed
Katherine Hermann
David Mayo
Michael Curtis Mozer
Venue: ICLR 2023 poster (2023)
Keywords: 1/f noise
Kalman filter
neural network
learning theory
optimizers

Abstract: Natural environments have temporal structure at multiple timescales. This property is reflected in biological learning and memory but typically not in machine learning systems. We advance a multiscale learning method in which each weight in a neural network is decomposed as a sum of subweights with different learning and decay rates. Thus knowledge becomes distributed across different timescales, enabling rapid adaptation to task changes while avoiding catastrophic interference. First, we prove previous models that learn at multiple timescales, but with complex coupling between timescales, are equivalent to multiscale learning via a reparameterization that eliminates this coupling. The same analysis yields a new characterization of momentum learning, as a fast weight with a negative learning rate. Second, we derive a model of Bayesian inference over $1/f$ noise, a common temporal pattern in many online learning domains that involves long-range (power law) autocorrelations. The generative side of the model expresses $1/f$ noise as a sum of diffusion processes at different timescales, and the inferential side tracks these latent processes using a Kalman filter. We then derive a variational approximation to the Bayesian model and show how it is an extension of the multiscale learner. The result is an optimizer that can be used as a drop-in replacement in an arbitrary neural network architecture. Third, we evaluate the ability of these methods to handle nonstationarity by testing them in online prediction tasks characterized by $1/f$ noise in the latent parameters. We find that the Bayesian model significantly outperforms online stochastic gradient descent and two batch heuristics that rely preferentially or exclusively on more recent data. Moreover, the variational approximation performs nearly as well as the full Bayesian model, and with memory requirements that are linear in the size of the network.


Published as a conference paper at ICLR 2023
LEARNING IN TEMPORALLY STRUCTURED ENVIRONMENTS
Matt Jones, 1.2 Tyler R. Scott,' 1 Mengye Ren, Gamaleldin ElSayed, Katherine Hermann," David Mayo, 1.1 Michael C. Mozer! BBrinn Team, Google Research University of Colorado NYU 4MIT mej@colorado.. edu dmayo2 @mit. edu mengyelcs eyu ftylerscc, gamaleldin, he hermannk, mcmozer @google.com
ABSTRACT
Natural environments have temporal structure at multiple timescales. This prop- erty is reflected in biological learning and memory but typically not in machine learning systems. We advance multiscale learning method in which each weight in a neural network is decomposed as a sum of subweights with different learning and decay rates. Thus knowledge becomes distributed across different timescales, enabling rapid adaptation to task changes while avoiding catastrophic interfer- ence. First, we prove previous models that learn at multiple timescales, but with complex coupling between timescales, are equivalent to multiscale learning via a reparameterization that eliminates this coupling. The same analysis yields a new characterization of momentum learning, as a fast weight with a negative learning rate. Second, we derive model of Bayesian inference over 1/ʃ noise, a com- mon temporal pattern in many online learning domains that involves long-range (power law) autocorrelations. The generative side of the model expresses 1/ f/f noise sum of diffusion processes at different timescales, and the inferential side tracks these latent processes using a Kalman filter. We then derive a varia- tional approximation to the Bayesian model and show how is extension of the multiscale learner. The result is an optimizer that can be used as a drop-in replace- ment an arbitrary neural network architecture. Third, evaluate the ability of these methods to handle nonstationarity by testing them in online prediction tasks characterized by 1/f noise in the latent parameters. We find that the Bayesian model significantly outperforms online stochastic gradient descent and two batch heuristics that rely preferentially or exclusively on more recent data. Moreover, the variational approximation performs nearly as well as the full Bayesian model, and with memory requirements that are linear in the size of the network.
INTRODUCTION
Many online tasks facing both biological and artificial intelligence systems involve changes in data distribution over time. Natural environments exhibit correlations at wide range of timescales, pattern variously referred as self-similarity, power-law correlations, and 1/f noise (Keshner) 1982). This is in stark contrast with the iid environments assumed by many machine learning (ML) methods, and with diffusion or random-walk environments that exhibit only short-range correla- tions. Moreover, biological learning systems are well-tuned to the temporal statistics of natural environments, as seen in phenomena of human cognition including power laws in learning (Ander-
son 1982), power-law forgetting (Wixted & Ebbesen, 1997), long-range sequential effects (Wilder al. 2013). , and spacing effects (Anderson & Schooler, 1991 Cepeda al., 2008). An important goal is to incorporate similar inductive biases into ML systems for online or continual learning.
This paper analyzes a framework for learning in temporally structured environments, multiscale learning, which for neural networks (NNs) can be implemented mon explanation for self-similar temporal structure in nature is that arises from a mixture of events at various timescales. Indeed, many generative models of 1/f noise involve summing independent stochastic processes with varying time constants (Eliazar & Klafter, 2009). Accordingly, the mul- tiscale optimizer comprises multiple learning processes operating in parallel different timescales.
new kind of optimizer. A com-

---

Published as a conference paper at ICLR 2023
In NN, every weight w is replaced by family of subweights Wej, each with its own learning rate and decay rate, that sum to determine the weight as a whole. Learning at multiple timescales is a key idea in several theories in neuroscience, including conditioning (Staddon et al. (002), learning 2016), memory (Howard & Kahana 2002 Mozer et al. [2009), and motor control 2007 and has also been exploited in ML. (Hinton & Plaut 1987; Rusch et al.)
(Benna & Fusi
(Kording al.
2022). The multiscal learner isolates and simplifies this idea, by assuming knowledge at different timescales evolves independently and that credit assignment follows gradient descent.
The first part of this paper (Sections 2 and proves three other models are formally equivalent to instances of the multiscale optimizer: new variant of fast weights (cf. [a et al.) 2016; Hinton & Plaut 1987). the model synapse of Benna & Fusi (2016). and momentum learning (Rumelhar et al., 1986 Qian 1999). The insight behind these proofs is that each of these models can be written in terms of linear update rule with diagonalizable transition matrix. Thus the eigenvectors of this matrix correspond to states that evolve independently. By writing the state of the model as a mixture of eigenvectors, we effect coordinate transformation that exactly yields the multiscale optimizer. These results imply that the complicated coupling among timescales assumed by some models can be superfluous. They also provide new perspective on momentum learning, with implications for how and when it is beneficial and how it interacts with nonstationarity in the task environment.
In Section 4 we provide a normative grounding for multiscale learning in terms of Bayesian infer- ence over T/ʃ noise. Our starting point is a generative model of 1/f noise as a sum of diffusion processes at different timescales. Exact Bayesian inference with respect to this generative process is possible using a Kalman filter (KF) that tracks the component processes jointly (Kording et al.) 2007). When learning a single environmental parameter 0, such as mean reward for some action in a bandit task, this amounts to modeling (()) = Σx= (t), where each 2; is a diffusion process with a different characteristic timescale Tú. and doing joint inference over
(Z1..... 2n).
We then generalize this approach to an arbitrary statistical model, h (xx 0), where æ is the input and e T is a parameter vector to be estimated. For instance, h might be a NN with parameters O. Our Bayesian model places a 1/f prior on 0 (as a stochastic process), by assuming O(t) Σm= (t) for diffusion processes 2; E IR" m with characteristic timescales 73. We then do approximate inference over the joint state Z = (=1, (Zi,.... Zp), using an extended Kalman filter (EKF) that linearizes h by calculating its Jacobian at each step (Singhal & Wu, 1989 Puskorius & Feldkamp 2003). Next, we derive a variational approximation to the EKF that constrains the covariance matrix to be diagonal, and show how it extends the multiscale optimizer. Specifically, writing Wj and Wij as the current mean estimates of 0; and zzy (for weight and time scale i), the variational update each wij follows that of the multiscale optimizer, with additional machinery for determining decay rates based on т; and adapting learning rates based on the current prior variance (t).
In Section 5] we test our methods in online prediction and classification tasks with nonstationary distributions. In online learning, nonstationarity often manifests as poorer generalization perfor- mance on future data versus held-out data from within the training interval. Common solutions are to train on a window of fixed length (to exclude "stale" data) or to use stochastic gradient descent (SGD) with fixed learning rate and weight decay, which leads older observations to have less influ- ence (Ditzler et al., 2015). Here, we demonstrate that performance can be significantly improved by retaining all data and using learning model that accounts for the temporal structure of the en- vironment, We introduce nonstationarity in our simulations by varying the latent data-generating parameters according to 1/f noise. Thus important caveat is the task domains are matched to the Bayesian model. Notwithstanding, we test robustness by using a different set of timescales for task generation versus learning (Section 5.I). generative process that mismatches the NN archi- tecture (Section 5.2). and a construction of I/f noise that differs from the sum-of of-diffusion process the model assumes (Section [.3). Results show the Bayesian methods (KF and EKF) outperform windowing and online SGD, as well as a novel heuristic of training the network on all past data with gradients weighted by recency. We also find the variational approximation performs nearly as well as the full model (Section 5.1 and scales well to multilayer NN trained on real data (Section 5.3
MULTISCALE OPTIMIZER
Assume a statistical model (t) = ((x()), wt)) and loss function C(y. 3). where x(t) is the input on step i, (t) is the parameter estimate, (t) is the model output, and y(() is the target output. In
2

---

Published as a conference paper at ICLR 2023
Toy illustration of fast weights. A sin- gle weight w (blue) with constant input (r 1) predicts target signal T (black) with square loss C (T- The weight a sum of subweights Walow (yellow) and whast (red). Initial learning is rapid, due to Whatt Because of decay and the shared error signal, knowledge is gradually transferred to Walow while Wiast returns to zero. When the task switches (trial 151). Wistt enables rapid adaptation while long-term knowledge is preserved in Walow Thus the model recovers quickly on the second re- versal (compare blue curve beginning on trials I vs 156). The general multiscale optimizer extends this idea to an array of faster and slower weights.
Figure 1:
Weight / Target
0.5
-.5
-1
100
150
NN, (t) is the vector of current weights. (Under the Bayesian framing in Section 4 is the mean estimate of the optimal parameters 0.) For exposition, we assume the weights are updated by SGD, w(t + 1) = wt) a ((t))((((()) (t)),
(1)
and we henceforth abbreviate the gradient as (t) L. However, the following approach can be natu- rally composed with other optimizers, such as extensions of SGD Hebbian learning, by replacing
-wd L with the appropriate update term.
The multiscale optimizer is motivated by the assumption that, in online learning tasks, the true or optimal parameters change over time, on multiple timescales. Accordingly, it expands each weight into a sum of subweights, wy = Σ Wiy, each with a different learning rate a and decay rate % Here j indexes weights in w, and indexes timescales. The subweights evolve according to:
wiy(t + 1) = 7,4r (1) (1).
Each Wiy has characteristic timescale π == (og Note that (1) = Ou) so one think of the gradient for wy being apportioned among the subweights (with total learning rate a = Σ o;), or equivalently of each subweight following its own gradient.
2.1
FAST WRIGHTS
A potentially important special case of multiscale learning arises with two timescales, Wslow + Whatt- We assume "ssow = (no decay) and Orfsst > Oslow- Thus each Walow.j can be thought as the original weight, which is augmented by Wjatt.j- second channel between the same neurons that both learns and decays rapidly. The fast weight enables the system to adapt quickly to distribution shifts while resisting catastrophic forgetting (Figure D.
This model is conceptually similar to the fast weights approach of Ba et al. (2016) and Hinton & Plaut (1987). In that work, the weights are updated by different mechanism (Hebbian learning) than the primary weights, and they act as a memory of recent hidden states in a recurrent network. In the present conception, fast weights optimize the same loss as the primary weights, only with differ- ent temporal properties, and they act as memory for recent learning signals (e.g., loss gradients). Thus they are perhaps better suited for handling distribution shifts of the sort considered here.
EQUIVALENCE RESULTS
3.1 BENNA-FUSI SYNAPSE
Benna & Fusils (s (0166) model synapse is designed to capture how biochemical mechanisms in real synapses implement a cascading hierarchy of timescales, and has been adopted in ML for continual reinforcement learning (Kaplanis et al.) 2018 2019). We consider a single weight w a network, suppressing the index j. The Benna-Fusi model assumes that the information in is maintained in a dd hierarchy of variables Wi, WI,.... ....... un. each dynamically coupled to its immediate neighbors: un (1) 91 (u2(t) ui (1)) Owie
Ci(ui(t Ca(ul(t+ + ut())) - (1 un (f)) + 9k (1)+1(t) up (t))
(3)
(4)
3

---

Published as a conference paper at ICLR 2023
for 2 I l I n, with In. O. The external behavior of the synapse comes from in alone (i.e., w = (1), while U2m act as stores with progressively longer timescales.
This update rule can be rewritten as
(5) with transition matrix T determined by the coefficients in Equations 3 and and external signal d(t) defined by di (t) Dwott) and d2.n. 0. It can be shown that the transition matrix is diagonalizable, T = VAV-1, with eigenvalues Aa A < (see Appendix A). We can further enforce V1. 1, for purpose explained below. We refer to the eigenvectors (columns V.) as modes the system, because they are preserved over time up to a scalar. That is, the initial state is proportional mode i, then in the absence of external signal (d 0), the system will remain in that mode, decaying exponentially with rate factor A;:
u(t =
Te(t)) d(t),
(0) x V, = vy :((u(t)= (u
(6) linear combination of modes, w V. = Vw.
In general, any state be written uniquely Therefore, reparameterizing the model as w == V- u yields the simplified update equation:
w(t + 1) = Aw(() d(t)
Because is diagonal, there is no cross-talk between the modes, unlike in the original dynamics. Thus we have derived an instance of the multiscale opti- mizer, with subweights (t), decay rates X; and learning rates Ci [V- The assumption above, V1 implies w t1 Σ W/s so the models agree on the external behavior of the weight as whole. Figure iillustrate the translation between the two models.
where V-id(() = 1 --
3.2 MOMENTUM LEARNING
The standard rationale for momentum learning is smooth updates over time, so that oscillations along directions of high curvature cancel out while progress can be made in directions with consis- tent gradients (Rumelhart et al. 1986). To simplify notation, we again focus on a single weight w in the network, suppressing the index j. The momentum is defined as an exponentially filtered running average of gradients, with weight update determined by current momentum:
g(t + 1) 3g(t) (1 (1- [
(8)
w(t + w(() ng((l + 1).
This formulation is equivalent to one in which the update Aw(() = w(t + 1) w(() includes portion of the previous update: (w( = (c) + 3Aw(( - 1), with a = 1y(1 - -))
Paralleling the analysis in Section 3.1 we write the state of the momentum optimizer as [w,g]] and Equations X and
9 to obtain the update rule:
w (l +
1
n
w(t)
]
+
n -B) Dwtt) (1 B)
(10)
3
((
The transition matrix has eigenvectors O], with eigenvalue 1, and [1, 9p Now use this eigenbasis to define a reparameterization:
1-81T with eigenvalue B.
w
Welow
(11)
g
Wfast
Substitution into Equation 10
yields the reparameterized update rule:
Welow (++1
Wfsst (t+ ]
1 0 Watow (1)
]
dw,
(12)
0 8
Wfast (t)
7
Thus we recover the fast-weight optimizer, with decay Taat = and learning rates Oslow = 1) and
Ofast = - The negative fast learning rate is perhaps surprising but can be understood as follows: When Efast < 0. the subweights learn opposite directions, with the latent knowledge in Wslow overshooting the observable knowledge in w Walow Whast As Wfast decays toward 0, w catches up to Walow, so that the model appears to continue learning from past input, just as it would with momentum. This analysis highlights the contrasting rationales of these two methods: Learning multiple timescales is motivated by an expectation of positive autocorrelation in the environment, whereas momentum is effective at smoothing out negative autocorrelation in the gradient signal.
4

---

Published as a conference paper at ICLR 2023
0.35
0.15
Au,Vi
0.-
-21
616-
02
42
12
Figure 2: Translation between the model of Benna & Fusi (016)) and the multiscale optimizer works by decomposing the state of the former model into modes, independently. which correspond to subweights in the multiscale optimizer. A: All modes for a de- fault Benna-Fuss model with eight variables (n 8). B: An arbitrary initial state of the model. C: Unique cigen-decomposition of the state in Figure 2 Implied values of the corresponding multi- scale optimizer's subweights can be read off the values of the curves A 1. D: Decay of the individual modes or subweights for 1000 steps (with no external input) at rates given by their eigen- values. Reconstruction of the final state exactly matches the result iterating the Benna-Fut update (dotted arrow from Figure 3)) F: Decomposition of unit impulse to (e.g., loss gradient, shown as grey bar) as weighted sum of modes. Learning rates for the corresponding subweights, Aw; can be read off as the values of the curves l =1 (because Vix 1).
ingen-pateerns of activation that decay
4 BAYESIAN MULTISCALE OPTIMIZER
We turn now to normative analysis of learning at multiple timescales, based on Bayesian inference over 1/f noise. The Bayesian model introduced here assumes that the latent parameters O governing the observed data in some learning task vary over time according to 1/f noise. When the statistical model (ax, 0) linear 0, exact Bayesian inference possible with KF that maintains posterior over an expanded representation of f. When the model is nonlinear, approximate Bayesian inference achieved by an EKF that uses a linear approximation of h. We then show that a variational approximation of the KF EKF, in which the posterior covariance matrix is constrained to be diagonal, yields an extension of the multiscale optimizer that adapts its learning rates online by tracking uncertainty.
4.1 GENERATIVE MODEL FOR 1/f NOISE
Let (t) an decay rate and diffusion rate of, defined by the following stochastic differential equation:
process (i.c., diffusion with decay), with timescale or inverse
dæp -T, dt o, dW.
(13)
Here (!t) is standard Wiener process (Brownian motion). As a Gaussian process, 2; has kernel (t),(t) + s)] & e --ll/T,, implying exponentially decaying autocorrelations. However, a super- position of such processes at different timescales can have qualitatively different properties (Eliazar
5

---

Published as a conference paper at ICLR 2023
& Klafter 2009). In particular, consider
<(t) Σ (t),,
(14)
where - U² and = wー for chosen 1, and is an integer such that Tn, is very large. We show in Appendix that has power-law (i.e., long range) autocorrelations, (E((())((+++)) x |s| 11 for < Tn. and accordingly power spectrum that is well-approximatee by f/ for frequencies f > T Moreover, m independent copies of this process constitute mn-dimensional 1/f noise, due to the rotational invariance of multidimensional Ornstein-Ullenbeck processes.
This construction affords a flexible generative model of nonstationarity in a variety of online learn- ing domains, by applying it the latent parameters governing the relationships among observable variables. Assume we receive observations xtt), yt) that we wish to model with a statistical model h that is parameterized by 0 IRm
yt) (t),
(15)
For example, h may be a NN with weights 0, input æ, and target output y. The generative side of our Bayesian model posits latent variables 2; (i = 1,..., ni) such that each z; is an Ornstein- -ubeenn-h process in R with timescale Ti, and these processes sum to determine the original parameters:
(t)) = (t).
(16)
These assumptions imply that O follows 1/f process, and they entail an expanded state represen- tation, Z (Z1...., 2n) E R""", that enables efficient inference as described in Section 4.2
4.2 INFERENCE OVER f/f NOISE VIA EXTENDED KALMAN FILTER
We consider Bayesian methods that adopt the construction in Section 4.1 a generative model to account for nonstationarity. Equations 13 and 14 describe a linear dynamic system with state Z (Z1...., EE R". If & is directly observed at discrete intervals, then optimal Bayesian online prediction of each Et)) based on all preceding observations can be implemented by a KF over Z
(Kording et 2007 (see Appendix D)
We extend this approach to arbitrary statistical models with nonstationarity in their latent parame- ters, as in Equations [5] and [6] When h linear in 0 (and hence Z), such as in the regression task and 1-layer perceptron model in Section 5.1] exact inference is possible with standard KF (Appendix DD. For general h, such as multilayer NN, we use an EKF, The EKF makes local linear approximation of / based on its Jacobian, the matrix of gradients of predictions ÿ with spect to 0 (Appendix D. We use Ollivier's (2018) generalization of the EKF that replaces Gaussian observation noise with any exponential family p(ylğ). which is better suited for modeling discrete outcomes such as the classification tasks of Sections 5.2 and
4.3 VARIATIONAL APPROXIMATION
Finally, we derive a variational approximation of the EKF that extends the multiscale optimizer and affords efficient implementation in large NNs (Appendix F. As is standard, the EKF maintains an iterative prior over the latent state based on all previous observations:
p( (1) ct, yes)) ~ (l) S(())
(17)
The mean, (t), is the vector of current subweight estimates in the network, while S(t) captures their joint uncertainty and hence determines updates (as a preconditioner on the gradient).
We use variational inference to approximate the distribution in Equation by one in which S(1) m Š((), where átt) is constrained to be diagonal, written as Š(() diag (())) This reduces the complexity from O(m²n²) to O (mmkk (the size of the Jacobian, where k is the size of the output layer). The simplest case is a KF that tracks a single 1/f variable, with no inputs or latent variables. That is, g(t) E(t) e R² as in Equation 14 Appendix F.1 derives the variational update rule
wi(t + 1) e w/t) (t) (vi(t) (t)),
(18)

---

Published as a conference paper at ICLR 2023
where ye(t) (t()) -dwL (i.e., square loss), and the learning rates are given by
a;(t) e e11/) st(t
(19)
For the EKF with a general nonlinear model hx 0), slight extension of the variational approxi- mation, derived Appendix F3 provides the following update:
Wiy (t++)) e w (t) E 1/ s (t)0
(20)
Here diag(s")) is the diagonal variational approximation of the posterior variance after observing y(),, and É is the negative loglikelihood in the EKF's Gaussian approximate output distribution. Importantly, the update rule for the variance uses the diagonal of the precision matrix but can be calculated without matrix inversion, which is relevant to scaling up to large networks.
Thus the variational method amounts decomposing every weight in the network sum of sub- weights, wy = E, Wij. that learn independently according to their individual gradients, with decay rates and learning rates coming from (t). This is special case within the family of multiscale optimizers, with additional machinery adapt the learning rates based on current uncertainty.
SIMULATION TESTS
5.1 REGRESSION TASK
As simple demonstration, we created an online linear regression task with 10 features (including bias term). in which the true weights 33 varied over time according to 1/f noise using Equation 16 The outcome was generated The corresponding predictive model distinguish weight estimates (w) from true parameters (3).. We model the data using the perceptron and compare methods for optimization, using square loss, L = (y-g))
x B (no noise term was needed because ,3 is inherently noisy). one-layer perceptron, which we write as ý = w to
We tested two baseline training methods, representing common heuristi practices with nonstationary Parisi et al. 2019). First, we tested a batch learning method that uses fixed memory horizon H. To produce prediction step t, the batch learner fits the perceptron to trials t t- through t- F shows performance U-shaped: accuracy suffers with short horizons because of sampling error, but it also suffers from longer horizons because older observations less valid. Second, we tested SGD, in which the weights are updated once after each step 4, based on xxt) and y(t). Figure 3B shows performance is best with an intermediate learning rate, which roughly corresponds to assuming the environment changes on a single characteristic timescale (see Appendix[]) C).
data (Ditzler et al. 2015; Parisi
As applied to the perceptron, the Bayesian 1/f model described in Section 4.2 decomposes the weight for each feature into subweights, Wy = Σ, Wijs and tracks the wij jointly with KF (generalizing Dayan & Kakade 2000). The subweights combine to predict the outcome on each step: 5(1) Eij xx (z(())un(()) Relative blind prediction (guessing the overall mean of y on
A
Mean squared (E)
(2)
Meen squared
Mean squared
22
8
Batch
SGD
Discounting
Variational
Bayes
130
200
0.86 Learning rate (a)
Memory
horizon (Ef)
Figure 3: Regression task with 1/f dynamics. A: Batch optimization.
B: Stochastic gradient de-
scent. C: Model comparison.

---

Published as a conference paper at ICLR 2023
every trial), this exact Bayesian solution explains 54.0% more variance in the outcome than the best batch learner, and 25.7% more than the best parameterization of SGD (Figure 3C). Moreover, the variational model that constrains the KF covariance matrix to be diagonal (see Appendix (F2) performs nearly (96.9%) as well as the full Bayesian model.
Finally, we tested discounting method, similar to windowed batch optimization except all past trials were used for training, discounted by a function of lag. Because the 1/f environment has power- law correlations, we weighted each observation i k by k and optimized a. This method also significantly outperforms windowed batch and SGD, showing that accounting for an environment's autocorrelation function can achieve much of the advantage the Bayesian approach. Nevertheless, the full KF and variational methods still outperform discounting, by 5.5% and 2.2% respectively.
5.2 LINEAR CLASSIFICATION TASK
Next, we investigated an online 10-way classification task with 10 features (including a bias term). The data were generated by first sampling the class, y(t) softmax(e(t))) and then sampling the feature vector, (t)) ~ N (Pu(()) I). The logits ej and the feature-class means jag were inde- pendently sampled from 1/f processes (using Equation 16]. so that there was nonstationarity in both p(y) and p(xlly). Scaling of e and j was chosen equate to maximum possible performance based on perfect knowledge of either one alone (both yielding average loss L ~ 1). For the predictive model, we used one-layer perceptron with a softmax output layer, ÿ softmax( W), where W is a matrix of learnable feature-class weights. We assumed cross-entropy loss, L(y, 9) - log Dy-
The batch method trained the network on trials t H through t t- until convergence before pre- dicting (t). Weight decay was included for regularization, optimized for each value of H. Figure 4A shows a U-shaped pattern of performance, reflecting the tradeoff between sampling error and stale data. We also used weight decay with SGD, optimized for each learning rate. Figure F again shows a U-shaped pattern performance. The variational EKF for this task and model is derived in Appendix F3 We applied & regularization to the prior each time step, on par with the SGD and batch optimization methods. Figure 4C shows the variational method outperforms the other two.
5.3 MNIST CLASSIFICATION
Finally, we tested our methods on classifying a stream of handwritten MNIST digits (LeCun et al. 2010). We created nonstationary online learning task by sampling an example from the 10- way MNIST training set on each time step according to class logits that followed 1/f noise (Figure 5 For predictive model, we used a convolutional neural network (CNN) with two convolution layers followed by two dense layers, with 824458 parameters. This experiment provides a more stringent test of the present methods in several ways. First, it tests whether the EKF's linear approximation and the variational diagonal- -variance approximation perform well on multilayer NN, and whether the algorithm is efficient with a moderately large number of parameters. Second, the predictive model the optimizer is doing inference over is unrelated to how the images and labels were generated. Third, the sampling procedure for 1/f noise did not accord with the additive generative process
A
C
Mean loss (C)
Monn. loss (Z)
18
Mean
1.5
100
14
02
Batch
SGD
Variational
Memory borlson (H)
Learning rate (a)
Figure 4: Synthetic classification task with 1/f dynamics. Loss is negative boglliklihood. A:
Windowed batch optimization. B: Stochastic gradient descent. C: Model comparison.
8

---

Published as a conference paper at ICLR 2023
A
iid
9.3
Optimizer
50D
Class frequencies per block
Variecion EXT
100
Cumulativv Error Rate
0.20
Blocks
1/f
Alarks 10 1/f
18 500
Blocks 100
300
lld
Training Sequence
1/
Figure 5: A: Sample frequencies in blocks of consecutive time steps for all 10 MNIST classes (indicated with a unique color per class). The 1/ʃ sequence exhibits long-range autocorrelations, with nearly the same pattern over blocks of 10 or 100. B: Model performance over sequence of 10000 examples (1/6 MNIST training set).
assumed by the Bayesian model (Section 4.1 but instead used spectral procedure described in Appendix G
The variational EKF was compared to SGD with momentum, in both 1/f and standard iid environ- ments. Hyperparameters for both methods (noise variance for EKF, learning and momentum rates for SGD) were optimized separately for the two environments. Because the models learn quickly, we evaluated them (within each replication) on a sequence comprising only a random subset of the MNIST training set. Performance was measured by top-1 error rate. This should be interpreted as generalization (i.e., test) performance, because each item was observed only once. The variational EKF outperforms SGD by 1.4% in the iid environment, because its tracking of uncertainty enables more effective gradient steps. However, its advantage jumps to 2.4% in the 1/ʃ environment, show- ing once again that is better able leverage dynamics multiple scales.
CONCLUSIONS
Our analytic and simulation results demonstrate how online learning performance in nonstationary environments can be improved by incorporating model of temporal structure. The Bayesian f/f model amounts distributing knowledge across multiple timescales, and the variational EKF en- ables approximate implementation in neural network using subweights with different learning and decay rates. The variational EKF extends the multiscale optimizer, which is closely related to pre- vious models in both neuroscience and MI,, and in some cases is equivalent to them despite being simpler in having no coupling between timescales.
We have implemented the variational EKF optimizer in JAX in format compatible with Optax. In the MNIST simulations of Section 5.3, we find our optimizer code (with timescales) is actually 1.6% faster than Optax's 's off-the-shelf SGD, in compute time per example. Note also that the mul- tiplexing of subweights is not expensive relative to current optimizers (e.g., Adam; Kingma & Ba 2013). which also store multiple variables for each weight.
In sum, the multiscale optimizer and variational EKF enjoy a combination of normative, heuristic, and biological justification, good performance, and computational efficiency. Our ongoing work aims to extend the theory in several ways. Chang et al. (2022) compare the present variational method the fully-decoupled EKF of Puskorius & Feldkamp (2003). Another possible variational method is to assume a block-diagonal matrix that maintains covariance information only between subweights (timescales) within each weight, so that computational complexity still scales linearly with network size. Finally, the present method is not limited to 1/f noise but generalizes to other power laws (c.g.. 1/ by appropriate choice of the timescales т; and noise variances o, in the generative model (see Appendix B). If, for example, data or theory is available bearing on the power spectrum the dynamics given domain, the optimizer could be tuned accordingly.

---

Published as a conference paper at ICLR 2023
REFERENCES
John R. Anderson. Acquisition of cognitive skill. Psychological Review, 89:369-400 0066994-6,-66 1982.
John R. Anderson and Lael J. Schooler. Reflections of the environment in memory. Psychological
Science, ((6)339-408,, 1991.
Jimmy Ba, Geoffrey E. Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin lonescu. Using fast weights attend the recent past. Advances neural information processing systems, 29, 2016.
Mareus K. Benna and Stefano Fusi. Computational principles synaptic memory consolidation. Nature Neuroscience, 19, 2016. doi: 10.1038/nn.4401. URLLattss://ei..or//11...
nn 4401
Nicholas J. Cepeda, Edward Vul, Doug Rohrer, John T. Wixted, and Harold Pashler. Spacing effects in learning: temporal ridgeline of optimal retention. Psychological Science, 19(11):105--110, 2008.
Peter G. Chang, Matt Jones, and Kevin P. Murphy. On diagonal approximations the extended Kalman filter for online training of Bayesian neural networks. In 14th Asian Conference on Machine Learning (ACML), Workshop Continual Lifelong Learning, 2022.
Peter Dayan and Sham Kakade. Explaining away in weight space. Advances in neural information processing systems, 13, 2000.
Gregory Ditzler, Manuel Roveri, Cesare Alippi, and Robi Polikar. Learning in nonstationary vironments: A survey. IEEE Computational Intelligence Magazine, 10(4): 2015. doi:
10.1109/MC1.22152247113
Iddo Eliazar and Joseph Klafter, A unified and universal explanation for lévy laws and 1/f noises. Proceedings of the National Academy of Sciences, (06(30):22511-1)) 2009.
Geoffrey E. Hinton and David C. Plaut. Using fast weights deblur old memories. In Proceedings of the 9th annual conference of the cognitive science society, pp. 177-186, 1987.
Mare W. Howard and Michael J. Kahana. A distributed representation of temporal context. Journal of mathematical psychology, 46(3):269-299) 2002.
Christos Kaplanis, Murray Shanahan, and Claudia Clopath. Continual reinforcement learning with complex synapses. In Proceedings of the 35th International Conference on Machine Learning, Stockholm, Sweden, 2018. PMLR 80.
Christos Kaplanis, Murray Shanahan, and Claudia Clopath. Policy consolidation for continual rein- forcement learning. In Proceedings of the 36th International Conference on Machine Learning, Long Beach, CA, 2019. PMLR 97.
Marvin S. Keshner. 1/f noise. Proceedings of the IEEE, 70(3):112211 1982.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Konrad P. Kording, Joshua Tenenbaum, and Reza Shadmehr. The dynamics of memory consequence of optimal adaptation to a changing body. Nature Neuroscience, 10(6):779-786, 2007. doi: 10. (0388/n1991..... URL https: http://oo..oogg////// 1038/ (033//n1100) nn1901
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available:
wtw.l//cwww.ooo/ee////msssssss 2010.
Michael C. Mozer, Harold Pashler, Nicholas Cepeda, Robert V. Lindsey, and Ed Vul. Predicting the optimal spacing of study: A multiscale context model of memory. Advances in neural information processing systems, 22, 2009.
Yann Ollivier. Online natural gradient as a kalman filter. Electronic Journal of Statistics, 12(2): 2930-2961, 2018.
10

---

Published as a conference paper at ICLR 2023
Genman I. Parisi, Ronald Kemker, Jose L. Part, Christopher Kanan, and Stefan Wermter. Con- tinual lifelong learning with neural networks: A review. Neural Networks, 113:54-71, 2019. ISSN 0893-6080. doi: http://dw..ao/////lo////eeeeeoo0
URL https: :ttps://wwww
sciencedi rect. com sc http:////c....aeaaaaaa.ccce///// 9300231
Gintaras V Puskorius and Lee A Feldkamp. Parameter-based kalman filter training: Theory and implementation. Simon Haykin (ed.), Kalman Filtering and Neural Networks, pp. 23-67.
John Wiley & Sons, Inc., 2003. URL https: http::://on.neeelbraayy wiley. com doi/IC.
IDG2/0471221546..
Ning Qian. the momentum term in gradient descent learning algorithms. Neural networks, 12 (1):145-151, 1999.
David E. Rumelhart, Geoffrey E. Hinton, and Ronald Williams. Learning internal representa- tions by error propagation. In D.E. Rumelhart and J.L. McClelland (eds.), Parallel distributed processing, Vol. 1, pp. 318-362. MIT Press, Cambridge, MA, 1986.
Konstantin Rusch, Siddhartha Mishra, N Benjamin Erichson, and Michael W Mahoney. Long expressive memory for sequence modeling. In International Conference on Learning Represen- tutions (ICLR), 2022.
Bernhard Schölkopf. Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning. arXiv preprint arXiv: 2012
Sharad Singhal and Lance Wu. Training multilayer perceptrons with the extended kalman algo- rithm. volume 1, 1989. URL httpsss/pproeeeednnss.s.
hwwp:is.eerrppp.p httpppsccpppppppppppp cctppppr//p9cc Paperr paf
JER Staddon, IM Chelaru, and Higa. Habituation, memory and the brain: The dynamics of interval timing. Behavioural Processes, 57(2-3):71-88, 2002.
Richard S Sutton. Gain adaptation beats least squares. In Proceedings of the 7th Yale workshop adaptive and learning systems, volume 161, pp. 166, 1992.
Matthew H. Wilder, Matt Jones, Alaa A. Ahmed, Tim Curran, and Michael C. Mozer. The persistent impact of incidental experience. Psychonomic bulletin & review, 20(6):1221-1233 2013.
John Wixted and Ebbe Ebbesen. Genuine power curves in forgetting: A quantitative analysis of individual subject forgetting functions. Memory & Cognition, 25(5):731-739, 1997.

---

Published as a conference paper at ICLR 2023
A DIAGONALIZATION OF BENNA-FUSI TRANSITION MATRIX
This section provides details on diagonalizing the transition matrix of the update rule for the Benna & Fusi ( model, and the corresponding reparameletization in terms of the eigenbasis.
First, reparameterize the Benna-Fusi model so that its transition matrix is symmetric, as follows. Recursively define
k = 1 1<k Sn
by
bE-1Ck Gh-1
(21)
and write the state of the Benna-Fusi synapse as
= (Vbbuu))sss
(22)
The update becomes
@1 + 1) = rm(()++ d(1)
(23)
with
-k- + gk.
KA =
(24)
Ck
T - =
9h-1
(25)
VO2-1C4
Symmetry of T implies it has orthonormal eigenbasis, {sb1,-.., 3, with corresponding eigen-
values , An. Because the scaling of cigenvectors arbitrary, we can enforce 1ú = 1 for all i.
To translate the eigenbasis back to u, define = diaggvVia (Vb) so that & Bu and = B-IPB.
is then easily verified that
V. B s
(26)
an eigenvector of T with eigenvalue AG. Therefore T VAV-1, as claimed in the main text. Note that the choice th,1 1 entails V1. 1 (because b1 1). Finally, v [rb..... v/n] is
invertible because it comprises basis, and therefore so is V B B-l. Therefore the reparameter- ization u + y well-defined.
B GENERATIVE MODEL FOR 1/f NOISE
Consider single Ornstein- Uhlenbech (OU) process =(t) described by Equation 1 with 1. The covariance function of 2 is given by
I'
Blz(t)a(t + s)]
(27)
-
(28)
Note that this expression decays exponentially with the lag s, yielding short-range autocorrelations. The power spectrum of 2, function of frequency f. is the Fourier transform of the covariance function:
Pa( = / _eee - 2rift
(29)
(30)
7-2 + (2x
where 2n f is angular frequency.
To define a generative model of 1/ʃ noise, we define a continuous mixture model over a family of OU processes (zr)errTT for some large T:
E(t) 1' 27 2,)dd.
(31)
12

---

Published as a conference paper at ICLR 2023
A
B
10"
10
Power
10
10
8000
10000
10"1
10
10"
10"1
Time Step
Frequency
Figure 6: A: Construction of 1/ʃ noise (black trajectory top) as of OU processes with different time constants (colored trajectories). Vertical offsets are applied as visual aid for discrim- inating the curves. B: The power spectrum of the aggregate trajectory in log-log coordinates. The red line has unit slope.
The power spectrum of & is then mixture over the component spectra P..:
(
4T -2 P2, (f)dr
(32)
T
-
dr
(33)
30
T-2 + (2
2
tan 1 fT)
(34)
which is approximately 1/f f 1/T, i.. for all but very low frequencies.
We next define a discrete approximation the continuum mixture model in Equation 31
E(t) Σ (t),
(35)
where 2; has timescale 7; and scaling parameter og, To approximate f/f spectrum, (f -.)- should scale as 4r; the squared weight density in Equation 31 (because power is additive and proportional to g²). For example, the could be arithmetically spaced with 0;x7 -1 Instead, we assume geometric spacing, with components defined by т, v² and a; 2pr, 1/2 where p² is hyperparameter determining overall steady- variance. Figure ililusstraess this construction and exemplifies the accuracy of the discrete approximation.
To define a generative model of 1/ʃ noise in IR",, we assume an independent copy of this process for each dimension. The resulting joint process is rotationally symmetric, property inherited from OU processes. That is, any linear combination Σ, esE, is a ld 1/f process.
C KALMAN FILTER FOR A SINGLE TIMESCALE
Start as in Appendix with single OU process with and assume observed unit intervals (EN) with Gaussian observation noise of variance 17². Assume conjugate iterative prior:
=(t) N (w(t),
(36)
The posterior after each observation is
(1)v(((
=((1)²² (t) +
a(t)uss N
(37)
s²()) + n?
13

---

Published as a conference paper at ICLR 2023
Evolving the process to time t + amounts decay by e- and accumulation of new noise. At each time E [t, + 1], variance from the noise appearing at time t' (i.e., from dt(?') in Equation
13 decays by a factor e 2(t+1111)// by time + 1. Therefore the total accumulated variance is:
--- 1/ T (1-e
(38)
2
Therefore the prior for the next observation is
z(t + 1)lysr N 1/r (t) ( (t)u(() (())(t)
-2/7 (²tr)
-
T (1-c-11))
{(1
(39)
(11)) +12
2
=N (w(t st++
(40)
We thus obtain the recursion
17 (() (t)v(() w(t + 1)= =
(41)
(t))+
= 1/T (w(t) + (t) w(t)))
(42)
This is temporal-difference learning rule, gradient descent on squared error (v(() w()))
with learning rate at()=
20) 20 This connection between temporal-difference learning and the
Kalman filter is well known (Sutton 1992).
The steady state for s² is given by the solution to s²(t) (² + 1). quadratic with one positive root. If (0) is initialized to this value. then a(t) will be constant. Note that this algorithm (and the extension to 1/ʃ noise in Appendix D generalizes to irregularly spaced observations, in which case one can derive how the optimal learning rate should vary on each step.
D KALMAN FILTER OVER 1/ʃ NOISE
Assume we receive sequence of observations unit time intervals, y(t for N, and we want to do online prediction under the assumption that is a 1/f noise process. Then we can make the generative assumption yt) E(t) with & defined as at the end of Appendix B) For simplicity and in contrast to Appendix we assume no observation noise, because the shortest timescales (21, etc.) already play this role.
Because y({ is conditionally independent of the history given (and indeed is fully determined by) the joint state (t) = (t), zn (f)), it suffices to compute the posterior for the latter. Write the iterative prior for Z as
z()) ~ N (I),
(43)
implying optimal (maximum- -ikelioo or least-squares) prediction of j(t) wt() The posterior after observing y(t) is the intersection of the prior with the hyperplane Σ, (t) = =t)::
S(t) ZUlyer ~ N (0(() + 1 (T)) (v(1) - i(1), (PS(())-P))) -1 P).
(44)
Here, 1 is the vector with all elements equal to 1, and = I 111
T is the orthogonal projector.
The prior for the next time step is then obtained by applying decay and adding variance from the noise accumulated over the intervening interval. Define D as the diagonal matrix of decay factors, -1/f,, and N as the diagonal matrix of added noise 2p e -2/r.), obtained by multiplying the RHS of Equation 38 by of = 4p²r, T3 (from Appendix the exact Bayesian model:
B Then we have the update equations for
( 11+ = D ((w(t)= + 1 S(t) (u(t) -0(1))
S(1)
(45)
and
-p 1P PD - N.
(46)
14

---

Published as a conference paper at ICLR 2023
Next, consider the perceptron in Section 5.1 where y 0. The 1/f generative model assumes 8, = Σ; žij for each J. where indexes timescales and indexes features, and 2y has timescale Tą and scaling parameter o, defined as above. The latent state is described by Z (2gg) We treat ij single composite index, that Z is vector. As above, write the iterative prior
( x ce Betr N (w(t),
(47)
Let X be multiplexed copy of the input æ, so that X; = z.. Assuming square loss, the optimal prediction for y(t) 0(t) = (t) (t). The posterior after observing y(t) is the intersection of the prior with the hyperplane X Z(t) y()):
(t) TSt DE.. Ust N ((())+ + X(I) (()) (1(1) - ;()),
((t)
S(t) + ((t)x((tT)) Px()),
(48)
where II-IXXX / X is the projector orthogonal X.
Generalizing the definitions above, let D and N be the diagonal matrices of decay factors and noise accumulation, Dij.is = e 1/%% and Nej.i = 2p² -2/m).. Applying these to Equation 48) to
obtain the prior for the next time step yields the update equations for the exact Bayesian model of the regression task:
((t 1) D w(f)+ ((I) X(t) (t)) (u(t) (((t))((0)) ((t))
(49)
D (t)P(0)) + ((()X(()))) P
(50)
Equation 49 exemplifies how the variance matrix, S(t), can be thought of as defining a precondi- tioner of the gradient, X (t)()(t) (())).
E EXTENDED KALMAN FILTER
Given a nonlinear model f(xx the 1/f EKF posits that the optimal parameters follow 1/f dy- namics according Equation 16 with expanded latent state (Z1,..., 2n,) T. It is convenient introduce the expanded model defined by h(., Z) h(x, Σ. 2,). The 1/f EKF maintains an iterative prior over as in Equation 47 and updates that prior by linearizing h about Z = (1):
h(x(t), Z) xe h((((t), (t)) - Jz (Z - w()).
(51)
Here, Jn = 0g Jn.
az the Jacobian matrix of h, evaluated at Z (t). Note that Jj. is just n copies of
Following Ollivier 2018). we assume the observation is governed by some exponential family ()))(())) with vector of sufficient statistics T The model's output ş = h( w) is taken to encode the predicted mean parameter of that family: ý Ey-P(-)m(()) [T(y)] (this can be read as definition of the mapping sufficient statistics as a Gaussian,
(ū)). It then approximates the conditional distribution of the
P (T())) R N (g. Ry),
(52)
where Rg = Var (T(y)l)) is the conditional variance.
For example, when a classification model as Section 5.2 or 5.3 the output of the network is vector ý of class probabilities, and the sufficient statistics T Tvy) are one-hot vector. For numerical stability, we exclude the final element of ý and T which are determined by the other elements. The conditional outcome variance is given by
10((1
9.)
i=
Ryk,
(53)
+ j.
15

---

Published as a conference paper at ICLR 2023
Under the approximations of Equations and 52
the posterior is given by the standard KF formula:
Z(t ysi ~ N (w(())+ S()J) (w(()) (J,S() Rg(o)) (T(x(t) 9()),
st) - S())_) (T)))) + Rgy))) J Sit)(t)
(54)
Applying decay (D) and accumulated noise (NN) as in Appendix D to obtain the prior for the next time step yields the update equations for the 1/f EKF:
wit + 1) D (w(t) + 5 (J,S(()J + Roco) (T(v()))- - -(())
(55)
S + 1) D (5() - S()))) (J,S())) + ( D ++N. N.
(56)
Although we apply the EKF to feedforward NNs in this paper, we note that the approach naturally generalizes to other probabilistic causal models relating the observed variables. Thus it offers means to model nonstationarity that covers all of the traditional forms of distribution shift, following the causal framework of (Schölkopf et al. (2012). Consider classification task, with input features xxt) E R' and class labels y(t) E Zu- Under a generative causal model where æ depends on y, label shift can be modeled by // noise in the distribution p(y), for example y ~ softmax (e) with & given by Equation 16 Manifestation shift can be modeled by 1/ʃ noise in the distributions p(ælyy) for example (t) lv(t) ~ ~ (Mu(e): Σ) with My given by Equation 16 for each y. Likewise, under a discriminative causal model where y depends on æ, covariate shift can be modeled by //ʃ noise in p(ær, and concept shift can be modeled by 1/ʃ noise in p(y)
F VARIATIONAL APPROXIMATION
This section derives variational approximations of the KF and EKF models, with S(() 2 5() where (tt) == diag(s" s² (t)). Thus the resulting algorithms need only track the individual variance terms in (t) rather than the full covariance matrix. Moreover, the update equations (63 64 68 69 74 77
all avoid matrix inversion-ene even though this might initially seem necessary from Equation 57
property that may be relevant for efficient scaling.
F.1 UNCUED INFERENCE
We begin with the simple Bayesian f/f model in Equations and 46 where there are predictors and the model merely tracks an observable yt). Given an arbitrary Gaussian distribution, the varia- tional approximation (in the sense of minimizing Kullback-Leibler divergence) by another Gaussian with diagonal covariance matrix obtained by taking the diagonal of the original distribution's precision matrix. Thus in the present case we have
$7 3 +
((++ (
(57)
where S(1+1) is given by Equation 46 with S --1(1) replaced by Š-1(() (the inductive assumption).
To calculate 5-1(1 + 1) under this assumption, first observe the identity
s282
P diag (s~3) P APP+11 -P111)-1 p= diag (s²)
(s
(58)
Σ5
Therefore Equation 46 becomes
S(t ((++10 (t)+1)= D (D)(()))))
(59)
(t)((
bb T
= diag(a)
(60)
where diag(a) = D)(()))) + N, , b b= (t) and c = Ex (t). We next use the identity
bb -
1
-1/2 bb) 1a-1/2
(daag (a)
diag
I
diag
(61)
Va
Va
c
Σ
16

---

Published as a conference paper at ICLR 2023
implying the diagonal elements are
diag(()) bb T
82
+
d,
(62)
Σx 0,
Combining Equations 57
60 and 62 gives the final form of the variational update:
(s((t)e 1/7, 4p² sinh 1 5
)
2
(t(+
(63)
(t(() 4p2 sinh + 82 st())
with
(t) sinh 1
Ω = Σ (t) + 4p² sinh
75
(64)
This update the variance converges exponentially to a unique fixed point. Numerical simula- tions confirm that, in this limit, sf is larger for smaller т;, meaning faster learning rates for shorter timescales. If the prior variances are initialized at the fixed point then they are constant throughout learning.
By substituting the diagonal matrix S(t) for St)), the update for the mean in Equation 45 simplifies to
(t (t + 1) 1/ (t) (t) (i(t) v())
(65)
where j(t) v(t is the loss gradient (assuming square loss), and the learning rates given by
ar (1) e 1/
(66)
Σ, (1)
That is, the subweights learn independently according to their gradients, with different decay rates and learning rates. Thus we have recovered an extension of the multiscale optimizer, with an addi- tional mechanism that adapts the learning rates on each time step (via 8).
F.2 KALMAN FILTER
To derive the variational KF for the regression task, we apply the analysis of Section F.] to the KF
update Equations 49 and 50
Paralleling the derivation of Equation [9) the variance update in Equation 50 (i.e., before applying the diagonal variational approximation) can be written
(D(x(1) (D()) s()))
1(1+1) D
(67)
(1)
Paralleling the derivation of Equations 63 and 64 the variational update comes out to
(1)e 1/T. 40² sinh ) n
(22 + 4p2 sinh +, Ω + +0++2((08) 27 11
(68)
with
Ω 4p (1) (2))((()) sinh +
(69)
1/ (t) + 4p² sinh 1
By substituting the diagonal matrix (1) for S(1), the mean update in Equation 49 simplifies to wty tt 1) //T (t) (ty (0(t) v(t))
(70)
where (f(1) - v()) the loss gradient for wij. and the learning rates are given by
1/
ai Eny (t)s,
(71)
Thus we have recovered an extension of the multiscale optimizer, with an additional mechanism that adapts the learning rates on each time step (via and s).
17

---

Published as a conference paper at ICLR 2023
F.3 EXTENDED KALMAN FILTER
To derive a closed-form variational approximation for the general EKF, such as for the classification models in Sections 5.2 and 5.3 turns that we need apply the variational approximation to
the posterior in Equation 54 rather than to the iterative prior in Equation 56 Using Woodbury's identity, the posterior variance can be rewritten as
S(I) -(1))) (JAS()) + Roo) JR5(() R +
(72)
The form on the RHS is convenient because it is in terms of precision, allowing read off the diagonal entries directly. That is, the variational approximation for the posterior variance is diag(s" with
=)= (1 R J1 41.
+ +s²)²(t)
(73)
Here was have used the inductive assumption S(t) = diag(s?(t)) Applying the transition from posterior on step to prior on step L + 1, we obtain the variance update for the variational model, replacing Equation 56
s (t + = (()) + NNy.y.
(74)
Woodbury's 's identity also enables the mean update from Equation 55 to be rewritten, as
(tt + 1) = D [J Ryln Jn + +--(()) -2 (s- (s-1(1) + J Ry (T((()) - D(t) + J.w(()))
(75)
Substituting the gradient of the EKF's approximate likelihood (denoted L) from Equation 52
Own J R (ū(() T(v())),
(76)
yields
((1) wit + 1)= ((()) (J/ R RO) + + (1) 2w
(77)
The preconditioner on the gradient here is the posterior variance (see Equation 72). which we have approximated di)g(s"?(t). Although not necessarily entailed by the variational approximation, we can consider applying the same approximation in updating the mean. This yields
Wiy (1 + 1) - w/T (1) - e (1)))-11)3((()00) (100
(78)
which once again extends the multiscale optimizer by adapting its learning rate to current uncer- tainty.
G IMPLEMENTATION DETAILS
were sampled using the generative model in Appendix That is, the data-generating process matched the generative assumptions of the Bayesian model in both of these cases. We used 20 timescales, geometrically spaced from T1 to T20 1000, as illustrated in Figure [6A. Each component OU process was run for 10r, burn-in steps to ensure stationarity. The regression task was run for 10k trials, and the linear classification task for 1000 trials. The 1/f power spectrum was confirmed with log-log plot (see Figure
Latent parameters for the synthetic tasks in Sections 5.1 and 5.2
In the regression task of Section [.II the first feature features were sampled independently as (0 I) on each time step.
constant bias term, I1 1. The other 9
For the linear classification task of Section 5.2 class logits were sampled from 1// processes and then multiplied by 0.886 before entering into softmax to determine class probabilities. Feature- class means were fixed at 1 for the first feature (i.e., bias term) and were sampled from mutually independent 1/f processes for features 2-10, multiplied by 0.224. These scaling factors were chosen so that perfect knowledge of either the prior probabilities the feature-class means every trial
18

---

Published as a conference paper at ICLR 2023
would yield ideal-observer performance of L Rs 1. Perfect knowledge of both would yield L ~ 0.35. These were merely guidelines for equating prior and likelihood information, as perfect knowledge of either source of information is not possible even with optimal model of the dynamics.
In both Sections 5.1 3.2 all Bayesian and variational models assumed 10 timescales, geometri- T10 800. This deliberate deviation from the data-generating process
cally spaced from (see the 20 timescales listed above) provided a mild test of robustness, specifically the hypothesis that it the aggregate f/ʃ character of the environment and of the model that matters, not the choice of component timescales used to approximate that character.
For the MNIST classification task in Section 5.3 the class on each time step sampled by softmax applied to class logits that varied across steps according to 1/f noise. The item on that time step was than sampled without replacement from all members of that class in the MNIST training set. The logit sequences were generated as follows. First, for each class, we sampled Standard Gaussian random vector of length equal to the total number of time steps (10k). Then we applied discrete Fourier transform, multiplied the result by f-1/2, and finally applied the inverse Fourier transform. Thus the resulting sequence had 1/f power spectrum. For the iid environment, the logits were sampled as white noise (constant power spectrum), by sampling Standard Gaussian vector as above and using it unaltered.
19


Reviews:
[{"review_id": "qRajqbYZHf", "reviewer": "Reviewer_5UF5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a gradient-based learning scheme that deals with changing parameters during learning. For this, the authors give a context of their algorithm to the classical momentum learning strategy and a multi-timescale model from neuroscience. The algorithm is derived using a Bayesian inference framework, where the posterior mean recovers the newly proposed learning method. Here, a continuous-time linear state space model is used to model the time-varying parameters. For the case of a supervised learning problem, the authors consider a nonlinear model for the time-varying parameters.\nBayesian inference is carried out by the authors by a variational approximation of the extended Kalman filter algorithm. The learning algorithm is tested on two synthetic problems, with time-varying parameters. For both linear regression and classification problems, the proposed algorithm leads to a smaller average loss. ", "summary_of_the_review": "The paper presents an interesting topic. The resulting algorithm is easy to implement and is probably useful for learning problems with time-varying parameters. The derivations are not very clear to me and the paper, but the results look reasonable. An application to a real-world learning problem would have been very helpful to show its applicability. ", "strength_and_weaknesses": "Overall, I think the paper presents an interesting new algorithm, for a problem that is probably very relevant in many applications. I am not very used to the style of writing, where the authors are probably coming a bit more from the direction of neuroscience. For me, the paper was pretty hard to follow. However, the results look reasonable even though the derivations are derived in a bit nonstandard way. The presented two empirical case studies are fair and are showing promising results. However, I would have hoped for an application where this strategy is tested for a real-world problem, to show its applicability.\n\nA sore point for me is the derivation using the power spectral analysis. I think this paper would have been much easier to understand and interpret would it had been derived using standard Bayesian inference techniques. For example, the definition of the $1/f$ noise process is very strange to me. To illustrate my confusion: equation (13) is a linear stochastic differential equation, for which the time-point-wise marginal densities are distributed multivariate Gaussian, see, e.g., [1]. Since the sum of the components of a multivariate Gaussian random variable in equation (14) is still Gaussian, see, e.g., https://web.ipac.caltech.edu/staff/fmasci/home/astro_refs/SumOfCorrelatedRVs.pdf , it implies that $\\xi(t)$ is a Gaussian random variable, where the mean and variance parameter can be computed in closed form. Maybe the authors can explain, why the $1/f$ definition is needed.\n\nStrengths:\n- Interesting Setup\n- Relation to well-known algorithms is given\n- A new learning algorithm, which is easy to implement\n- Convincing empirical results for some synthetic problems under the modeling assumption\n\nWeaknesses:\n- The derivations are hard to follow\n- The $1/f$ noise definition is not clear to me\n- No real-world application given in the paper\n\n[1] Särkkä, Simo, and Arno Solin. Applied stochastic differential equations. Vol. 10. Cambridge University Press, 2019.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "k7eNIrzutS", "reviewer": "Reviewer_8p8i", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a generalized multiscale learning framework for learning in temporally structured environments. The basic idea is that each weight in a NN is decomposed into a sum of subweights, each updated with different learning and decay rates, with all components evolving independently as compared to previous models which involved complex coupling between timescales. They show that these previous models can be shown to be equivalent to the multiscale learner, with no need for these couplings across timescales. They also show that momentum learning can be characterized within the multiscale learner framework as well. Lastly they derive an equivalent model of bayesian inference over 1/f noise and show that that model is also a special case of the multiscale learner framework.", "summary_of_the_review": "I believe this is a paper worth appearing at ICLR. It presents a well thought out and well presented unifying and simplifying framework for learning in temporally structured environments.", "strength_and_weaknesses": "Strengths:\n- The paper is well written, thorough in theory and background information (it does require reading the appendix to fully grasp the derivations).\n- Simplifying and unifying various timescale learning methods into a single framework is a significant contribution.\n- I specially enjoyed figures 1 and 2. They do a great job of summarizing the whole paper in to two clear, concise and simple figures + captions.\n\nWeaknesses:\n- I'm hard pressed to find one. I understand this is a mostly theoretical paper but it would still have been great to see experiments on more than the simulated cases.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "iqMXi4T0-L", "reviewer": "Reviewer_x4Pe", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors make the connection between the framework of multiscale learning and models of fast weights, momentum, and synaptic modelling. The authors then go on to formulate a model of multiscale learning in terms of Bayesian inference over 1/f noise, which has been reported in prior work to match many real-world distributions. The authors test out their model on synthetic online prediction and classification tasks while showing a variational approximation can retain most of the benefits of the full model. ", "summary_of_the_review": "Ultimately, I am on the fence about this paper. I definitely see some good parts about it, but ultimately feel that each strength is somewhat undercut by a corresponding weakness in the current submission. I lean towards rejection at the current time because I feel that the current paper is somewhat incomplete. If large scale experiments were provided validating the proposed method, I would think it could make a very nice contribution to the conference and I would be firmly voting for accept. However, the current experiments are well below the standards of this conference, and not enough justification is provided to make readers understand why this is a logical ending point for the work. ", "strength_and_weaknesses": "Strengths: \n- It is nice that the paper contextualized prior approaches within a common framework.\n- I appreciate the synthetic experiments as they help build an understanding of what the model is doing.\n- It is nice to show the effectiveness of the variational approach, hinting that the model may showcase scalability.\n\nWeaknesses: \n- I did not find the insight about these models being special cases of a common more general framework to be particularly surprising or enlightening. It really only matters, in my view, if it results in a better alternative to the specific models from prior work considered. \n- The experiments are extremely limited and the barriers for scaling are not even clearly discussed. The authors end with the comment \"the next step will be to investigate how well these methods scale up to large networks, and whether they confer an advantage in real online-learning domains.\" It feels to me that this is such an obvious step based on the discourse that the current paper comes across as incomplete. If the authors could present experiments showcasing this, I think this paper could have significant impact. However, if this is not possible, I do not believe there will be much interest in the perspective highlighted here. The experiments are so synthetic and toy relative to the typical standards of this conference that I feel that the lack of any experiments on a meaningful scale requires far more justification than provided within the submitted draft. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "VHSYQGnHFDq", "reviewer": "Reviewer_6H7B", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Temporally structured environments involve multiple learning processes operating at different time-scales. The paper provides a unifying view of learning in these temporally structured environments. The authors begin by formalizing the multi-scale learning setup, where the weights are decomposed into sub-weights, each characterized by a timescale and learning rate and decay. The authors first establish fast weights as an instantiation of the multiscale optimizer with 2 timescales. Next, the authors establish the equivalence between the multiscale optimizer (in the configuration of a fast-weight model) and the Benna-Fusi model for real synapses, as well as momentum learning. An interesting observation from the analysis is that momentum can be characterized as a fast weight with a negative learning rate. The authors then present a Bayesian inference view for a multiscale learner, adopting the $1/f$ model, and describe exact inference with a Kalman filter and approximate inference with extended Kalman filters and variational approximation for extended Kalman filter. Finally the authors validate the multiscale learning framework with some experiments on synthetic tasks. ", "summary_of_the_review": "To summarize, the paper presents a unifying perspective for multiscale optimization in some non-iid settings. The paper makes useful connections to existing ideas and proposes a novel Bayesian formulation for multiscale optimizers. The empirical evaluations seem lacking given certain assumptions made in the formalization. More realistic experiments would be helpful in addressing this shortcoming. Nonetheless, the conceptual and theoretical contributions of the paper are substantial, so I lean towards acceptance. ", "strength_and_weaknesses": "**Strengths**:\n\n- The paper tackles an important problem of learning in non-i.i.d settings, in the specific case where multiple learning processes operate at different timescales. The authors present a unifying framework of multiscale optimizers for learning in such settings. The framework encompasses well studied approaches of momentum learning and the Benna-Fusi model, providing novel insights for momentum learning.  \n- The Bayesian formulation provides a simple and straightforward approach for implementing the multiscale optimizer. \n- The synthetic simulations provide _some_ evidence on the applicability of the multiscale optimizer in practical settings. (I discuss this a bit more in the weakness)\n\n**Weaknesses**\n\n- The paper studies a particular setting where the multi-scale weights are additive. This however seems like a somewhat limiting assumption to me. The authors do not discuss this choice in the paper and it is not clear whether this assumption holds true in realistic settings.\n- While the simulations provide a sanity check for the approach, they leave quite a bit to be desired. As mentioned in the first point, the assumption of weights of the weights of the different timescales being additive might not hold in any realistic setting. Experiments on more realistic settings would be helpful to address this concern. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster