Non-IID Quantum Federated Learning with One-shot Communication Complexity

We begin by setting up a typical decentralized quantum dataset. Classical datasets can be regarded as a special case of it, where the data samples are orthogonal to each other.

In federated learning tasks, we are given a set of clients $\{C_{i}\}_{i=1}^{n}$, and we can use a set of orthonormal basis elements $\{\ket{C_{i}}\in\mathcal{H}_{C}\}_{i=1}^{n}$ to denote them. Each of the clients has access to its own dataset containing $N_{C_{i}}$ samples $D_{i}=\{\ket{\psi_{j}^{C_{i}}}\}_{j=1}^{N_{C_{i}}}$ from the data Hilbert space $\mathcal{H}_{x}$. The dataset can be statistically represented by the density matrix $\rho_{x}^{C_{i}}=N_{C_{i}}^{-1}\sum_{j}\ket{\psi_{j}^{C_{i}}}\bra{\psi_{j}^{C_{i}}}$. Note that the density matrices of different clients may vary dramatically (i.e. non-IID). For example, in a multi-class classification problem, each client may have only seen samples from two or three classes.

Since the construction of entanglement between macroscopic objects from distant places is in general very difficult and is not expected to be realized in the near future [2], we focus on the situation where there is no entanglement among the clients. Then the joint density matrix of both the data and the clients is represented by $\rho=\sum_{i=1}^{n}\rho_{x}^{C_{i}}\otimes p_{C_{i}}\ket{C_{i}}\bra{C_{i}}$, where $p_{C_{i}}=N_{C_{i}}/N$ is the statistical weight of client $C_{i}$, and $N=\sum_{i}N_{C_{i}}$ is the total number of samples. For a classical dataset, the joint density matrix $\rho$ is diagonal, and its matrix elements correspond to the joint probability distribution $p(x,C_{i})$.

The averaged data density matrix is obtained by tracing out the clients: $\rho_{x}=\text{Tr}_{C}(\rho)=\sum_{i}p_{C_{i}}\rho_{x}^{C_{i}}.$ Let $P_{C_{i}}=\ket{C_{i}}\bra{C_{i}}$ be the projector onto the subspace of client $C_{i}$. Then we can introduce the conditional density matrix [30, 31] $\rho_{x|C_{i}}={P_{C_{i}}\rho P_{C_{i}}}/{\text{Tr}(\rho P_{C_{i}})}=\rho_{x}^{C_{i}}\otimes\ket{C_{i}}\bra{C_{i}}$, which characterizes the data of a given client $C_{i}$. Tracing out the clients recovers the local dataset: $\text{Tr}_{C}(\rho_{x|C_{i}})=\rho_{x}^{C_{i}}$. We summarize these concepts in Figure 1. For supervised learning tasks such as classification, there will also be a label $y$ associated with each data sample $\ket{\psi}$. We have omitted the labels here for simplicity.

Here we briefly review the quantum version of FedAvg [16]. In a general supervised quantum machine learning problem, we aim to find a quantum channel $\mathcal{M}$ that takes an input state $\ket{\psi}$ and transform it into the desired result (e.g. which class the image belongs to). We achieve this by tuning the variational parameters $w$ of the channel $\mathcal{M}_{w}$ to minimize the average of some loss function $\mathcal{L}$ on a given training dataset $D$:

where $\ket{\psi}$ is the input and $y$ is the corresponding label. We use gradient descent with learning rate $\eta$ to iteratively solve this optimization problem: at time step $t$,

For decentralized data, the total dataset is divided into local datasets from multiple clients: $D=\cup_{i}D_{i}$. Then we can decompose the update rule of Equation (2) into three steps: (1) local updates at time step $t$ for each client $C_{i}$,

(2) global average for every $T$ steps, e.g. at time step $mT,m\in\mathbb{N}$,

