Technical Summary: Distributed Prediction under Heterogeneity with Unidentifiable Parameter

1. Problem Statement

The paper addresses the fundamental challenge of predicting a response $Y$ based on covariates $x$ in distributed settings where the underlying structural parameters are unidentifiable. This scenario arises in sufficient dimension reduction (SDR) contexts, such as 1-bit compressed sensing, transformation models, and single-index models, where the goal is to recover a low-dimensional structure $E(Y|x) = E(Y|x^T\beta)$.

The problem is compounded by three intertwined barriers:

1. Parameter Non-Identifiability: In these models, the parameter $\beta$ is only identifiable up to its column space (direction), not its magnitude. Standard Euclidean metrics fail to capture similarity between parameters that share the same column space but differ in scale. Quantifying similarity requires projection matrices $P(\beta) = \beta(\beta^T\beta)^{-1}\beta^T$, which introduces severe nonconvexity into the optimization landscape.
2. Data Heterogeneity: In distributed systems (e.g., federated learning), local nodes often possess distinct structural parameters. Existing aggregation methods relying on Euclidean distances are ineffective for non-identifiable parameters.
3. Communication Overhead: Distributed training is constrained by bandwidth limitations. Standard iterative algorithms that exchange parameters frequently incur prohibitive communication costs, especially when dealing with high-dimensional data.

2. Methodology

The authors propose a novel distributed semiparametric framework called InvexDR (Invex Dimensional Reduction) to systematically resolve these challenges.

2.1. Regularized Optimization with Trace Penalty

To handle heterogeneity, the method formulates a global objective that aggregates local empirical least squares losses with a trace-similarity penalty.

• Local Loss: Each node $j$ minimizes a semiparametric least squares loss derived from a Newton-Raphson update, estimating the link function $m(\cdot)$ and its gradient nonparametrically.
• Trace Penalty: Instead of Euclidean distance, the penalty term $P_j(\alpha_j) = -\sum_{l \neq j} \text{tr}\{P(\alpha_j)P(\beta_l)\}$ is used. This penalty maximizes the inner product of projection matrices, effectively aligning the column spaces of local parameters to pursue homogeneity without enforcing rigid geometric constraints (like unit spheres) that induce nonconvexity.

2.2. Invex Relaxation

The trace penalty renders the objective function nonconvex. To overcome this, the authors introduce an invex relaxation:

• Augmentation: The parameter space is expanded from $\beta \in \mathbb{R}^{p \times d}$ to an augmented matrix $A \in \mathbb{R}^{(p+d) \times md}$.
• Transformation: The original nonconvex problem is reformulated into a constrained optimization problem over $A$. The authors prove that this relaxed objective is $\eta$-invex.
• Significance of Invexity: Unlike standard nonconvex problems, an invex function guarantees that every stationary point satisfying the Karush-Kuhn-Tucker (KKT) conditions is a global minimum. This ensures that gradient-based algorithms converge to the global optimum regardless of initialization, bypassing the local minima traps typical of nonconvex SDR problems.

2.3. Communication-Efficient Distributed Algorithm

The algorithm employs a bi-level architecture:

• Outer Loop: Updates nonparametric components (link functions and gradients) using a Newton-Raphson procedure.
• Inner Loop: Performs multi-step local updates on the invex relaxation. Nodes execute $K$ local gradient descent steps before communicating.
• Communication Strategy: Instead of exchanging raw parameters, nodes share structural proxies (aggregated projection matrices). The global synchronization of these proxies occurs only every $R$ local steps, significantly reducing communication overhead to $O(pK^{1/2})$.

3. Key Contributions

1. Adaptive Homogeneity Pursuit: The integration of a trace-similarity penalty with Newton-Raphson least squares allows for the adaptive aggregation of heterogeneous, non-identifiable parameters without imposing restrictive geometric constraints (e.g., $\|\beta\|=1$).
2. Resolution of Nonconvexity: The proposed invex relaxation transforms the inherently nonconvex problem into one where KKT conditions are sufficient for global optimality. This provides a stable optimization landscape that guarantees convergence from arbitrary initializations.
3. Theoretical Guarantees:
   • Statistical Optimality: The estimator achieves a minimax-optimal convergence rate of $O_p(N^{-1/2} + H \wedge n^{-1/2})$, where $N$ is the total sample size, $n$ is the local sample size, and $H$ quantifies heterogeneity. This rate matches the fundamental limits established for identifiable parameters.
   • Model-Free Prediction: A non-asymptotic upper bound for model-free prediction error is established, showing the method achieves superior predictive accuracy compared to local-only estimators.
   • Algorithmic Convergence: The paper proves stable convergence of the distributed algorithm and bounds the communication cost, demonstrating that multi-step local updates do not compromise statistical consistency.

4. Results

Numerical Simulations

Extensive simulations across varying sample sizes, network scales, and heterogeneity levels (controlled by rotation angles) demonstrate:

• Superiority: InvexDR consistently outperforms local baselines (MAVE, NR-B, NR-O) and pooled distributed methods.
• Adaptivity: The method adaptively leverages latent similarities; performance gains increase as homogeneity increases.
• Robustness: Unlike pooled methods that degrade under high heterogeneity or excessive node counts, InvexDR maintains stability. It avoids the performance degradation seen in constrained methods (e.g., block identity or orthogonal constraints) which fail to capture the true parameter structure.

Real-World Application

The framework was applied to the eICU Collaborative Research Database to predict the Remaining Length of Stay (RLOS) for ICU patients across multiple hospitals.

• Setup: Hospitals acted as distributed nodes with strict privacy constraints preventing data pooling.
• Performance: InvexDR achieved the lowest mean squared error (MSE) and standard deviation compared to isolated and constrained distributed variants.
• Impact: The method demonstrated robust predictive accuracy even with severe local sample size imbalances, validating its utility in authentic distributed healthcare settings.

5. Significance and Claims

The paper claims to provide a robust, communication-efficient framework for a class of problems previously considered intractable in distributed settings due to the combination of non-identifiability and nonconvexity.

• Theoretical Breakthrough: By establishing that the relaxed problem is invex, the authors claim to have solved the "optimization barrier" inherent in dimension reduction, ensuring global convergence without the need for careful initialization or restrictive parameter constraints.
• Statistical Efficiency: The method achieves minimax optimality, proving that distributed learning can recover the same statistical efficiency as centralized learning even when parameters are unidentifiable and data is heterogeneous.
• Practical Utility: The successful application to multi-center medical data demonstrates that the framework is not merely theoretical but offers tangible improvements in real-world scenarios where data privacy and heterogeneity are paramount.

The authors position this work as a critical step toward enabling collaborative learning in domains where structural parameters are inherently unidentifiable, such as compressed sensing, econometrics, and bio-statistics.