Alzheimer’s disease (AD) is a progressive neurodegenerative disorder, with its prevalence increasing annually (AuthorAnonymous, 2023). Approximately 50 million individuals worldwide are affected by dementia with approximately 60%–70% being AD cases, and this figure is projected to rise to 152 million by 2050 (Santiago et al., 2019; Zhang et al., 2021). AD is a multifactorial condition manifested through various traits, such as cognitive decline and functional changes (Löffler et al., 2014; Jabir et al., 2021). Genetic risk prediction models have been developed for various AD-related traits, but these models usually only focus on one trait, ignoring their inter-relationships (Jung et al., 2020; Zhu et al., 2024). Although each AD-related trait provides valuable information on the genetic risk of AD, none of them alone can capture the full complexity of the disease and a comprehensive model that can jointly model multiple traits is needed.

Cognitive and functional changes commonly observed in AD patients can be assessed using several tools, including the Mini-Mental State Examination (MMSE), Montreal Cognitive Assessment (MoCA), Clinical Dementia Rating-Sum of Boxes (CDRSB), Alzheimer’s Disease Assessment Scale-Cognitive Subscale 13 (ADAS13), and the Functional Activities Questionnaire (FAQ). MMSE and MoCA assess general cognitive impairment, with MoCA being more sensitive in detecting early AD (Pinto et al., 2019; Duc et al., 2020). ADAS13 and CDRSB are designed for tracking AD progression. However, ADAS13 measures the severity of cognitive symptoms (Bucholc et al., 2019), whereas CDRSB assesses both cognitive and functional domains, offering a more comprehensive view of how AD affects a patient’s daily life (Cullen et al., 2020). FAQ focuses on assessing functional ability in daily activities (Petersen et al., 2021). Neuroimaging is also used in AD diagnosis and monitoring (Besson et al., 2015; Winer et al., 2018). For example, florbetapir (AV45) detects the amyloid-beta plaque in the brain (Mattson, 2004; Johnson et al., 2013). Fluorodeoxyglucose (FDG) measures brain glucose metabolism and identifies regions of hypometabolism (de Paula Faria et al., 2022). While AD assessment tools provide valuable information, none of them alone can be treated as a gold standard for AD diagnosis, especially for early-stage cases. For example, although amyloid plaque is a hallmark feature of AD, some individuals with such manifestations never develop into AD (Reinitz et al., 2022).

PET-imaging, cognitive, and functional changes provide confirmatory and complementary information regarding AD risk. Simultaneous modelling of them can leverage information across traits, which facilitates the detection of new biomarkers and improves the overall prediction. However, existing prediction models mainly focus on a single trait. For example, traditional models such as gBLUP build separate prediction models for each trait (de Los Campos et al., 2013). This trait-specific focus persists even within the deep learning domain. For example, Duc et al. (2020) employed a single task learning (STL) model to automatically diagnose AD (Duc et al., 2020). Liu et al. (2022) introduced an interpretable STL model to assess the risk of AD based on high-dimensional genomic data, where PET imaging outcomes were predicted separately (Liu et al., 2022). Notably, some studies have explored modeling multiple tasks simultaneously. For example, a classic multi-task model, Multi-Lasso, has shown promising performance when applied on SNPs data (Wang et al., 2012; Bee et al., 2024). This method, originally proposed by Obozinski et al. (2006), applies joint sparse regularization through ℓ ₂,₁-norm across tasks, enabling feature sharing among related tasks. It has been applied in several studies (Wang et al., 2012; Bee et al., 2024), including the recent application on Alzheimer’s Disease Neuroimaging Initiative (ADNI) dataset (Cheng et al., 2019). These studies indicate the potential of multi-task modeling strategies in the context of AD prediction, although their practical application remains limited.

