Decentralized Federated Learning with Model Caching on Mobile Agents

Federated learning (FL) is a type of distributed machine learning (ML) that prioritizes data privacy [mcmahan2017communication]. The traditional FL involves a central server that connects with a large number of agents. The agents retain their data and do not share them with the server. During each communication round, the server sends the current global model to the agents, and a small subset of agents are chosen to update the global model by running stochastic gradient descent (SGD) [robbins1951stochastic] for multiple iterations on their local data. The central server then aggregates the updated parameters to obtain the new global model. FL is a natural fit for the emerging Internet-of-Things (IoT) systems, where each IoT device not only can sense its surrounding environment to collect local data, but also is equipped with computation resources for local model training, as well as communication interfaces to interact with a central server for model aggregation. Many IoT devices are mobile, ranging from mobile phones/tablets, autonomous cars/drones, to self-navigating robots, etc. In the recent research efforts on smart connected vehicles, there has been a focus on integrating vehicle-to-everything (V2X) networks with Machine Learning (ML) tools and distributed decision making [barbieri2022decentralized], particularly in the area of computer vision tasks such as traffic light and signal recognition, road condition sensing, intelligent obstacle avoidance, and intelligent road routing, etc. With FL, vehicles locally train deep ML models and upload only the model parameters (i.e., neural network weights and biases) to the central server. This approach not only reduces bandwidth consumption, as the size of model parameters is much smaller than the size of raw image/video data, but also leverages computing power on vehicles, and protects user privacy.

However, FL on mobile agents still faces communication and computation challenges. The movements of mobile agents, especially at high speed, lead to fast-changing channel conditions on the wireless connections between mobile agents and the central FL server, resulting in long latency in FL [niknam2020federated]. Battery-powered mobile agents also have limited power budget for long-range wireless communications. Meanwhile, the non-i.i.d data distributions on mobile agents make it difficult for local models to converge. As a result, FL on mobile agents to obtain an optimal global model remains an open challenge. Decentralized federated learning (DFL) has emerged as a potential solution where local model aggregations are conducted between neighboring mobile agents using local device-to-device (D2D) communications with high bandwidth, low latency and low power consumption [10251949] . Preliminary studies have demonstrated that decentralized FL algorithms have the potential to significantly reduce the high communication costs associated with centralized FL. However, blindly applying model aggregation algorithms, such as FedSDG [shokri2015privacy] and FedAvg [mcmahan2017communication], developed for centralized FL to distributed FL cannot achieve fast convergence and high model accuracy [liu2022distributed].

Motivated by DTN, we propose delay-tolerant DFL communication and aggregation enabled by model caching on mobile agents. To realize DTN-like model spreading, each mobile agent stores not only its own local model, but also local models received from other agents in the recent history. Whenever it meets another agent, it transfers its own model as well as the cached models to the agent through high-speed D2D communication. Local model aggregation on an agent works on all its cached models, mimicking a local parameter server. Compared with DFL without caching, DTN-like model spreading can push local models faster and more evenly to the whole network; aggregating all cached models can facilitate more balanced learning than pairwise model aggregation. While DFL model caching sounds promising, it also faces a new challenge of model staleness: a cached model from an agent is not the current model on that agent, with the staleness determined by the mobility patterns, as well as the model spreading and caching algorithms. Using stale models in model aggregation may slow down or even deviate model convergence.

The key challenge we want to address in this paper is how to design cached model spreading and aggregation algorithms to achieve fast convergence and high accuracy in DFL on mobile agents. Towards this goal, we make the following contributions:

1. 1.

   We develop a new decentralized FL framework that utilizes model caching on mobile agents to realize delay-tolerant model communication and aggregation;

2. 2.

   We theoretically analyze the convergence of aggregation with cached models, explicitly taking into account the model staleness;

3. 3.

   We design and compare different model caching algorithms for different DFL and mobility scenarios.

4. 4.

   We conduct a detailed case study on vehicular network to systematically investigate the interplay between agent mobility, cache staleness, and convergence of model aggregation. Our experimental results demonstrate that our cached DFL converges quickly and significantly outperforms DFL without caching.

Similar to the standard federated learning problem, the overall objective of mobile DFL is to learn a single global statistical model from data stored on tens to potentially millions of mobile agents. The overall goal is to find the optimal model weights $x^{*}\in\mathbb{R}^{d}$ to minimize the global loss function:

where $N$ denotes the total number of mobile agents, and each agent has its own local dataset, i.e., $\mathcal{D}^{i}\neq\mathcal{D}^{j},\forall i\neq j$. And $z^{i}$ is sampled from the local data $\mathcal{D}^{i}$.

All agents participate in DFL training over $T$ global epochs. At the beginning of the $t^{th}$ epoch, agent $i$’s local model is $x_{i}(t)$. After $K$ steps of SGD to solve the following optimization problem with a regularized loss function:

agent $i$ obtains an updated local model $\tilde{x}_{i}(t)$. Meanwhile, during the $t^{th}$ epoch, driven by their mobility patterns, each agent meets and exchanges models with other agents. Other than its own model, agent $i$ also stores models it received from other agents encountered in the recent history in its local cache $\mathcal{C}_{i}(t)$. When two agents meet, they not only exchange their own local models, but also share their cached models with each others to maximize the efficiency of DTN-like model spreading. The models received by agent $i$ will be used to update its model cache $\mathcal{C}_{i}(t)$, following different cache update algorithms, such as LRU update method (Algorithm 2) or Group-based LRU update method, which will be described in details later. As the cache size of each agent is limited, it is important to design an efficient cache update rule in order to maximize the caching benefit.

After cache updating, each agent conducts local model aggregation using all the cached models with customized aggregation weights $\{\alpha_{j}\in(0,1)\}$ to get the updated local model $x_{i}(t+1)$ for epoch $t+1$. In our simulation, we take the aggregation weight as $\alpha_{j}=(n_{j}/\sum_{j\in\mathcal{C}_{i}(t)}n_{j})$, where $n_{j}$ is the number of samples on agent $j$.

The whole process repeats until the end of $T$ global epochs. The detailed algorithm is shown in Algorithm 1. $z^{i}_{k}$ are randomly drawn local data samples on agent $i$ for the $k$-th local update, and $\eta$ is the learning rate.

We now highlight the key ideas and challenges behind our convergence proof.

Step 1: Similar to Theorem 1 in DBLP:journals/corr/abs-1903-03934, we bound the expected cost reduction after $K$ steps of local updates on the $j$-th agent, $\forall j\in[N]$, as

Step 2: For any epoch $t$, we find the index $M(t)$ of the agent whose model is the “worst", i.e., $M(t)=\arg\max_{j\in[N]}\{F(x_{j}(t))\}$, and the “worst" model on all agents over the time period of $[t-\tau_{max}+1,t]$ as

Step 3: We bound the cost reduction of the “worst" model at epoch $t+1$ from the “worst" model in the time period of $[t-\tau_{max}+1,t]$, i.e., the worst possible model that can be stored in some agent’s cache at time $t$, as:

Step 4: We iteratively construct a time sequence $\{T^{\prime}_{0},T^{\prime}_{1},T^{\prime}_{2},...,T^{\prime}_{N_{T}}\}\subseteq\{0,1,...,T-1\}$ in the backward fashion so that

Step 5: Applying inequality (3.1) at all time instances $\{T^{\prime}_{0},T^{\prime}_{1},T^{\prime}_{2},...,T^{\prime}_{N_{T}}\}$, after T global epochs, we have,

Step 6: With the results in (3.1) and by leveraging Theorem 1 in yang2021achieving, Algorithm 1 converges to a critical point after $T$ global epochs.

We conduct experiments using three standard FL benchmark datasets: MNIST [deng2012mnist], FashionMNIST [xiao2017fashionmnistnovelimagedataset], CIFAR-10 [krizhevsky2009learning] on $N=100$ vehicles, with CNN, CNN and ResNet-18 [he2016deep] as models respectively. We design three different data distribution settings: non-i.i.d, i.i.d, Dirichlet distribution method. In extreme non-i.i.d, we use a setting similar to su2022boost, data points in training set are sorted by labels and then evenly divided into 200 shards, with each shards containing 1–2 labels out of 10 labels. Then 200 shards are randomly assigned to 100 vehicles unevenly: 10% vehicles receive 4 shards, 20% vehicles receive 3 shards, 30% vehicles receive 2 shards and the rest 40% receive 1 shards. For i.i.d, we randomly allocate all the training data points to 100 vehicles. For Dirichlet distribution, we follow the setting in xiong2024deprl, to take a heterogeneous allocation by sampling $p_{i}\sim Dir_{N}(\pi)$, where $\pi$ is the parameter of Dirichlet distribution. We take $\pi=0.5$ in our following experiments.

We use SGD as the optimizer and set the initial learning rate $\eta=0.1$ for all experiments, and use learning rate scheduler named ReduceLROnPlateau from PyTorch, to automatically adjust the learning rate for each training.

Algorithm 2, which we name as LRU method for convenience, is the basic cache update method we proposed. Basically, the LRU updating rule aims to fetch the most-up-to-date models and keep as many of them in the cache as possible. At line 13, 18, the metric for making choice among models is the timestamp of models, which is defined as the epoch time when the model was received from the original vehicle, rather than received from the cache. What’s more, to fully utilize the caching mechanism, a vehicle not only fetches other vehicles’ own trained models, but also models in their caches. For instance, at epoch $t$, vehicle $i$ can directly fetch model $x_{j}(t)$ from vehicle $j$, at line 7, and also fetch models $x_{k}(\tau)\in\mathcal{C}_{j}(t)$ from the cache of vehicle $j$, at line 8. This way, each vehicle can not only fetch the models directly from its neighbours, but also indirectly from its neighbours’ neighbours, thus boosting the spreading of the underlying data information from different vehicles, and improving the DFL convergence speed, especially with heterogeneous data distribution. Additionally, at line 2, before updating cache, models with staleness $t-\tau\geq\tau_{max}$ will be removed from each vehicle’s cache.

