MT-LLE: Multi-Task Locally Linear Embedding for Interpretable Disease Modeling from Longitudinal Omics Data

Abstract

Constructing interpretable disease models from longitudinal omics data is a central challenge in precision medicine. The goal is a low-dimensional representation in which a patient's position encodes their molecular state and clinical severity, and along which disease progression can be read directly. Existing dimensionality reduction methods (e.g., UMAP, Variational Autoencoders) fall short of this goal: they optimize a single generic objective and are blind to clinical labels and to the temporal ordering of measurements. Consequently, trajectory inference is typically applied after the fact to an embedding that was never optimized to reveal progression, decoupling the representation from disease dynamics. Manifold learning offers a natural route to such representations, and we build on Locally Linear Embedding (LLE) to preserve the local geometry of the omics data (i.e., keeping molecularly similar patients close together in the low-dimensional space). Geometry alone, however, yields a space that is faithful to molecular similarity yet uninformative about clinical severity and progression. We therefore recast the problem as multi-task learning: MT-LLE jointly optimizes five objectives: geometric reconstruction, supervised organization by clinical stage, embedding and phenotype forecasting, and clustering. Because naively combining such heterogeneous objectives induces gradient conflicts that distort the molecular geometry, an embedded reinforcement learning agent dynamically schedules their weights during training, establishing global geometry before refining clinical boundaries. Across two independent Chronic Obstructive Pulmonary Disease (COPD) cohorts (SPIROMICS and COPDGene), MT-LLE deliberately relaxes exact geometric reconstruction, by a modest margin, in exchange for substantial gains in clinical structure. On held-out patients, a linear model reads disease severity (GOLD stage, 0--4) from the MT-LLE embedding 35--40\% more accurately than from standard dimensionality reduction (0.53 vs.\ 0.38 F1-Macro). The gap is starker for progression: forecasting a patient's next-visit severity from their trajectory reaches 0.38 F1-Macro, while unsupervised baselines sit near zero (0.06--0.09 F1-Macro), a temporal signal those methods fail to capture. To test whether the reinforcement learning agent earns its place, we compared it against a fixed schedule that imposes the same ordering of objectives but cannot adapt during training; the learned agent outperforms it by 13--20\% across clinical metrics, showing the gains come from adapting the weights to how training unfolds, not merely from ordering the objectives correctly, and at no cost to geometric fidelity. Beyond these quantitative gains, the manifold supports complementary analyses that surface structure invisible to standard staging: static phenotyping isolates subjects with active molecular pathology despite preserved lung function; trajectory inference maps two mechanistically distinct progression axes (inflammatory fibrosis and pan-immune activation); and kinematic analysis of each patient's speed and acceleration identifies subjects whose molecular trajectories accelerate ahead of detectable spirometric decline.