Optimal Query Selection Using Multi-Armed Bandits

Abstract

Query selection for latent variable estimation is conventionally performed by opting for observations with low noise or optimizing information theoretic objectives related to reducing the level of estimated uncertainty based on the current best estimate. In these approaches, typically the system makes a decision by leveraging the current available information about the state. However, trusting the current best estimate results in poor query selection when truth is far from the current estimate, and this negatively impacts the speed and accuracy of the latent variable estimation procedure. We introduce a novel sequential adaptive action value function for query selection using the multi-armed bandit (MAB) framework which allows us to find a tractable solution. For this adaptive-sequential query selection method, we analytically show: (i) performance improvement in the query selection for a dynamical system, (ii) the conditions where the model outperforms competitors. We also present favorable empirical assessments of the performance for this method, compared to alternative methods, both using Monte Carlo simulations and human-in-the-loop experiments with a brain computer interface (BCI) typing system where the language model provides the prior information.

I. Introduction

Recursive state estimation contributes a key role in signal processing and system identification. The recursive paradigm is often used to extract information from model parameters or the states of a dynamic system in real time, given noisy observations. On the other hand, Bayesian methods are valuable decision making approaches, since they take into account a variety of prior knowledge about the system due to the experience and previous observations (history). In stochastic dynamic systems, to estimate the state variables, maximum a posteriori (MAP) inference is commonly used. To estimate the state with a high confidence (usually pre-defined), the system probes the environment through multiple recursions of sequences of queries, which reduces the rate of state estimation convergence. Therefore, the queries need to be designed specifically to optimize both the speed and the accuracy of the state estimation. Query selection/optimization in the recursive state estimation is often performed by greedy selections using: (i) expected posterior maximization [1]; (ii) Fisher information-based approaches [2], [3], and (iii) information theory-based approaches such as entropy minimization or maximum mutual information (MMI) [4], [5], [6], [7]. It is shown that all these approaches for optimum sequence design through query selection lead to the selection of N-best queries with respect to the current belief [8], [9].

In estimation problems, the system may have access to an additional knowledge called prior information on the state of the system which can improve the estimation process. However, imprecise prior information may lead to incorrect posterior beliefs given the same set of observations. Accordingly, choosing the N-best queries by trusting the current belief does not always offer the best query optimization. In dynamic systems, the prior information about the environment can be adversarial due to the transition noise, observation noise, change of environment distributions and being outdated; hence resulting in longer decision cycles or the wrong state estimation [10], [11]. Many applications involving state estimation, system identification or sequential decision making using prior information encounter these challenges: recommender systems [12], coimnunication networks [13], [14], radar systems [15] and clinical studies [16], [17]. The common problem in all of these applications is that the misleading information can extremely impact the final decision or estimation. Another category of methods to overcome the misleading information, is variance based methods [15], [18], [19] that can be extended using Fisher information [20] to either explore or exploit. The main drawback of these methods is that they only commit to either explore or exploit for the query selection, which leads to the same solution provided by the N-best method [8], [9].

In this paper, we propose an information theoretic query selection to discard the ambiguity (exploitation) in the state estimation while also measuring the credibility of the prior information (exploration). The proposed objective function is a linear combination of exploration and exploitation. Moreover, we reformulate the query selection as a multi-armed bandit (MAB) problem. We denote this framework as MAB based on State-Measurement-MI and State-Posterior-Momentum for RBSE. The MAB framework is a well-studied approach to formulate the learning process under available observations [21], [22], [23], [24]. MAB has been proposed for decision making, predictive entropy search for sequential action selection and estimation applications [25], [26], [27], [28], [29], [30]. Such applications consider sequential selection based on MAB, in which the approaches can be considered to repeatedly make choices among elements of a finite set of state elements [28], [25]. Such a formulation enables us to analytically demonstrate that the proposed policy for query selection containing exploration and exploitation performs at least as good as the methods that only rely on the exploitation of the current belief. MAB framework also enables the formulation of subset query selection optimization as a tractable problem especially through sequential selection with theoretical guarantees [31], [21].

The novel contributions of this paper can be summarized in: (i) introducing a new action-value function for query selection using changes in the posterior to encourage exploration in the MAB setting, (ii) providing a short-term policy evaluation to demonstrate that the proposed method has theoretical guarantees under certain assumptions, and (iii) evaluating the proposed method in an actual human-in-the-loop typing scheme employing a language-model-assisted Electroencephalogram (EEG)- based BCI typing system called RSVP Keyboard. Because of space limitations, we present the proofs of analytical propositions in the supplementary material. The system code is under revision and the current version can be accessed at https://github.com/BciPy [32].

