Task2Vec Readiness: Diagnostics for Federated Learning Performance from Pre-Training Embeddings

ICLR 2026 Conference Submission22170 Authors

Keywords: Federated Learning, Task2Vec, Readiness Index, Client Heterogeneity, Performance Prediction

TL;DR: We propose Task2Vec Readiness, a pre-training diagnostic that uses task embeddings as a proxy for federated learning performance.

Abstract:

Federated Learning (FL) has emerged as a central paradigm for collaborative model training under privacy and communication constraints. However, its performance may be hindered by client heterogeneity, data imbalance, and federation size. Despite significant algorithmic advances in optimization and aggregation, practitioners still lack principled tools to predict, before training, whether a given federation is likely to succeed. This absence of pre-training diagnostics forces costly trial-and-error experimentation and slows both research progress and deployment.

We introduce \emph{Task2Vec Readiness}, a framework that leverages Task2Vec embeddings to derive a quantitative readiness index for federated learning. Our approach transforms each client’s data distribution into a fixed-dimensional embedding via Fisher Information geometry, and then evaluates unsupervised metrics of federation structure. In particular, we measure cohesion (average cosine similarity among client embeddings), dispersion (average distance from the federation centroid), and density (RBF-kernel similarity over pairwise Euclidean distances). These metrics can be computed before training, and together they form a readiness profile that anticipates how well the federation can support collaborative optimization. The key novelty lies in repurposing task embeddings from transfer learning into a diagnostic signal for distributed training under heterogeneity.

We conduct extensive experiments across four benchmark datasets (CIFAR-10, FEMNIST, PathMNIST, BloodMNIST), with client counts ranging from 10 to 20 and Dirichlet non-IID partitions spanning $\alpha \in {0.05,\dots,5.0}$. Correlation analyses consistently reveal significant Pearson and Spearman coefficients between readiness metrics and final model performance, frequently exceeding $0.9$ across dataset$\times$client conditions. This validates readiness as a robust proxy for FL performance. Importantly, the results hold across different sources of heterogeneity, indicating that readiness captures structural properties of federations rather than dataset-specific artifacts.

Our contributions are twofold: (i) we introduce the first readiness index for federated learning based on task embeddings and (ii) we demonstrate its predictive validity across diverse datasets and heterogeneity regimes. By moving the focus from post-hoc evaluation to pre-training diagnostics, our framework not only provides a new lens for understanding federated learning dynamics, but also offers a practical tool for improving performance, guiding client selection, and enhancing the efficiency and reliability of federated training at scale.

Primary Area: other topics in machine learning (i.e., none of the above)

Submission Number: 22170