Jiahao Ai

Hi! I’m Jiahao Ai, an undergraduate in Statistics at Peking University with a second major in Computer Science. I will join the Ph.D. program in Statistics and Data Science at the Wharton School, University of Pennsylvania in Fall 2025.

I have been fortunate to receive guidance from Prof. Zhimei Ren, Prof. Yuxin Chen, and Prof. Ruixun Zhang during my undergraduate years.Their mentorship has shaped my thinking and nurtured my passion for research.

My research interests lie at the intersection of uncertainty quantification, robust statistical inference, optimization & online learning, and reinforcement learning.

selected publications

ICML

Not all distributional shifts are equal: Fine-grained robust conformal inference

Jiahao Ai and Zhimei Ren

In Proceedings of the 41st International Conference on Machine Learning, 21–27 jul 2024

Abs arXiv PDF Code Poster

We introduce a fine-grained framework for uncertainty quantification of predictive models under distributional shifts. This framework distinguishes the shift in covariate distributions from that in the conditional relationship between the outcome (Y) and the covariates (X). We propose to reweight the training samples to adjust for an identifiable shift in covariate distribution while protecting against the worst-case conditional distribution shift bounded in an f-divergence ball. Based on ideas from conformal inference and distributionally robust learning, we present an algorithm that outputs (approximately) valid and efficient prediction intervals in the presence of distributional shifts. As a use case, we apply the framework to sensitivity analysis of individual treatment effects with hidden confounding. The proposed methods are evaluated in simulations and four real data applications, demonstrating superior robustness and efficiency compared with existing benchmarks.