PaperlayerView exact revision on arXiv

Source-linked AI summary

Fed-ReMasker: Federated Tabular Imputation under Feature-Level Missingness

Ioannis Papathanail, Rooholla Poursoleymani, Lubnaa Abdur Rahman, Stavroula Georgia Mougiakakou

arXiv:2609.28105v1cs.LGcs.AI

TL;DR

Multi-center tabular collaboration faces restrictions on sharing raw records, alongside missing values and features absent at some clients. Fed-ReMasker adapts a masked autoencoder for federated imputation. In the homogeneous benchmark it is lowest-error in most value- and feature-level scenarios, and it remains robust under client heterogeneity with FedAvg.

  • Problem

    Multi-center data collaboration must contend with raw-record sharing restrictions and missing tabular values, including features clients never observe locally.

    Source: §I · p. 1

  • Method

    Fed-ReMasker adapts the ReMasker masked autoencoder to federated learning so clients can impute features never observed locally.

    Source: §I · p. 1

  • Results

    Fed-ReMasker is lowest-error in 93.2% of value-level and 96.7% of feature-level homogeneous scenarios, and remains robust under client heterogeneity with standard FedAvg.

    Source: §I · p. 2, A. Homogeneous benchmark · p. 6

  • Takeaways & Limitations

    Standard FedAvg performs competitively with more elaborate optimization strategies for federated tabular imputation.

    Source: §I · p. 2

    The study does not evaluate formal privacy mechanisms, communication or computation costs, or effects on downstream predictive or statistical analyses.

    Source: D. Limitations and Future Work · p. 9

Abstract

from arXiv · show

Multi-center clinical studies and biomedical research collaborations increasingly seek to utilize data across centers to build models that generalize beyond any single center. This creates two distinct challenges: data protection regulations may restrict the sharing of raw patient data across institutions, while centers may collect only partially overlapping sets of features under different protocols. Federated learning enables collaborative model training without centralizing raw data. However, existing federated imputation methods rarely evaluate feature-level missingness, in which entire features are unobserved at some centers. To address this setting, we adapt the ReMasker masked autoencoder to federated learning (Fed-ReMasker), enabling centers to impute features never observed locally by leveraging knowledge learned across collaborating centers. We evaluate Fed-ReMasker in a benchmark spanning synthetic datasets with linear and nonlinear relationships and real-world tabular datasets, including clinical data. The benchmark varies the number of centers, the missingness ratios, and client heterogeneity. Fed-ReMasker achieves the lowest imputation error in 93.2% of value-level and 96.7% of feature-level scenarios in the homogeneous benchmark. It also remains robust to client heterogeneity using simple federated averaging, outperforming all baselines in all 36 value-level scenarios and each baseline in at least 35 of 36 feature-level scenarios, and comes within 3.0% on average of a centralized model trained on the pooled data.

I. INTRODUCTION

The paper addresses feature-level missingness in federated tabular data, where some centers have no local observations of certain features. It introduces Fed-ReMasker and benchmarks it against federated imputation methods across missingness settings and client heterogeneity.

Motivation

  • Multi-center collaborations seek to use data across centers while data-protection rules, legal restrictions, or institutional policies may prevent sharing raw records.Federated learning supports shared-model training while records remain local.Source: §I · p. 1
  • Feature-level missingness occurs when a center has no local observations of a feature, leaving no local data to learn its distribution or relationships with other variables.This differs from value-level missingness, where some values of a feature remain locally available.Source: §I · p. 1

Scope

  • The study focuses on MCAR missingness at both value and feature levels, with an additional data-dependent feature-missingness condition based on clustering clients in feature space.The authors describe the latter as a robustness assessment beyond MCAR.Source: §I · p. 1

