Exact Decomposition of Quantum Channels for Non-IID Quantum Federated Learning

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. We omit the labels for simplicity here.

In federated learning tasks, we are given a set of clients $\{C_i{\}}_{i=1}^n$, and we use $|C_i⟩\in H_{C_i}$ to denote the $i$th client. Each of the clients has access to its own dataset containing $N_{C_i}$ samples $D_i=\{|{\psi }_j^{C_i}⟩{\}}_{j=1}^{N_{C_i}}$ from the data Hilbert space $H_x$. The dataset can be statistically represented by the density matrix ${\rho }_x^{C_i}=N_{C_i}^{-1}{\sum }_j|{\psi }_j^{C_i}⟩⟨{\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.

If we assume that there is no entanglement among 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}|C_i⟩⟨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. In a classical dataset, this corresponds 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\left(\rho \right)={\sum }_i{p}_{C_i}{\rho }_x^{C_i}.$ Let $P_{C_i}=|C_i⟩⟨C_i|$ be the projector onto the subspace of client $C_i$. Then we can introduce the conditional density matrix [swart2020introduction] ${\rho }_{x|C_i}=P_{C_i}\rho P_{C_i}/\text{Tr}\left(\rho P_{C_i}\right)={\rho }_x^{C_i}\otimes |C_i⟩⟨C_i|$, which characterizes the data of a given client $C_i$. Tracing out the clients recovers the local dataset: ${\text{Tr}}_C\left({\rho }_{x|C_i}\right)={\rho }_x^{C_i}$. We summarize these concepts in Figure 1.

Here we briefly review the quantum version of FedAvg [mcmahan2017comm]. In a general supervised quantum machine learning problem, we aim to find a quantum channel $M$ that takes an input state $|\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 $M_w$ to minimize the average of some loss function $L$ on a given training dataset $D$:

where $|\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 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 [li2021quantum, xia2021quantumfed, chen2021federated, chehimi2022quantum, yun2022slimmable]. 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 [zhao2018federated, hsieh2020non], FedAvg faces the problem of non-IID quagmire, in which its performance deteriorate when the data of different clients are non-IID. Does this phenomenon also exist in the quantum regime?

We follow the steps of [zhao2018federated] and quantify the performance difference of FedAvg and the centralized case by the weight divergence $\Delta =\parallel w^{\left(f\right)}-w^{\left(c\right)}\parallel$, where $w^{\left(f\right)}$ and $w^{\left(c\right)}$ 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) [rubner2000earth] between the label distribution of client $C_i$ $p^{\left(i\right)}(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\parallel p^{\left(i\right)}(y=k)-p(y=k)\parallel$. Then, as a direct quantum extension to the Proposition 3.1 in [zhao2018federated], we have the following proposition:

Therefore, EMD is indeed a source of the weight divergence when $T>1$. This provides theoretical explanation for the existence of non-IID quagmire in quantum federated learning. The detailed proof is provided in Appendix A for completeness. 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: $L(w,|\psi ⟩,y)=L\left(M_w\right({\sigma }_x^{\psi }),y)$, where ${\sigma }_x^{\psi }=|\psi ⟩⟨\psi |$.

We begin by noting that when the data are decentralized, each client $C_i$ can still train a local channel $M_i$ on its own data $D_i$:

Apparently, the original problem, Equation (1), can be decomposed as

We can saturate the whole minimization problem if we saturate each individual sub problems, which is exactly what $M_i$ does. Therefore, if we can construct a global channel $M$, such that it coincides with all the local channels $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 $M$ is well defined. For example, suppose we have two identical samples $|{\psi }^{C_i}⟩=|{\psi }^{C_j}⟩=|\psi ⟩$ with the same label $y$, but are from different clients $C_i\ne C_j$. Then, in order to fulfill the local minimization problems, Equation (6), we must have $M_i(|\psi ⟩⟨\psi |)=M_j(|\psi ⟩⟨\psi |)$. This can be rephrased backwards, as demanding that the local channels $M_i$ coincide with the global channel $M$ on condition that the input data ${\sigma }_x^{\psi }$ comes from its own dataset ${\rho }_x^{C_i}$. That is,

where $P_x^{\psi }=|\psi ⟩⟨\psi |$ denotes the projector.

In order 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 $D_i$, which is trained to output the probability density of the input state within the local dataset:

In fact, the local channels $M_i$ and density estimators $D_i$ are enough to give an explicit construction of the global channel $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 [liu2022federated]. As a result, for any input state ${\sigma }_x$, if we randomly apply a local channel $M_i$ with probability $q_i$, the result will be statistically the same as the global channel $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 $M_i$ and a local density estimator $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 $M_i$ and density estimators $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 $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 $M_i$ and the density estimators $D_i$. For channels, we can use classical, quantum or hybrid algorithms at our will [bishop2006pattern, goodfellow2016deep, li2021recent]. For density estimators, we can use classical ones like Gaussian mixture models [bishop2006pattern] and normalizing flows [rezende2015variational], or quantum ones such as classical shadow tomography [huang2020predicting] and quantum state diagonalization algorithms [larose2019variational, xin2021experimental]. 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 [liu2022federated].

Classical Data & Quantum Channel: To perform quantum channels on classical data, one need 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 from another. Nevertheless, quantum density estimators offer a potential exponential speed-up. So there’s a trade-off between accuracy and efficiency.

Quantum Data & Quantum Channel: For quantum data, we need a quantum density estimator to estimate $\text{Tr}\left({\rho }_x^{C_i}{P}_x^{\psi }\right)$. This task can be solved by the classical shadow tomography, which is proved to saturate the information-theoretic bound [huang2020predicting].

Compared to qFedAvg, the proposed framework qFedInf shares some merits with its classical counterpart [liu2022federated]. 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 [zhou2020distilled, kasturi2020fusion]. Moreover, since no gradient information is transmitted, it’s automatically immune to the attacks based on gradient inversion [zhu2019deep, geiping2020inverting]. Finally, the density estimators can capture possible data heterogeneity, providing a new way to perform federated learning with non-IID data.

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 directly to machine learning tasks beyond classification.

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 [chehimi2022quantum, li2021quantum]. Therefore, to better illustrate the performance differences of qFedInf and qFedAvg, we devise a highly heterogeneous federated dataset based on 8 classes (“0” through “7”) from the MNIST hand written digits dataset. To produce heterogeneity, we adopt the star structure and cycle-m structure settings from [liu2022federated] 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 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:

where $\left\{x_j\right\}$ is the pixel values of the image and $\{|j⟩\}$ denotes the computational basis.

To perform inference, we measure the expectation values 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 [kingma2017adam] 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 [mitarai2018quantum, li2021quantum, li2017hybrid].

All the numerical simulations are conducted with JAX [jax2018github] and TensorCircuit [zhang2022tensorcircuit] 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\ge 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 a 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 need to perform a binary classification, and thus in practice, we find that circuits with only a few layers suffice to achieve good performance. For a fair 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 stays 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.

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 drops to only $\sim 75\%$. This suggests that, putting the federated settings aside, qFedInf can utilize the collective knowledge of many small models to achieve the capability of a large model. It’s a quantum analog of mixture of experts [masoudnia2014mixture] and ensemble learning [sagi2018ensemble].