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

TL;DR

This paper introduces Task2Vec-based readiness indices to predict federated learning performance before training, using unsupervised metrics from client embeddings, validated across multiple datasets and heterogeneity levels.

Contribution

It presents a novel pre-training diagnostic method for federated learning based on Task2Vec embeddings, enabling better anticipation of federation outcomes.

Findings

High correlation (often >0.9) between readiness indices and final performance.

Readiness metrics reliably predict federated learning success across datasets and heterogeneity levels.

Task2Vec-based diagnostics can guide client selection and improve federated learning strategies.

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 $\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$\times$client 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.