Contributions

  • Fed-ReMasker adapts the ReMasker masked autoencoder to let clients impute features never observed locally, with the largest gains on entirely missing features.The work presents it as the lowest-error method among existing federated imputation methods.Source: §I · pp. 1–2
  • The benchmark systematically evaluates both missingness levels, including features entirely unobserved at individual clients, across client counts, synthetic and real datasets, and four heterogeneity conditions.It is presented as the first systematic benchmark of this setting, to the authors’ knowledge.Source: §I · p. 2
  • Fed-ReMasker outperforms existing federated imputation methods across datasets under both value- and feature-level missingness, while standard FedAvg performs competitively under client heterogeneity.The authors also compare FedAvg, FedProx, and FedAdam and quantify federation's cost against a centralized reference.Source: §I · p. 2

II. RELATED WORK

Prior work includes centralized imputation methods and federated approaches to incomplete data, but their aims and evaluations differ. The paper positions its benchmark around imputing features entirely unobserved at individual clients.

A. Centralized Imputation

  • Centralized methods range from classical approaches to deep generative and masked-autoencoding models, including ReMasker, which reconstructs masked observed values during training and imputes genuinely missing values at inference.These methods assume centrally accessible data.Source: A. Centralized Imputation · p. 2

B. Federated Imputation

  • Other federated methods address heterogeneous feature spaces through representation learning or downstream prediction rather than data-level imputation.Examples include FLIC, DARN, and FedFeatGen.Source: B. Federated Imputation · p. 2
  • Federated methods address tabular imputation through generative, personalized-aggregation, or adversarial approaches, but mainly consider value-level or heterogeneous missingness rather than features wholly unobserved at clients.Fed-MIWAE, Cafe, and FedImpute are examples of these approaches.Source: B. Federated Imputation · p. 2
  • FedHF-Impute uses a global feature graph, but its evaluation excludes positions in features unobserved at a client and uses a fixed four-client setting with 60% feature availability per client.The paper contrasts this with evaluation across client counts, missingness levels, and synthetic and real datasets.Source: B. Federated Imputation · p. 2
  • Cross-site imputation addresses variables unrecorded at some sites by transporting regression coefficients, but targets downstream valid inference and handles one unrecorded variable at a time.The authors say it is not directly comparable using NRMSE.Source: B. Federated Imputation · p. 2

C. Federated Optimization

The paper formulates federated imputation with a shared feature schema and clients that may lack entire features. Fed-ReMasker trains locally and aggregates models so that information from observing clients can support imputation elsewhere.

C. Federated Optimization

  • FedAvg, FedProx, and FedAdam are compared to assess whether specialized optimization improves federated imputation under client heterogeneity.FedAvg weights client updates by sample size; FedProx limits client drift, while FedAdam applies a stateful server optimizer.Source: C. Federated Optimization · p. 2

A. Problem Formulation

  • The horizontal federated setting uses a shared global feature schema, while clients hold local datasets and may have entire features unobserved.Only model parameters are communicated to the central server for aggregation.Source: A. Problem Formulation · p. 2
  • Feature-level missingness means a feature is entirely unobserved at a client, unlike value-level missingness, where only individual entries are absent.The study estimates missing entries against their unobserved ground truth and treats clustered feature-distribution skew as data-dependent missingness.Source: A. Problem Formulation · p. 3

B. Fed-ReMasker

  • Fed-ReMasker keeps ReMasker’s architecture and masked-reconstruction objective, trains a local model at each client, and aggregates parameters with sample-size-weighted FedAvg.The resulting global model is used by each client to impute both missing values and entirely missing features.Source: B. Fed-ReMasker · p. 3

C. Federated Optimization Strategies

The paper compares sample-size-weighted FedAvg with FedProx and FedAdam as alternative federated optimization strategies. The alternatives modify either local training or server-side update handling.

FedAvg

  • FedAvg sets the next global model to a sample-size-weighted average of client models, equivalently applying the aggregated client update as a unit server step.Client weights are p_k = N_k/N.Source: C. Federated Optimization Strategies · p. 3

