One-shot Collaborative Data Distillation

Let $\mathcal{D}\overset{\Delta}{=}\{(x_{i},y_{i})\}_{i=1}^{|\mathcal{D}|}$ be the data set that needs to be distilled, where $x_{i}\in\mathcal{X}$ denotes the input features and $y_{i}\in\mathcal{Y}$ is the label for $x_{i}$. Throughout, the notation $d\sim\mathcal{D}$ refers to a data point $d$ selected uniformly at random from the set $\mathcal{D}$. Given a data budget $n\in\mathrm{Z}^{+}$, a data distillation technique aims to synthesize a high-fidelity summary $\mathcal{S}\overset{\Delta}{=}\{(\tilde{x_{i}},\tilde{y_{i}})\}_{i=1}^{n}$ such that $n\ll|\mathcal{D}|$. The small distilled dataset should achieve a comparable generalization performance to the large original dataset.

For a given learning algorithm $\Phi_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$, with parameterization $\theta$, the empirical risk $\mathcal{R}$ on parameterization $\theta$ and input data $\mathcal{D}$ is defined as

where $l$ is a loss function. A training algorithm for $\Phi$ attempts to find $\theta$ that minimizes $\mathcal{R}$.

Thus, the objective is to extract the knowledge from $\mathcal{D}$ and transfer it to the synthetic set $\mathcal{S}$, such that the model trained on $\mathcal{S}$ should achieve comparable generalization to the model trained on $\mathcal{D}$.

With compact synthetic sets, collaborative data distillation aims to reduce the communication overhead in distributed machine-learning applications at a minimal cost in terms of fidelity.

Meta-learning based methods [1, 33, 37, 46] treat $\mathcal{S}$ as a hyperparameter, which is updated by a meta (outer) algorithm and a base (inner) algorithm solves a conventional learning problem with respect to the synthetic dataset. The objective can thus be formulated as a bi-level optimization:

subject to

The inner loop, which optimizes parameters on the synthetic data, can be realized through gradient descent or kernel regression. The objective function can be defined as the meta-loss $\mathcal{L}(\mathcal{S})=\mathcal{R}(\mathcal{D},\theta^{\mathcal{S}})$. Consequently, the synthetic data can be updated as $\mathcal{S}=\mathcal{S}-\alpha\nabla_{\mathcal{S}}\mathcal{L}(\mathcal{S})$ for learning rate $\alpha$.

Data matching [2, 44] aims to align the byproducts of model training on real and synthetic data. The synthetic data learns to mimic the influence of real data on model training. The objective function of data matching can be summarized as follows:

subject to

where $Q$ is a distance function, $\eta$ is the learning rate, and $\phi$ maps a dataset to informative spaces such as gradient, parameter, and feature spaces. For example, the map $\phi(\mathcal{S},\theta)=\nabla_{\theta}\mathcal{R}(\mathcal{S};\theta)$ [44] equates to matching the gradients, with respect to $\theta$, of the observed empirical risk, and $\mathcal{S}$ is optimized to mimic the gradients observed on $\mathcal{D}$ during model training. In practice, the full dataset might be replaced with a batch to save memory and facilitate faster convergence.

The underlying bi-level optimization of prior approaches is often expensive in terms of computation and memory. To mitigate these costs, distribution matching [36, 43] aims to solve a correlated proxy task that restricts optimization to a single level and improves scalability. Instead of matching the quality of the models generated by $\mathcal{D}$ and $\mathcal{S}$, distribution matching attempts to match the underlying distributions of $\mathcal{D}$ and $\mathcal{S}$. The assumption here is that datasets with the same distribution will lead to similarly trained models.

Distribution matching uses a collection of random parametric encoders to embed data into low-dimensional latent spaces. Distance metrics can then be used to compute the distribution mismatch between $\mathcal{D}$ and $\mathcal{S}$. Formally, given a set of encoders $\mathcal{E}=\{\psi_{i}:\mathcal{X}\rightarrow\mathcal{X}_{i}\}$, the optimization objective, under maximum mean discrepancy, is:

where $\mathcal{D}^{y}=\{x\mid(x,y)\in\mathcal{D}\}$. This objective, for a given $\psi\in\mathcal{E}$, can be approximated with the following empirical loss:

where $D^{y}\subset\mathcal{D}^{y}$ denotes a batch of real data and $S^{y}\subset\mathcal{S}^{y}$ denotes a batch of synthetic data. Typically $\mathcal{E}$ is generated randomly and each $\psi\in\mathcal{E}$ has the same network architecture.

The collaborative distillation process begins with the server initializing a set of synthetic data $\mathcal{S}$. This can be achieved with a random initialization or by asking clients to transmit local data distillations: $\mathcal{S}=\cup_{i=1}^{K}\mathcal{S}_{i}$. With the synthetic data initialized, it is then updated iteratively. For the meta-learning and data-matching frameworks, each iteration has the following steps:

1. 1.

   The server retrieves a network $\theta$. The network parameters could be generated randomly or retrieved from a cache.

2. 2.

   The network $\theta$ is updated. The update could be based on either $\mathcal{R}(\mathcal{D};\theta)$ or $\mathcal{R}(\mathcal{S};\theta)$. For the former, clients conduct federated learning. The latter objective can be executed on the server.

3. 3.

   $\theta$ is broadcast to clients.

4. 4.

   The loss function $\mathcal{L}(\mathcal{S})$ is computed. The clients collaborate to compute $\Phi_{\theta}(D)$, for batch $D\subset\mathcal{D}$.

