Provable Sparse Inversion and Token Relabel Enhanced One-shot Federated Learning with ViTs

TL;DR

FedMITR is a novel federated learning framework that enhances one-shot training with sparse inversion and token relabeling for ViTs, improving data quality and model generalization in non-IID scenarios.

Contribution

The paper introduces FedMITR, combining sparse model inversion and token relabeling strategies, with theoretical analysis and empirical results showing significant performance improvements.

Findings

FedMITR outperforms existing methods in various non-IID federated learning settings.

Sparse model inversion reduces gradient instability caused by background noise.

Token relabeling decreases gradient variance, enhancing model generalization.

Abstract

One-Shot Federated Learning, where a central server learns a global model in a single communication round, has emerged as a promising paradigm. However, under extremely non-IID settings, existing data-free methods often generate low-quality data that suffers from severe semantic misalignment with ground-truth labels. To overcome these issues, we propose a novel Federated Model Inversion and Token Relabel (FedMITR) framework, which trains the global model by fully exploiting all patches of synthetic images. Specifically, FedMITR employs sparse model inversion during data generation, selectively inverting semantic foregrounds while halting the inversion of uninformative backgrounds. To address semantically meaningless tokens that hinder ViT predictions, we implement a differentiated strategy: patches with high information density utilize generated pseudo-labels, while patches with low information density are relabeled via ensemble models for robust distillation. Theoretically, our analysis based on algorithmic stability reveals that Sparse Model Inversion eliminates gradient instability arising from background noise, while Token Relabel effectively reduces gradient variance, collectively guaranteeing a tighter generalization bound. Empirically, extensive experimental results demonstrate that FedMITR substantially outperforms existing baselines under various settings.