In recent years, Multi-task learning (MTL) has been widely applied in the field of deep learning as an effective strategy to improve model performance. It has been successfully applied to model multiple correlated outcomes, particularly in typical deep learning scenarios such as natural language processing (Zhang et al., 2020; Li et al., 2022) and image classification (Liu et al., 2021). Current deep MTL approaches can be broadly categorized into hard and soft parameter sharing models. Hard parameter sharing models use shared layers across tasks, typically sharing all layers except the last to learn a common representation while capturing task-specific characteristics (Vandenhende et al., 2022). Recent advances enable these models to automatically determine which layers to share. For example, the Fully Adaptive Feature Sharing method dynamically widens layers based on similarities among tasks (Lu et al., 2017). The Multilinear Relationship Network discovers inter-task relationships and alleviates the dilemma of negative transfer by jointly training transferable features (Long et al., 2017). The Task Affinity Grouping optimizes layer sharing by branching the network according to inter-task affinity scores (Fifty et al., 2021). Unlike hard sharing, soft parameter sharing models allow each task to maintain independent parameters and control the levels of sharedness using additional parameters. They offer greater flexibility, but at the cost of increased computational demands. The cross-stitch network represents a classic example, where additional parameters are introduced in the cross-stitch unit to ascertain the optimal degree of sharedness (Misra et al., 2016). AD-related traits, such as cognitive scores, functional assessments, and neuroimaging findings are interconnected. MTL approaches, especially those computationally efficient hard parameter sharing models, have great potential to enhance generalization, learning efficiency and overall prediction accuracy for genetic risk predictions. However, existing MTLs have been rarely used in such applications, partially due to the low signal-to-noise ratio and unclear levels of genetic relatedness.

We here developed an auto-branch multi-task learning model for the prediction analyses of multiple correlated traits using genetic data. Our method can distinguish and integrate commonalities and unique characteristics across multiple traits, leading to improved prediction performance, as measured by both Pearson correlation and root mean squared error (RMSE). In the following sections, we first provided the technical details of our method and then conducted extensive simulation studies to evaluate its performance. Finally, we built genetic risk prediction models for multiple AD-related traits, including cognitive and functional assessments and PET imaging outcomes, using data sourced from the ADNI (Mueller et al., 2005).

Our method is developed using the idea originally proposed by Fifty et al. (2021) in the analyses of facial image dataset CelebA (Liu et al., 2015) and computer vision dataset Taskonomy (Zamir et al., 2018). In hard parameter sharing models, the gradient update of one task can influence others. If the gradient update of one task reduces the loss of another, a “synergistic effect” between the two tasks is observed. Conversely, an “antagonistic effect” occurs when the update negatively impacts the other task. Jointly training can enhance model performance for synergistic tasks by leveraging their positive correlations, but it may reduce performance for antagonistic tasks. We proposed to quantify the “synergistic effect” among correlated traits and branch the network where traits are considered antagonistic. Specifically, we first constructed a hard parameter sharing model with all layers except the last shared to predict multiple traits using genetic data. We then quantified trait similarities and grouped traits using the inter-trait affinity (Fifty et al., 2021). Finally, we branched the hard parameter sharing model for traits that are deemed antagonistic. Unlike Fifty et al. (2021), who used separate models for traits in each group, we proposed to use a hard-parameter sharing strategy and branch the network when phenotypes are “antagonistic.” This is mainly because correlated phenotypes are likely to have shared genetic determinants (Badré and Pan, 2023), and we hypothesized that this can be exploited to improve predictions. The overview of our workflow is in Figure 1.

