Communication Efficient Distributed Learning for
Kernelized Contextual Bandits

We consider an unknown reward function $f$ that lies in a RKHS, denoted as $\mathcal{H}$, such that the reward can be equivalently written as

where $\mathbf{\theta}_{\star}\in\mathcal{H}$ is an unknown parameter, and $\phi:\mathbb{R}^{d}\rightarrow\mathcal{H}$ is a known feature map associated with $\mathcal{H}$. We assume $\eta_{t}$ is zero-mean $R$-sub-Gaussian conditioned on $\sigma\bigl{(}(\mathbf{x}_{s},\eta_{s})_{s\in\mathcal{N}_{t-1}(i_{t})},\mathbf{x}_{t}\bigr{)},\forall t$, which denotes the $\sigma$-algebra generated by client $i_{t}$’s previously pulled arms and the corresponding noise. In addition, there exists a positive definite kernel $k(\cdot,\cdot)$ associated with $\mathcal{H}$, and we assume $\forall\mathbf{x}\in\mathcal{A}$ that, $\lVert\mathbf{x}\rVert_{k}\leq L$ and $\lVert f\rVert_{k}\leq S$ for some $L,S>0$.

The goal of the learning system is to minimize the cumulative (pseudo) regret for all $N$ clients, i.e., $R_{NT}=\sum_{t=1}^{NT}r_{t}$, where $r_{t}=\max_{\mathbf{x}\in\mathcal{A}_{t}}\phi(\mathbf{x})^{\top}\mathbf{\theta}_{\star}-\phi(\mathbf{x}_{t})^{\top}\mathbf{\theta}_{\star}$. Meanwhile, the learning system also wants to keep the communication cost $C_{NT}$ low, which is measured by the total number of scalars being transferred across the system up to time step $NT$.

For each round $l\in[T]$, when client $i\in[N]$ interacts with the environment, i.e., the $t$-th interaction between the learning system and the environment where $t=N(l-1)+i$, it chooses arm $\mathbf{x}_{t}\in\mathcal{A}_{t}$ based on the UCB of the mean estimator (line 5):

where $\hat{\mu}_{t,i}(\mathbf{x})$ and $\hat{\sigma}^{2}_{t,i}(\mathbf{x})$ denote client $i$’s local estimated mean reward for arm $\mathbf{x}\in\mathcal{A}$ and its variance, and $\alpha_{t-1,i}$ is a carefully chosen scaling factor to balance exploration and exploitation (see Lemma 1 for proper choice).

