arXiv subjects

Tianxi Cai

Publications and source records attributed to Tianxi Cai.

At least 19 recordsLinked to original sources

arXiv2026

HALO: Hierarchical Evidence Selection and Reinforcement Learning for Allocating Budgeted Label Optimization from Longitudinal Clinical Records

Learning from extensive longitudinal clinical records--spanning repeated encounters, evolving conditions, and heterogeneous documentation--under a finite labeling budget requires solving two coupled problems: deciding what evidence to keep and deciding which patients deserve higher-fidelity labels. Prior work on active learning and multi-fidelity label acquisition mainly addresses the second problem, while treating the evidence presented to the acquisition policy as fixed. We present HALO, a Hierarchical Evidence Selection and Reinforcement Learning pipeline for Allocating Budgeted Label Optimization from Longitudinal Clinical Records. This framework jointly learns what to read and what label source to acquire under budget. HALO first selects informative evidence from longitudinal records and then uses a budget-aware policy to allocate free-, silver-, and gold-fidelity label sources. For frozen policies, our theory characterizes asymmetric acquired-label channels and gives a budget-constrained finite-sample gold-risk bound for a common downstream logistic model. In a heart-failure case study, HALO attained the highest MedAlign acquired-label F1 at six of eight target budgets. A downstream classifier trained on HALO-acquired labels attained the highest external MGB F1 at budgets 15-60 and remained within 0.01 of the best comparator at budgets 75-120.

cs.CL↗

arXiv2026

Learning Sequential Decisions from Multiple Sources via Group-Robust Markov Decision Processes

Learning a common policy from multiple source environments requires accounting for cross-source heterogeneity, and the offline setting introduces an additional challenge of limited data coverage. We formulate this problem as a distributionally robust Markov decision process under a shared linear MDP representation with source-specific transition and reward components. The uncertainty set is constructed through feature-wise mixtures of these components, yielding a $d$-rectangular structure that captures cross-source heterogeneity and admits a robust Bellman recursion. Building on this formulation, we develop Group-Robust Pessimistic Value Iteration (GR-PEVI), an offline policy-learning algorithm that combines feature-wise worst-case aggregation with a data-dependent pessimism penalty. Under a robust partial coverage condition, we establish finite-sample suboptimality guarantees for GR-PEVI and extend the analysis to clustered sources, characterizing the trade-off between the statistical gain from pooling similar sources and the approximation error induced by within-cluster heterogeneity. Simulation studies demonstrate the effectiveness of GR-PEVI, with empirical behavior consistent with the theoretical analysis. We apply GR-PEVI to a multi-source electronic health record study of multiple sclerosis treatment decisions to illustrate its use in a clinical setting.

stat.ME↗

arXiv2026

Toward Global Large Language Models in Medicine

Despite continuous advances in medical technology, the global distribution of health care resources remains uneven. The development of large language models (LLMs) has transformed the landscape of medicine and holds promise for improving health care quality and expanding access to medical information globally. However, existing LLMs are primarily trained on high-resource languages, limiting their applicability in global medical scenarios. To address this gap, we constructed GlobMed, a large multilingual medical dataset containing more than 800,000 entries spanning 20 languages, including seven low-resource languages. Building on this, we established GlobMed-Bench, which systematically assesses 56 state-of-the-art proprietary and open-weight LLMs across 12 of the 20 languages represented in GlobMed, comprising over 20,000 independent experiments and more than 32 million model responses. This large-scale evaluation revealed significant performance disparities across languages, particularly for low-resource languages. Additionally, we introduced GlobMed-LLMs, a suite of multilingual medical LLMs trained across the 12 benchmark languages, with parameters ranging from 1.7B to 8B. GlobMed-LLMs achieved an average performance improvement of nearly 40% relative to baseline LLMs, with a maximum performance improvement of more than threefold in low-resource languages. These resources provide an important foundation for advancing the equitable development and application of LLMs globally, enabling broader language communities to benefit from technological advances.

cs.CL↗

arXiv2026

Prior-Informed Generalizable Low-Rank Representation Learning via Regularized Adversarial Reweighting

