Understanding Head Geometry and Dynamics in Federated Regression through a Natural Solution Selection Rule: An Unconstrained Feature Model Analysis

Chuang Ma; Tomoyuki Obuchi

cs.LGSep 30, 2026

Understanding Head Geometry and Dynamics in Federated Regression through a Natural Solution Selection Rule: An Unconstrained Feature Model Analysis

Authors: Chuang Ma, Tomoyuki Obuchi

Organizations: Kyoto University · NII LLMC · RIKEN AIP

Abstract

In federated averaging, local objectives can admit multiple optimal heads, making the aggregate depend on which heads clients return. We study this ambiguity in federated multivariate regression with private backbones and a shared linear head, using an unconstrained feature model (UFM) that treats training-sample features as free variables. We introduce a natural selection rule: each client returns the optimal head closest to the broadcast head. We show that global minimization with a vanishing proximal penalty on the head realizes this rule. When the clients' optimal Gram matrices and the initial shared Gram matrix are positive definite, the shared Gram matrix follows a closed recursion and converges to the unique Bures-Wasserstein barycenter of the clients' optimal Gram matrices. Even with this alignment, the limit generally differs from the centralized optimal Gram matrix. We decompose this gap into three positive-semidefinite terms arising from differences in client target means, covariance heterogeneity, and averaging the aligned heads. A correction based on a one-time exchange of target means and covariances recovers the centralized optimal Gram matrix in one round under exact local optimization and the same selection rule. We verify these results numerically in the UFM and test its predictions on five tabular and five image regression datasets using deep networks with feature regularization and long local training. In these experiments, ordinary training approaches the predicted barycenter, while a weak proximal penalty improves endpoint agreement and yields trajectories that closely follow the predicted Gram dynamics. The correction moves the final Gram matrices close to the centralized UFM prediction.

Figures & tables

Figure 1: Federated averaging of a shared head under the selection rule. In the UFM, local optimality fixes the Gram matrix of a client’s head but not the head itself. Our selection rule chooses the optimal head closest to the broadcast head. Averaging these heads over communication rounds drives the Gram matrix of the shared head to the BW barycenter of the clients’ optimal Gram matrices. Inset: the gap to centralized training and its one-round correction from shared target moments.

Figure 2: UFM numerical checks on one fixed instance (

C=P=2

M=3

\lambda_{H}=\lambda_{W}=0.1

; Appendix D.1 ). Ellipses

\{G^{1/2}u:\left\lVert u\right\rVert_{2}=1\}

represent Gram matrices

G

, and dots trace trained ones.

E

counts full-batch local steps, and

\gamma

scales the moment correction. Errors below

10^{-14}

are drawn at the axis floor.

Figure 3: Proximal training with and without moment correction (

\rho=10^{-3}

). Colors indicate the reference Gram; solid and dashed curves denote training without and with correction. Curves show seed means

\pm 1

standard deviation. Below, outlines show

G_{\star}

and

G_{\mathrm{cen}}

, and markers show final Grams for seed 0. Figure 9 gives the other five settings.

Figure 4: Gram evolution on MPIIGaze. (a) Seed-0 Gram matrices under proximal training. (b) Relative errors to

G_{\star}

without correction (green) and to

G_{\mathrm{cen}}

with correction (graphite gray), with five-seed means

\pm 1

standard deviation.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Selection by a vanishing proximal regularizer (Lemma 1 ). Relative Gram error

e_{G}(\rho)

(circles) and selection error

e_{\Pi}(\rho)

(squares) in six

(C,P)

settings, with medians and interquartile ranges over 128 instances. The gray line has slope one. At

\rho=0

, the circle marks exact optimal heads and the squares mark unregularized L-BFGS returns from random starts.

Figure 6: Convergence to the BW barycenter (Theorem 1 ). Median relative error of the shared Gram matrix to

G_{\star}

over 16 instances in each of six problem sizes. Dashed: exact selection. Colored circles: proximal returns at

\rho=10^{-1},\ldots,10^{-6}

. Red squares: unregularized joint training from the broadcast head.

Figure 7: The three terms of the gap

G_{\mathrm{cen}}-G_{\star}

(Proposition 1 ). Blue, orange and green identify

\mathcal{M}_{\mu}

\mathcal{M}_{\Sigma}

and

\mathcal{M}_{A}

