Multi-task Learning for Brain Network Analysis in the ABCD study

Authors: xuan  kan
hejie  cui
keqi  han
ying  guo
carl  yang
Venue: IEEE_org_EMBS_BHI_2024_Conference
Published: 2024-06-13 07:23:42

Abstract: The Adolescent Brain Cognitive Development study provides a rich data resource for exploring the associations between brain connectome (network) and cognitive, personality, and mental health measures in adolescents. To leverage this rich dataset, we propose a novel multi-task learning framework that predicts these measures from multi-view brain network data using a graph transformer architecture. Our approach learns shared representations across tasks while allowing for task-specific predictions, improving performance compared to single-task learning. Ablation studies reveal the importance of our proposed techniques of Batch-Wise Loss Balancing and Target Standardization in ensuring successful multi-task learning. Furthermore, we develop innovative visualization techniques based on integrated gradients to interpret the learned task correlations and identify influential brain network edges for each task. Our findings contribute to understanding the complex relationships between brain connectome and behavioral outcomes, highlighting the potential of multi-task learning in this domain. The implementation is available at https://github.com/Wayfear/MTML-ABCD/.

Main paper: 
 i. introduction: 
 Adolescent Brain Cognitive Development (ABCD) study [1] is the largest and most long-term study of brain development and child health in the US. It provides a vast brain development dataset in a diverse population, including functional magnetic resonance imaging data (fMRI) and abundant biological and behavioral survey results. This dataset offers an opportunity to explore the relationship between intricate brain connections and various behavioral data [2] - [7] .Leveraging the potential of neuroimaging data, recent studies have shown a growing trend of using brain networks derived from fMRI to predict various clinical outcomes and individual behaviors with different models [8] - [16] . Researchers have also developed innovative approaches to analyze these models and uncover potential correlations between functional brain networks and predicted outcomes [17] - [23] . For example, Kawahara et al [9] introduced BrainNetCNN, a convolutional neural network designed to predict cognitive and motor developmental outcome scores from brain networks. Similarly, Li et al [10] proposed a graph neural network model to predict clinical targets and discovered task-specific neurological biomarkers, demonstrating the effectiveness of graph-based approaches in capturing meaningful patterns in brain networks. Chen et al [24] further extended this line of research by training kernel regression models for 36 tasks and analyzing task relationships based on the learned model. These studies highlight the potential of leveraging brain networks to gain insights into the underlying neural mechanisms associated with various clinical outcomes and individual behaviors.Multi-task learning (MTL) [25] - [28] has emerged as a promising approach for improving the generalization abilities of predictive models by enabling multiple learning tasks to share their knowledge. In the context of brain network analysis in ABCD, where there is a diverse range of prediction targets, MTL can be particularly beneficial. By training several tasks simultaneously, MTL allows for a more native capture of task correlations, potentially leading to improved individual task performance. This contrasts with the approach taken by Chen et al [24] , which builds individual models for each behavior task. By leveraging the shared representations learned across tasks, MTL can enhance the model's ability to uncover the underlying relationships between brain networks and various behavioral measures, resulting in more accurate and generalizable predictions.Our study employs a shared bottom multi-task learning (MTL) approach for brain network analysis using the ABCD dataset, predicting multiple behavioral outcomes from restingstate functional connectivity. This differs from Xiao et al.'s work [29] , which uses a manifold regularized MTL model to predict IQ scores from two task-based fMRI paradigms. While both leverage MTL for fMRI analysis, our study addresses a broader range of prediction tasks and uses resting-state data, potentially capturing more general brain organization patterns. In contrast to Marquand et al.'s Bayesian MTL framework for multi-subject decoding [30] , our approach focuses on learning shared representations across tasks rather than subjects. We use a neural network instead of a Bayesian framework, which may offer computational efficiency advantages but potentially at the cost of explicit uncertainty quantification. Like ours, Huang et al.'s study uses resting-state fMRI data for MTL, but employs a Multi-cluster Multi-gate Mixture-of-Experts (M-MMOE) model for joint diagnosis of ASD and ADHD [31] . Our shared bottom model provides a simpler, more interpretable approach, while the M-MMOE allows for more complex task relationships. Additionally, our study encompasses a broader range of behavioral measures beyond just ASD and ADHD diagnosis, potentially offering insights into a wider spectrum of brain-behavior relationships.In this work, we propose a novel MTL framework that jointly trains 35 tasks using multi-view functional brain networks from 6,682 samples in the ABCD study. We employ the Brain Network Transformer [32] as the backbone model, which converts a brain network into a graph-level embedding. This embedding is then fed into task-specific fully connected networks for each prediction target. By learning shared representations across tasks while still allowing for task-specific predictions, our approach aims to leverage the commonalities between tasks and improve overall prediction performance. Our main contributions are summarized as follows:• We propose a novel MTL framework for predicting various measures from multi-view brain network data using a graph transformer architecture. Our approach learns shared representations across tasks while allowing for task-specific predictions, improving performance compared to single-task learning. Besides, the ablation study shows the effectiveness of two key training strategies during Multi-Task training. • We conduct extensive experiments on the ABCD dataset, including 35 tasks categorized into three domains: cogni-tion, personality, and mental health. We demonstrate the impact of MTL on different types of tasks. • We develop innovative visualization techniques based on integrated gradients to interpret the learned task correlations and identify influential brain network edges, contributing to a better understanding of the complex relationships between brain connectome and behavioral outcomes.Let D = {(X (i) , Y (i) )} n i=1 be a dataset consisting of n samples. For each sample i, the inputX (i) = {X (i) 1 , X (i) 2 , . . . , X (i)v } represents a collection of v brain networks, each derived from a distinct fMRI task (e.g., restingstate, stop-signal task, and N-Back). These brain networks, denoted by X (i) j ∈ R M ×M , capture the functional connectivity between M brain regions. The corresponding prediction targetY (i) = {y (i) 1 , y (i) 2 , . . . , y (i)t } is a set of t behavioral measures, such as cognitive scores, personality traits, or mental health indicators, associated with the i-th subject. In short, given this multi-view, multi-task dataset, our goal is to develop a predictive model that leverages the complementary information from the v brain networks to simultaneously predict the t behavioral outcomes.Figure 1 shows the proposed multi-task learning framework for predicting behavioral outcomes from brain networks. Shared Representation Learning. The footstone of our framework is Brain Network Transformer (BNT) [32] , which serves as the shared backbone model. BNT is designed to process individual views of brain networks, denoted as X (i) j , where j ∈ {1, . . . , v} indexes the view and i ∈ {1, . . . , n} indexes the sample. For each view j, BNT learns a hidden representation embedding h (i) j = BNT(X (i) j ). These view-specific embeddings capture the patterns present in the corresponding brain networks derived from distinct fMRI tasks. To obtain a comprehensive representation of each sample, we concatenate the view-specific embeddings h (i) j from all v views, resulting in a sample-level embeddingH (i) = v j=1 h (i) j, where denotes the concatenation operation. This sample-level embedding integrates the multifaceted information captured across different fMRI views, providing a holistic representation of each individual's brain connectivity patterns. Task-specific Prediction. To achieve multi-task learning, we employ a separate Multi-Layer Perceptron (MLP) for each task k ∈ {1, . . . , t}. These task-specific MLPs take the samplelevel embedding H (i) as input and predict the corresponding behavioral outcomeŷ(i) k = MLP k (H (i) ). By leveraging dedicated MLPs for each task, our framework allows for task-specific adaptations while benefiting from the shared representation learned by the BNT.View1: Rest-State 𝑋 ! (#) Brain Network Transformer Brain Network Embedding 𝐻 (#) … Predict Working Memory Score 𝑦 " ! (#) FCN 1 Predict Vocabulary Score 𝑦 " % (#) FCN 2 Predict Lack-of-planning Score 𝑦 " & (#) FCN 3 Predict Anxious Depressed Score 𝑦 " ' (#) FCN 4 Predict Psychosis Severity Score 𝑦 " ( (#) FCN t … View2: SST 𝑋 % (#) View3: N-Back 𝑋 & (#) View Embedding ℎ ! (#) View Embedding ℎ % (#) View Embedding ℎ & (#)Fig. 1 . Overview of our multi-task learning framework for predicting various measures from multi-view brain networks. Given a set of brain networks{X (i) 1 , X (i) 2 , . . . , X (i)v } derived from different views for the i-th subject, the Brain Network Transformer generates a unified brain network embedding H (i) . This embedding is then fed into task-specific FCN to predict the corresponding target scores {ŷ(i) 1 , ŷ(i) 2 , . . . , ŷ(i)t } for various measures. The entire framework is trained end-to-end using multi-task learning, allowing for the sharing of knowledge across tasks while still enabling task-specific predictions.The entire framework is trained end-to-end using a multitask learning approach. However, training a model to simultaneously predict multiple behavioral outcomes presents challenges due to the diverse characteristics and varying scales of the prediction targets. We introduce two key strategies to address these challenges and ensure effective training: Batch-Wise Loss Balancing and Target Standardization. Batch-Wise Loss Balancing. In multi-task learning, tasks with larger loss values can potentially dominate the training process, hindering the model's ability to learn from all tasks equally. To mitigate this issue, we develop batch-wise loss balancing that adaptively adjusts the weight of each task's loss within a training batch. Let L k denote the loss associated with task k, where k ∈ {1, . . . , t}. We compute the balanced loss Lk for each task as follows:Lk = L kLk , where Lk is the average loss value for task k over the current batch of samples. By normalizing each task's loss to have an average value of 1, we ensure that all tasks contribute equally to the overall optimization process. The total balanced loss L total is then calculated as the sum of all individual balanced task losses: L total = t k=1 Lk . Target Standardization. Another challenge in multi-task learning is the varying scales of the target variables across different tasks. To address this issue, we employ a target standardization preprocessing step. For a regression task k, where k ∈ {1, . . . , t}, we standardize the training labels y (i) k to have zero mean and unit variance:y (i) k = y (i) k -µ k σ k, where µ k and σ k denote the mean and standard deviation of the training labels for task k, respectively. This normalization brings all tasks to a similar scale, facilitating the model's ability to learn from them concurrently. During the validation and testing phase, we apply the inverse of this normalization process to transform the predicted labels ŷ(i) k back to their original scale. Using the mean µ k and standard deviation σ k computed from the training set, we perform the following operation:ŷ(i) k = ŷ(i) k • σ k + µ k .By standardizing the targets during training and reversing the normalization during inference, we ensure that the model can effectively learn from tasks with different scales while producing predictions in the original target domain.The following section describes our experimental methodology and findings, evaluating our multi-task learning framework for brain network analysis using the ABCD dataset. We detail our data, model architecture, training procedures, and evaluation metrics. We then present results comparing our multi-task approach to single-task baselines across various behavioral measures, followed by an ablation study on our key training strategies. Finally, we analyze task correlations learned by our model, providing insights into brain-behavior relationships in adolescent development.We use the Adolescent Brain Cognitive Development (ABCD) dataset [1] , which includes fMRI and behavioral data from a large cohort of children. We utilize restingstate and task-based fMRIs (stop-signal task [33] and N-Back task [34] ) for brain network construction based on the HCP 360 ROI atlas [35] . In this study, we aim to predict 35 distinct labels, which span across 15 neurocognitive ability [36] , 9 impulsivity-related personality and 11 mental health assessments [37] , as detailed in Table I . The experimental dataset includes 6,682 samples with quality control procedures and filtering samples with incomplete fMRI and behavioral data.We employ a Brain Network Transformer as the shared model backbone, which consists of 3 transformer layers with 4 attention heads and an output dimension matching the number of nodes (360) in the brain network. The transformer layers are followed by task-specific 3-layer MLP branches with activation functions, each responsible for making predictions for one of the tasks. Since all tasks are regression tasks, we use Mean Squared Error (MSE) Loss as the loss function for all tasks. The model, which can predict 35 tasks simultaneously, has a total of 22.52 million parameters. We randomly split To evaluate our multi-task learning model's performance on the 35 regression tasks from the ABCD dataset, we use two metrics: Mean Squared Error (MSE) and R-squared (R 2 ). MSE measures the average squared difference between predicted and actual values, with lower values indicating better performance. R 2 , on the other hand, is used to compare performance across tasks with varying label scales. R 2 values range from -∞ to 1, where negative values indicate worse performance than using the target variable's mean, zero indicates equivalence to using the mean, and positive values suggest the model captures useful information from brain networks and the prediction beats the mean of the target. Thus, R 2 can be used to evaluate a task's predictability. All reported results are averaged over 5 runs with different random seeds.The overall results of our multi-task learning model on the ABCD dataset are shown in Table II . From the table, we can obtain 3 key insights: (1) Single-task performance: The Single-Task column reveals that for all Personality and Mental Health tasks, these models' R 2 is below 0.03, indicating that there is limited predictive power when using brain networks to predict these labels. In contrast, for the Cognition tasks, except for the Visuospatial Reaction Time task, all other 14 tasks have an R 2 greater than or equal to 0.03. This suggests that the model can capture useful information from the brain networks and outperform predictions based solely on the mean of the target variable for these Cognition tasks; (2) Multi-task learning benefits: By comparing the Single-Task performance column with the Multi-Task performance column, we observe that multi-task training improves the performance of almost all tasks that already exhibit predictive power in the singletask setting. However, the Personality and Mental Health tasks that were unpredictable in the single-task setting remain unpredictable in the multi-task setting, indicating that these tasks cannot effectively leverage useful information from other tasks during joint training; (3) Impact of removing unpredictable tasks: To further investigate the influence of the unpredictable Personality and Mental Health tasks on the overall model performance, we conducted an additional experiment where we removed these tasks during multi-task training. The results of this experiment are shown in the Multi-CogTask column. Interestingly, we observe that by excluding these unpredictable tasks, the performance of the remaining tasks drops compared to the multi-task setting that includes all tasks. This finding suggests that labeling information from Personality and Mental Health tasks is still helpful for other tasks, even though these tasks themselves remain unpredictable.We investigate the effectiveness of our two key training strategies in our multi-task learning model: Batch-Wise Loss Balancing and Target Standardization. We compare the performance of our full model with three ablated versions:(1) without Batch-Wise Loss Balancing, (2) without Target Standardization, and (3) without both strategies. The results in Fig. 2 show that removing Batch-Wise Loss Balancing leads to a slight decrease in performance across all 14 predictable tasks, while removing Target Standardization causes a significant drop. When both strategies are removed, the model fails to learn any meaningful information, resulting in negative R² values for all tasks. This study demonstrates the importance of these training strategies in enabling successful multi-task learning for brain network analysis.In this section, we visualize the task correlations learned by our multi-task learning model using the integrated gradients method [38] and compare them with the inherent correlations between task labels. Finally, we visualize these important edges for prediction in different tasks.Algorithm 1 illustrates the process to obtain the task-level correlation matrix C and edge importance G k for each task k. The algorithm uses these symbols: M for the trained multitask learning model, n s for the number of samples selected from the test set, X (i) and Y (i) for the input brain networks and target labels of sample i, respectively, v for the number of views (e.g. rest-state or N-Back task), M for the number of nodes in each brain network, IGj,e,k for the integrated gradients of edge e and task k in view j for sample i, t for the total number of tasks, Ī G k j for the average edge importance of view j and task k, g k j for the vectorized edge importance of view j and task k, Ĝk for the concatenated edge importance across views for task k, G k for the averaged edge importance across views for task k, and C for the correlation matrix between task-level edge importance vectors. for each view j ∈ {1, . . . , v} do 5:for each edge e ∈ {1, . . . , M 2 } do 6:for each task k ∈ {1, . . . , t} do G k = 1 v v j=1 Ī G k j▷ Average IG across views for task k 18: end for19: Compute correlation matrix C ∈ R t×t between { Î G k } t k=1The correlation matrix C (Fig. 3 Left) reveals the similarities between tasks based on the brain network edges that the model deems important for making predictions. In other words, tasks that have a high positive correlation in this matrix rely on similar brain network edges for their predictions, suggesting that the model has learned shared patterns for these tasks. This matrix provides insights into the model's internal representation and how it leverages the brain network information for different tasks.On the other hand, the label correlation matrix (Fig. 3 Right) represents the inherent relationships between the task variables themselves, independent of the model's learning process. This matrix shows the correlations between the target variables of different tasks, revealing the intrinsic similarities or differences among them. By comparing these two matrices, we can assess the model's ability to capture meaningful task relationships from the brain network data. The results highlight two key findings: (a) Strong positive correlations among the 14 most predictable tasks: The model learns shared patterns for these tasks, as evidenced by the high positive correlations in the integrated gradients-based task correlation matrix. This suggests that the model has successfully identified common brain network edges that are predictive of these tasks, aligning with their inherent similarities revealed in the label correlation matrix; (b)Distinct patterns for personality and mental health task groups: Within each of these groups, tasks exhibit strong positive correlations in the integrated gradients-based task correlation matrix, indicating that the model learns shared patterns for tasks within the same group. However, the correlations between these two groups are weak, suggesting that the model learns distinct patterns for personality and mental health tasks. This finding implies that the model has captured the unique brain network edges that are relevant for each group of tasks, reflecting their underlying differences.These observations demonstrate that our multi-task learning model effectively captures both the shared patterns among similar tasks and the distinct patterns between different task groups. By learning these task relationships from the brain network data, the model can leverage the commonalities among tasks to improve its predictions while still maintaining the ability to capture task-specific patterns. This highlights the power of our multi-task learning approach in uncovering meaningful relationships and utilizing them for predictions.Furthermore, the alignment between the integrated gradients-based task correlation matrix and the label correlation matrix validates the model's interpretability. The model's learned task relationships, as revealed by the important brain network edges, match the inherent similarities and differences among the task variables. This interpretability strengthens the trustworthiness of our model and its potential for application in clinical and research settings.To further investigate the task-specific patterns learned by our multi-task learning model, we visualize the most important edges for a subset of tasks in Figure 4 . These edges are determined based on the integrated gradients G k obtained for each task k. Specifically, we select the top 0.05% of edges with the highest integrated gradient magnitudes and visualize them on a brain network template using BrainNet Viewer [39] . The node colors in the visualizations represent different functional modules, while the edge colors and thickness indicate the magnitude of the integrated gradients.Figure 4 showcases the important edges for six tasks: two from the Cognition group (Vocabulary and Reading), two from the Personality group (Drive and Fun Seeking), and two from the Mental Health group (Rule-breaking Behavior and Aggressive Behavior). By comparing the visualizations within and across task groups, we observe several interesting patterns. First, tasks within the same group tend to have similar important edges. For example, the Vocabulary and Reading tasks, both belonging to the Cognition group, share many common edges in the visualized brain networks. Similarly, the Drive and Fun Seeking tasks from the Personality group also exhibit overlapping important edges. This observation aligns with the within-group correlations depicted in the integrated gradients-based task correlation matrix (Figure 3 , Left).On the other hand, when comparing tasks from different groups, we notice distinct patterns in the important edges. For instance, the important edges for the Vocabulary task (Cognition group) differ significantly from those of the Drive task (Personality group). This finding suggests that our multitask learning model captures task-specific patterns while also learning shared representations within task groups.The visualization of important edges provides valuable insights into the brain network regions and connections that the model relies on for making predictions in each task.We proposed a novel multiple-task learning framework for predicting cognitive, personality, and mental health measures from brain networks using the ABCD dataset. Our approach effectively captures meaningful relationships across tasks and improves prediction performance compared to single-task learning. Through experiments, we demonstrated the importance of two training strategies and provided deep task correlation analysis by the integrated gradient method, presenting a significant step towards understanding the complex relationships between brain connectome and behavioral outcomes.This research was partially supported by the National Science Foundation under Award Numbers 2319449 and 2312502, as well as the National Institute of Health under Award Numbers R01MH105561 and R01MH118771. 
 
 Appendix:  



Reviews: review: title: title: It is developed and implemented a multi-task learning framework for brain network analysis using the ABCD dataset. This framework predicts cognitive, personality, and mental health measures by capturing meaningful inter-task relationships and outperforming single-task learning models.
paper_summary: paper_summary: In this work, it is proposed a novel multi-task learning framework utilizing the ABCD dataset to predict cognitive, personality, and mental health measures from brain networks. This approach captures meaningful inter-task relationships and outperforms single-task learning models.  It is demonstrated the importance of two key training strategies and employed integrated gradient methods for deep task correlation analysis.
strength_weakness: strength_weakness: The authors recognize that the correlations between the two groups are weak, indicating that the model learns distinct patterns for personality and mental health tasks. This suggests that the model has identified unique brain network connections relevant to each group of tasks, highlighting their underlying differences.
questions: questions: Noting to point out.
title: title: Multi-task Learning for Brain Network Analysis in the ABCD study
paper_summary: paper_summary: This work presents a multi-task learning framework for predicting 35 tasks using fMRIs acquired from the Adolescent Brain Cognitive Development (ABCD) study. Transformers were employed as the backbone for shared feature learning, while multiple layer perceptrons were used for individual task prediction. Additionally, a visualization technique was developed for the interpretation of learned task correlations.
strength_weakness: strength_weakness: While the application to the ABCD study is new, the multi-task learning approach presented in this work appears to be quite standard, limiting the overall innovation.
questions: questions: none
title: title: Multi-task Learning for Brain Network Analysis in the ABCD study
paper_summary: paper_summary: In this paper, authors propose to study multi-task learning (using a graph-based architecture) to predict cognition, personality and mental health outcomes from functionnal magnetic resonance imaging data of the Adolescent Brain Cognitive Development study.
strength_weakness: strength_weakness: The paper may lack of positioning according to the current state of art on multitask learning from fMRI and comparison/discussion with other techniques.

There is some redundancies in the model architecture part regarding the introduction of the Brain Network Transformer (BNT). I would suggest to merge "Brain Network Transformer" and "Shared Representation Learning" paragraph

It may be more relevant to put author's methodological contributions (resumed in the strenghts part) in evidence into the method section. 

Minor comments :
- Add a sentence to explain how scores were defined in the ABCD study.
- Remove "our ablation study results in Section III-E show the necessity of incorporating them during training"
- Add an introductive paragraph to explain the content of the experiments section
- Precise on which approach was based the R^2 highlighting in Table II
questions: questions: - Why Chen et al. target 36 tasks while 35 tasks are investigated in this work?
- How was performed the Single-Task strategy?