BAM: Bayes Augmented with Memory

1 INTRODUCTION . The ability of an agent to continuously modulate its belief while interacting with a non-stationary environment is a hallmark of intelligence and has garnered a lot of attention in recent years ( Zhang et al. , 2020 ; Ebrahimi et al. , 2020 ; Xie et al. , 2020 ) . The Bayesian framework enables online learning by providing a principled way to incorporate new observations into an agent ’ s model of the world ( Jaynes , 2003 ; Gelman et al. , 2013 ) . Through the use of Bayes ’ theorem , the agent can combine its own ( subjective ) a priori knowledge with data to achieve an updated belief encoded by the posterior distribution . The Bayesian framework is a particularly appealing option for online learning because Bayes ’ theorem is closed under recursion , enabling continuous updates in what is commonly referred to as the recursive Bayes method ( Wakefield , 2013 ) . As an example , suppose the agent first observes a batch of data , D1 , and then later observes another batch of data , D2 . We can express the agent ’ s posterior distribution over the world , where the world is represented by θ , as p ( θ|D1 , D2 ) = p ( D2|θ ) p ( θ|D1 ) p ( D2|D1 ) , ( 1 ) where p ( D2|D1 ) = ∫ p ( D2|θ ) p ( θ|D1 ) dθ . ( 2 ) Equation 1 demonstrates the elegance and simplicity of recursive Bayes : at time t , the agent recycles its previous posterior , p ( θ|D < t ) , where D < t = { D1 , · · · , Dt−1 } , into its current prior and then combines it with a newly observed batch of data , Dt , to obtain an updated posterior , p ( θ|D≤t ) . ∗corresponding authors At first glance , it would appear that a naive application of recursive Bayes would suffice for most online learning tasks . However , the recursive Bayes method relies on the assumption that the world is stationary , i.e . D1 , D2 , · · · are all independent and identically distributed . When this assumption is violated , recursive Bayes can fail catastrophically . As an illustration , consider the law of total variance : Var ( θ|D < t ) = E [ Var ( θ|D < t , Dt ) ∣∣D < t ] + Var ( E [ θ|D < t , Dt ] ∣∣D < t ) . ( 3 ) Since both terms on the right hand side are positive , equation 3 reveals that in expectation , the variance of the posterior decreases as more data is seen regardless of the actual distribution of Dt , i.e . Var ( θ|D < t ) ≥ E [ Var ( θ|D < t , Dt ) ∣∣D < t ] . ( 4 ) In fact , for some models equation 4 is true with probability 1 ; we demonstrate examples in Appendix A . Thus , if the parameters of the environment , θ , were to change , the variance of the posterior would still decrease , becoming more certain of a potentially obsolete parameter estimate . Modeling the environment as stationary when it is actually changing also keeps the learning speed of the agent artificially low , as tighter posteriors prevent large jumps in learning . This is the opposite of what an intelligent agent should do in such an event : if the environment changes , we would expect the agent ’ s uncertainty and learning speed to increase in response . As was elegantly stated by Monton ( 2002 ) , the problem with naive use of recursive Bayes is that `` Such a Bayesian never forgets . '' Previous approaches on enabling recursive Bayes to work in non-stationary settings have primarily focused on forgetting past experience either through the use of changepoint detection ( Adams & MacKay , 2007 ; Li et al. , 2021 ) , or by exponentially weighting past experiences ( Moens , 2018 ; Moens & Zénon , 2019 ; Masegosa et al. , 2020 ) . While empirically successful , their focus on forgetting the past means that revisited states are treated as novel . In this work we take an alternative approach to online Bayesian learning in non-stationary environments by endowing an agent with an explicit memory module . Crucially , the addition of a memory buffer equips the agent with the ability to modulate its uncertainty by choosing what past experiences to both forget and remember . We call our approach Bayes with Adaptive Memory ( BAM ) and demonstrate its wide applicability and effectiveness on a number of non-stationary learning tasks . 2 BAYES WITH ADAPTIVE MEMORY . The generative model is assumed to evolve according to θt ∼ p ( θt|θt−1 , t ) , ( 5 ) Dt ∼ pt ( D ) ≡ p ( D|θt ) , ( 6 ) where equation 5 is the latent dynamics that dictate the evolution of the environment parameters , θt , and equation 6 is the likelihood whose parametric form is fixed throughout time , i.e . pt ( D ) = N ( θt , σ2 ) . Equations 5 and 6 define a state-space model , which allows one to infer θt through Bayesian filtering ( Särkkä , 2013 ) p ( θt|D≤t ) ∝ p ( Dt|θt ) p ( θt|D < t ) , ( 7 ) p ( θt|D < t ) = ∫ p ( θt|θt−1 , t ) p ( θt−1|D < t ) dθt−1 . ( 8 ) The parameterization of equations 5 and 6 dictate the tractability of equations 7 and 8 . If a priori an agent knew that equation 5 is a linear dynamical system with additive white Gaussian noise and equation 6 is also Gaussian whose conditional mean is a linear function of θt , then the Kalman filter can be used ( Kalman , 1960 ) . For more complicated latent dynamics and/or likelihood models , methods such as particle filtering ( Doucet & Johansen , 2009 ) and unscented Kalman filtering ( Julier & Uhlmann , 1997 ) can be used . Crucially , Bayesian filtering methods assume that the latent dynamics governed by equation 5 are known ; however , this is rarely the case in practice . Instead of making assumptions on the parametric form of equation 5 , we take a different approach . In BAM , the agent maintains a memory buffer , D < t , that stores previous observations of the environment . At time t the agent obtains a new batch of data , Dt ∼ pt ( D ) . How should the agent combine the newly observed data , Dt , with its stored memory , D < t , to update its belief as encoded by the posterior distribution ? In recursive Bayes , the posterior distribution is computed according to1 p ( θt|Dt , D < t ) ∝ p ( Dt|θt ) p ( θt|D < t ) , ( 9 ) p ( θt|D < t ) ∝ p ( θt ) t−1∏ j=1 p ( Dj |θt ) , ( 10 ) where we refer to p ( θt ) as the base prior . Equation 10 allows us to interpret recursive Bayes as the agent constructing a dynamic prior , p ( θt|D < t ) , using all the experiences stored in its memory buffer . This works under the stationarity assumption ; when this assumption is violated , the application of Bayes ’ theorem can lead to confidently wrong results as the `` distance '' between pi ( D ) and pj ( D ) can be vast . An alternative is for the agent to completely forget all of its past experiences p ( θt|Dt ) ∝ p ( Dt|θt ) p ( θt ) . ( 11 ) While equation 11 may be viable in situations where Dt is sufficiently informative , it is wasteful when experiences in the memory buffer may help infer θt . BAM dynamically finds a middle ground between these two extremes of remembering ( equation 10 ) and forgetting ( equation 11 ) everything by allowing the agent to choose which data to use from its memory buffer to construct the prior . Specifically , the agent is endowed with a time-dependent readout weight , Wt = [ wt,1 , wt,2 , · · · , wt , t−1 ] where wt , j ∈ [ 0 , 1 ] . Given a new datum Dt , BAM constructs its posterior according to p ( θt|Dt , D < t , Wt ) ∝ p ( θt ) p ( Dt|θt ) t−1∏ j=1 p ( Dj |θt ) wt , j . ( 12 ) We can rewrite equation 12 as p ( θt|Dt , D < t , Wt ) = p ( Dt|θt ) p ( θt|D < t , Wt ) p ( Dt|D < t , Wt ) , ( 13 ) where p ( θt|D < t , Wt ) ∝ p ( θt ) t−1∏ j=1 p ( Dj |θt ) wt , j , ( 14 ) and p ( Dt|D < t , Wt ) = ∫ p ( Dt|θt ) p ( θt|D < t , Wt ) dθt . ( 15 ) The prior construction in equation 14 is akin to recursive Bayes , but now the agent can dynamically and adaptively change its prior by using the readout weights , Wt , to weigh the importance of previous experience where at the extreme , it can choose to completely forget a previous experience , wt , j = 0 , or fully remember it , wt , j = 1 . For simplicity , we restrict the readout weights to be binary , i.e . wt , j ∈ { 0 , 1 } . The combination of a memory buffer , D < t , with a time-dependent readout weight , Wt , allows BAM to generalize many previously proposed approaches . By setting wt,1 = wt,2 = · · · = wt , t−1 = 1 , we recover recursive Bayes ( equation 10 ) . By setting wt,1 = wt,2 = · · · = wt , t−1 = α , where 0 ≤ α ≤ 1 we recover the power priors approach of Ibrahim et al . ( 2015 ) . By setting wt , j = αt−1−j , where 0 ≤ α ≤ 1 , we recover exponential forgetting ( Moens , 2018 ; Moens & Zénon , 2019 ; Masegosa et al. , 2020 ) . Lastly , by setting a particular subset of the readout weights to be 0 , we recover Bayesian unlearning ( Nguyen et al. , 2020 ) . The ability to adaptively change its prior implies that BAM can increase/decrease its uncertainty as the situation demands ; subsequently , this modulates the agent ’ s learning speed . Using variance as a proxy for uncertainty , one would expect that the variance of the prior used in BAM ( equation 14 ) is always at least as large as the variance of the prior used in recursive Bayes ( equation 10 ) . We formalize this for the case of binary readout weights in the following proposition . Proposition 1 . Let p ( θ|D < t , Wt ) be the prior used by BAM , defined in equation 14 and let p ( θ|D < t ) be the recursive Bayes prior , defined in equation 13 . Then E [ Var ( θ|D < t , Wt ) ∣∣Wt ] ≥ E [ Var ( θ|D < t ) ] , ∀Wt ∈ { 0 , 1 } t−1 . ( 16 ) Proof . Proof is in Appendix B . 1Recursive Bayes is equivalent to Bayesian filtering when p ( θt|θt−1 , t ) = δ ( θt = θt−1 ) .

Summary: The paper aims to solve the downside of the posterior shrinkage of Bayesian online learning when applied in a non-stationary environment. The previously posterior (a.k.a., the current prior) could be misspecified for current observations with the posterior shrinkage due to the environment change. To solve this misspecification problem, the proposed method constructs an adaptive prior that is correctly specified for current observations. This adaptive prior distribution is constructed by selecting the previous relevant data samples, which requires storing the whole history. The method is evaluated on various analytical experiments.