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: This paper proposes a learning framework called Bayes Augmented with Memory (BAM), where a recursive adaptation of the Bayes formula is augmented with selection variables $W$ to allow an agent to adaptively choose past experiences to forget. The original recursive Bayes formula assumes the stationarity of the data generating distribution, so keep all the past data for updating the posteriors for newly arriving data, and thus always tends to decrease the variance of the posteriors as it is meant to be. However, this can severely fail when the stationarity assumption is violated; when an agent encounters a change in the environment, it should somehow discard the past experience and quickly adapt to new data. BAM explicitly models this procedure with the binary selection variables, and by greedily optimizing those selection variables at each entrance of new data, quickly adapt to the change of environment and still leverages the past experiences. Various experiments with non-stationary environments demonstrate the usefulness of the proposed framework.