HYPERBOLIC MODEL AGGREGATION FOR FEDERATED LEARNING IN FMRI

Abstract

The privacy-sensitive nature of clinical data often limits the use of machine learning in medical imaging applications, particularly for modalities with high acquisition costs such as functional MRI (fMRI). Federated learning mitigates data-sharing barriers by training site-specific models locally and aggregating weights centrally into a server model. However, small and heterogeneous per-site samples in medical imaging heighten the need for robust model-aggregation strategies. In this work, we introduce a federated aggregation scheme based on hyperbolic geometry to provide a robust and flexible approach to federated model weight integration. The proposed scheme is plug-and-play for standard federated learning loops. Empirically, our method improves stability and accuracy across multi-site fMRI data from ABIDE I, yielding more consistent convergence versus methods based on Euclidean mean and median. Codes are publicly available at https://github.com/Jiyao96/FedHAvg.

1. INTRODUCTION

Recent large public fMRI datasets, such as the Human Connectome Project (HCP) [1], have provided ample data to study brain function in healthy populations. However, collecting sizable homogeneous cohorts for specific neurological disorders within narrow age ranges remains challenging.

To address this, federated learning (FL) [2] enables knowledge aggregation from small multi-site datasets while preserving the privacy of raw data. In each standard FL communication round, clients train locally and transmit model parameters or gradients rather than raw subject-level data to a central server. The server aggregates these updates to form a global model, which is then broadcast back to all clients (Fig. 1). Without direct sharing of raw data, the global model learns robust and generalizable parameters applicable across all clients.

Fig. 1.

Federated Learning

Several studies have applied FL to multi-site fMRI data. Li et al. [3] utilized vectorized upper-triangular entries of the functional connectome (FC) as input features. Building on this, Wang et al. [4] and Zhang et al. [5] incorporated multimodal inputs, including FCs from multiple region of interest (ROI) parcellations, demographic variables, and cognitive test scores. Despite differences in inputs and architectures, all three studies employed FedAvg [2] to aggregate local weights for updating the global model. While FedAvg [2] is a widely used aggregation baseline, it is susceptible to a small number of corrupted or highly biased updates under non-IID client drift. Prior medianbased work [6] mitigates this by discarding outliers. However, when the number of participating clients is small, median-based estimators can be unstable and high-variance, reducing their practical reliability in medical imaging applications.

Hyperbolic geometry has been leveraged across multiple aspects of FL. HyperFed [7] projects features into a shared hyperbolic embedding space to facilitate clustering. FlatLand [8] assigns client-specific curvatures to capture site heterogeneity and enable per-client specialization. However, applying hyperbolic transformation directly to server-side weight aggregation remains unexplored to our knowledge. In this paper, we propose Federated Hyperbolic Averaging (FedHAvg), a geometry-aware robust aggregation method for FL. Through hyperbolic transform of local weights before averaging, extreme updates are compressed towards the boundary of the Poincaré ball. The curvature serves as a tunable robustness parameter to control the strength of compression to the target application.

2. DATA

For this study, we use the PCP-preprocessed resting-state fMRI from the multi-site ABIDE I [9] dataset consisting of autism spectrum disorder (ASD) and healthy control (HC) subjects. For the FL setup, we select data from the top 4 sites with the largest number of subjects (NYU, UM, UCLA, USM). Following common practice [10], we merge subsites acquired at the same institution under identical rs-fMRI protocols and treat each institution as a single federated client (UM1 merged with UM2 and UCLA1 merged with UCLA2). For quality control, we exclude scans with a mean frame distance larger than 0.15mm. The total and per-site number of subjects are summarized in Table 1. We employ the PCP-preprocessed fMRIs parcellated into 116 ROIs following the AAL atlas [11]. For each subject, we use the flattened upper triangle of the FC matrix as input, yielding 116×115/2 = 6,670 features. Subjects at each site are stratified by diagnosis and partitioned into five approximately equal folds. For simplicity in evaluation, we employ a synchronized 5-fold cross-validation across sites.

Table 1.

Total and per site number of subjects in ABIDE