and (3) broadcast the averaged weights to all clients as the initial parameters for the next iteration. This is the basic protocol of qFedAvg, the common ground of all the existing quantum federated learning algorithms [17, 18, 19, 20, 21]. It recovers the centralized update rule when $T=1$. However, in practice, this cannot be achieved since we use mini-batch training strategies.

As pointed out in [22, 23], FedAvg faces the problem of non-IID quagmire, in which its performance deteriorates when the data of different clients are non-IID. Does this phenomenon also exist in the quantum regime?

We follow the steps of [22] and quantify the performance difference of FedAvg and the centralized case by the weight divergence $\Delta=\|w^{(f)}-w^{(c)}\|$, where $w^{(f)}$ and $w^{(c)}$ are the weights given by qFedAvg and the centralized case respectively. On the other hand, the level of non-IID can be quantified by the earth mover’s distance (EMD) [32] between the label distribution of client $C_{i}$, $p^{(i)}(y=k)=\sum_{y\in D_{i}}\delta_{y,k}/N_{C_{i}}$, and the centralized distribution $p(y=k)=\sum_{y\in D}\delta_{y,k}/N_{C}$: $\text{EMD}_{i}=\sum_{k}\|p^{(i)}(y=k)-p(y=k)\|$. Then, as a direct quantum extension to Proposition 3.1 in [22], we have the following proposition:

Therefore, EMD is indeed a source of the weight divergence when $T>1$. This provides a theoretical explanation for the existence of non-IID quagmire in quantum federated learning, similar to its classical counterpart [22]. We provide a detailed proof of Proposition 1 in Appendix A. Numerical experiments are conducted in Section III.3 to empirically illustrate this phenomenon in quantum federated learning.

Now that we have identified the non-IID quagmire in qFedAvg, we aim to find a different approach to tackle Equation (1) when the data are decentralized. We consider the case where the loss function is a function of the output of the quantum channel and the label: $\mathcal{L}(w,\ket{\psi},y)=\mathcal{L}(\mathcal{M}_{w}(\sigma_{x}^{\psi}),y)$, where $\sigma_{x}^{\psi}=\ket{\psi}\bra{\psi}$.

We begin by noting that when the data are decentralized, each client $C_{i}$ can still train a local channel $\mathcal{M}_{i}$ on its own data $D_{i}$:

To make use of the local channels, we can decompose the original problem, Equation (1), as

We can solve the whole minimization problem if we can solve each individual sub-problems simultaneously, which is exactly what $\mathcal{M}_{i}$ does. Therefore, if we can construct a global channel $\mathcal{M}$, such that it coincides with all the local channels $\mathcal{M}_{i}$ on their own data, i.e.

then the goal is achieved. ²²2The fact that identical samples have identical labels guarantees that the global channel $\mathcal{M}$ is well defined. For example, suppose we have two identical samples $\ket{\psi^{C_{i}}}=\ket{\psi^{C_{j}}}=\ket{\psi}$ with the same label $y$, but are from different clients $C_{i}\neq C_{j}$. Then, in order to fulfill the local minimization problems, Equation (6), we must have $\mathcal{M}_{i}(\ket{\psi}\bra{\psi})=\mathcal{M}_{j}(\ket{\psi}\bra{\psi})$. This can be rephrased backwards, as demanding that the local channels $\mathcal{M}_{i}$ coincide with the global channel $\mathcal{M}$ on condition that the input data $\sigma_{x}^{\psi}$ comes from its own dataset $D_{i}$, which is statistically represented by $\rho_{x}^{C_{i}}$. That is,

where $P_{x}^{\psi}=\ket{\psi}\bra{\psi}$ denotes the projector. ³³3To avoid confusion, we insist on using different notations for $P_{x}^{\psi}=\sigma_{x}^{\psi}=\ket{\psi}\bra{\psi}$ to emphasize their differences in physical meaning: we use $\sigma_{x}^{\psi}$ to denote the quantum state that you can load into your circuit, while we use $P_{x}^{\psi}$ to denote the projection operator.

