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

2026-04-12

AI Summary

This work addresses the challenge that federated learning performance is highly sensitive to client data heterogeneity, yet lacks effective pre-training prediction methods. The study introduces Task2Vec embeddings into federated learning for the first time and proposes an unsupervised readiness metric that quantifies client task alignment by measuring the cohesion, dispersion, and density of task embeddings. This metric enables informed client selection prior to training and offers both interpretability and practical utility. Extensive experiments across multiple datasets under the FedAvg framework demonstrate a strong correlation between the proposed metric and final model performance, with Pearson and Spearman correlation coefficients consistently exceeding 0.9.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Natural Language Processing: Safety and Robustness
• Humans and AI: Human-in-the-loop Machine Learning

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web evaluation methodologies and metrics
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

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 $α\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.

Problem

• Federated Learning
• Heterogeneity
• Pre-Training Diagnostics
• Task2Vec
• Readiness Assessment

Innovation

• Task2Vec
• Federated Learning
• Readiness Index
• Pre-training Diagnostics
• Heterogeneity