Sheaf-Based Federated Representation Learning

# Sheaf-Based Federated Representation Learning
Source: [https://arxiv.org/html/2608.10016](https://arxiv.org/html/2608.10016)
Mario Edoardo Pandolfo &Leonardo Di Nino &Sergio Barbarossa& Paolo Di Lorenzo

###### Abstract

Heterogeneous federated systems require agents to learn and exchange informative representations despite differences in data distributions, sensing modalities, model architectures, latent dimensionalities, and local learning objectives\. To address this challenge, we propose Sheaf\-based Federated Representation Learning \(SFRL\), a general framework that jointly optimizes local objectives with a manifold\-constrained geometric alignment regularizer based on learnable sheaf restriction maps\. Unlike most existing approaches, SFRL does not assume a shared global latent space\. Instead, global consistency emerges from the alignment of neighboring latent representations through orthogonal transformations and isometric embeddings\. This alignment is enforced by a quadratic gluing regularizer induced by the sheaf Laplacian, whose learnable restriction maps adapt the geometry to the observed data\. The penalty is evaluated on a small set of shared pilot samples, ensuring scalability and communication efficiency\. We develop a decentralized algorithm for solving SFRL, termed Sheaf\-FRL, which alternates between gradient updates of the local models and closed\-form Procrustes updates of the edge\-wise restriction maps\. We further establish convergence of Sheaf\-FRL to first\-order stationary points in both deterministic and stochastic settings\. As an application, we consider a cooperative classification task in the context of semantic communication, under model and data heterogeneity\. Our results show that Sheaf\-FRL outperforms baseline approaches in terms of local and post\-communication classification accuracy across different levels of local distribution shift and exhibits greater robustness to latent\-space dimensionality compression\.

\*\*footnotetext:These authors contributed equally\. All authors are with Sapienza University of Rome, Italy\. Correspondence to:gabriele\.dacunto@uniroma1\.itandenrico\.grimaldi@uniroma1\.it\.## 1Introduction

Modern machine learning systems increasingly operate in decentralized and heterogeneous environments, where data, models, sensing modalities, and computational resources are distributed across multiple agents, devices, or institutions\. In this context, we adopt the term*agent*to denote any physical or artificial entity performing a task in an environment\. Privacy constraints, regulatory requirements, and communication limitations often prevent centralized data aggregation, motivating*federated learning*\(FL\), where agents optimize local models on private data while exchanging limited information rather than raw samples\(Kairouz et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib35)\)\.