Modern learning problems increasingly leverage data from heterogeneous sources to improve robustness and generalizability under distribution shift, often with high-dimensional parameters admitting low-rank structure. Existing group distributionally robust optimization (group-DRO) approaches optimize against adverse source mixtures, but extending this principle to low-rank matrix estimation with a broad class of losses leads to a challenging nonconvex-concave problem for which methods and theory remain limited. We propose Prior-Informed generalizable Low-rank representation Learning via Adversarial Reweighting (PILLAR), which smooths the group-DRO objective by regularizing adversarial source weights toward a user-specified reference distribution while learning a shared low-rank representation. By tempering extreme adversarial reweighting, PILLAR improves estimation stability and reduces the potential conservativeness of group-DRO. Using Burer-Monteiro factorization, we optimize the nonconvex-concave objective by bifactored gradient descent with mirror ascent. We establish linear convergence to a statistical neighborhood of the population saddle point and derive an error bound that separates low-rank estimation from adversarial reweighting. We further show that PILLAR improves on vanilla group DRO when the latter's excess conservativeness dominates source heterogeneity and prior misspecification. Simulations and an electronic health record study of suicide-attempt risk demonstrate improved worst-source and transported predictive accuracy.

stat.ME↗

arXiv2026

Enhancing Spectral Embedding through Robust and Flexible Knowledge Transfer in Electronic Health Records

We propose a spectral-based, unsupervised representation learning framework to derive low-dimensional embeddings for clinical concepts and patients in rare disease cohorts from electronic health records, where data are high-dimensional but sample sizes are limited. To overcome this challenge, we incorporate a knowledge matrix extracted from a broader population that shares a partially overlapping subspace with the rare-disease cohort. Our method departs from existing approaches by relaxing restrictive one-to-one signal-alignment assumptions between the latent data matrix and knowledge matrix, allowing more flexible and realistic forms of structured sharing. We introduce a novel two-step spectral embedding procedure: first, we identify and remove irrelevant components from the knowledge matrix; then, we apply a projection-based method to separately recover shared and heterogeneous components. Simulations and an analysis of a real-world multiple sclerosis cohort show that the proposed method outperforms competing approaches, particularly in challenging scenarios where shared signals are weak and only partially aligned, as is common in rare-disease data.

stat.ML↗

arXiv2026

Guided Uncertainty-Aware Robust Domain Transfer

Artificial intelligence models deployed as clinical decision support tools often suffer substantial performance degradation over time and across institutions due to covariate shift, concept drift, and cross-system heterogeneity. Fully retraining complex models is frequently infeasible, particularly in EHR settings where labeled outcome data are scarce and regulatory constraints limit model modification. We propose GUARD (Guided and Uncertainty-Aware Robustness to Domain shift), a unified statistical framework for updating an existing deployed model using limited labeled target data and multiple heterogeneous source populations while guarding against future distributional shifts. GUARD formulates robust multi-source transfer learning as a semi-supervised, adversarial optimization problem anchored at the current model, and employs double machine learning with cross-fitting to obtain debiased, efficient estimates that leverage abundant unlabeled target covariates. It then constructs an uncertainty-aware, guided group DRO estimator that combines source and target information while accounting for sampling variability in the source models. Our framework jointly provides principled robustness to domain shift, efficient semi-supervised estimation, and valid statistical inference for multi-source transfer of clinical prediction models. We demonstrate the utility of GUARD via extensive simulation experiments and a real world application to predicting future disease activity in rheumatoid arthritis from EHR data spanning more than a decade. GUARD recalibrates a pretrained model with only a small number of labeled records per year and remains accurate over multi-year horizons where target-only and standard transfer estimators degrade sharply.

stat.ME↗

arXiv2026

Guided Adversarial Robust Transfer Learning with Source Mixing

Transfer learning is a critical technique that enables the application of knowledge gained from existing tasks or domains to improve performance on a new one, reducing the need for extensive data and training in each new context. Many existing transfer learning methods rely on leveraging information from source populations closely resembling the target population. However, this approach often overlooks valuable knowledge that may be present in different yet potentially related auxiliary samples. When dealing with a limited amount of target data and multiple source data, we introduce a novel approach, Guided Adversarial Robust Transfer (GART) learning, that breaks free from strict similarity constraints. GART is designed to optimize the most adversarial loss with respect to a collection of source mixture distributions that guarantee excellent prediction performances for the target data. We establish the closed form of the population GART and show that the GART estimator achieves a faster convergence rate than the model fitted with the target data. Our simulation studies suggest that GART outperforms existing transfer learning methods, attaining higher robustness and accuracy. We highlight GART's predictiveness and robustness by applying it to form genetic prediction models of high-density lipoprotein cholesterol using multi-institutional biobank-linked electronic health records data.

cs.LG↗

arXiv2026

Generalized Linear Markov Decision Process