5. 5.

   Synthetic data is updated based on the gradient of the loss function.

At each iteration, model parameters are broadcast to clients, and clients send losses incurred on real data to the server. For large models, this can create a communication overhead that compromises the benefits of collaborative distillation. In other words, similar to federated learning, the framework involves multiple rounds of communication that broadcast model parameters to clients.

This limitation can be overcome with distribution matching. As distribution matching does not require model training, Step 2 from the framework can be removed. Further, as embeddings are initialized randomly, they can be broadcast with a random seed. Thus, distribution matching can be implemented in a collaborative learning environment without explicitly sharing network parameters.

The goal of Collaborative Distribution Matching (CollabDM) is to compute the loss function in Equation (3) for each embedding $\psi\in\mathcal{E}$. The gradient of the loss is used to update the synthetic dataset stored on the server. As the loss function is calculated over the global dataset $\mathcal{D}=\cup_{i=1}^{K}\mathcal{D}_{i}$, the updates are able to capture the global dynamics of the data. Equation  (3) can be split into two components: the embeddings on real data, which are computed at clients, and the embeddings on synthetic data, which are computed at the server. As $\mathcal{E}$ is fixed, it can be broadcast to clients prior to the distillation process. This allows each client to compute the mean embeddings on real data, one for each iteration of server training, in a single batch and complete their share of the collaboration in a single round of communication. A high-level overview of the procedure is presented in Figure 1. We now outline the full algorithm (provided in Algorithm 1) in more detail.

Let $|\mathcal{E}|=T$ denote the number of rounds required to distill synthetic data through distribution matching. In CollabDM, for each future training round $t\in\{1,\ldots,T\}$, the server begins by selecting a random seed $\alpha_{t}$ to encode a lower-dimension embedding and selecting a subset of clients $Z_{t}\subset\{1,\ldots,K\}$ to participate in the round. A batch of seeds $A_{k}=\{\alpha_{t}\mid k\in Z_{t}\}$ is then broadcast to each client $k$. Once the clients receive embedding seeds from the server, client training begins. Each client has two roles. First, the client performs a local data distillation to produce $\mathcal{S}_{k}$. Any data distillation technique could be used. Local distillations will be used to initialize the synthetic data at the server. Second, the client computes their contribution to each objective function. For each embedding $\alpha_{t}\in A_{k}$ and label $y\in\mathcal{Y}$, the client selects a batch of real data $D_{t,k}^{y}\subset\mathcal{D}_{k}^{y}$, of size $B=|D_{t,k}^{y}|$, and computes the mean of the embeddings on the batch:

The collection of sums

is then sent to the server, along with $\mathcal{S}_{k}$. This concludes the client’s role in CollabDM. Thus, in a single round of communication, the client receives $A_{k}$ and, subsequently, transmits $(L_{k},\mathcal{S}_{k})$.

The server can now complete data distillation through distribution matching. The synthetic data is initialized through the local distillations $\mathcal{S}=\cup_{k=1}^{K}\mathcal{S}_{k}$. The server then iterates through the embeddings in $\mathcal{E}$. For each embedding $\alpha_{t}$, using the client computations on real data $\cup_{k\in Z_{t}}L_{t,k}$, the loss function $\mathcal{L}$ of Equation (1) is evaluated. The synthetic data is then updated with the gradient of the loss with respect to $\mathcal{S}$:

As the embeddings on real data are constant with respect to $\mathcal{S}$, the gradients can be computed at the server. Therefore, a global data distillation is achieved without further communication with clients.

There are a number of optimizations that can be applied to the distribution matching objective to improve the utility of the synthetic data [19, 29, 36, 45]. These optimizations can be adapted to CollabDM. Notably, the synthetic data variables can be parameterized in a more efficient manner to allow for the distillation of more instances (for the same memory budget) and an enhanced representation of the real data. To achieve this, we adopt a technique called partition and expand [43]. For partition parameter $l$, each image $s\in\mathcal{S}$ is partitioned into $l\times l$ mini-samples, and each mini-sample is then expanded to the input data dimensions using differentiable augmentation:

Thus, the number of features extracted from $\mathcal{S}$ is increased without changing the storage budget.

Attack detection on 5G network data is a motivating application for our technique. 5G networks are decentralized by design and cater to a range of verticals. It is a setting in which data generation is innately distributed and heterogeneous. In our test setting, data collection is split across two base stations, which act as the clients in CollabDM. For all experiments, the amount of data transmitted is limited to $9$MB per client.

We first look at the impact of the number of images per class on attack classification. The results are presented in Figure 3. Remarkably, at just 1 image-per-class, CollabDM-pae achieves $89\%$ testing accuracy. This represents the distillation of 12,995 images of network traffic into 5 informative images, which, with the partition-and-expand technique, contain enough information to allow the network to not only distinguish between benign and malicious flows but also classify concrete attacks. This suggests that different attacks have highly distinct profiles. In addition, at just 10 images-per-class, CollabDM-pae achieves peak testing accuracy at $99\%$.

To verify the generalizability of the global synthetic sets, we conduct cross-architecture experiments. The results are presented in Table 3. Synthetic data is learned on one architecture and evaluated on a separate architecture. Each synthetic set contains 10 images per class. Results indicate that the synthetic sets generalize well, with, at worst, only a small drop in accuracy when moving to new architectures. These results promote the use of data distillation for data sharing in 5G networks, with very small global synthetic sets available for machine learning applications at different locations in the network.