We conducted extensive simulations to assess the performance of our method. As outlined in the method, we pre-set the number of branches to be 2, 3, and 4, and then found the trait partition accordingly. We compared our model with a typical hard parameter sharing model where all parameters except the last are shared (denoted as HPS), a single-task model with each trait modeled independently (denoted as STL), and Multi-lasso that is a classic method for modeling multiple traits (Cheng et al., 2019). As we aim to improve prediction for multiple AD-related traits (i.e., FDG, AV45, FAQ, CDRSB, ADAS13, MMSE, and MoCA), we make our simulation settings like the ADNI dataset. Specifically, we directly extracted genomic data from ADNI and simulated seven phenotypes based on causal variants, which are harbored on six randomly selected genes and comprised of 10% of the total variants. We randomly split the samples into training, validation, and testing sets with a ratio of 8:1:1. For all models, we set 2 hidden layers with 128 and 32 hidden nodes, respectively (Supplementary Figure S1).

Since all simulated phenotypes are continuous, we used mean squared error (MSE) as the loss function during the model training process for all deep learning models. We evaluated the prediction performance based on the testing set and reported the average prediction Pearson correlations and root of MSE (RMSE) based on 100 Monte Carlo simulations under each setting. Pearson correlation captures the linear consistency between predicted and true values and is widely used in continuous phenotype prediction. RMSE complements this by quantifying the average prediction error in the same unit as the original data, providing a comprehensive assessment of model performance.

We are interested in predicting multiple AD-related traits, including cognitive scores (i.e., MMSE, MoCA, ADAS13), functional assessments (i.e., FAQ, CDRSB), and neuroimaging findings (i.e., AV45 and FDG), using genetic data obtained from ADNI. ADNI is a comprehensive longitudinal study aimed at identifying biomarkers associated with AD and enhancing its clinical diagnosis and early intervention. As all seven phenotypes are quantitatively measured continuous traits, they were modeled as regression tasks in our study.

