Semantic-Aware Task Clustering

• The paper introduces semantic-aware task clustering, a novel method that partitions tasks based on representational similarity using metrics like Jensen–Shannon divergence.
• It details an algorithmic framework in federated learning that mitigates negative transfer by confining parameter sharing within semantically-cohesive clusters.
• Empirical results demonstrate faster convergence and an average accuracy increase of up to 7.18 percentage points compared to unclustered methods.

Semantic-aware task clustering refers to algorithmic methodologies that partition a collection of learning tasks into groups based on semantic similarity, as inferred from information-theoretic or representational criteria, rather than naive feature proximity or superficial data characteristics. In distributed and federated learning environments, and in heterogeneous multitask or few-shot scenarios, semantic-aware clustering mitigates negative information transfer by ensuring that model parameter sharing or knowledge transfer is confined to semantically-cohesive clusters. Distinct from classical clustering, these methods operate in low-dimensional semantic spaces or utilize hierarchical structure to delineate task relations, enabling more robust cooperative learning and adaptation across diverse or non-IID (non-identically distributed) task regimes (Razlighi et al., 24 Jan 2026, &&&1&&&).

1. Foundations of Semantic-aware Task Clustering

Semantic-aware clustering methods emerge from the need to distinguish between positive and negative task transfer in distributed multitask systems. Traditional approaches—often based on high-dimensional input or feature space similarity—fail in heterogeneous scenarios because they do not capture semantic congruence at the task level. Semantic-aware clustering leverages intrinsic properties of the underlying task distributions or learned representations, providing clusters where intra-group knowledge transfer is statistically and semantically justified (Razlighi et al., 24 Jan 2026).

The framework formalizes the similarity between tasks using a metric in the semantic space, most notably the Jensen–Shannon (JS) divergence between the empirical distributions of semantic variables associated with each task. The normalized semantic-similarity score is then used to threshold and relate tasks for cluster formation.

2. Mathematical Formulation and Clustering Criteria

In semantic-aware task clustering for federated multi-task semantic communication, consider $N$ tasks (e.g., users or satellites), each associated with a semantic random variable $Z_i$ over a finite alphabet $\mathcal{Z}$. Each task locally computes an empirical probability mass function (pmf) $\pi_i(c)$ for $c \in \mathcal{Z}$.

The pairwise similarity between tasks $i$ and $j$ is defined as: $\omega_{ij} = 1 - D_{\mathrm{JS}}(\pi_i \|\ \pi_j)$ where $D_{\mathrm{JS}}$ denotes the JS divergence. This $\omega_{ij}$ attains its maximum $(=1)$ for identical distributions, and minimum $(=0)$ when distributions are maximally dissimilar.

Clustering seeks a partition $\{\mathcal{T}_1,\dots,\mathcal{T}_K\}$ maximizing intra-cluster similarity: $$\text{maximize} \quad \sum_{k=1}^K \sum_{i,j\in\mathcal{T}_k} \omega_{ij}$$ However, practical implementation uses a binary threshold $\tau$ to establish edges:

• $r_{ij} = 1$ iff $\omega_{ij} \geq \tau$, else $0$. Transitivity $(r_{ij}+r_{jk} \leq r_{ik}+1)$ is enforced to prevent overlapping clusters. The graph induced by $R$ is partitioned into connected components, yielding the disjoint clusters (Razlighi et al., 24 Jan 2026).

3. Algorithmic Implementation in Federated Learning

The clustering process comprises:

• Local estimation on each device of $\pi_i(c)$.
• Communication of $\pi_i$ to the parameter server.
• Central computation of $\omega_{ij}$ for all pairs.
• Construction of the binary relation matrix $R$ and enforcement of transitivity.
• Extraction of clusters as connected components of the associated undirected graph.

In each federated learning round:

• Devices update local model parameters (e.g., semantic encoder $\theta_i$) via task-specific objectives (e.g., InfoMax for maximizing $I(Y_i;Z_i)$).
• Encoders within each cluster are averaged: $$\bar{\theta}_k = \frac{1}{|\mathcal{T}_k|} \sum_{i\in\mathcal{T}_k} \theta_i$$
• Devices receive updated cluster-specific encoders, facilitating intra-cluster parameter sharing only (Razlighi et al., 24 Jan 2026).

4. Empirical and Practical Considerations

Experimental evaluation in low Earth orbit (LEO) satellite networks, wherein each satellite executes a distinct classification task on MNIST, demonstrates key properties:

• Semantic-aware clustering preserves or improves accuracy for semantically related tasks (e.g., achieving $>98\%$ and $97\%$ accuracy on closely related tasks in 15 rounds), while unrelated tasks remain unaffected in performance.
• In contrast, unclustered federated averaging degrades the accuracy of unrelated tasks (e.g., dropping from $94\%$ to $85\%$), due to negative transfer (Razlighi et al., 24 Jan 2026).
• Clustered federated learning converges faster, requiring $\sim30\%$ fewer rounds to reach 90% accuracy, and yields average accuracy increases of up to $+7.18$ percentage points over unclustered FL.

Key hyperparameters include choice of threshold $\tau$ (typically selected between min and max observed $\omega_{ij}$), network shape (e.g., two-layer encoders/decoders), and training parameters consistent with distributed regimes. The computational complexity is dominated by the $O(N^2f)$ pairwise similarity computation, but is practical for moderate $N$ and small $|\mathcal{Z}|$ (Razlighi et al., 24 Jan 2026).

5. Extensions: Hierarchical and Self-supervised Task Clustering

In few-shot text classification, a related paradigm organizes tasks hierarchically using learned task embeddings. For example, the self-supervised hierarchical task clustering (SS-HTC) method extracts a task embedding $g_{\rm in}$ via average pooling over label-augmented example encodings (e.g., BERT outputs), supervised by a Label-Oriented Masked Language Modeling loss. Tasks are organized in an $L$-level soft hierarchy, with cluster assignments computed via Gaussian-kernel softmax over embedding distances; embeddings at higher levels aggregate lower-level task information using gated sums. This hierarchy induces a multi-granular soft clustering, capturing both coarse and fine-grained inter-task relations (Zha et al., 2022).

Cluster-specific adapters (akin to FiLM layers) then warp feature representations for each cluster, and downstream prediction leverages these for prototype-based few-shot learning.

The result is dynamic, semantics-driven task grouping, with cluster assignments reflecting label and data semantics rather than arbitrary data source distinctions. Empirical analysis on multiple public benchmarks for few-shot classification yields substantial gains over non-clustered models (e.g., +9.81% in 1-shot, +5.19% in 5-shot), confirming the value of semantic granularity.

6. Theoretical and Practical Implications

Both the information-theoretic and hierarchical representational approaches to semantic-aware clustering address the fundamental problem of destructive negative transfer in distributed or multitask regimes. By explicitly measuring or disentangling semantic affinity, these methods offer principled, efficient task partitioning that scales to heterogeneous, real-world settings (e.g., satellite networks, multi-source few-shot learning).

A plausible implication is that such clustering will become foundational to scalable federated learning in non-IID environments, enabling constructive transfer and communication-efficient coordination across diverse application domains.

7. Relation to Broader Research and Future Directions

Semantic-aware task clustering occupies an intersection of federated learning, information theory, transfer learning, and few-shot adaptation. Techniques differ from traditional clustering by prioritizing semantic abstraction and minimizing reliance on high-dimensional data features, a critical distinction in privacy-sensitive or communication-constrained settings.

Future research may further integrate these approaches with domain adaptation and continual learning, investigate richer semantic representations (beyond empirical pmfs or LLM embeddings), and develop scalable algorithms for larger and more diverse task sets (Razlighi et al., 24 Jan 2026, Zha et al., 2022).