OnDev-LCT: On-Device Lightweight Convolutional Transformers towards federated learning

TL;DR

Federated learning on edge devices demands compact, compute-efficient vision models that cope with non-IID data. The paper introduces OnDev-LCT, a lightweight convolutional transformer that embeds image-specific inductive biases via an LCT tokenizer (depthwise separable convolutions and residual linear bottlenecks) and a ViT-like encoder with multi-head self-attention, followed by SeqPool for classification. Across centralized and FL experiments on CIFAR, EMNIST/FEMNIST, and ImageNet-32, OnDev-LCT variants outperform several lightweight baselines while using far fewer parameters and MACs, and show robustness to data heterogeneity and communication constraints. These results highlight the practicality of OnDev-LCT for on-device vision tasks in federated settings and point to its potential for real-world edge AI deployments.

Abstract

Federated learning (FL) has emerged as a promising approach to collaboratively train machine learning models across multiple edge devices while preserving privacy. The success of FL hinges on the efficiency of participating models and their ability to handle the unique challenges of distributed learning. While several variants of Vision Transformer (ViT) have shown great potential as alternatives to modern convolutional neural networks (CNNs) for centralized training, the unprecedented size and higher computational demands hinder their deployment on resource-constrained edge devices, challenging their widespread application in FL. Since client devices in FL typically have limited computing resources and communication bandwidth, models intended for such devices must strike a balance between model size, computational efficiency, and the ability to adapt to the diverse and non-IID data distributions encountered in FL. To address these challenges, we propose OnDev-LCT: Lightweight Convolutional Transformers for On-Device vision tasks with limited training data and resources. Our models incorporate image-specific inductive biases through the LCT tokenizer by leveraging efficient depthwise separable convolutions in residual linear bottleneck blocks to extract local features, while the multi-head self-attention (MHSA) mechanism in the LCT encoder implicitly facilitates capturing global representations of images. Extensive experiments on benchmark image datasets indicate that our models outperform existing lightweight vision models while having fewer parameters and lower computational demands, making them suitable for FL scenarios with data heterogeneity and communication bottlenecks.

Figures (11)

• Figure 1: Performance comparison on CIFAR-10 dataset in the centralized scenario. The comparison is performed among small model variants, i.e., #Params $<$ 1.3M, while constraining the computational budget within $0.1 \times 10^9$ MACs. Each bubble's area is proportional to the model size.
• Figure 2: The vanilla federated averaging (FedAvg) framework. A typical four-step procedure is iterated until convergence.
• Figure 3: An overview of On-Device Lightweight Convolutional Transformer (OnDev-LCT). The middle row depicts the architecture of an LCT tokenizer, visualizing its core components in the bottom row. The structure of an LCT encoder is illustrated on the right.
• Figure 4: Detailed data partitions on CIFAR-10 dataset using Dirichlet distribution. Smaller $\beta$ corresponds to higher data heterogeneity. Best viewed in color.
• Figure 5: t-SNE visualizations of feature embeddings on CIFAR-10 test set learned by centralized models. Best viewed in color.