🤖 machine learning

Federated Imputation under Heterogeneous Feature Spaces

The paper proposes FedHF-Impute, a federated imputation framework that utilizes a shared global feature graph to enable indirect cross-client knowledge transfer and significantly improve accuracy in heterogeneous feature spaces where clients observe only partially overlapping feature subsets.

Original authors: Imane Hocine, Chaimaa Medjadji, Sylvain Kubler, Gregoire Danoy, Yves Le Traon

Published 2026-05-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Imane Hocine, Chaimaa Medjadji, Sylvain Kubler, Gregoire Danoy, Yves Le Traon

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a group of doctors trying to build a single, super-smart medical diagnosis tool. They want to work together without ever sharing their actual patient records (to protect privacy). This is the world of Federated Learning.

However, there's a big problem: Not every doctor has the same tools.

Doctor A has a thermometer, a blood pressure cuff, and a stethoscope.
Doctor B has a thermometer and an X-ray machine, but no blood pressure cuff.
Doctor C has a stethoscope and a blood test kit, but no thermometer.

In the old way of doing things (standard Federated Learning), everyone tries to average their answers. But if Doctor A is trying to guess a patient's blood pressure and Doctor B never measures blood pressure, they can't really help each other. It's like trying to solve a puzzle where half the pieces are missing from everyone's box, and the boxes don't even match up. The result is a confused, inaccurate model.

The Solution: FedHF-Impute

The authors of this paper created a new system called FedHF-Impute. Think of it as a "Universal Translator" for medical tools.

Here is how it works, using simple analogies:

1. The "Feature Graph" (The Master Map)

Before the doctors start training, they send the server a simple list: "I have a thermometer and a stethoscope." They don't send patient data.
The server uses this list to draw a Master Map (called a Feature Graph). This map connects tools that usually go together.

Analogy: The map knows that "Blood Pressure" and "Heart Rate" are best friends. Even if Doctor B never sees "Blood Pressure," the map knows that if Doctor B sees a high "Heart Rate," it can guess what the "Blood Pressure" might be based on what Doctor A knows about Heart Rate.

2. The "Message Passing" (The Relay Race)

Instead of just averaging numbers, the system lets information flow along the lines of the Master Map.

Analogy: Imagine a relay race. Doctor A runs with the "Heart Rate" baton. They pass it to the "Blood Pressure" station on the map. Even though Doctor B never held the "Blood Pressure" baton, they receive the message from the "Heart Rate" station. Now, Doctor B can make a smart guess about Blood Pressure using the clues from Heart Rate.
This allows doctors to learn from each other's tools, even if they never use the exact same tools on the same patient.

3. The "Self-Teaching" (The Practice Exam)

Since the doctors can't share real patient data, how do they know if they are getting better?

Analogy: The system plays a game of "Hide and Seek" on the data they do have. The computer hides a few known numbers (like hiding a blood pressure reading that is already there) and asks the model to guess it. If the model guesses correctly, it gets a point. This happens over and over until the model becomes a master at filling in the blanks.

What Happened When They Tested It?

The researchers tested this on three different "puzzles" (datasets):

Air Quality: Measuring pollution (moderate difficulty).
PhysioNET: Patient vital signs (medium difficulty).
SECOM: A massive factory with hundreds of sensors (very high difficulty).

The Results:

The Old Way: When the "doctors" had different tools, the old methods failed miserably, especially on the big factory dataset. It was like trying to build a house with mismatched bricks.
FedHF-Impute: This new method was a champion.
- On the Air Quality and PhysioNET puzzles, it did as well as or better than the best existing methods.
- On the SECOM (factory) puzzle, it was a massive winner. It improved accuracy by nearly 27% compared to the next best method. It was so good that it actually beat the "Centralized" method (where all data is pooled in one place), which is usually the gold standard.

The Bottom Line

The paper claims that FedHF-Impute is a breakthrough because it stops treating "missing tools" as a dead end. Instead, it uses a smart map to connect the dots between different tools, allowing decentralized teams to build a powerful, accurate model without ever sharing their private data.

It's like turning a group of people with different flashlights into a single, powerful searchlight that can see in the dark, even though no single person has the whole beam.

Technical Summary: FedHF-Impute for Federated Imputation under Heterogeneous Feature Spaces

1. Problem Statement

Federated Learning (FL) typically assumes that all participating clients share an identical feature space (aligned schemas), enabling direct parameter aggregation via methods like FedAvg. However, in realistic tabular settings, clients often operate under heterogeneous feature spaces due to differing measurement pipelines, sensor configurations, or institutional constraints. In these scenarios, clients observe only partially overlapping subsets of features.

The paper distinguishes between two types of data absence:

Temporary Missingness: A feature is part of a client's schema but has missing entries (standard missing data).
Structural Unavailability: A feature is never measured by a specific client because it is absent from their local schema.

Existing federated imputation methods primarily address temporary missingness and implicitly assume aligned schemas. When applied to heterogeneous settings, standard parameter-averaging methods (e.g., FedAvg) fail because they cannot transfer information across weakly overlapping or disjoint feature groups, leading to ill-posed aggregation and degraded model performance.

