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

arXiv: https://arxiv.org/abs/2609.39464 · HTML: https://mlsanity.com/paper/2609.39464/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 43 of its 58 references that are in the map.
- First posted on arXiv 30 Sep 2026.

## Builds on (most cited first, 30 of 43)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Think Locally, Act Globally: Federated Learning with Local and Global Representations](https://mlsanity.com/paper/2001.01523/) (2020)
- [No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data](https://mlsanity.com/paper/2106.05001/) (2021)
- [A Geometric Analysis of Neural Collapse with Unconstrained Features](https://mlsanity.com/paper/2105.02375/) (2021)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Exploring Deep Neural Networks via Layer-Peeled Model: Minority Collapse in Imbalanced Training](https://mlsanity.com/paper/2101.12699/) (2021)
- [Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path](https://mlsanity.com/paper/2106.02073/) (2021)
- [Appearance-Based Gaze Estimation in the Wild](https://mlsanity.com/paper/1504.02863/) (2015)
- [Neural collapse with unconstrained features](https://mlsanity.com/paper/2011.11619/) (2020)
- [FedSplit: An algorithmic framework for fast federated optimization](https://mlsanity.com/paper/2005.05238/) (2020)
- [Extended Unconstrained Features Model for Exploring Deep Neural Collapse](https://mlsanity.com/paper/2202.08087/) (2022)
- [On the Optimization Landscape of Neural Collapse under MSE Loss: Global Optimality with Unconstrained Features](https://mlsanity.com/paper/2203.01238/) (2022)
- [Imbalance Trouble: Revisiting Neural-Collapse Geometry](https://mlsanity.com/paper/2208.05512/) (2022)
- [FedExP: Speeding Up Federated Averaging via Extrapolation](https://mlsanity.com/paper/2301.09604/) (2023)
- [Visual Representation Learning Does Not Generalize Strongly Within the Same Domain](https://mlsanity.com/paper/2107.08221/) (2021)
- [FedAvg with Fine Tuning: Local Updates Lead to Representation Learning](https://mlsanity.com/paper/2205.13692/) (2022)
- [No Fear of Classifier Biases: Neural Collapse Inspired Federated Learning with Synthetic and Fixed Classifier](https://mlsanity.com/paper/2303.10058/) (2023)
- [On the emergence of simplex symmetry in the final and penultimate layers of neural network classifiers](https://mlsanity.com/paper/2012.05420/) (2020)
- [Convergence and Accuracy Trade-Offs in Federated Learning and Meta-Learning](https://mlsanity.com/paper/2103.05032/) (2021)
- [Towards Understanding and Mitigating Dimensional Collapse in Heterogeneous Federated Learning](https://mlsanity.com/paper/2210.00226/) (2022)
- [Deep Neural Collapse Is Provably Optimal for the Deep Unconstrained Features Model](https://mlsanity.com/paper/2305.13165/) (2023)
- [Neural Collapse for Unconstrained Feature Model under Cross-entropy Loss with Imbalanced Data](https://mlsanity.com/paper/2309.09725/) (2023)
- [Perturbation Analysis of Neural Collapse](https://mlsanity.com/paper/2210.16658/) (2022)

## Abstract (arXiv)

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.