To facilitate further discussion, for time step $t\in[NT]$, we denote the sequence of time indices for the data points that have been used to update client $i$’s local estimate as $\mathcal{D}_{t}(i)$, which include both data points collected locally and those shared by the other clients. If the clients never communicate, $\mathcal{D}_{t}(i)=\mathcal{N}_{t}(i),\forall t,i$; otherwise, $\mathcal{N}_{t}(i)\subset\mathcal{D}_{t}(i)\subseteq[t]$, with $\mathcal{D}_{t}(i)=[t]$ recovering the centralized setting, i.e., each new data point collected from the environment immediately becomes available to all the clients in the learning system. The design matrix and reward vector for client $i$ at time step $t$ are denoted by $\mathbf{X}_{\mathcal{D}_{t}(i)}=[\mathbf{x}_{s}]_{s\in\mathcal{D}_{t}(i)}^{\top}\in\mathbb{R}^{|\mathcal{D}_{t}(i)|\times d},\mathbf{y}_{t,i}=[y_{s}]_{s\in\mathcal{D}_{t}(i)}^{\top}\in\mathbb{R}^{|\mathcal{D}_{t}(i)|}$, respectively. By applying the feature map $\phi(\cdot)$ to each row of $\mathbf{X}_{\mathcal{D}_{t}(i)}$, we obtain $\mathbf{\Phi}_{\mathcal{D}_{t}(i)}\in\mathbb{R}^{|\mathcal{D}_{t}(i)|\times p}$, where $p$ is the dimension of $\mathcal{H}$ and is possibly infinite. Since the reward function is linear in $\mathcal{H}$, client $i$ can construct the Ridge regression estimator $\hat{\theta}_{t,i}=(\mathbf{\Phi}_{\mathcal{D}_{t}(i)}^{\top}\mathbf{\Phi}_{\mathcal{D}_{t}(i)}+\lambda I)^{-1}\mathbf{\Phi}_{\mathcal{D}_{t}(i)}^{\top}\mathbf{y}_{t,i}$, where $\lambda>0$ is the regularization coefficient. This gives us the estimated mean reward and variance in primal form for any arm $\mathbf{x}\in\mathcal{A}$, i.e., $\hat{\mu}_{t,i}(\mathbf{x})=\phi(\mathbf{x})^{\top}\mathbf{A}_{t,i}^{-1}\mathbf{b}_{t,i}$ and $\hat{\sigma}_{t,i}(\mathbf{x})=\sqrt{\phi(\mathbf{x})^{\top}\mathbf{A}_{t,i}^{-1}\phi(\mathbf{x})}$, where $\mathbf{A}_{t,i}=\mathbf{\Phi}_{\mathcal{D}_{t}(i)}^{\top}\mathbf{\Phi}_{\mathcal{D}_{t}(i)}+\lambda\mathbf{I}$ and $\mathbf{b}_{t,i}=\mathbf{\Phi}_{\mathcal{D}_{t}(i)}^{\top}\mathbf{y}_{t,i}$. Then using the kernel trick, we can obtain their equivalence in the dual form that only involves entries of the kernel matrix, and avoids directly working on $\mathcal{H}$ which is possibly infinite:

where $\mathbf{K}_{\mathcal{D}_{t}(i)}(\mathbf{x})=\mathbf{\Phi}_{\mathcal{D}_{t}(i)}\phi(\mathbf{x})=[k(\mathbf{x}_{s},\mathbf{x})]^{\top}_{s\in\mathcal{D}_{t}(i)}\in\mathbb{R}^{|\mathcal{D}_{t}(i)|}$, and $\mathbf{K}_{\mathcal{D}_{t}(i),\mathcal{D}_{t}(i)}=\mathbf{\Phi}_{\mathcal{D}_{t}(i)}^{\top}\mathbf{\Phi}_{\mathcal{D}_{t}(i)}=[k(\mathbf{x}_{s},\mathbf{x}_{s^{\prime}})]_{s,s^{\prime}\in\mathcal{D}_{t}(i)}\in\mathbb{R}^{|\mathcal{D}_{t}(i)|\times|\mathcal{D}_{t}(i)|}$.

To reduce the regret in future interactions with the environment, the $N$ clients need to collaborate via communication, and a carefully designed communication protocol is essential in ensuring the communication efficiency. In prior works like DisLinUCB Wang et al. (2019), after each round of interaction with the environment, client $i$ checks whether the event $\{(|\mathcal{D}_{t}(i)|-|\mathcal{D}_{t_{\text{last}}}(i)|)\log(\frac{\det(\mathbf{A}_{t,i})}{\det(\mathbf{A}_{t_{\text{last}},i})})>D\}$ is true, where $t_{\text{last}}$ denotes the time step of last global synchronization. If true, a new global synchronization is triggered, such that the server will require all clients to upload their sufficient statistics since $t_{\text{last}}$, aggregate them to compute $\{\mathbf{A}_{t},\mathbf{b}_{t}\}$, and then synchronize the aggregated sufficient statistics with all clients, i.e., set $\{\mathbf{A}_{t,i},\mathbf{b}_{t,i}\}=\{\mathbf{A}_{t},\mathbf{b}_{t}\},\forall i\in[N]$.

Using kernel trick, we can obtain an equivalent event-trigger in terms of the kernel matrix,

