Kaiyi Zhang

I am a Staff Data Scientist at Walmart Global Tech in Sunnyvale, CA. I received my Ph.D. in Computer Science from Virginia Tech in 2024, where I worked with Dr. Sharath Raghvendra in the Algorithms Lab. My doctoral research focused on optimal transport and combinatorial algorithms for machine learning. Before that, I earned an M.S. in Mining and Mineral Engineering from Virginia Tech in 2019, advised by Dr. Bahareh Nojabaei, where I focused on computational simulation of multiphase flow in complex reservoir systems.

My work has been published in leading machine learning venues, including ICML, NeurIPS, and ICLR.

Selected Projects

A New Robust Partial p-Wasserstein-Based Metric for Comparing Distributions

Sharath Raghvendra, Pouyan Shirzadian, Kaiyi Zhang (Alphabetical authorship)

ICML, 2024

TL;DR: Robust partial Wasserstein metrics that align distributions without prior assumptions, unlocking reliable PU-learning and covariate-shift comparison.

Paper / Slides / Poster / Code

Computing all Optimal Partial Transports

Abhijeet Phatak, Sharath Raghvendra, Chittaranjan Tripathy, Kaiyi Zhang (Alphabetical authorship)

ICLR, 2023

TL;DR: Enumerates every optimal partial transport solution, providing trustworthy alignment tools for imbalanced datasets and label-shift scenarios.

Paper / Slides / Code

A Combinatorial Algorithm for Approximating the Optimal Transport in the Parallel and MPC Settings

Nathaniel Lahn, Sharath Raghvendra, Kaiyi Zhang (Alphabetical authorship)

NeurIPS, 2023

TL;DR: Graph-based approximation that outperforms Sinkhorn on GPUs and massively parallel clusters, with proven bounds and linear scalability.

Paper / Code

Academic Service

• Reviewer for NeurIPS, ICML, ICLR, and IJCAI.

Teaching

• CS4104 Data and Algorithm Analysis (Virginia Tech)
• CS3214 Computer Systems (Virginia Tech)
• CS5526 Data Analytics II (Virginia Tech)

Professional Experience

For detailed industry experience, please visit my LinkedIn profile.