throughout; circles, squares and triangles mark the terms in (a–e) and (h–j), while in (g) they distinguish

M=2,4,8

. Curves are medians and bands are interquartile ranges. (a, b) Trace of each term when the client means or the client covariances are varied with strength

s

while the other is held fixed. The vertical axis is linear below

10^{-3}

so that exact zeros are visible. (c–e) Share of the gap trace carried by each term against the number of targets

C

, the number of clients

M

and a lower bound

\kappa

on the fully active margin. (f) Smallest eigenvalue of each term over 512 instances, sorted; numerical zeros are drawn at

10^{-16}

. (g)

\Tr\mathcal{M}_{A}

against the pairwise dispersion

D

, with the bounds

D

and

2D

of ( 17 ). (h–j) Smallest eigenvalue of each term over the sweeps of (c–e); dashed curves show the worst of the 256 instances in each setting.

Figure 8: One-round recovery of the centralized Gram matrix (Proposition 2 ). Median relative error of the shared Gram matrix to

G_{\mathrm{cen}}

over 32 instances in each of six problem sizes, under exact selection. Red squares: no correction; the plateau is the gap of Proposition 1 . Dashed: the correction

C_{m}

of ( 13 ). Colored: the scaled correction

\gamma C_{m}

, with deficit

1-\gamma

from

10^{-1}

10^{-5}

Dataset	Input dimension	Target dimension	Clients	Training samples	Test samples
Tabular datasets
Beijing	8	2	4	9,000	—
Swimmer	8	2	4	4,000	—
Hopper	11	3	4	4,000	—
HalfCheetah	17	6	4	4,000	—
Walker	17	6	4	4,000	—

Appendix

Table 1: Datasets and sample counts for the main DNN experiments. Test error is reported only for the image datasets (Appendix F.7 ).

$\checkmark$ : experiment repeated with five random seeds $\{0,1,2,3,4\}$
	Without correction		With moment correction
Setting	Ordinary	Proximal	Ordinary	Proximal
Tabular datasets
Beijing	$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$
Swimmer	$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$
Hopper	$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$

Appendix

Table 2: Main DNN experiments. Coverage by setting and procedure in Appendix F.1 .

$\checkmark$ : experiment repeated with three random seeds $\{0,1,2\}$
	F.2	F.3	F.4	F.6							F.7
Setting	Means	$E$	$\rho$	CNN	CNN $+C_{m}$	RN34	Data	Full $E=4000$	Full $E=400$	Full $+C_{m}$	Predict
Tabular datasets
Beijing	$\checkmark$	$\checkmark$	$\checkmark$	—	—	—	—	$\checkmark$	$\checkmark$	$\checkmark$	—
Swimmer	$\checkmark$	$\checkmark$	$\checkmark$	—	—	—	—	$\checkmark$	$\checkmark$	$\checkmark$	—
Hopper	$\checkmark$	$\checkmark$	$\checkmark$	—	—	—	—	$\checkmark$	$\checkmark$	$\checkmark$	—

Appendix

Table 3: Supplementary DNN experiments. Coverage by setting and condition; sweep columns give coverage at every parameter value.

	Without correction			With moment correction
Training	Endpoint $e_{F}$	Trajectory $e_{F}$	$\delta_{\rm last}$	Round 1 $d$	Round 1 $e_{F}$	Round 8 $e_{F}$
Ordinary	$1.63\times 10^{-2}$	$3.44\times 10^{-2}$	$1.85\times 10^{-1}$	$8.35\times 10^{-2}$	$3.08\times 10^{-1}$	$5.06\times 10^{-3}$
Proximal	$2.90\times 10^{-4}$	$3.39\times 10^{-4}$	$4.47\times 10^{-3}$	$4.21\times 10^{-3}$	$2.29\times 10^{-2}$	$7.08\times 10^{-5}$

Appendix

Table 4: Selection and moment correction. Medians over the same 50 dataset–seed combinations. Without correction, endpoint errors compare with

G_{\star}

and trajectory errors with the exact iteration from each run’s

G_{0}

. All corrected errors compare with

G_{\mathrm{cen}}

. Bold marks the smaller value in each column.

Figure 9: Proximal training in the other five settings. The design matches Figure 3 ; six-dimensional settings use coordinates 1–3. All runs use

\rho=10^{-3}