Site	Total	ASD		HC
NYU	163	66	(40.5%)	97	(59.5%)
UM	104	41	(39.4%)	63	(60.6%)
UCLA	72	33	(45.8%)	39	(54.2%)
USM	50	31	(62.0%)	19	(38.0%)
Total	389	171	(44.0%)	218	(56.0%)

3. METHOD

3.1. Notation and Problem Definition

We consider $K$ clients indexed by $𝒦=\{1,\dots ,K\}$. Each client $k$ holds a local dataset ${𝒟}_k={\left\{\left(x_{\mathit{\text{ki}}},y_{\mathit{\text{ki}}}\right)\right\}}_{i=1}^{n_k}$ with $n_k$ samples. The server maintains a global model $w\in 𝒲$, where $𝒲={\mathbb{R}}^d$. In FL training, we initialize the global model weights ${\mathbf{w}}^0$ and iteratively update the global model through local client training and global server aggregation (Algorithm 1). In the client update, returning the gradient or the updated weights is equivalent for strategies based on Euclidean mean or median.

For our ASD versus HC binary classification application, the global model is a binary classifier $f$ that maps the input $x$ to binary labels given the federated trained weights $\mathbf{w}$. In this paper, we investigate the server-side aggregation function while using a baseline MultiLayer Perceptron (MLP) classifier.

$$f(x,\mathbf{w})\to \{0,1\}$$

3.2. Federated Hyperbolic Averaging (FedHAvg)

We leverage the closed form expression of geometric operations in the Poincaré ball to enable stable and vectorized implementations of transformations. In hyperbolic geometry, the $d$-dimensional Poincaré ball model with constant sectional curvature $-c(c>0)$ is defined as

$${\mathbb{D}}_c^d=\left\{x\in {\mathbb{R}}^d:\Vert x\Vert <1/\sqrt{c}\right\}.$$

At the origin, the exponential and logarithmic maps provide forward and backward transformations between the Euclidean space ${\mathbb{R}}^d$ and the Poincaré ball ${\mathbb{D}}_c^d$. Using $\gamma >0$ as a scalar to regularize the vector norm, the closed-form exponential map at the origin can be described as

