English

Task2vec Readiness: Diagnostics for Federated Learning from Pre-Training Embeddings

Machine Learning 2026-04-14 v1 Artificial Intelligence

Abstract

Federated learning (FL) performance is highly sensitive to heterogeneity across clients, yet practitioners lack reliable methods to anticipate how a federation will behave before training. We propose readiness indices, derived from Task2Vec embeddings, that quantifies the alignment of a federation prior to training and correlates with its eventual performance. Our approach computes unsupervised metrics -- such as cohesion, dispersion, and density -- directly from client embeddings. We evaluate these indices across diverse datasets (CIFAR-10, FEMNIST, PathMNIST, BloodMNIST) and client counts (10--20), under Dirichlet heterogeneity levels spanning α{0.05,,5.0}\alpha \in \{0.05,\dots,5.0\} and FedAVG aggregation strategy. Correlation analyses show consistent and significant Pearson and Spearman coefficients between some of the Task2Vec-based readiness and final performance, with values often exceeding 0.9 across dataset×\timesclient configurations, validating this approach as a robust proxy for FL outcomes. These findings establish Task2Vec-based readiness as a principled, pre-training diagnostic for FL that may offer both predictive insight and actionable guidance for client selection in heterogeneous federations.

Keywords

Cite

@article{arxiv.2604.10849,
  title  = {Task2vec Readiness: Diagnostics for Federated Learning from Pre-Training Embeddings},
  author = {Cristiano Mafuz and Rodrigo Silva},
  journal= {arXiv preprint arXiv:2604.10849},
  year   = {2026}
}
R2 v1 2026-07-01T12:05:21.503Z