FedAdam

  • FedAdam leaves local training unchanged and applies a stateful Adam-based server optimizer to the aggregated client update.Its server learning rate and the FedProx proximal coefficient are selected as described in the experimental setup.Source: C. Federated Optimization Strategies · p. 3

FedProx

  • FedProx adds a proximal penalty to each client’s local objective to limit client drift; when its coefficient is zero, it reduces to FedAvg.The penalty strength is controlled by μ.Source: C. Federated Optimization Strategies · p. 3

IV. EXPERIMENTAL SETUP

The benchmark combines three filtered real-world tabular datasets with synthetic datasets spanning linear and nonlinear relationships. Preprocessing statistics use observed values only, while masked ground truth is reserved for evaluation.

  • The benchmark includes 9,249 Codon samples with 64 features, 9,313 PhysioNet patient records with 19 numeric features, and 36,758 NHANES participants with 32 numeric features.Codon retains genomic DNA entries and complete rows; PhysioNet uses each patient's last measurement per parameter; NHANES merges 2003–2018 survey cycles.Source: A. Datasets · p. 3
  • Six synthetic datasets span linear and neural-network-induced nonlinear relationships at dimensions 20, 50, and 100, with 50,000 samples each and additive Gaussian noise coefficient 0.4.The datasets are generated using the Causal Discovery Toolbox, following.Source: A. Datasets · pp. 3–4
  • Preprocessing aggregates feature-wise means, minima, and ranges from observed client values; masked ground truth is used only for evaluation, and the benchmark provides no formal privacy guarantees.The experiments omit secure aggregation and differential privacy, evaluating decentralized imputation without raw-data sharing instead.Source: A. Datasets · p. 4

B. Missingness Simulation and Federated Setup

The setup varies value- and feature-level missingness across client counts, then adds controlled heterogeneity conditions on real datasets. Every feature remains observed by at least one client.

B. Missingness Simulation and Federated Setup

1) Homogeneous client scenarios:

  • 324 homogeneous scenarios span six synthetic datasets, three real datasets, four feature-missingness ratios, three value-missingness ratios, and three client counts.Clients are randomly partitioned with shared MFR and MVR; each feature is retained by at least one client, and one fixed partition seed is used.Source: 1) Homogeneous client scenarios · p. 4

2) Heterogeneous client scenarios:

  • The heterogeneity study fixes MFR = 0.2, MVR = 0.3, and K = 5, and compares quantity skew, missingness skew, and feature-distribution skew across three partition seeds.Quantity skew uses Gini coefficients 0.2 and 0.35, yielding four conditions in total; the study includes 36 heterogeneous scenarios and nine homogeneous reference scenarios.Source: 2) Heterogeneous client scenarios · p. 4
  • FedProx and FedAdam hyperparameters are selected by leave-one-dataset-out validation, minimizing mean overall NRMSE on the other two real datasets and applying the selected value unchanged.The search grid for both parameters is {0.001, 0.01, 0.1, 1.0}.Source: 2) Heterogeneous client scenarios · p. 4

3) Centralized reference:

The study compares federated methods with a pooled-data ReMasker reference, alongside four imputation baselines. Evaluation separates missingness types and summarizes paired performance across scenarios and datasets.

3) Centralized reference:

  • The centralized ReMasker uses pooled client data and the same architecture, optimizer settings, and 300 data passes; pooling turns feature-level missingness into value-level missingness.The reference covers 126 scenarios across nine datasets, with additional three-seed runs on real datasets for feature-distribution skew and the homogeneous reference.Source: 3) Centralized reference · pp. 4–5

C. Baselines

  • Fed-ReMasker is compared with global-mean imputation, Cafe, Fed-MIWAE [13], and FedHF-Impute [16].For entirely missing local features, Cafe averages models from observing clients, Fed-MIWAE can use the global model, and FedHF-Impute uses its global feature graph.Source: C. Baselines · p. 5

