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

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.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

More in Machine Learning

• arXiv preprint · 30 September 2026

  ReGain: Restoring Subject Fidelity in Personalization on Synthetic Images

  Shubhang Bhatnagar, Ishan Bhatnagar, Viraj Shah and 1 other
• arXiv preprint · 1 October 2026

  SmoothOperator: Enhancing Representations for Fine-grained Open-set Recognition via Modulated Label Smoothing

  Thiru Thillai Nadarasar Bahavan, Yu Xia, Sachith Seneviratne and 1 other
• arXiv preprint · 30 September 2026

  Code to Control: Synthesizing Parameterized Reactive Controllers

  Zergham Ahmed, Joshua B. Tenenbaum, Chris Bates and 1 other
• arXiv preprint · 1 October 2026

  Fixing a Model That Learned Worse Cancer Means Lower Risk: Monotonic Constraints in Bladder Cancer Recurrence Prediction

  Saram Abbas, David Thomas, Naeem Soomro and 3 others
• arXiv preprint · 30 September 2026

  Molecular Property Prediction under Structural Shift with Tabular Foundation Models

  Jinmo Lee, Dooho Lee, Minho Jeong and 1 other

All new papers in Machine Learning