Research

Information Borrowing across Heterogeneous Data Sources

Calibrating auxiliary information from AI-generated or observational data

Generation-Powered Inference for Distribution-valued Outcomes Yijiao Zhang and Hongzhe Li. arXiv preprint, 2026. [arXiv] Abstract

Abstract

Modern generative models increasingly produce distribution-valued outputs, such as predicted cellular responses to genetic perturbations in single-cell genomics. While these models provide valuable auxiliary information, they are inherently imperfect, creating a need for statistical methods that leverage their predictions without relying on their correctness. We propose generation-powered inference (GPI), a general framework for improving inference on distribution-valued parameters using auxiliary generative models. Focusing on Wasserstein barycenters and related distributional functionals, we introduce a function-valued bridge representation that transforms inference in the nonlinear Wasserstein space into estimation of a mean function in a Hilbert space, enabling an augmented estimation framework analogous to prediction-powered inference. We develop a family of GPI estimators with optimal information borrowing, establish consistency, asymptotic normality, and simultaneous confidence bands, and derive valid inference for linear functionals and Wasserstein distances. Simulation studies demonstrate efficiency gains over labeled-data-only methods and robust performance under generative model misspecification. We illustrate the proposed framework using a Perturb-seq study of K562 cells, where synthetic perturbation responses generated by the State foundation model are used to improve inference for pathway-level consensus gene expression distributions associated with perturbations of the 40S ribosome module.

A Data Fusion Method for Quantile Treatment Effects Yijiao Zhang and Zhongyi Zhu. Statistica Sinica, 35, 981–1002, 2025. [Paper] [arXiv] Abstract

Abstract

With the increasing availability of datasets, developing data fusion methods to leverage the strengths of different datasets to draw causal effects is of great practical importance to many scientific fields. In this paper, we consider estimating the quantile treatment effects using small validation data with fully-observed confounders and large auxiliary data with unmeasured confounders. We propose a Fused Quantile Treatment effects Estimator (FQTE) by integrating the information from two datasets based on doubly robust estimating functions. We allow for the misspecification of the models on the dataset with unmeasured confounders. Under mild conditions, we show that the proposed FQTE is asymptotically normal and more efficient than the initial QTE estimator using the validation data solely. By establishing the asymptotic linear forms of related estimators, convenient methods for covariance estimation are provided. Simulation studies demonstrate the empirical validity and improved efficiency of our fused estimators. We illustrate the proposed method with an application.

Adaptive information transfer across related sources

Transfer Learning for High-dimensional Quantile Regression via Convolution Smoothing Yijiao Zhang and Zhongyi Zhu. Statistica Sinica, 35, 939–958, 2025. [Paper] [arXiv] Abstract

Abstract

This paper studies the high-dimensional quantile regression problem under the transfer learning framework, where possibly related source datasets are available to make improvements on the estimation or prediction based solely on the target data. In the oracle case with known transferable sources, a smoothed two-step transfer learning algorithm based on convolution smoothing is proposed and the L1/L2 estimation error bounds of the corresponding estimator are also established. To avoid including non-informative sources, we propose to select the transferable sources adaptively and establish its selection consistency under regular conditions. Monte Carlo simulations as well as an empirical analysis of gene expression data demonstrate the effectiveness of the proposed procedure.

Transfer Learning for High-dimensional Quantile Regression with Distribution Shift Ruiqi Bai, Yijiao Zhang†, Hanbo Yang, and Zhongyi Zhu†. arXiv preprint, 2026. [arXiv] Abstract

Abstract

Transfer learning can improve estimation in a target study by borrowing information from related source studies, but its effectiveness can deteriorate when the source and target distributions differ. We study high-dimensional quantile regression under three forms of distribution shift: parameter shift, covariate shift, and residual distribution shift. We introduce a transferability criterion and a transfer framework designed to remain useful under these simultaneous discrepancies. Nonasymptotic estimation bounds characterize the gains from informative sources, while source-detection results establish conditions under which transferable studies can be identified consistently. We also develop an orthogonal debiasing procedure for valid coefficient inference after transfer, leading to sharper asymptotic results. Simulation studies and real-data analyses demonstrate the benefits of the proposed procedures across a range of heterogeneous transfer settings.

