Feature-Aware One-Shot Federated Learning via Hierarchical Token Sequences

TL;DR

This paper tackles the challenge of one-shot federated learning under non-IID data by introducing FALCON, a framework that generates hierarchical token sequences from pretrained encoders using hierarchical scale encoding (HSE) and models their distribution with a multi-scale autoregressive (M-AR) transformer. It further enhances global learning through a distillation-guided training regime at the server, leveraging an ensemble of local classifiers to supervise the global model. The approach yields strong, robust accuracy improvements across natural and medical image datasets, with notable gains over strong OSFL baselines and improved efficiency compared to diffusion- or GAN-based data generation. The work demonstrates the practicality of combining feature-aware representations, structured synthetic data, and knowledge distillation to enable effective, privacy-preserving one-shot federated classification in diverse non-IID settings.

Abstract

One-shot federated learning (OSFL) reduces the communication cost and privacy risks of iterative federated learning by constructing a global model with a single round of communication. However, most existing methods struggle to achieve robust performance on real-world domains such as medical imaging, or are inefficient when handling non-IID (Independent and Identically Distributed) data. To address these limitations, we introduce FALCON, a framework that enhances the effectiveness of OSFL over non-IID image data. The core idea of FALCON is to leverage the feature-aware hierarchical token sequences generation and knowledge distillation into OSFL. First, each client leverages a pretrained visual encoder with hierarchical scale encoding to compress images into hierarchical token sequences, which capture multi-scale semantics. Second, a multi-scale autoregressive transformer generator is used to model the distribution of these token sequences and generate the synthetic sequences. Third, clients upload the synthetic sequences along with the local classifier trained on the real token sequences to the server. Finally, the server incorporates knowledge distillation into global training to reduce reliance on precise distribution modeling. Experiments on medical and natural image datasets validate the effectiveness of FALCON in diverse non-IID scenarios, outperforming the best OSFL baselines by 9.58% in average accuracy.

Figures (6)

• Figure 1: Illustration of the impact of locally specialized encoders on generalization. Colors indicate different locally specialized encoders; shapes indicate samples from non-IID datasets. (1) Global classifier training: Features from locally specialized encoders are aggregated to train a global classifier. (2) Local inference: When a client's specialized encoder encounters data that is non-IID with respect to its local data, the extracted features cannot be correctly recognized by the global classifier.
• Figure 2: The framework of FALCON. Each client (1) performs hierarchical scale encoding (HSE), (2) trains a local M-AR transformer generator (3) and classifier, and uploads synthetic token sequences and the classifier to the server for (4) distillation-guided global classifier training.
• Figure 3: Illustration of hierarchical scale encoding (HSE).
• Figure 4: Architecture of the M-AR Transformer Generator.
• Figure 5: Computational cost (bars, GFLOPs per sample) and average classification accuracy (dots, %) for OSFL methods involving sample generation.