Offline reinforcement learning for longitudinal studies often faces two linked challenges: rewards may be binary or bounded, and reward observations may be available only for a subset of trajectories or time points even when the corresponding state-action-next-state histories are available. Linear Markov decision process methods are tractable because Bellman backups remain linear, but they require linear rewards and do not indicate how transition-only observations should be used. We introduce GRASP-MDP, Generalized Reward And Semi-supervised Pessimism for Markov Decision Processes, a reward-transition separated framework addressing both issues through one Bellman decomposition. It preserves linear transition dynamics while modeling reward means through generalized linear models. Although this breaks the usual linear Bellman form, the backup remains explicit as a nonlinear reward component plus a linear continuation component, yielding a nonlinear-plus-linear Bellman-complete class. In the resulting recursion, observed rewards estimate the generalized reward component, while all available transitions estimate the continuation component. The resulting pessimistic value iteration controls the two estimation errors separately without reward imputation. Finite-sample guarantees show that transition-only observations reduce transition-estimation error while reward uncertainty remains governed by observed rewards. Simulations and a multiple sclerosis electronic health record application illustrate the empirical benefit of retaining transition-only observations.

stat.ML↗

arXiv2026

Efficient Modeling of Surrogates to Improve Multi-source High-dimensional Integrative Regression

Surrogate variables play an important role in various fields due to the scarcity or absence of gold standard labels. We develop a novel approach named SASH for Surrogate-Assisted and data-Shielding High-dimensional integrative regression. It is a semi-supervised approach that efficiently leverages sizable unlabeled samples with error-prone surrogate outcomes from multiple local sites, to improve the learning accuracy of the small gold-labeled data. To facilitate stable and efficient knowledge extraction from the surrogates, our method first obtains a preliminary supervised estimator, and then uses it to assist training a regularized single index model (SIM) for the surrogates. Interestingly, through a chain of convex and properly penalized sparse regressions that approximate the SIM loss with bias-correction, our method avoids the local minima issue of the SIM, and fully eliminates the impact of the preliminary estimator's excessive error. In addition, it protects individual-level information through the aggregation of summary statistics from local sites, leveraging a similar idea of bias-corrected approximation. Through simulation studies, we demonstrate that our method outperforms existing approaches. Finally, we apply our method to develop a genetic risk model for type II diabetes using large-scale data sets from UK and Mass General Brigham biobanks, where only a small fraction of subjects in one site are labeled through chart reviewing.

stat.ME↗

arXiv2026

Calibrated Estimation and Inference for Semiparametric Regression Models

We consider a broad class of semiparametric regression models in which the conditional distribution of the response takes the form $f\{Y|\boldsymbol{x}^T\boldsymbolβ+m(z),ϕ\}$, known up to a parametric component $\boldsymbolβ$ of diverging dimension $p$, a smooth function $m(\cdot)$, and a dispersion parameter $ϕ$. The existing literature on such models has focused on semiparametric efficiency for $\boldsymbolβ$, treating $ϕ$ and $m(\cdot)$ as nuisances and largely ignoring finite-sample bias. Yet this bias can be substantial, particularly when $p$ is large relative to $n$ or the dispersion is high, and it can seriously undermine inference for $\boldsymbolβ$; moreover, $ϕ$ is often of direct scientific interest. We therefore propose SABRE, a general calibration framework for semiparametric estimation and inference, which calibrates an initial estimator against its model-implied expectation under a tractable parametric approximation to the semiparametric model. For generalized partially linear models, we show that SABRE reduces the bias of both $\boldsymbolβ$ and $ϕ$, accommodates a diverging parameter dimension without sparsity, and preserves the first-order variance and semiparametric efficiency of the initial estimator; the joint construction also improves estimation and inference for $m(\cdot)$. Simulation studies and an application to Alzheimer's disease genetics association analysis demonstrate the empirical effectiveness of SABRE in reducing bias and improving inference.

stat.ME↗

arXiv2026

Multimodal domain adaptation under label shift and blockwise missing modalities

Multimodal domain adaptation uses labeled source datasets to predict outcomes in an unlabeled target population. Here, different sources may observe different subsets of modalities and have distribution shifts from the target. In this scenario, the cross-modal patterns used for alignment depend on the outcome distribution, so aligning unadjusted source data can misrepresent the target. Nevertheless, the distribution shift of the outcome cannot be measured directly when labels are unavailable in the target sample and the missing modality blocks prevent simple pooling. We propose a reference-anchored domain adaptation method. A reference modality observed in every source and the target is used to estimate the target outcome distribution and reweight source observations before alignment. The auxiliary modalities are then mapped to a common, target-defined representation obtained by canonical correlation analysis (CCA) in the target and reproduced in each source by ridge-regression maps. Outcome information is transferred through density-ratio models on the aligned representation. When gold-standard source labels are sparse, a surrogate-label-assisted approach is developed to enable robust domain adaptation. We establish a unified-rotation match-up result for the target CCA and source ridge maps and consistency of the target conditional outcome distribution. Simulations and a renal cell carcinoma (RCC) application show improved calibration and stable prediction under distributional shift and blockwise missing modalities.