Figure 10: Ordinary training with and without moment correction. The design matches Figures 3 and 9 , with

\rho=0

. Curves show seed means

\pm 1

standard deviation; geometry uses seed 0. Both variants retain the feature and head penalties.

Figure 11: Direct comparison of ordinary and proximal training. Green curves measure training without correction against

G_{\star}

; graphite gray curves measure training with moment correction against

G_{\mathrm{cen}}

. Within each color, crosses and filled diamonds denote ordinary and proximal training; both use solid lines. Curves use the same five seeds per setting as Figures 3 and 10 , with means and one-standard-deviation bands. All panels share their axis limits.

Figure 12: Removing the mean contribution leaves a residual gap. Stacked bars separate the predicted mean contribution

\Tr(\mathcal{M}_{\mu})

from

\Tr(\mathcal{M}_{\Sigma}+\mathcal{M}_{A})

before and after centering each client’s targets. Both conditions use the original theoretical gap trace as denominator, computed per seed before averaging the components. Open markers show all three paired seeds, joined by thin lines; pale shading spans their range. Their observed gaps are

\Tr(G_{\mathrm{cen}}-G_{8})

before centering and

\Tr(G_{\mathrm{within}}-G_{8})

afterward, on the same scale.

Figure 13: Local epochs and the endpoint. Direction error

d(G_{8},G_{\star})

versus local epochs

E

, with shared logarithmic axes. Curves are medians over seeds

\{0,1,2\}

and bands span these seeds. Top row: ResMLP; bottom row: ResNet-18.

Proximal weight $\rho$	Distance to selected heads $\delta_{\rm last}$	Relative original-objective gap	$e_{F}(G_{8},G_{\star})$
$0$	$1.91\times 10^{-1}$	$1.43\times 10^{-5}$	$1.71\times 10^{-2}$
$10^{-5}$	$1.19\times 10^{-1}$	$1.49\times 10^{-5}$	$7.38\times 10^{-3}$
$10^{-4}$	$3.73\times 10^{-2}$	$1.53\times 10^{-5}$	$9.14\times 10^{-4}$
$10^{-3}$	$5.33\times 10^{-3}$	$4.05\times 10^{-5}$	$3.08\times 10^{-4}$
$10^{-2}$	$3.31\times 10^{-2}$	$1.40\times 10^{-3}$	$4.17\times 10^{-3}$
$10^{-1}$	$1.32\times 10^{-1}$	$1.81\times 10^{-2}$	$4.26\times 10^{-2}$

Appendix

Table 5: Proximal weight and local fitting. Medians at the final round over the same 30 dataset–seed combinations at each weight. Relative gaps use the original loss without the proximal regularizer and first take the median across clients. Bold marks the smallest value in each column.

Figure 14: Weight of the proximal regularizer. Endpoint relative matrix error

e_{F}(G_{8},G_{\star})

with

E=4000

. In each setting, curves are medians over seeds

\{0,1,2\}

and pale bands span these seeds. The zero-weight baseline is separate from the positive logarithmic axis; dashed guides mark

\rho=10^{-3}

	Without correction ( $n=30$ )			Corrected ( $n=18$ )
Training	$\delta_{\rm last}$	Trajectory $e_{F}$	Endpoint $e_{F}$	Round 1 $e_{F}$	Round 8 $e_{F}$
Ordinary	$1.87\times 10^{-1}$	$6.75\times 10^{-2}$	$1.86\times 10^{-2}$	$5.20\times 10^{-1}$	$7.07\times 10^{-3}$
Proximal	$6.35\times 10^{-3}$	$3.63\times 10^{-4}$	$3.24\times 10^{-4}$	$2.38\times 10^{-2}$	$7.06\times 10^{-5}$
Aligned	$1.54\times 10^{-4}$	$1.12\times 10^{-4}$	$1.10\times 10^{-4}$	$8.38\times 10^{-5}$	$6.27\times 10^{-5}$

Appendix

Table 6: Alignment, proximal training and ordinary training. Medians over 30 matched dataset–seed combinations without correction and 18 with moment correction. The distance

\delta_{\rm last}

and trajectory error are defined in Appendix E.4 . Without correction, endpoint error compares

G_{8}

with

G_{\star}

; corrected errors compare

G_{t}

with

G_{\mathrm{cen}}

. Bold marks the smallest value in each column.