Data were downloaded from the ADNI website (https://www.adni.loni.usc.edu/). We excluded individuals without genomic data or have missing phenotypes. Only autosome SNPs were considered in our analysis. We adopted a candidate gene approach, where 57 AD susceptibility genes identified based on existing literature were included (Supplementary Table S18). For quality control, we excluded SNPs if they met any of the subsequent criteria: 1) missing rate >1%; 2) minor allele frequency (MAF) < 5%; 3) Hardy-Weinberg equilibrium test with p -value <1 × 10⁻⁶; 4) linkage disequilibrium (LD) analysis with an LD threshold of R² > 0.9. Additionally, individuals with missing rate of SNPs larger than 1% were excluded. All selected SNPs were directly genotyped in the ADNI dataset. For individuals who passed the quality control, we utilized the plink2R package to impute the missing for the included SNPs. A total of 463 participants and 3,797 SNPs harbored on the 57 AD susceptibility genes were included in the final analyses. The distributions of seven phenotypes are shown in Supplementary Figure S2.

We randomly split the data into an 8:1:1 ratio for training, validation and testing to mitigate overfitting, and repeated the random sampling 100 times for robustness. Given that all seven traits are AD-related and some even focus on similar aspect of AD (e.g., cognitive changes), it is unlikely they are fully independent. Likewise, while these traits provide complementary insights into AD, they are unlikely to share identical genetic causes. Therefore, for our method, we set the pre-specified number of branches to be 2, 3, and 4, excluding cases where all traits are independent or share identical causes. For comparison, we included the Multi-Lasso model, the hard sharing model HPS, where all layers are shared except for the last, and the STL model, where each trait is trained separately. We kept the network architecture consistent with that used in the simulations. We reported Pearson correlation and RMSE for all methods and further used Wilcoxon signed-rank test to compare our method to the others.

The prediction performance, measured by Pearson correlation and RMSE, is illustrated in Figures 2, 3, respectively. Our auto-branch method performs the best when the pre-specified number of branches is set to three, though the performance differences across various branch numbers are minimal. This method demonstrates superior prediction across multiple traits, not only in terms of higher average Pearson correlations but also by maintaining lower average RMSEs. Among all compared methods, Multi-Lasso demonstrates weaker performance in several key traits (Table 5). Specifically, when the number of branches is set to three, the average Pearson correlations for FDG, AV45, FAQ, CDRSB, and ADAS13 show relative improvements of 29.45%, 33.96%, 98.60%, 17.48%, and 7.51% over Multi-Lasso, with corresponding absolute increases of 0.045, 0.059, 0.061, 0.022, and 0.011, respectively. In addition, Multi-Lasso yields consistently higher RMSEs, particularly for FAQ and ADAS13, reflecting its limited capacity to model complex inter-trait relationships. The Wilcoxon signed-rank test indicated that the increase is statistically significant for FDG, AV45, and FAQ (Supplementary Table S19).

Regardless of branch number, our method consistently performs similarly to or better than HPS (i.e., similar or higher Pearson correlations, and similar or lower RMSE). Specifically, the improvements in average Pearson correlations for FDG, AV45, FAQ, CDRSB, ADAS13, and MMSE with a three-branch setup are 18.56%, 31.64%, 32.61%, 32.43%, 36.21%, and 24.79%, with corresponding absolute increases of 0.031, 0.056, 0.030, 0.036, 0.042, and 0.030, respectively. The Wilcoxon signed-rank test indicated that the increase is statistically significant for AV45, CDRSB, and ADAS13 (Supplementary Table S19).

Compared to the STL models, our method consistently performs similarly to or better than STL (i.e., similar or higher Pearson correlations and consistently lower RMSE). The improvements in average Pearson correlations for FDG, AV45, FAQ, CDRSB, ADAS13, MMSE, and MoCA are 76.18%, 29.07%, 35.78%, 20.36%, 21.62%, 19.03%, and 3.22%, respectively, with corresponding absolute increases of 0.086, 0.052, 0.032, 0.025, 0.028, 0.024, and 0.004. The Wilcoxon signed-rank test indicated a statistically significant increase for FDG, AV45, FAQ, and MMSE (Supplementary Table S19).

Our analyses suggest that the seven AD-related traits neither share identical genetic causes nor are completely independent. PET-imaging traits FDG and AV45 benefit substantially from joint modeling, indicating cognitive and function tests provide auxiliary information. Information transfer between PET-imaging outcomes and cognitive as well as function tests help learn a better representation, leading to the improvement of prediction. In summary, our method improves prediction performance for most of the traits regardless of the pre-specified number of branches, highlighting its robustness and potential for broad practical applications, especially for phenotypes where identifying latent patterns is essential.

To further investigate our model, we calculated the predictive feature importance score for each gene using a permutation-based approach proposed by Liu et al. (2022). The basic rationale is that if a gene is predictive, then the model accuracies with and without it would differ significantly. Following the procedure proposed by Liu et al. (2022), we assessed the importance of each gene by quantifying the difference in accuracies while accounting for variability. Specifically, we first calculated the Pearson correlation between predicted and observed outcomes using the original data with the already trained auto-branch model. We then recalculated the Pearson correlations after randomly shuffling the SNPs located within the gene of interest, while preserving the genetic structure (e.g., linkage disequilibrium), to generate a null distribution. To empirically estimate the variance, we repeated the permutation process 100 times. The standardized difference in Pearson correlations between the original and shuffled data (also called the predictive feature importance score by Liu et al.) reflects the gene’s contribution to predictive performance, with a larger difference indicating greater predictive importance. We then calculated the predictive feature importance score, which follows an asymptotic normal distribution under the null as shown by Liu et al. (2022), to evaluate the predictive importance of each gene. Genes were ranked based on their predictive importance scores, and we focused on those with scores greater than 1.645, corresponding to a 5% significance threshold under the asymptotic normal distribution. It is important to note that we did not intend to perform hypothesis testing, rather we chose a cut-off value (i.e., 1.645) and focused on genes with predictive feature importance score larger than this cut-off. Therefore, no multiple testing correction was applied. Table 6 presents the genes identified as having a probability greater than 75% of being significantly predictive at the 5% level for at least one trait, and the details for each gene are provided in Supplementary Tables S20–S22, which correspond to models with 2, 3, and 4 pre-specified numbers of branches, respectively. APOC1, APOE, and TOMM40 demonstrate stable and significant predictive power. Even though the most significant genes are similar among phenotypes, their predictive power varies considerably (Figure 4), which explains the superior performance of our branching method over HPS. In our model with three branches, APOE plays a crucial role in predicting FDG and AV45, with its removal leading to average Pearson correlations decreases of 0.10 and 0.14, respectively. For FAQ, CDRSB, ADAS13, MMSE and MoCA, the impact is small-to-moderate, with Pearson correlations reduced by 0.06, 0.08, 0.08, 0.08, and 0.04 respectively.

In this study, we proposed an efficient and robust auto-branch multi-task learning method for simultaneously predicting multiple correlated traits. Using total inter-task affinity, which quantifies the impact of gradient updates from one trait on the others, our method automatically determines the best partition of traits to enable efficient information transfer among similar traits, thereby enhancing prediction performance.

Through simulations, we found that our method has similar or better performance than that of the hard parameter sharing model where shared layers are pre-specified (Zhao et al., 2020). Our auto-branch model identifies the optimal phenotype partitioning that maximizes overall inter-trait affinity. Phenotypes grouped together share layers, while those in separate groups are assigned distinct branches. This data-driven approach to sharing layers enables efficient information transfer among inherently similar phenotypes and significantly reduces the risk of negative transfer. Therefore, our method can facilitate the capture of complex patterns in the data, which includes not only the shared representations but also the uniqueness of each phenotype. Our auto-branch method can also outperform the single-task models, when phenotypes shared moderate levels of common causes. This is mainly because single-task models fail to utilize information from auxiliary tasks and the efficient sample sizes are much less than those in the multi-task settings. Note that we designed our prediction model using a branch network architecture as opposed to training separate HPSs for each group of phenotypes. The rationale for such a design lies in the fact that most disease-related traits have shared genetic architecture (Badré and Pan, 2023), and by allowing all disease-related phenotypes to have some shared layers to enable the efficient modeling of these common patterns. In the unlikely event that the underlying genetic architectures differ significantly among traits in different groups, our proposed inter-trait affinity measure can still guide trait groupings. Separate HPS models can then be applied to each trait group for predictions (Supplementary Figure S3). This network structure facilitates capturing the unique characteristics of each group, making it particularly powerful for analyzing phenotypes that are largely distinct (Supplementary Tables S23, S24). Nevertheless, we recommend using branch network structure as illustrated in Figure 1 for genetic risk prediction of multiple correlated traits in most practical applications.

Our proposed method offers significant advantages in predicting seven AD-related phenotypes, including neuroimaging findings, cognitive scores, and functional assessments. For example, in the prediction of FDG, our method with 3 pre-determined branches increases the average Pearson correlations by 29.45%, 18.56% and 76.18% for Multi-Lasso, HPS, and STL, with corresponding absolute increase of 0.059, 0.031 and 0.086, respectively. Similarly, in the prediction of AV45, we have observed an increase of 33.96%, 31.64% and 29.07% for Multi-Lasso, HPS, and STL, with corresponding absolute increase of 0.061, 0.056 and 0.052, respectively. Notably, MoCA did not show significant improvement across the models, which may suggest weaker correlations with other traits or higher noise in its measurements that impacted model performance.

To verify the reliability and effectiveness of our proposed model, we conducted a comparative analysis with several related studies. For instance, Zhu et al. (2016) incorporated SNPs from the top 10 APOE-related genes into their MMSE prediction model and achieved a Pearson correlation of 0.150. In contrast, our method attained a higher correlation of 0.165 on the same task, demonstrating improved predictive performance. Hongmyeong-eup (2015) selected 39 SNPs, identified from approximately 1.5 million candidates as being closely associated with Alzheimer’s disease (AD) progression, and achieved a Pearson correlation of 0.400 in MMSE prediction. To ensure a fair comparison, we also constructed a model using SNPs from three well-established AD-associated genes—APOC1, APOE, and TOMM40—and obtained a correlation of 0.390, indicating that our approach performs comparably when using similarly strong genetic signals. In addition, Hao et al. (2016) reported Pearson correlation values ranging from 0.02 to 0.25 and 0.03–0.23 for traditional and improved methods on AV45 and FDG phenotypes, respectively. Our model also falls within these performance ranges, suggesting comparable accuracy in these tasks. In summary, our proposed method exhibited comparable predictive capabilities to existing state-of-the-art methods in multiple AD-related phenotypes, further demonstrating its potential in modeling Alzheimer’s disease progression.

These performance gains can be largely attributed to the underlying design of our method. Our method is constructed within a deep learning framework, effectively parsing complex trait relationships, particularly in the context of multi-gene co-regulation. Additionally, it can dynamically determine shared layers and thus is more powerful in managing intricate relationships among traits. Evidence suggests that these AD-related traits reflect different aspects of AD, and they are neither entirely correlated nor completely independent. For example, cognitive decline and neuroimaging abnormalities often exhibit strong genetic correlations (Zhang et al., 2021), but they do not share identical genetic causes. Therefore, methods like Multi-Lasso that rely on fixed feature selection across all tasks are unable to capture these nuanced relationships, as they do not allow for the dynamic identification of shared and unique factors between traits. Similarly, both STL and HPS models are unlikely to achieve the optimal performance, as STL fails to exploit inter-task correlations and HPS pre-specifies shared layers that lead to limited flexibility. Therefore, for the prediction of AD-related traits, our auto-branch multi-task learning method can leverage shared signals among genes or biological pathways to gain a more comprehensive understanding of these phenotypes while allowing each phenotype to have their unique characteristics, and thus offers greater flexibility as compared to Multi-Lasso, STL, and HPS.

We found that the APOC1, APOE, and TOMM40 exhibit stable and significant predictive abilities for all seven AD-related phenotypes, but their predictive capabilities vary substantially across traits, indicating HPS that forces most of the model parameter the same is unlikely to work well. All highly predictive genes are significantly associated with AD. For example, APOE4*ε4 regulates neuronal metabolism and epigenetics, and is involved in the pathological processes of AD (Prasad and Rao, 2018). APOC1 influences AD development through its role in cholesterol metabolism (Leduc et al., 2010), with the APOC1 H2 allele potentially acting synergistically with APOE to increase the risk of cognitive decline (Zhou et al., 2014). TOMM40 is strongly linked to APOE and contributes to the pathological changes in AD (Tasaki et al., 2019), including the formation of neurofibrillary tangles and neuritic plaques. TOMM40 regulates oxidative stress and mitochondrial function, and is associated with late-onset AD (Roses, 2010). We noticed that the predictive ability of APOE for FDG and AV45 is especially significant, aligning with findings from previous research (Prasad and Rao, 2018). Future studies are needed to further decipher additional factors that contribute to the variability of each phenotype.

While this study offers valuable insights into the effectiveness of multi-task learning, several limitations remain. Due to the NP-hard nature of the problem, we pre-specified the optimal number of branches and then found the corresponding partitions. Future research can improve this by determining the optimal number of branches using a data-driven manner (e.g., cross-validation). Although our method outperforms baseline approaches, factors such as limited sample size, low-to-median heritability, and trait heterogeneity may still contribute to variability and should be further explored in future work. Additionally, the findings presented in this study are based on simulated data and ADNI dataset. Future work should validate these results across diverse datasets and applications.

In summary, we developed an efficient auto-branch multi-task learning framework for the prediction analyses of multiple correlated phenotypes. It can dynamically branch the network to allow for efficient information transfer and improve the overall prediction. Our method is available at https://github.com/jiaqi69/TAB.