Jiaxuan You’s research while affiliated with Stanford University and other places

Objective: To determine whether graph neural network based models of electronic health records can predict specialty consultation care needs for endocrinology and hematology more accurately than the standard of care checklists and other conventional medical recommendation algorithms in the literature. Methods: Demand for medical expertise far outstrips supply, with tens of millions in the US alone with deficient access to specialty care. Rather than potentially months long delays to initiate diagnostic workup and medical treatment with a specialist, referring primary care supported by an automated recommender algorithm could anticipate and directly initiate patient evaluation that would otherwise be needed at subsequent a specialist appointment. We propose a novel graph representation learning approach with a heterogeneous graph neural network to model structured electronic health records and formulate recommendation/prediction of subsequent specialist orders as a link prediction problem. Results: Models are trained and assessed in two specialty care sites: endocrinology and hematology. Our experimental results show that our model achieves an 8% improvement in ROC-AUC for endocrinology (ROC-AUC = 0.88) and 5% improvement for hematology (ROC-AUC = 0.84) personalized procedure recommendations over prior medical recommender systems. These recommender algorithm approaches provide medical procedure recommendations for endocrinology referrals more effectively than manual clinical checklists (recommender: precision = 0.60, recall = 0.27, F1-score = 0.37) vs. (checklist: precision = 0.16, recall = 0.28, F1-score = 0.20), and similarly for hematology referrals (recommender: precision = 0.44, recall = 0.38, F1-score = 0.41) vs. (checklist: precision = 0.27, recall = 0.71, F1-score = 0.39). Conclusion: Embedding graph neural network models into clinical care can improve digital specialty consultation systems and expand the access to medical experience of prior similar cases.

AutoML has demonstrated remarkable success in finding an effective neural architecture for a given machine learning task defined by a specific dataset and an evaluation metric. However, most present AutoML techniques consider each task independently from scratch, which requires exploring many architectures, leading to high computational cost. Here we propose AutoTransfer, an AutoML solution that improves search efficiency by transferring the prior architectural design knowledge to the novel task of interest. Our key innovation includes a task-model bank that captures the model performance over a diverse set of GNN architectures and tasks, and a computationally efficient task embedding that can accurately measure the similarity among different tasks. Based on the task-model bank and the task embeddings, we estimate the design priors of desirable models of the novel task, by aggregating a similarity-weighted sum of the top-K design distributions on tasks that are similar to the task of interest. The computed design priors can be used with any AutoML search algorithm. We evaluate AutoTransfer on six datasets in the graph machine learning domain. Experiments demonstrate that (i) our proposed task embedding can be computed efficiently, and that tasks with similar embeddings have similar best-performing architectures; (ii) AutoTransfer significantly improves search efficiency with the transferred design priors, reducing the number of explored architectures by an order of magnitude. Finally, we release GNN-Bank-101, a large-scale dataset of detailed GNN training information of 120,000 task-model combinations to facilitate and inspire future research.

Objective To determine whether graph neural network based models of electronic health records can predict specialty consultation care needs for endocrinology and hematology more accurately than the standard of care checklists and other conventional medical recommendation algorithms in the literature. Methods Demand for medical expertise far outstrips supply, with tens of millions in the US alone with deficient access to specialty care. Rather than potentially months long delays to initiate diagnostic workup and medical treatment with a specialist, referring primary care supported by an automated recommender algorithm could anticipate and directly initiate patient evaluation that would otherwise be needed at subsequent a specialist appointment. We propose a novel graph representation learning approach with a heterogeneous graph neural network to model structured electronic health records and formulate recommendation/prediction of subsequent specialist orders as a link prediction problem. Results Models are trained and assessed in two specialty care sites: endocrinology and hematology. Our experimental results show that our model achieves an 8% improvement in ROC-AUC for endocrinology (ROC-AUC=0.88) and 5% improvement for hematology (ROC-AUC=0.84) personalized procedure recommendations over prior medical recommender systems. These recommender algorithm approaches provide medical procedure recommendations for endocrinology referrals more effectively than manual clinical checklists (recommender: precision=0.60, recall=0.27, F1-score=0.37) vs. (checklist: precision=0.16, recall=0.28, F1-score=0.20), and similarly for hematology referrals (recommender: precision=0.44, recall=0.38, F1-score=0.41) vs. (checklist: precision=0.27, recall=0.71, F1-score=0.39). Conclusion Embedding graph neural network models into clinical care can improve digital specialty consultation systems and expand the access to medical experience of prior similar cases.

Despite the success of automated machine learning (AutoML), which aims to find the best design, including the architecture of deep networks and hyper-parameters, conventional AutoML methods are computationally expensive and hardly provide insights into the relations of different model design choices. To tackle the challenges, we propose FALCON, an efficient sample-based method to search for the optimal model design. Our key insight is to model the design space of possible model designs as a design graph, where the nodes represent design choices, and the edges denote design similarities. FALCON features 1) a task-agnostic module, which performs message passing on the design graph via a Graph Neural Network (GNN), and 2) a task-specific module, which conducts label propagation of the known model performance information on the design graph. Both modules are combined to predict the design performances in the design space, navigating the search direction. We conduct extensive experiments on 27 node and graph classification tasks from various application domains, and an image classification task on the CIFAR-10 dataset. We empirically show that FALCON can efficiently obtain the well-performing designs for each task using only 30 explored nodes. Specifically, FALCON has a comparable time cost with the one-shot approaches while achieving an average improvement of 3.3% compared with the best baselines.

