Title: NEURAL LINEAR BANDITS: OVERCOMING CATASTROPHIC FORGETTING THROUGH LIKELIHOOD MATCHING

Abstract: We study neural-linear bandits for solving problems where both exploration and representation learning play an important role. Neural-linear bandits leverage the representation power of deep neural networks and combine it with efficient exploration mechanisms, designed for linear contextual bandits, on top of the last hidden layer. Since the representation is being optimized during learning, information regarding exploration with "old" features is lost. Here, we propose the first limited memory neural-linear bandit that is resilient to this catastrophic forgetting phenomenon. We perform simulations on a variety of real-world problems, including regression, classification, and sentiment analysis, and observe that our algorithm achieves superior performance and shows resilience to catastrophic forgetting.

1. INTRODUCTION:

Deep neural networks (DNNs) can learn representations of data with multiple levels of abstraction and have dramatically improved the state-of-the-art in speech recognition, visual object recognition, object detection and many other domains such as drug discovery and genomics (LeCun et al., 2015; Goodfellow et al., 2016) . Using DNNs for function approximation in reinforcement learning (RL) enables the agent to generalize across states without domain-specific knowledge, and learn rich domain representations from raw, high-dimensional inputs (Mnih et al., 2015; Silver et al., 2016) .

Nevertheless, the question of how to perform efficient exploration during the representation learning phase is still an open problem. The -greedy policy (Langford & Zhang, 2008) is simple to implement and widely used in practice (Mnih et al., 2015) . However, it is statistically suboptimal. Optimism in the Face of Uncertainty (Abbasi-Yadkori et al., 2011; Auer, 2002, OFU) , and Thompson Sampling (Thompson, 1933; Agrawal & Goyal, 2013, TS) use confidence sets to balance exploitation and exploration. For DNNs, such confidence sets may not be accurate enough to allow efficient exploration. For example, using dropout as a posterior approximation for exploration does not concentrate with observed data (Osband et al., 2018) and was shown empirically to be insufficient (Riquelme et al., 2018) . Alternatively, pseudo-counts, a generalization of the number of visits, were used as an exploration bonus (Bellemare et al., 2016; Pathak et al., 2017) . Inspired by tabular RL, these ideas ignore the uncertainty in the value function approximation in each context. As a result, they may lead to inefficient confidence sets (Osband et al., 2018) .