Tabular datasets
	Fraction of the objective removed by alignment
Setting	Common backbone init.	Independent backbone init.
Beijing	15.7%	15.4%
Swimmer	10.1%	10.5%
Hopper	17.7%	18.1%
HalfCheetah	23.6%	23.7%

Appendix

Table 7: Cost of mismatched coordinates in one aggregation. The same first-round local models are averaged twice, with the heads as trained and after alignment to the broadcast head, giving training objectives ( 2 )

L_{\rm raw}

and

L_{\rm al}

. Each entry is the median over five seeds of

(L_{\rm raw}-L_{\rm al})/L_{\rm raw}

, the fraction of the objective that alignment removes. Columns give the initialization of the private backbones; all clients start from the same head. Tabular settings use ResMLP and image settings use ResNet-18.

Training / architecture	To $G_{\star}$	To $G_{\mathrm{cen}}$	Closer to $G_{\star}$
Architecture and data
Residual CNN	$1.45\times 10^{-2}$	$5.24\times 10^{-1}$	100%
Residual CNN + correction	$5.05\times 10^{-1}$	$1.93\times 10^{-3}$	0%
ResNet-34	$1.29\times 10^{-2}$	$5.19\times 10^{-1}$	100%
Larger data, $E=1000$	$9.55\times 10^{-3}$	$5.63\times 10^{-1}$	100%
Full-model averaging

Appendix

Table 8: Architecture, data size and full-model averaging. Median direction errors over all runs in each comparison. Closer to

G_{\star}

is the percentage of runs closer in direction to

G_{\star}

than to

G_{\mathrm{cen}}

; corrected runs target

G_{\mathrm{cen}}

. Architecture comparisons use

E=4000

unless specified; other training conditions are described in the text.

Figure 15: Local training and moment correction under full-model averaging. Each panel shows one tabular dataset with ResMLP and compares the same three seeds across ordinary training with

E=400

, ordinary training with

E=4000

, and corrected training with

E=4000

. Colors distinguish direction error to

G_{\star}

and

G_{\mathrm{cen}}

. Curves show medians; pale bands span the three seeds. The first transition changes the local budget; the second adds correction at a fixed budget.

Training	$e_{F}(G_{8},G_{\star})$	$e_{F}(G_{8},G_{\mathrm{cen}})$	Train loss change	Test loss change
Ordinary	$1.37\times 10^{-2}$	$5.04\times 10^{-1}$	–	–
Proximal	$1.54\times 10^{-4}$	$4.96\times 10^{-1}$	$-0.2\%$	$-0.1\%$
Proximal + moment correction	$6.23\times 10^{-1}$	$7.03\times 10^{-5}$	$-2.0\%$	$-1.5\%$

Appendix

Table 9: Geometry and prediction under one training recipe. Medians over the 15 image dataset–seed combinations. Geometry columns are relative matrix errors. The loss is the MSE, and loss changes are paired against ordinary training before taking medians. Negative changes indicate lower loss, which is better.

Figure 16: Paired test-MSE changes on the five image datasets. Values are

100(\mathrm{MSE}/\mathrm{MSE}_{\mathrm{O}}-1)

, with the same seed’s ordinary-training MSE as denominator for all procedures; O, P and C denote ordinary training, proximal training, and proximal training with moment correction. Curves show means over seeds

\{0,1,2\}

, and bands span one sample standard deviation; the dashed line is zero. Negative values indicate lower loss, which is better.

Figure 17: Numerical checks of finite-step profiled dynamics. (a) BW distances from

G_{\eta,E}

G_{\mathrm{within}}

(circles, solid) and

G_{\star}

(squares, dashed), normalized by

d_{\mathrm{BW}}(G_{\mathrm{within}},G_{\star})

; medians over three instances with

C=3

and

M=4

. (b) The same distances plotted against

\eta(E-1)

; the dotted guide has slope one. (c) Communication rounds needed to reach the stopping criterion. (d) Measured displacements divided by the predictions

\eta(E-1)\Delta_{i}

in 24 commuting configurations, colored by

\eta

. (e) Absolute errors in the first-order prediction

1-2h\theta_{\min}

for the Jacobian spectral radius; medians and interquartile ranges within each instance family over 780 stability records. The dotted guide has slope two.

Explore similar work

Sep 9, 2026cs.LG

Beyond Conventional Federated Learning via High-Order Regularization

