EFTViT: Efficient Federated Training of Vision Transformers with Masked Images on Resource-Constrained Edge Devices

Federated learning is a decentralized machine learning approach that enhances privacy by training models directly on client devices, only transmitting model parameters to a central server. Most studies focus on addressing data heterogeneity [22, 17, 11, 20] and privacy protection [2, 27, 5] in FL. For instance, FedProx [22] adds a proximal term to optimize the local updates for addressing data heterogeneity. Regarding privacy protection, Asad et al. [2] apply homomorphic encryption to FL, enabling clients to encrypt their local models using private keys. Shi et al. [27] propose a FL method with differential privacy (DP). However, these works rely on the ideal assumption that clients have sufficient resources to handle model training process.

Federated learning approaches on resource-constrained clients can be categorized into federated distillation (FD) [12, 15, 19, 31] and partial training (PT) [7, 1]. FD methods focus on aggregating knowledge from heterogeneous client models to a server model. For instance, FedGKT [12] trains small models on clients and periodically transfers their knowledge to a large server model via knowledge distillation. PT methods divide a global model into smaller sub-models that can be locally trained on resource-constrained clients. For instance, HeteroFL [7] randomly selects sub-models from the global model to distribute to clients. However, these methods adapt model size to clients’ capacities, rather than fundamentally addressing the computational burden of client-side training.

When dealing with transformer-based complex models, Parameter-Efficient Fine-Tuning (PEFT) [36, 16, 14] provides a practical solution for efficiently adapting pre-trained models across the various downstream tasks, which can reduce storage and computation costs by fixing most pre-trained parameters and fine-tuning only a small subset [10]. Several studies [29, 38] have explored using different PEFT techniques to assess performance improvements and resource savings in federated systems. However, the limited fine-tuning of parameters in PEFT inevitably constrains the adaptability of pre-trained models to new tasks, potentially resulting in suboptimal performance in federated systems with data heterogeneity.

We employ supervised classification tasks distributed across $K$ clients to formulate our problem. Each client $k$ possesses a dataset $D_{k}:=(X_{k},Y_{k})$, where $X_{k}\in\mathbb{R}^{N_{k}\times d_{k}}$ denotes the data samples and $Y_{k}\in\mathbb{R}^{N_{k}\times c_{k}}$ represents their corresponding labels. Here, $N_{k}$ represents the number of data points, $d_{k}$ denotes the input dimension, and $c_{k}$ indicates the number of classes.

As illustrated in Figure 3, EFTViT employs hierarchical training across clients and the central server to enable privacy-preserving and efficient collaborative learning. Each client includes a local module with $M$ Transformer layers, a shared global module with $N$ Transformer layers, and a classification head. The local module and classification head are trained on each client with unmasked image patches $X_{p}$, enabling efficient local training and generating patch features that represent local knowledge. To safeguard data distribution privacy, a median sampling strategy is applied on each client to create a balanced patch features (BPF) dataset before uploading to the server. The global module is then trained on the server using the BPF dataset from clients to effectively learn global representations for all tasks. Finally, the server transmits the updated global module parameters back to clients for next training round.

To enable efficient local training on resource-constrained clients, we present a patch-wise optimization strategy. Firstly, each input image is divided into a sequence of regular, non-overlapping patches, which are randomly masked at a ratio $r_{m}$. The remaining unmasked patches, denoted as $X_{p}$, are then used to train our framework. We define the patch features obtained by the local module on the client $k$ $\mathcal{M}_{k}(\phi_{k};\cdotp)$ as $H_{p}^{k}=\mathcal{M}_{k}(\phi_{k};X_{p}^{k})$, where $X_{p}^{k}=Mask\_D(X_{k})$ and $Mask\_D(X_{k})$ is the operation of randomly masking image patches from $X_{k}$ and discarding the selected patches. To preserve patch ordering for ViTs, the positional embeddings [28] of the remaining patches are retained. This is inspired by the internal redundancy of images and reduces the amount of data that the model needs to process, thereby lowering computational complexity. Additionally, these patch features $H_{p}^{k}$ make it pretty challenging to reconstruct the original images since they are encoded from a very small portion of each image, inherently providing EFTViT with a content privacy advantage. Notably, the entire images are adopted for the inference on each client.