where $D>0$ denotes the predefined threshold value. If event $\mathcal{U}_{t}(D)$ is true (line 7), a global synchronization is triggered (line 7-10), where the local datasets of all $N$ clients are synchronized to $\{(\mathbf{x}_{s},y_{s})\}_{s\in[t]}$. We should note that the transfer of raw data $(\mathbf{x}_{s},y_{s})$ is necessary for the update of the kernel matrix and reward vector in line 6 and line 10, which will be used for arm selection at line 5. This is an inherent disadvantage of kernelized estimation in distributed settings, which, as we mentioned in Section 2, is also true for the existing distributed kernelized bandit algorithm Dubey et al. (2020). Lemma 1 below shows that in order to obtain the optimal order of regret, DisKernelUCB incurs a communication cost linear in $T$ (proof given in the appendix), which is expensive for an online learning problem.

For each round $l\in[T]$, when client $i\in[N]$ interacts with the environment, i.e., the $t$-th interaction between the learning system and the environment where $t:=N(l-1)+i$, instead of using the UCB for the exact estimator in Eq (1), client $i$ chooses arm $\mathbf{x}_{t}\in\mathcal{A}_{t}$ that maximizes the UCB for the approximated estimator (line 5):

where $\tilde{\mu}_{t-1,i}(\mathbf{x})$ and $\tilde{\sigma}_{t-1,i}(\mathbf{x})$ are approximated using Nyeström method, and the statistics used to compute these approximations are much more efficient to communicate as they scale with the maximum information gain $\gamma_{NT}$ instead of $T$.

Specifically, Nyström method works by projecting some original dataset $\mathcal{D}$ to the subspace defined by a small representative subset $\mathcal{S}\subseteq\mathcal{D}$, which is called the dictionary. The orthogonal projection matrix is defined as

We then take eigen-decomposition of $\mathbf{K}_{\mathcal{S},\mathcal{S}}=\mathbf{U}\mathbf{\Lambda}\mathbf{U}^{\top}$ to rewrite the orthogonal projection as $\mathbf{P}_{\mathcal{S}}=\mathbf{\Phi}_{\mathcal{S}}^{\top}\mathbf{U}\mathbf{\Lambda}^{-1/2}\mathbf{\Lambda}^{-1/2}\mathbf{U}^{\top}\mathbf{\Phi}_{\mathcal{S}}$, and define the Nyström embedding function

which maps the data point $\mathbf{x}$ from $\mathbb{R}^{d}$ to $\mathbb{R}^{|\mathcal{S}|}$.

Therefore, we can approximate the Ridge regression estimator in Section 3.2 as $\tilde{\theta}_{t,i}=\tilde{\mathbf{A}}_{t,i}^{-1}\tilde{\mathbf{b}}_{t,i}$, where $\tilde{\mathbf{A}}_{t,i}=\mathbf{P}_{\mathcal{S}}\mathbf{\Phi}_{\mathcal{D}_{t}(i)}^{\top}\mathbf{\Phi}_{\mathcal{D}_{t}(i)}\mathbf{P}_{\mathcal{S}}+\lambda\mathbf{I}$, and $\tilde{\mathbf{b}}_{t,i}=\mathbf{P}_{\mathcal{S}}\mathbf{\Phi}_{\mathcal{D}_{t}(i)}^{\top}\mathbf{y}_{\mathcal{D}_{t}(i)}$, and thus the approximated mean reward and variance in Eq (3) can be expressed as $\tilde{\mu}_{t,i}(\mathbf{x})=\phi(\mathbf{x})^{\top}\tilde{\mathbf{A}}_{t,i}^{-1}\tilde{\mathbf{b}}_{t,i}$ and $\tilde{\sigma}_{t,i}(\mathbf{x})=\sqrt{\phi(\mathbf{x})^{\top}\tilde{\mathbf{A}}_{t,i}^{-1}\phi(\mathbf{x})}$, and their kernelized representation are (see appendix for detailed derivation)