stat.ME↗

arXiv2026

Unveiling Invariant and Transferable Latent Factors Across Heterogeneous Environments via ATLAS

This paper considers a multi-environment factor model in which high-dimensional covariates are collected from heterogeneous environments, with auxiliary labels available in a subset of these environments. The joint distribution of the covariates may vary across environments, whereas the latent structure is decomposed into invariant factors with shared loadings and heterogeneous factors with environment-specific loadings. Such a model is motivated by transfer learning and latent factor regression, where one seeks stable low-dimensional representations for both interpretation and robust out-of-sample prediction of the response $Y$. Leveraging the invariance principle, we show that the invariant and heterogeneous factors are disentangled under a minimal structural condition. Based on this, we propose ATLAS, an Auxiliary-label and invariance-guided Transfer via Latent Alignment across heterogeneous environmentS. ATLAS is a unified procedure that leverages the invariance principle to separate aligned invariant and unaligned heterogeneous factors, and further exploits supervision from auxiliary labels to extract prediction-invariant and transferable factors from those unaligned heterogeneous factors. ATLAS yields near-oracle performance for downstream latent factor regression, enables transferable prediction in new environments through the full latent signal when auxiliary labels are available, and reduces to robust invariant-factor-only prediction otherwise. We establish sharp non-asymptotic error bounds for recovering invariant and heterogeneous factors, identifying all the response-invariant factors, and estimating the invariant signal in $Y$.

math.ST↗

arXiv2026

Domain Adaptation Targeting Heterogeneous and Imbalanced Subgroups

Domain adaptation enables generalizable and efficient data-driven research. However, existing work has largely focused on domain adaptation for some intrinsically homogeneous target cohort, overlooking inherent heterogeneity within the target, which can exacerbate biases and unfairness in the presence of subgroups with imbalanced sample sizes. We develop a novel domain adaptation framework that addresses more complicated target data that consists of heterogeneous and data-sparse subgroups and lacks gold-standard label observations. Our method simultaneously handles high-dimensionality, covariate shift, and outcome model heterogeneity by combining a model-assisted debiasing step used for covariate shift correction with an adaptive knowledge-guided sparsification procedure used to mitigate the issue of sample disparity. We also introduce a new model selection strategy to avoid negative knowledge transfer in the absence of labels in the target data. Our method is theoretically justified for being robust to nuisance model misspecification and adaptive to heterogeneity between the subgroups. Numerical experiments and two real-world applications, including genetic risk modeling of type II diabetes and prediction of mutation-induced protein stability changes, demonstrate the practical advantages of our method.

stat.ME↗

arXiv2026

Contrastive Learning on Multimodal Analysis of Electronic Health Records

Electronic health record (EHR) systems capture a wealth of multimodal clinical data, encompassing both structured clinical codes and unstructured clinical notes. Yet, many EHR-focused studies have traditionally examined these modalities in isolation or combined them using simplistic methods, overlooking the intrinsic synergy between them. In reality, these modalities are deeply interconnected, each containing clinically relevant and complementary information that, when integrated effectively, can provide a more comprehensive understanding of patient health. Despite the success of multimodal contrastive learning in vision-language applications, its potential remains under-explored in multimodal EHR, particularly in terms of theoretical understanding. To support statistical analysis of multimodal EHR data, we propose a multimodal feature embedding generative model and design a multimodal contrastive loss to learn EHR feature representations. Our theoretical analysis demonstrates the effectiveness of multimodal learning over single-modality learning and connects the solution of the loss function to the singular value decomposition of a pointwise mutual information matrix. This connection leads to a privacy-preserving algorithm tailored for multimodal EHR representation learning. Simulation studies show that the proposed algorithm performs well under a variety of configurations. We further validate its clinical utility using real-world EHR data.

stat.ML↗

arXiv2026

Spherical Mixture Integration for Latent Embedding Alignment across Multi-Source Feature Spaces