A central challenge in such systems is that agents may learn representations that are useful locally but incompatible across the network\. Existing FL methods commonly rely on parameter averaging, global regularization, or shared encoders, and therefore implicitly assume that features learned by different agents are directly comparable\(McMahan et al\.,[2017](https://arxiv.org/html/2608.10016#bib.bib45); Scardapane & Di Lorenzo,[2017](https://arxiv.org/html/2608.10016#bib.bib53); Li et al\.,[2020b](https://arxiv.org/html/2608.10016#bib.bib41); Karimireddy et al\.,[2020](https://arxiv.org/html/2608.10016#bib.bib36); Collins et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib11); Li et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib40); Dong & Voiculescu,[2021](https://arxiv.org/html/2608.10016#bib.bib19)\)\. This assumption can fail when agents differ in their data distributions, sensing modalities, architectures, latent dimensions, or local objectives\. Independently trained encoders may produce representations that encode equivalent task\-relevant information while differing by rotations, reflections, or more general isometric embeddings\. Such ambiguity reflects the inherent*gauge freedom*of representation learning objectives, whereby equivalent solutions may differ by transformations that preserve geometric structure\(Wang & Isola,[2020](https://arxiv.org/html/2608.10016#bib.bib65)\)\. Consequently, forcing all agents into a single shared latent space can be unnecessarily restrictive and may degrade local representation quality or hinder knowledge transfer\.

These observations motivate*semantic alignment*: rather than requiring agents to adopt identical latent coordinates, one should learn how their local representation spaces are related\. This perspective is supported by recent work on relative representations, which shows that independently learned latent spaces can support zero\-shot communication when their relational structure is preserved\(Moschella et al\.,[2023](https://arxiv.org/html/2608.10016#bib.bib48)\)\. Semantic alignment can support heterogeneous federated representation learning across supervised, semi\-supervised, and self\-supervised settings, as well as cooperative tasks in which agents exploit representations received from their neighbors\. For example, it is relevant to multimodal sensing, collaborative robotics, distributed channel charting, transfer across heterogeneous models, and task\-aware semantic communication, where latent information generated by one agent must remain useful to another despite differences in their internal representations\.

A natural mathematical framework for these settings is a*cellular sheaf*on graph, hereafter referred to as a*network sheaf*\(Curry,[2014](https://arxiv.org/html/2608.10016#bib.bib12)\)\. A network sheaf associates a local vector space with each node and linear restriction maps with graph edges, thereby encoding how information in neighboring spaces should be compared or transported\. Recent work in graph signal processing and graph neural networks has shown that sheaf\-based models can represent vector\-valued data coupled through local consistency constraints\(Hansen & Ghrist,[2019](https://arxiv.org/html/2608.10016#bib.bib30); Di Nino et al\.,[2024](https://arxiv.org/html/2608.10016#bib.bib15);[2026](https://arxiv.org/html/2608.10016#bib.bib16); Bodnar et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib7); Barbero et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib5)\)\. This provides a natural abstraction for federated representation learning: each agent maintains its own latent space, while learnable edge\-wise maps encode the geometric relations required to align representations across neighboring agents\. Global consistency then emerges from local compatibility relations, rather than from a shared latent space\.

Contributions\.We propose a sheaf\-based framework for federated representation learning that enables semantic alignment across heterogeneous agents without requiring a shared latent space\. Specifically,*\(i\)*we formulate federated representation learning over a learnable network sheaf, where agent\-specific latent spaces are related through edge\-wise geometric transport maps;*\(ii\)*we introduce a quadratic gluing penalty induced by a learnable sheaf Laplacian, evaluated on a small set of reference representations, hereinafter*pilots*, to promote alignment while limiting communication;*\(iii\)*we employ orthogonal and Stiefel restriction maps to capture rotations, reflections, and isometric embeddings across latent spaces of possibly different dimensions;*\(iv\)*we develop a fully decentralized alternating algorithm, termed Sheaf\-FRL, that combines local gradient updates with closed\-form Procrustes updates of the transport maps, and establish convergence to first\-order stationary points in deterministic and stochastic settings; and*\(v\)*we demonstrate that Sheaf\-FRL improves over baseline approaches in learning and representation transfer for collaborative supervised classification in semantic communication, under model and data heterogeneity\. The code implementing the proposed method and reproducing all the experiments is publicly available\.111[https://github\.com/SPAICOM/sheaf\-based\-federated\-representation\-learning\.git](https://github.com/SPAICOM/sheaf-based-federated-representation-learning.git)

## 2Related Work

In this section, we review the most relevant literature on federated learning, representation learning, and geometric alignment, and position our work with respect to these directions\.

Federated learning under heterogeneity\.Federated learning \(FL\) trains a shared model through repeated local stochastic gradient updates followed by either server\-side aggregation\(McMahan et al\.,[2017](https://arxiv.org/html/2608.10016#bib.bib45)\)or neighbor\-based combination over a communication graph\(Scardapane & Di Lorenzo,[2017](https://arxiv.org/html/2608.10016#bib.bib53)\)\. A large body of work studies the impact of statistical and system heterogeneity, which can lead to client drift and unstable convergence, and proposes stabilizing mechanisms such as proximal regularization \(FedProx\)\(Li et al\.,[2020b](https://arxiv.org/html/2608.10016#bib.bib41)\)and control variates \(SCAFFOLD\)\(Karimireddy et al\.,[2020](https://arxiv.org/html/2608.10016#bib.bib36)\)\. More broadly, these approaches aim to improve optimization stability while maintaining a*single global model*\. In contrast, we focus on settings where enforcing a shared latent representation is inherently restrictive\.

Federated multi\-task learning\.Federated multi\-task learning \(FMTL\) extends FL by training client\-specific models while exploiting relationships across tasks or clients\(Smith et al\.,[2017](https://arxiv.org/html/2608.10016#bib.bib59)\)\. Early formulations build upon regularized multi\-task learning objectives, where task relationships are captured through shared covariance or relational structures\(Evgeniou & Pontil,[2004](https://arxiv.org/html/2608.10016#bib.bib21); Evgeniou et al\.,[2005](https://arxiv.org/html/2608.10016#bib.bib22); Jacob et al\.,[2008](https://arxiv.org/html/2608.10016#bib.bib33); Zhang & Yeung,[2010](https://arxiv.org/html/2608.10016#bib.bib69); Zhang & Yang,[2017](https://arxiv.org/html/2608.10016#bib.bib68)\)\. In the federated setting, this leads to approaches that jointly optimize local models and their coupling structure via alternating or primal\-dual procedures\(Smith et al\.,[2017](https://arxiv.org/html/2608.10016#bib.bib59); Jaggi et al\.,[2014](https://arxiv.org/html/2608.10016#bib.bib34); Ma et al\.,[2015](https://arxiv.org/html/2608.10016#bib.bib43)\)\. A prominent class of methods leverages graph\-based regularization, where a predefined or learned graph encodes similarities among clients and enforces smoothness of model parameters across the network\. In particular, Laplacian\-based approaches penalize discrepancies between neighboring models through total variation terms\(Dinh et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib18); Issaid et al\.,[2025](https://arxiv.org/html/2608.10016#bib.bib32)\)\. While these methods capture task relatedness, they operate in the*parameter space*and typically assume homogeneous representations\. In contrast, our approach operates in the*representation space*and models relationships between clients up to learnable transformations, enabling alignment across heterogeneous latent spaces and accommodating different architectures or dimensions\.

Multimodal Federated Learning\.Multimodal Federated Learning \(MFL\)\(Che et al\.,[2023](https://arxiv.org/html/2608.10016#bib.bib9)\)extends classical FL to settings where clients possess data from multiple modalities \(e\.g\., text, images, signals\), enabling richer representations while preserving privacy\. Existing MFL approaches are typically categorized based on data heterogeneity into congruent \(shared modalities and feature spaces\) and incongruent settings, the latter including vertical, transfer, and hybrid scenarios with partially overlapping or distinct modalities\. In general, MFL optimizes a global objective by aggregating client\-specific multimodal losses, often defined as weighted combinations of modality\-wise contributions\. However, most methods assume a shared latent space and rely on centralized or weakly decentralized protocols\. Recent graph\-based approaches such as CoMFed\(Badi et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib2)\)address heterogeneity by aligning compressed latent prototypes across clients, but still require a predefined common embedding space and enforce similarity via projections\. In contrast, our sheaf\-based framework does not assume a global latent space; instead, it enables decentralized learning in which each client maintains its own representation, and cross\-client consistency is enforced through edge\-dependent geometric alignments without information loss, naturally accommodating multimodal and heterogeneous settings\.

Federated Representation Learning\.Our problem formulation is closely related Federated Representation Learning \(FRL\), which aims to learn transferable representations across clients\. In the supervised setting, existing approaches either enforce consistency via auxiliary losses on latent embeddings\(Li et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib40)\), typically assuming a shared global model, or adopt a split architecture with a common encoder and client\-specific heads\(Collins et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib11); McLaughlin & Su,[2024](https://arxiv.org/html/2608.10016#bib.bib44)\), still requiring a shared latent dimensionality\. More recent methods based on prototype alignment\(Tan et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib62); Tran et al\.,[2024](https://arxiv.org/html/2608.10016#bib.bib64)\)partially address heterogeneity but continue to rely on a common embedding space and often centralized training\. In contrast, our framework allows each client to learn its own latent space and enforces consistency in a decentralized manner through structured transport maps, so that classical FRL methods can be interpreted as a special case in which all latent spaces share the same dimensionality and are constrained to coincide via identity alignments\. Beyond supervised settings, unsupervised and self\-supervised approaches operating on unlabeled data aim to improve representation transferability across clients, often leveraging contrastive learning frameworks\(Zhuang et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib70); Han et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib29); Zhang et al\.,[2023](https://arxiv.org/html/2608.10016#bib.bib66); Miao & Koyuncu,[2024](https://arxiv.org/html/2608.10016#bib.bib47); Ghalkha et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib26)\)\. Our proposed framework is not limited to any specific local loss or model architecture; thus, it is also amenable to the latter settings\.

Representation alignment, synchronization, and manifold methods\.Alignment up to orthogonal transformations arises naturally when representations are identifiable only up to rotations or reflections, with the orthogonal Procrustes problem providing a classical closed\-form solution\(Schönemann,[1966](https://arxiv.org/html/2608.10016#bib.bib54)\)\. We leverage this solution in a federated and decentralized setting, both under homogeneous latent spaces \(cf\. App\.[D](https://arxiv.org/html/2608.10016#A4)\) and under stringent privacy constraints \(cf\. App\.[F](https://arxiv.org/html/2608.10016#A6)\)\. Related ideas arise in group synchronization and connection\-Laplacian methods\(Singer,[2011](https://arxiv.org/html/2608.10016#bib.bib57); Singer & Wu,[2012](https://arxiv.org/html/2608.10016#bib.bib58); Bandeira et al\.,[2013](https://arxiv.org/html/2608.10016#bib.bib3); Thunberg et al\.,[2017](https://arxiv.org/html/2608.10016#bib.bib63)\), which enforce consistency across pairwise transformations; unlike these approaches, where node signals and edge relations are given or observed, possibly with noise, we jointly learn latent representations at the nodes and alignment maps on the edges\. To accommodate heterogeneous latent dimensions, we further consider semi\-orthogonal maps on the Stiefel manifold, enabling isometric embeddings across representation spaces\. This connects our formulation to optimization on matrix manifolds\(Absil et al\.,[2008](https://arxiv.org/html/2608.10016#bib.bib1); Edelman et al\.,[1998](https://arxiv.org/html/2608.10016#bib.bib20)\)and federated optimization under manifold constraints\(Li & Ma,[2022](https://arxiv.org/html/2608.10016#bib.bib39)\); however, these works focus on optimizing model parameters subject to geometric constraints, rather than learning alignment maps as a component of representation learning\.

Positioning of this paper\.Overall, our framework generalizes across all paradigms discussed above, which differ mainly in*what*is assumed to be shared among agents—a single global model, a coupled set of parameters, a common multimodal embedding, or a shared latent space—rather than in the underlying goal of exploiting relatedness across agents\. Indeed, methods introduced within one paradigm, such as prototype alignment in FRL or graph\-based regularization in FMTL, are often applicable to the others as well\. Rather than enforcing any of these forms of sharing, we learn edge\-wise geometric transport maps that relate neighboring latent spaces and promote their semantic compatibility\. In the homogeneous setting, our formulation recovers the single global model of classical FL as a limiting case, together with orthogonal alignment, synchronization, and connection\-Laplacian models\. More generally, it extends graph\-based FMTL and Laplacian regularization through a*learnable sheaf Laplacian*acting on latent representations rather than model parameters\. Additionally, it relaxes the shared, predefined embedding space assumed by MFL approaches such as CoMFed, and recovers FRL methods as the special case in which all latent spaces coincide via identity alignments\. Finally, pilot\-restricted gluing enables communication\-efficient alignment, while manifold\-constrained transport maps admit decentralized optimization through local model updates and closed\-form edge\-wise Procrustes steps\.

## 3Our Sheaf\-based Representation Learning Framework

![Refer to caption](https://arxiv.org/html/2608.10016v1/x1.png)Figure 1:An illustration of our sheaf\-based representation learning framework\. In this example, we consider a supervised setting and equip local agents with task\-specific heads\.[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1)illustrates the proposed sheaf\-based representation learning framework\. Let𝒢=\(𝒩,ℰ\)\\mathcal\{G\}=\(\\mathcal\{N\},\\mathcal\{E\}\)denote a finite undirected graph, where\|𝒩\|=N\|\\mathcal\{N\}\|=Nis the number of nodes\. Each node corresponds to an agent, while edges encode structural relationships among agents\. For example, an edge\(i,j\)∈ℰ\(i,j\)\\in\\mathcal\{E\}may capture similarities in tasks, data distributions, modalities, or model architectures, thereby enabling beneficial cooperation between agentsiiandjjthrough a physical communication channel\.

Local agents\.Each agenti∈𝒩i\\in\\mathcal\{N\}observes a local dataset𝒯i\\mathcal\{T\}\_\{i\}comprisingMiM\_\{i\}samples inℝpi\\mathbb\{R\}^\{p\_\{i\}\}, drawn from an underlying distributionPiP\_\{i\}, as illustrated in[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1)\. We denote thenn\-th sample by𝐱in\\mathbf\{x\}\_\{i\}^\{n\}, withn∈\[Mi\]n\\in\[M\_\{i\}\]\. Depending on the local learning paradigm,𝒯i\\mathcal\{T\}\_\{i\}may additionally include labels, side information, paired observations, or locally generated views of the samples\. Each agent learns a local representation model comprising a neural encoder

f𝝋i:ℝpi→ℱ​\(i\),f\_\{\\bm\{\\varphi\}\_\{i\}\}:\\mathbb\{R\}^\{p\_\{i\}\}\\to\\mathcal\{F\}\(i\)\\,,\(1\)where𝝋i\\bm\{\\varphi\}\_\{i\}is the set of learnable parameters, andℱ​\(i\)≅ℝdi\\mathcal\{F\}\(i\)\\cong\\mathbb\{R\}^\{d\_\{i\}\}is the latent representation space of agentii\. When required by the local task, the encoder is followed by a personalized head parametrized by𝝍i\\bm\{\\psi\}\_\{i\}, i\.e\.,

h𝝍i:ℱ​\(i\)→𝒴i,h\_\{\\bm\{\\psi\}\_\{i\}\}:\\mathcal\{F\}\(i\)\\to\\mathcal\{Y\}\_\{i\}\\,,\(2\)where𝒴i\\mathcal\{Y\}\_\{i\}is an agent\-dependent output space\. We let𝜽i\\bm\{\\theta\}\_\{i\}denote the collection of local trainable parameters, including𝝋i\\bm\{\\varphi\}\_\{i\}and, when present,𝝍i\\bm\{\\psi\}\_\{i\}\. Each agenti∈𝒩i\\in\\mathcal\{N\}aims at minimizing an arbitrary local learning objectiveℒi​\(𝜽i\)\\mathcal\{L\}\_\{i\}\(\\bm\{\\theta\}\_\{i\}\)based on the dataset𝒯i\\mathcal\{T\}\_\{i\}\. For example, in the supervised setting,

ℒi​\(𝜽i\)=1Mi​∑n=1Miℓi​\(h𝝍i∘f𝝋i​\(𝐱in\)\)\.\\mathcal\{L\}\_\{i\}\(\\bm\{\\theta\}\_\{i\}\)=\\frac\{1\}\{M\_\{i\}\}\\sum\_\{n=1\}^\{M\_\{i\}\}\\ell\_\{i\}\\\!\\left\(h\_\{\\bm\{\\psi\}\_\{i\}\}\\circ f\_\{\\bm\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{n\}\)\\right\)\\,\.\(3\)Other choices naturally accommodate semi\-supervised, unsupervised, or self\-supervised learning\.

Communication over the learnable latent network sheaf\.To formalize the relationships, and thus the communication, among the latent representation spacesℱ​\(i\)\\mathcal\{F\}\(i\)introduced above, we adopt the framework of network sheaves\(Bredon,[1997](https://arxiv.org/html/2608.10016#bib.bib8); Curry,[2014](https://arxiv.org/html/2608.10016#bib.bib12)\)\. Intuitively, thoseℱ​\(i\)\\mathcal\{F\}\(i\)spaces will be structured so as to encode consistency relations between the latent representations of neighboring agents\. Specifically, a*network sheaf*ℱ\\mathcal\{F\}on𝒢\\mathcal\{G\}, valued in the category of finite\-dimensional real vector spaces and linear maps, consists of the following assignments\. To each nodei∈𝒩i\\in\\mathcal\{N\}, it associates adid\_\{i\}\-dimensional real vector spaceℱ​\(i\)\\mathcal\{F\}\(i\), being the learnable latent representation space of agentiiand referred to asnode stalk\. In this work we consider the general case where node stalksℱ​\(i\)\\mathcal\{F\}\(i\),i∈𝒩i\\in\\mathcal\{N\}, have different dimensionalitydid\_\{i\}\. For instance, in[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1), the vector spaceℱ​\(i\)\\mathcal\{F\}\(i\)is three\-dimensional, while the vector spaceℱ​\(j\)\\mathcal\{F\}\(j\)is two\-dimensional\. For each node stalk, the valuation corresponds to a latent embedding, e\.g\., those vectors𝐳i=f𝝋i​\(𝐱i\)\\mathbf\{z\}\_\{i\}=f\_\{\\bm\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}\)and𝐳j=f𝝋j​\(𝐱j\)\\mathbf\{z\}\_\{j\}=f\_\{\\bm\{\\varphi\}\_\{j\}\}\(\\mathbf\{x\}\_\{j\}\)in[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1)\. The collection of latent representations across the network, viz\.𝐳=\{𝐳i\}i∈𝒩\\mathbf\{z\}=\\\{\\mathbf\{z\}\_\{i\}\\\}\_\{i\\in\\mathcal\{N\}\}, forms a 0\-cochain\. Accordingly, the0\-cochain space isC0​\(𝒢,ℱ\)=⨁i∈𝒩ℱ​\(i\)C^\{0\}\(\\mathcal\{G\},\\mathcal\{F\}\)=\\bigoplus\_\{i\\in\\mathcal\{N\}\}\\mathcal\{F\}\(i\)\. Similarly, to each edgeei​j=\(i,j\)∈ℰe\_\{ij\}=\(i,j\)\\in\\mathcal\{E\}, the network sheafℱ\\mathcal\{F\}assigns adi​jd\_\{ij\}\-dimensional real vector spaceℱ​\(ei​j\)\\mathcal\{F\}\(e\_\{ij\}\), representing a shared space encoding pairwise compatibility between agentsiiandjj\. The latter space is callededge stalk\. The11\-cochain is a collection𝐲=\{𝐲ei​j\}ei​j∈ℰ\\mathbf\{y\}=\\\{\\mathbf\{y\}\_\{e\_\{ij\}\}\\\}\_\{e\_\{ij\}\\in\\mathcal\{E\}\}with𝐲ei​j∈ℱ​\(ei​j\)\\mathbf\{y\}\_\{e\_\{ij\}\}\\in\\mathcal\{F\}\(e\_\{ij\}\), and the corresponding11\-cochain space isC1​\(𝒢,ℱ\)=⨁ei​j∈ℰℱ​\(ei​j\)C^\{1\}\(\\mathcal\{G\},\\mathcal\{F\}\)=\\bigoplus\_\{e\_\{ij\}\\in\\mathcal\{E\}\}\\mathcal\{F\}\(e\_\{ij\}\)\. Additionally, for each incidence relationi→ei​ji\\to e\_\{ij\}, whereiiis a vertex ofei​je\_\{ij\},ℱ\\mathcal\{F\}specifies a linearrestriction map

ℱi→ei​j:ℱ​\(i\)→ℱ​\(ei​j\),\\mathcal\{F\}\_\{i\\to e\_\{ij\}\}:\\mathcal\{F\}\(i\)\\to\\mathcal\{F\}\(e\_\{ij\}\)\\,,which relates the local latent representation at nodeiito the corresponding edge space and enables the communication among latent spaces\. In this work, we focus on learnable restriction maps suitable for dealing with the well\-known geometric\-invariance of latent representations\(Moschella et al\.,[2023](https://arxiv.org/html/2608.10016#bib.bib48)\), and that preserve distances and angles\. These maps are orthogonal transformations and Stiefel matrices, where the orthogonal and Stiefel manifolds are

O​\(d\)≔\{𝐎∈ℝd×d∣𝐎⊤=𝐎−1\},andSt​\(d,k\)≔\{𝐕∈ℝd×k∣𝐕⊤​𝐕=𝐈k,k<d\}\.\\mathrm\{O\}\(\{d\}\)\\coloneqq\\\{\\mathbf\{O\}\\\!\\in\\\!\\mathbb\{R\}^\{d\\times d\}\\,\\mid\\,\\mathbf\{O\}^\{\\top\}\\\!=\\\!\\mathbf\{O\}^\{\-1\}\\\}\\,,\\quad\\text\{and\}\\quad\\mathrm\{St\}\(\{d\},\{k\}\)\\coloneqq\\\{\\mathbf\{V\}\\\!\\in\\\!\\mathbb\{R\}^\{d\\times k\}\\,\\mid\\,\\mathbf\{V\}^\{\\top\}\\mathbf\{V\}\\\!=\\\!\\mathbf\{I\}\_\{k\},\\,k<d\\\}\\,\.\(4\)As we consider structured restriction maps that preserve the information but do not increase latent perturbations, for each edgeei​j=\(i,j\)∈ℰe\_\{ij\}=\(i,j\)\\in\\mathcal\{E\}withdi\>djd\_\{i\}\>d\_\{j\}, we setℱ​\(ei​j\)=ℱ​\(i\)≅ℝdi\\mathcal\{F\}\(e\_\{ij\}\)=\\mathcal\{F\}\(i\)\\cong\\mathbb\{R\}^\{d\_\{i\}\}, as illustrated in[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1)\. This impliesℱi→ei​j=𝐎j​i∈O​\(di\)\\mathcal\{F\}\_\{i\\to e\_\{ij\}\}\\\!=\\\!\\mathbf\{O\}\_\{ji\}\\in\\mathrm\{O\}\(\{d\_\{i\}\}\)andℱj→ei​j=𝐕i​j∈St​\(di,dj\)\\mathcal\{F\}\_\{j\\to e\_\{ij\}\}\\\!=\\\!\\mathbf\{V\}\_\{ij\}\\\!\\in\\\!\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)\. Further discussion on the benefits of choosing the edge stalk dimensionalitydi​j=max⁡\(di,dj\)d\_\{ij\}=\\max\(d\_\{i\},d\_\{j\}\), as opposed to compressing, is provided in App\.[A](https://arxiv.org/html/2608.10016#A1)\.

Heterogeneous latent space geometry\.In this setting, the encodersf𝝋if\_\{\\bm\{\\varphi\}\_\{i\}\}are trained independently, on different data, and possibly with different architectures\. Then, it is reasonable to expect the induced latent spacesℱ​\(i\)\\mathcal\{F\}\(i\)to exhibit heterogeneous geometry, in the sense of different scales and anisotropies\(Pandolfo et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib51); Svendsen et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib61)\)\. This may occur even when dimensionalitiesdid\_\{i\}coincide\. However, the above geometric restriction maps are effective for measuring inconsistency only when the latent representations𝐳i\\mathbf\{z\}\_\{i\}and𝐳j\\mathbf\{z\}\_\{j\}adhere to a common geometric reference on the edge stalk\. Thus, to make representations from distinct node stalks comparable once transported, each encoder is assumed to include a final whitening layer\(Zhang et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib67)\)that approximately normalizes the latent representations to zero mean and identity covariance\. This normalization reduces the node\-specific geometry to a common Euclidean reference\. However, the geometry ofℱ​\(i\)\\mathcal\{F\}\(i\)is key to minimizing the local learning objectiveℒi​\(𝜽i\)\\mathcal\{L\}\_\{i\}\(\\bm\{\\theta\}\_\{i\}\)\. Hence, to let the local head exploit the local geometry, we equip the latter with an initial coloring layer that is assumed to perform the left inverse map of the normalization applied by the whitening layer of the encoder\.

Centralized formulation under global alignment\.The sheaf induces a coboundary operator𝚫:C0​\(𝒢;ℱ\)→C1​\(𝒢;ℱ\)\\mathbf\{\\Delta\}:C^\{0\}\(\\mathcal\{G\};\\mathcal\{F\}\)\\to C^\{1\}\(\\mathcal\{G\};\\mathcal\{F\}\)defined edgewise as

\(𝚫​𝐳\)ei​j=𝐎j​i​𝐳i−𝐕i​j​𝐳j,ei​j=\(i,j\)∈ℰ\.\(\\mathbf\{\\Delta\}\\mathbf\{z\}\)\_\{e\_\{ij\}\}=\\mathbf\{O\}\_\{ji\}\\mathbf\{z\}\_\{i\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\},\\qquad e\_\{ij\}=\(i,j\)\\in\\mathcal\{E\}\.\(5\)EquippingC0C^\{0\}andC1C^\{1\}with canonical Euclidean inner products⟨𝐳,𝐳′⟩=∑i∈𝒩𝐳i⊤​𝐳i′\\langle\\mathbf\{z\},\\mathbf\{z\}^\{\\prime\}\\rangle=\\sum\_\{i\\in\\mathcal\{N\}\}\\mathbf\{z\}\_\{i\}^\{\\top\}\\mathbf\{z\}\_\{i\}^\{\\prime\}and⟨𝐲,𝐲′⟩=∑ei​j∈ℰ𝐲ei​j⊤​𝐲ei​j′\\langle\\mathbf\{y\},\\mathbf\{y\}^\{\\prime\}\\rangle=\\sum\_\{e\_\{ij\}\\in\\mathcal\{E\}\}\\mathbf\{y\}\_\{e\_\{ij\}\}^\{\\top\}\\mathbf\{y\}\_\{e\_\{ij\}\}^\{\\prime\}, the adjoint𝚫⊤\\mathbf\{\\Delta\}^\{\\top\}is well defined\. The associated network sheaf Laplacian is

𝐋ℱ=𝚫⊤​𝚫,\\mathbf\{L\}\_\{\\mathcal\{F\}\}=\\mathbf\{\\Delta\}^\{\\top\}\\mathbf\{\\Delta\}\\,,\(6\)which is symmetric positive semidefinite\. A valuation𝐳⋆∈C0​\(𝒢;ℱ\)\\mathbf\{z\}^\{\\star\}\\in C^\{0\}\(\\mathcal\{G\};\\mathcal\{F\}\)belonging toker⁡\(𝐋ℱ\)\\ker\(\\mathbf\{L\}\_\{\\mathcal\{F\}\}\)is said*global section*and it can be thought of as an assignment of latent representations to the node stalks that does not break local rules\. Indeed, from[Eq\.6](https://arxiv.org/html/2608.10016#S3.E6)it is clear that𝚫​𝐳⋆=𝟎\\bm\{\\Delta\}\\mathbf\{z\}^\{\\star\}=\\mathbf\{0\}which implies

𝐎j​i​𝐳i⋆=𝐕i​j​𝐳j⋆for each​ei​j∈ℰ;\\mathbf\{O\}\_\{ji\}\\mathbf\{z\}\_\{i\}^\{\\star\}=\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}^\{\\star\}\\,\\quad\\text\{for each \}e\_\{ij\}\\in\\mathcal\{E\}\\,;\(7\)where𝐳i⋆∈ℝdi\\mathbf\{z\}\_\{i\}^\{\\star\}\\in\\mathbb\{R\}^\{d\_\{i\}\}and𝐳j⋆∈ℝdj\\mathbf\{z\}\_\{j\}^\{\\star\}\\in\\mathbb\{R\}^\{d\_\{j\}\}are the components of𝐳⋆\\mathbf\{z\}^\{\\star\}at nodeiiandjj, respectively\. Thus, a plausible strategy would be to constrain the learned latent representations across the network to be globally consistent with respect to the learned network sheafℱ\\mathcal\{F\}\. For simplicity, assume the existence of a set of shared sample indices across agents, and letM≤mini∈𝒩⁡MiM\\leq\\min\_\{i\\in\\mathcal\{N\}\}M\_\{i\}be the number of such comparable samples\. Then, according to the previous rationale we can pose the following centralized problem formulation

min\{𝜽i=\(𝝋i,𝝍i\)\}i∈𝒩\{𝐎j​i∈O​\(di\),𝐕i​j∈St​\(di,dj\)\}ei​j∈ℰ\\displaystyle\\min\_\{\\begin\{subarray\}\{c\}\\\{\\bm\{\\theta\}\_\{i\}=\(\\bm\{\\varphi\}\_\{i\},\\bm\{\\psi\}\_\{i\}\)\\\}\_\{i\\in\\mathcal\{N\}\}\\\\ \\\{\\mathbf\{O\}\_\{ji\}\\in\\mathrm\{O\}\(\{d\_\{i\}\}\),\\,\\mathbf\{V\}\_\{ij\}\\in\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)\\\}\_\{e\_\{ij\}\\in\\mathcal\{E\}\}\\end\{subarray\}\}∑i∈𝒩ℒi​\(𝜽i\)\\displaystyle\\sum\_\{i\\in\\mathcal\{N\}\}\\mathcal\{L\}\_\{i\}\(\\bm\{\\theta\}\_\{i\}\)\(8\)s\.t\.𝐳n=\{f𝝋i​\(𝐱in\)\}i∈𝒩∈ker⁡\(𝐋ℱ\),∀n∈\[M\];\\displaystyle\\mathbf\{z\}^\{n\}=\\\{f\_\{\\bm\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{n\}\)\\\}\_\{i\\in\\mathcal\{N\}\}\\in\\ker\(\\mathbf\{L\}\_\{\\mathcal\{F\}\}\),\\qquad\\forall n\\in\[M\]\\,;where the network sheafℱ\\mathcal\{F\}is learned through its restriction maps under a known communication topology\. However, the constraints in Prob\. \([8](https://arxiv.org/html/2608.10016#S3.E8)\) might be too restrictive in practice\. Enforcing exact compatibility across heterogeneous agents is generally unrealistic, and, moreover, satisfying the constraint for all samples would in principle require a centralized solution, as it couples all node representations simultaneously\. This limitation motivates the need for a more flexible and decentralized formulation that can accommodate inconsistencies while still promoting coherence across agents\.

## 4Towards a Scalable and Decentralized Formulation

The discussed limitations of Prob\. \([8](https://arxiv.org/html/2608.10016#S3.E8)\) in realistic settings motivate the introduction of a relaxed, distributed formulation\. To obtain a tractable formulation suitable for distributed implementation, we replace the hard constraints in Prob\. \([8](https://arxiv.org/html/2608.10016#S3.E8)\) with a soft penalty based on the sheaf total variation, which promotes geometric consistency across agents without enforcing exact agreement:

𝒯​𝒱​\(𝐳\)≔‖𝚫​𝐳‖22=𝐳⊤​𝐋ℱ​𝐳=∑ei​j∈ℰ‖𝐎j​i​𝐳i−𝐕i​j​𝐳j‖22\.\\mathcal\{TV\}\(\\mathbf\{z\}\)\\coloneqq\\norm\{\\bm\{\\Delta\}\\mathbf\{z\}\}\_\{2\}^\{2\}=\\mathbf\{z\}^\{\\top\}\\mathbf\{L\}\_\{\\mathcal\{F\}\}\\mathbf\{z\}=\\sum\_\{e\_\{ij\}\\in\\mathcal\{E\}\}\\norm\{\\mathbf\{O\}\_\{ji\}\\mathbf\{z\}\_\{i\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}\}\_\{2\}^\{2\}\\,\.\(9\)From[Eq\.7](https://arxiv.org/html/2608.10016#S3.E7), it is easy to see that constraining the learned𝐳\\mathbf\{z\}to be a global section is equivalent to zeroing[Eq\.9](https://arxiv.org/html/2608.10016#S4.E9)\. Additional discussion on the geometric consistency promoted by the action of the restriction maps, as well as the use of the total variation penalty, is provided in App\.[A](https://arxiv.org/html/2608.10016#A1)\.

Semantic embedding and orientation\.We can recast the local terms in[Eq\.9](https://arxiv.org/html/2608.10016#S4.E9)as

‖𝐎j​i​𝐳i−𝐕i​j​𝐳j‖22=\(a\)‖𝐳i−𝐕i​j​𝐳j‖22;\\norm\{\\mathbf\{O\}\_\{ji\}\\mathbf\{z\}\_\{i\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}\}\_\{2\}^\{2\}\\stackrel\{\{\\scriptstyle\(a\)\}\}\{\{=\}\}\\norm\{\\mathbf\{z\}\_\{i\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}\}\_\{2\}^\{2\}\\,;\(10\)where in\(a\)\(a\)we exploit the transitive action of the orthogonal group on the Stiefel manifold, thus re\-parameterizing accordingly𝐕i​j=𝐎j​i⊤​𝐕i​j\\mathbf\{V\}\_\{ij\}=\\mathbf\{O\}\_\{ji\}^\{\\top\}\\mathbf\{V\}\_\{ij\}\.[Eq\.10](https://arxiv.org/html/2608.10016#S4.E10)highlights that perfect alignment, and thus null total variation, corresponds to the case in which the higher\-dimensional latent space is an isometric embedding of a lower\-dimensional latent space\. While in general this condition does not always hold \(see the discussion in App\.[A](https://arxiv.org/html/2608.10016#A1)on geometric consistency\), the restriction maps in our work adheres to the*semantic embedding principle*\(D’Acunto et al\.,[2025](https://arxiv.org/html/2608.10016#bib.bib13);[2026](https://arxiv.org/html/2608.10016#bib.bib14)\): informally, from a probabilistic perspective, the semantics of the lower\-dimensional node stalkℱ​\(j\)\\mathcal\{F\}\(j\)is preserved when embedded into the higher\-dimensionalℱ​\(i\)\\mathcal\{F\}\(i\)\.

Another distinguishing feature of our framework is that the restriction map reparameterization in[Eq\.10](https://arxiv.org/html/2608.10016#S4.E10)induces a natural orientation for𝒢\\mathcal\{G\}which we exploit in the sequel to reduce the communication cost\. Specifically, as exemplified in[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1), we endow𝒢\\mathcal\{G\}with the*embedding orientation*, that is, from nodes with lower\-dimensional latent spaces to those with higher\-dimensional ones, accordingly to the direction in which semantics is preserved\. Further details on the special case of latent representation spaces with homogeneous dimensionality are provided in App\.[A](https://arxiv.org/html/2608.10016#A1)\.

Incoming and outgoing embedding maps\.By using[Eq\.10](https://arxiv.org/html/2608.10016#S4.E10), the total variation in[Eq\.9](https://arxiv.org/html/2608.10016#S4.E9)can be rewritten as

𝒯​𝒱​\(𝐳\)=∑ei​j∈ℰ‖𝐳i−𝐕i​j​𝐳j‖22;\\mathcal\{TV\}\(\\mathbf\{z\}\)=\\sum\_\{e\_\{ij\}\\in\\mathcal\{E\}\}\\norm\{\\mathbf\{z\}\_\{i\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}\}\_\{2\}^\{2\}\\,;\(11\)where for the edgeei​j=\(i,j\)e\_\{ij\}=\(i,j\), nodeiiis the head andjjthe tail, in accordance with the relationdi\>djd\_\{i\}\>d\_\{j\}\. Ties occurring whendi=dj=dd\_\{i\}=d\_\{j\}=dare instead broken arbitrarily: since in this case𝐕i​j=𝐕j​i⊤∈O​\(d\)\\mathbf\{V\}\_\{ij\}=\\mathbf\{V\}\_\{ji\}^\{\\top\}\\in\\mathrm\{O\}\(\{d\}\), the edge contribution in[Eq\.11](https://arxiv.org/html/2608.10016#S4.E11)is invariant to the chosen orientation\. Throughout, we therefore adopt the convention that forei​je\_\{ij\}, the first indexiiis always the head\.

Now, exploiting the orientation fixed above, for eachi∈𝒩i\\in\\mathcal\{N\}, we can distinguish between two sets of neighbors\. Specifically,𝒩​\(i\)−\\mathcal\{N\}\(i\)^\{\-\}collects the neighborsjjfor whichiiis the head, and𝒩​\(i\)\+\\mathcal\{N\}\(i\)^\{\+\}those for whichiiis the tail, so thatdj≤did\_\{j\}\\leq d\_\{i\}on𝒩​\(i\)−\\mathcal\{N\}\(i\)^\{\-\}anddj≥did\_\{j\}\\geq d\_\{i\}on𝒩​\(i\)\+\\mathcal\{N\}\(i\)^\{\+\}, with equality holding only on tied edges\. The overall neighborhood is thus𝒩​\(i\)=𝒩​\(i\)−∪𝒩​\(i\)\+\\mathcal\{N\}\(i\)=\\mathcal\{N\}\(i\)^\{\-\}\\cup\\mathcal\{N\}\(i\)^\{\+\}\. Taking advantage of this distinction, starting from[Eq\.11](https://arxiv.org/html/2608.10016#S4.E11), we have

𝒯​𝒱​\(𝐳\)\\displaystyle\\mathcal\{TV\}\(\\mathbf\{z\}\)=∑i∈𝒩∑j∈𝒩​\(i\)−‖𝐳i−𝐕i​j​𝐳j‖22\\displaystyle=\\sum\_\{i\\in\\mathcal\{N\}\}\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\}\}\\norm\{\\mathbf\{z\}\_\{i\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}\}\_\{2\}^\{2\}\(12\)=12​∑i∈𝒩\[∑j∈𝒩​\(i\)−‖𝐳i−𝐕i​j​𝐳j‖22\+∑j∈𝒩​\(i\)\+‖𝐳j−𝐕j​i​𝐳i‖22\];\\displaystyle=\\dfrac\{1\}\{2\}\\sum\_\{i\\in\\mathcal\{N\}\}\\left\[\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\}\}\\norm\{\\mathbf\{z\}\_\{i\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}\}\_\{2\}^\{2\}\+\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\+\}\}\\norm\{\\mathbf\{z\}\_\{j\}\-\\mathbf\{V\}\_\{ji\}\\mathbf\{z\}\_\{i\}\}\_\{2\}^\{2\}\\right\]\\,;which highlights that the contribution of nodeiito the total variation is made of two terms with different geometric meanings:

𝒯​𝒱​\(𝐳\)\|i=12​∑j∈𝒩​\(i\)−‖𝐳i−𝐕i​j​𝐳j‖22⏟Incoming Embedding\+12​∑j∈𝒩​\(i\)\+‖𝐳j−𝐕j​i​𝐳i‖22⏟Outgoing Embedding\.\\mathcal\{TV\}\(\\mathbf\{z\}\)\|\_\{i\}=\\dfrac\{1\}\{2\}\\underbrace\{\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\}\}\\norm\{\\mathbf\{z\}\_\{i\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}\}\_\{2\}^\{2\}\}\_\{\\text\{Incoming Embedding\}\}\+\\dfrac\{1\}\{2\}\\underbrace\{\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\+\}\}\\norm\{\\mathbf\{z\}\_\{j\}\-\\mathbf\{V\}\_\{ji\}\\mathbf\{z\}\_\{i\}\}\_\{2\}^\{2\}\}\_\{\\text\{Outgoing Embedding\}\}\\,\.\(13\)Leveraging[Eq\.13](https://arxiv.org/html/2608.10016#S4.E13), for each nodeii, the matrices𝐕i​j\\mathbf\{V\}\_\{ij\}in the first term are referred to as the*incoming*embedding maps, while the matrices𝐕j​i\\mathbf\{V\}\_\{ji\}in the second term are referred to as the*outgoing*embedding maps\. As detailed in the sequel, the expression in[Eq\.13](https://arxiv.org/html/2608.10016#S4.E13)is useful for computing local updates and reducing the communication cost in our proposed decentralized algorithm\.

Scalable gluing penalty and decentralized problem\.For scalability and computational aspects, to enforce geometric consistency via the sheaf total variation, we only use a subset of reference samples chosen among those𝐳in\\mathbf\{z\}\_\{i\}^\{n\}\(pilots\), for eachi∈𝒩i\\in\\mathcal\{N\}\. Let

𝒜⊂\{1,…,Mmin\},\|𝒜\|=K≪Mmin,\\mathcal\{A\}\\subset\\\{1,\\dots,M\_\{\\mathrm\{min\}\}\\\},\\qquad\|\\mathcal\{A\}\|=K\\ll M\_\{\\mathrm\{min\}\},denote fixed pilot indices, whereMmin=mini⁡MiM\_\{\\mathrm\{min\}\}=\\min\_\{i\}M\_\{i\}\. For the sake of exposition, we assume the datasets\{𝒯i\}\\\{\\mathcal\{T\}\_\{i\}\\\}are aligned\. For each agentii, define the pilot feature matrix

𝐀i​\(𝝋i\)=\[f𝝋i​\(𝐱ik\)\]k∈𝒜=\[𝐳ik\]k∈𝒜∈ℝdi×K\.\\mathbf\{A\}\_\{i\}\(\\boldsymbol\{\\varphi\}\_\{i\}\)=\\Big\[f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\\Big\]\_\{k\\in\\mathcal\{A\}\}=\\left\[\\mathbf\{z\}\_\{i\}^\{k\}\\right\]\_\{k\\in\\mathcal\{A\}\}\\in\\mathbb\{R\}^\{d\_\{i\}\\times K\}\.\(14\)Different strategies can be used to select the pilot set𝒜\\mathcal\{A\}, trading off semantic reliability, geometric coverage, and communication efficiency; we discuss several of such strategies in App\.[B](https://arxiv.org/html/2608.10016#A2)\. The usage of pilots also highlights that we enforce a weaker notion of alignment among latent spaces\. In other words, while a global agreement may occur on these pilots, the local latent representations for the remaining samples may not, in principle, be perfectly alignable via the learned restriction maps\. This is consistent with the discussion above, particularly with the fact that we neither assume nor aim for a single global latent representation\.

Now, exploiting[Eq\.12](https://arxiv.org/html/2608.10016#S4.E12)and introducing a penalty hyperparamerλ\>0\\lambda\>0, our considered gluing penalty decouples as

ℛ𝒜​\(\{𝝋i\},\{𝐕i​j\},\{𝐕j​i\}\)=∑i=1Nℛ𝒜\|i​\(𝝋i,\{𝐕i​j\},\{𝐕j​i\}\);\\mathcal\{R\}\_\{\\mathcal\{A\}\}\(\\\{\\bm\{\\varphi\}\_\{i\}\\\},\\\{\\mathbf\{V\}\_\{ij\}\\\},\\\{\\mathbf\{V\}\_\{ji\}\\\}\)=\\sum\_\{i=1\}^\{N\}\\mathcal\{R\}\_\{\\mathcal\{A\}\}\|\_\{i\}\(\\bm\{\\varphi\}\_\{i\},\\\{\\mathbf\{V\}\_\{ij\}\\\},\\\{\\mathbf\{V\}\_\{ji\}\\\}\)\\,;\(15\)where

ℛ𝒜\|i​\(𝝋i,\{𝐕i​j\},\{𝐕j​i\}\)\\displaystyle\\mathcal\{R\}\_\{\\mathcal\{A\}\}\|\_\{i\}\(\\bm\{\\varphi\}\_\{i\},\\\{\\mathbf\{V\}\_\{ij\}\\\},\\\{\\mathbf\{V\}\_\{ji\}\\\}\)=λi2​K\[∑j∈𝒩​\(i\)−‖𝐀i​\(𝝋i\)−𝐕i​j​𝐀j​\(𝝋j\)‖F2\\displaystyle=\\dfrac\{\\lambda\_\{i\}\}\{2K\}\\Bigg\[\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\}\}\\norm\{\\mathbf\{A\}\_\{i\}\(\\bm\{\\varphi\}\_\{i\}\)\-\\mathbf\{V\}\_\{ij\}\\mathbf\{A\}\_\{j\}\(\\bm\{\\varphi\}\_\{j\}\)\}\_\{\\mathrm\{F\}\}^\{2\}\(16\)\+∑j∈𝒩​\(i\)\+‖𝐀j​\(𝝋j\)−𝐕j​i​𝐀i​\(𝝋i\)‖F2\];\\displaystyle\\quad\+\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\+\}\}\\norm\{\\mathbf\{A\}\_\{j\}\(\\bm\{\\varphi\}\_\{j\}\)\-\\mathbf\{V\}\_\{ji\}\\mathbf\{A\}\_\{i\}\(\\bm\{\\varphi\}\_\{i\}\)\}\_\{\\mathrm\{F\}\}^\{2\}\\Bigg\]\\,;withλi=λ/di\\lambda\_\{i\}=\\lambda/d\_\{i\}andλ\>0\\lambda\>0\. At this point, we are ready to pose the sheaf\-based federated representation learning problem\.

###### Problem 1\(Sheaf\-based Federated Representation Learning\)\.

Given \(i\) an undirected graph𝒢=\(𝒩,ℰ\)\\mathcal\{G\}=\(\\mathcal\{N\},\\mathcal\{E\}\)endowed with an orientation induced by semantic embeddings, \(ii\) a collection ofNNlocal training datasets\{𝒯i\}i∈𝒩\\\{\\mathcal\{T\}\_\{i\}\\\}\_\{i\\in\\mathcal\{N\}\}drawn from unknown distributionsPiP\_\{i\}, and \(iii\) a set of reference pilot indices𝒜\\mathcal\{A\}, the goal of sheaf\-based federated representation learning \(SFRL\) is to learn, for each nodeii, a local model𝛉i=\(𝛗i,𝛙i\)\\bm\{\\theta\}\_\{i\}=\(\\bm\{\\varphi\}\_\{i\},\\bm\{\\psi\}\_\{i\}\)together with incoming and outgoing embeddings\{𝐕i​j∈St​\(di,dj\)\}\(i,j\)∈ℰ\\\{\\mathbf\{V\}\_\{ij\}\\in\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)\\\}\_\{\(i,j\)\\in\\mathcal\{E\}\}and\{𝐕j​i∈St​\(dj,di\)\}\(j,i\)∈ℰ\\\{\\mathbf\{V\}\_\{ji\}\\in\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\)\\\}\_\{\(j,i\)\\in\\mathcal\{E\}\}, respectively\. These embedding matrices align the local latent spaceℱ​\(i\)\\mathcal\{F\}\(i\), using reference pilots indexed by𝒜\\mathcal\{A\}, with those of the neighbors in𝒩​\(i\)\\mathcal\{N\}\(i\)\. This is achieved by solving in a decentralized manner

min\{𝜽i=\(𝝋i,𝝍i\)\}\{𝐕i​j∈St​\(di,dj\)\}\{𝐕j​i∈St​\(dj,di\)\}∑i∈𝒩ℒi​\(𝜽i\)\+ℛ𝒜​\(𝝋i,\{𝐕i​j\},\{𝐕j​i\}\);\\min\_\{\\begin\{subarray\}\{c\}\\\{\\boldsymbol\{\\theta\}\_\{i\}=\(\\bm\{\\varphi\}\_\{i\},\\bm\{\\psi\}\_\{i\}\)\\\}\\\\ \\\{\\mathbf\{V\}\_\{ij\}\\in\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)\\\}\\\\ \\\{\\mathbf\{V\}\_\{ji\}\\in\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\)\\\}\\end\{subarray\}\}\\quad\\sum\_\{i\\in\\mathcal\{N\}\}\\mathcal\{L\}\_\{i\}\(\\bm\{\\theta\}\_\{i\}\)\+\\mathcal\{R\}\_\{\\mathcal\{A\}\}\(\\bm\{\\varphi\}\_\{i\},\\\{\\mathbf\{V\}\_\{ij\}\\\},\\\{\\mathbf\{V\}\_\{ji\}\\\}\)\\,;\(SFRL\)while preserving agents’ privacy and reducing communication costs over the network\.

Prob\. \([SFRL](https://arxiv.org/html/2608.10016#S4.Ex3)\) is a regularized nonconvex optimization problem, which couples smooth local neural parameters optimization with manifold\-constrained edge transports\. It can be interpreted as a scalable decentralized relaxation of the centralized Prob\. \([8](https://arxiv.org/html/2608.10016#S3.E8)\)\.

## 5The Sheaf\-FRL Algorithm

We solve Prob\. \([SFRL](https://arxiv.org/html/2608.10016#S4.Ex3)\) in a decentralized manner, by adopting an alternating minimization strategy\. The latter consists of the following two updates\.

### 5\.1Step 1: Isometric Embedding Update

Recall that𝒩​\(i\)⊆𝒩∖\{i\}\\mathcal\{N\}\(i\)\\subseteq\\mathcal\{N\}\\setminus\\\{i\\\}denotes the set of neighbors of nodeii, which splits into𝒩​\(i\)−\\mathcal\{N\}\(i\)^\{\-\}and𝒩​\(i\)\+\\mathcal\{N\}\(i\)^\{\+\}, collecting the neighborsjjfor whichiiis the head and the tail, respectively\. Given the neural parameters\{𝜽it=\(𝝋it,𝝍it\)\}\\\{\\bm\{\\theta\}\_\{i\}^\{t\}=\(\\bm\{\\varphi\}\_\{i\}^\{t\},\\bm\{\\psi\}\_\{i\}^\{t\}\)\\\}updated at communication roundtt, each nodei∈𝒩i\\in\\mathcal\{N\}updates its alignment map with its neighborsj∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)\. Let

𝐀it=𝐀i​\(𝝋it\)=\[f𝝋it​\(𝐱ik\)\]k∈𝒜\\mathbf\{A\}\_\{i\}^\{t\}=\\mathbf\{A\}\_\{i\}\\left\(\\bm\{\\varphi\}\_\{i\}^\{t\}\\right\)=\\Big\[f\_\{\\bm\{\\varphi\}^\{t\}\_\{i\}\}\\left\(\\mathbf\{x\}^\{k\}\_\{i\}\\right\)\\Big\]\_\{k\\in\\mathcal\{A\}\}denote the pilot feature matrix at communication roundtt\. Since the local gluing penalty in[Eq\.16](https://arxiv.org/html/2608.10016#S4.E16)separates over the oriented edges, each nodeiican update its alignment maps by solving local problems involving only its neighborsj∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)\. In particular, for eachei​j∈ℰe\_\{ij\}\\in\\mathcal\{E\}, whereiiis the head andjjthe tail, nodeiioptimizes the corresponding incoming and outgoing embedding maps associated with its neighborhood\. This leads to two symmetric alignment problems corresponding to the incoming and outgoing embedding terms\. Specifically, we consider

𝐕i​jt\\displaystyle\\mathbf\{V\}\_\{ij\}^\{t\}=arg​min𝐕i​j∈St​\(di,dj\)⁡‖𝐀it−𝐕i​j​𝐀jt‖F2,∀j∈𝒩​\(i\)−,\\displaystyle=\\operatorname\*\{arg\\,min\}\_\{\\mathbf\{V\}\_\{ij\}\\in\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)\}\\;\\;\\norm\{\\mathbf\{A\}\_\{i\}^\{t\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{A\}\_\{j\}^\{t\}\}\_\{\\mathrm\{F\}\}^\{2\}\\,,\\quad\\forall\\,j\\in\\mathcal\{N\}\(i\)^\{\-\}\\,,\(P2a\)𝐕j​it\\displaystyle\\mathbf\{V\}\_\{ji\}^\{t\}=arg​min𝐕j​i∈St​\(dj,di\)⁡‖𝐀jt−𝐕j​i​𝐀it‖F2,∀j∈𝒩​\(i\)\+\.\\displaystyle=\\operatorname\*\{arg\\,min\}\_\{\\mathbf\{V\}\_\{ji\}\\in\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\)\}\\;\\;\\norm\{\\mathbf\{A\}\_\{j\}^\{t\}\-\\mathbf\{V\}\_\{ji\}\\mathbf\{A\}\_\{i\}^\{t\}\}\_\{\\mathrm\{F\}\}^\{2\}\\,,\\quad\\forall\\,j\\in\\mathcal\{N\}\(i\)^\{\+\}\\,\.\(P2b\)The objective of Prob\. \([P2a](https://arxiv.org/html/2608.10016#S5.Ex5)\) \(and analogously of \([P2b](https://arxiv.org/html/2608.10016#S5.Ex6)\)\) can be recast as:

‖𝐀it−𝐕i​j​𝐀jt‖F2=tr⁡\(𝐀it​𝐀it⊤\)\+tr⁡\(𝐀jt​𝐀jt⊤\)−2​tr⁡\(𝐀it​𝐀jt⊤​𝐕i​j⊤\)\.\\norm\{\\mathbf\{A\}\_\{i\}^\{t\}\-\\mathbf\{V\}\_\{ij\}\\mathbf\{A\}\_\{j\}^\{t\}\}\_\{\\mathrm\{F\}\}^\{2\}=\\tr\{\\mathbf\{A\}\_\{i\}^\{t\}\\mathbf\{A\}\_\{i\}^\{\{t\}^\{\\top\}\}\}\+\\tr\{\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{j\}^\{\{t\}^\{\\top\}\}\}\-2\\tr\{\\mathbf\{A\}\_\{i\}^\{t\}\\mathbf\{A\}\_\{j\}^\{\{t\}^\{\\top\}\}\\mathbf\{V\}\_\{ij\}^\{\\top\}\}\\,\.\(17\)Minimizing[Eq\.17](https://arxiv.org/html/2608.10016#S5.E17)thus corresponds to maximizing the third term\. Denoting by𝐔​𝚺​𝐖⊤\\mathbf\{U\}\\bm\{\\Sigma\}\\mathbf\{W\}^\{\\top\}the thin SVD of the cross\-covariance term𝐀it​𝐀jt⊤\\mathbf\{A\}\_\{i\}^\{t\}\\mathbf\{A\}\_\{j\}^\{\{t\}^\{\\top\}\}, the solution to \([P2a](https://arxiv.org/html/2608.10016#S5.Ex5)\) is given by

𝐕i​jt=𝐔𝐖⊤,where𝐔∈St​\(di,dj\)​and​𝐖∈O​\(dj\)\.\\mathbf\{V\}\_\{ij\}^\{t\}=\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}\\,,\\quad\\text\{where\}\\quad\\mathbf\{U\}\\in\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)\\text\{ and \}\\mathbf\{W\}\\in\\mathrm\{O\}\(\{d\_\{j\}\}\)\\,\.\(18\)Similarly, the solution to \([P2b](https://arxiv.org/html/2608.10016#S5.Ex6)\) is obtained from the thin SVD of𝐀jt​𝐀it⊤\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{i\}^\{\{t\}^\{\\top\}\}\. For each nodei∈𝒩i\\in\\mathcal\{N\}, the updates in[Eq\.18](https://arxiv.org/html/2608.10016#S5.E18)are computed locally once its neighborsj∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)share their feature evaluations\{𝐀jt\}\\\{\\mathbf\{A\}\_\{j\}^\{t\}\\\}\. Importantly, the orientation is exploited byiito distinguish the edges for which it acts as the head, in which case the incoming embedding map𝐕i​jt\\mathbf\{V\}\_\{ij\}^\{t\}must be computed via \([P2a](https://arxiv.org/html/2608.10016#S5.Ex5)\), from those for which it acts as the tail, where the outgoing embedding map𝐕j​it\\mathbf\{V\}\_\{ji\}^\{t\}is updated via \([P2b](https://arxiv.org/html/2608.10016#S5.Ex6)\)\.

### 5\.2Step 2: Gradient Update of Neural Parameters

In Prob\. \([SFRL](https://arxiv.org/html/2608.10016#S4.Ex3)\), each agentiiupdates its local parameters𝜽i\\bm\{\\theta\}\_\{i\}through a gradient step that minimizes the global objective with respect to𝜽i\\bm\{\\theta\}\_\{i\}, while keeping the neighboring parameters\{𝜽jt\}j∈𝒩​\(i\)\\\{\\bm\{\\theta\}\_\{j\}^\{t\}\\\}\_\{j\\in\\mathcal\{N\}\(i\)\}fixed at the current iteration\. The regularization term, instead, depends exclusively on the encoder parameters𝝋i\\bm\{\\varphi\}\_\{i\}\. Therefore, starting from[Eq\.16](https://arxiv.org/html/2608.10016#S4.E16), for each agenti∈𝒩i\\in\\mathcal\{N\}we obtain

∇𝝋iℛ𝒜\|i=\\displaystyle\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\mathcal\{R\}\_\{\\mathcal\{A\}\}\|\_\{i\}=λiK\[∑j∈𝒩​\(i\)−∑k∈𝒜\(∇𝝋if𝝋i\(𝐱ik\)\)⊤\(f𝝋i\(𝐱ik\)−𝐕i​jf𝝋j\(𝐱jk\)\)\\displaystyle\\frac\{\\lambda\_\{i\}\}\{K\}\\Bigg\[\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\}\}\\sum\_\{k\\in\\mathcal\{A\}\}\\left\(\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\right\)^\{\\top\}\\left\(f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\-\\mathbf\{V\}\_\{ij\}f\_\{\\boldsymbol\{\\varphi\}\_\{j\}\}\\left\(\\mathbf\{x\}\_\{j\}^\{k\}\\right\)\\right\)\(19\)−∑j∈𝒩​\(i\)\+∑k∈𝒜\(∇𝝋if𝝋i\(𝐱ik\)\)⊤𝐕j​i⊤\(f𝝋j\(𝐱jk\)−𝐕j​if𝝋i\(𝐱ik\)\)\]\\displaystyle\-\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\+\}\}\\sum\_\{k\\in\\mathcal\{A\}\}\\left\(\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\right\)^\{\\top\}\\mathbf\{V\}\_\{ji\}^\{\\top\}\\left\(f\_\{\\boldsymbol\{\\varphi\}\_\{j\}\}\\left\(\\mathbf\{x\}\_\{j\}^\{k\}\\right\)\-\\mathbf\{V\}\_\{ji\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\\right\)\\Bigg\]=\(a\)\\displaystyle\\stackrel\{\{\\scriptstyle\(a\)\}\}\{\{=\}\}λiK​∑j∈𝒩​\(i\)∑k∈𝒜\(∇𝝋if𝝋i​\(𝐱ik\)\)⊤​\(f𝝋i​\(𝐱ik\)−𝐕i​j​f𝝋j​\(𝐱jk\)\);\\displaystyle\\frac\{\\lambda\_\{i\}\}\{K\}\\sum\_\{j\\in\\mathcal\{N\}\(i\)\}\\sum\_\{k\\in\\mathcal\{A\}\}\\left\(\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\right\)^\{\\top\}\\left\(f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\-\\mathbf\{V\}\_\{ij\}f\_\{\\boldsymbol\{\\varphi\}\_\{j\}\}\\left\(\\mathbf\{x\}\_\{j\}^\{k\}\\right\)\\right\)\\,;where∇𝝋if𝝋i​\(𝐱ik\)∈ℝ\|𝝋i\|×di\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\in\\mathbb\{R\}^\{\|\\boldsymbol\{\\varphi\}\_\{i\}\|\\times d\_\{i\}\}denotes the Jacobian of the feature mapf𝝋if\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}with respect to the parameter vector𝝋i\\boldsymbol\{\\varphi\}\_\{i\}, evaluated at𝐳ik\\mathbf\{z\}\_\{i\}^\{k\}; and in\(a\)\(a\), forj∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}, we used that𝐕i​j=𝐕j​i⊤\\mathbf\{V\}\_\{ij\}=\\mathbf\{V\}\_\{ji\}^\{\\top\}and𝐕i​j​𝐕j​i=𝐈di\\mathbf\{V\}\_\{ij\}\\mathbf\{V\}\_\{ji\}=\\mathbf\{I\}\_\{d\_\{i\}\}, which follow from the orthonormality of the columns of𝐕j​i∈St​\(dj,di\)\\mathbf\{V\}\_\{ji\}\\in\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\)\. This shows that both incoming and outgoing contributions admit a unified expression\. Furthermore, let

𝐫i=\[∇𝝋iℛ𝒜\|i𝟎\|𝝍i\|\]\\mathbf\{r\}\_\{i\}=\\begin\{bmatrix\}\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\mathcal\{R\}\_\{\\mathcal\{A\}\}\|\_\{i\}\\\\ \\boldsymbol\{0\}\_\{\|\\bm\{\\psi\}\_\{i\}\|\}\\end\{bmatrix\}\(20\)be the contribution of the sheaf\-based regularization to the gradient of the local neural parameters𝜽i\\bm\{\\theta\}\_\{i\}\. Then, for each agenti∈𝒩i\\in\\mathcal\{N\}, the neural parameters are updated locally via

𝜽it\+1=𝜽it−ηθ​\(∇𝜽iℒi​\(𝜽i\)\+𝐫it\),\\boldsymbol\{\\theta\}\_\{i\}^\{t\+1\}=\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\-\\eta\_\{\\theta\}\\left\(\\nabla\_\{\\boldsymbol\{\\theta\}\_\{i\}\}\\mathcal\{L\}\_\{i\}\(\\bm\{\\theta\}\_\{i\}\)\+\\mathbf\{r\}^\{t\}\_\{i\}\\right\),\(21\)whereηθ\>0\\eta\_\{\\theta\}\>0is the stepsize, and𝐫it\\mathbf\{r\}^\{t\}\_\{i\}is the gradient term in \([20](https://arxiv.org/html/2608.10016#S5.E20)\) evaluated at timett\.

Gradient computations in \([21](https://arxiv.org/html/2608.10016#S5.E21)\) are fully decentralized\. Specifically, each agentiican evaluate the gradient with respect to its local parameters𝜽i\\bm\{\\theta\}\_\{i\}using only locally available information together with the pilots’ latent representations\{𝐀jt\}j∈𝒩​\(i\)\\\{\\mathbf\{A\}\_\{j\}^\{t\}\\\}\_\{j\\in\\mathcal\{N\}\(i\)\}received from its neighboring agents\. Since these pilot representations are already exchanged during Step 1 to update the restriction maps, no additional communication is required to evaluate the regularization gradient in \([19](https://arxiv.org/html/2608.10016#S5.E19)\)\. The resulting decentralized algorithm, termed*Sheaf\-FRL*, is summarized in Appendix[D](https://arxiv.org/html/2608.10016#A4)for both the heterogeneous and homogeneous settings\. A detailed convergence analysis, covering both deterministic and stochastic settings, is presented in Appendix[E](https://arxiv.org/html/2608.10016#A5)\. Finally, the proposed framework naturally accommodates accelerated first\-order methods: momentum\-based and adaptive schemes, such as Nesterov acceleration or Adam, can be directly incorporated into the local parameter updates \([21](https://arxiv.org/html/2608.10016#S5.E21)\) without affecting the decentralized structure of the algorithm\.

## 6Collaborative supervised classification with semantic communication

As an application, we test our proposed Sheaf\-FRL algorithm in a collaborative and decentralized classification task in the setting of semantic communication, where agents exchange compressed latent representations of data\(Gündüz et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib28); Barbarossa et al\.,[2023](https://arxiv.org/html/2608.10016#bib.bib4); Strinati et al\.,[2024](https://arxiv.org/html/2608.10016#bib.bib60); Pandolfo et al\.,[2025](https://arxiv.org/html/2608.10016#bib.bib50); Grimaldi et al\.,[2025](https://arxiv.org/html/2608.10016#bib.bib27)\)\. Each agent minimizes a local cross\-entropy lossℒi​\(𝜽i\)≔CE​\(𝜽i\)\\mathcal\{L\}\_\{i\}\(\\bm\{\\theta\}\_\{i\}\)\\coloneq\\mathrm\{CE\}\(\\bm\{\\theta\}\_\{i\}\), while the representation alignment across agents is promoted via the sheaf regularization term\. At a high level, the goal for the agents is to improve their classification accuracy on both*\(i\)*their private local latent representations, and*\(ii\)*the latent representations received by their correspondingneighboring agentsj∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)\. Accordingly, as customary in this application setting, we monitor*average private*and*average communication accuracies*\. Denote by𝒯itest\\mathcal\{T\}\_\{i\}^\{\\text\{test\}\}the test set for agentiiconsisting ofMitestM\_\{i\}^\{\\text\{test\}\}pairs\{\(𝐱itest,n,yitest,n\)\}n∈Mitest\\\{\(\\mathbf\{x\}\_\{i\}^\{\\text\{test\},n\},y\_\{i\}^\{\\text\{test\},n\}\)\\\}\_\{n\\in M\_\{i\}^\{\\text\{test\}\}\}\. Then,

Aipriv\\displaystyle A\_\{i\}^\{\\text\{priv\}\}=1Mitest∑n=1Mitest𝟏\[argmaxc∈𝒴\[h𝝍i∘f𝝋i\(𝐱itest,n\)\]c=yitest,n\]\\displaystyle=\\dfrac\{1\}\{M\_\{i\}^\{\\text\{test\}\}\}\\sum\_\{n=1\}^\{M\_\{i\}^\{\\text\{test\}\}\}\\boldsymbol\{1\}\[\\arg\\max\_\{c\\in\\mathcal\{Y\}\}\[h\_\{\\bm\{\\psi\}\_\{i\}\}\\circ\{\\color\[rgb\]\{0,0,0\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{0,0,0\}f\_\{\\bm\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{\\text\{test\},n\}\)\}\]\_\{c\}=y\_\{i\}^\{\\text\{test\},n\}\]\\,\(22\)Aicomm\\displaystyle A\_\{i\}^\{\\text\{comm\}\}=1\|𝒩​\(i\)\|\(∑j∈𝒩​\(i\)−1Mjtest∑n=1Mjtest𝟏\[argmaxc∈𝒴\[h𝝍i∘𝐕i​j∘f𝝋j\(𝐱jtest,n\)\]c=yjtest,n\]\+\\displaystyle=\\dfrac\{1\}\{\|\\mathcal\{N\}\(i\)\|\}\\Bigg\(\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\}\}\\dfrac\{1\}\{M\_\{j\}^\{\\text\{test\}\}\}\\sum\_\{n=1\}^\{M\_\{j\}^\{\\text\{test\}\}\}\\boldsymbol\{1\}\[\\arg\\max\_\{c\\in\\mathcal\{Y\}\}\[h\_\{\\bm\{\\psi\}\_\{i\}\}\\circ\\mathbf\{V\}\_\{ij\}\\circ\{\\color\[rgb\]\{0,\.5,\.5\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{0,\.5,\.5\}f\_\{\\bm\{\\varphi\}\_\{j\}\}\(\\mathbf\{x\}\_\{j\}^\{\\text\{test\},n\}\)\}\]\_\{c\}=y\_\{j\}^\{\\text\{test\},n\}\]\+∑j∈𝒩​\(i\)\+1Mjtest∑n=1Mjtest𝟏\[argmaxc∈𝒴\[h𝝍i∘𝐕j​i⊤∘f𝝋j\(𝐱jtest,n\)\]c=yjtest,n\]\)\.\\displaystyle\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\+\}\}\\dfrac\{1\}\{M\_\{j\}^\{\\text\{test\}\}\}\\sum\_\{n=1\}^\{M\_\{j\}^\{\\text\{test\}\}\}\\boldsymbol\{1\}\[\\arg\\max\_\{c\\in\\mathcal\{Y\}\}\[h\_\{\\bm\{\\psi\}\_\{i\}\}\\circ\\mathbf\{V\}\_\{ji\}^\{\\top\}\\circ\{\\color\[rgb\]\{0,\.5,\.5\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{0,\.5,\.5\}f\_\{\\bm\{\\varphi\}\_\{j\}\}\(\\mathbf\{x\}\_\{j\}^\{\\text\{test\},n\}\)\}\]\_\{c\}=y\_\{j\}^\{\\text\{test\},n\}\]\\Bigg\)\\,\.
Model heterogeneity\.We considerN=15N=15agents, each consisting of*\(i\)*a CNN local encoder and*\(ii\)*an MLP local decoder\. This choice of qualitative architectural similarity is driven by the task homogeneity of the application setting\. Heterogeneity across models is enforced by varying architectural characteristics \(encoder and decoder widths\) and the dropout hyper\-parameter \(cf\.[Tab\.2](https://arxiv.org/html/2608.10016#A7.T2)\)\.

Data heterogeneity\.We use MNIST as a global dataset𝒯\\mathcal\{T\}with label space𝒴=\{0,…,9\}\\mathcal\{Y\}=\\\{0,\\dots,9\\\}\. We first set aside a global pilot set𝒜\\mathcal\{A\}, a uniformly sampled10%10\\%subset of𝒯\\mathcal\{T\}, held out before any agent\-level partitioning;𝒜\\mathcal\{A\}is shared, unmodified, across all agents and used exclusively to align/communicate their representations\. The remaining90%90\\%is partitioned disjointly across theNNagents to create heterogeneity: each agentiihas a target class set𝒞i⊂𝒴\\mathcal\{C\}\_\{i\}\\subset\\mathcal\{Y\}, and its local label distribution is a convex combination of two uniform distributions, the first over its target classes and the second over all classes, viz\.

Pi​\(y\)≔s​Unif​\{𝒞i\}\+\(1−s\)​Unif​\{𝒴\},∀i∈\[N\];P\_\{i\}\(y\)\\coloneq s\\,\\mathrm\{Unif\}\\\{\\mathcal\{C\}\_\{i\}\\\}\+\(1\-s\)\\,\\mathrm\{Unif\}\\\{\\mathcal\{Y\}\\\}\\,,\\qquad\\forall i\\in\[N\]\\,;\(23\)wheres∈\[0,1\]s\\in\[0,1\]is the distribution\-shift strength\. Samples of each class are split across agents in proportion to\{Pi​\(y\)\}i=1N\\\{P\_\{i\}\(y\)\\\}\_\{i=1\}^\{N\}\. Each agent’s local pool is then split80%/10%/10%80\\%/10\\%/10\\%into train/val/test, preserving its label skew across all three splits\.

Network topology\.Agents communicate over an undirected, unweighted graph𝒢≔\(𝒩,ℰ\)\\mathcal\{G\}\\coloneqq\(\\mathcal\{N\},\\mathcal\{E\}\), with density level40%40\\%, built to be connected and to reflect the degree of class overlap across agents, as described below\. Each agenti∈\[N\]i\\in\[N\]is assigned to a distinct node, so that\|𝒩\|=N\|\\mathcal\{N\}\|=N\. The edge setℰ\\mathcal\{E\}is constructed from the target\-class sets\{𝒞i\}i∈\[N\]\\\{\\mathcal\{C\}\_\{i\}\\\}\_\{i\\in\[N\]\}as follows\. For every pair\(i,j\)\(i,j\), letwi​j≔\|𝒞i∩𝒞j\|w\_\{ij\}\\coloneqq\|\\mathcal\{C\}\_\{i\}\\cap\\mathcal\{C\}\_\{j\}\|denote the number of target classes shared by agentsiiandjj, collected into the weight vector𝐰\\mathbf\{w\}\. These weights define the fully\-connected, weighted graph𝒢~≔\(𝒩,ℰ~,𝐰\)\\widetilde\{\\mathcal\{G\}\}\\coloneqq\(\\mathcal\{N\},\\widetilde\{\\mathcal\{E\}\},\\mathbf\{w\}\), withℰ~≔\{\(i,j\)∈\[N\]×\[N\]∣i≠j\}\\widetilde\{\\mathcal\{E\}\}\\coloneqq\\\{\(i,j\)\\in\[N\]\\times\[N\]\\,\\mid\\,i\\neq j\\\}\. We then extract the maximum\-weight spanning tree𝒢^≔\(𝒩,ℰ^,𝐰^\)\\hat\{\\mathcal\{G\}\}\\coloneqq\(\\mathcal\{N\},\\hat\{\\mathcal\{E\}\},\\hat\{\\mathbf\{w\}\}\)from𝒢~\\widetilde\{\\mathcal\{G\}\}, and initializeℰ\\mathcal\{E\}with the \(unweighted\) edges ofℰ^\\hat\{\\mathcal\{E\}\}, which guarantees connectivity\. We next add edges fromℰ~∖ℰ^\\widetilde\{\\mathcal\{E\}\}\\setminus\\hat\{\\mathcal\{E\}\}, in decreasing order of weight, until the target40%40\\%density is reached\. Finally, each edge is assigned an orientation, consistent with the orientation rule in[Sec\.4](https://arxiv.org/html/2608.10016#S4)\.

Baselines\.This heterogeneous setup precludes direct comparison with classical federated methods based on model or submodel averaging\(Scardapane & Di Lorenzo,[2017](https://arxiv.org/html/2608.10016#bib.bib53); Diao et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib17); Li et al\.,[2020a](https://arxiv.org/html/2608.10016#bib.bib38)\), or shared representations\(Tan et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib62); Setayesh et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib55); Collins et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib11); Liang et al\.,[2020](https://arxiv.org/html/2608.10016#bib.bib42)\)\. Indeed, these methods assume homogeneous architectures or shared latent spaces\. We therefore compare against methods explicitly designed for heterogeneous settings, namely*ComFed*\(Badi et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib2)\)and*Sheaf\-FMTL*\(Issaid et al\.,[2025](https://arxiv.org/html/2608.10016#bib.bib32)\)\. We also consider a*non\-cooperative*baseline in which models are trained independently and their latent spaces are aligned only after training, providing a natural lower bound\. All baselines use the same communication network topology as Sheaf\-FRL\. Except for ComFed, which learns alignment maps during training, all baselines require post\-hoc alignment to enable semantic communication\. Specifically, for each edge\(i,j\)\(i,j\)we solve \([P2a](https://arxiv.org/html/2608.10016#S5.Ex5)\) and \([P2b](https://arxiv.org/html/2608.10016#S5.Ex6)\) in closed\-form via[Eq\.18](https://arxiv.org/html/2608.10016#S5.E18), fitting isometric maps between whitened latent spaces using the same pilot sets as in Sheaf\-FRL\. In our experiments, the whitening for these baselines is pursued through classical zero\-phase component analysis\(Kessy et al\.,[2018](https://arxiv.org/html/2608.10016#bib.bib37)\)fitted on the latent representations from the private training set, independently for each agent\.

![Refer to caption](https://arxiv.org/html/2608.10016v1/x2.png)Figure 2:Average private accuracy \(left\) and average communication accuracy \(right\) against the distribution\-shift strengthss, for the collaborative supervised classification task with1515agents\. Markers denote the simple average across agents for the private accuracy, and the degree\-weighted average for the communication accuracy, to account for the imbalance in the agents’ degree distribution\.Results\.[Fig\.2](https://arxiv.org/html/2608.10016#S6.F2)shows the metrics in[Eq\.22](https://arxiv.org/html/2608.10016#S6.E22)against the distribution\-shift parameterss, which controls the level of data heterogeneity according to[Eq\.23](https://arxiv.org/html/2608.10016#S6.E23)\. Overall, Sheaf\-FRL consistently outperforms all baselines on both metrics\. Looking at the average communication accuracy, the performance gap widens asssincreases, suggesting that our framework is more effective at learning latent spaces that are*\(i\)*well aligned across agents, thus enhancing transferability, and*\(ii\)*more relevant to the downstream task, thus being semantically meaningful\. Point*\(ii\)*is further supported by the average private accuracy, where Sheaf\-FRL still compares favorably with the baselines, although the margin is less pronounced\. Among the baselines, ComFed performs poorly in terms of average communication accuracy, while being comparable to the others in private accuracy\. This suggests that the joint learning of general, unstructured alignment maps and compressed representations of a single shared latent space is not well suited for transferability\. Sheaf\-FMTL, in turn, shows no consistent improvement in communication accuracy over the non\-cooperative baseline\. This stems from a limitation in how it parameterizes the restriction maps: since Sheaf\-FMTL connects parameter spaces rather than latent representation spaces, it inherits their much higher dimensionality\. The resulting memory overhead requires applying a high compression factor to the edge stalks, which negatively affects the diffusion of information across the network \(Appendix[G](https://arxiv.org/html/2608.10016#A7)for further details\)\.

### 6\.1Robustness to semantic compression

An important feature in semantic communication is robustness to latent\-space compression\(Gündüz et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib28); Barbarossa et al\.,[2023](https://arxiv.org/html/2608.10016#bib.bib4)\)\. This can be understood as the ability of methods to extract semantically meaningful information to be communicated\. To investigate this aspect, we restrict our focus to two agents, sayiiandjj, connected by an edge\(i,j\)\(i,j\), and set the level of distribution shift tos=0\.7s=0\.7\. Then, we monitor the two metrics in[Eq\.22](https://arxiv.org/html/2608.10016#S6.E22)across different bottleneck dimensions\. Specifically, although the overall architectures ofiiandjjare different \(Appendix[G](https://arxiv.org/html/2608.10016#A7)for further details\), they share the same latent space dimensionality,d=di=djd=d\_\{i\}=d\_\{j\}\. Consequently, in this case, we can include among the baselines those methods suitable for heterogeneous architectures sharing the same latent space dimensionality\. Specifically, we consider theFedProto\(Tan et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib62)\)andFedMuscle\(Setayesh et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib55)\)methods, although their communication protocols are centralized and therefore not directly applicable to a decentralized setting\.

![Refer to caption](https://arxiv.org/html/2608.10016v1/x3.png)Figure 3:Average private accuracy \(left\) and average communication accuracy \(right\) against the bottleneck dimension at distribution\-shift strengths=0\.7s=0\.7, for the robustness\-to\-compression case study\. Markers denote the simple average across the two agents\.Results\.[Fig\.3](https://arxiv.org/html/2608.10016#S6.F3)shows the results for all the considered methods\. Overall, Sheaf\-FRL achieves a better trade\-off between average classification accuracy \(private and communication\) and bottleneck dimensionality than the baselines, with the performance gap increasing as the bottleneck dimension decreases, i\.e\., in the high\-compression regime\. As expected, all methods suffer as the compression increases; however, the performance of Sheaf\-FRL degrades much more gracefully\. Notably, it preserves high accuracy even at a bottleneck dimension of1616, where we approach the minimum number of dimensions required by the network to discriminate the1010MNIST classes\. With respect to the baselines, their performance is comparable to that of the non\-cooperative method, occasionally showing even lower communication accuracy despite employing the same post\-training alignment strategy\. The results confirm the poor transferability of the representations learned under these collaboration schemes: an agent’s latent space is not readily reusable by another, regardless of whether a post\-training alignment pipeline is applied\. Moreover, in terms of average communication accuracy, ComFed achieves the lowest performance among the considered methods despite its learnable alignment maps, consistent with the results in[Fig\.2](https://arxiv.org/html/2608.10016#S6.F2)\.

## 7Conclusions and Future Work

We introduced*Sheaf\-based Federated Representation Learning*, a framework that enables semantic alignment among heterogeneous agents without imposing a shared global latent space\. By modeling agent\-specific representations as sections of a learnable network sheaf and enforcing consistency through a quadratic gluing penalty induced by the sheaf Laplacian, SFRL accommodates differences in data distributions, sensing modalities, architectures, and latent dimensionalities, while remaining compatible with supervised, semi\-supervised, and self\-supervised learning paradigms\. We developed Sheaf\-FRL, a fully decentralized alternating algorithm that couples local gradient updates with closed\-form Procrustes updates of the edge\-wise restriction maps, and established its convergence to first\-order stationary points in both deterministic and stochastic settings\. Applied to a cooperative classification task in semantic communication under model and data heterogeneity, Sheaf\-FRL outperformed existing baselines in terms of local and post\-communication classification accuracy across different levels of distribution shift\. Additionally, it proved to be more robust to latent\-space compression\. Notably, this is achieved while evaluating the gluing penalty on only a small set of shared pilots, keeping communication overhead low\. These results support the sheaf\-theoretic perspective as a principled and flexible foundation for representation alignment in heterogeneous federated systems\. Future work includes extending the application of SFRL to multi\-modal, semi\-supervised, and self\-supervised settings, as well as improving the efficiency of Sheaf\-FRL from both an algorithmic and a communication\-strategy perspective—for instance by replacing full pilot representations with prototype pilots—to further reduce communication cost\.

#### Acknowledgments

The work was supported by the SNS JU project 6G\-GOALS\(Strinati et al\.,[2024](https://arxiv.org/html/2608.10016#bib.bib60)\)under the EU’s Horizon program Grant Agreement No 101139232, and by Huawei Technology France SASU under Grant N\. Tg20250616041\.

## References

- Absil et al\. \(2008\)P\.\-A\. Absil, Robert Mahony, and Rodolphe Sepulchre\.*Optimization Algorithms on Matrix Manifolds*\.Princeton University Press, 2008\.
- Badi et al\. \(2026\)Mohamed Badi, Chaouki Ben Issaid, and Mehdi Bennis\.Communication\-efficient and robust multi\-modal federated learning via latent\-space consensus\.*IEEE Wireless Communications Letters*, pp\. 1–1, 2026\.doi:10\.1109/LWC\.2026\.3676885\.
- Bandeira et al\. \(2013\)Afonso S\. Bandeira, Amit Singer, and Daniel A\. Spielman\.A Cheeger inequality for the graph connection Laplacian\.*SIAM Journal on Matrix Analysis and Applications*, 34\(4\):1611–1630, 2013\.doi:10\.1137/120875338\.
- Barbarossa et al\. \(2023\)Sergio Barbarossa, Danilo Comminiello, Eleonora Grassucci, Francesco Pezone, Stefania Sardellitti, and Paolo Di Lorenzo\.Semantic communications based on adaptive generative models and information bottleneck\.*IEEE Communications Magazine*, 61\(11\):36–41, 2023\.
- Barbero et al\. \(2022\)Federico Barbero, Cristian Bodnar, Haitz Sáez de Ocáriz Borde, Michael Bronstein, Petar Veličković, and Pietro Liò\.Sheaf neural networks with connection Laplacians\.In*Topological, Algebraic and Geometric Learning Workshops 2022*, pp\. 28–36\. PMLR, 2022\.
- Beck \(2017\)Amir Beck\.*First\-order methods in optimization*\.SIAM, 2017\.
- Bodnar et al\. \(2022\)Cristian Bodnar, Francesco Di Giovanni, Benjamin Paul Chamberlain, Pietro Liò, and Michael M\. Bronstein\.Neural sheaf diffusion: A topological perspective on heterophily and oversmoothing in GNNs\.In*Advances in Neural Information Processing Systems \(NeurIPS\)*, 2022\.URL[https://dblp\.org/rec/conf/nips/BodnarGCLB22](https://dblp.org/rec/conf/nips/BodnarGCLB22)\.
- Bredon \(1997\)Glen E Bredon\.*Sheaf theory*, volume 170\.Springer Science & Business Media, 1997\.
- Che et al\. \(2023\)Liwei Che, Jiaqi Wang, Yao Zhou, and Fenglong Ma\.Multimodal federated learning: A survey\.*Sensors*, 23\(15\):6986, 2023\.
- Chung et al\. \(2014\)Fan Chung, Wenbo Zhao, and Mark Kempton\.Ranking and sparsifying a connection graph\.*Internet Mathematics*, 10\(1\-2\):87–115, 2014\.
- Collins et al\. \(2021\)Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai\.Exploiting shared representations for personalized federated learning\.In*Proceedings of the 38th International Conference on Machine Learning \(ICML\)*, volume 139 of*PMLR*, 2021\.
- Curry \(2014\)Justin Michael Curry\.*Sheaves, cosheaves and applications*\.PhD thesis, University of Pennsylvania, 2014\.
- D’Acunto et al\. \(2025\)Gabriele D’Acunto, Fabio Massimo Zennaro, Yorgos Felekis, and Paolo Di Lorenzo\.Causal abstraction learning based on the semantic embedding principle\.In*Forty\-second International Conference on Machine Learning*, 2025\.URL[https://openreview\.net/forum?id=J16AIOkjjY](https://openreview.net/forum?id=J16AIOkjjY)\.
- D’Acunto et al\. \(2026\)Gabriele D’Acunto, Paolo Di Lorenzo, and Sergio Barbarossa\.Networks of causal abstractions: A sheaf\-theoretic framework\.*arXiv preprint arXiv:2509\.25236*, 2026\.
- Di Nino et al\. \(2024\)Leonardo Di Nino, Sergio Barbarossa, and Paolo Di Lorenzo\.Learning sheaf Laplacian optimizing restriction maps\.In*2024 58th Asilomar Conference on Signals, Systems, and Computers*, pp\. 59–63\. IEEE, 2024\.
- Di Nino et al\. \(2026\)Leonardo Di Nino, Gabriele D’Acunto, Sergio Barbarossa, and Paolo Di Lorenzo\.Learning the structure of connection graphs\.In*ICASSP 2026\-2026 IEEE International Conference on Acoustics, Speech and Signal Processing \(ICASSP\)*, pp\. 76–80\. IEEE, 2026\.
- Diao et al\. \(2021\)Enmao Diao, Jie Ding, and Vahid Tarokh\.HeteroFL: Computation and communication efficient federated learning for heterogeneous clients\.In*International Conference on Learning Representations \(ICLR\)*, 2021\.
- Dinh et al\. \(2022\)Canh T Dinh, Tung T Vu, Nguyen H Tran, Minh N Dao, and Hongyu Zhang\.A new look and convergence rate of federated multitask learning with Laplacian regularization\.*IEEE Transactions on Neural Networks and Learning Systems*, 35\(6\):8075–8085, 2022\.
- Dong & Voiculescu \(2021\)Nanqing Dong and Irina Voiculescu\.Federated contrastive learning for decentralized unlabeled medical images\.In*International Conference on Medical Image Computing and Computer\-Assisted Intervention*, pp\. 378–387\. Springer, 2021\.
- Edelman et al\. \(1998\)Alan Edelman, Tomás A\. Arias, and Steven T\. Smith\.The geometry of algorithms with orthogonality constraints\.*SIAM Journal on Matrix Analysis and Applications*, 20\(2\):303–353, 1998\.doi:10\.1137/S0895479895290954\.
- Evgeniou & Pontil \(2004\)Theodoros Evgeniou and Massimiliano Pontil\.Regularized multi–task learning\.In*Proceedings of the Tenth ACM SIGKDD International Conference on Knowledge Discovery and Data Mining*, pp\. 109–117, 2004\.
- Evgeniou et al\. \(2005\)Theodoros Evgeniou, Charles A Micchelli, Massimiliano Pontil, and John Shawe\-Taylor\.Learning multiple tasks with kernel methods\.*Journal of Machine Learning Research*, 6\(4\), 2005\.
- Fiorellino et al\. \(2024\)Simone Fiorellino, Claudio Battiloro, Emilio Calvanese Strinati, and Paolo Di Lorenzo\.Dynamic relative representations for goal\-oriented semantic communications\.In*Proc\. IEEE EUSIPCO*, pp\. 2107–2111, 2024\.
- Fiorellino et al\. \(2026\)Simone Fiorellino, Claudio Battiloro, Emilio Calvanese Strinati, and Paolo Di Lorenzo\.Frame\-based zero\-shot semantic channel equalization for AI\-native communications\.*IEEE Transactions on Cognitive Communications and Networking*, 2026\.
- Fredrikson et al\. \(2015\)Matt Fredrikson, Somesh Jha, and Thomas Ristenpart\.Model inversion attacks that exploit confidence information and basic countermeasures\.In*Proceedings of the 22nd ACM SIGSAC Conference on Computer and Communications Security \(CCS\)*, pp\. 1322–1333, 2015\.
- Ghalkha et al\. \(2026\)Abdulmomen Ghalkha, Zhuojun Tian, Chaouki Ben Issaid, and Mehdi Bennis\.SheafAlign: A sheaf\-theoretic framework for decentralized multimodal alignment\.*IEEE Communications Letters*, 2026\.
- Grimaldi et al\. \(2025\)Enrico Grimaldi, Mario Edoardo Pandolfo, Gabriele D’Acunto, Sergio Barbarossa, and Paolo Di Lorenzo\.Learning network sheaves for AI\-native semantic communication\.In*2025 59th Asilomar Conference on Signals, Systems, and Computers*, pp\. 1692–1696\. IEEE, 2025\.
- Gündüz et al\. \(2022\)Deniz Gündüz, Zhijin Qin, Inaki Estella Aguerri, Harpreet S Dhillon, Zhaohui Yang, Aylin Yener, Kai Kit Wong, and Chan\-Byoung Chae\.Beyond transmitting bits: Context, semantics, and task\-oriented communications\.*IEEE Journal on Selected Areas in Communications*, 41\(1\):5–41, 2022\.
- Han et al\. \(2022\)Sungwon Han, Sungwon Park, Fangzhao Wu, Sundong Kim, Chuhan Wu, Xing Xie, and Meeyoung Cha\.Fedx: Unsupervised federated learning with cross knowledge distillation\.In*European Conference on Computer Vision*, pp\. 691–707\. Springer, 2022\.
- Hansen & Ghrist \(2019\)Jakob Hansen and Robert Ghrist\.Toward a spectral theory of cellular sheaves\.*Journal of Applied and Computational Topology*, 3\(4\):315–358, 2019\.
- Hüttebräucker et al\. \(2024\)Tomás Hüttebräucker, Simone Fiorellino, Mohamed Sana, Paolo Di Lorenzo, and Emilio Calvanese Strinati\.Relative representations of latent spaces enable efficient semantic channel equalization\.In*Proc\. IEEE GLOBECOM*, 2024\.
- Issaid et al\. \(2025\)Chaouki Ben Issaid, Praneeth Vepakomma, and Mehdi Bennis\.Tackling feature and sample heterogeneity in decentralized multi\-task learning: A sheaf\-theoretic approach\.*Transactions on Machine Learning Research*, 2025\.ISSN 2835\-8856\.URL[https://openreview\.net/forum?id=JlPq0LmApB](https://openreview.net/forum?id=JlPq0LmApB)\.
- Jacob et al\. \(2008\)Laurent Jacob, Jean\-Philippe Vert, and Francis Bach\.Clustered multi\-task learning: A convex formulation\.*Advances in Neural Information Processing Systems*, 21, 2008\.
- Jaggi et al\. \(2014\)Martin Jaggi, Virginia Smith, Martin Takáč, Jonathan Terhorst, Sanjay Krishnan, Thomas Hofmann, and Michael I Jordan\.Communication\-efficient distributed dual coordinate ascent\.*Advances in Neural Information Processing Systems*, 27, 2014\.
- Kairouz et al\. \(2021\)Peter Kairouz, Brendan H\. Mcmahan, et al\.*Advances and Open Problems in Federated Learning*, volume 14 of*Foundations and Trends® in Machine Learning*\.2021\.doi:10\.1561/2200000083\.URL[https://inria\.hal\.science/hal\-02406503](https://inria.hal.science/hal-02406503)\.
- Karimireddy et al\. \(2020\)Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh\.SCAFFOLD: Stochastic controlled averaging for federated learning\.In*Proceedings of the 37th International Conference on Machine Learning*, PMLR, pp\. 5132–5143, 2020\.
- Kessy et al\. \(2018\)Agnan Kessy, Alex Lewin, and Korbinian Strimmer\.Optimal whitening and decorrelation\.*The American Statistician*, 72\(4\):309–314, 2018\.
- Li et al\. \(2020a\)Ang Li, Jingwei Sun, Binghui Wang, Lin Duan, Sicheng Li, Yiran Chen, and Hai Li\.LotteryFL: Personalized and communication\-efficient federated learning with lottery ticket hypothesis on non\-iid datasets\.*arXiv preprint arXiv:2008\.03371*, 2020a\.
- Li & Ma \(2022\)Jiaxiang Li and Shiqian Ma\.Federated learning on Riemannian manifolds\.*arXiv preprint arXiv:2206\.05668*, 2022\.
- Li et al\. \(2021\)Qinbin Li, Bingsheng He, and Dawn Song\.Model\-contrastive federated learning\.In*Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\)*, 2021\.
- Li et al\. \(2020b\)Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith\.Federated optimization in heterogeneous networks\.In*Proceedings of the Third Conference on Machine Learning and Systems \(MLSys\)*, 2020b\.
- Liang et al\. \(2020\)Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis\-Philippe Morency\.Think locally, act globally: Federated learning with local and global representations\.*arXiv preprint arXiv:2001\.01523*, 2020\.
- Ma et al\. \(2015\)Chenxin Ma, Virginia Smith, Martin Jaggi, Michael Jordan, Peter Richtárik, and Martin Takác\.Adding vs\. averaging in distributed primal\-dual optimization\.In*International Conference on Machine Learning*, pp\. 1973–1982\. PMLR, 2015\.
- McLaughlin & Su \(2024\)Connor J McLaughlin and Lili Su\.Personalized federated learning via feature distribution adaptation\.*Advances in Neural Information Processing Systems*, 37:77038–77059, 2024\.
- McMahan et al\. \(2017\)Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas\.Communication\-efficient learning of deep networks from decentralized data\.In*Proceedings of the 20th International Conference on Artificial Intelligence and Statistics \(AISTATS\)*, volume 54, pp\. 1273–1282\. PMLR, 2017\.
- Melis et al\. \(2019\)Luca Melis, Congzheng Song, Emiliano De Cristofaro, and Vitaly Shmatikov\.Exploiting unintended feature leakage in collaborative learning\.In*IEEE Symposium on Security and Privacy \(S&P\)*, 2019\.
- Miao & Koyuncu \(2024\)Runxuan Miao and Erdem Koyuncu\.Contrastive and non\-contrastive strategies for federated self\-supervised representation learning and deep clustering\.*IEEE Journal of Selected Topics in Signal Processing*, 18\(6\):1070–1084, 2024\.
- Moschella et al\. \(2023\)Luca Moschella, Valentino Maiorca, Marco Fumero, Antonio Norelli, Francesco Locatello, and Emanuele Rodolà\.Relative representations enable zero\-shot latent space communication\.In*The Eleventh International Conference on Learning Representations*, 2023\.URL[https://openreview\.net/forum?id=SrC\-nwieGJ](https://openreview.net/forum?id=SrC-nwieGJ)\.
- Nesterov \(2004\)Yurii Nesterov\.*Introductory Lectures on Convex Optimization: A Basic Course*, volume 87 of*Applied Optimization*\.Springer, 2004\.
- Pandolfo et al\. \(2025\)Mario Edoardo Pandolfo, Simone Fiorellino, Emilio Calvanese Strinati, and Paolo Di Lorenzo\.Latent space alignment for AI\-native mimo semantic communications\.In*2025 International Joint Conference on Neural Networks \(IJCNN\)*, pp\. 1–8\. IEEE, 2025\.
- Pandolfo et al\. \(2026\)Mario Edoardo Pandolfo, Enrico Grimaldi, Lorenzo Marinucci, Leonardo Di Nino, Simone Fiorellino, Sergio Barbarossa, and Paolo Di Lorenzo\.SEMASIA: A large\-scale dataset of semantically structured latent representations\.*arXiv preprint arXiv:2605\.09485*, 2026\.
- Robbins & Siegmund \(1971\)Herbert Robbins and David Siegmund\.A convergence theorem for non negative almost supermartingales and some applications\.pp\. 233–257\. Academic Press, 1971\.
- Scardapane & Di Lorenzo \(2017\)Simone Scardapane and Paolo Di Lorenzo\.A framework for parallel and distributed training of neural networks\.*Neural Networks*, 91:42–54, 2017\.
- Schönemann \(1966\)Peter H\. Schönemann\.A generalized solution of the orthogonal Procrustes problem\.*Psychometrika*, 31\(1\):1–10, 1966\.
- Setayesh et al\. \(2026\)Mehdi Setayesh, Mahdi Beitollahi, Yasser H\. Khalil, and Hongliang Li\.Toward enhancing representation learning in federated multi\-task settings\.In*The Fourteenth International Conference on Learning Representations*, 2026\.URL[https://openreview\.net/forum?id=nIOIfHHYzk](https://openreview.net/forum?id=nIOIfHHYzk)\.
- Shokri et al\. \(2017\)Reza Shokri, Marco Stronati, Congzheng Song, and Vitaly Shmatikov\.Membership inference attacks against machine learning models\.In*IEEE Symposium on Security and Privacy \(S&P\)*, pp\. 3–18\. IEEE, 2017\.
- Singer \(2011\)Amit Singer\.Angular synchronization by eigenvectors and semidefinite programming\.*Applied and Computational Harmonic Analysis*, 30\(1\):20–36, 2011\.
- Singer & Wu \(2012\)Amit Singer and H\-T Wu\.Vector diffusion maps and the connection Laplacian\.*Communications on Pure and Applied Mathematics*, 65\(8\):1067–1144, 2012\.
- Smith et al\. \(2017\)Virginia Smith, Chao\-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar\.Federated multi\-task learning\.*Advances in Neural Information Processing Systems*, 30, 2017\.
- Strinati et al\. \(2024\)Emilio Calvanese Strinati, Paolo Di Lorenzo, et al\.Goal\-oriented and semantic communication in 6G AI\-native networks: The 6G\-GOALS approach\.In*Proceedings of EuCNC/6G Summit*, pp\. 1–6, 2024\.doi:10\.1109/EuCNC/6GSummit60053\.2024\.10597087\.
- Svendsen et al\. \(2026\)Oscar Thorsted Svendsen, Nikolaj Holst Jakobsen, Fabian Mager, and Hiba Nassar\.Improving relative representations with learned anchors and whitened inner products\.*arXiv preprint arXiv:2605\.30596*, 2026\.
- Tan et al\. \(2022\)Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang\.Fedproto: Federated prototype learning across heterogeneous clients\.In*Proceedings of the AAAI Conference on Artificial Intelligence*, 2022\.
- Thunberg et al\. \(2017\)Johan Thunberg, Florian Bernard, and Jorge Gonçalves\.Distributed methods for synchronization of orthogonal matrices over graphs\.*Automatica*, 80:243–252, 2017\.
- Tran et al\. \(2024\)Trung\-Kien Tran, Ha\-Phuong Tran, Thi\-Lan Le, and Thanh\-Hai Tran\.Fedntproto: A prototype\-based approach for personalized federated learning\.In*2024 International Conference on Multimedia Analysis and Pattern Recognition \(MAPR\)*, pp\. 1–6\. IEEE, 2024\.
- Wang & Isola \(2020\)Tongzhou Wang and Phillip Isola\.Understanding contrastive representation learning through alignment and uniformity on the hypersphere\.In*International Conference on Machine Learning*, pp\. 9929–9939\. PMLR, 2020\.
- Zhang et al\. \(2023\)Fengda Zhang, Kun Kuang, Long Chen, Zhaoyang You, Tao Shen, Jun Xiao, Yin Zhang, Chao Wu, Fei Wu, Yueting Zhuang, and Xiaolin Li\.Federated unsupervised representation learning\.*Frontiers of Information Technology & Electronic Engineering*, 24\(8\):1181–1193, 2023\.
- Zhang et al\. \(2021\)Shengdong Zhang, Ehsan Nezhadarya, Homa Fashandi, Jiayi Liu, Darin Graham, and Mohak Shah\.Stochastic whitening batch normalization\.In*Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp\. 10978–10987, 2021\.
- Zhang & Yang \(2017\)Yu Zhang and Qiang Yang\.Learning sparse task relations in multi\-task learning\.In*Proceedings of the AAAI Conference on Artificial Intelligence*, volume 31, 2017\.
- Zhang & Yeung \(2010\)Yu Zhang and Dit\-Yan Yeung\.A convex formulation for learning task relationships in multi\-task learning\.In*Proceedings of the Twenty\-Sixth Conference on Uncertainty in Artificial Intelligence*, pp\. 733–742, 2010\.
- Zhuang et al\. \(2021\)Weiming Zhuang, Xin Gan, Yonggang Wen, Shuai Zhang, and Shuai Yi\.Collaborative unsupervised visual representation learning from decentralized data\.In*Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp\. 4912–4921, 2021\.

## Appendix AAdditional discussion of our framework

This section elaborates on two salient features of our proposed framework, namely the edge stalk dimensionality and the enforced geometric consistency\. We also briefly discuss the homogeneous setting case where all node stalks have the same dimensionality\.

Embedding vs\. compression\.Choosing the edge stalk dimensionalitydi​j=max⁡\(di,dj\)d\_\{ij\}=\\max\(d\_\{i\},d\_\{j\}\)avoids introducing an information bottleneck when aligning heterogeneous latent representations\. In contrast, as done for instance in\(Badi et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib2)\), enforcingdi​j≤min⁡\(di,dj\)d\_\{ij\}\\leq\\min\(d\_\{i\},d\_\{j\}\)would require projecting both representations onto a shared lower\-dimensional subspace, implicitly assuming the existence of a common latent structure and potentially discarding node\-specific information\. Moreover, such a projection only enforces alignment on the retained subspace, leaving the remaining components unconstrained\. As a result, the learned restriction maps capture consistency only up to an information loss, potentially leading to inconsistencies when transferring latent representations across models\. Our choice instead pursues a geometric notion of consistency across representations, and preserves the semantics of local representations\. Importantly, this also makes our framework well\-suited for downstream transfer of latent representations among pre\-trained models, where alignment based on compression may result in performance degradation on downstream local tasks\.

On the geometric consistency\.To better understand the action of the restriction maps as well as the role of the total variation penalty, let us considerℱ​\(i\)→ℱ​\(ei​j\)←ℱ​\(j\)\\mathcal\{F\}\(i\)\\rightarrow\\mathcal\{F\}\(e\_\{ij\}\)\\leftarrow\\mathcal\{F\}\(j\)in[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1)\. When transported to the edge stalkℱ​\(ei​j\)⊆ℝ3\\mathcal\{F\}\(e\_\{ij\}\)\\subseteq\\mathbb\{R\}^\{3\}, the latent representation𝐳i∈ℱ​\(i\)\\mathbf\{z\}\_\{i\}\\in\\mathcal\{F\}\(i\)is rotated and reflected by the orthogonal map𝐎j​i\\mathbf\{O\}\_\{ji\}\. Instead, the latent representation𝐳j\\mathbf\{z\}\_\{j\}is embedded via the Stiefel map𝐕i​j\\mathbf\{V\}\_\{ij\}into the higher\-dimensional edge stalk\. Interestingly, the imageIm​\(𝐕i​j\)\\mathrm\{Im\}\\left\(\\mathbf\{V\}\_\{ij\}\\right\)of𝐕i​j\\mathbf\{V\}\_\{ij\}defines a two\-dimensional subspace embedded inℱ​\(ei​j\)\\mathcal\{F\}\(e\_\{ij\}\)\. Thus, a perfect alignment between the latent representations𝐳i\\mathbf\{z\}\_\{i\}and𝐳j\\mathbf\{z\}\_\{j\}implies that𝐳i\\mathbf\{z\}\_\{i\}is a rotated and reflected version of a vector𝐲j→ei​j=𝐕i​j​𝐳j\\mathbf\{y\}\_\{j\\to e\_\{ij\}\}=\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}belonging toIm​\(𝐕i​j\)\\mathrm\{Im\}\\left\(\\mathbf\{V\}\_\{ij\}\\right\)\. When this occurs for all the latent representations inℱ​\(i\)\\mathcal\{F\}\(i\), it means that the latter representations lie on a linear two\-dimensional subspace ofℱ​\(i\)\\mathcal\{F\}\(i\), as depicted in[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1)\.

In[Fig\.1](https://arxiv.org/html/2608.10016#S3.F1), we notice that there is only a partial alignment between𝐳i\\mathbf\{z\}\_\{i\}and𝐳j\\mathbf\{z\}\_\{j\}when they are transported ontoℱ​\(ei​j\)\\mathcal\{F\}\(e\_\{ij\}\)\. Indeed, even though the projectionProj​\(𝐎j​i​\(𝐳i\)\)\\mathrm\{Proj\}\(\\mathbf\{O\}\_\{ji\}\(\\mathbf\{z\}\_\{i\}\)\)of𝐎j​i​\(𝐳i\)\\mathbf\{O\}\_\{ji\}\(\\mathbf\{z\}\_\{i\}\)ontoIm​\(𝐕i​j\)\\mathrm\{Im\}\\left\(\\mathbf\{V\}\_\{ij\}\\right\)exactly matches𝐕i​j​𝐳j\\mathbf\{V\}\_\{ij\}\\mathbf\{z\}\_\{j\}, there is still an orthogonal component that makes non\-null the Euclidean distance between the two transported vectors\. The latter component can be better visualized inℱ​\(i\)\\mathcal\{F\}\(i\), where we see that𝐳i\\mathbf\{z\}\_\{i\}does not lie on the subspace individuated by the transport of𝐲j→ei​j∈Im​\(𝐕i​j\)\\mathbf\{y\}\_\{j\\to e\_\{ij\}\}\\in\\mathrm\{Im\}\\left\(\\mathbf\{V\}\_\{ij\}\\right\)\. Thus, the total variation term acts as a regularizer that, for allei​j∈ℰe\_\{ij\}\\in\\mathcal\{E\}, enforces the zeroing of components orthogonal toIm​\(𝐕i​j\)\\mathrm\{Im\}\\left\(\\mathbf\{V\}\_\{ij\}\\right\)\.

Notably, the latter discussion highlights that within our proposed framework we do not assume or leverage the existence a*single*global latent representation, as instead pursued in recent works\(Setayesh et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib55); Badi et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib2)\)\. Indeed, in our framework a global representation may eventually emerge from the alignment of local representations over the network sheaf\.

##### Homogeneous setting\.

Whenℱ​\(i\)=ℱ​\(e\)≅ℝd\\mathcal\{F\}\(i\)=\\mathcal\{F\}\(e\)\\cong\\mathbb\{R\}^\{d\}for alli∈𝒩i\\in\\mathcal\{N\}ande∈ℰe\\in\\mathcal\{E\}, the local total variation term in[Eq\.10](https://arxiv.org/html/2608.10016#S4.E10)reduces to

‖𝐳i−𝐎i​j​𝐳j‖22,\\norm\{\\mathbf\{z\}\_\{i\}\-\\mathbf\{O\}\_\{ij\}\\mathbf\{z\}\_\{j\}\}\_\{2\}^\{2\}\\,,\(24\)where𝐎i​j∈O​\(d\)\\mathbf\{O\}\_\{ij\}\\in\\mathrm\{O\}\(\{d\}\)\. In this case, alignment between neighboring representations is achieved up to an orthogonal transformation, i\.e\., rotations and reflections of the latent space\. This setting connects[Eq\.9](https://arxiv.org/html/2608.10016#S4.E9)to the Dirichlet energy of the connection Laplacian\(Chung et al\.,[2014](https://arxiv.org/html/2608.10016#bib.bib10)\), and more generally to synchronization problems over the orthogonal group\(Singer & Wu,[2012](https://arxiv.org/html/2608.10016#bib.bib58)\)\. In particular, the transport maps satisfy𝐎j​i=𝐎i​j⊤\\mathbf\{O\}\_\{ji\}=\\mathbf\{O\}\_\{ij\}^\{\\top\}, and the two variation terms in[Eq\.12](https://arxiv.org/html/2608.10016#S4.E12)coincide\. In contrast, in the heterogeneous setting we allow semi\-orthogonal maps on the Stiefel manifold, which act as isometric embeddings and are only left\-invertible\. Finally, since all node stalks share the same dimension, the orientation of𝒢\\mathcal\{G\}does not affect the model and can be chosen arbitrarily\.

## Appendix BPilot Selection Strategies

The alignment of latent spaces relies on a set of*semantic pilots*: inputs whose correspondence across agents is known a priori, so thatf𝝋i​\(𝐱ik\)f\_\{\\bm\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)andf𝝋j​\(𝐱jk\)f\_\{\\bm\{\\varphi\}\_\{j\}\}\(\\mathbf\{x\}\_\{j\}^\{k\}\)can be treated as two latent representations of the same underlying sample\. In practice, pilots arise from public datasets, calibration signals, standardized inputs, or an overlapping portion of the local datasets, and we assume that a common index set\{1,…,M\}\\\{1,\\ldots,M\\\}has been agreed upon\.

The choice of𝒜⊂\{1,…,M\}\\mathcal\{A\}\\subset\\\{1,\\ldots,M\\\}, with\|𝒜\|=K≪M\|\\mathcal\{A\}\|=K\\ll M, governs the trade\-off between communication cost, computational efficiency, geometric coverage, and alignment quality\. In the sequel, we outline the most widely used practical selection criteria\. Additionally, a normalization step, applied to the resulting pilot matrices for balancing directions in their corresponding latent spaces, is discussed at the end of the section\.

Original pilots\.Each agent transmits the embeddingsf𝝋i​\(𝐱ik\)f\_\{\\bm\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)of the shared inputs themselves, wherek∈𝒜k\\in\\mathcal\{A\}\. Taking𝒜=\{1,…,M\}\\mathcal\{A\}=\\\{1,\\ldots,M\\\}evaluates the penalty exactly and is the most reliable option, but scales linearly inMMin both communication and Procrustes cost; subsampling to a budgetK≪MK\\ll Minstead yields an unbiased Monte Carlo estimator of the full alignment energy\.

*Unsupervised selection:*Pilots are drawn uniformly, which requires no side information but may under\-represent low\-density regions of the representation space and lead to ill\-conditioned𝐀i\\mathbf\{A\}\_\{i\}\. Clustering the representation space and drawing a fixed number of pilots per cluster recovers the same stratification that labels provide in the supervised case, at the cost of one clustering pass, and guarantees that sparsely populated regions are represented\.

*Supervised selection:*Labels are exploited since they induce a natural stratification, and a fixed number of pilots is drawn per class, guaranteeing balanced coverage even for smallKK\.

Prototype pilots\.The shared set is partitioned intoKKgroups and each pilot is defined as the group average \(centroid\),

𝐟¯i,c=1\|𝒮c\|​∑k∈𝒮cf𝝋i​\(𝐱ik\),c=1,…,K;\\bar\{\\mathbf\{f\}\}\_\{i,c\}=\\frac\{1\}\{\|\\mathcal\{S\}\_\{c\}\|\}\\sum\_\{k\\in\\mathcal\{S\}\_\{c\}\}f\_\{\\bm\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\),\\qquad c=1,\\ldots,K\\,;which preserves semantic structure while averaging out per\-sample noise\. When groups are large, the centroid need not be computed over all their members: averaging over a subset𝒮c′⊂𝒮c\\mathcal\{S\}\_\{c\}^\{\\prime\}\\subset\\mathcal\{S\}\_\{c\}gives an unbiased estimate at a fraction of the local forward passes, trading estimation variance for computational cost\.

*Unsupervised partitioning:*The partition is obtained by clustering the representation space \(e\.g\., viakk\-means\), with the resulting index assignment broadcast once and reused thereafter\(Hüttebräucker et al\.,[2024](https://arxiv.org/html/2608.10016#bib.bib31); Fiorellino et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib24)\)\.

*Supervised partitioning:*Class labels provide the partition directly, recovering prototype\-based federated schemes such as FedProto\(Tan et al\.,[2022](https://arxiv.org/html/2608.10016#bib.bib62)\)\.

Ensuring correspondence and adaptive selection\.To ensure that the transmitted quantities are in correspondence, both original and prototype pilots require the selection to be replicated across agents\. Specifically:*\(i\)*the sampled indices in the first case, and*\(ii\)*the partition and the subsets used to estimate the centroids in the second\. Furthermore, the selection can be made adaptive by prioritizing the samples that currently exhibit the largest cross\-client mismatch, biasing which inputs are transmitted or which groups are refined so as to concentrate the alignment effort on geometrically inconsistent regions\.

Dynamic selection\.Any of the above criteria can be applied per iteration, letting𝒜t\\mathcal\{A\}\_\{t\}vary over time and turning the penalty into a stochastic approximation of the full alignment energy\. A practical recipe combines a small persistent subset \(for stability of the transport maps\) with a periodically refreshed remainder \(for coverage\), which empirically improves generalization of the learned maps beyond the sampled pilots\(Fiorellino et al\.,[2024](https://arxiv.org/html/2608.10016#bib.bib23)\)\.

Parseval normalization\.Independently of how𝒜\\mathcal\{A\}is selected, the conditioning of the Procrustes step depends on the geometry of the pilot matrix𝐀i=\[f𝝋i​\(𝐱ik\)\]k∈𝒜∈ℝdi×K\\mathbf\{A\}\_\{i\}=\\big\[f\_\{\\bm\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\big\]\_\{k\\in\\mathcal\{A\}\}\\in\\mathbb\{R\}^\{d\_\{i\}\\times K\}, withK<diK<d\_\{i\}in the regime of interest\. To mitigate degeneracy, we can normalize pilots via

𝐀~i=𝐀i​\(𝐀i⊤​𝐀i\)−1/2,\\widetilde\{\\mathbf\{A\}\}\_\{i\}=\\mathbf\{A\}\_\{i\}\(\\mathbf\{A\}\_\{i\}^\{\\top\}\\mathbf\{A\}\_\{i\}\)^\{\-1/2\}\\,,the orthogonal polar factor of𝐀i\\mathbf\{A\}\_\{i\}: this enforces𝐀~i⊤​𝐀~i=𝐈K\\widetilde\{\\mathbf\{A\}\}\_\{i\}^\{\\top\}\\widetilde\{\\mathbf\{A\}\}\_\{i\}=\\mathbf\{I\}\_\{K\}, i\.e\., a Parseval frame of the pilot subspace\(Fiorellino et al\.,[2026](https://arxiv.org/html/2608.10016#bib.bib24)\)\. This transformation balances directions in latent space and stabilizes the Procrustes updates\.

## Appendix CIdentifiability of the Transport Maps

The edge\-wise transports are computed via \(semi\-\)orthogonal Procrustes problems

𝐕i​j∈arg⁡min𝐕∈𝒞i​j⁡‖𝐀i−𝐕𝐀j‖F2,\\mathbf\{V\}\_\{ij\}\\in\\arg\\min\_\{\\mathbf\{V\}\\in\\mathcal\{C\}\_\{ij\}\}\\left\\\|\\mathbf\{A\}\_\{i\}\-\\mathbf\{V\}\\mathbf\{A\}\_\{j\}\\right\\\|\_\{F\}^\{2\},\(25\)where𝒞i​j=O​\(d\)\\mathcal\{C\}\_\{ij\}=\\mathrm\{O\}\(\{d\}\)in the homogeneous case and𝒞i​j=St​\(di,dj\)\\mathcal\{C\}\_\{ij\}=\\mathrm\{St\}\(d\_\{i\},d\_\{j\}\)fordi\>djd\_\{i\}\>d\_\{j\}\. We summarize conditions under which the minimizer is unique\.

Homogeneous case \(di=dj=dd\_\{i\}=d\_\{j\}=d\)\.Let𝐌i​j≔𝐀i​𝐀j⊤∈ℝd×d\\mathbf\{M\}\_\{ij\}\\coloneqq\\mathbf\{A\}\_\{i\}\\mathbf\{A\}\_\{j\}^\{\\top\}\\in\\mathbb\{R\}^\{d\\times d\}\. An optimal solution is𝐕i​j⋆=𝐔𝐖⊤\\mathbf\{V\}\_\{ij\}^\{\\star\}=\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}, where𝐌i​j=𝐔​𝚺​𝐖⊤\\mathbf\{M\}\_\{ij\}=\\mathbf\{U\}\\mathbf\{\\Sigma\}\\mathbf\{W\}^\{\\top\}is an SVD\(Schönemann,[1966](https://arxiv.org/html/2608.10016#bib.bib54)\)\. The solution is unique if𝐌i​j\\mathbf\{M\}\_\{ij\}is full rank and has simple singular values\. Non\-uniqueness arises only when𝐌i​j\\mathbf\{M\}\_\{ij\}is rank\-deficient or has repeated singular values, which induce rotational ambiguity within singular subspaces\. In particular, ifrank​\(𝐀i\)=rank​\(𝐀j\)=d\\mathrm\{rank\}\(\\mathbf\{A\}\_\{i\}\)=\\mathrm\{rank\}\(\\mathbf\{A\}\_\{j\}\)=d\(e\.g\.,K≥dK\\geq dwith non\-degenerate pilots\), then𝐌i​j\\mathbf\{M\}\_\{ij\}is generically full rank and𝐕i​j⋆\\mathbf\{V\}\_\{ij\}^\{\\star\}is uniquely defined up to measure\-zero degeneracies\.

Stiefel case \(di\>djd\_\{i\}\>d\_\{j\}\)\.Let𝐌i​j=𝐀i​𝐀j⊤∈ℝdi×dj\\mathbf\{M\}\_\{ij\}=\\mathbf\{A\}\_\{i\}\\mathbf\{A\}\_\{j\}^\{\\top\}\\in\\mathbb\{R\}^\{d\_\{i\}\\times d\_\{j\}\}and𝐌i​j=𝐔​𝚺​𝐖⊤\\mathbf\{M\}\_\{ij\}=\\mathbf\{U\}\\mathbf\{\\Sigma\}\\mathbf\{W\}^\{\\top\}a thin SVD, with𝐔∈ℝdi×dj\\mathbf\{U\}\\in\\mathbb\{R\}^\{d\_\{i\}\\times d\_\{j\}\},𝐖∈ℝdj×dj\\mathbf\{W\}\\in\\mathbb\{R\}^\{d\_\{j\}\\times d\_\{j\}\}\. A minimizer is

𝐕i​j⋆=𝐔𝐖⊤∈St​\(di,dj\)\.\\mathbf\{V\}\_\{ij\}^\{\\star\}=\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}\\in\\mathrm\{St\}\(d\_\{i\},d\_\{j\}\)\\,\.\(26\)Uniqueness holds ifrank​\(𝐌i​j\)=dj\\mathrm\{rank\}\(\\mathbf\{M\}\_\{ij\}\)=d\_\{j\}and its singular values are simple\. Crucially,𝐕i​j\\mathbf\{V\}\_\{ij\}is identifiable only on the subspace spanned by𝐀j\\mathbf\{A\}\_\{j\}: ifrank​\(𝐀j\)=r<dj\\mathrm\{rank\}\(\\mathbf\{A\}\_\{j\}\)=r<d\_\{j\}, the solution is underdetermined on the orthogonal complement\. Hence, identifiability requiresrank​\(𝐀j\)=dj\\mathrm\{rank\}\(\\mathbf\{A\}\_\{j\}\)=d\_\{j\}\(typically ensured whenK≥djK\\geq d\_\{j\}and pilots provide sufficient geometric coverage\), in which case𝐕i​j⋆\\mathbf\{V\}\_\{ij\}^\{\\star\}is generically unique up to measure\-zero degeneracies\.

## Appendix DPseudocode of the Sheaf\-FRL Algorithm

In this section, we provide additional details on the proposed decentralized algorithm, reported in Alg\.[1](https://arxiv.org/html/2608.10016#alg1)\. The method follows an alternating minimization strategy, where each communication round consists of local feature extraction, decentralized alignment updates, and gradient\-based optimization of the neural parameters\.

Algorithm 1Sheaf\-FRL0:Oriented

𝒢=\(𝒩,ℰ\)\\mathcal\{G\}=\(\\mathcal\{N\},\\mathcal\{E\}\), pilots

𝒜\\mathcal\{A\},

λ\>0\\lambda\>0,

ηθ\>0\\eta\_\{\\theta\}\>0, iterations

TT
1:Initialize

\{𝜽i0=\(𝝋i0,𝝍i0\)\}i∈𝒩\\\{\\boldsymbol\{\\theta\}\_\{i\}^\{0\}=\(\\boldsymbol\{\\varphi\}\_\{i\}^\{0\},\\boldsymbol\{\\psi\}\_\{i\}^\{0\}\)\\\}\_\{i\\in\\mathcal\{N\}\}
2:for

t=0,…,T−1t=0,\\dots,T\-1do

3:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

4:

𝐀it←\[f𝝋it​\(𝐱ik\)\]k∈𝒜\\mathbf\{A\}\_\{i\}^\{t\}\\leftarrow\[f\_\{\\boldsymbol\{\\varphi\}\_\{i\}^\{t\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\]\_\{k\\in\\mathcal\{A\}\}
5:endfor

6:Each node

iibroadcasts

𝐀it\\mathbf\{A\}\_\{i\}^\{t\}to those

j∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)
7:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

8:for

j∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)do

9:if

j∈𝒩​\(i\)−j\\in\\mathcal\{N\}\(i\)^\{\-\}then

10:

\[𝐔,𝚺,𝐖⊤\]←thinSVD​\(𝐀it​𝐀jt⊤\)\[\\mathbf\{U\},\\mathbf\{\\Sigma\},\\mathbf\{W\}^\{\\top\}\]\\leftarrow\\mathrm\{thinSVD\}\(\\mathbf\{A\}\_\{i\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t^\{\\top\}\}\)
11:

𝐕i​jt←𝐔𝐖⊤\\mathbf\{V\}\_\{ij\}^\{t\}\\leftarrow\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}
12:else

13:

\[𝐔,𝚺,𝐖⊤\]←thinSVD​\(𝐀jt​𝐀it⊤\)\[\\mathbf\{U\},\\mathbf\{\\Sigma\},\\mathbf\{W\}^\{\\top\}\]\\leftarrow\\mathrm\{thinSVD\}\(\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t^\{\\top\}\}\)
14:

𝐕j​it←𝐔𝐖⊤\\mathbf\{V\}\_\{ji\}^\{t\}\\leftarrow\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}
15:endif

16:endfor

17:endfor

18:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

19:

𝜽it\+1←𝜽it−ηθ​\(∇𝜽iℒi​\(𝜽it\)\+𝐫it\)\\boldsymbol\{\\theta\}\_\{i\}^\{t\+1\}\\leftarrow\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\-\\eta\_\{\\theta\}\\left\(\\nabla\_\{\\boldsymbol\{\\theta\}\_\{i\}\}\\mathcal\{L\}\_\{i\}\(\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\)\+\\mathbf\{r\}^\{t\}\_\{i\}\\right\)
20:endfor

21:endfor

##### Heterogeneous case\.

Alg\.[1](https://arxiv.org/html/2608.10016#alg1)implements a fully decentralized procedure in which agents iteratively update both their local representations and the alignment maps across the network\. At each iterationtt, every agentiifirst computes its pilot feature matrix𝐀it\\mathbf\{A\}\_\{i\}^\{t\}using the current encoder \(lines 3–5\)\. These pilot representations provide a compact summary of the local latent space and constitute the only information exchanged across agents\. The pilot features are then broadcast to neighboring nodes \(line 6\)\. Upon receiving\{𝐀jt\}j∈𝒩​\(i\)\\\{\\mathbf\{A\}\_\{j\}^\{t\}\\\}\_\{j\\in\\mathcal\{N\}\(i\)\}, each agent locally updates the alignment maps associated with its incident edges \(lines 7–17\)\. In particular, for each neighborjj, agentiisolves a small Procrustes problem based on the cross\-covariance of the corresponding pilot features\. The orientation of the graph determines whether the update corresponds to an incoming map𝐕i​jt\\mathbf\{V\}\_\{ij\}^\{t\}or an outgoing map𝐕j​it\\mathbf\{V\}\_\{ji\}^\{t\}, ensuring consistency with the underlying sheaf structure\. Finally, each agent performs a gradient\-based update of its local parameters𝜽i\\boldsymbol\{\\theta\}\_\{i\}\(lines 18–20\), combining the gradient of the local loss with the contribution of the sheaf\-based regularization\. This step enforces alignment of neighboring representations while preserving local task objectives\. Overall, the algorithm alternates between closed\-form alignment updates and gradient\-based parameter updates\. All computations are carried out locally, and communication is limited to low\-dimensional pilot features, resulting in an efficient and scalable decentralized learning procedure\.

Algorithm 2Sheaf\-FRL \(homogeneous case\)0:Oriented

𝒢=\(𝒩,ℰ\)\\mathcal\{G\}=\(\\mathcal\{N\},\\mathcal\{E\}\), pilots

𝒜\\mathcal\{A\},

λ\>0\\lambda\>0,

ηθ\>0\\eta\_\{\\theta\}\>0, iterations

TT
1:Initialize

\{𝜽i0=\(𝝋i0,𝝍i0\)\}i∈𝒩\\\{\\boldsymbol\{\\theta\}\_\{i\}^\{0\}=\(\\boldsymbol\{\\varphi\}\_\{i\}^\{0\},\\boldsymbol\{\\psi\}\_\{i\}^\{0\}\)\\\}\_\{i\\in\\mathcal\{N\}\}and

\{𝐎j​i0=𝐈d\}j∈𝒩​\(i\),i∈𝒩\\\{\\mathbf\{O\}\_\{ji\}^\{0\}=\\mathbf\{I\}\_\{d\}\\\}\_\{j\\in\\mathcal\{N\}\(i\),i\\in\\mathcal\{N\}\}
2:for

t=0,…,T−1t=0,\\dots,T\-1do

3:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

4:

𝐀it←\[f𝝋it​\(𝐱ik\)\]k∈𝒜\\mathbf\{A\}\_\{i\}^\{t\}\\leftarrow\[f\_\{\\boldsymbol\{\\varphi\}\_\{i\}^\{t\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\]\_\{k\\in\\mathcal\{A\}\}
5:endfor

6:Each node

iiexchanges

𝐎j​it​𝐀it\\mathbf\{O\}\_\{ji\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t\}with all

j∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)
7:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

8:for all

j∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}do

9:

\[𝐔,𝚺,𝐖⊤\]←SVD​\(𝐎i​jt​𝐀jt​𝐀it⊤\)\[\\mathbf\{U\},\\mathbf\{\\Sigma\},\\mathbf\{W\}^\{\\top\}\]\\leftarrow\\mathrm\{SVD\}\\big\(\\mathbf\{O\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t^\{\\top\}\}\\big\)
10:

𝐎j​it←𝐔𝐖⊤\\mathbf\{O\}\_\{ji\}^\{t\}\\leftarrow\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}
11:endfor

12:endfor

13:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

14:

𝜽it\+1←𝜽it−ηθ​\(∇𝜽iℒi​\(𝜽it\)\+𝐫it\)\\boldsymbol\{\\theta\}\_\{i\}^\{t\+1\}\\leftarrow\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\-\\eta\_\{\\theta\}\\left\(\\nabla\_\{\\boldsymbol\{\\theta\}\_\{i\}\}\\mathcal\{L\}\_\{i\}\(\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\)\+\\mathbf\{r\}\_\{i\}^\{t\}\\right\)
15:endfor

16:endfor

Homogeneous case\.When the latent dimensions are homogeneous, i\.e\.,di=dd\_\{i\}=dfor alli∈𝒩i\\in\\mathcal\{N\}, the alignment maps reduce to orthogonal transformations𝐎i​j∈O​\(d\)\\mathbf\{O\}\_\{ij\}\\in\\mathrm\{O\}\(\{d\}\), for each\(i,j\)∈ℰ\(i,j\)\\in\\mathcal\{E\}\. In this setting, as discussed in[Remark4](https://arxiv.org/html/2608.10016#Thmremark4), the local computational cost of the alignment step can be reduced without increasing communication\. Alg\.[2](https://arxiv.org/html/2608.10016#alg2)presents a simplified Sheaf\-FRL protocol tailored to this scenario\. For each agentii, we initialize to the identity the maps used to transport its pilot representations toward its neighbors, i\.e\.,𝐎j​i0=𝐈d\\mathbf\{O\}\_\{ji\}^\{0\}=\\mathbf\{I\}\_\{d\}for allj∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)\. At each iterationtt, each agent first computes its pilot representations𝐀it\\mathbf\{A\}\_\{i\}^\{t\}and sends to each neighborj∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)the transported features𝐎j​it​𝐀it\\mathbf\{O\}\_\{ji\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t\}\. Notably, since we exploit the reparameterization in[Eq\.24](https://arxiv.org/html/2608.10016#A1.E24), whenj∈𝒩−​\(i\)j\\in\\mathcal\{N\}^\{\-\}\(i\), we have𝐎j​it=𝐎j​i0=𝐈d\\mathbf\{O\}\_\{ji\}^\{t\}=\\mathbf\{O\}\_\{ji\}^\{0\}=\\mathbf\{I\}\_\{d\}for allt=0,…,T−1t=0,\\ldots,T\-1, so that the transmitted features coincide with the original ones and no additional communication cost is incurred\. From a local computation perspective, since transported pilot representations are communicated instead of raw ones, agentiidoes not need to solve the orthogonal Procrustes problem arising in \([P2a](https://arxiv.org/html/2608.10016#S5.Ex5)\) for evaluating the alignment and the gradient required for the local parameter update\. As a consequence, only the SVDs forj∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}are required, effectively halving the computational cost of the alignment step, as specified in lines 8\-10 of Alg\.[2](https://arxiv.org/html/2608.10016#alg2)\. Overall, this variant preserves the communication efficiency of the general algorithm while reducing the computational burden \(cf\.[Tab\.1](https://arxiv.org/html/2608.10016#A6.T1)\)\.

## Appendix EConvergence Analysis

In this section we establish convergence guarantees for the proposed Sheaf\-FRL algorithm in both deterministic and stochastic settings\. We first recall the classical descent result for smooth functions\.

###### Lemma E\.1\(Descent Lemma,Nesterov,[2004](https://arxiv.org/html/2608.10016#bib.bib49), Prop\. 1\.2\.3,Beck,[2017](https://arxiv.org/html/2608.10016#bib.bib6), Sec\. 2\.1\)\.

Letffbe differentiable withLL\-Lipschitz gradient\. Then for any𝐱\\mathbf\{x\}and anyη\>0\\eta\>0,

f​\(𝐱−η​∇f​\(𝐱\)\)≤f​\(𝐱\)−η​\(1−L​η2\)​‖∇f​\(𝐱\)‖2\.f\(\\mathbf\{x\}\-\\eta\\nabla f\(\\mathbf\{x\}\)\)\\leq f\(\\mathbf\{x\}\)\-\\eta\\Big\(1\-\\frac\{L\\eta\}\{2\}\\Big\)\\\|\\nabla f\(\\mathbf\{x\}\)\\\|^\{2\}\.

### E\.1Deterministic Setting

Let

J​\(𝜽,𝐕\)=∑i∈𝒩ℒi​\(𝜽i\)\+ℛ𝒜​\(\{𝝋i\}i∈𝒩,\{𝐕i​j\}\(i,j\)∈ℰ,\{𝐕j​i\}\(i,j\)∈ℰ\)J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}\)=\\sum\_\{i\\in\\mathcal\{N\}\}\\mathcal\{L\}\_\{i\}\(\\bm\{\\theta\}\_\{i\}\)\+\\mathcal\{R\}\_\{\\mathcal\{A\}\}\\\!\\left\(\\\{\\bm\{\\varphi\}\_\{i\}\\\}\_\{i\\in\\mathcal\{N\}\},\\\{\\mathbf\{V\}\_\{ij\}\\\}\_\{\(i,j\)\\in\\mathcal\{E\}\},\\\{\\mathbf\{V\}\_\{ji\}\\\}\_\{\(i,j\)\\in\\mathcal\{E\}\}\\right\)\(27\)in \([SFRL](https://arxiv.org/html/2608.10016#S4.Ex3)\) where, with slight abuse of notation,𝜽=\{𝜽i\}i∈𝒩=\{\(𝝋i,𝝍i\)\}i∈𝒩\\boldsymbol\{\\theta\}=\\\{\\boldsymbol\{\\theta\}\_\{i\}\\\}\_\{i\\in\\mathcal\{N\}\}=\\\{\(\\bm\{\\varphi\}\_\{i\},\\bm\{\\psi\}\_\{i\}\)\\\}\_\{i\\in\\mathcal\{N\}\}, and𝐕=\{𝐕i​j\}\(i,j\)∈ℰ∪\{𝐕j​i\}\(i,j\)∈ℰ\\mathbf\{V\}=\\\{\\mathbf\{V\}\_\{ij\}\\\}\_\{\(i,j\)\\in\\mathcal\{E\}\}\\cup\\\{\\mathbf\{V\}\_\{ji\}\\\}\_\{\(i,j\)\\in\\mathcal\{E\}\}\. For each edge\(i,j\)∈ℰ\(i,j\)\\in\\mathcal\{E\}, the matrices𝐕i​j∈St​\(di,dj\)\\mathbf\{V\}\_\{ij\}\\in\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)and𝐕j​i∈St​\(dj,di\)\\mathbf\{V\}\_\{ji\}\\in\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\), whereSt​\(di,dj\)\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)denotes a Stiefel manifold as defined in[Eq\.4](https://arxiv.org/html/2608.10016#S3.E4)\. The feasible set is therefore𝒱:=\(∏\(i,j\)∈ℰSt​\(di,dj\)\)×\(∏\(i,j\)∈ℰSt​\(dj,di\)\)\\mathcal\{V\}:=\\left\(\\prod\_\{\(i,j\)\\in\\mathcal\{E\}\}\\mathrm\{St\}\(\{d\_\{i\}\},\{d\_\{j\}\}\)\\right\)\\times\\left\(\\prod\_\{\(i,j\)\\in\\mathcal\{E\}\}\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\)\\right\)\. The proposed method alternates between an exact minimization with respect to𝐕\\mathbf\{V\}and a gradient step with respect to𝜽\\boldsymbol\{\\theta\}, and can be compactly cast as:

𝐕t\\displaystyle\\mathbf\{V\}^\{t\}∈arg⁡min𝐕∈𝒱⁡J​\(𝜽t,𝐕\),\\displaystyle\\in\\arg\\min\_\{\\mathbf\{V\}\\in\\mathcal\{V\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}\),\(28\)𝜽t\+1\\displaystyle\\boldsymbol\{\\theta\}^\{t\+1\}=𝜽t−η​∇𝜽J​\(𝜽t,𝐕t\)\\displaystyle=\\boldsymbol\{\\theta\}^\{t\}\-\\eta\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\(29\)We introduce the following assumptions\.

###### Assumption 1\.

The objectiveJ​\(𝛉,𝐕\)J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}\)is bounded from below, with lower bound

Jinf≔inf𝜽,𝐕J​\(𝜽,𝐕\)\.J\_\{\\inf\}\\coloneqq\\inf\_\{\\boldsymbol\{\\theta\},\\mathbf\{V\}\}J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}\)\\,\.

###### Assumption 2\.

For every fixed𝐕∈𝒱\\mathbf\{V\}\\in\\mathcal\{V\},J​\(𝛉,𝐕\)J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}\)is differentiable and its gradient isLL\-Lipschitz:

‖∇𝜽J​\(𝜽1,𝐕\)−∇𝜽J​\(𝜽2,𝐕\)‖≤L​‖𝜽1−𝜽2‖\.\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}\_\{1\},\\mathbf\{V\}\)\-\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}\_\{2\},\\mathbf\{V\}\)\\\|\\leq L\\\|\\boldsymbol\{\\theta\}\_\{1\}\-\\boldsymbol\{\\theta\}\_\{2\}\\\|\.

###### Assumption 3\.

At each iteration,

𝐕t∈arg⁡min𝐕∈𝒱⁡J​\(𝜽t,𝐕\),\\mathbf\{V\}^\{t\}\\in\\arg\\min\_\{\\mathbf\{V\}\\in\\mathcal\{V\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}\),i\.e\.,𝐕t\\mathbf\{V\}^\{t\}is an exact minimizer ofJJgiven𝛉t\\boldsymbol\{\\theta\}^\{t\}\.

###### Assumption 4\.

The feasible set𝒱\\mathcal\{V\}is compact\.

Under[1](https://arxiv.org/html/2608.10016#Thmassumption1)\-[4](https://arxiv.org/html/2608.10016#Thmassumption4), the following convergence result holds\.

###### Theorem E\.2\.

Suppose Assumptions[1](https://arxiv.org/html/2608.10016#Thmassumption1)–[4](https://arxiv.org/html/2608.10016#Thmassumption4)hold and choose0<η≤1/L0<\\eta\\leq 1/L\. Then:

1. 1\.Descent, Convergence, and Stationarity\.The sequence\{J​\(𝜽t,𝐕t\)\}\\\{J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\}is monotonically decreasing and therefore convergent\. Moreover,limt→∞‖∇𝜽J​\(𝜽t,𝐕t\)‖=0\\lim\_\{t\\to\\infty\}\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|=0, and every accumulation point\(𝜽⋆,𝐕⋆\)\(\\boldsymbol\{\\theta\}^\{\\star\},\\mathbf\{V\}^\{\\star\}\)is a block\-wise first\-order stationary point\.
2. 2\.Sublinear stationarity rate\.For allT≥0T\\geq 0, min0≤t≤T⁡‖∇𝜽J​\(𝜽t,𝐕t\)‖2≤2​\(J​\(𝜽0,𝐕0\)−Jinf\)η​\(T\+1\)\.\\min\_\{0\\leq t\\leq T\}\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|^\{2\}\\leq\\frac\{2\\big\(J\(\\boldsymbol\{\\theta\}^\{0\},\\mathbf\{V\}^\{0\}\)\-J\_\{\\inf\}\\big\)\}\{\\eta\(T\+1\)\}\\,\.\(30\)

###### Proof\.

\(i\) Descent, Convergence, and Stationarity\. Fix𝐕t\\mathbf\{V\}^\{t\}and considerJ​\(𝜽,𝐕t\)\.J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}^\{t\}\)\.By Assumption[2](https://arxiv.org/html/2608.10016#Thmassumption2),J​\(𝜽,𝐕t\)J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}^\{t\}\)hasLL\-Lipschitz gradient\. Applying Lemma[E\.1](https://arxiv.org/html/2608.10016#A5.Thmtheorem1)withη≤1/L\\eta\\leq 1/Lyields

J​\(𝜽t\+1,𝐕t\)≤J​\(𝜽t,𝐕t\)−η2​‖∇𝜽J​\(𝜽t,𝐕t\)‖2\.J\(\\boldsymbol\{\\theta\}^\{t\+1\},\\mathbf\{V\}^\{t\}\)\\leq J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\-\\frac\{\\eta\}\{2\}\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|^\{2\}\.\(31\)By exact minimization of𝐕\\mathbf\{V\}at iterationt\+1t\+1\(Assumption[3](https://arxiv.org/html/2608.10016#Thmassumption3)\),J​\(𝜽t\+1,𝐕t\+1\)≤J​\(𝜽t\+1,𝐕t\)J\(\\boldsymbol\{\\theta\}^\{t\+1\},\\mathbf\{V\}^\{t\+1\}\)\\leq J\(\\boldsymbol\{\\theta\}^\{t\+1\},\\mathbf\{V\}^\{t\}\)\. Combining this inequality with the previous bound yieldsJ​\(𝜽t\+1,𝐕t\+1\)≤J​\(𝜽t,𝐕t\)J\(\\boldsymbol\{\\theta\}^\{t\+1\},\\mathbf\{V\}^\{t\+1\}\)\\leq J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\), which shows that the sequence\{J​\(𝜽t,𝐕t\)\}\\\{J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\}is monotonically decreasing\. SinceJJis bounded below \(Assumption[1](https://arxiv.org/html/2608.10016#Thmassumption1)\), the sequence therefore converges\.

Summing \([31](https://arxiv.org/html/2608.10016#A5.E31)\) overttyields∑t=0∞‖∇𝜽J​\(𝜽t,𝐕t\)‖2<∞,\\sum\_\{t=0\}^\{\\infty\}\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|^\{2\}<\\infty,which implies

limt→∞‖∇𝜽J​\(𝜽t,𝐕t\)‖=0\.\\lim\_\{t\\to\\infty\}\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|=0\.Let\(𝜽tk,𝐕tk\)\(\\boldsymbol\{\\theta\}^\{t\_\{k\}\},\\mathbf\{V\}^\{t\_\{k\}\}\)be a convergent subsequence\. By continuity of∇𝜽J\\nabla\_\{\\boldsymbol\{\\theta\}\}J, it holds∇𝜽J​\(𝜽⋆,𝐕⋆\)=0\.\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{\\star\},\\mathbf\{V\}^\{\\star\}\)=0\.Moreover, since𝐕t\\mathbf\{V\}^\{t\}minimizesJ​\(𝜽t,⋅\)J\(\\boldsymbol\{\\theta\}^\{t\},\\cdot\)over the compact set𝒱\\mathcal\{V\}, passing to the limit gives

𝐕⋆∈arg⁡min𝐕∈𝒱⁡J​\(𝜽⋆,𝐕\)\.\\mathbf\{V\}^\{\\star\}\\in\\arg\\min\_\{\\mathbf\{V\}\\in\\mathcal\{V\}\}J\(\\boldsymbol\{\\theta\}^\{\\star\},\\mathbf\{V\}\)\.Thus every accumulation point is block\-wise first\-order stationary\.

\(ii\) Sublinear Stationarity Rate\.From the descent inequality \([31](https://arxiv.org/html/2608.10016#A5.E31)\), summing fromt=0t=0toTTand telescoping yields

η2​∑t=0T‖∇𝜽J​\(𝜽t,𝐕t\)‖2≤J​\(𝜽0,𝐕0\)−J​\(𝜽T\+1,𝐕T\)\.\\frac\{\\eta\}\{2\}\\sum\_\{t=0\}^\{T\}\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|^\{2\}\\leq J\(\\boldsymbol\{\\theta\}^\{0\},\\mathbf\{V\}^\{0\}\)\-J\(\\boldsymbol\{\\theta\}^\{T\+1\},\\mathbf\{V\}^\{T\}\)\.SinceJ​\(𝜽T\+1,𝐕T\)≥JinfJ\(\\boldsymbol\{\\theta\}^\{T\+1\},\\mathbf\{V\}^\{T\}\)\\geq J\_\{\\inf\}\(Assumption[1](https://arxiv.org/html/2608.10016#Thmassumption1)\),

∑t=0T‖∇𝜽J​\(𝜽t,𝐕t\)‖2≤2​\(J​\(𝜽0,𝐕0\)−Jinf\)η\.\\sum\_\{t=0\}^\{T\}\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|^\{2\}\\leq\\frac\{2\\big\(J\(\\boldsymbol\{\\theta\}^\{0\},\\mathbf\{V\}^\{0\}\)\-J\_\{\\inf\}\\big\)\}\{\\eta\}\.Finally, usingmin0≤t≤T⁡at≤1T\+1​∑t=0Tat\\min\_\{0\\leq t\\leq T\}a\_\{t\}\\leq\\frac\{1\}\{T\+1\}\\sum\_\{t=0\}^\{T\}a\_\{t\}withat=‖∇𝜽J​\(𝜽t,𝐕t\)‖2a\_\{t\}=\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|^\{2\}, we obtain

min0≤t≤T⁡‖∇𝜽J​\(𝜽t,𝐕t\)‖2≤2​\(J​\(𝜽0,𝐕0\)−Jinf\)η​\(T\+1\),\\min\_\{0\\leq t\\leq T\}\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\|^\{2\}\\leq\\frac\{2\\big\(J\(\\boldsymbol\{\\theta\}^\{0\},\\mathbf\{V\}^\{0\}\)\-J\_\{\\inf\}\\big\)\}\{\\eta\(T\+1\)\},which proves \([30](https://arxiv.org/html/2608.10016#A5.E30)\)\. ∎

### E\.2Stochastic Setting

We consider the same deterministic objectiveJ​\(𝜽,𝐕\)J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}\)in \([27](https://arxiv.org/html/2608.10016#A5.E27)\), and extend the𝜽\\boldsymbol\{\\theta\}\-update to the stochastic \(mini\-batch\) setting while keeping the𝐕\\mathbf\{V\}\-block*exactly*minimized with respect to the full objective\. Let\{Φt\}t≥0\\\{\\Phi\_\{t\}\\\}\_\{t\\geq 0\}be the natural filtration generated by the iterates and the sampling randomness up to timett\. At iterationtt, we first update the alignment variables by exact minimization,

𝐕t∈arg⁡min𝐕∈𝒱⁡J​\(𝜽t,𝐕\),\\mathbf\{V\}^\{t\}\\in\\arg\\min\_\{\\mathbf\{V\}\\in\\mathcal\{V\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}\),\(32\)and then form a mini\-batch stochastic gradient estimator𝐠t≡𝐠​\(𝜽t,𝐕t;ξt\),\\mathbf\{g\}^\{t\}\\equiv\\mathbf\{g\}\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\};\\xi\_\{t\}\),whereξt\\xi\_\{t\}denotes the mini\-batch sampling randomness\. The stochastic update of the neural parameters is

𝜽t\+1=𝜽t−ηt​𝐠t\.\\boldsymbol\{\\theta\}^\{t\+1\}=\\boldsymbol\{\\theta\}^\{t\}\-\\eta\_\{t\}\\,\\mathbf\{g\}^\{t\}\.\(33\)
To analyze the convergence behavior of the stochastic scheme, we impose the following assumptions on the stochastic gradient estimator\.

###### Assumption 5\(Unbiased stochastic gradient\)\.

For allt≥0t\\geq 0,

𝔼​\[𝐠t∣Φt\]=∇𝜽J​\(𝜽t,𝐕t\)\.\\mathbb\{E\}\[\\mathbf\{g\}^\{t\}\\mid\\Phi\_\{t\}\]=\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\.

###### Assumption 6\(Bounded conditional variance\)\.

There existsσ2<∞\\sigma^\{2\}<\\inftysuch that for allt≥0t\\geq 0,

𝔼\[∥𝐠t−∇𝜽J\(𝜽t,𝐕t\)∥2\|Φt\]≤σ2\.\\mathbb\{E\}\\\!\\left\[\\big\\\|\\mathbf\{g\}^\{t\}\-\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\big\\\|^\{2\}\\ \\middle\|\\ \\Phi\_\{t\}\\right\]\\leq\\sigma^\{2\}\.

Under these assumptions, we establish convergence guarantees for the stochastic variant of the proposed method in the following sections\.

#### E\.2\.1Convergence Guarantee with Optimized Step Size

###### Theorem E\.3\(Stochastic convergence rate\)\.

Assume[Assumptions1](https://arxiv.org/html/2608.10016#Thmassumption1),[2](https://arxiv.org/html/2608.10016#Thmassumption2),[3](https://arxiv.org/html/2608.10016#Thmassumption3)and[4](https://arxiv.org/html/2608.10016#Thmassumption4), and the stochastic gradient conditions in[Assumptions5](https://arxiv.org/html/2608.10016#Thmassumption5)and[6](https://arxiv.org/html/2608.10016#Thmassumption6)\. Let\{\(𝛉t,𝐕t\)\}t≥0\\\{\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\}\_\{t\\geq 0\}be generated by \([32](https://arxiv.org/html/2608.10016#A5.E32)\)–\([33](https://arxiv.org/html/2608.10016#A5.E33)\)\. DefineΔ0:=J​\(𝛉0,𝐕0\)−Jinf\\Delta\_\{0\}:=J\(\\boldsymbol\{\\theta\}^\{0\},\\mathbf\{V\}^\{0\}\)\-J\_\{\\inf\}, whereJinf:=inf𝛉,𝐕J​\(𝛉,𝐕\)J\_\{\\inf\}:=\\inf\_\{\\boldsymbol\{\\theta\},\\mathbf\{V\}\}J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}\)\. Then, for any horizonT≥0T\\geq 0, choosing the constant step size

ηt≡η:=min⁡\{1L,2​Δ0L​σ2​\(T\+1\)\}\\eta\_\{t\}\\equiv\\eta:=\\min\\\!\\left\\\{\\frac\{1\}\{L\},\\sqrt\{\\frac\{2\\Delta\_\{0\}\}\{L\\sigma^\{2\}\(T\+1\)\}\}\\right\\\}yields the bound

min0≤t≤T⁡𝔼​\[‖∇𝜽J​\(𝜽t,𝐕t\)‖2\]≤2​2​L​σ2​Δ0T\+1\+2​L​Δ0T\+1\.\\min\_\{0\\leq t\\leq T\}\\mathbb\{E\}\\\!\\left\[\\left\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\right\\\|^\{2\}\\right\]\\leq 2\\sqrt\{\\frac\{2L\\sigma^\{2\}\\,\\Delta\_\{0\}\}\{T\+1\}\}\\;\+\\;\\frac\{2L\\Delta\_\{0\}\}\{T\+1\}\.In particular, the method attains the stochastic convergence rate

min0≤t≤T⁡𝔼​\[‖∇𝜽J​\(𝜽t,𝐕t\)‖2\]=𝒪​\(1T\)\.\\min\_\{0\\leq t\\leq T\}\\mathbb\{E\}\\\!\\left\[\\left\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\right\\\|^\{2\}\\right\]=\\mathcal\{O\}\\\!\\left\(\\frac\{1\}\{\\sqrt\{T\}\}\\right\)\.

###### Proof\.

For brevity, define

Jt:=J​\(𝜽t,𝐕t\),∇t:=∇𝜽J​\(𝜽t,𝐕t\),𝐠t:=𝐠​\(𝜽t,𝐕t;ξt\)\.J\_\{t\}:=J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\),\\qquad\\nabla\_\{t\}:=\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\),\\qquad\\mathbf\{g\}^\{t\}:=\\mathbf\{g\}\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\};\\xi\_\{t\}\)\.Fixt≥0t\\geq 0\. By Assumption[2](https://arxiv.org/html/2608.10016#Thmassumption2), for theLL\-smooth map𝜽↦J​\(𝜽,𝐕t\)\\boldsymbol\{\\theta\}\\mapsto J\(\\boldsymbol\{\\theta\},\\mathbf\{V\}^\{t\}\), we have for any vector𝐮\\mathbf\{u\},

J​\(𝜽t\+𝐮,𝐕t\)≤J​\(𝜽t,𝐕t\)\+∇t⊤𝐮\+L2​‖𝐮‖2\.J\(\\boldsymbol\{\\theta\}^\{t\}\+\\mathbf\{u\},\\mathbf\{V\}^\{t\}\)\\leq J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\+\\nabla\_\{t\}^\{\\top\}\\mathbf\{u\}\+\\frac\{L\}\{2\}\\\|\\mathbf\{u\}\\\|^\{2\}\.Applying this with𝐮=−η​𝐠t\\mathbf\{u\}=\-\\eta\\,\\mathbf\{g\}^\{t\}yields

J​\(𝜽t\+1,𝐕t\)≤J​\(𝜽t,𝐕t\)−η​∇t⊤𝐠t\+L​η22​‖𝐠t‖2\.J\(\\boldsymbol\{\\theta\}^\{t\+1\},\\mathbf\{V\}^\{t\}\)\\leq J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\-\\eta\\,\\nabla\_\{t\}^\{\\top\}\\mathbf\{g\}^\{t\}\+\\frac\{L\\eta^\{2\}\}\{2\}\\\|\\mathbf\{g\}^\{t\}\\\|^\{2\}\.\(34\)By exact minimization in𝐕\\mathbf\{V\}\(Assumption[3](https://arxiv.org/html/2608.10016#Thmassumption3)\),

J​\(𝜽t\+1,𝐕t\+1\)≤J​\(𝜽t\+1,𝐕t\)\.J\(\\boldsymbol\{\\theta\}^\{t\+1\},\\mathbf\{V\}^\{t\+1\}\)\\leq J\(\\boldsymbol\{\\theta\}^\{t\+1\},\\mathbf\{V\}^\{t\}\)\.\(35\)Combining \([34](https://arxiv.org/html/2608.10016#A5.E34)\)–\([35](https://arxiv.org/html/2608.10016#A5.E35)\) gives

Jt\+1≤Jt−η​∇t⊤𝐠t\+L​η22​‖𝐠t‖2\.J\_\{t\+1\}\\leq J\_\{t\}\-\\eta\\,\\nabla\_\{t\}^\{\\top\}\\mathbf\{g\}^\{t\}\+\\frac\{L\\eta^\{2\}\}\{2\}\\\|\\mathbf\{g\}^\{t\}\\\|^\{2\}\.\(36\)
Taking conditional expectation w\.r\.t\.Φt\\Phi\_\{t\}and using unbiasedness \(Assumption[5](https://arxiv.org/html/2608.10016#Thmassumption5)\) gives

𝔼​\[∇t⊤𝐠t∣Φt\]=∇t⊤𝔼​\[𝐠t∣Φt\]=‖∇t‖2\.\\mathbb\{E\}\\\!\\left\[\\nabla\_\{t\}^\{\\top\}\\mathbf\{g\}^\{t\}\\mid\\Phi\_\{t\}\\right\]=\\nabla\_\{t\}^\{\\top\}\\mathbb\{E\}\[\\mathbf\{g\}^\{t\}\\mid\\Phi\_\{t\}\]=\\\|\\nabla\_\{t\}\\\|^\{2\}\.Moreover, expanding the second moment and using the bounded conditional variance \(Assumption[6](https://arxiv.org/html/2608.10016#Thmassumption6)\) yields

𝔼​\[‖𝐠t‖2∣Φt\]\\displaystyle\\mathbb\{E\}\\\!\\left\[\\\|\\mathbf\{g\}^\{t\}\\\|^\{2\}\\mid\\Phi\_\{t\}\\right\]=𝔼​\[‖𝐠t−∇t\+∇t‖2∣Φt\]\\displaystyle=\\mathbb\{E\}\\\!\\left\[\\\|\\mathbf\{g\}^\{t\}\-\\nabla\_\{t\}\+\\nabla\_\{t\}\\\|^\{2\}\\mid\\Phi\_\{t\}\\right\]=‖∇t‖2\+𝔼​\[‖𝐠t−∇t‖2∣Φt\]\+2​∇t⊤𝔼​\[𝐠t−∇t∣Φt\]\\displaystyle=\\\|\\nabla\_\{t\}\\\|^\{2\}\+\\mathbb\{E\}\\\!\\left\[\\\|\\mathbf\{g\}^\{t\}\-\\nabla\_\{t\}\\\|^\{2\}\\mid\\Phi\_\{t\}\\right\]\+2\\,\\nabla\_\{t\}^\{\\top\}\\mathbb\{E\}\[\\mathbf\{g\}^\{t\}\-\\nabla\_\{t\}\\mid\\Phi\_\{t\}\]≤‖∇t‖2\+σ2,\\displaystyle\\leq\\\|\\nabla\_\{t\}\\\|^\{2\}\+\\sigma^\{2\},since𝔼​\[𝐠t−∇t∣Φt\]=𝟎\\mathbb\{E\}\[\\mathbf\{g\}^\{t\}\-\\nabla\_\{t\}\\mid\\Phi\_\{t\}\]=\\mathbf\{0\}by Assumption[5](https://arxiv.org/html/2608.10016#Thmassumption5)\. Substituting these two identities into the conditional expectation of \([36](https://arxiv.org/html/2608.10016#A5.E36)\) yields

𝔼​\[Jt\+1∣Φt\]≤Jt−η​\(1−L​η2\)​‖∇t‖2\+L​η22​σ2\.\\mathbb\{E\}\[J\_\{t\+1\}\\mid\\Phi\_\{t\}\]\\leq J\_\{t\}\-\\eta\\Big\(1\-\\frac\{L\\eta\}\{2\}\\Big\)\\\|\\nabla\_\{t\}\\\|^\{2\}\+\\frac\{L\\eta^\{2\}\}\{2\}\\sigma^\{2\}\.\(37\)Assuming0<η≤1/L0<\\eta\\leq 1/Limplies1−L​η2≥121\-\\frac\{L\\eta\}\{2\}\\geq\\frac\{1\}\{2\}, hence

𝔼​\[Jt\+1∣Φt\]≤Jt−η2​‖∇t‖2\+L​η22​σ2\.\\mathbb\{E\}\[J\_\{t\+1\}\\mid\\Phi\_\{t\}\]\\leq J\_\{t\}\-\\frac\{\\eta\}\{2\}\\\|\\nabla\_\{t\}\\\|^\{2\}\+\\frac\{L\\eta^\{2\}\}\{2\}\\sigma^\{2\}\.\(38\)
Taking total expectation in \([38](https://arxiv.org/html/2608.10016#A5.E38)\) and summing fromt=0t=0toTTgives

𝔼​\[JT\+1\]≤J0−η2​∑t=0T𝔼​‖∇t‖2\+L​η22​σ2​\(T\+1\)\.\\mathbb\{E\}\[J\_\{T\+1\}\]\\leq J\_\{0\}\-\\frac\{\\eta\}\{2\}\\sum\_\{t=0\}^\{T\}\\mathbb\{E\}\\\|\\nabla\_\{t\}\\\|^\{2\}\+\\frac\{L\\eta^\{2\}\}\{2\}\\sigma^\{2\}\(T\+1\)\.Rearranging and using the lower bound𝔼​\[JT\+1\]≥Jinf\\mathbb\{E\}\[J\_\{T\+1\}\]\\geq J\_\{\\inf\}\(Assumption[1](https://arxiv.org/html/2608.10016#Thmassumption1)\) yields

η2​∑t=0T𝔼​‖∇t‖2≤Δ0\+L​η22​σ2​\(T\+1\),Δ0:=J0−Jinf\.\\frac\{\\eta\}\{2\}\\sum\_\{t=0\}^\{T\}\\mathbb\{E\}\\\|\\nabla\_\{t\}\\\|^\{2\}\\leq\\Delta\_\{0\}\+\\frac\{L\\eta^\{2\}\}\{2\}\\sigma^\{2\}\(T\+1\),\\qquad\\Delta\_\{0\}:=J\_\{0\}\-J\_\{\\inf\}\.\(39\)Dividing byη​\(T\+1\)\\eta\(T\+1\)and usingmin0≤t≤T⁡at≤1T\+1​∑t=0Tat\\min\_\{0\\leq t\\leq T\}a\_\{t\}\\leq\\frac\{1\}\{T\+1\}\\sum\_\{t=0\}^\{T\}a\_\{t\}gives

min0≤t≤T⁡𝔼​‖∇t‖2≤2​Δ0η​\(T\+1\)\+L​η​σ2\.\\min\_\{0\\leq t\\leq T\}\\mathbb\{E\}\\\|\\nabla\_\{t\}\\\|^\{2\}\\leq\\frac\{2\\Delta\_\{0\}\}\{\\eta\(T\+1\)\}\+L\\eta\\,\\sigma^\{2\}\.\(40\)
We now optimize the right\-hand side inη\\etaunder the constraintη≤1/L\\eta\\leq 1/L\. Let

η⋆:=2​Δ0L​σ2​\(T\+1\)\.\\eta^\{\\star\}:=\\sqrt\{\\frac\{2\\Delta\_\{0\}\}\{L\\sigma^\{2\}\(T\+1\)\}\}\.Ifη⋆≤1/L\\eta^\{\\star\}\\leq 1/L, choosingη=η⋆\\eta=\\eta^\{\\star\}in \([40](https://arxiv.org/html/2608.10016#A5.E40)\) yields

min0≤t≤T⁡𝔼​‖∇t‖2≤2​2​L​σ2​Δ0T\+1\.\\min\_\{0\\leq t\\leq T\}\\mathbb\{E\}\\\|\\nabla\_\{t\}\\\|^\{2\}\\leq 2\\sqrt\{\\frac\{2L\\sigma^\{2\}\\Delta\_\{0\}\}\{T\+1\}\}\.If insteadη⋆\>1/L\\eta^\{\\star\}\>1/L, we chooseη=1/L\\eta=1/Lin \([40](https://arxiv.org/html/2608.10016#A5.E40)\) to obtain

min0≤t≤T⁡𝔼​‖∇t‖2≤2​L​Δ0T\+1\+σ2\.\\min\_\{0\\leq t\\leq T\}\\mathbb\{E\}\\\|\\nabla\_\{t\}\\\|^\{2\}\\leq\\frac\{2L\\Delta\_\{0\}\}\{T\+1\}\+\\sigma^\{2\}\.Both cases are covered by the choice

η=min⁡\{1L,2​Δ0L​σ2​\(T\+1\)\}\.\\eta=\\min\\left\\\{\\frac\{1\}\{L\},\\ \\sqrt\{\\frac\{2\\Delta\_\{0\}\}\{L\\sigma^\{2\}\(T\+1\)\}\}\\right\\\}\.Moreover, with this choice we have the unified bound

min0≤t≤T⁡𝔼​‖∇t‖2≤2​2​L​σ2​Δ0T\+1\+2​L​Δ0T\+1,\\min\_\{0\\leq t\\leq T\}\\mathbb\{E\}\\\|\\nabla\_\{t\}\\\|^\{2\}\\leq 2\\sqrt\{\\frac\{2L\\sigma^\{2\}\\Delta\_\{0\}\}\{T\+1\}\}\+\\frac\{2L\\Delta\_\{0\}\}\{T\+1\},where the second term is redundant whenη=η⋆\\eta=\\eta^\{\\star\}but remains valid\. This proves Theorem[E\.3](https://arxiv.org/html/2608.10016#A5.Thmtheorem3)\. ∎

#### E\.2\.2Almost\-sure Convergence with Diminishing Step Sizes

We establish an almost\-sure convergence guarantee under a Robbins–Monro step\-size schedule\. Throughout, let\{Φt\}t≥0\\\{\\Phi\_\{t\}\\\}\_\{t\\geq 0\}denote the natural filtration and let the iterates be generated by \([32](https://arxiv.org/html/2608.10016#A5.E32)\)–\([33](https://arxiv.org/html/2608.10016#A5.E33)\)\.

###### Assumption 7\(Robbins–Monro step sizes\)\.

The step sizes satisfy0<ηt≤1/L0<\\eta\_\{t\}\\leq 1/Lfor allt≥0t\\geq 0, and

∑t=0∞ηt=∞,∑t=0∞ηt2<∞\.\\sum\_\{t=0\}^\{\\infty\}\\eta\_\{t\}=\\infty,\\qquad\\sum\_\{t=0\}^\{\\infty\}\\eta\_\{t\}^\{2\}<\\infty\.

###### Theorem E\.4\(Almost\-sure convergence\)\.

Assume[Assumptions1](https://arxiv.org/html/2608.10016#Thmassumption1),[2](https://arxiv.org/html/2608.10016#Thmassumption2),[3](https://arxiv.org/html/2608.10016#Thmassumption3)and[4](https://arxiv.org/html/2608.10016#Thmassumption4), and[Assumptions5](https://arxiv.org/html/2608.10016#Thmassumption5)and[6](https://arxiv.org/html/2608.10016#Thmassumption6)and[Assumption7](https://arxiv.org/html/2608.10016#Thmassumption7)\. Let\{\(𝛉t,𝐕t\)\}t≥0\\\{\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\}\_\{t\\geq 0\}be generated by \([32](https://arxiv.org/html/2608.10016#A5.E32)\)–\([33](https://arxiv.org/html/2608.10016#A5.E33)\)\. Then:

1. 1\.The sequence\{J​\(𝜽t,𝐕t\)\}\\\{J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\\}converges almost surely to a finite random variable\.
2. 2\.Moreover, ∑t=0∞ηt​‖∇𝜽J​\(𝜽t,𝐕t\)‖2<∞almost surely\.\\sum\_\{t=0\}^\{\\infty\}\\eta\_\{t\}\\left\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\right\\\|^\{2\}<\\infty\\qquad\\text\{almost surely\}\.
3. 3\.Consequently, lim inft→∞‖∇𝜽J​\(𝜽t,𝐕t\)‖=0almost surely\.\\liminf\_\{t\\to\\infty\}\\left\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\)\\right\\\|=0\\qquad\\text\{almost surely\}\.

###### Proof\.

We begin by recalling the Robbins–Siegmund lemma, which provides a standard convergence result for quasi\-supermartingale sequences\.

###### Lemma E\.5\(Robbins–Siegmund,Robbins & Siegmund,[1971](https://arxiv.org/html/2608.10016#bib.bib52)\)\.

Let\{Xt\}t≥0\\\{X\_\{t\}\\\}\_\{t\\geq 0\}be a nonnegative adapted sequence with respect to a filtration\{Φt\}t≥0\\\{\\Phi\_\{t\}\\\}\_\{t\\geq 0\}\. Assume there exist nonnegativeΦt\\Phi\_\{t\}\-measurable sequences\{at\}t≥0\\\{a\_\{t\}\\\}\_\{t\\geq 0\}and\{bt\}t≥0\\\{b\_\{t\}\\\}\_\{t\\geq 0\}such that

𝔼​\[Xt\+1∣Φt\]≤Xt−at\+bta\.s\. for all​t,\\mathbb\{E\}\[X\_\{t\+1\}\\mid\\Phi\_\{t\}\]\\leq X\_\{t\}\-a\_\{t\}\+b\_\{t\}\\qquad\\text\{a\.s\. for all \}t,and∑t=0∞bt<∞\\sum\_\{t=0\}^\{\\infty\}b\_\{t\}<\\inftyalmost surely\. ThenXtX\_\{t\}converges almost surely to a finite random variable and∑t=0∞at<∞\\sum\_\{t=0\}^\{\\infty\}a\_\{t\}<\\inftyalmost surely\.

We now show that the sequence generated by the algorithm satisfies the conditions of this lemma\. Define

Jt:=J​\(𝜽t,𝐕t\),∇t:=∇𝜽J​\(𝜽t,𝐕t\),Xt:=Jt−Jinf≥0\.J\_\{t\}:=J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\),\\qquad\\nabla\_\{t\}:=\\nabla\_\{\\boldsymbol\{\\theta\}\}J\(\\boldsymbol\{\\theta\}^\{t\},\\mathbf\{V\}^\{t\}\),\\qquad X\_\{t\}:=J\_\{t\}\-J\_\{\\inf\}\\geq 0\.
From the one\-step bound \(cf\. \([38](https://arxiv.org/html/2608.10016#A5.E38)\)\), for allt≥0t\\geq 0,

𝔼​\[Jt\+1∣Φt\]≤Jt−ηt2​‖∇t‖2\+L​ηt22​σ2\.\\mathbb\{E\}\[J\_\{t\+1\}\\mid\\Phi\_\{t\}\]\\leq J\_\{t\}\-\\frac\{\\eta\_\{t\}\}\{2\}\\\|\\nabla\_\{t\}\\\|^\{2\}\+\\frac\{L\\eta\_\{t\}^\{2\}\}\{2\}\\sigma^\{2\}\.SubtractingJinfJ\_\{\\inf\}from both sides yields

𝔼​\[Xt\+1∣Φt\]≤Xt−at\+bt,at:=ηt2​‖∇t‖2,bt:=L​ηt22​σ2\.\\mathbb\{E\}\[X\_\{t\+1\}\\mid\\Phi\_\{t\}\]\\leq X\_\{t\}\-a\_\{t\}\+b\_\{t\},\\qquad a\_\{t\}:=\\frac\{\\eta\_\{t\}\}\{2\}\\\|\\nabla\_\{t\}\\\|^\{2\},\\quad b\_\{t\}:=\\frac\{L\\eta\_\{t\}^\{2\}\}\{2\}\\sigma^\{2\}\.
By Assumption[7](https://arxiv.org/html/2608.10016#Thmassumption7)we have∑t=0∞ηt2<∞\\sum\_\{t=0\}^\{\\infty\}\\eta\_\{t\}^\{2\}<\\infty, hence∑t=0∞bt<∞\\sum\_\{t=0\}^\{\\infty\}b\_\{t\}<\\infty\. Therefore all the conditions of Lemma[E\.5](https://arxiv.org/html/2608.10016#A5.Thmtheorem5)are satisfied\. Applying the lemma implies thatXtX\_\{t\}converges almost surely to a finite random variable, and

∑t=0∞ηt​‖∇t‖2<∞a\.s\.\\sum\_\{t=0\}^\{\\infty\}\\eta\_\{t\}\\\|\\nabla\_\{t\}\\\|^\{2\}<\\infty\\qquad\\text\{a\.s\.\}
Finally, since∑t=0∞ηt=∞\\sum\_\{t=0\}^\{\\infty\}\\eta\_\{t\}=\\infty\(Assumption[7](https://arxiv.org/html/2608.10016#Thmassumption7)\), the above summability implies

lim inft→∞‖∇t‖=0a\.s\.\\liminf\_\{t\\to\\infty\}\\\|\\nabla\_\{t\}\\\|=0\\qquad\\text\{a\.s\.\}This proves Theorem[E\.4](https://arxiv.org/html/2608.10016#A5.Thmtheorem4)\. ∎

## Appendix FTrade\-off between Privacy, Communication, and Local Computation

Our proposed algorithmic solution in[Sec\.5](https://arxiv.org/html/2608.10016#S5)leverages sharing of raw latent representations to lower the communication cost, according to the objective in[Prob\.1](https://arxiv.org/html/2608.10016#Thmproblem1)\. This is compatible with standard regulatory constraints on privacy, for instance when latent representations do not allow direct identification of individuals or when appropriate safeguards such as aggregation, anonymization, or secure communication protocols are in place, as commonly assumed in federated learning systems\(Kairouz et al\.,[2021](https://arxiv.org/html/2608.10016#bib.bib35); McMahan et al\.,[2017](https://arxiv.org/html/2608.10016#bib.bib45)\)\. However, as highlighted in[Remark3](https://arxiv.org/html/2608.10016#Thmremark3), in some circumstances more stringent regulatory constraints on privacy impair the sharing of raw latent representations\. For example, this occurs when latent representations may still encode sensitive or identifiable information and are therefore susceptible to reconstruction or inference attacks, such as model inversion\(Fredrikson et al\.,[2015](https://arxiv.org/html/2608.10016#bib.bib25)\), membership inference\(Shokri et al\.,[2017](https://arxiv.org/html/2608.10016#bib.bib56)\), or attribute inference\(Melis et al\.,[2019](https://arxiv.org/html/2608.10016#bib.bib46)\), making their direct exchange incompatible with strict privacy requirements\. Our proposed approach naturally extends to this more stringent setting\. Instead of sharing raw latent representations, agents exchange their transported counterparts, i\.e\., latent features mapped through the restriction maps associated with the edges\. This limits the direct exposure of local features\. However, this enhanced privacy comes at a cost: it requires additional local computation and increases the communication overhead, as detailed below\.

Concerning the local computational cost, each agenti∈𝒩i\\in\\mathcal\{N\}needs to encode its latent representations toward both sets of neighbors𝒩​\(i\)−\\mathcal\{N\}\(i\)^\{\-\}and𝒩​\(i\)\+\\mathcal\{N\}\(i\)^\{\+\}\. To this end, we adopt the formulation in[Eq\.10](https://arxiv.org/html/2608.10016#S4.E10)*before reparameterization*, where both orthogonal and Stiefel maps explicitly appear, allowing consistent encoding across heterogeneous neighbors\. Specifically,iimust communicate𝐕j​it​𝐀it∈ℝdj×K\\mathbf\{V\}\_\{ji\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t\}\\in\\mathbb\{R\}^\{d\_\{j\}\\times K\}to neighborsj∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}, where𝐕j​it∈St​\(dj,di\)\\mathbf\{V\}\_\{ji\}^\{t\}\\in\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\)is the outgoing embedding matrix, and𝐎j​it​𝐀it∈ℝdi×K\\mathbf\{O\}\_\{ji\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t\}\\in\\mathbb\{R\}^\{d\_\{i\}\\times K\}to neighborsj∈𝒩​\(i\)−j\\in\\mathcal\{N\}\(i\)^\{\-\}, where𝐎j​it∈O​\(di\)\\mathbf\{O\}\_\{ji\}^\{t\}\\in\\mathrm\{O\}\(\{d\_\{i\}\}\)is an orthogonal transformation\. Hence, in this setting we cannot exploit the reparameterized form in[Eq\.10](https://arxiv.org/html/2608.10016#S4.E10), which avoids the explicit use of orthogonal maps\. In detail, after receiving the encoded latent representations from its neighbors, namely

𝐕i​jt​𝐀jt​from​j∈𝒩​\(i\)−and𝐎i​jt​𝐀jt​from​j∈𝒩​\(i\)\+;\\mathbf\{V\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\\text\{ from \}j\\in\\mathcal\{N\}\(i\)^\{\-\}\\quad\\text\{and\}\\quad\\mathbf\{O\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\\text\{ from \}j\\in\\mathcal\{N\}\(i\)^\{\+\}\\,;\(41\)under stringent privacy constraints, agentiimust solve the following problems to update its restriction maps:

𝐎j​it\\displaystyle\\mathbf\{O\}\_\{ji\}^\{t\}=arg​min𝐎j​i∈O​\(di\)‖𝐎j​i​𝐀it−𝐕i​jt​𝐀jt‖F2,,∀,j∈𝒩\(i\)−,\\displaystyle=\\operatorname\*\{arg\\,min\}\_\{\\mathbf\{O\}\_\{ji\}\\in\\mathrm\{O\}\(\{d\_\{i\}\}\)\}\\;\\norm\{\\mathbf\{O\}\_\{ji\}\\mathbf\{A\}\_\{i\}^\{t\}\-\\mathbf\{V\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\}\_\{\\mathrm\{F\}\}^\{2\},,\\quad\\forall,j\\in\\mathcal\{N\}\(i\)^\{\-\}\\,,\(P2c\)𝐕j​it\\displaystyle\\mathbf\{V\}\_\{ji\}^\{t\}=arg​min𝐕j​i∈St​\(dj,di\)‖𝐎i​jt​𝐀jt−𝐕j​i​𝐀it‖F2,,∀,j∈𝒩\(i\)\+\.\\displaystyle=\\operatorname\*\{arg\\,min\}\_\{\\mathbf\{V\}\_\{ji\}\\in\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\)\}\\;\\norm\{\\mathbf\{O\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\-\\mathbf\{V\}\_\{ji\}\\mathbf\{A\}\_\{i\}^\{t\}\}\_\{\\mathrm\{F\}\}^\{2\},,\\quad\\forall,j\\in\\mathcal\{N\}\(i\)^\{\+\}\\,\.\(P2d\)
Unlike the case of standard privacy constraints, where forj∈𝒩​\(i\)−j\\in\\mathcal\{N\}\(i\)^\{\-\}we can only update the incoming restriction map𝐕i​j\\mathbf\{V\}\_\{ij\}and solve Prob\. \([P2a](https://arxiv.org/html/2608.10016#S5.Ex5)\) via a thin SVD at a computational cost of𝒪​\(di​dj2\)\\mathcal\{O\}\(d\_\{i\}d\_\{j\}^\{2\}\)withdj≤did\_\{j\}\\leq d\_\{i\}, Prob\. \([P2c](https://arxiv.org/html/2608.10016#A6.Ex44)\) is a canonical orthogonal Procrustes problem\. Its solution is obtained via a classical SVD of𝐕i​jt​𝐀jt​𝐀it⊤\\mathbf\{V\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t^\{\\top\}\}at a cost of𝒪​\(di3\)\\mathcal\{O\}\(d\_\{i\}^\{3\}\), thus leading to an increase in the local computational burden\. Conversely, similarly to Problem \([P2b](https://arxiv.org/html/2608.10016#S5.Ex6)\), forj∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}the update of the outgoing restriction map𝐕j​i\\mathbf\{V\}\_\{ji\}can be computed via a thin SVD of𝐎i​jt​𝐀jt​𝐀it⊤\\mathbf\{O\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t^\{\\top\}\}at a cost of𝒪​\(dj​di2\)\\mathcal\{O\}\(d\_\{j\}d\_\{i\}^\{2\}\), where in this casedj≥did\_\{j\}\\geq d\_\{i\}\.

Please notice that, more stringent privacy constraints affect only the update of the restriction maps\. Indeed, the local gradient computation in[Eq\.19](https://arxiv.org/html/2608.10016#S5.E19)remains unchanged:

∇𝝋iℛ𝒜\|i=\\displaystyle\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\mathcal\{R\}\_\{\\mathcal\{A\}\}\|\_\{i\}=λiK\[∑j∈𝒩​\(i\)−∑k∈𝒜\(∇𝝋if𝝋i\(𝐱ik\)\)⊤𝐎j​i⊤\(𝐎j​if𝝋i\(𝐱ik\)−𝐕i​jf𝝋j\(𝐱jk\)\)\\displaystyle\\frac\{\\lambda\_\{i\}\}\{K\}\\Bigg\[\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\}\}\\sum\_\{k\\in\\mathcal\{A\}\}\\left\(\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\right\)^\{\\top\}\\mathbf\{O\}\_\{ji\}^\{\\top\}\\left\(\\mathbf\{O\}\_\{ji\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\-\\mathbf\{V\}\_\{ij\}f\_\{\\boldsymbol\{\\varphi\}\_\{j\}\}\\left\(\\mathbf\{x\}\_\{j\}^\{k\}\\right\)\\right\)\(42\)−∑j∈𝒩​\(i\)\+∑k∈𝒜\(∇𝝋if𝝋i\(𝐱ik\)\)⊤𝐕j​i⊤\(𝐎i​jf𝝋j\(𝐱jk\)−𝐕j​if𝝋i\(𝐱ik\)\)\]\\displaystyle\-\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\+\}\}\\sum\_\{k\\in\\mathcal\{A\}\}\\left\(\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\right\)^\{\\top\}\\mathbf\{V\}\_\{ji\}^\{\\top\}\\left\(\\mathbf\{O\}\_\{ij\}f\_\{\\boldsymbol\{\\varphi\}\_\{j\}\}\\left\(\\mathbf\{x\}\_\{j\}^\{k\}\\right\)\-\\mathbf\{V\}\_\{ji\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\\right\)\\Bigg\]=\\displaystyle=λiK\[∑j∈𝒩​\(i\)−∑k∈𝒜\(∇𝝋if𝝋i\(𝐱ik\)\)⊤\(f𝝋i\(𝐱ik\)−𝐎j​i⊤𝐕i​jf𝝋j\(𝐱jk\)\)\\displaystyle\\frac\{\\lambda\_\{i\}\}\{K\}\\Bigg\[\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\}\}\\sum\_\{k\\in\\mathcal\{A\}\}\\left\(\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\right\)^\{\\top\}\\left\(f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\-\\mathbf\{O\}\_\{ji\}^\{\\top\}\\mathbf\{V\}\_\{ij\}f\_\{\\boldsymbol\{\\varphi\}\_\{j\}\}\\left\(\\mathbf\{x\}\_\{j\}^\{k\}\\right\)\\right\)−∑j∈𝒩​\(i\)\+∑k∈𝒜\(∇𝝋if𝝋i\(𝐱ik\)\)⊤\(𝐕j​i⊤𝐎i​jf𝝋j\(𝐱jk\)−f𝝋i\(𝐱ik\)\)\]\\displaystyle\-\\sum\_\{j\\in\\mathcal\{N\}\(i\)^\{\+\}\}\\sum\_\{k\\in\\mathcal\{A\}\}\\left\(\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\right\)^\{\\top\}\\left\(\\mathbf\{V\}\_\{ji\}^\{\\top\}\\mathbf\{O\}\_\{ij\}f\_\{\\boldsymbol\{\\varphi\}\_\{j\}\}\\left\(\\mathbf\{x\}\_\{j\}^\{k\}\\right\)\-f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\\right\)\\Bigg\]=\(a\)\\displaystyle\\stackrel\{\{\\scriptstyle\(a\)\}\}\{\{=\}\}λiK​∑j∈𝒩​\(i\)∑k∈𝒜\(∇𝝋if𝝋i​\(𝐱ik\)\)⊤​\(f𝝋i​\(𝐱ik\)−𝐎j​i⊤​𝐕i​j​f𝝋j​\(𝐱jk\)\);\\displaystyle\\frac\{\\lambda\_\{i\}\}\{K\}\\sum\_\{j\\in\\mathcal\{N\}\(i\)\}\\sum\_\{k\\in\\mathcal\{A\}\}\\left\(\\nabla\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\\right\)^\{\\top\}\\left\(f\_\{\\boldsymbol\{\\varphi\}\_\{i\}\}\\left\(\\mathbf\{x\}\_\{i\}^\{k\}\\right\)\-\\mathbf\{O\}\_\{ji\}^\{\\top\}\\mathbf\{V\}\_\{ij\}f\_\{\\boldsymbol\{\\varphi\}\_\{j\}\}\\left\(\\mathbf\{x\}\_\{j\}^\{k\}\\right\)\\right\)\\,;where in\(a\)\(a\), for thosej∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}we used that for a network sheaf𝐕j​i⊤​𝐎i​j=𝐎j​i⊤​𝐕i​j\\mathbf\{V\}\_\{ji\}^\{\\top\}\\mathbf\{O\}\_\{ij\}=\\mathbf\{O\}\_\{ji\}^\{\\top\}\\mathbf\{V\}\_\{ij\}\. Importantly,𝐎j​i⊤​𝐕i​j\\mathbf\{O\}\_\{ji\}^\{\\top\}\\mathbf\{V\}\_\{ij\}is exactly the reparameterization exploited in[Eq\.10](https://arxiv.org/html/2608.10016#S4.E10), thus showing the equivalence between[Eq\.19](https://arxiv.org/html/2608.10016#S5.E19)and[Eq\.42](https://arxiv.org/html/2608.10016#A6.E42)\.

Regarding the communication cost, sharing transported latent representations increases the communication overhead toward neighborsj∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}\. Under standard privacy constraints, agentiitransmits𝐀it∈ℝdi×K\\mathbf\{A\}\_\{i\}^\{t\}\\in\\mathbb\{R\}^\{d\_\{i\}\\times K\}, whereas in the more stringent setting it must send𝐕j​i​𝐀it∈ℝdj×K\\mathbf\{V\}\_\{ji\}\\mathbf\{A\}\_\{i\}^\{t\}\\in\\mathbb\{R\}^\{d\_\{j\}\\times K\}, withdj≥did\_\{j\}\\geq d\_\{i\}\. Thus, the communicated payload scales withdjd\_\{j\}, and strictly increases wheneverdj\>did\_\{j\}\>d\_\{i\}\.

Our modified Sheaf\-FRL protocols under stringent privacy regulatory constraints for both the heterogeneous and homogeneous cases are given in[Algs\.3](https://arxiv.org/html/2608.10016#alg3)and[4](https://arxiv.org/html/2608.10016#alg4), respectively\. Additionally,[Tab\.1](https://arxiv.org/html/2608.10016#A6.T1)details the privacy vs\. communication and computational costs trade\-off\.

Table 1:Trade\-off between privacy and communication and local computation cost per agent\. The local computation refers to the update of the restriction maps, i\.e\., the part influenced by more stringent privacy constraints\. To aid comparison, for the heterogeneous case we highlight inbluethe larger ofdid\_\{i\}anddjd\_\{j\}\(either one whendi=djd\_\{i\}=d\_\{j\}\) when considering eitherj∈𝒩​\(i\)−j\\in\\mathcal\{N\}\(i\)^\{\-\}orj∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}\. Inred, we instead highlight we instead highlight the \(potential\) increase in cost implied by more stringent privacy constraints\.Non\-identifiability and privacy guarantees\.[Tab\.1](https://arxiv.org/html/2608.10016#A6.T1)highlights a clear trade\-off\. Under stricter privacy requirements, the proposed protocol incurs higher communication and local computational costs, since agents exchange transported latent representations and must explicitly update the corresponding restriction maps\. The benefit, however, is a stronger form of privacy: the communicated messages no longer reveal the latent pilots in their original coordinates\. More precisely, when only transported pilots are shared, an observer that does not know the underlying restriction maps cannot uniquely reconstruct the original latent representations, but can only recover geometric information that is invariant under isometries, as formalized in the following proposition\.

###### Proposition F\.1\(Non\-identifiability of latent pilots under unknown maps\)\.

Let𝐀∈ℝd×K\\mathbf\{A\}\\in\\mathbb\{R\}^\{d\\times K\}be a matrix of latent pilot representations, and let𝐘=𝐓𝐀\\mathbf\{Y\}=\\mathbf\{T\}\\mathbf\{A\}, where𝐓\\mathbf\{T\}is an unknown restriction map satisfying𝐓⊤​𝐓=𝐈d\\mathbf\{T\}^\{\\top\}\\mathbf\{T\}=\\mathbf\{I\}\_\{d\}\. This includes both the homogeneous case𝐓∈O​\(d\)\\mathbf\{T\}\\in\\mathrm\{O\}\(\{d\}\)and the heterogeneous case𝐓∈St​\(D,d\)\\mathbf\{T\}\\in\\mathrm\{St\}\(\{D\},\{d\}\)withD≥dD\\geq d\. Then the latent pilots𝐀\\mathbf\{A\}are not uniquely identifiable from𝐘\\mathbf\{Y\}without knowledge of𝐓\\mathbf\{T\}\. In particular,𝐀\\mathbf\{A\}is determined only up to an orthogonal transformation\.

###### Proof\.

Let𝐐∈O​\(d\)\\mathbf\{Q\}\\in\\mathrm\{O\}\(\{d\}\)and𝐀′=𝐐𝐀\\mathbf\{A\}^\{\\prime\}=\\mathbf\{Q\}\\mathbf\{A\}\. Then𝐘=𝐓𝐀=𝐓𝐐⊤​𝐀′\\mathbf\{Y\}=\\mathbf\{T\}\\mathbf\{A\}=\\mathbf\{T\}\\mathbf\{Q\}^\{\\top\}\\mathbf\{A\}^\{\\prime\}, and\(𝐓𝐐⊤\)⊤​\(𝐓𝐐⊤\)=𝐈d\(\\mathbf\{T\}\\mathbf\{Q\}^\{\\top\}\)^\{\\top\}\(\\mathbf\{T\}\\mathbf\{Q\}^\{\\top\}\)=\\mathbf\{I\}\_\{d\}\. Thus𝐓𝐐⊤\\mathbf\{T\}\\mathbf\{Q\}^\{\\top\}is a valid map, implying that𝐀\\mathbf\{A\}is identifiable only up to orthogonal transformations\. ∎

Algorithm 3Sheaf\-FRL under stringent privacy regulatory constraints0:Oriented

𝒢=\(𝒩,ℰ\)\\mathcal\{G\}=\(\\mathcal\{N\},\\mathcal\{E\}\), pilots

𝒜\\mathcal\{A\},

λ\>0\\lambda\>0,

ηθ\>0\\eta\_\{\\theta\}\>0, iterations

TT
1:Initialize

\{𝜽i0=\(𝝋i0,𝝍i0\)\}i∈𝒩\\\{\\boldsymbol\{\\theta\}\_\{i\}^\{0\}=\(\\boldsymbol\{\\varphi\}\_\{i\}^\{0\},\\boldsymbol\{\\psi\}\_\{i\}^\{0\}\)\\\}\_\{i\\in\\mathcal\{N\}\},

\{𝐎j​i0∈O​\(di\)\}j∈𝒩​\(i\)−,i∈𝒩\\\{\\mathbf\{O\}\_\{ji\}^\{0\}\\in\\mathrm\{O\}\(\{d\_\{i\}\}\)\\\}\_\{j\\in\\mathcal\{N\}\(i\)^\{\-\},\\,i\\in\\mathcal\{N\}\},

\{𝐕j​i0∈St​\(dj,di\)\}j∈𝒩​\(i\)\+,i∈𝒩\\\{\\mathbf\{V\}\_\{ji\}^\{0\}\\in\\mathrm\{St\}\(\{d\_\{j\}\},\{d\_\{i\}\}\)\\\}\_\{j\\in\\mathcal\{N\}\(i\)^\{\+\},\\,i\\in\\mathcal\{N\}\}
2:for

t=0,…,T−1t=0,\\dots,T\-1do

3:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

4:

𝐀it←\[f𝝋it​\(𝐱ik\)\]k∈𝒜\\mathbf\{A\}\_\{i\}^\{t\}\\leftarrow\[f\_\{\\boldsymbol\{\\varphi\}\_\{i\}^\{t\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\]\_\{k\\in\\mathcal\{A\}\}
5:endfor

6:Each node

iibroadcasts

𝐎j​it​𝐀it\\mathbf\{O\}\_\{ji\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t\}to those

j∈𝒩​\(i\)−j\\in\\mathcal\{N\}\(i\)^\{\-\}and

𝐕j​it​𝐀it\\mathbf\{V\}\_\{ji\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t\}to

j∈𝒩​\(i\)\+j\\in\\mathcal\{N\}\(i\)^\{\+\}
7:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

8:for

j∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)do

9:if

j∈𝒩​\(i\)−j\\in\\mathcal\{N\}\(i\)^\{\-\}then

10:

\[𝐔,𝚺,𝐖⊤\]←SVD​\(𝐕i​jt​𝐀jt​𝐀it⊤\)\[\\mathbf\{U\},\\mathbf\{\\Sigma\},\\mathbf\{W\}^\{\\top\}\]\\leftarrow\\mathrm\{SVD\}\(\\mathbf\{V\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t^\{\\top\}\}\)
11:

𝐎j​it←𝐔𝐖⊤\\mathbf\{O\}\_\{ji\}^\{t\}\\leftarrow\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}
12:else

13:

\[𝐔,𝚺,𝐖⊤\]←thinSVD​\(𝐎i​jt​𝐀jt​𝐀it⊤\)\[\\mathbf\{U\},\\mathbf\{\\Sigma\},\\mathbf\{W\}^\{\\top\}\]\\leftarrow\\mathrm\{thinSVD\}\(\\mathbf\{O\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t^\{\\top\}\}\)
14:

𝐕j​it←𝐔𝐖⊤\\mathbf\{V\}\_\{ji\}^\{t\}\\leftarrow\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}
15:endif

16:endfor

17:endfor

18:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

19:

𝜽it\+1←𝜽it−ηθ​\(∇𝜽iℒi​\(𝜽it\)\+𝐫it\)\\boldsymbol\{\\theta\}\_\{i\}^\{t\+1\}\\leftarrow\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\-\\eta\_\{\\theta\}\\left\(\\nabla\_\{\\boldsymbol\{\\theta\}\_\{i\}\}\\mathcal\{L\}\_\{i\}\(\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\)\+\\mathbf\{r\}\_\{i\}^\{t\}\\right\)
20:endfor

21:endfor

Algorithm 4Sheaf\-FRL under stringent privacy regulatory constraints \(homogeneous case\)0:Oriented

𝒢=\(𝒩,ℰ\)\\mathcal\{G\}=\(\\mathcal\{N\},\\mathcal\{E\}\), pilots

𝒜\\mathcal\{A\},

λ\>0\\lambda\>0,

ηθ\>0\\eta\_\{\\theta\}\>0, iterations

TT
1:Initialize

\{𝜽i0=\(𝝋i0,𝝍i0\)\}i∈𝒩\\\{\\boldsymbol\{\\theta\}\_\{i\}^\{0\}=\(\\boldsymbol\{\\varphi\}\_\{i\}^\{0\},\\boldsymbol\{\\psi\}\_\{i\}^\{0\}\)\\\}\_\{i\\in\\mathcal\{N\}\},

\{𝐎j​i0∈O​\(d\)\}j∈𝒩​\(i\),i∈𝒩\\\{\\mathbf\{O\}\_\{ji\}^\{0\}\\in\\mathrm\{O\}\(\{d\}\)\\\}\_\{j\\in\\mathcal\{N\}\(i\),i\\in\\mathcal\{N\}\}
2:for

t=0,…,T−1t=0,\\dots,T\-1do

3:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

4:

𝐀it←\[f𝝋it​\(𝐱ik\)\]k∈𝒜\\mathbf\{A\}\_\{i\}^\{t\}\\leftarrow\[f\_\{\\boldsymbol\{\\varphi\}\_\{i\}^\{t\}\}\(\\mathbf\{x\}\_\{i\}^\{k\}\)\]\_\{k\\in\\mathcal\{A\}\}
5:endfor

6:Each node

iiexchanges

𝐎j​it​𝐀it\\mathbf\{O\}\_\{ji\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t\}with all

j∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)
7:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

8:for all

j∈𝒩​\(i\)j\\in\\mathcal\{N\}\(i\)do

9:

\[𝐔,𝚺,𝐖⊤\]←SVD​\(𝐎i​jt​𝐀jt​𝐀it⊤\)\[\\mathbf\{U\},\\mathbf\{\\Sigma\},\\mathbf\{W\}^\{\\top\}\]\\leftarrow\\mathrm\{SVD\}\\big\(\\mathbf\{O\}\_\{ij\}^\{t\}\\mathbf\{A\}\_\{j\}^\{t\}\\mathbf\{A\}\_\{i\}^\{t^\{\\top\}\}\\big\)
10:

𝐎j​it←𝐔𝐖⊤\\mathbf\{O\}\_\{ji\}^\{t\}\\leftarrow\\mathbf\{U\}\\mathbf\{W\}^\{\\top\}
11:endfor

12:endfor

13:for all

i∈𝒩i\\in\\mathcal\{N\}in paralleldo

14:

𝜽it\+1←𝜽it−ηθ​\(∇𝜽iℒi​\(𝜽it\)\+𝐫it\)\\boldsymbol\{\\theta\}\_\{i\}^\{t\+1\}\\leftarrow\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\-\\eta\_\{\\theta\}\\left\(\\nabla\_\{\\boldsymbol\{\\theta\}\_\{i\}\}\\mathcal\{L\}\_\{i\}\(\\boldsymbol\{\\theta\}\_\{i\}^\{t\}\)\+\\mathbf\{r\}\_\{i\}^\{t\}\\right\)
15:endfor

16:endfor

## Appendix GSupervised classification experimental details

All experiments run on MNIST with the label\-shift\-induced data heterogeneity of[Eq\.23](https://arxiv.org/html/2608.10016#S6.E23)\. Every agent trains with stochastic gradient descent \(learning rate10−210^\{\-2\}, momentum0\.90\.9, weight decay5⋅10−45\\cdot 10^\{\-4\}\) and batch size6464; the fifteen\-agent runs additionally clip gradient norms at1\.01\.0\. A held\-out shared pilot set \(10%10\\%of the training pool, restricted to the classes observed by every agent so that no encoder is asked to embed out\-of\-distribution samples\) serves two purposes: it provides the pilots exchanged during sheaf\-based training, and it is the data on which whitening layers of Sheaf\-FRL are trained, and \(post\-training\) alignment maps are fitted\. Reported accuracies are computed on the private test sets, different from both training and pilot data\.

Two\-agent heterogeneous pair \([Fig\.3](https://arxiv.org/html/2608.10016#S6.F3)\)\.The two agents are drawn from a parametric CNN family indexed by the bottleneck dimensiond∈\{16,32,64,128,256,512\}d\\in\\\{16,32,64,128,256,512\\\}, matched so that the last block of both encoders outputs the samedd\([Tab\.3](https://arxiv.org/html/2608.10016#A7.T3)\)\. Each convolutional block consists of a convolution, batch normalization, ReLU, and max\-pooling; the final feature map is globally pooled, so the latent dimension equals the width of the last block\. Agent0uses a dropout of0\.30\.3while agent11uses a dropout of0\.10\.1\. Agent0’s target classes are\{4,…,9\}\\\{4,\\dots,9\\\}and agent11’s are\{0,…,5\}\\\{0,\\dots,5\\\}, at a fixed shift strengths=0\.7s=0\.7; the graph has two nodes corresponding to the agents and a single edge connecting them\. Training lasts2020epochs with the common optimizer above\.

Fifteen\-agent network \([Fig\.2](https://arxiv.org/html/2608.10016#S6.F2)\)\.[Tab\.2](https://arxiv.org/html/2608.10016#A7.T2)lists the fifteen architectures and the target classes assigned to each agent\. All encoders use batch normalization; the latent \(stalk\) dimensiondid\_\{i\}equals the width of the last convolutional block\. The proposed methods run for a total of100100epochs\.

Table 2:The fifteen heterogeneous CNN agents of the multi\-agent benchmark \(Section[6](https://arxiv.org/html/2608.10016#S6)\)\. Encoder: output channels of each convolutional block; head: hidden widths of the MLP classifier;did\_\{i\}: latent \(stalk\) dimension\. Total:3\.413\.41M parameters\.Table 3:The shared two\-agent architecture family used at every bottleneck dimensionddin Fig\.[3](https://arxiv.org/html/2608.10016#S6.F3); both agents’ encoders are matched to end at widthdd\. Encoder/head: output channels of each convolutional block / hidden widths of the MLP classifier\. Dropout is0\.30\.3for agent 0 and0\.10\.1for agent 1\.Selection of the regularization coefficient\.For Sheaf\-FRL, as well as Sheaf\-FMTL, ComFed, FedProto, and FedMuscle baselines, we need to find the best regularization coefficient\. The applied strategy in each architecture configuration under consideration—either the bottleneck dimension in the two\-agent sweep, or the distribution\-shift strength in the fifteen\-agent benchmark—is a grid search for the bestλ\\lambdavalue in the\[10−5,1\]\[10^\{\-5\},1\]interval\.

Method\-specific hyperparameters\.Beyond the regularization coefficient, every baseline also carries its own additional, fixed hyperparameters that are untouched by the cross\-validation above\. Sheaf\-FMTL introduces two further constants, the local\-regularization weightγ=10−3\\gamma=10^\{\-3\}and the personalization rateη=10−2\\eta=10^\{\-2\}, both held fixed across every shift strength and every bottleneck dimension\. For further information to support the selection process of the Sheaf\-FMTL hyperparameterγ\\gammarefer to the next paragraph\. ComFed instead learns, for every agentii, an unconstrained linear projection matrix𝐏i∈ℝr×di\\mathbf\{P\}\_\{i\}\\in\\mathbb\{R\}^\{r\\times d\_\{i\}\}mapping that agent’s localdid\_\{i\}\-dimensional latent space—the right, input\-side dimension of𝐏i\\mathbf\{P\}\_\{i\}—into a shared space of dimensionrr, the left, output\-side dimension common to every agent, in which class prototypes from all agents are aligned\. We fixr=64r=64in the fifteen\-agent benchmark\. In the two\-agent bottleneck sweep, where every configuration already forces both agents to share a common latent dimensiondd, we instead setproj​\_​dim=d\\mathrm\{proj\\\_dim\}=dat each bottleneck value, so that ComFed’s shared space coincides with the bottleneck dimension already under study, rather than further compressing or expanding it\. FedProto maintains a single global prototype per class, updated as an exponential moving average of the agents’ local per\-class embeddings with momentum0\.90\.9\. FedMuscle organizes training into repeating communication rounds, each consisting of four local epochs of ordinary task\-loss training followed by one additional epoch in which a contrastive alignment loss, with temperature0\.10\.1, pulls every agent’s representations toward the network average\.

Memory footprint and choice ofγ\\gammafor Sheaf\-FMTL\.Let the number of parameters of agentsiiandjjbe respectivelydid\_\{i\}anddjd\_\{j\}, then the restriction maps matrix of Sheaf\-FMTL𝐏i​j\\mathbf\{P\}\_\{ij\}acts directly in parameter space, and its size isdi​j×did\_\{ij\}\\times d\_\{i\}withdi​j=max⁡\(1,⌊γ⋅min⁡\(di,dj\)⌋\)d\_\{ij\}=\\max\(1,\\lfloor\\gamma\\cdot\\min\(d\_\{i\},d\_\{j\}\)\\rfloor\)\. The aggregate storage across all edges grows asO​\(γ​∑\(i,j\)∈ℰmin⁡\(di,dj\)​di\)O\(\\gamma\\sum\_\{\(i,j\)\\in\\mathcal\{E\}\}\\min\(d\_\{i\},d\_\{j\}\)\\,d\_\{i\}\), with the hyperparameterγ\\gammacontrolling how fast the memory usage scales with the architecture sizes\. With the1515\-agent heterogeneous ensemble of models of[Tab\.2](https://arxiv.org/html/2608.10016#A7.T2)and the class\-overlap graph of[Sec\.6](https://arxiv.org/html/2608.10016#S6), this places a hard ceiling on the values ofγ\\gammathat fit in the machine memory \(GPU RAM\), well before any accuracy trade\-off becomes relevant\.[Fig\.4](https://arxiv.org/html/2608.10016#A7.F4)reports a blue curve providing the theoretical storage occupied by the restriction maps𝐏i​j\\mathbf\{P\}\_\{ij\}and model parameters in the real1515\-agent setup, and the green markers then represent the empirical peak resident\-memory measurements from isolated single\-epoch instantiations of Sheaf\-FMTL at escalatingγ\\gamma\. We useγ=10−3\\gamma=10^\{\-3\}throughout the1515\-agent benchmark \([Tab\.2](https://arxiv.org/html/2608.10016#A7.T2)\)—a margin below the maximum available memory on the considered machine that accounts for the full evaluation protocol rather than a single run\.

![Refer to caption](https://arxiv.org/html/2608.10016v1/x4.png)Figure 4:Aggregate memory footprint of the Sheaf\-FMTL restriction maps\{𝐏i​j\}\\\{\\mathbf\{P\}\_\{ij\}\\\}as a function ofγ\\gamma, for the1515\-agent architecture and communication graph of[Sec\.6](https://arxiv.org/html/2608.10016#S6)\(log\-log axes\)\. Solid curve: exact storage in the half\-precision format used at rest\. Dashed line: available system memory\. Markers: measured peak resident memory of an isolated, single\-epoch Sheaf\-FMTL run at eachγ\\gamma\(crosses mark values at which the run was terminated for exceeding available memory\)\. The operating point used throughout the paper \(γ=10−3\\gamma=10^\{\-3\}\) sits below the single\-run ceiling, leaving headroom for the repeated model instantiations of the full evaluation protocol \(hyperparameter search preceding each reported run\)\.