D. Evaluation Metrics

  • NRMSE is reported for value-level, feature-level, and all missing positions; scenario-level scores are averaged unweighted by dataset, with heterogeneity experiments also reporting equal-client macro NRMSE.Feature-level NRMSE includes only scenarios with MFR > 0; masked values are retained solely as evaluation ground truth and are not used in training or preprocessing.Source: D. Evaluation Metrics · p. 5
  • Paired differences are calculated as NRMSEbaseline − NRMSEFed-ReMasker, with positive values favoring Fed-ReMasker; t-intervals use dataset-level mean differences.Repeated partitioning is assessed separately in the heterogeneity study; feature-level comparisons cover 243 scenarios with MFR > 0.Source: D. Evaluation Metrics · p. 5

E. Implementation Details

Fed-ReMasker uses a fixed federated training budget and a specified masked-autoencoder architecture and optimizer configuration.

  • All methods train for 60 federated rounds with five local epochs per round; Fed-ReMasker uses an 8-layer encoder, 4-layer decoder, 64-dimensional embedding, and four attention heads.Clients use shuffled batches of 64 with AdamW, learning rate 10^-3, weight decay 0.05, cosine decay, four warm-up rounds, and a 10^-6 floor.Source: E. Implementation Details · p. 5

V. RESULTS AND DISCUSSION

Fed-ReMasker performs best across the homogeneous benchmark’s value- and feature-level missingness scenarios. Its masked-autoencoding design supports imputing features absent locally by transferring information through cross-client aggregation.

A. Homogeneous benchmark

  • Fed-ReMasker remains lowest on weakly correlated PhysioNet data and relatively stable as feature missingness and client count increase.The mean absolute pairwise feature correlation for PhysioNet is 0.085; Cafe degrades as missing features increase.Source: A. Homogeneous benchmark · pp. 5–6, Figure 1 · p. 6
  • Fed-ReMasker achieves the lowest NRMSE on every dataset, averaging 0.061 for value-level and 0.065 for feature-level missingness.Table I reports unweighted averages across datasets and scenario-level settings.Source: A. Homogeneous benchmark · p. 6, Table I · p. 7
  • Fed-MIWAE’s feature-level NRMSE is 0.203, worse than Mean imputation at 0.165 on all nine datasets.For value-level missingness, Fed-MIWAE’s NRMSE is 0.118 versus 0.144 for FedHF-Impute and 0.161 for Mean.Source: A. Homogeneous benchmark · p. 6
  • Fed-ReMasker is the lowest-error method in 93.2% of value-level and 96.7% of feature-level scenarios.Across the remaining scenarios, its deficit is at most 0.0058 NRMSE for value-level and 0.0044 for feature-level positions.Source: A. Homogeneous benchmark · p. 6
  • Self-attention captures inter-feature relationships, while masking and re-masking objectives support imputing entirely unobserved features using information transferred from clients that observe them.The passage identifies these as properties of the centralized ReMasker carried into the federated setting.Source: A. Homogeneous benchmark · pp. 6–7

B. Robustness to client heterogeneity

Fed-ReMasker retains its advantage under the evaluated client heterogeneity conditions. Feature-distribution skew particularly affects feature-level imputation, while FedAvg remains an effective optimization strategy in this benchmark.