Federated clients that perform several local optimization steps can return parameter displacements with widely different magnitudes. The quadratic regularization of FedProx grows linearly with displacement and therefore offers limited control over the contrast between ordinary and unusually large client movements. We here introduce HiFedProx, which replaces the quadratic penalty with a scale-matched power-type regularizer indexed by

p\geq2

. All powers have the same regularization-gradient magnitude at a reference displacement

R

, while every

p>2

gives a weaker response below

R

and a stronger response above it. An exact affine reference calculation shows that increasing

p

compresses relative displacement disparities, although very large powers approach fixed-radius behavior and increase local curvature. HiFedProx combines this geometry with finite-budget stochastic client optimization and same-minibatch Armijo backtracking. In paired five-seed experiments on a frozen 60-writer FEMNIST subset, a common-parameter study over

p\in\{2,3,4,5,6,7,8\}

shows similar clean-training performance but substantial gains under composite stress. The lowest moderate- and severe-stress losses occur at

p=7

and

p=6

, improving over

p=2

11.44\%

and

23.16\%

, respectively. Although displacement-tail ratios continue to decrease through

p=8

, predictive performance peaks in an intermediate range and Armijo trial cost increases with

p

. These results indicate that the exponent should be calibrated rather than maximized. In our experiments,

p=5

7

provides the most useful range.

Alireza Kabgani, Masoud Ahookhosh

Department of Mathematics, University of Antwerp, Antwerp, Belgium

Federated Learning Regularization

Ask PDF

Nov 30, 2025cs.LG

FDRMFL: Multimodal Federated Feature Extraction Model Based on Information Maximization and Contrastive Learning

We propose FDRMFL, a task-driven multimodal feature extraction framework for federated regression under non-IID data distributions. Extracting predictive features from high-dimensional multimodal inputs is particularly challenging in this setting: data cannot leave each client, local samples are scarce and heterogeneously distributed, and unsupervised dimensionality reduction discards task-relevant information while federated training introduces representation drift across communication rounds. FDRMFL addresses these challenges through a unified four-term local objective: MSE prediction loss, a correlation-based mutual information surrogate that preserves dependence between the fused representation and the continuous target, a symmetric KL penalty that aligns cross-modal latent distributions before fusion, and an InfoNCE-style contrastive loss that anchors local representations to the global consensus. Experiments on three synthetic and two real-world near-infrared spectroscopy datasets under non-IID federated partitions, with comprehensive ablation and sensitivity analyses, demonstrate that each component contributes to the framework's effectiveness. FDRMFL reduces mean MSE by 33.8% relative to the best traditional baseline (PCA) and by 43.0% relative to VAE in simulation, and attains the lowest overall mean MSE among six federated algorithms including FedAvg, FedProx, MOON, SCAFFOLD, and FedBN.

Haozhe Wu

School of Computing and Data Science, The University of Hong Kong, Hong Kong SAR, China

Multimodal Fusion Cross-Modal Contrastive Loss Contrastive Learning Extraction

Ask PDF

May 12, 2026cs.LG

FedOUI: OUI-Guided Client Weighting for Federated Aggregation

Federated learning usually aggregates client updates using dataset size or gradient-level criteria, while overlooking internal signals about how each client model is organizing its input space during training. We introduce FedOUI, a simple aggregation rule based on the Overfitting-Underfitting Indicator (OUI), an activation-based and label-free metric. Each participating client sends its local update together with a OUI value computed on a fixed probe batch, and the server estimates the round-wise OUI distribution to assign lower weights to structurally atypical clients through a smooth reweighting rule. We evaluate FedOUI on CIFAR-10 under strong non-IID partitioning and noisy-client conditions, comparing it with FedAvg, FedProx, and a gradient-alignment baseline. The clearest gains appear under strong heterogeneity, where OUI-based weighting improves aggregation quality while remaining lightweight and interpretable. These results show that internal activation structure can provide useful information for federated aggregation beyond client size and gradient geometry.

Alberto Fernández-Hernández, Jose I. Mestre, Cristian Pérez-Corral +3

Universitat Politècnica de València, Valencia, Spain · Universitat Jaume I, Castelló de la Plana, Spain · Openchip & Software Technologies S.L., Spain

Federated Learning Fedavg Balanced Learning Overfitting Clients

Ask PDF