To evaluate the performance of DFedCache, we compare the DFL with LRU Cache, Centralized FL (CFL) and DFL on MNSIT, FashionMNIST, CIFAR-10 with three distributions: non-i.i.d, i.i.d, Dirichlet with $\pi=0.5$. For LRU update, we take the cache size as 10 for MNIST and FashionMNIST, and 3 for CIFAR-10 and $\tau_{max}=5$ based on practical considerations. Given a speed of $13.89m/s$ and and a communication distance of $100m$, the communication window of two agents driving in opposite directions could be limited. Additionally, considering the size of chosen models and communication overhead, the above cache sizes were chosen. From the results in Fig. 2, we can see that DFedCache boosts the convergence and outperforms non-caching method (DFL) and gains performance much closer to CFL in all the cases, especially in non-i.i.d. scenarios.

Then we evaluate the performance gains with different cache sizes from 1 to 30 and $\tau_{max}=10$, on MNIST and FashionMNIST in Fig. 3. LRU can benefit more from larger cache sizes, especially in non-i.i.d scenarios, as aggregation with more cached models gets closer to CFL and training with global data distribution.

As mentioned in the previous sections, one drawback of model caching is introducing stale models into aggregation, so it is very important to choose a proper staleness tolerance $\tau_{max}$. First, we statistically calculate the relation between the average number and the average age of cached models at different $\tau_{max}$ from 1 to 20, when epoch time is 30s, 60s, 120s, with unlimited cache size, in Table 2. We can see that, with the fixed epoch time, as the $\tau_{max}$ increases, the average number and average age of cached models increase, approximately linearly. It’s not hard to understand, as every epoch each agent can fetch a limited number of models directly from other agents. Increasing the staleness tolerance $\tau_{max}$ will increase the number of cached models, as well as the age of cached models. What’s more, we can see that the communication frequency or the moving speed of agents will also impact the average age of cached models, as the faster an agent moves, the more models it can fetch within each epoch, which we will further discuss later.

In practice, vehicle mobility patterns and local data distributions may naturally form groups. For example, a vehicle may mostly move within its home area, and collect data specific to that area. Vehicles within the same area meet with each others frequently, but have very similar data distributions. A fresh model of a vehicle in the same area is not as valuable as a slightly older model of a vehicle from another area. So model caching should not just consider model freshness, but should also take into account the coverage of group-based data distributions. For those scenarios, we develop a Group-based (GB) caching algorithm as Algorithm 3. More specifically, knowing that there are $m$ distribution groups, one should maintain balanced presences of models from different groups. One straightforward extension of the LRU caching algorithm is to partition the cache into $m$ sub-caches, one for each group. Each sub-cache is updated by local models from its associated group, using the LRU policy.

We now conduct a case study for group-based cache update. As shown in Fig 1, the whole Manhattan road network is divided into 3 areas, downtown, mid-town and uptown. Each area has 30 area-restricted vehicles that randomly moves within that area, and 3 or 4 free vehicles that can move into any area. We set 4 different area-related data distributions: Non-overlap, 1-overlap, 2-overlap, 3-overlap. $n$-overlap means the number of shared labels between areas is $n$. We use the same non-i.i.d shards method in the previous section to allocate data points to the vehicles in the same area. On each vehicle, we evenly divide the cache for the three areas. We evaluate our proposed GB cache method on FashionMNIST. As shown in Fig. 6, while vanilla LRU converges faster at the very beginning, it cannot outperform DFL at last. However, the GB update method can solve the problem of LRU update and outperform DFL under different overlap settings.

In general, the convergence rate of DFedCache outperforms non-caching method (DFL), especially for non-i.i.d. distributions. Larger cache size and smaller $\tau_{max}$ make DFedCache closer to the performance of CFL. Empirically, there is a trade-off between the age and number of cached models, it is critical to choose a proper $\tau_{max}$ to control the staleness. What’s more, the choice of $\tau_{max}$ should take into account the diversity data distributions on agents. In general, with non-i.i.d data distributions, the benefits of increasing the number of cached models can outweigh the damages caused by model staleness; while when data distributions are close to i.i.d, it is better to use a small number of fresh models than a large number of stale models. Similarly conclusions can also be drawn from the results of area-restricted vehicles. What’s more, in a system of moving agents, the mobility will also have big impact on the training, usually higher speed and communication frequency improve the model convergence and accuracy.