II. Preliminaries

In the framework of the state estimation problem, We refer σ as the (unknown) state which is an element of a finite set $\mathcal{A}$. The system (learner) proceeds with the estimation through a sequential decision making process containing sequences indexed by s of multiple trials indexed by i. We denote a list of variables with {:}, for example Φ_0:s represents a sequence of variables from 0 to s. The query and evidence sets at sequence s are denoted by ${\mathrm{\Phi }}_s\triangleq \{{\varphi }_s^1,\dots ,{\varphi }_s^K\}$ and ${\epsilon }_s\triangleq \{{\epsilon }_s^1,\dots ,{\epsilon }_s^K\}$, respectively. Here, $K\in \mathbb{N}$ denotes number of trials in a sequence. We use the query class definition presented in [8] and assume each σ has a corresponding query defined with the class conditional representation. Therefore, without loss of generality, it is assumed that all of the observations are noisy and comes from two unimodal probability distributions conditioned on state and query tuples. Assuming all trials are independent and the current observation is only function of the current query and independent of the task history, ${\mathcal{H}}_s\triangleq \{{\epsilon }_{1:s},{\mathrm{\Phi }}_{1:s}\}$, the posterior probability at time s can be expressed as:

$$p(\sigma \mid {\epsilon }_s,{\mathrm{\Phi }}_s,{\mathcal{H}}_{s-1})=p(\sigma \mid {\mathcal{H}}_0)\prod _{j=1}^s\prod _{i=1}^{t_i}\frac{p({\epsilon }_j^i\mid \sigma ,{\varphi }_j^i)}{p({\epsilon }_j^i\mid {\varphi }_j^i)}$$

where $p(\sigma \mid {\mathcal{H}}_0)$ is a prior information. Using maximum a posteriori (MAP) estimation [16] the learner attempts to estimate σ with re-occuring evidence collection. Based on the collected evidence if a decision is not possible, the system decides on a subset of queries for the upcoming sequence to improve its confidence. Accordingly, the query selection process is formulated by the following optimization:

$${\mathrm{\Phi }}_s=\text{arg}\underset{\mathrm{\Phi }}{\text{max}}{q}_s({\epsilon }_{1:s-1},{\mathrm{\Phi }}_{1:s-1},\mathrm{\Phi })$$ (1)

where q_s denotes the objective term, which we call it action-value function. Following querying, evidence ε_s is observed and accordingly the posterior is updated. In next section we propose an action value function that balances exploration and exploitation.

III. Method and Analysis

By imposing the MAB settings to the context of the state estimation problem, each query can be represented as an arm of a MAB. This reformulation allows us to solve the subset selection problem through a greedy approach by optimizing the action value function for each arm with theoretical guarantees [31]. Therefore, selection of arms (queries in each sequence) with highest action value one by one allows us to form the query in a computationally efficient way. In MAB formulation, for the design of upcoming sequence, we assume that multiple arms are pulled according to the state posterior probability that depends on the task history. The goal is then to define an objective that specifies the subset of queries (arms) to be picked at each step. Independency between trails allows us to perform optimization per trial at each sequence. Therefore, set optimization in (1) reduces to single query selection.

Conventionally, query selection is achieved through MMI [6], [5], which is equivalent to entropy minimization when the evidence ε_s corresponding to the sequence being designed is not observed; and hence, σ is independent from ϕ_s. Query selection using mutual information can be written as the following policy:

$$\begin{array}{c}{\varphi }_s^i=\text{arg}\underset{\varphi }{\max }I(\sigma ,{\epsilon }_s^i\mid \varphi ,{\mathcal{H}}_{s-1})\hfill \\ =\text{arg}\underset{\varphi }{\max }-H(\sigma \mid {\epsilon }_s^i,\varphi ,{\mathcal{H}}_{s-1})\hfill \end{array}$$ (2)

In this paper, we consider three different action-value functions for the MAB formulation: (i) mutual information objective (2) (ii) history-based objective (iii) combination of (i) and (ii).

We introduce a new term called Momentum that is function of posterior changes across sequences, such that:

$$m(\varphi \mid {\mathcal{H}}_j)={\mathbf{E}}_{p(\sigma \mid {\mathcal{H}}_{j-1})}\left[\log p(\sigma \mid {\epsilon }_j^i,\varphi ,{\mathcal{H}}_{j-1})-\log p(\sigma \mid {\mathcal{H}}_{j-1})\right]{\mathbb{1}}_{\varphi }(\sigma )$$ (3)

