Uplink Scheduling in Federated Learning: an Importance-Aware Approach via Graph Representation Learning

Over the last years, different metrics for FL client scheduling have been proposed and discussed. In [14], by leveraging the concept of Age of Information (AoI), a metric termed age of update (AoU) is introduced, which takes into account the staleness of the received parameters. In [10, 11, 14, 15] channel conditions experienced by different clients are considered during the scheduling decision. In [15], a channel prediction algorithm based on Gaussian process regression is incorporated into the scheduling process when dealing with imperfect channel state information. Authors from [11], instead, exploit both diversity in multiuser channels and diversity in the importance of the edge devices’ local learning updates. Importance, in this case, is measured by the local parameter update’s gradient divergence, which must be reported to the PS downstream of the training of all client devices. A similar training-feedback metric of importance is introduced in [10], where the significance of the model updates at the devices is captured by the L2-norm of the model update.
Under the frequentist setting, training data constitute a fundamental part of the inductive bias of a model. In this context, in line with the idea of importance-aware communications, knowledge about data distribution among devices would suffice for driving proper FL scheduling decisions. However, as previously stated, this approach is unfeasible in FL systems. Nevertheless, clients exhibit relations and correlations in a network setting, especially in the context of massive Internet-of-Things (IoT) and ultra-dense 6G networks. Here, we argue that the scheduling of FL client devices is a problem with a rich relational structure and, as a consequence, there is a need to tackle this problem effectively by taking node correlations into account. Relations among clients, which relate to local data distribution too, can be learned and inferred by encompassing network geometry and relational representation learning, while at the same time preserving users’ privacy. Graphs, generally, are a representation that supports arbitrary relational structures, and computations over graphs afford a strong relational inductive bias [16]. By considering networks of clients as graphs, we introduce this bias in the clients’ scheduling process by leveraging Unsupervised Graph Representation Learning (UGRL). As results show, this effectively makes up for the impossibility of selecting users based on their data, while at the same time aiming for an energy and computationally-efficient scheduling protocol.

Let us consider a system comprised of one Access Point (AP) co-located with a PS and multiple client devices with local data and computation capabilities, as depicted in Fig. 2. Physically running on the PS, there are two software entities responsible for model aggregation and radio resource management: the Federated Aggregator and the Graph Scheduler, respectively described in the following subsections.

Each client $m\in\mathcal{M}$ holds a local data set $D_{m}=\left\{\mathbf{x}_{m}\in\mathcal{R}^{d},\mathbf{y}_{m}\in\mathcal{R}\right\}$ with cardinality $|D_{m}|$, such that $\sum_{m\in\mathcal{M}}|D_{m}|=|D|$, and is equipped with a single isotropic antenna. On the PS side, we consider an antenna with a directive gain of 15 dBi. All clients are randomly uniformly distributed within a radius $R$ of the PS. The PS broadcasts the global model to the selected clients with a transmit power of 15 dBm, while the latter send their local updates with a transmit power of 10 dBm.
As further detailed in the next subsections, we consider the two cases of non-i.i.d. spatially correlated data and spatially correlated/uncorrelated non-i.i.d clusters of i.i.d. data. Both cases are artificially reproduced in our experiments by spatially distributing MNIST digit labels among neighbor clients.
Within the network area, we consider an Orthogonal Frequency Division Multiple Access (OFDMA) scheme with perfect equalization, where $M=|\mathcal{M}|$ clients share the same spectrum and can be assigned one of the $K<M$ set of orthogonal sub-channels for model parameters transmission. Moreover, we assume a slow fading propagation model, where each model transmission from device $m\in\mathcal{M}$ to the PS is shorter than the channel coherence time $T_{c}$:

where $T_{D,m}$ is the model transmission time for client $m$. With the aforementioned assumptions, the channel impulse response $h_{m,PS}(f,t)$ from device $m$ to the PS loses its time and frequency dependency within a block transmission duration $T_{D,m}$ (2):

where $\mathbf{H}_{M+1}$ is the channel matrix of dimension $M+1$.

The goal of the Federated Aggregator is to learn a ML model by offloading and aggregating the training to the set $\mathcal{M}$ of distributed clients with local data. The federated training process involves a number of iterations, namely communication rounds, until convergence. Each client $m\in\mathcal{M}$, upon receiving a global model $\bm{\theta}_{PS}$ from the PS at the beginning of a new round, executes multiple Stochastic Gradient Descent (SGD) updates to minimize the model’s loss function with respect to its local dataset (3):

where $F_{m}$ is the $m$-th client loss function and $f(\bm{\theta},\{x_{i},y_{i}\})$ indicates the task-dependent loss (e.g., Mean Square Error (MSE), categorical cross-entropy, etc.) for every training example $\{x_{i},y_{i}\}$. At every $j$-th local SGD iteration, each client $m$ updates its local model according to (4):

