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

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.

• The paper introduces FedMITR, a server-side one-shot federated learning method that combines attention-guided sparse model inversion with token relabeling to improve synthetic data and distillation under non-IID client data.
• Sparse inversion suppresses low-attention background patches, while token relabeling applies reliable pseudo-labels to informative tokens and ensemble-based soft targets to uncertain regions, reducing noisy gradients and improving stability.
• FedMITR outperforms leading baselines by up to 8.92 percentage points under moderate heterogeneity and achieves 45.26% versus 22.49% accuracy on Mini-ImageNet at extreme heterogeneity, using only one communication round.

Motivation and problem setting

One-shot federated learning (FL) restricts client–server interaction to a single communication round, in which clients transmit pre-trained local models $\{f_i(\cdot;\theta_i)\}_{i=1}^N$ and the server must aggregate them into a global model $f_S(\cdot;\theta_S)$. This eliminates the communication overhead, connection-drop fragility, and interception risk of multi-round FL, but it also removes the server's ability to refine the aggregate iteratively. Under severe non-IID data heterogeneity, parameter averaging is ineffective, so existing methods rely on data-free knowledge distillation (DFKD): a generator is trained on the server to synthesize pseudo-data that fools the ensemble $E_S(x)=\sum_i w_i f_i(x)$, and the ensemble is then distilled into the global model.

The paper identifies a structural weakness in this pipeline: under strong heterogeneity, local models encode divergent knowledge, so generators trained against the ensemble produce low-fidelity samples with severe semantic misalignment between synthesized images and their assigned labels. A t-SNE analysis shows that features of traditionally generated data lack class boundaries, whereas inversion-based synthesis yields more separable clusters. The proposed framework, FedMITR (Federated Model Inversion and Token Relabel), replaces generator training with direct model inversion on ViTs and exploits the token structure of transformers to salvage even low-quality synthetic patches (2605.10748).

The FedMITR method

FedMITR is entirely server-side: clients perform standard local training and upload models once; no auxiliary datasets, extra transmissions, or client-side modifications are required. Server training alternates two stages.

Sparse model inversion. Rather than training a generator, the server synthesizes class-conditional images by optimizing random noise inputs against each local model. The inversion objective combines a classification loss enforcing the sampled label $\hat{y}$, a negatively scaled Jensen–Shannon divergence between the local model and the current server model for diversity, and total-variation plus $\ell_2$ regularization on the input. Unlike dense inversion, which optimizes all patches—including backgrounds that converge to high-frequency noise—the method uses the attention weights $\bm{a}_\text{cls}$ from the previous iteration as an information-density indicator and masks out the lowest-attention fraction $r$ of patches, halting their inversion.

Token relabel. During distillation, tokens are treated differentially according to density. High-information-density patches $\hat{x}_h$ are trained with the generated pseudo-label via cross-entropy; low-density patches $\hat{x}_l$, whose pseudo-labels are unreliable under heterogeneity, are relabeled by the ensemble and matched through KL divergence:

$$\mathcal{L}_\text{TR}=\mathcal{L}_\text{KD}+\lambda_1\mathcal{L}_\text{CLS}(\theta_S(\hat{x}_h),\hat{y})+\lambda_2\mathcal{L}_\text{KL}(E_S(\hat{x}_l),f_S(\hat{x}_l)).$$

The two stages iterate, progressively improving both synthetic data quality and the global model.

Theoretical analysis

The generalization argument proceeds through algorithmic stability. By Hardt et al.'s bound, a $f_S(\cdot;\theta_S)$0-uniformly stable algorithm has generalization gap controlled by $f_S(\cdot;\theta_S)$1, and for SGD with decaying step sizes $f_S(\cdot;\theta_S)$2 increases monotonically with the gradient Lipschitz constant $f_S(\cdot;\theta_S)$3. The paper therefore proves $f_S(\cdot;\theta_S)$4 relative to dense inversion, via two mechanisms:

• Sparsity as gradient truncation. For the value-projection gradient $f_S(\cdot;\theta_S)$5, masking background tokens zeroes their contribution exactly, eliminating the "random walk" induced by fitting hard labels to noise.
• Relabeling as variance reduction. Under an assumption that inverted background regions are statistically orthogonal to the label direction (so $f_S(\cdot;\theta_S)$6), replacing hard targets with soft ensemble outputs reduces gradient variance per Menon et al.'s distillation analysis, i.e., $f_S(\cdot;\theta_S)$7.