To enhance privacy in EFTViT, we propose a median sampling strategy to generate a balanced patch features dataset $D_{H}^{k}$ on each client. It aims to ensure that the generated patch features on each client contain an equal number of samples for each class, thereby preventing the leakage of statistical information or user preferences when uploaded to the central server. Imbalanced data distribution on clients is a common issue in federated learning, and the median, being less sensitive to extreme values, is well-suited for addressing this challenge. Our median sampling strategy uses the median of class sample counts on each client to differentiate between minority and majority classes. It then applies oversampling to increase samples of minority classes and downsampling to reduce samples of majority classes. Specifically, for minority class samples, all patch features generated across multiple local training epochs are retained, whereas for majority class samples, only patch features from the final epoch are preserved. Next, downsampling is applied to reduce the number of samples in each class to the median. Empirically, we find that increasing the sampling threshold adds to computation costs but does not significantly improve final performance.

To effectively reduce the computational burden on clients without compromising performance, we propose a new hierarchical training strategy for ViTs that minimizes the number of trainable parameters on the clients. As aforementioned, our ViT models comprise a collection of lightweight local modules, a shared large global module and a classification head.

The overall workflow of our EFTViT is shown in  Algorithm 1 and  Algorithm 2. At the start of each round $t$, the server will randomly choose a proportion $P$ from $K$ clients to participate in training. Each client updates the parameters of its global module and classification head with those received from the server, and then initiates local training. The median sampling is applied to patch features $H_{p}^{k}$ to obscure local data distribution and produce a balanced dataset. The detailed process is presented in Algorithm 1.

The server receives the balanced patch features dataset $D_{H}^{k}$ from $P\cdot K$ clients to update the global dataset $D_{H}$, storing new client data and updating existing client data. This dataset is used to train the global module $\mathcal{M}_{S}(w;\cdot)$ and classification head $\mathcal{M}_{S}(\theta;\cdot)$, with the updated parameters $w$ and $\theta$ sent back to clients upon completion of training. The process is elaborated in Algorithm 2.

Data Content Privacy. Contrary to previous beliefs, recent studies show that exchanging intermediate features during federated learning training is safer than sharing gradients. This is because attackers only have access to evolving feature maps rather than the final, fully trained maps, making data reconstruction attacks more challenging [12, 35, 39, 41]. Furthermore, EFTViT uploads patch features corresponding to 25% of the image area, controlled by the masking rate $r_{m}$, which makes recovering the original image highly challenging, even if theoretically feasible. The masking rate can be further increased to enhance data content privacy, if necessary.

To comprehensively evaluate EFTViT, we conduct experiments on two widely used federated learning datasets, CIFAR-10 [18] and CIFAR-100 [18], as well as a more challenging datasets, UC Merced Land-Use [34], for remote sensing. CIFAR-10 and CIFAR-100 datasets each contain 60,000 color images. CIFAR-10 is organized into 10 classes, with 6,000 images per class (5,000 for training and 1,000 for testing), while CIFAR-100 has 100 classes, with 600 images per class (500 for training and 100 for testing). UC Merced Land-Use dataset contains 21 land-use classes, e.g., agricultural, forest, freeway, beach, and other classes, each with 100 images (80 for training and 20 for testing). We partition samples to all clients following a Dirichlet distribution $Dir_{N}(\beta)$ with a concentration parameter $\beta$, setting $\beta=\{0.1,1\}$ to simulate high or low levels of heterogeneity.

We use ViT-B [8] pre-trained on ImageNet-21K [26] as the backbone of our framework. The input images are resized to $224\times 224$ with a patch size of $16\times 16$. During training, data augmentation techniques such as random cropping, flipping, and brightness adjustment are applied. Following federated learning practices, we set the number of clients to 100, with a client selection ratio $P=0.1$. The AdamW optimizer is used with an initial learning rate of $5\times 10^{-5}$, weight decay of 0.05, and a cosine annealing learning rate schedule with warm-up. We use a batch size of 32 for both training and testing on each client. All experiments are conducted on a single NVIDIA GeForce RTX 3090 GPU. In each round, clients train for 5 epochs locally, while the server performs an additional 2 epochs. The framework is trained for a total of 200 rounds, requiring approximately 24 hours.

Given the lack of studies on training ViTs on resource-constrained clients, we adapt the FEDBFPT approach [33], originally designed for natural language processing tasks, as a strong baseline, which progressively optimizes the shallower layers while selectively sampling deeper layers to reduce resource consumption. To establish additional baselines, we adapt several well-known PEFT methods to our federated learning setup: (a) Fed-Head: trains only the head layer parameters; (b) Fed-Bias: applies bias-tuning [36], focusing on training only the bias terms; (c) Fed-Prompt: incorporates prompt-tuning [16], adding trainable prompt embeddings to the input; and (d) Fed-LoRA: integrates LoRA-tuning [14] by adding the LoRA module to the query and value layers. These methods use FedAVG [24] for parameter aggregation. Otherwise, our method and the baseline methods share the same settings in the federated learning scenario.

We conduct extensive ablation experiments to investigate the key components of our approach.