$${\text{exp}}_0^{\left(c\right)}(\gamma v)=\left\{\begin{array}{ll}\frac{\text{tanh}(\sqrt{c}\gamma \Vert v\Vert )}{\sqrt{c}\Vert v\Vert }v,& v\ne 0,\\ 0,& v=0,\end{array}\right.$$

Similarly, the closed-form logarithmic map can be described as

$$\frac{1}{\gamma }{\text{log}}_0^{\left(c\right)}(y)=\left\{\begin{array}{ll}\frac{1}{\gamma \sqrt{c}}\text{artanh}(\sqrt{c}\Vert y\Vert )\frac{y}{\Vert y\Vert },& y\ne 0,\\ 0,& y=0.\end{array}\right.$$

According to the above definition, we propose two aggregation strategies: one utilizing the extrinsic mean and the other based on the gyro-barycentric mean.

3.2.1. Extrinsic Mean

In the Poincaré ball, the extrinsic mean follows a project-average-retract construction that maps to the manifold with ${\text{exp}}_0$, takes the Euclidean weighted average, and maps back to the tangent space by ${\text{log}}_0$. In our implementation, we use an unweighted Euclidean average for simplicity. With ${\Delta }_k$ denoting the local updates, the resulting extrinsic mean is

$${\mu }_{\text{ext}}={\text{log}}_0^{\left(c\right)}\left(\frac{1}{K}\sum _{k=1}^K{\text{exp}}_0^{\left(c\right)}\left({\Delta }_k\right)\right).$$

3.2.2. Gyro-barycentric Mean

While the extrinsic mean is robust and computationally efficient, it is neither coordinate-invariant nor equivariant under hyperbolic isometries. As an intrinsic alternative, we can aggregate weights using the gyro-barycentric mean.

By definition, the closed-form unweighted gyro-barycentric mean of points in the Poincaré ball is computed via Möbius addition ${\oplus }_c$ and Möbius scalar multiplication ${\otimes }_c$ as

$${\mu }_{\text{gbar}}=\underset{k=1}{\overset{K}{⨁}}\frac{1}{K}{\otimes }_c{\Delta }_k,$$

where the closed-form Möbius operations are

$$x{\oplus }_{c}y=\frac{(1+2c⟨x,y⟩+c\Vert y{\Vert }^2)x+(1-c\Vert x{\Vert }^2)y}{1+2c⟨x,y⟩+c^2\Vert x{\Vert }^2\Vert y{\Vert }^2},$$

$$r{\otimes }_{c}x=\left\{\begin{array}{ll}\frac{\text{tanh}\left(r\text{artanh}\right(\sqrt{c}\Vert x\Vert \left)\right)}{\sqrt{c}\Vert x\Vert }x,& x\ne 0,\\ 0,& x=0.\end{array}\right.$$

The gyro-barycentric mean not only yields a chart-free and geometric isometry-friendly averaging of client parameters, but also offers less attenuation of large-norm updates compared to the extrinsic mean.

In implementation, to prevent division-by-zero and instability of ${\text{log}}_0$ near the ball boundary, we add a small $ϵ=1\times {10}^{-12}$ to all denominators and clamp argument of artanh by enforcing a minimum margin from the boundary $d_{\mathit{\text{min}}}=1\times {10}^{-5}$.

4. EXPERIMENTS

4.1. Implementation Details

Our global model is a four-layer MLP with LayerNorm [12] and ReLU activation. BatchNorm is excluded to prevent mixing client-specific feature statistics following [13]. We train the global model for 100 federated rounds, with 1 local epoch per client in each round. Each local client is trained with a batch size 16 using SGD optimizer with a learning rate of 0.0001, momentum of 0.9, and weight decay of 0.001. All models are trained for the ASD/HC classification with cross-entropy loss.

4.2. Per-Site Training with Cross-Site Evaluation

To assess the ability of single-source models to be applied on other sites, we train a model on one site and evaluate it on all other sites (Table 2). Overall, cross-site adaptation is weak, with models trained on larger sites generalizing slightly better. Training on NYU and evaluating on UCLA yields the best cross-site accuracy (0.667), underscoring substantial cross-site heterogeneity in the ABIDE dataset.

Table 2.

Accuracy in ASD/HC classification task from single-site training and inference on other sites.

Site	NYU	UM	UCLA	USM
NYU-trained	-	0.606	0.667	0.520
UM-trained	0.663	-	0.597	0.480
UCLA-trained	0.650	0.529	-	0.400
USM-trained	0.491	0.452	0.486	-

4.3. Federated Learning

We evaluate the performance of the baselines and our geometry-aware federated model aggregation approaches (Fig. 2).

Fig. 2.

Baseline and geometry-aware federated model aggregation methods. In a), FedAvg and FedAvg-weighted correspond to mean-based aggregation [2]. In b), Trimmed-mean and Coordinate-wise Median are median-based aggregation [6]. c) describes two variants of the proposed FedHAvg approaches (FedHAvg-ext and FedHAvg-gbar).

4.3.1. Baselines

We assess FL performance with the following four baseline model aggregation methods:

• Federated Averaging (FedAvg) [2]: Compute the global parameters as the unweighted Euclidean average of all the client parameters.

• FedAvg-weighted [2]: Sample-size–weighted Euclidean average of client parameters. Each client’s contribution is proportional to its local training set size. In our experiments this provided only marginal improvement over the unweighted mean. Hence, for simplicity and clearer comparisons, we default to unweighted variants for other eligible methods.

• Coordinate-wise Median [6]: For each parameter channel, set the global value to the median of the client parameter values at that coordinate.

• Trimmed-mean [6]: Sort the $K$ client values for each parameter channel, discard the lowest and highest $\lfloor \alpha K\rfloor$ values with $\alpha =0.25$, and average the remaining $K-2\lfloor \alpha K\rfloor$ values. This yields robustness to a bounded fraction of corrupted or extreme updates, a softer version of the medianonly aggregation.