To combine these local channels into a global one, we also need to capture the information of the local datasets $\rho_{x}^{C_{i}}$. To achieve this, we introduce the local density estimator $\mathcal{D}_{i}$, which is trained to output the probability density of the input state within the local dataset:

In fact, the local channels $\mathcal{M}_{i}$ and density estimators $\mathcal{D}_{i}$ are enough to give an explicit construction of the global channel $\mathcal{M}$. This is provided by the following theorem (proved in Appendix B):

We note that the classical special case of this theorem was first introduced in [36]. As a result, for any input state $\sigma_{x}$, if we randomly apply a local channel $\mathcal{M}_{i}$ with probability $q_{i}$, the result will be statistically the same as the global channel $\mathcal{M}$. This fact leads to the following framework for quantum federated learning.

Theorem 1 provides a framework for quantum federated learning. The specific protocol goes as follows. Firstly, each client $C_{i}$ trains a local channel $\mathcal{M}_{i}$ and a local density estimator $\mathcal{D}_{i}$ with its own data $\rho_{x}^{C_{i}}$. This step is completely distributed and concludes the whole training phase. Secondly, the trained channels $\mathcal{M}_{i}$ and density estimators $\mathcal{D}_{i}$ are sent to the server for inference according to Equation (11). That is, when a new input $\sigma_{x}^{\psi}$ comes, the server computes $q_{i}$ via the density estimators. Then it randomly loads the parameters from $\mathcal{M}_{i}$ with probability $q_{i}$ and gathers the outcomes. We call this framework quantum federated inference, or qFedInf for short. The detailed algorithms are summarized in Algorithm 1 and 2.

In practice, there is a wide range of choices for the channels $\mathcal{M}_{i}$ and the density estimators $\mathcal{D}_{i}$. For channels, we can use classical, quantum, or hybrid algorithms at our will [37, 1, 38]. For density estimators, we can use classical ones like Gaussian mixture models [37] and normalizing flows [39], quantum-inspired ones like [40], or quantum ones such as classical shadow tomography [41] and quantum state diagonalization algorithms [42, 43]. Each kind of density estimator comes with a different training strategy. We can also classify the possible scenarios based on the classical/quantum nature of the data and the channels:

Classical Data & Classical Channel: This is a purely classical problem already discussed in [36].

Classical Data & Quantum Channel: To apply quantum channels to classical data, one needs to choose an appropriate encoding scheme, e.g. amplitude encoding or gate encoding, to encode the data into quantum states. This gives rise to the problem of whether to use a classical density estimator on the original data or a quantum one on the encoded states. Note that before encoding, different samples are orthogonal to each other. However, the encoded quantum samples are in general overlapped, meaning that there’s a possibility of mistaking one sample for another. Nevertheless, quantum density estimators offer a potential exponential speed-up. So there’s a trade-off between accuracy and efficiency.

Quantum Data & Classical Channel: In general, trying to apply classical channels to process quantum data is not very efficient, as the tomography and representation of a quantum state cost exponentially large classical computational resources. Nevertheless, proposals [44] have been made to use the classical shadows [41] of a quantum state as the input to a classical machine learning algorithm. We defer to future works to investigate the performance of these proposals in a federated learning context.

Quantum Data & Quantum Channel: For quantum data, we need a quantum density estimator to estimate $\text{Tr}(\rho_{x}^{C_{i}}P_{x}^{\psi})$. This task can be solved by classical shadow tomography, which is proved to saturate the information-theoretic bound [41].

Compared to qFedAvg, the proposed framework qFedInf shares some merits with its classical counterpart [36]. It’s one-shot in the sense that only one communication between the server and each client is required. Meanwhile, there’s no need for a global public dataset or data synthesis/distillation, which are the common ground of existing one-shot algorithms [28, 29]. Moreover, since no gradient information is transmitted, it’s automatically immune to attacks based on gradient inversion [24, 25]. Finally, the density estimators can capture possible data heterogeneity, providing a new way to perform federated learning with non-IID data.