Graph Neural Networks (GNNs) have been successfully applied to many real-world static graphs. However, the success of static graphs has not fully translated to dynamic graphs due to the limitations in model design, evaluation settings, and training strategies. Concretely, existing dynamic GNNs do not incorporate state-of-the-art designs from static GNNs, which limits their performance. Current evaluation settings for dynamic GNNs do not fully reflect the evolving nature of dynamic graphs. Finally, commonly used training methods for dynamic GNNs are not scalable. Here we propose ROLAND, an effective graph representation learning framework for real-world dynamic graphs. At its core, the ROLAND framework can help researchers easily repurpose any static GNN to dynamic graphs. Our insight is to view the node embeddings at different GNN layers as hierarchical node states and then recurrently update them over time. We then introduce a live-update evaluation setting for dynamic graphs that mimics real-world use cases, where GNNs are making predictions and being updated on a rolling basis. Finally, we propose a scalable and efficient training approach for dynamic GNNs via incremental training and meta-learning. We conduct experiments over eight different dynamic graph datasets on future link prediction tasks. Models built using the ROLAND framework achieve on average 62.7% relative mean reciprocal rank (MRR) improvement over state-of-the-art baselines under the standard evaluation settings on three datasets. We find state-of-the-art baselines experience out-of-memory errors for larger datasets, while ROLAND can easily scale to dynamic graphs with 56 million edges. After re-implementing these baselines using the ROLAND training strategy, ROLAND models still achieve on average 15.5% relative MRR improvement over the baselines.

One of the most important factors that contribute to the success of a machine learning model is a good training objective. Training objective crucially influences the model's performance and generalization capabilities. This paper specifically focuses on graph neural network training objective for link prediction, which has not been explored in the existing literature. Here, the training objective includes, among others, a negative sampling strategy, and various hyperparameters, such as edge message ratio which controls how training edges are used. Commonly, these hyperparameters are fine-tuned by complete grid search, which is very time-consuming and model-dependent. To mitigate these limitations, we propose Adaptive Grid Search (AdaGrid), which dynamically adjusts the edge message ratio during training. It is model agnostic and highly scalable with a fully customizable computational budget. Through extensive experiments, we show that AdaGrid can boost the performance of the models up to $1.9\%$ while being nine times more time-efficient than a complete search. Overall, AdaGrid represents an effective automated algorithm for designing machine learning training objectives.

AI is undergoing a paradigm shift with the rise of models (e.g., BERT, DALL-E, GPT-3) that are trained on broad data at scale and are adaptable to a wide range of downstream tasks. We call these models foundation models to underscore their critically central yet incomplete character. This report provides a thorough account of the opportunities and risks of foundation models, ranging from their capabilities (e.g., language, vision, robotics, reasoning, human interaction) and technical principles (e.g., model architectures, training procedures, data, systems, security, evaluation, theory) to their applications (e.g., law, healthcare, education) and societal impact (e.g., inequity, misuse, economic and environmental impact, legal and ethical considerations). Though foundation models are based on conventional deep learning and transfer learning, their scale results in new emergent capabilities, and their effectiveness across so many tasks incentivizes homogenization. Homogenization provides powerful leverage but demands caution, as the defects of the foundation model are inherited by all the adapted models downstream. Despite the impending widespread deployment of foundation models, we currently lack a clear understanding of how they work, when they fail, and what they are even capable of due to their emergent properties. To tackle these questions, we believe much of the critical research on foundation models will require deep interdisciplinary collaboration commensurate with their fundamentally sociotechnical nature.

Message passing Graph Neural Networks (GNNs) provide a powerful modeling framework for relational data. However, the expressive power of existing GNNs is upper-bounded by the 1-Weisfeiler-Lehman (1-WL) graph isomorphism test, which means GNNs that are not able to predict node clustering coefficients and shortest path distances, and cannot differentiate between different d-regular graphs. Here we develop a class of message passing GNNs, named Identity-aware Graph Neural Networks (ID-GNNs), with greater expressive power than the 1-WL test. ID-GNN offers a minimal but powerful solution to limitations of existing GNNs. ID-GNN extends existing GNN architectures by inductively considering nodes’ identities during message passing. To embed a given node, ID-GNN first extracts the ego network centered at the node, then conducts rounds of heterogeneous message passing, where different sets of parameters are applied to the center node than to other surrounding nodes in the ego network. We further propose a simplified but faster version of ID-GNN that injects node identity information as augmented node features. Alto- gether, both versions of ID-GNN represent general extensions of message passing GNNs, where experiments show that transforming existing GNNs to ID-GNNs yields on average 40% accuracy improvement on challenging node, edge, and graph property prediction tasks; 3% accuracy improvement on node and graph classification benchmarks; and 15% ROC AUC improvement on real-world link prediction tasks. Additionally, ID-GNNs demonstrate improved or comparable performance over other task-specific graph networks.