where $1_{\varphi }(\sigma )$ denotes the indicator function which equates 1 if ϕ = σ Since $m(\varphi \mid {\mathcal{H}}_j)$ is the summation of probability displacement multiplied by the probability mass along axes of the state space, we call it Momentum. Additionally, $m(\varphi \mid {\mathcal{H}}_0)=0,\forall \varphi$, in words, without collecting any evidence we can not infer the trend of a particular state in estimation. Accordingly, for the history-based approach, the objective is defined as the average of Momentum as follows.

$$M(\varphi \mid {\mathcal{H}}_{s-1})=\frac{1}{s-1}\sum _{j=1}^{s-1}m(\varphi \mid {\mathcal{H}}_j){\mathbb{1}}_{{\mathrm{\Phi }}_j}(\varphi )$$ (4)

For this approach, the query selection policy becomes:

$${\varphi }_s^i=\text{arg}\underset{\varphi }{\max }M(\varphi \mid {\mathcal{H}}_{s-1})$$ (5)

We present a new action-value function for the query selection based on combination of mutual information and history-based objectives in (4) to balance exploration and exploitation. Accordingly, the action-value function and policy can be defined as:

$$q_s^i(\varphi )=I(\sigma ,{\epsilon }_s^i\mid \varphi ,{\mathcal{H}}_{s-1})+\lambda M(\varphi \mid {\mathcal{H}}_{s-1}),\lambda \ge 0$$ (6)

$$\begin{array}{c}{\varphi }_s^i=\text{arg}\underset{\varphi }{\max }{q}_s(\varphi )\hfill \\ =\text{arg}\underset{\varphi }{\max }-H(\sigma \mid {\epsilon }_s^i,\varphi ,{\mathcal{H}}_{s-1})+\lambda M(\varphi \mid {\mathcal{H}}_{s-1})\hfill \end{array}$$ (7)

where, λ is a tuning parameter that balances MMI and Momentum-based policies. The objective function (7) can be written as (8) by replacing the entropy term as shown in our previous work [8].

$$\begin{array}{c}{\varphi }_s^i=\text{arg}\underset{\varphi }{\max }{\mathbf{E}}_{p(\sigma \mid {\mathcal{H}}_{s-1)}}{\mathbf{E}}_{p({\epsilon }_s^i\mid \sigma ,\varphi )}[\log p({\epsilon }_s^i\mid \sigma ,\varphi )\phantom{]}\hfill \\ -\phantom{[}\log p({\epsilon }_s^i\mid \varphi )]+\frac{\lambda }{s-1}\sum _{j=1}^{s-1}m(\varphi \mid {\mathcal{H}}_j){\mathbb{1}}_{{\mathrm{\Phi }}_j}(\varphi )\hfill \end{array}$$ (8)

For the same given task history, we show that when policy in (8) is used in MAB formulation, the target state has higher probability to be chosen to appear in the query subset compared to other policies in (2) and (5). Here we propose the analysis of the correctness of the statement. To save space we use $I(a,{\epsilon }_s^i\mid \varphi =a,{\mathcal{H}}_{s-1})=I_s(a)$ and $M(\varphi =b\mid {\mathcal{H}}_{s-1})=M_s(a)$ notations for the following lemmas.

Lemma 1. Given $a,b\in \mathcal{A}$ where a ≠ b and λ ≥ 0, if $\exists {\mathcal{H}}_{s-1}$ s.t. $p(a\mid {\mathcal{H}}_{s-1})<p(b\mid {\mathcal{H}}_{s-1})$, then

$$p\left(I_s(a)+\lambda M_s(a)>I_s(b)+\lambda M_s(b)\right)\ge p\left(I_s(a)>I_s(b)\right)$$

Lemma 1 shows that the probability of a (assuming to be the target state) having a higher action-value than b is larger when policy (8) is used instead of (2). Although the probability of a given, the task history is lower than the probability of b given the task history. This means that even if a is less likely than b according to the prior information and observations, using the proposed policy, a has more chance to appear in the query subset compared to policy in (2).

Lemma 2. Given $a,b\in \mathcal{A}$ where a ≠ b and λ ≥ 0, if $\exists {\mathcal{H}}_{s-1}$ s.t. $p(a\mid {\mathcal{H}}_{s-1})>p(b\mid {\mathcal{H}}_{s-1})$, then

$$p\left(I_s(a)+\lambda M_s(a)>I_s(b)+\lambda M_s(b)\right)\ge p\left(M_s(a)>M_s(b)\right)$$