2. Methodology: FedHF-Impute

The authors propose FedHF-Impute, a framework that decouples global feature semantics from local client schemas. The core innovation is the use of a shared global feature graph to propagate information across statistically related features, even when those features are never jointly observed on any single client.

Key Components

Global Feature Graph Construction:
- The server collects feature metadata (identifiers) and pairwise correlation statistics (e.g., Pearson correlations) from clients. No raw data is shared.
- A weighted graph $G = (V, E, w)$ is constructed where nodes represent global features and edges represent statistical dependencies.
- Edge weights are derived from the top- $k$ correlated neighbors. This graph is shared with all clients once before training begins and remains fixed.
GNN-Based Imputer Architecture:
- The model operates over the feature dimension rather than the sample dimension.
- For each sample, the input includes the feature vector (with missing values zeroed), an observation mask, and an availability mask (indicating structural presence/absence).
- The architecture uses a Graph Neural Network (GNN) with message-passing layers. Each feature node aggregates signals from its neighbors in the global graph, allowing the model to infer values for unavailable features based on correlations with available ones.
- Initial node representations are formed by concatenating feature embeddings with scalar observations, masks, and availability indicators, processed through an MLP.
Self-Supervised Training Objective:
- Since ground-truth values for missing data are unknown, the model uses a block-masking strategy.
- During local training, a subset of observed and available features is randomly masked (corrupted).
- The model is trained to reconstruct these masked values using a Mean Squared Error (MSE) loss, effectively learning the feature dependencies required for imputation.
Federated Protocol:
- The framework adheres to standard FL protocols (FedAvg). Clients perform local stochastic gradient descent and upload updated model parameters.
- The server aggregates parameters based on local dataset sizes.
- Crucially, the method supports indirect cross-client knowledge transfer: a client can learn to impute a feature it never observes by leveraging the global graph's propagation of statistical strength from other clients who do observe that feature.

3. Key Contributions

Decoupling Structural Unavailability from Missingness: The framework explicitly models the difference between a feature being missing (data issue) and being structurally unavailable (schema issue), a distinction often ignored in prior FL work.
Global Feature Graph for Heterogeneous Schemas: It introduces a mechanism to share statistical dependencies via a global graph, enabling information flow between features that are never co-observed locally.
Compatibility: The approach maintains standard federated communication protocols (parameter exchange only) without requiring raw data sharing or complex alignment transformations during training.
Comprehensive Evaluation: The method is evaluated against both centralized and federated baselines across three distinct datasets with varying dimensionality and missingness patterns.

4. Experimental Results

The authors evaluated FedHF-Impute on three datasets: AirQuality (13 features), PhysioNET (42 features), and SECOM (590 features). In the experiments, a partial schema overlap was simulated where each of the 4 clients held a random 60% subset of features.

Performance Gains:
- SECOM (High-Dimensional): FedHF-Impute achieved an RMSE of 0.8205, significantly outperforming all federated baselines (ranging from 1.17 to 4.05) and centralized competitors (best centralized MICE: 1.1231). This represents a 26.9% improvement over the best federated baseline and a 27% improvement over the best centralized method.
- AirQuality: Achieved an RMSE of 0.7074, outperforming all federated baselines and approaching the best centralized method (MICE: 0.6756).
- PhysioNET: Achieved an RMSE of 0.9363, comparable to the best federated baseline (Fed-FeatGen: 0.9339) and the best centralized method (DAE: 0.9500).
Stability: The method demonstrated stable convergence across communication rounds, whereas federated baselines exhibited high variance and instability, particularly in high-dimensional settings.
Scalability: The advantage of the graph-based approach was observed to grow with feature dimensionality and schema heterogeneity, suggesting that feature-level modeling becomes increasingly critical as local per-feature estimation becomes ill-conditioned.

5. Significance and Claims

The paper claims that FedHF-Impute addresses a critical gap in federated learning: the inability to effectively impute data when clients possess structurally different feature sets.

Overcoming Schema Mismatch: The authors assert that simple parameter averaging (FedAvg) is insufficient for heterogeneous feature spaces. By explicitly modeling feature dependencies via a graph, the method enables "indirect cross-client knowledge transfer," allowing clients to infer values for features they never observe.
Privacy Preservation: The framework preserves data locality, as only model parameters and scalar validation metrics are exchanged. The graph construction relies only on metadata and correlation statistics, not raw records.
Superiority in Complex Settings: The results suggest that for high-dimensional, heterogeneous tabular data, structured feature-level modeling (via graph propagation) is superior to both naive federated averaging and centralized imputation methods that cannot leverage distributed correlations.

The authors conclude that explicitly modeling feature-level structure is essential for effective imputation in next-generation federated systems facing schema heterogeneity. They note that while the current work treats graph construction as a pre-training step under standard trust assumptions, future work could incorporate differential privacy for the graph construction phase.