Covariate-Elaborated Robust Partial Information Transfer with Conditional Spike-and-Slab Prior Ruqian Zhang*, Yijiao Zhang*, Juan Shen, Zhongyi Zhu, and Annie Qu. Journal of the American Statistical Association, to appear, 2026. [Paper] Abstract

Abstract

The popularity of transfer learning stems from the fact that it can borrow information from useful auxiliary datasets. Existing statistical transfer learning methods usually adopt a global similarity measure between the source data and the target data, which may lead to inefficiency when only partial information is shared. In this paper, we propose a novel Bayesian transfer learning method named “CONCERT” to allow robust partial information transfer for high-dimensional data analysis. A conditional spike-and-slab prior is introduced in the joint distribution of target and source parameters for information transfer. By incorporating covariate-specific priors, we can characterize partial similarities and integrate source information collaboratively to improve the performance on the target. In contrast to existing work, the CONCERT is a one-step procedure which achieves variable selection and information transfer simultaneously. We establish variable selection consistency, as well as estimation and prediction error bounds for CONCERT. Our theory demonstrates the covariatespecific benefit of transfer learning. To ensure the scalability of the algorithm, we adopt the variational Bayes framework to facilitate implementation. Extensive experiments and two real data applications showcase the validity and advantages of CONCERT over existing cutting-edge transfer learning methods.

Optimal Transport based Cross-Domain Integration for Heterogeneous Data Yubai Yuan, Yijiao Zhang, Babak Shahbaba, Norbert Fortin, Keiland Cooper, Qing Nie, and Annie Qu. Journal of the American Statistical Association, 120(551), 1449–1462, 2025. [Paper] Abstract

Abstract

Detecting dynamic patterns shared across heterogeneous datasets is a critical yet challenging task in many scientific domains, particularly within the biomedical sciences. Systematic heterogeneity inherent in diverse data sources can significantly hinder the effectiveness of existing machine learning methods in uncovering shared underlying dynamics. Additionally, practical and technical constraints in real-world experimental designs often limit data collection to only a small number of subjects, even when rich, time-dependent measurements are available for each individual. These limited sample sizes further diminish the power to detect common dynamic patterns across subjects. In this article, we propose a novel heterogeneous data integration framework based on optimal transport to extract shared patterns in the conditional mean dynamics of target responses. The key advantage of the proposed method is its ability to enhance discriminative power by reducing heterogeneity unrelated to the signal. This is achieved through the alignment of extracted domain-shared temporal information across multiple datasets from different domains. Our approach is effective regardless of the number of datasets and does not require auxiliary matching information for alignment. Specifically, the method aligns longitudinal data from heterogeneous datasets within a common latent space, capturing shared dynamic patterns while leveraging temporal dependencies within subjects. Theoretically, we establish generalization error bounds for the proposed data integration approach in supervised learning tasks, highlighting a novel tradeoff between data alignment and pattern learning. Additionally, we derive convergence rates for the barycentric projection under Gromov-Wasserstein and fused Gromov-Wasserstein distances. Numerical studies on both simulated data and neuroscience applications demonstrate that the proposed data integration framework substantially improves prediction accuracy by effectively aggregating information across diverse data sources and subjects. Supplementary materials for this article are available online, including a standardized description of the materials available for reproducing the work.

Exploiting Heterogeneity for Causal and Invariant Learning

Causal Network Discovery from Interventional Count Data with Latent Linear DAGs Yijiao Zhang and Hongzhe Li. arXiv preprint, 2026. [arXiv] Abstract

Abstract