Lemma 2 represents the case where the prior knowledge is supporting a rather than being adversarial. It shows that compared to using (4) as the action-value function, when a (unknown state) has higher probability given the task history compared to b, using action value function (6) has higher probability to choose a. These two lemmas show that the proposed query selection policy provides a balance between (2) and (5), and accordingly between the adversarial and supporting prior information. This balance is achieved through the selection of λ. Detailed proofs of Lemmas are provided in the Supplementary Materials.

The query λ should be updated dynamically such that the emphasis on mutual information component of the proposed function is increased with the number of sequences; i.e., the λ value should be decreased as the number of sequences is increasing [33]. We introduce a theorem that defines upper and lower bounds for the λ value to satisfy to guarantee that including the target state in the selected subset has higher probability when the proposed policy in (8) is used in MAB formulation compared to other policies (2) and (5).

Theorem 1. Let $\sigma \in \mathcal{A}$ be the target state and $\mid \mathcal{A}\mid$ represent the size of the finite set $\mathcal{A}$. Consider three query selection policies as follows:

$$\begin{array}{cc}\hfill {\varphi }_{i,s}^{{\pi }_1}& =\text{arg}\underset{\varphi }{\max }I(\sigma ,{\epsilon }_s^i\mid \varphi ,{\mathcal{H}}_{s-1})\hfill \\ \hfill {\varphi }_{i,s}^{{\pi }_2}& =\text{arg}\underset{\varphi }{\max }M(\varphi \mid {\mathcal{H}}_{s-1})\hfill \\ \hfill {\varphi }_{i,s}^{{\pi }_3}& =\text{arg}\underset{\varphi }{\max }I(\sigma ,{\epsilon }_s^i\mid \varphi ,{\mathcal{H}}_{s-1})+\lambda M(\varphi \mid {\mathcal{H}}_{s-1}).\hfill \end{array}$$

If p(ϕ^π3 = σ) ≥ p(ϕ^πi = σ) for i = 1, 2, then ∃λ that satisfies

$$\begin{array}{cc}\hfill & \frac{2(s-1)\mid \mathcal{A}\mid (d_{(p_s^{{\varphi }^{{\pi }_1}},U)}^2-d_{(p_s^{{\varphi }^{{\pi }_3}},U)}^2)}{\sum _{j=1}^{s-1}[d_{p_{i,j}^{{\varphi }^{{\pi }_3}}}^2-d_{p_{i,j}^{{\varphi }^{{\pi }_1}}}^2]{\mathbb{1}}_{\varphi }(\sigma )}\le \lambda \le \hfill \\ \hfill & \frac{(s-1)(d_{(p_s^{{\pi }_3},U)}^2-d_{(p_s^{{\pi }_2},U)}^2)}{2\sum _{j=1}^{s-1}[\log (\frac{p(\sigma \mid {\epsilon }_j^i,{\varphi }^{{\pi }_2},{\mathcal{H}}_{j-1})}{p(\sigma \mid {\epsilon }_j^i,{\varphi }^{{\pi }_3},{\mathcal{H}}_{j-1})})]}\hfill \end{array}$$ (9)

where

$$\begin{array}{c}{d}_{(p_{i,s}^{\pi },U)}^2=\sum _{a\in \mathcal{A}}{\left(p(a\mid {\epsilon }_s^i,{\varphi }^{\pi },{\mathcal{H}}_{s-1})-\frac{1}{\mid \mathcal{A}\mid }\right)}^2\hfill \\ d_{p_{i,s}^{\pi }}^2=\mid p(\sigma \mid {\mathcal{H}}_{s-1})-p(\sigma \mid {\epsilon }_s^i,{\varphi }^{\pi },{\mathcal{H}}_{s-1}){\mid }^2.\hfill \end{array}$$

This theorem shows the existence of the parameter λ makes the joint objective optimal. The proof of this theorem, which uses Pinsker’s Theorem [34] and the results of Lemmas 1 and 2 can be found in the Supplementary Materials.

IV. Results

To assess the performance of the proposed query selection method, an language-model-assisted EEG-based BCI typing system called RSVP Keyboard^™ [35] has been used as a test framework. Ten healthy participants (six females and four males), 20-35 years old were recruited under IRB-130107 protocol approved by Northeastern University. A DSI-24 Wearable Sensing EEG Headset was used for data acquisition, at a sampling rate of 300 Hz with active dry electrodes. EEG signals were acquired from 20 sensors according to International 10-20 System locations: Fp1, Fp2, Fz, F3, F4, F7, F8, Cz, C3, C4, T3, T4, T5, T6, P3, P4, O1, O2, A1 and A2. All participants were asked to perform two sessions including calibration and Copy Phrase. During calibration, the users were asked to attend to predefined target symbols within randomly ordered sequences to enable the system learn the class conditional EEG evidence distributions. Here, calibration session contains 100 sequences; each sequence includes five trials (letters); and one trial in each sequence is the target symbol which is displayed on the screen prior to each sequence. The time interval between trials is 200 ms. Optimal parameters for both target and non-target class distributions were learned using the calibration data which are used in Copy Phrase task. In Copy Phrase, participants were tasked to write a missing word in a pre-defined phrase using the system (a total of 6 words with various difficulties based on LM were typed). In addition to Copy Phrase, we also use the calibration data from each participant for BCI performance simulation. We present both simulation and realtime experiment results.

