[Paper Review] Multi-Task Learning with Incomplete Data for Healthcare

Xin Jiang, Saba Emrani+2 more|arXiv (Cornell University)|Jul 6, 2018

Domain Adaptation and Few-Shot Learning15 references3 citations

TL;DR

This paper proposes a robust multi-task learning framework with plug-in covariance matrix estimators to handle incomplete data in healthcare applications, avoiding bias from traditional imputation. By integrating LASSO and graph regularization within a unified optimization, the method improves model estimation and prediction accuracy—especially under high missing data rates—demonstrated on Alzheimer’s disease progression data.

ABSTRACT

Multi-task learning is a type of transfer learning that trains multiple tasks simultaneously and leverages the shared information between related tasks to improve the generalization performance. However, missing features in the input matrix is a much more difficult problem which needs to be carefully addressed. Removing records with missing values can significantly reduce the sample size, which is impractical for datasets with large percentage of missing values. Popular imputation methods often distort the covariance structure of the data, which causes inaccurate inference. In this paper we propose using plug-in covariance matrix estimators to tackle the challenge of missing features. Specifically, we analyze the plug-in estimators under the framework of robust multi-task learning with LASSO and graph regularization, which captures the relatedness between tasks via graph regularization. We use the Alzheimer's disease progression dataset as an example to show how the proposed framework is effective for prediction and model estimation when missing data is present.

Motivation & Objective

• Address the challenge of missing features in multi-task learning for healthcare data, where traditional imputation distorts data covariance and reduces sample size.
• Develop a framework that integrates missing data handling directly into the learning process, avoiding separate imputation steps that introduce bias.
• Improve model generalization and prediction accuracy in multi-task settings with incomplete data, particularly when missing data rates are high.
• Demonstrate the effectiveness of the proposed method on real-world Alzheimer’s disease progression data with high missingness.
• Provide a scalable and robust solution using plug-in estimators for covariance matrices within a LASSO and graph regularization framework.

Proposed method

• Use plug-in estimators for covariance matrices to handle missing features directly within the multi-task learning optimization, avoiding data imputation.
• Formulate the learning problem as a regularized optimization: minimize prediction error, L1 (LASSO) sparsity, and graph-based regularization (Frobenius norm of W·R).
• Incorporate graph regularization via matrix R to encode task relationships, where connected tasks share information through shared parameters.
• Apply projected gradient descent to solve the optimization problem efficiently, ensuring convergence and stability.
• Estimate the autocovariance matrix Γ using plug-in estimators that account for missing data, preserving data structure.
• Adapt the framework to various multi-task learning formulations by replacing the objective function while retaining the covariance estimation approach.

Experimental results

Research questions

• RQ1 Can plug-in covariance estimators improve model estimation and prediction in multi-task learning when data contains high levels of missing features?
• RQ2 How does the proposed robust multi-task learning framework compare to standard imputation methods (e.g., mean imputation, matrix factorization) in terms of prediction accuracy and model stability?
• RQ3 Does graph regularization enhance performance when combined with plug-in covariance estimation in the presence of missing data?
• RQ4 How does the framework perform as missing data rates increase, particularly in real-world medical datasets with complex correlation structures?
• RQ5 Can the proposed method maintain low bias and high accuracy without relying on separate data imputation or matrix completion steps?

Key findings

• The R-LGR1 method outperformed mean-imputation and matrix factorization (MF-LGR) across all missing data percentages in synthetic experiments, showing the lowest NMSE for both model and autocovariance estimation.
• At 40% missing data, R-LGR1 achieved an NMSE of 0.007535 for model estimation, significantly lower than mean-impute (0.007661) and MF-LGR (0.007684).
• On the Alzheimer’s disease dataset, R-LGR1 achieved the lowest prediction error (RMSE = 7.538e-04) on original data (9% missing), outperforming mean-impute (7.596e-04) and MF-LGR (7.670e-04).
• As missing data increased to 49%, R-LGR1 maintained the lowest error (7.660e-04), while mean-impute and MF-LGR showed increasing error, indicating superior robustness.
• R-LGR1 showed consistent improvement over baselines across all missing data levels, with the largest relative gain observed at higher missingness (e.g., 39% and 49%), confirming its stability.
• The method preserved data structure better than imputation-based approaches, as evidenced by lower NMSE on the autocovariance matrix estimator Γ, which is critical for accurate inference.