Jiaxuan Wang
Me at Grand Canyon, 2014
Hobbies
basketball (20+ years); guitar (4 years); violin (20+ years);
Skills
C++ (4 years); python (10+ years); Java (1 year); PyTorch (6 years);
I'm an AI research scientist at GE HealthCare. I work on improving
perinatal care using machine learning. I'm interested in making AI
interpretable and safe. I received PhD in computer science from the
University of Michigan in April 2022, advised by
professor Jenna Wiens.
My research focuses on model interpretability and model robustness to distribution shift. In addition, I'm interested in a wide range of topics including time series analysis, non-convex optimization, reinforcement learning, and sports analytics. My work is usually motivated by application in healthcare. Here's a copy of my PhD thesis.
I spent most of my pre-college years in Beijing. For the past 9 years, I've been living in Ann Arbor. I completed my undergrad at the Unisersity of Michigan as a computer science major and math minor in 2017, and started the PhD program in Fall 2017. In undergrad, I worked with professor Jia Deng to augment CNN with rotation invariant filters (2015-2016). I also worked as a machine learning intern at Bloomberg L.P. (2016) on parsing and ranking natural language query for financial charts, under the supervision of Dr. Konstantine Arkoudas and Dr. Srivas Prasad. Summer of 2020, I worked as a research intern in the Adaptive Systems and Interaction Group at Microsoft Research, mentored by Scott Lundberg, working on unifying Shapley value based model interpretation methods with a causal graph. After PhD graudation (2022), I joined Meta as a research scientist to protect user privacy. As a team, we worked on building a reinforcement learning agent to prevent adversaries from misusing user data (that is the model decides actions to apply on users/IP addresses suspected to be malicious). The setting was formulated as a contextual bandit problem (context being user/ip level features and actions include block/no response/kill sessions and etc.) and was solved using thompson sampling with baysian linear regression. From scratch, I set up a feature attribution pipeline and a dashboard that enables the team to understand and defend the model’s logic. In addition, the tool makes it easy to track both covariate and conditional shifts, which are often the culprits of degrading model performances. In addition, I worked on improving content recomendataion at Facebook, with a focus on conforming to personalized integrity preferences (i.e., reducing content that the user may give negative feedback such as xout or show less). I have experience tuning the full production ML stack: from problem definition and opportunity sizing, data gathering, model training, value model tuning (e.g., combining multiple objectives into a single number), to model monitoring through alerts and feature/model refresh. The job also provided me ample experience with LLM finetuning, model distillation, and machine translation.
In 2024, I joined GE HealthCare as a senior AI research scientist. I work on improving perinatal heath outcomes by developing AI models that interpret fetal heart rate signals and provide multimodal AI chat functionalities to help clinicians make decisions. As a young father myself, helping babies and mothers is a cause close to my heart.
Learning Concept Credible Models for Mitigating Shortcuts
TL;DR: Mitigating shortcuts with partial knowledge on relevant concepts and extending credible model to the image domain
Jiaxuan Wang, Sarah Jabbour, Maggie Makar, Michael Sjoding, Jenna Wiens
Proceedings of the 36th Conference on Neural Information Processing Systems (NeurIPS), 2022
[paper][code][mobile friendly paper]
Shapley Flow: A Graph-based Approach to Interpreting Model Predictions
TL;DR: Don't choose between true to the model or true to the data: do both and more with a system level view.
Jiaxuan Wang, Jenna Wiens, Scott Lundberg
Proceedings of the 24th International Conference on Artificial Intelligence and Statistics (AISTATS), 2021
[paper][code]
AdaSGD: Bridging the gap between SGD and Adam
TL;DR: Speed of Adam and performance of SGD may be achieved by adapting a single learning rate.
Jiaxuan Wang, Jenna Wiens
arxiv preprint, 2020
[paper][code]
Relaxed Weight Sharing: Effectively Modeling Time-Varying Relationships in Clinical Time-Series
TL;DR: Reducing temporal conditional shift using multi task learning by treating each time step as a separate task.
Jeeheh Oh*, Jiaxuan Wang*, Shengpu Tang, Michael Sjoding, Jenna Wiens
Machine Learning for Healthcare, 2019
[paper][code]
Learning Credible Models
TL;DR: Learning models that are accurate and comply with human intuition so that they don't use proxy variables.
Jiaxuan Wang, Jeeheh Oh, Haozhu Wang, Jenna Wiens
ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2018
[paper][code]
The Advantage of Doubling: A Deep Reinforcement Learning Approach to Studying the Double Team in the NBA
TL;DR: It is a bad idea to double team star playes like LeBron who can pass and score.
Jiaxuan Wang*, Ian Fox*, Jonathan Skaza, Nick Linck, Satinder Singh, Jenna Wiens
MIT Sloan Sports Analytics Conference, 2018
[paper][code]
Learning to Exploit Invariances in Clinical Time-Series Data using Sequence Transformer Networks
TL;DR: A learnable data preprocessing module for time series data in healthcare.
Jeeheh Oh, Jiaxuan Wang, and Jenna Wiens
Machine Learning for Healthcare, 2018
[paper][code]
HICO: A Benchmark for Recognizing Human-Object Interactions in Images
TL;DR: A new image dataset focusing on who did what.
Yu-Wei Chao, Zhan Wang, Yugeng He, Jiaxuan Wang, Jia Deng
International Conference on Computer Vision (ICCV) 2015
[paper][data][code]
Using feature attribution to debug and monitor distribution shift for a production ML system
TL;DR: A case study to help debug and montior for distribution shift on a production ML system, demonstrating the use of feature attribution both with and without human in the loop.
Jiaxuan Wang
2023
[paper][slides]
My research focuses on model interpretability and model robustness to distribution shift. In addition, I'm interested in a wide range of topics including time series analysis, non-convex optimization, reinforcement learning, and sports analytics. My work is usually motivated by application in healthcare. Here's a copy of my PhD thesis.
I spent most of my pre-college years in Beijing. For the past 9 years, I've been living in Ann Arbor. I completed my undergrad at the Unisersity of Michigan as a computer science major and math minor in 2017, and started the PhD program in Fall 2017. In undergrad, I worked with professor Jia Deng to augment CNN with rotation invariant filters (2015-2016). I also worked as a machine learning intern at Bloomberg L.P. (2016) on parsing and ranking natural language query for financial charts, under the supervision of Dr. Konstantine Arkoudas and Dr. Srivas Prasad. Summer of 2020, I worked as a research intern in the Adaptive Systems and Interaction Group at Microsoft Research, mentored by Scott Lundberg, working on unifying Shapley value based model interpretation methods with a causal graph. After PhD graudation (2022), I joined Meta as a research scientist to protect user privacy. As a team, we worked on building a reinforcement learning agent to prevent adversaries from misusing user data (that is the model decides actions to apply on users/IP addresses suspected to be malicious). The setting was formulated as a contextual bandit problem (context being user/ip level features and actions include block/no response/kill sessions and etc.) and was solved using thompson sampling with baysian linear regression. From scratch, I set up a feature attribution pipeline and a dashboard that enables the team to understand and defend the model’s logic. In addition, the tool makes it easy to track both covariate and conditional shifts, which are often the culprits of degrading model performances. In addition, I worked on improving content recomendataion at Facebook, with a focus on conforming to personalized integrity preferences (i.e., reducing content that the user may give negative feedback such as xout or show less). I have experience tuning the full production ML stack: from problem definition and opportunity sizing, data gathering, model training, value model tuning (e.g., combining multiple objectives into a single number), to model monitoring through alerts and feature/model refresh. The job also provided me ample experience with LLM finetuning, model distillation, and machine translation.
In 2024, I joined GE HealthCare as a senior AI research scientist. I work on improving perinatal heath outcomes by developing AI models that interpret fetal heart rate signals and provide multimodal AI chat functionalities to help clinicians make decisions. As a young father myself, helping babies and mothers is a cause close to my heart.
Publications
* denotes equal contributionLearning Concept Credible Models for Mitigating Shortcuts
TL;DR: Mitigating shortcuts with partial knowledge on relevant concepts and extending credible model to the image domain
Jiaxuan Wang, Sarah Jabbour, Maggie Makar, Michael Sjoding, Jenna Wiens
Proceedings of the 36th Conference on Neural Information Processing Systems (NeurIPS), 2022
[paper][code][mobile friendly paper]
Shapley Flow: A Graph-based Approach to Interpreting Model Predictions
TL;DR: Don't choose between true to the model or true to the data: do both and more with a system level view.
Jiaxuan Wang, Jenna Wiens, Scott Lundberg
Proceedings of the 24th International Conference on Artificial Intelligence and Statistics (AISTATS), 2021
[paper][code]
AdaSGD: Bridging the gap between SGD and Adam
TL;DR: Speed of Adam and performance of SGD may be achieved by adapting a single learning rate.
Jiaxuan Wang, Jenna Wiens
arxiv preprint, 2020
[paper][code]
Relaxed Weight Sharing: Effectively Modeling Time-Varying Relationships in Clinical Time-Series
TL;DR: Reducing temporal conditional shift using multi task learning by treating each time step as a separate task.
Jeeheh Oh*, Jiaxuan Wang*, Shengpu Tang, Michael Sjoding, Jenna Wiens
Machine Learning for Healthcare, 2019
[paper][code]
Learning Credible Models
TL;DR: Learning models that are accurate and comply with human intuition so that they don't use proxy variables.
Jiaxuan Wang, Jeeheh Oh, Haozhu Wang, Jenna Wiens
ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2018
[paper][code]
The Advantage of Doubling: A Deep Reinforcement Learning Approach to Studying the Double Team in the NBA
TL;DR: It is a bad idea to double team star playes like LeBron who can pass and score.
Jiaxuan Wang*, Ian Fox*, Jonathan Skaza, Nick Linck, Satinder Singh, Jenna Wiens
MIT Sloan Sports Analytics Conference, 2018
[paper][code]
Learning to Exploit Invariances in Clinical Time-Series Data using Sequence Transformer Networks
TL;DR: A learnable data preprocessing module for time series data in healthcare.
Jeeheh Oh, Jiaxuan Wang, and Jenna Wiens
Machine Learning for Healthcare, 2018
[paper][code]
HICO: A Benchmark for Recognizing Human-Object Interactions in Images
TL;DR: A new image dataset focusing on who did what.
Yu-Wei Chao, Zhan Wang, Yugeng He, Jiaxuan Wang, Jia Deng
International Conference on Computer Vision (ICCV) 2015
[paper][data][code]
Technical reports
* denotes equal contributionUsing feature attribution to debug and monitor distribution shift for a production ML system
TL;DR: A case study to help debug and montior for distribution shift on a production ML system, demonstrating the use of feature attribution both with and without human in the loop.
Jiaxuan Wang
2023
[paper][slides]