To evaluate the empirical performance of the proposed query selection we first utilized simulation. For our simulations we used conditional evidence distributions which are learned in the calibration session. More details about the simulation framework can be found in [35] study. During simulations, through Monte Carlo simulations samples from these conditional distributions were drawn to type ‘O’ and ‘C’ for each simulation in the phrase ‘IT_OCCURRED_RANDOMLY’. These letters are chosen because they have different difficulties to be typed based on the language model (prior information). The number of Monte-Carlo simulations is chosen to be 500. Figure 1 shows the simulation results. Figure 1 presents the typing performance for two users with different calibration performance which is quantified by the area under the receiver operating characteristics curve (AUC); AUC_U1 = 0.82 and AUC_u7 = 0.67. The bar plot next to each plot shows the prior information provided by the language model (LM) at the beginning of a decision cycle. For instance, it can be seen that the LM probability for ‘O’ is very low and it is not quite likely to start a word with this letter. Accordingly, MMI method highly influenced by the LM prior, needs more sequences to estimate the target letter. In the early sequences of the decision process, the Momentum-based approach on average is faster than MMI to pick the intended letter for the query subset. Although, due to noise in the EEG evidence and miss-classification of observations, Momentum gets close to zero and does not pick the intended state for the query subset. Overall, the proposed method outperforms the other methods. However, when there is a likely letter like ‘C’, MMI and the proposed method perform similarly. By comparing the simulation results of user 1 (lower AUC) with user 7 (higher AUC), it can be seen that all of the query selections are faster for the user 7 with higher calibration performance. It can also be observed that for user 1 even when there is more overlap between class conditional distributions (because of low AUC), the proposed method can estimate the target letter quite fast. On the other hand, it is more difficult for Momentum-based approach to capture the target letter. Moreover, MMI also requires more number of sequences.

Fig. 1:

Probability of the letter completion for 500 Monte-Carlo simulations for typing three target symbols in phrase ‘occurred’. Intended symbols contain: (a)‘O’ and (b)‘C’ in the target phrase. Simulation results are reported for two users with different calibration performances. User 1 with AUC = 0.67 has lower performance than user 7 with AUC = 0.82. Bar plots show the LM prior probability over all typing symbols.

As described above, the participants also attended Copy Phrase sessions after the calibration. Each participant attended four Copy Phrase tasks with different query selection methods. The order of the tasks were randomly assigned for each participant to avoid the learning impact.

Figure 2 shows the average performance of all the query selection methods for all users including a statistical test results for Copy Phrase sessions. We reported the query selection performance in terms of two measures: accuracy in typing a letter correctly (ATL) that is the total number of correctly typed letters divided by the total number of typed letters and the information transfer rate (ITR) [36]. ITR summarizes the accuracy and speed into a single metric and it is commonly used to measure BCI performance. These results show that the proposed method outperforms the other methods both in terms of speed and accuracy. We used the Wilcoxon signed-rank test as a non-parametric statistical hypothesis test to perform a paired-comparison between the proposed method and the other query selection methods. The proposed method significantly enhanced the ITR compared to the other methods. Our statistical analysis also shows that the proposed method significantly improved the ATL compared to MMI and random query selection.

Fig. 2:

Average of information transfer rate for four query selection methods. All of the results belongs to 10 users attending the copy phrase task in RSVP Keyboard experiment. p corresponds to the Wilcoxon signed-rank test.

V. Conclusion

Being motivated by the MAB framework, we proposed a tractable solution to the subset query optimization for recursive Bayesian state estimation to enhance the estimation speed and accuracy. More specifically, a new action-value function was introduced for query selection, which uses a linear combination of mutual information and a momentum term which is a function of logarithmic changes of the posterior probability across sequences. We have also presented a bound for the action-value tuning parameter, which guarantees that the proposed query selection policy outperforms the others. An BCI typing system has been used as a test framework to assess the performance of the proposed method. Our results for both simulation and the human-in-the-loop experiment showed that the proposed method outperforms the alternative methods as shown by analytical results.