Linear models, on the other hand, are considered more stable and provide accurate uncertainty estimates but require substantial feature engineering to achieve good results. Additionally, they are known to work in practice only with "medium-sized" inputs (with around 1, 000 features) due to numerical issues. A natural attempt at getting the best of both worlds is to learn a linear exploration policy on top of the last hidden layer of a DNN, which we term the neural-linear approach. In RL, this approach was shown to refine the performance of DQNs (Levine et al., 2017) and improve exploration when combined with TS (Azizzadenesheli et al., 2018) and OFU (O'Donoghue et al., 2018; Zahavy et al., 2018a) . For contextual bandits, Riquelme et al. (2018) showed that neural-linear TS achieves superior performance on multiple data sets.

A practical challenge for neural-linear bandits is that the representation (the activations of the last hidden layer) change after every optimization step, while the features are assumed to be fixed over time when used by linear contextual bandits. Riquelme et al. (2018) tackled this problem by storing the entire data set in a memory buffer and computing new features for all the data after each DNN learning phase. The authors also experimented with a bounded memory buffer, but observed a significant decrease in performance due to catastrophic forgetting (Kirkpatrick et al., 2017) , i.e., a loss of information from previous experience. In this work, we propose a neural-linear bandit that uses TS on top of the last layer of a DNN (Fig. 1 )1 . Key to our approach is a novel method to compute priors whenever the DNN features change that makes our algorithm resilient to catastrophic forgetting. Specifically, we adjust the moments of the likelihood of the reward estimation conditioned on new features to match the likelihood conditioned on old features. We achieve this by solving a semi-definite program (Vandenberghe & Boyd, 1996, SDP) to approximate the covariance and using the weights of the last layer as prior to the mean.

We present simulation results on several real-world and simulated data sets, including classification and regression, using Multi-Layered Perceptrons (MLPs). Our findings suggest that using our method to approximate priors improves performance when memory is limited. Finally, we demonstrate that our neural-linear bandit performs well in a sentiment analysis data set where the input is given in natural language (of size R 8k ) and we use a Convolution Neural Network (CNNs). In this regime, it is not feasible to use a linear method due to computational problems. To the best of our knowledge, this is the first neural-linear algorithm that is resilient to catastrophic forgetting due to limited memory.

2. BACKGROUND:

Algorithm 1 TS for linear contextual bandits

∀i ∈ [1.., N ], set B i = I d , μi = 0 d , f i = 0 d for t = 1, 2, . . . , do ∀i ∈ [1.., N ], sample μi from N ( μi , v 2 B -1 i ) Play arm a(t) := arg max i b(t) T μi Observe reward r t Update: B a(t) = B a(t) + b(t)b(t) T f a(t) = f a(t) + b(t)r t , μa(t) = B -1 a(t) f a(t) end for

The stochastic, contextual (linear) multiarmed bandit problem. There are N arms (actions). At time t a context vector b(t) ∈ R d , is revealed. The history at time t is defined to be H t-1 = {b(τ ), a(τ ), r a(τ ) (τ ), τ = 1, ..., t-1}, where a(τ ) denotes the arm played at time τ . The contexts b(t) are assumed to be realizable, i.e., the reward for arm i at time t is generated from an (unknown) distribution s

.t. E [r i (t)|b(t), H t-1 ] = E [r i (t)|b(t)] = b(t) T µ i , where {µ i ∈ R d } N

i=1 are fixed but unknown parameters. An algorithm for this problem needs to choose at every time t an arm a(t) to play, with the knowledge of history H t-1 and current context b(t). Let a * (t) denote the optimal arm at time t, i.e. a * (t) = arg max i b(t) T µ i , and let ∆ i (t) the difference between the mean rewards of the optimal arm and of arm i at time t, i.e.,

∆ i (t) = b(t) T µ a * (t) -b(t) T µ i . The objective is to minimize the total regret R(T ) = T t=1 ∆ a(t)

, where the time horizon T is finite.

TS for linear contextual bandits. Thompson sampling is an algorithm for online decision problems where actions are taken sequentially in a manner that must balance between exploiting what is known to maximize immediate performance and investing to accumulate new information that may improve future performance (Russo et al., 2018; Lattimore & Szepesvári, 2018) . For linear contextual bandits, TS was introduced in (Agrawal & Goyal, 2013, Alg. 1) .

Suppose that the likelihood of reward r i (t), given context b(t) and parameter µ i , were given by the pdf of Gaussian distribution N (b(t) T µ i , ν 2 ), and let

B i (t) = B 0 i + t-1 τ =1 b(τ )b(τ ) T 1 i=a(τ ) , μi (t) = B -1 i (t) t-1 τ =1 b(τ )r a(τ ) (τ )1 i=a(τ )

, where 1 is the indicator function. Given a Gaussian prior for arm i at time t, N (μ i (t), v 2 B -1 i (t)), the posterior distribution at time t + 1 is given by, P r(μ i |r i (t)) ∼ P r(r i (t

EQUATION

At each time step t, the algorithm generates samples {μ i (t)} N i=1 from the posterior distribution N (μ i (t), v 2 B -1 i (t)), plays the arm i that maximizes b(t) T µ i (t) and updates the posterior. TS is guaranteed to have a total regret at time T that is not larger than O(d 3/2 √ T ), which is within a factor of √ d of the information-theoretic lower bound for this problem. It is also known to achieve excellent empirical results (Lattimore & Szepesvári, 2018) .

Although that TS is a Bayesian approach, the description of the algorithm and its analysis are priorfree, i.e., the regret bounds will hold irrespective of whether or not the actual reward distribution matches the Gaussian likelihood function used to derive this method (Agrawal & Goyal, 2013) .

Bayesian Linear Regression. A different mechanism, based on Bayesian Linear Regression, was proposed by Riquelme et al. (2018) . Here, the noise parameter ν (Alg. 1) is replaced with a prior belief that is being updated over time. The prior for arm i at time t is given by P r(μ i , ν2

i ) = P r(ν 2 i )P r(μ i |ν 2 i ), where P r(ν 2 i ) is an inverse-gamma distribution Inv-Gamma(a i (t), b i (t)), and the conditional prior density P r(

μ i |ν 2 i ) is a normal distribution, P r(μ i |ν 2 i ) ∝ N μi (t), ν2 i B i (t) -1 .

For Gaussian likelihood, the posterior distribution at time

τ = t + 1 is, Pr(ν i ) = Inv-Gamma (a i (τ ), b i (τ )) and Pr(μ i |ν i ) = N μi (τ ), ν 2 i B i (τ ) -1

, where:

EQUATION

The problem with this approach is that the marginal distribution of µ i is heavy tailed (multi-variate t-student distribution, see O' Hagan & Forster (2004) , page 246, for derivation), and does not satisfy the necessary concentration bounds for exploration in (Agrawal & Goyal, 2013; Abeille et al., 2017) .

Thus, in order to analyze the regret of this approach, new analysis has to be derived, which we leave to future work. Empirically, this update scheme was shown to convergence to the true posterior and demonstrated excellent empirical performance (Riquelme et al., 2018) . This can be explained by the fact that the mean of the noise parameter ν i , given by Eν i = bi(t) ai(t)-1 , is decreasing to zero with time, which may compensate for the lack of shrinkage due to the heavy tail distribution.

3. LIMITED MEMORY NEURAL-LINEAR TS:

Our algorithm, as depicted in Fig. 1 , is composed of four main components: (1) A DNN that takes a raw context as an input and is trained to predict the reward of each arm; (2) An exploration mechanism that uses the last layer activations of the DNN as features and performs linear TS on top of them; (3) A memory buffer that stores previous experience; (4) A likelihood matching mechanism that uses the memory buffer and the DNN to account for changes in representation. We now explain how each of these components works; code can be found in (link).

To derive our algorithm we make the assumption that all the representations that are produced by the DNN are realizable. That is, for each representation there exist a different linear coefficients vector (e.g. µ for φ, β for ψ,) such that the expected reward is linear in the features. Explicitly, this means that for representations φ, ψ it holds that

E[r i (t)|φ(t)] = φ(t) T µ i = ψ(t) T β i = E[r i (t)|ψ(t)].

While the realizability assumption is standard in the existing literature on contextual multi-armed bandits (Chu et al., 2011; Abbasi-Yadkori et al., 2011; Agrawal & Goyal, 2013) , it is quite strong and may not be realistic in practice. We further discuss these assumptions in the discussion paragraph below and in Section 5.

1. Representation. Our algorithm uses a DNN, denoted by D, that takes the raw context b(t) ∈ R d as its input. The network has N outputs that correspond to the estimation of the reward of each arm; given context b(t) ∈ R d , D(b(t)) i denotes the estimation of the reward of the i-th arm.

Using a DNN to predict the reward of each arm allows our algorithm to learn a nonlinear representation of the context. This representation is later used for exploration by performing linear TS on top of the last hidden layer activations. We denote the activations of the last hidden layer of D applied to this context as φ(t) = LastLayerActivations(D(b(t))), where φ(t) ∈ R g . The context b(t) represents raw measurements that can be high dimensional (e.g., image or text), where the size of φ(t) is a design parameter that we choose to be smaller (g < d) . This makes contextual bandit algorithms practical for such data sets. Moreover, φ(t) can potentially be linearly realizable (even if b(t) is not) since a DNN is a global function approximator (Barron, 1993) and the last layer is linear.

1.1 Training. Every L iterations, we train D for P mini-batches. Training is performed by sampling experience tuples {b(τ ), a(τ ), r a(τ ) (τ )} from the replay buffer E (details below) and minimizing the mean squared error (MSE),

EQUATION

where r a(τ ) is the reward that was received at time τ after playing arm a(τ ) and observing context b(τ ) (similar to Riquelme et al. (2018) ). Notice that only the output of arm a(τ ) is differentiated.

We emphasize that the DNN, including the last layer, are trained end-to-end to minimize Eq. ( 3).

Algorithm 2 Limited Memory Neural-linear TS

Set ∀i ∈ [1, .., N ] : Φ 0 i = I d , μi = µ 0 i = 0 d , Φ i = 0 dxd , f i = 0 d , a i = a 0 i , b i = b 0 i Initialize Replay Buffer E, and DNN D Define φ(t) ← LastLayerActivations(D(b(t))) for t = 1, 2, . . . , do Observe b(t), evaluate φ(t) Posterior sampling: ∀i ∈ [1, .., N ], sample: νi (t) ∼ Inv-Gamma (a i (t), b i (t)) μi (t) ∼ N μi , νi (t) 2 (Φ 0 i + Φ i ) -1 Play arm a(t) := arg max i φ(t) T μi (t) Observe reward r t Store {b(t), a(t), r t } in E if E is full then

Remove the first tuple in E with a = a(t) (round robin) end if Bayesian linear regression update: Solve for Φ 0 i using Eq. ( 6) with

Φ a(t) = Φ a(t) + φ(t)φ(t) T , f a(t) = f a(t) + φ(t) T r t μa(t) = (Φ 0 a(t) + Φ a(t) ) -1 Φ 0 a(t) µ 0 a(t) + f a(t) a a(t) = a a(t) + 1 2 R 2 a(t) = R 2 a(t) + r 2 t b a(t) = b 0 a(t) + 1 2 R 2 a(t) + (µ 0 a(t) ) T B 0 a(t) µ 0 a(t) -μT a(t) B a(t) μa(t) if (t mod L) = 0 then for ∀i ∈ [1, ..,

E i φ , E i φ old , Φ old i Set µ 0 i ← LastLayerWeights(D) i Φ i = ni j=1 φ i j (φ i j ) T , f i = ni j=1 (φ i j ) T r j .

end for end if end for 2. Exploration. Since our algorithm is performing training in phases (every L steps), exploration is performed using a fixed representation φ (D has fixed weights between training phases). At each time step t, the agent observes a raw context b(t) and uses the DNN D to produces a feature vector φ(t). The features φ(t) are used to perform linear TS, similar to Algorithm 1, but with two key differences.

First, we introduce a likelihood matching mechanism that accounts for changes in representation (see 4. below for more details). Second, we follow the Bayesian linear regression equations, as suggested in (Riquelme et al., 2018) , and perform TS while updating the posterior both for µ, the mean of the estimate, and ν, its variance. This is done in the following manner. We begin by sampling a weight vector μi for each arm i ∈ 1..N, from the posterior by following two steps. First, the variance ν2

i is sampled from Inv-Gamma (a i , b i ) . Then, the weight vector μi is sampled, from N μi , ν2

i (Φ 0 i + Φ i ) -1 . Once we sampled a weight vector for each arm, we choose to play arm a(t) = arg max i φ(t) T μi , and observe reward r a(t) (t). This is followed by a posterior update step, based on Eq. ( 2):

Φ a(t) = Φ a(t) + φ(t)φ(t) T , f a(t) = f a(t) + φ(t) T r t , R 2 a(t) = R 2 a(t) + r a(t) (t) 2 (4) μa(t) = (Φ a(t) + Φ 0 a(t) ) -1 (Φ 0 a(t) µ 0 a(t) + f a(t) ), a a(t) = a a(t) + 1 2 , b a(t) = b 0,a(t) + 1 2 R 2 a(t) + µ T 0,a(t) Φ 0,a(t) µ 0,a(t) -μa(t) (t) T Φ a(t) (t)μ a(t) (t) .

The exploration mechanism is responsible for choosing actions; it does not change the weights of the DNN.

3. Memory buffer. After an action a(t) is played at time t, we store the experience tuple {b(t), a(t), r a(t) (t)} in a finite memory buffer of size n that we denote by E. Once E is full, we remove tuples from E in a round robin manner, i.e., we remove the first tuple in E with a = a(t).

4. Likelihood matching. Before each learning phase, we evaluate the features of D on the replay buffer. Let E i be a subset of memory tuples in E at which arm i was played, and let n i be its size. We denote by E i φ old ∈ R ni×g a matrix whose rows are feature vectors that were played by arm i. After a learning phase is complete, we evaluate the new activations on the same replay buffer and denote the equivalent set by E i φ ∈ R ni×g . Our approach is to summarize the knowledge that the algorithm has gained from exploring with the features φ old into priors on the new features Φ 0 i , µ 0 i . Once these priors are computed, we restart the linear TS algorithm using the data that is currently available in the replay buffer. For each arm i, let φ i j = (E i φ ) j be the j-th row in E i φ and let r j be the corresponding reward, we set

Φ i = ni j=1 φ i j (φ i j ) T , f i = ni j=1 (φ i j ) T r j .

We now explain how we compute Φ 0 i , µ 0 i . Recall that under the realizability assumption we have that

E[r i (t)|φ(t)] = φ(t) T µ i = φ old (t) T µ old i = E[r i (t)|ψ(t)].

Thus, the likelihood of the reward is invariant to the choice of representation , i.e. N (φ(t) T µ i , ν 2 ) ∼ N (φ old (t) T µ old i , ν 2 ). For all i, define the estimator of the reward as θ i (t) = φ(t) T μi (t), and its standard deviation s t,i = φ(t) T Φ i (t) -1 φ(t) (see (Agrawal & Goyal, 2013) for derivation). By definition of μi (t), marginal distribution of each θ i (t) is Gaussian with mean φ i (t) T μi (t) and standard deviation ν i s t,i . The goal is to match the likelihood of the reward estimation θ i (t) given the new features to be the same as with the old features.

4.1. Approximation of the mean µ 0:

i : Recall that the realizability assumption implies a linear connection between µ old i , µ 0 i , i.e., (µ old i ) T φ old = (µ 0 i ) T φ, thus, we can solve a linear set of equations and get a linear mapping from µ 0 i to µ old i :

(µ 0 i ) T = (µ old i ) T E i φ old (E i φ ) -1 .

(5)

In addition to the realizability assumption, for Eq. ( 5) to hold the matrix E i φ must be invertible. In practice, we found that a different solution that is based on using the DNN weights performed better. Recall that the DNN is trained to minimize the MSE (Eq. ( 3)). Thus, given the new features φ, the weights of the last layer of the DNN make a good prior for µ 0 i . This approach was shown empirically to make a good approximation (Levine et al., 2017) , as the DNN was optimized online by observing all the data (and is therefore not limited to the current replay buffer).

4.2 Approximation of the variance s j,i :. For each arm i, our algorithm receives as input the sets of new and old features E i φ , E i φ old ; denote the elements in these sets by {φ old j , φ j } ni j=1 . In addition, the algorithm receives the correlation matrix Φ old i . Notice that due the nature of our algorithm, Φ old i holds information on contexts that are not available in the replay buffer. The goal is to find a correlation matrix,Φ 0 i , for the new features that will hold the same information on past context as Φ old i . I.e., we want to find

Φ 0 i such that ∀i ∈ [1..N ], j ∈ [1..n i ] s 2 j,i . = (φ old j ) T (Φ old i ) -1 φ old j = φ T j (Φ 0 i ) -1 φ j .

Using the cyclic property of the trace, this is equivalent to finding Φ 0 i , s.t. ∀j ∈ [1, .., n i ], s 2 j,i = Trace (Φ 0 i ) -1 φ j φ T j . Next, we define X i to be a vector of size n i in the vector space of d × d symetric matrices, with its j-th element X j,i to be the matrix φ j φ T j . Notice that (Φ 0 i ) -1 is constrained to be semi positive definite (being a correlation matrix), thus, the solution can be found by solving an SDP (Eq. ( 6)). Note that Trace(X T j,i (Φ 0 i ) -1 ) is an inner product over the vector space of symmetric matrices, known as the Frobenius inner product. Thus, the optimization problem is equivalent to a linear regression problem in the vector space of PSD matrices. In practice, we use cvxpy (Diamond & Boyd, 2016) to solve for all actions i ∈ [1..N ] :

EQUATION

)

Discussion. The correctness of our algorithm follows from the proof of (Agrawal & Goyal, 2013) .

To see this, recall that we match the moments of the reward estimate θ i (t) after every time that the representation changes. Assuming that we solve Eq. ( 5) and Eq. ( 6) precisely, then the reward estimation given the new features have precisely the same moments and distribution as with the old features. Since the distribution of the estimate did not change, its concentration and anti-concentration bounds do not change, and the proof in (Agrawal & Goyal, 2013 ) can be followed.

The problem is, that in general, we cannot guarantee to solve Eq. ( 5) and Eq. ( 6) exactly. We will soon show that under the realizability assumption, in addition to an invertibility assumption, it is possible to choose an analytical solution for the priors µ 0 , Φ 0 that guarantees an exact solution. However, these conditions may be too strong and not realistic. We describe this scenario to highlight the existence of a scenario (and conditions) in which our algorithm is optimal; we hope to relax them in future work.

We emphasize here that if these conditions do not hold, then our algorithm is only an approximation, without theoretical guarantees. In the next section we justify using our algorithm through thorough experimentation.

For simplicity, we consider a single arm. Assume that m past observations, which we denote by E m φ old , were used to learn estimators (μ m ) old , Φ old m using BLR (Eq. ( 2)). Due to the limited memory, some of these measurements are not available in the replay buffer, and all of the information regarding them is summarized in (μ m ) old , Φ old m . In addition, we are given a replay buffer of size n, that is used to produce (before and after the training) new and old feature matrices E n φ old , E n φ . We also denote by R n the reward vector (using data from the replay buffer) and by R m the reward vector that corresponds to features E m φ old which is not available in the replay buffer. Recall that the realizeability assumption implies that the features φ and φ old are linear mappings of the raw context b, i.e., φ = A φ b, φ old = A φ old b. Under the assumption that all the relevant matrices are invertible, we use Eq. ( 5) to find a prior for µ 0 , i.e., we set µ T 0 = (μ old m ) T E n φ old (E n φ ) -1 . In addition, for the covariance matrix, we set

(Φ 0 ) -1 = (E n φ ) -1 E n φ old (Φ old m ) -1 (E n φ old ) T ((E n φ ) T ) -1

, which is a solution to Eq. ( 6).

In addition, we get that if the relevant matrices are invertibele, then Φ 0 = Φ m , and that Φ 0 µ 0 = E m φ R m (see the supplementary for derivation). Plugging these estimates as priors in the Bayesian linear regression equation we get the following solution for μ :

μ = (Φ n + Φ 0 ) -1 (Φ 0 µ 0 + E n φ R n ) = (Φ n + Φ m ) -1 (E m φ R m + E n φ R n ), i.

e., we got the linear regression solution for µ as if we were able to evaluate the new features φ on the entire data, while having a finite memory buffer and changing features! 3.1 COMPUTATIONAL COMPLEXITY Solving the SDP. Recall that the dimension of the last layer is g < d where d is the dimension of the raw features, and the size of the buffer is n. Following this notation, when solving the SDP, we optimize over matrices in R g×g that are subject to n equality constraints. We refer the reader to Vandenberghe & Boyd (1996) for an excellent survey on the complexity of solving SDPs. Here, we will refer to interior-point methods. The number of iterations required to solve an SDP to a given accuracy grows with problem size as O(g 0.5 ). Each iteration involves solving a least-squares problem of dimension g 2 . If, for example, this least-squares problem is solved with a projected gradient descent method, then the time complexity for finding an -optimal solution is g 2 / , and the computational complexity of each gradient step is ng 2 (matrix-vector multiplications). Vandenberghe & Boyd experimented with solving SDPs of different sizes and observed that it takes almost the same amount of iterations to solve them. In addition, they found that SDP algorithms converge much faster than the worst-case theoretical bounds.

In our case, g, the size of the last layer, was fixed to be 50 in all the experiments. Thus, although the dimension of the raw features d varies in size across data sets, the complexity of the solving the SDP is fixed. The dependence of the computational complexity on the buffer size n is at most linear (CVXPY exploits sparsity structure of the matrix to enhance computations); we didn't encounter a significant changes in computation time when changing the buffer size in the range of 200 -2000. It took us 10 -30 seconds on a standard "MacBook Pro" to solve a single SDP.

Dependence on T . The full memory approach results in computational complexity of O(T 2 ) and memory complexity of O(T ) where T is the number of contexts seen by the algorithm. This is because it is estimating the TS posterior using the complete data every time the representation changes. On the other hand, the limited memory approach uses only the memory buffer to estimate the posterior but additionally solves an SDP. This gives a memory complexity of O(1) and computational complexity of O(T ).

Dependence on A. The computational complexity is linear in the number of actions (we solve an SDP for each action). There is a large variety of problems where this is not an issue (as in our experiments). However, if the problem of interest has many discrete actions, our approach may not be useful.

To summarize, our method is more efficient than the full memory baseline in problems with big data (large T ). Nevertheless, our method requires to solve an SDP (every L iterations), which is computationally prohibitive in general. We deal with this issue by restricting the size of the last layer to be small (g = 50), for which solving the SDP is reasonable.

4. EXPERIMENTS:

We begin this section by testing the resilience of our method to catastrophic forgetting. We present an ablative analysis of our approach and show that the prior on the covariance is crucial. Then, we present results for using MLPs on ten real-world data sets, including a high dimensional natural language data on a task of sentiment analysis (all of these data sets are publicly available through the UCI Machine Learning Repository). Additionally, in the supplementary material, we use synthetic data to test and visualize the ability of our algorithm to learn nonlinear representations during exploration. In all the experiments we used the same hyperparameters (as in (Riquelme et al., 2018 )) for the model, and the same network architecture (an MLP with a single hidden layer of size 50). The only exception is with the text CNN (details below). The size of the memory buffer is set to be 100 per action. We use the Shuttle Statlog data set (Newman et al., 2008) , a real world, nonlinear data set. Each context is composed of 9 features of a space shuttle flight, and the goal is to predict the state of the radiator of the shuttle (the reward). There are k = 7 possible actions, and if the agent selects the right action, then reward 1 is generated. Otherwise, the agent obtains no reward (r = 0).

4.1. CATASTROPHIC FORGETTING:

We experimented with the following algorithms: (1) Linear TS (Agrawal & Goyal, 2013 , Algorithm 1) using the raw context as a feature, with an additional uncertainty in the variance (Riquelme et al., 2018) . ( 2) Neural-Linear TS (Riquelme et al., 2018) . (3) Our neural-linear TS algorithm with limited memory. (4) An ablative version of (3) that calculates the prior only for the mean, similar to (Levine et al., 2017) .

(5) An ablative version of (3) that does not use prior calculations. Algorithms 3-5 make an ablative analysis for the limited memory neural-linear approach. As we will see, adding each one of the priors improves learning and exploration. Fig. 2 shows the performance of each of the algorithms in this setup. We let each algorithm run for 4000 steps (contexts) and average each algorithm over 10 runs. The x-axis corresponds to the number of contexts seen so far, while the y-axis measures the instantaneous regret. All the neural-linear methods retrained the DNN every L = 400 steps for P = 800 mini-batches.

First, we can see that the neural linear method (2nd row) outperforms the linear one (1st row), suggesting that this data set in nonlinear. We can also see that our approach to computing the priors allows the limited memory algorithm (3rd row) to perform almost as good as the neural linear algorithm without memory constraints (2nd row).

In the last two rows we can see a version of the limited memory neural linear algorithm that does not calculate the prior for the covariance matrix (4th row), and a version that does not compute priors at all (5th row). Both of these algorithms suffer from "catastrophic forgetting" due to limited memory. Intuitively, the covariance matrix holds information regarding the number of contexts that were seen by the agent and are used by the algorithm for exploration. When no such prior is available, the agent explores sub-optimal arms from scratch every time the features are modified (every L = 400 steps, marked by the x-ticks on the graph). Indeed, we observe "peaks" in the regret curve for these algorithms (rows 4&5); this is significantly reduced when we compute the prior on the covariance matrix (3rd row), making the limited memory neural-linear bandit resilient to catastrophic forgetting.

4.2. REAL WORLD DATA:

We evaluate our approach on several (10) real-world data sets; for each data set, we present the cumulative reward achieved by the algorithms, detailed above, averaged over 50 runs. Each run was performed for 5000 steps.

Linear vs. nonlinear data sets: The results are divided into two groups, linear and nonlinear data sets. The separation was performed post hoc, based on the results achieved by the full memory methods, i.e., the first group consists of five data sets on which Linear TS (Algorithm 1) outperformed Neural-Linear TS (Algorithm 2), and vice versa. We observed that most of the linear datasets consisted of a small number of features that were mostly categorical (e.g., the mushroom data set has 22 categorical features that become 117 binary features). The DNN based methods performed better when the features were dense and high dimensional. Linear data sets: Since there is no apriori reason to believe that real world data sets should be linear, we were surprised that the linear method made a competitive baseline to DNNs. To investigate this further, we experimented with the best reported MLP architecture for the covertype data set (taken from Kaggle). Linear methods were reported (link) to achieve around 60% test accuracy. This number is consistent with our reported cumulative reward (3000 out 5000). Similarly, DNNs achieved around 60% accuracy, which indicates that the Covertype data set is indeed relatively linear. However, when we measure the cumulative reward, the deep methods take initial time to learn, which can explain the slightly worst score. One particular architecture (MLP with layers 54-500-800-7) was reported to achieve 68%; however, we didn't find this architecture to yield better cumulative reward.

Similarly, for the Adult data set, linear and deep classifiers were reported to achieve similar results (link) (around 84%), which is again equivalent to our cumulative reward of 4000 out of 5000. A specific DNN was reported to achieve 90% test accuracy but did not yield improvement in cumulative reward. These observations can be explained by the different loss function that we optimize or by the partial observably of the bandit problem (bandit feedback). Alternatively, competitions tend to suffer from overfitting in model selection (see the "reusable holdout" paper for more details (Dwork et al., 2015) ). Regret, on the other hand, is less prune to model overfitting, because the model is evaluated at each iteration, and because we shuffle the data at each run.

Limited memory: Looking at Table 1 we can see that on eight out of ten data sets, using the prior computations (Algorithm 3), improved the performance of the limited memory Neural-Linear algorithms. On four out of ten data sets (Mushroom, Financial, Statlog, Epileptic), Algorithm 3 even outperformed the unlimited Neural-Linear algorithm (Algorithm 2).

Limited memory neural linear vs. linear: as linear TS is an online algorithm it can store all the information on past experience using limited memory. Nevertheless, in four (out of five) of the nonlinear data sets the limited memory TS (Algorithm 3) outperformed Linear TS (Algorithm 1).

Our findings suggest that when the data is indeed not linear, than neural-linear bandits beat the linear method, even if they must perform with limited memory. In this case, computing priors improve the performance and make the algorithm resilient to catastrophic forgetting.

4.3. SENTIMENT ANALYSIS FROM TEXT USING CNNS:

We use the "Amazon Reviews: Unlocked Mobile Phones" data set, which contains reviews of unlocked mobile phones sold on "Amazon.com". The goal is to find out the rating (1 to 5 stars) of each review using only the text itself. We use our model with a Convolutional Neural Network (CNN) that is suited to NLP tasks (Kim, 2014; Zahavy et al., 2018b) . Specifically, the architecture is a shallow word-level CNN that was demonstrated to provide state-of-the-art results on a variety of classification tasks by using word embeddings, while not being sensitive to hyperparameters (Zhang & Wallace, 2015) . We use the architecture with its default hyper-parameters (Github) and standard pre-processing (e.g., we use random embeddings of size 128, and we trim and pad each sentence to a length of 60). The only modification we made was to add a linear layer of size 50 to make the size of the last hidden layer consistent with our previous experiments.

-greedy Neural-Linear Neural-Linear Limited Memory 2963.9 ± 68.5 3155.6 ± 34.9 3143.9 ± 33.5

Figure 3 : Cumulative reward on Amazon review's Since the input is in R 7k (60 × 128), we did not include a linear baseline in these experiments as it is impractical to do linear algebra (e.g., calculate an inverse) in this dimension. Instead, we focused on comparing our final method with the full memory neural linear TS and both prior computations with an -greedy baseline. We experimented with 10 values of , ∈ [0.1, 0.2, ..., 1] and report the results for the value that performed the best (0.1). Looking at Fig. 3 we can see that the limited memory version performs almost as good as the full memory, and better than the -greedy baseline.

5. DISCUSSION:

We presented a neural-linear contextual bandit algorithm that is resilient to catastrophic forgetting and demonstrated its performance on several real-world data sets. Our algorithm showed comparable results to a previous method that stores all the data in a replay buffer. The method requires to solve an SDP, which is computationally prohibitive in general. Thus, we restricted the size of the last layer to be small, such that solving the SDP is feasible.

To design our algorithm, we assumed that all the representations that are produced by the DNN are realizable. In practice, the features that are learned in the first iterations are nearly realizable and further iterations improve them. We hope to relax these assumptions in future work.

Our algorithm presented excellent performance on multiple real-world data sets. Moreover, its performance did not deteriorate due to the changes in the representation and the limited memory. We believe that our findings make an important step towards solving problems where both exploration and representation learning play an important role. The realizability assumption gives us a method to compute µ 0 for the new features :

EQUATION

Similarly, using the analytically solution to the SDP, we get

EQUATION

Using Eq. ( 7) and Eq. ( 8) and rearranging we get that

Φ 0 µ 0 = = E n φ (E n φ old ) -1 E m φ old R m = (A φ b n )(A φ old b n ) -1 A φ old b m R m = A φ b m b -1 m A -1 φ old A φ old b m R m = A φ b m R m = E m φ R m .

Similarly, for Φ 0 we get that

Φ 0 = E n φ (E n φ old ) -1 (Φ old m )((E n φ old ) T ) -1 (E n φ ) T = (A φ b n )(A φ old b n ) -1 (Φ old m )(b -1 n A -1 φ old ) T (A φ b n ) T = A φ A -1 φ old A ψ m i=1 b i b T i A T φ old (A T φ old ) -1 A T φ = A φ m i=1 b i b T i A T φ = Φ m .

:

Image credits: bandit (bottom), Microsoft research; confidence ellipsoid (right), OriginLab.

We will provide a detailed comparison of the neural-linear algorithms and priors later in this section.

Figures and Tables:

Figure 1: Neural-Linear contextual Thompson sampling with limited memory.

Figure 2: Catastrophic forgetting

Cumulative reward of TS algorithms on 10 real world data sets. The context dim d and the size of the action space A are reported for each data set. The mean result and standard deviation of each algorithm is reported for 50 runs.

Table content: <table><tr><td/><td/><td/><td colspan="2">Full memory</td><td colspan="3">Limited memory, Neural-Linear</td></tr><tr><td>Name</td><td>d</td><td>A</td><td>Linear</td><td>Neural-Linear</td><td>Both Priors</td><td>µ Prior</td><td>No Prior</td></tr><tr><td/><td/><td/><td/><td colspan="2">Linear Data Sets</td><td/><td/></tr><tr><td>Mushroom</td><td colspan="2">117 2</td><td>11022 ± 774</td><td>10880 ± 853</td><td>10923 ± 839</td><td>9442 ± 1351</td><td>7613 ± 1670</td></tr><tr><td>Financial</td><td>21</td><td>8</td><td>4588 ± 587</td><td>4389 ± 584</td><td>4597 ± 597</td><td>4311 ± 598</td><td>4225 ± 594</td></tr><tr><td>Jester</td><td>32</td><td>8</td><td>14080 ± 2240</td><td>12819± 2135</td><td>9624 ± 2186</td><td>10996 ± 2013</td><td>11114 ± 2050</td></tr><tr><td>Adult</td><td>88</td><td>2</td><td>4066.1 ± 11.03</td><td>4010.0 ± 22.19</td><td>3943.0 ± 54.29</td><td>3839.5 ± 17.63</td><td>3608.2 ± 34.94</td></tr><tr><td>Covertype</td><td>54</td><td>7</td><td>3054 ± 557</td><td>2898 ± 545</td><td>2828 ± 593</td><td>2347 ± 615</td><td>2334 ± 603</td></tr><tr><td/><td/><td/><td/><td colspan="2">Nonlinear Data Sets</td><td/><td/></tr><tr><td>Census</td><td colspan="2">377 9</td><td>1791.5 ± 39.47</td><td>2135.5 ± 51.47</td><td>2023.16 ± 37.3</td><td>1873 ± 757</td><td>1943.83 ± 84.2</td></tr><tr><td>Statlog</td><td>9</td><td>7</td><td>4483 ± 353</td><td>4781 ± 274</td><td>4825 ± 305</td><td>4681 ± 285</td><td>4623 ± 276</td></tr><tr><td>Epileptic</td><td colspan="2">178 5</td><td>1202.9 ± 34.68</td><td>1706.9 ± 41.26</td><td>1716.8 ± 60.44</td><td>1572.9 ± 48.66</td><td>1411.0 ± 33.43</td></tr><tr><td>Smartphones</td><td colspan="2">561 6</td><td>3085.8 ± 24.64</td><td>3643.5 ± 64.89</td><td>2660.4 ± 84.72</td><td>3064.5 ± 55.06</td><td>2851.6 ± 58.77</td></tr><tr><td>Scania Trucks</td><td colspan="2">170 2</td><td>4691.8 ± 7.23</td><td>4784.7 ± 6.05</td><td>4742.0 ± 33.0</td><td>4698.0 ± 13.06</td><td>4470.4 ± 37.9</td></tr></table>

Cumulative reward on the wheel bandit APPENDIX B ANALYSIS B.1 DERIVATION OF AUXILIARY RESULTS FOR THE SANITY CHECK

Table content: <table><tr><td/><td>Linear</td><td>Neural-Linear Lim-ited Memory</td></tr><tr><td>δ=0.5</td><td>737.44 ± 3.04</td><td>899.72 ± 12.79</td></tr><tr><td>δ=0.3</td><td>735.37 ± 2.58</td><td>781.09 ± 11.34</td></tr><tr><td>δ=0.1</td><td>735.51 ± 2.59</td><td>751.75 ± 3.6</td></tr></table>

Peer review (iclr2020):
This paper adapts Bayesian linear regression to the setting of a limited memory replay buffer. The idea is to calibrate the prior mean and variance when the neural representation of context is updated. Overall the paper is well written and explained clearly. Some experiments are provided to show that the proposed method is able to achieve a performance competitive to Bayesian linear regression with infinite memory. The result of this paper is interesting. But I am not sure if the current experimental results are convincing enough to justify the significance of the proposed method. 1. The results in Section 4.2 seem to be following the setting in Riquelme 2018. These datasets are all in a supervised learning setting. It is a bit disappointing that the proposed method is not tested on RL datasets. 2. No other baseline is provided in the experiments for comparison. a. There are other methods in the literature to overcome catastrophic forgetting of neural networks, e.g regularizing the update of the network. How would that be compared to the proposed method? b. What about other methods, like [1]? 3. Most of the experiment details are missing. For example, how is the reward defined in section 4.3? What is the overhead in computation in practice, especially for the SDP? Other comments: 1. Why would solving a SDP require only O(g^{0.5}) in section 3.1? 2. In the discussion in section 3, even if equation (5) and (6) can be exactly solved, how does the heavy tailed problem mentioned in section 2 been solved? [1] Elmachtoub, Adam N., et al. *A practical method for solving contextual bandit problems using decision trees.* arXiv preprint arXiv:1706.04687 (2017).