where $\alpha(t)$ denotes the learning rate scheduled by the PS for communication round $t$ and $\mathbf{g}_{m}:=\nabla(F_{m}(\bm{\theta}_{m},D_{m})$ is the $m$-th client’s gradient of the local model’s weights. Once the training is terminated, each client selected for scheduling must forward its local model $\bm{\theta}_{m}(t)$ to the PS, which will update the global model upon aggregation of all received clients’ models. Here, we refer to FedAvg algorithm [1], for which the aggregation is a weighted sum described by (5):

where we denote by $\mathcal{K}$ ($|\mathcal{K}|=K$) the set of scheduled clients for communication. For model evaluation, we consider a separate centralized test set $D_{test,PS}$ locally residing on the PS.

The Graph Scheduler is the entity responsible for the procedures depicted in Fig. 3: dynamic graph creation, UGRL, and distance maximization.

• •

  Sensing phase and reporting: a sensing phase is performed by all client devices, and it should be repeated with a periodicity that depends on environment dynamicity. Following, all nodes $v$ report a list of sensed neighbors $N_{t}(v)$ (Fig. 2) to the graph scheduler.

• •

  Dynamic graph formation: the scheduler builds a dynamic graph from each node’s list of sensed neighbors. In our approach, we retain the strongest-K neighbors, as it allows to adapt to the density of nodes’ deployment. Note that alternative approaches, such as retaining adjacencies based on a received power threshold, might also be used.

• •

  UGRL: graph representation learning involves the transformation via an encoding function $f:(\mathcal{G}_{|\mathcal{M}|})\rightarrow\mathbb{R}^{Tx|\mathcal{M}|}$ from the graph-structured node representation $\mathcal{G}_{|\mathcal{M}|}$, to a vectorial space of T-dimensional node embeddings. In this procedure, node embeddings from $\mathcal{G}_{|\mathcal{M}|}$ are efficiently computed in an unsupervised way with the use of random walks procedures. In our scenario, we make use of Node2Vec algorithm [17]. Further details are discussed in section IV.

• •

  Distance maximization: the final step involves the computation of the scheduling sequence $\mathbf{U}_{t+1}$. This is based on the distance maximization between node embeddings retained in a context window of tunable dimension. Additional details are provided in section IV.

The intuition behind the proposed approach is that the clients’ data distribution reflects the structural relation of nodes in a graph. Consequently, this method does not apply to the trivial case of i.i.d. data. Vice versa, it is possible to think of a plethora of applications and AI-driven network procedures foreseen for 6G in which this condition holds: localization, tracking, integrated sensing and communication, channel estimation and measures of a physical quantity from a sensors network is just a non-exhaustive list of examples where clients’ data are non-i.i.d. and spatially correlated. Another vertical of great importance for future 6G networks is Industrial Internet-of-Things (IIoT), where typically nodes are arranged into clusters, and nodes within a cluster might hold similar kinds of measurements (e.g., monitoring sensors inside automatic machines in a warehouse). To this end, we consider in this work typical kinds of client data distributions that find use in many real-world 6G applications:

• •

  Non-i.i.d. and spatially correlated data distribution (Fig. 4(a)).

• •

  Spatially correlated/uncorrelated clusters of i.i.d. data (Fig. 4(b), 4(c)).

For benchmarking purposes, we reproduced the three scenarios described above in our simulations by distributing MNIST data arranged by labels to a set of randomly distributed clients. According to the considered scenario, clients are distributed, at the beginning of every new simulation, a random number of examples according to: a) their respective position to the PS (for the case depicted in Fig. 4(a)), or b) their belonging cluster (for the case depicted in Fig. 4(b), 4(c)).

Once $f(\mathcal{G})$ is obtained, clients’ scheduling relies on the distance maximization of the nodes in the embedding space. This has the effect of increasing data heterogeneity of the client devices during consecutive communication rounds. The distance among a pair of nodes $(v,u)$ in the graph $\mathcal{G}$ is evaluated as the normalized dot product (cosine similarity) $s(v,u)$ of their node embeddings $\mathbf{z}_{v},\mathbf{z}_{u}$ (8):

To introduce memory of the past actions in the scheduling process, nodes scheduled in previous communication rounds are stored in a context window $W_{L}$ of tunable length $L=\max(|W_{L}|)$. Accordingly, the similarity scores of the nodes $v\in W_{L}$ are computed and collected in a matrix $\mathbf{S}_{v,u}$ of dimension $|W_{L}|\times(M-|W_{L}|)$:

A scheduling decision is finally determined by equation (10).

The latter is equivalent to summing all elements of the matrix $\mathbf{S}_{v,u}$ by column, and selecting the next scheduled client as the argument corresponding to the column holding the minimum sum value, i.e., the node with maximum distance with respect to all previously scheduled nodes in $W_{L}$.