Multi-institutional electronic health record (Multi-EHR) data have emerged as a powerful resource for developing predictive models to support clinical decisions and for generating reliable real-world evidence. By aggregating information from diverse patient populations and institutions, they enhance the robustness and generalizability of models and findings. However, analyzing multi-EHR remains challenging because disparate institutions rarely map all data elements to common ontologies, and raw EHR codes are often overly granular and institution-specific, fragmenting representations of the same clinical concept. Hence, integrative analysis must overcome two key hurdles: harmonizing codes with the same clinical meaning (synonymy), and aligning institutional feature spaces. To address these challenges, we propose SMILE, a Spherical Mixture Integration for Latent Embedding alignment across multi-source feature spaces, where embeddings from heterogeneous sources serve as privacy-preserving summaries of clinical concepts and sparse relational pairs provide weak supervision. Synonymy is modeled via a mixture of von Mises-Fisher distributions, yielding unified representations of semantically equivalent raw codes. We develop a composite quasi-likelihood estimator with non-asymptotic error bounds for the latent representations and mixture mean directions and consistent synonym-cluster recovery, quantifying the gains from integrating multiple sources and knowledge-graph information. Simulations and a multi-institutional EHR application demonstrate improved alignment and synonym clustering.

stat.ME↗

arXiv2026

Optimally taming biases in black-box models for efficient semiparametric estimation

Modern semiparametric estimation often relies on flexible black-box machine learning methods to estimate nuisance functions, raising a fundamental question: how do nuisance estimation errors propagate into inference for low-dimensional target parameters? The dominant paradigm, exemplified by double machine learning (DML), yields error bounds in which nuisance estimation errors enter multiplicatively. While widely adopted, it remains unclear whether this multiplicative-rate dependence is optimal for black-box models. In this paper, we start by revisiting the partial linear model $Y = μ_0(X)+T\cdotβ_0+\varepsilon$ under a structure-agnostic setting, where the nuisance function $μ_0$ is estimated using a generic machine learning model, with approximation error $δ^a_μ$ and stochastic error $δ_μ^s$. We show that the standard DML rate is not optimal in the regime where the auxiliary function $\mathbb{E}[T|X=x]$ cannot be consistently estimated. We propose a new estimator for $β_0$ that achieves a sharper rate of $n^{-1/2}+δ^a_μ+(δ_μ^s)^2$ and establish a matching lower bound demonstrating its optimality. Our results reveal a new principle: the first-order stochastic error of nuisance estimation can be eliminated without imposing any additional assumptions. This also leads to a revised tuning strategy favoring under-smoothing, where $δ^a_μ\asymp(δ_μ^s)^2$, rather than the classical bias-variance trade-off $δ^a_μ\asymp δ_μ^s$. Under mild additional conditions, the estimator is asymptotically normal with minimal asymptotic variance. The proposed method extends to a broad class of semi-parametric linear functional estimation problems, including average treatment effect estimation. Our results imply that popular orthogonal score methods in semiparametric estimation with black-box nuisance learners can be substantially improved.

math.ST↗

arXiv2026

Cost-optimal Sequential Testing via Doubly Robust Q-learning

Clinical decision-making often involves selecting tests that are costly, invasive, or time-consuming, motivating individualized, sequential strategies for what to measure and when to stop ascertaining. We study the problem of learning cost-optimal sequential decision policies from retrospective data, where test availability depends on prior results, inducing informative missingness. Under a sequential missing-at-random mechanism, we develop a doubly robust Q-learning framework for estimating optimal policies. The method introduces path-specific inverse probability weights that account for heterogeneous test trajectories and satisfy a normalization property conditional on the observed history. By combining these weights with auxiliary contrast models, we construct orthogonal pseudo-outcomes that enable unbiased policy learning when either the acquisition model or the contrast model is correctly specified. We establish oracle inequalities for the stage-wise contrast estimators, along with convergence rates, regret bounds, and misclassification rates for the learned policy. Simulations demonstrate improved cost-adjusted performance over weighted and complete-case baselines, and an application to a prostate cancer cohort study illustrates how the method reduces testing cost without compromising predictive accuracy.

stat.ML↗

arXiv2026

Representation learning to advance multi-institutional studies with electronic health record data from US and France

The widespread adoption of electronic health records has created new opportunities for translational clinical research, yet this promise remains constrained by fragmented data across privacy-siloed institutions and substantial heterogeneity in local coding practices. While privacy-preserving collaborative learning allows institutions to work together without sharing patient-level data, it does not address inconsistencies in how clinical concepts are represented across sites. We introduce a graph-based framework that addresses this gap by treating data harmonization as a scalable representation learning problem. Rather than relying on fixed standards or manual mappings, the framework integrates institution-specific summary statistics from health records, curated biomedical knowledge graphs, and semantic information derived from large language models to learn a shared semantic space. This joint learning approach aligns diverse, site-specific vocabularies while preserving patient privacy. Evaluated across seven institutions and two languages, the framework provides a robust, data-centric foundation for training and deploying clinical models across heterogeneous healthcare systems.

cs.AI↗