Though we have mainly focused on supervised learning throughout this paper, we also note that Theorem 1 holds for generic quantum channels, which is the most general form of quantum information processing. So the proposed framework qFedInf may be applied to machine learning tasks beyond classification. In Appendix C we provide an example of applying qFedInf to perform quantum generative learning [10, 11, 12].

Finally, we give a brief discussion on the complexity of the proposed framework. Suppose that the number of clients is $N_{\text{client}}$, the number of iterations needed for training is $N_{\text{iter}}$, and the number of parameters in each of the quantum channels $\mathcal{M}_{i}$ and density estimators $\mathcal{D}_{i}$ are $N_{\text{channel}}$ and $N_{\text{density}}$. Then the communication complexity [15] of the proposed qFedInf is $\mathcal{O}(N_{\text{client}}(N_{\text{channel}}+N_{\text{density}}))$. In comparison, the communication complexity of qFedAvg is $\mathcal{O}(N_{\text{client}}N_{\text{iter}}N_{\text{channel}})$, which is much less efficient when $N_{\text{iter}}$ is large. As for the circuit complexity [45, 46], it depends on the specific choice of the quantum circuit used as the channels. If the channels have a quantum speed-up, so will qFedInf.

In this section, we discuss the connection between the proposed framework qFedInf and mixture of experts (MoE) [47, 48, 49], which is an important strategy in ensemble learning. The idea of ensemble learning is to combine several models to make a joint prediction [50, 51]. Different models are expected to compensate for each other, leading to better performance. MoE, as a special kind of ensemble learning method, consists of two parts: a set of functions serving as the judge that decides the relative weight of different models, called the gating function; and an ensemble of specialized models, called the experts, each expected to perform well only on a subset of the total input space. In practice, MoE, along with other ensemble approaches such as bagging and boosting, has given birth to many state-of-the-art solutions to a wide range of machine learning problems [50, 51, 47].

In hindsight, we can see that qFedInf also consists of two parts: a set of density estimators $\mathcal{D}_{i}$ that decides the probabilistic weight of different local channels $\mathcal{M}_{i}$, as in Equation (11); and a set of local channels $\mathcal{M}_{i}$, each performing well only on the client’s own data. In the language of MoE, the density estimators are the gating functions, and the local channels correspond to the experts. Therefore, qFedInf is exactly an MoE working in the one-shot federated learning context. In Section III, numerical experiments show that qFedInf can indeed combine the knowledge of many weak classifiers (shallow circuits only capable of binary classification) to achieve the capability of large models.

We observe that the existing quantum federated learning algorithms can already achieve high accuracies on binary classification tasks with synthetic and common classical/quantum datasets [20, 17]. Therefore, to better illustrate the performance differences between qFedInf and qFedAvg, we devise a highly heterogeneous federated dataset based on 8 classes (“0” through “7”) from the MNIST handwritten digits dataset. To produce heterogeneity, we adopt the star structure and cycle-m structure settings from [36] as follows.

Star structure: Each client only has access to the data of two classes, with one of the classes fixed. That is, client $C_{i},i=1,\cdots,7$ only has access to the data of digit “0” and “$i$”.

Cycle-m structure: Each client only has access to the data of $m$ classes in a cyclic way. Client $C_{i},i=1,\cdots,7$ has access to the data of digit “$i$”, …, “$i+m-1$ module 8”.

Datasets of the same structures are also prepared for the Fashion-MNIST dataset, which is composed of images of daily objects and is regarded as a harder version of MNIST.

We parameterize the quantum classifier as a quantum circuit of $l$ layers. Each layer contains a set of controlled-NOT gates on adjacent qubits, followed by a parameterized rotation on each qubit. The detailed circuit is shown in Figure 2.