The increasing availability of interventional data offers new opportunities for causal discovery, with gene perturbation studies providing a prominent example. Such data are typically count-valued and subject to substantial measurement error arising from technical variability and latent state heterogeneity. Motivated by these challenges, we study identification and estimation in latent linear structural causal models for interventional count data. We propose a latent linear Gaussian directed acyclic graph (DAG) model with Poisson measurement error that explicitly separates the latent causal structure from the observed counts. Under a mean-shift intervention design, we establish population-level identifiability of the latent causal DAG. Building on these identification results, we develop an estimation procedure based on sparse inverse matrix estimation and provide theoretical guarantees on estimation error and finite-sample causal discovery. Simulation studies and applications to Perturb-seq data demonstrate the practical effectiveness of the proposed method.

Bayesian Environment Invariant Regression Ruqian Zhang, Juan Shen, and Yijiao Zhang†. arXiv preprint, 2026. [arXiv] Abstract

Abstract

The availability of data from multiple heterogeneous environments has motivated methods that remain reliable under distributional shifts. When the joint distribution of response and predictors varies across environments, the response may still depend on a subset of predictors through an invariant mechanism. Existing methods typically assess candidate invariant sets through pooled stability criteria, treating environmental variation as nuisance. In this paper, we propose a Bayesian framework that explicitly separates a shared response mechanism from environment-specific or response-dependent associations, exploiting heterogeneity as evidence for structure learning. A competitive spike-and-slab prior is designed to force each predictor to compete between invariant and non-invariant spurious effects. Under a tractable working model, we establish invariant model selection consistency and posterior contraction for invariant coefficients. We further study the presence of irrelevant predictors, characterize posterior concentration on an equivalent invariant class, and introduce a post-selection refinement that consistently recovers the minimal invariant model. Simulations and a real application illustrate the robustness and finite-sample efficiency of the proposed method.

Individualized Learning with Shared Structure

Individualized Dynamic Mediation Analysis Using Latent Factor Models Yijiao Zhang, Yubai Yuan, Yuexia Zhang, Zhongyi Zhu, and Annie Qu. Journal of the American Statistical Association, to appear, 2026. [Paper] Abstract

Abstract

Mediation analysis plays a crucial role in causal inference as it can investigate the pathways through which treatment influences outcome. Most existing mediation analysis assumes that mediation effects are static and homogeneous within populations. However, mediation effects usually change over time and exhibit significant heterogeneity among individuals in many real-world applications. Additionally, the mediation mechanism can be complicated and involves non-sparse, making mediator selection particularly challenging. To address these issues, we propose an individualized dynamic mediation analysis method for mediator selection. Our approach can identify the significant mediators at the population level while capturing the time-varying and heterogeneous mediation effects at the individual level via varying-coefficient structural equation models. Another advantage of our method is that we allow the presence of unmeasured time-varying confounders that induce the heterogeneous mediation effects. We provide asymptotic results for the proposed estimator and selection consistency for significant mediators. Extensive simulation studies and an application to a DNA methylation study demonstrate the effectiveness and advantages of our method.

Syntax-Guided Diffusion Language Models with User-Integrated Personalization Ruqian Zhang, Yijiao Zhang, Juan Shen, Zhongyi Zhu, and Annie Qu. The Annals of Applied Statistics, to appear, 2026. [arXiv] Abstract

Abstract

Information from related source studies can often enhance the findings of a target study. However, the distribution shift between target and source studies can severely impact the efficiency of knowledge transfer. In the high-dimensional regression setting, existing transfer approaches mainly focus on the parameter shift. In this paper, we focus on the high-dimensional quantile regression with knowledge transfer under three types of distribution shift: parameter shift, covariate shift, and residual shift. We propose a novel transferable set and a new transfer framework to address the above three discrepancies. Non-asymptotic estimation error bounds and source detection consistency are established to validate the availability and superiority of our method in the presence of distribution shift. Additionally, an orthogonal debiased approach is proposed for statistical inference with knowledge transfer, leading to sharper asymptotic results. Extensive simulation results as well as real data applications further demonstrate the effectiveness of our proposed procedure.

* Joint first author; † corresponding author.