B. Robustness to client heterogeneity

  • Fed-ReMasker beats every baseline in all 36 value-level scenarios and in at least 35 of 36 feature-level scenarios, including Cafe in all 36.Fig. 2 compares NRMSE across the homogeneous reference and four heterogeneity conditions on three real datasets.Source: B. Robustness to client heterogeneity · p. 7, Figure 2 · p. 8
  • Under quantity skew with Gini = 0.35, Fed-ReMasker changes by only 0.0003 between weighting schemes.The largest shift across all methods is 0.0062; client sizes differ by 6.8× at this Gini level.Source: B. Robustness to client heterogeneity · p. 7
  • FedAvg achieves the best or tied-best NRMSE in 29 of 30 cells across optimization strategies.Table II reports mean ± standard deviation across three seeds for each dataset and condition.Source: B. Robustness to client heterogeneity · p. 7, Table II · p. 8
  • Feature-distribution skew increases feature-level error by 0.025 on Codon, 0.014 on PhysioNet, and 0.005 on NHANES.Value-level error under feature-distribution skew increases by at most 0.004; under missingness skew, feature-level error falls by 0.020 on Codon and 0.014 on PhysioNet.Source: B. Robustness to client heterogeneity · p. 7
  • Feature-level positions rely on relationships learned from clients that observe the feature, unlike value-level positions reconstructable from local correlations.This distinction is consistent with the greater feature-level sensitivity to feature-distribution skew.Source: B. Robustness to client heterogeneity · p. 7
  • FedAdam is closer to FedAvg than FedProx on all three datasets, while feature-level error reaches ±0.028 seed standard deviation on skewed PhysioNet.On PhysioNet, the largest difference between strategies in any cell is 0.004, within that feature-level seed standard deviation.Source: B. Robustness to client heterogeneity · p. 7
  • Sample-size-weighted FedAvg is retained as the default because the shared masked-reconstruction objective may already align client updates.The authors report that proximal regularization and server-side adaptivity provide little benefit under the evaluated conditions.Source: B. Robustness to client heterogeneity · p. 8

C. Centralized vs. Federated ReMasker

Fed-ReMasker’s error averages 3.0% above the pooled centralized model, with a larger gap at higher client counts. The evaluation also identifies limits in missingness mechanisms, data scope, and deployment outcomes.

C. Centralized vs. Federated ReMasker

  • The relative NRMSE gap between Fed-ReMasker and centralized ReMasker grows with client count, from 1.2% at K = 3 to 3.2% at K = 5 and 7.2% at K = 10.Source: C. Centralized vs. Federated ReMasker · p. 8
  • Under feature-distribution skew, both centralized and federated errors rise versus the homogeneous reference, while no overall advantage from centralizing is detected.In Fig. 3, homogeneous client partitions overlap almost completely, whereas feature-distribution-skew partitions are visibly separated.Source: C. Centralized vs. Federated ReMasker · p. 9, Figure 3 · p. 9

D. Limitations and Future Work

  • The benchmark mainly tests MCAR missingness, excludes MAR and MNAR value-level mechanisms and systematic site-specific feature availability, and uses datasets with at least 9,200 samples.Quantity, missingness, and feature-distribution skew are evaluated separately, leaving their combined effect unknown.Source: D. Limitations and Future Work · p. 9

C. Centralized vs. Federated ReMasker

D. Limitations and Future Work

  • The study does not evaluate formal privacy mechanisms, communication or computation costs, downstream effects of imputation, or categorical features.The homogeneous grid also uses a single client-partition seed.Source: D. Limitations and Future Work · p. 9

VI. CONCLUSION

Fed-ReMasker adapts ReMasker for federated tabular imputation when features are entirely absent at some centers. It achieves the lowest NRMSE on every dataset for both value- and feature-level missing positions, averaging 3.0% above pooled centralized ReMasker.

VI. CONCLUSION

  • 3.0% above pooled centralized ReMasker on average, Fed-ReMasker achieved the lowest NRMSE on every dataset for both value- and feature-level missing positions.The largest gains were on entirely missing features.Source: §VI · pp. 9–10
  • Fed-ReMasker adapts the ReMasker masked autoencoder to federated tabular imputation, addressing a setting where entire features are absent at some centers.The benchmark spans value- and feature-level missingness, synthetic and real datasets, client counts, and client heterogeneity.Source: §VI · p. 9
  • By imputing locally absent features that are partially observed elsewhere, Fed-ReMasker supports collaborative analysis across centers without centralizing raw data.Source: §VI · p. 10

Related papers

Loading 2609.28105v1…