where $\mathbf{Z}_{\mathcal{D}_{t}(i);\mathcal{S}}\in\mathbb{R}^{|\mathcal{D}_{t}(i)|\times|\mathcal{S}|}$ is obtained by applying $z(\cdot;\mathcal{S})$ to each row of $\mathbf{X}_{\mathcal{D}_{t}(i)}$, i.e., $\mathbf{Z}_{\mathcal{D}_{t}(i);\mathcal{S}}=\mathbf{\Phi}_{\mathcal{D}_{t}(i)}\mathbf{P}_{\mathcal{S}}^{1/2}$. We can see that the computation of $\tilde{\mu}_{t,i}(\mathbf{x})$ and $\tilde{\sigma}_{t,i}(\mathbf{x})$ only requires the embedded statistics: matrix $\mathbf{Z}_{\mathcal{D}_{t}(i);\mathcal{S}}^{\top}\mathbf{Z}_{\mathcal{D}_{t}(i);\mathcal{S}}\in\mathbb{R}^{|\mathcal{S}|\times|\mathcal{S}|}$ and vector $\mathbf{Z}_{\mathcal{D}_{t}(i);\mathcal{S}}^{\top}\mathbf{y}_{\mathcal{D}_{t}(i)}\in\mathbb{R}^{|\mathcal{S}|}$, which, as we will show later, makes joint kernelized estimation among $N$ clients much more efficient in communication.

After obtaining the new data point $(\mathbf{x}_{t},y_{t})$, client $i$ immediately updates both $\tilde{\mu}_{t-1,i}(\mathbf{x})$ and $\tilde{\sigma}_{t-1,i}(\mathbf{x})$ using the newly collected data point $(\mathbf{x}_{t},y_{t})$, i.e., by projecting $\mathbf{x}_{t}$ to the finite dimensional RKHS spanned by $\mathbf{\Phi}_{\mathcal{S}_{t_{\text{last}}}}$ (line 6). Recall that, we use $\mathcal{N}_{t}(i)$ to denote the sequence of indices for data collected by client $i$, and denote by $\mathcal{D}_{t}(i)$ the sequence of indices for data that has been used to update client $i$’s model estimation $\tilde{\mu}_{t,i}$. Therefore, both of them need to be updated to include time step $t$.

With the approximated estimator, the size of message being communicated across the learning system is reduced. However, a carefully designed event-trigger is still required to minimize the total number of global synchronizations up to time $NT$. Since the clients can no longer evaluate the exact kernel matrices in Eq (2), we instead use the event-trigger in Eq (4), which can be computed using the approximated variance from last global synchronization as,

Similar to Algorithm 1, if Eq (4) is true, global synchronization is triggered, where both the dictionary and the embedded statistics get updated. During synchronization, each client first samples a subset $\mathcal{S}_{t}(i)$ from $\mathcal{N}_{t}(i)$ (line 8) using Ridge Leverage Score sampling (RLS) (Calandriello et al., 2019, 2020), which is given in Algorithm 3, and then sends $\{(\mathbf{x}_{s},y_{s})\}_{s\in\mathcal{S}_{t}(i)}$ to the server. The server aggregates the received local subsets to construct a new dictionary $\{(\mathbf{x}_{s},y_{s})\}_{s\in\mathcal{S}_{t}}$, where $\mathcal{S}_{t}=\cup_{i=1}^{N}\mathcal{S}_{t}(i)$, and then sends it back to all $N$ clients (line 9). Finally, the $N$ clients use this updated dictionary to re-compute the embedded statistics of their local data, and then synchronize it with all other clients via the server (line 10-12).

Intuitively, in Algorithm 2, the clients first agree upon a common dictionary $\mathcal{S}_{t}$ that serves as a good representation of the whole dataset at the current time $t$, and then project their local data to the subspace spanned by this dictionary before communication, in order to avoid directly sending the raw data as in Algorithm 1. Then using the event-trigger, each client monitors the amount of new knowledge it has gained through interactions with the environment from last synchronization. When there is a sufficient amount of new knowledge, it will inform all the other clients to perform a synchronization. As we will show in the following section, the size of $\mathcal{S}_{t}$ scales linearly w.r.t. the maximum information gain $\gamma_{NT}$, and therefore it improves both the local computation efficiency on each client, and the communication efficiency during the global synchronization.