Combining these yields a strictly smaller Lipschitz constant, hence smaller stability parameter and a tighter high-probability generalization bound than dense inversion. Two caveats bear directly on this result's strength: the orthogonality assumption on background noise is asymptotic and heuristic rather than verified empirically, and the proof analyzes only the $f_S(\cdot;\theta_S)$8 gradient within MSA layers while treating the Lipschitz constant decomposition approximately ($f_S(\cdot;\theta_S)$9). The claim of strict inequality therefore rests on idealized assumptions about the statistics of inverted noise.

Empirical results

Experiments use 10 clients with DeiT/16 backbones (primarily Tiny, plus Base and ViT-Small ablations) on CIFAR10, CIFAR100, OfficeHome, and Mini-ImageNet, under Dirichlet ($E_S(x)=\sum_i w_i f_i(x)$0) and pathological ($E_S(x)=\sum_i w_i f_i(x)$1) partitions, against FedAvg, FedFTG, DENSE, Co-Boosting, and DeepInversion adapted to one-shot FL.

Key findings:

• Large gains at extreme heterogeneity. Under $E_S(x)=\sum_i w_i f_i(x)$2, FedMITR exceeds the best baseline (DeepInversion) by 3.20% on CIFAR10, 8.92% on OfficeHome, and 7.93% on Mini-ImageNet. At $E_S(x)=\sum_i w_i f_i(x)$3 the gaps widen dramatically: 45.26% vs. 22.49% on Mini-ImageNet and 24.05% vs. 11.88% on OfficeHome.
• Extreme label scarcity. With one class per client on CIFAR10, FedMITR reaches 13.98% versus 9.77% for DeepInversion; with 10 classes per client on Mini-ImageNet, 21.81% versus 13.59%.
• Ablations. Inversion with pseudo-labels alone ("Inversion + PL") outperforms inversion with plain KD, and full token relabeling adds further gains across all three backbone architectures; removing either loss component degrades accuracy substantially.
• Communication efficiency. On Mini-ImageNet under $E_S(x)=\sum_i w_i f_i(x)$4, FedMITR attains 68.21% with 218 MB total communication (one round), whereas FedAvg needs roughly 100 rounds to match it and SCAFFOLD uses 5450 MB over 50 rounds for 69.78%. Scaling to 50 clients, FedMITR retains 49.33% versus 30.40% for DeepInversion.
• Privacy observation. Visualizations show synthetic images bear little resemblance to training data, suggesting limited leakage of client-specific content—though the paper offers no formal privacy guarantee (e.g., differential privacy) for the inversion procedure.

Two results qualify the headline claims. First, at low heterogeneity ($E_S(x)=\sum_i w_i f_i(x)$5 on Mini-ImageNet), DeepInversion surpasses FedMITR; the authors attribute this to well-trained local models not benefiting from relabeling. Second, absolute accuracies remain modest in the most heterogeneous regimes (e.g., ~19% on CIFAR10 at $E_S(x)=\sum_i w_i f_i(x)$6 with DeiT-Tiny), indicating that one-shot aggregation itself remains fundamentally difficult regardless of synthesis strategy.

Limitations and open questions

The paper concedes several points explicitly. Performance degrades below DeepInversion when clients are already well-trained, so the relabeling mechanism appears beneficial only when pseudo-label quality is genuinely poor—an operating regime the paper does not characterize precisely. The theoretical guarantees depend on unverified assumptions (isotropic background noise, statistical orthogonality, bounded losses) and cover only part of the ViT parameter space. Hyperparameter sensitivity analyses show that too large a pseudo-label weight $E_S(x)=\sum_i w_i f_i(x)$7 or mask ratio $E_S(x)=\sum_i w_i f_i(x)$8 hurts performance, implying manual tuning matters. Open questions include whether the stability-based bound can be extended beyond $E_S(x)=\sum_i w_i f_i(x)$9 gradients, whether formal privacy certificates can accompany inversion-based synthesis, and how the approach behaves with heterogeneous client architectures or larger client counts where accuracy still declines sharply.

Conclusion

FedMITR reframes server-side data synthesis in one-shot FL around ViT-native primitives: attention-guided sparse inversion suppresses noisy background gradients, and density-aware token relabeling converts unreliable pseudo-labels into stable distillation signals. The stability-based theory offers a coherent account of why these choices tighten generalization bounds relative to dense inversion, and the empirical gains—up to nearly doubling baseline accuracy at extreme heterogeneity—are substantial. The contribution is strongest precisely where one-shot FL is hardest, though its advantages diminish as client data becomes more homogeneous and its theoretical guarantees rest on idealized assumptions about inverted noise statistics.