4.3.2. Federated Hyperbolic Averaging (FedHAvg)

As described in Section 3.2, we experiment with two geometry-aware federated model aggregation variants: FedHAvg-ext and FedHAvg-gbar, one using the extrinsic mean and the other using gyro–barycentric mean, respectively.

For aggregation methods involving a nonlinear hyperbolic projection, averaging client weights (${\Delta }_k=\mathbf{u}$) is not equivalent to averaging their updates (${\Delta }_k=\mathbf{u}-\mathbf{w}$). Therefore, we evaluate both weight-based and update-based aggregation, denoting the latter as “on $\Delta$”.

We empirically fix the curvature at $c=2$ for all variants. To harmonize scales of ${\ell }_2$ norms between weight-based and update-based inputs, we empirically set $\gamma =1$ for weight-based aggregation and $\gamma =100$ for update-based aggregation.

5. RESULTS AND DISCUSSION

Results are summarized in Table 3. The proposed geometry-aware aggregators secure the majority of top-2 performances. Exceptions occur only for the second-best Acc and F1 on NYU, and the best AUC on UCLA and USM. This shows our methods to outperform mean and median based model aggregation methods on the 4-site federated learning experiment on the ABIDE dataset. Through hyperbolic projection, client updates are implicitly reweighted according to geometric statistics of the update vectors. The resulting aggregator is more robust than FedAvg yet does not incur the higher variance associated with median-based schemes that rely on fewer effective clients.

Table 3.

Accuracy in ASD/HC classification task under federated learning approaches. Results are reported as mean of 5-fold cross-validation experiments measured by Accuracy (Acc) and Area-Under-the-ROC-Curve (AUC). Best performance in each column are bolded, second best underlined.

Method	NYU			UM			UCLA			USM
	Acc	F1	AUC	Acc	F1	AUC	Acc	F1	AUC	Acc	F1	AUC
FedAvg	0.708	0.631	0.731	0.652	0.603	0.733	0.657	0.539	0.670	0.635	0.651	0.749
FedAvg-weighted	0.695	0.597	0.744	0.661	0.604	0.730	0.643	0.499	0.708	0.635	0.651	0.717
Median	0.708	0.636	0.722	0.654	0.597	0.713	0.656	0.553	0.627	0.635	0.651	0.696
Trimmed-mean	0.720	0.653	0.733	0.652	0.599	0.726	0.656	0.552	0.649	0.635	0.651	0.722
FedHAvg-ext on Δ	0.726	0.669	0.726	0.661	0.615	0.721	0.629	0.522	0.633	0.675	0.702	0.733
FedHAvg-gbar on Δ	0.695	0.609	0.744	0.652	0.596	0.726	0.643	0.533	0.701	0.655	0.677	0.744
FedHAvg-ext	0.682	0.591	0.753	0.682	0.614	0.753	0.724	0.606	0.696	0.655	0.677	0.710
FedHAvg-gbar	0.695	0.602	0.757	0.662	0.626	0.736	0.705	0.604	0.692	0.672	0.712	0.675

Across all experimented FedHAvg variants, weight-based aggregation consistently surpasses the update-based alternatives (“on $\Delta$”). A plausible cause is the high variance in ${\ell }_2$ norms of the updates when it is computed as the difference between the locally trained weights and the global weights.

We observe no significant difference between extrinsic-mean and gyro-barycentric mean aggregation. In theory, the extrinsic mean should impose stronger shrinkage on outlying client weights than the gyro-barycenter, but this effect did not translate into a measurable performance gap in our setting.

6. CONCLUSION

We introduce federated hyperbolic averaging (FedHAvg), a family of geometry-aware aggregation schemes. By exploiting the expansion and compression properties of hyperbolic geometry, FedHAvg attenuates large-magnitude or outlier updates while preserving the directions of all client updates.

While our experiments demonstrate successful federated learning in a 4-site setting, scaling to all 17 ABIDE sites remains challenging, particularly for sites with small sample sizes or severe label imbalance. To incorporate more sites, our future work will pursue robustness improvements for federated learning from both the server-side aggregation and client-side specification perspectives.