To load the data into an 8-qubit quantum circuit, we interpolate and resize each image into a size of 16 $\times$ 16 and use amplitude encoding to transform it into a quantum state [3, 52]:

where $\{x_{j}\}$ are the pixel values of the image and $\{\ket{j}\}$ denote the computational basis. In experiments, amplitude encoding can be implemented using quantum random access memory [53, 54] or universal gate decomposition [46, 55, 56].

To perform inference, we measure the expectation values $\braket{Z_{k}}$ on each qubit, amplify the outcomes by a factor of 10, and fed them into the softmax function to predict the probability of each class. We use the Adam optimizer [57] of learning rate $10^{-2}$ and a batch size of $128$ to minimize the standard cross entropy loss function. The gradients used in the optimization can be computed via the parameter shift rule [58, 17, 59].

All the numerical simulations are conducted with JAX [60] and TensorCircuit [61] on one NVIDIA Tesla V100 GPU. The source code is available at https://github.com/JasonZHM/quantum-fed-infer.

We begin by demonstrating the non-IID quagmire of qFedAvg discussed in Section II.3 with numerical simulations. Specifically, we train and test qFedAvg on datasets of the cycle-m structure. The parameter $m$ serves as a good controller over the level of non-IID: as $m$ increases, each client will have access to more classes, and the level of non-IID will decrease.

We train the channel with $l=48$ and $m=2,3,4,5,6$ using qFedAvg for 5 epochs. The global synchronization frequency of qFedAvg is set to be one time per batch step. The resulting test accuracies on a test set of size 1024 are plotted in Figure 3. In line with the theoretical analysis, we find that the top-1 accuracy increases as the level of non-IID drops. When the data are highly heterogeneous ($m=2$), qFedAvg suffers from a loss of $\sim 4\%$ ($10\%$) in accuracy on MNIST (Fashion-MNIST) compared to the benchmark trained on the centralized data with the same circuit structure. Nevertheless, when the data heterogeneity is mild ($m\geq 5$), qFedAvg can achieve comparable performance compared with the centralized classifier. We also plot the test loss and accuracy curves on star structure in Figure 4. As expected, qFedAvg converges much slower to significantly lower accuracy.

To test the performance of qFedInf on non-IID data, we train and test qFedInf on the most heterogeneous settings, namely the star structure and cycle-2 structure, and compare it with qFedAvg and the centralized benchmark. In such settings, each client’s local classifier only needs to perform a binary classification, and thus in practice, we find that circuits with only a few layers suffice to achieve good performance. For comparison with the $l=48$ qFedAvg and the centralized benchmark, we choose $l=6$ for the local classifiers, so that the total number of variational parameters of the $8$ clients remains the same. We train the local classifiers for 5 epochs and plot the loss and accuracy curves in Figure 5. We adopt Gaussian mixture models with 5 modes as the local density estimators. The combined global model achieves a top-1 accuracy similar to the centralized benchmark and significantly higher than that of qFedAvg on both settings of both datasets. The detailed accuracies are listed in Table 1. Meanwhile, the performance of qFedInf is roughly unaffected by the number of classes per client, demonstrating its robustness against the level of non-IID.

We note that in the training process, qFedAvg requires a total number of 500 communication rounds, while qFedInf only needs one. Moreover, the local classifiers used in qFedInf are much shallower compared to those of qFedAvg and the centralized benchmark. If we change the $l$ of the centralized classifier to 6, its test accuracy will drop to only $\sim 75\%$. To rule out the possibility that barren plateau [62] causes this performance difference, we also test qFedInf with $l=48$, and the resulting performance is the same as $l=6$, which suggests that the barren plateau issue is not significant in our settings. Therefore, putting the federated settings aside, qFedInf can indeed utilize the collective knowledge of many small models to achieve the capability of a large model, in line with MoE and ensemble learning as discussed in Section II.6.