Distributionally Robust Learning for Multi-source Unsupervised Domain Adaptation
Extracted source text · arXiv 2309.02211
\documentclass[10zpt]{article}
\usepackage{magging}
\title{Distributionally Robust Machine Learning with Multi-source Data \footnote{The research of Z. Wang and Z. Guo was partly supported by the NSF grant DMS 2015373 and NIH grants R01GM140463 and R01LM013614.
P. B\"{u}hlmann received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation program (grant agreement No. 786461)
}} %via \\
%Interpretable Maximin Aggregation}
\author[1]{Zhenyu Wang}
\author[2]{Peter B\"{u}hlmann}
\author[1]{Zijian Guo}
\affil[1]{Department of Statistics, Rutgers University, USA}
\affil[2]{Seminar for Statistics, ETH Z\"urich, Switzerland}
% \author{Zhenyu Wang,\;\; Peter B\"{u}hlmann,\;\; Zijian Guo }
\date{}
\pgfplotsset{compat=1.17}
\begin{document}
\maketitle
\vspace{-10mm}
\begin{abstract}
Classical machine learning methods may lead to poor prediction performance when the target distribution differs from the source populations. This paper utilizes data from multiple sources and introduces a group distributionally robust prediction model defined to optimize an adversarial reward about explained variance with respect to a class of target distributions. Compared to classical empirical risk minimization, the proposed robust prediction model improves the prediction accuracy for target populations with distribution shifts. We show that our group distributionally robust prediction model is a weighted average of the source populations' conditional outcome models. We leverage this key identification result to robustify arbitrary machine learning algorithms, including, for example, random forests and neural networks. We devise a novel bias-corrected estimator to estimate the optimal aggregation weight for general machine-learning algorithms and demonstrate its improvement in the convergence rate. Our proposal can be seen as a distributionally robust federated learning approach that is computationally efficient and easy to implement using arbitrary machine learning base algorithms, satisfies some privacy constraints, and has a nice interpretation of different sources' importance for predicting a given target covariate distribution. We demonstrate the performance of our proposed group distributionally robust method on simulated and real data with random forests and neural networks as base-learning algorithms.
\end{abstract}
%\Peter{COMMENT: Is it for arbitrary ML algorithms for continuous $Y$? I am not sure...} \Zijian{Our algorihtm is for any machine learning. The outcome variable is not necessarily continuous. Our algorithm can also be applied to binary outcome. The only constraint is the loss function, instead of the outcome type.}
{\bf Key words}: federated learning; interpretable machine learning; distributionally robust random forests; distributionally robust deep neural network; minimax optimization
% \Zijian{The research of Z. Wang and Z. Guo were partly supported by the NSF grant DMS 2015373 and NIH grants R01GM140463 and R01LM013614.}
% \Peter{Grant information etc. is missing? In my case:\\PB received funding from the European
% Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme
% (grant agreement No. 786461).\\Also, we need to add author affiliations. In my case:\\
% Seminar for Statistics, ETH Z\"urich, Switzerland. }
%This closed-form identification enhances the interpretability of the distributionally robust algorithms and enables the construction of computationally efficient and privacy-preserving algorithms.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%% 1.0 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Introduction}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
A fundamental assumption for the success of machine learning algorithms is that the training and test data sets share the same generating distribution. However, in many applications, the underlying distribution of the test data may exhibit a shift from that of the training data, which might be due to changing environments or data being collected at different times and locations \citep{quinonero2008dataset, koh2021wilds, malinin2021shifts, nado2021uncertainty}. Such distributional shifts may lead to most machine learning algorithms having a poor or unstable prediction performance for the test data even when the algorithms are fine-tuned on the training data. It is critical yet challenging to construct a generalizable machine learning algorithm that guarantees excellent prediction performance even in the presence of distributional shifts.
Distributionally robust optimization (DRO) has proven effective in improving the prediction model's generalizability for unseen target populations; see \citet{rahimian2019distributionally, namkoong2017variance, sinha2017certifying,ben2013robust,bertsimas2018data,blanchet2019data,blanchet2019robust, gao2023distributionally,gao2022wasserstein,kuhn2019wasserstein} for examples. The main idea of DRO is to incorporate the distributional uncertainty of the test data into the optimization process, with the crucial operational step of minimizing the adversarial loss defined over a class of target distributions centered around the training data's empirical distribution.
Since DRO can be overly pessimistic in practice, the group DRO was proposed as a more accurate prediction model when analyzing the data collected from multiple sources \citep{sagawa2019distributionally, hu2018does}. In various real-world applications such as face recognition \citep{grother2011report}, language identification \citep{blodgett2016demographic}, and electronic health record (EHR) data analysis \citep{singh2022generalizability, rasmy2018study}, the training data is often collected from multiple sources, with group labels indicating sources of the data. The group DRO was defined to optimize the worst-group loss over a collection of target distributions generated from the multi-source populations \citep{sagawa2019distributionally,hu2018does}. Despite its superior empirical performance, the construction of more flexible distributionally robust prediction models with multi-source data and their related interpretation and statistical theory is largely lacking.
%To model this, we focus on the setting of unsupervised domain adaptation; that is, the target population only has the covariates observation but not the outcome . %\Zijian{Zhenyu, add references}. \Zhenyu{Done}
%%%%%%%%%%%%%%%%%%%%% 1.1 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Our results and contribution}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%Without expert knowledge about the target conditional outcome distribution, we assume that it can be seen as a combination of the source ones, or potentially close be close to such mixtures. This assumption is frequently made in modern applications. For instance, in speech recognition, \citet{liao2013speaker} developed an accurate acoustic model for a broader population, which can be viewed as a mixture of source groups. Similarly, in object recognition, \citet{torralba2011unbiased} considered that the target application may contain images resembling a subset of multiple image source databases.
In this paper, we consider the construction of distributionally robust prediction models with multi-source data and focus on the unsupervised domain adaption regime, where the outcome labels are observed across multiple sources but missing for the target domain. The unsupervised domain adaption is commonly encountered in different applications \citep{chen2019self,jiao2022learning, ouali2020overview} {, as it is typically costly to acquire labeled data}. %van2020survey}. %For example, in EHR analysis, there are often limited or even no outcome observations for the target population \citep{humbert2022strategies}. %\Zijian{We focus on the unsupervised domain adaption. }
We consider $L$ independent groups of source data $\left\{(\X{l}, \Y{l})\right\}_{1\leq l\leq L}$: the data $\{(\X{l}_i, \Y{l}_i)\}_{1\leq i\leq n_l}$ from the $l$-th group is assumed to be i.i.d. generated following the source population,
\begin{equation}
% \X{l}_i \stackrel{i.i.d.}{\sim} \PP{l}_X,
\X{l}_i \distras{iid} \PP{l}_X \quad \text{and} \quad
% \stackrel{iid}{\sim} \buildrel iid\over\sim
% \substack{iid\\\sim}
\quad
\Y{l}_i = \f{l}(\X{l}_i) + \eps{l}_i, \quad \text{for}\quad 1\leq i\leq n_l,
\label{eq: source-models}
\end{equation}
with $\f{l}(\X{l}_i) = \E[\Y{l}_i | \X{l}_i]$ and $\X{l}_i\in \R^p$ and $\Y{l}_i\in \R$ respectively representing the covariates and outcome. %\Peter{Perhaps say somethig about $\eps{l}_i$? They have mean 0. Are the independent of $X_i^{(l)}$ or just $E[\eps{l} | X^{(\ell)}] = 0$?}\Zijian{We simply require $E[\eps{l} | X^{(\ell)}] = 0$. so this is implied by $\f{l}(\X{l}_i) = \E[\Y{l}_i | \X{l}_i]$. should we still say some additional thing about $\eps{l}_i$?} \Peter{OK. Then it is fine.}
We allow the covariate distribution $\PP{l}_X$ and the conditional outcome model $\f{l}$ to vary with the group label $l$. %\Zhenyu{Zijian, Do we need to mention $\eps{l}$ difference here? or just say both covariates' and conditional outcome's distributions}
We use $\QQ=(\QQ_{X},\QQ_{Y|X})$ to denote the target population and {assume}
%consider
that we only have access to covariate observations $\{\XQ_j\}_{1\leq j\leq n_\QQ}\distras{iid} \QQ_{X}$ but no outcome observations. We allow $\QQ$ to differ from any of the source populations' distributions in \eqref{eq: source-models}.
Our algorithm
%target model
is a distributionally robust prediction model that guarantees excellent prediction performance for target populations differing from source populations. To achieve this, we define an uncertainty set as a collection of target populations generated as a mixture of source populations in \eqref{eq: source-models}. We then define the distributionally robust prediction model to optimize the worst-case reward (about the explained variance) evaluated over this uncertainty set. In practice, domain experts might have prior knowledge about the target population {or one wants to robustify within a certain region like a ball of a certain radius around the observed mixture of source distributions}. Our proposed framework can incorporate such prior information {or additional constraints} into constructing the uncertainty set. {It}
%The prior information
plays a critical role in balancing the robustness and predictiveness of the final robust prediction model.
As our main contribution, we establish in Theorem \ref{thm: identification} that our proposed distributionally robust prediction model is identified as a weighted average of the conditional outcome models $\{\f{l}\}_{1\leq l\leq L}$ in \eqref{eq: source-models}. Our obtained closed-form expression is instrumental in robustifying arbitrary machine learning algorithms, and in fact, it does not require re-training of machine learning algorithms to robustify them. This is a great computational advantage, together with an outstanding simplicity for implementation which does not need new computational tools or software. In more detail, for $1\leq l\leq L$, we leverage the $l$-th source data to construct a machine learning prediction model $\widehat{f}^{(l)}(\cdot)$, e.g., by random forests, boosting, and deep neural networks. We compute the data-dependent optimal aggregation weights based on $\{\widehat{f}^{(l)}(\cdot)\}_{1\leq l\leq L}$ and then aggregate the multi-source prediction models $\{\widehat{f}^{(l)}(\cdot)\}_{1\leq l\leq L}.$ {We simply need the machine learning outputs from the different sources $\{\widehat{f}^{(l)}(\cdot)\}_{1\leq l\leq L}.$ } In addition and importantly, we devise a novel bias-correction step to {better} estimate the {optimal} aggregation weights.
%instead of simply plugging in the machine-learning predicted values.
We refer to our proposed method as distributionally robust learning (DRL), including distributionally robust random forests and neural networks as important special cases. We establish the convergence rate of our proposed DRL algorithm and demonstrate that the novel bias correction step improves the convergence rate for certain regimes.
Our proposed DRL can be seen as a distributionally robust federated learning approach that is computationally efficient and privacy-preserving {as it is based only on the machine learning fits from the different sources $\{\widehat{f}^{(l)}(\cdot)\}_{1\leq l\leq L}$.} {Furthermore, it}
%Our proposed DRL
allows for
%the
distribution heterogeneity across source populations and tends to capture the shared association across different sources; see Figure \ref{fig: illus mag} for illustration. More importantly, our proposed DRL has a nice model interpretability: the {estimated}
optimal aggregation weights indicate how much each source data contributes to the constructed robust prediction model.
%interpretation of the importance of different sources for predicting {at a given target covariate distribution}.
%An essential feature of our proposal is its
We demonstrate with simulated and real data that our proposed DRL algorithm leads to a significantly better prediction performance than ERM for target populations with distributional shifts, and {our} DRL {which incorporates}
%incorporating
%the
prior information {or additional constraints} balances the robustness and predictiveness of the constructed model.
To summarize, the main contributions of this paper are as follows:
\begin{itemize}
\item We leverage Theorem \ref{thm: identification} and introduce a theoretically justified framework for constructing distributionally robust machine learning algorithms {in great generality, yet enjoying efficiency and simplicity in terms of computation and implementation}. In addition, our proposed DRL algorithm provides an important tool for robust federated learning {and exhibiting privacy-preserving aspects}. %Our proposal can be seen as a distributionally robust federated learning approach that is computationally efficient and has a nice interpretation of the importance of different sources for predicting {at a given target covariate distribution}. %The construction relies on our obtained closed-form identification result in Theorem \ref{thm: identification}, which enhances the interpretability of the constructed model and enables computationally efficient implementation of different machine learning algorithms.
\item We propose a novel bias-correction method in leveraging the machine learning prediction models to construct the optimal aggregation weights.
%This is one prominent example for our advocation: e
Even for machine learning algorithms achieving excellent prediction performance on a single source of data, {such an additional bias correction method leads to improvements for robust integration of machine learning in a multi-source context.}
%procedures may be needed to integrate machine learning for solving statistical tasks.
\end{itemize}
%\Zijian{Why do not we directly say that w introduce the debiaed estimator to improve the rate? There is no need that we criticize something and then propose sth.} Moreover, in Section \ref{sec: Theoretical}, we quantify the approximation errors for DRL estimators. We emphasize that the plugin estimator may suffer from slow convergence rate, particularly when the source conditional outcome models $\left\{\f{l}\right\}_{l\in[L]}$ have large scales \Zijian{I do not think anyone can understand the previous sentence}. Therefore, in Section \ref{sec: methodology}, we propose the bias correction technique integrated with machine learning models to significantly improve the convergence rate of DRL estimators. We also provide a series of experimental results on several real-world data applications, demonstrating the superior performance of our proposed DRL framework.
%This property empowers Group DRO with the ability to interpret the final robust model. By examining the computed aggregation weights, we can assess the importance of each source population.
%a prediction model defined to optimize the adversarial reward for a class of target populations.
%We define our targeted distributionally robust prediction model as the optimizer of the adversarial reward with respect to a collection of target populations.
%we define the uncertainty set as a collection of target populations generated as the mixture of multiple source populations.
%This paper focuses on the scenario involving multiple source populations and aims to derive a robust model that can adapt to any new mixture target population. \Zijian{I do not think that our framework is novel...We should emphasize on what is really novel.}We introduce a novel framework called Distributional Robust Learning (DRL) in Section \ref{sec: formulation}, building upon the formulations of Group DRO \citep{sagawa2019distributionally} and Maximin Effects \citep{meinshausen2015maximin}.
%%%%%%%%%%%%%%%%%%%%% 1.2 %%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Comparison to related works}
This section further discusses the connections and differences between our proposal and existing literature. We start with the group DRO and agnostic federated learning, which are most relevant to our framework.\\ %Then, we will discuss other related multi-source literature, including maximin aggregation, meta learning, other robust federated learning, and invariant causal learning.\\
\vspace{-4mm}
\noindent {\bf Group DRO and Agnostic Federated Learning. } %DRO proves effective in constructing generalizable prediction models when a distribution shift exists between source and target data. DRO was proposed to minimize the adversarial loss defined with respect to a neighborhood near the training data's empirical distribution
%\citep[e.g.]{rahimian2019distributionally, gao2023distributionally,gao2022wasserstein, namkoong2017variance, sinha2017certifying, namkoong2016stochastic, oren2019distributionally}.
When the group structure is available, Group DRO proves effective in constructing a generalizable prediction model, which is defined to minimize the loss over the worst-performing group during training \citep{sagawa2019distributionally, hu2018does, zhang2020coping, mohri2019agnostic}. As a particular type of group DRO, our proposed DRL has a few fundamental differences from the existing works on the group DRO \citep{sagawa2019distributionally, hu2018does, zhang2020coping} or agnostic federated learning \citep{mohri2019agnostic, deng2020distributionally}.
%Similarly, our proposed DRL also requires knowledge of the group structure in the source data. However, there are notable differences between our DRL framework and the group DRO works.
\begin{enumerate}
\vspace{-2mm}
\item {\bf Uncertainty set and loss/reward functions.} Compared to \cite{sagawa2019distributionally, hu2018does,mohri2019agnostic} {whose uncertainty set is specified with joint distribution on $(X,Y)$}, we introduce a different uncertainty set in defining the group DRO {with respect to the conditional outcome distribution on $Y|X$}, which is more useful in the popular unsupervised domain adaptation; see the detailed discussions in Section \ref{sec: discussion_loss}. Moreover, instead of considering a general loss or reward function, we focus on using the explained variance as the reward function. Interestingly, we demonstrate in Section \ref{sec: discussion_loss} that the standard squared loss function leads to a much less robust prediction model.
\vspace{-2mm}
\item {\bf Identification with closed-form solution.} Our identification strategy (via the closed-form) is distinct from existing works \citep{sagawa2019distributionally, hu2018does,mohri2019agnostic, deng2020distributionally}, which are mainly designed for gradient-based algorithms (e.g., neural networks). In contrast, our proposal can be used to robustify all machine learning algorithms, including random forests, boosting, and neural networks. Importantly, questions about the contribution of each source group to the final robust prediction model remain unanswered in group DRO or agnostic federated learning, while the aggregation weights in DRL provide a natural interpretation of the contribution of each source group to the final prediction model.
\vspace{-2mm}
\item {\bf Computation efficiency and privacy.} Our proposed DRL is computationally more efficient than the group DRO algorithms \citep{sagawa2019distributionally, hashimoto2018fairness}. Instead of co-training all source data sets simultaneously, our proposed DRL framework requires an accurate machine-learning prediction model computed based on each source data set and automatically provides a robust aggregation of these machine-learning prediction models.
Since we do not require the passing of individual data, our proposal is helpful in applications to federated learning with privacy constraints. {Furthermore, it has the advantage of simplicity: no specialized software is needed, and one can use any favored machine learning base algorithm.} %In Group DRO, the robust model is typically trained over the entire source data directly, which requires the transition of individual data among $L$ source groups. In contrast, the DRL framework allows for parallel training of models locally on each source group, followed by aggregation. As a result, DRL can be integrated into federated learning settings with privacy concerns since individual data does not need to be passed around.
\end{enumerate}
%\vspace{-2mm}
%In summary, while Group DRO and DRL both leverage group structures to improve robustness, our proposed DRL offers several advantages, including a closed-form solution, better interpretability, and potential privacy benefits in federated learning scenarios.\\
\noindent {\bf
%Comparison to
Maximin aggregation.} %we would like to mention that in the presence of group information, the maximin effects framework
\cite{meinshausen2015maximin, buhlmann2015magging, rothenhausler2016confidence, guo2020inference} studied the estimation and inference for the maximin effects in the linear models, a particular type of distributionally robust prediction model. In contrast, we focus on a general framework for constructing distributionally robust machine learning algorithms instead of restricting to the linear outcome models. Additionally, maximin effect is proposed to optimize the prediction performance over the worst group. In contrast, our DRL frameworks integrates the prior information about the target distribution and produces a less conservative prediction model. As an important intermediate step, we have devised a novel bias-correction idea for computing the optimal aggregation weights for generic machine-learning models.
\noindent {\bf
%Comparison to
Distributionally Robust Optimization (DRO).}
%To address the distributional shift between the source and target data,
A line of DRO works \citep{ben2013robust,namkoong2016stochastic,namkoong2017variance,duchi2021learning,sinha2017certifying,blanchet2019robust,blanchet2019quantifying,gao2022wasserstein,gao2023distributionally} introduced algorithms to minimize adversarial loss defined with respect to a predefined uncertainty set. These DRO methods are not designed for analyzing data with the group structure. As the main difference, we leverage the group structure to construct a group DRO and establish a closed-form identification.
\noindent {\bf
%Comparison to
Robust federated learning.}
%Our DRL can be seen as a federated learning (FL) approach, which allows for data heterogeneity across different source data sets and aims to train a model on a central server without sharing the individual-level data from all local sites to the central \citep{mcmahan2017communication}. As a growing research direction,
%In federated learning,
\citet{mohri2019agnostic,deng2020distributionally,reisizadeh2020robust,du2021fairness} leveraged the distributional robustness idea in the federated learning setting. This paper considers a different unsupervised domain adaptation regime and introduces a different distributionally robust prediction model to adapt to this new regime.
We have discussed the detailed differences from the most relevant works \citep{mohri2019agnostic,deng2020distributionally} and pointed out that our proposed DRL has a distinct identification strategy, which leads to privacy-preserving and more computationally efficient algorithms. As another direction, \citet{guo2023robust,maity2022meta} proposed to learn a prevailing model that is satisfied by the majority of multiple studies. In contrast, instead of assuming the existence of a prevailing model, our proposed DRL seeks to identify a robust prediction model satisfying the distributional robustness property. %with good prediction performance, given the training data from multiple studies.
% \Zhenyu{Add ``instead'', focus on different model.}
% Our DRL framework introduces a novel perspective to FL. Instead of assuming the existence of a prevailing model, we focus on the prediction robustness of our final model. In the unsupervised domain adaptation case, where the target covariates observations are accessible, the final robust model can be expressed as an aggregation of prediction models learned from each source client. This unique feature makes our DRL highly efficient by avoiding the back-and-forth communications while retaining the advantage of distributionally robustness.
% To improve communication, works like \citet{mcmahan2017communication,li2020federated,stich2018local,woodworth2020minibatch,deng2020distributionally} have propose and analyzed methods to alleviate the communication bottleneck in FL.
% Different robust meta and federated learning approaches have been proposed to learn a general federated learning model by accounting for the data heterogeneity. Firstly, some propose learning a model that performs uniformly well over a pre-specified uncertain class, where they adopt the idea of adversarial loss and train the model through min-max objectives \citep{reisizadeh2020robust,du2021fairness}. Among these works, the most relevant one to ours is the Agnostic federated learning \citep{mohri2019agnostic}, where both the known source group structure and the target distribution being a mixture of source distributions are considered. Building upon this work, another work \cite{deng2020distributionally} further decouples the updates of model parameters and mixing parameters to significantly reduce the communication frequency between the central server and local clients, thereby enhancing the learning process's efficiency.
% In contrast, our proposed DRL focuses on the unsupervised domain adaptation case, where we have access to the target covariates' observations and formulates it with a subtly different uncertainty set. By doing so, our DRL only needs to aggregate the localized fitted models at the final step, thereby alleviating the back and forth communication between the central server and local clients. This particular feature makes our DRL highly efficient while retaining the advantage of distributionally robustness.
% \noindent {\bf
% %Comparison to
% Invariant causal learning.}
% Another research area is to conduct causal learning through identifying the invariant relationship across multiple environments \citep{peters2016causal,rojas2018invariant, heinze2018invariant,fan2023environment,arjovsky2019invariant}. In contrast, our DRL framework is mainly designed for constructing robust prediction models for a target population with possible distributional shifts instead of exploring causal invariant relationships.
%assume that the conditional distribution of the outcome variable given a set of covariates remain invariant across multiple environments. Anchor regression \citep{rothenhausler2021anchor} relies on the presence of exogeneous anchor variables which generate heterogeneity then quantitatively relates the robustness and causality in the linear model settings.
% \Zhenyu{efficient aggregate, generalize, care about distributional robust model}
% \citet{peters2016causal,rojas2018invariant, heinze2018invariant,fan2023environment,arjovsky2019invariant}
% assume that the conditional distribution of the outcome variable given a subset of covariates remains invariant across multiple environments. While causal models aim to provide predictions that remain valid under arbitrarily intense interventions, which can lead to conservative predictions，
% \citep{rothenhausler2021anchor},
% our DRL framework is specifically designed for robust and accurate predictions in a distributionally-robust manner without explicitly exploring causal (invariant) relationships.
%\paragraph{Comparison to DRO.} To address the issue of target shifts, DRO methods have been proposed to incorporate uncertainty about the potential target distribution \citep{rahimian2019distributionally, gao2017distributional}. These methods minimize an adversarial loss within an uncertainty region. Various approaches \citep{namkoong2017variance, sinha2017certifying, namkoong2016stochastic, oren2019distributionally} have defined the uncertainty region in different ways. When group information is available, Group DRO emerged as a method to minimize the loss over the worst-performing group. However, these methods lack sufficient interpretation regarding the role played by each source group in the final robust model. For example, they cannot answer questions about how much the final model will change if one of the source groups is removed or a new source group is included. In our DRL formulation, the robust model takes the form of a weighted average of the source models. The computed weights naturally indicate the importance of each source in the final robust model. Moreover, while Group DRO may yield a useless robust model when naively minimizing over the worst group, our proposed DRO automatically filters out negative sources by assigning them a weight of zero, see Section \ref{sec: discussion_loss}. Additionally, our DRL framework can be integrated with any parametric or non-parametric machine learning models, including SVM, random forest, boosting, and neural networks. In contrast, existing Group DRO methods only support gradient-based neural networks due to their reliance on iterative gradient updates for model parameters and adversarial distribution.
%our framework accommodates various machine learning models, including random forests and boosting, not limited to neural networks like Group DRO, which relies on gradient-based algorithms. Users have the freedom to apply any machine learning model independently in each local source population and can train source models in parallel for efficiency.
%\paragraph{Comparison to Group DRO.}\Zijian{As an important generalization of the distributional robustness optimization,} Group DRO \citep{sagawa2019distributionally, hashimoto2018fairness} leverages group information to learn a robust model that minimizes the loss for the worst-performing group during training. Although these works show promising improvements in worst group performance compared to standard ERM, they face several fundamental limitations:
%\begin{enumerate}
% \item Data from multiple sources must be accessible simultaneously during training. However, in practice, industries like finance and healthcare have stringent data-privacy regulations that restrict data sharing across different sites \citep{topaloglu2021pursuit, long2020federated}. Recently, federated learning has emerged as a solution to train decentralized models while keeping data local on each client. However, the integration of Group DRO with federated learning remains less explored in the existing literature.
% \item There is a lack of sufficient interpretation regarding the final robust model. While existing Group DRO approaches aim to guarantee the worst group loss, they fail to answer specific questions such as the contribution of each source group to the final robust model and the generalization impact if one of these source groups is removed or a new source group is included. Prior studies \Zijian{should be citep?}\citet{schweikert2008empirical, rosenstein2005transfer} have highlighted the presence of negative transfer learning, which can significantly deteriorate generalization if models are naively fitted over data from all sources. Therefore, solely focusing on the worst group may lead to the learning of a seemingly robust model that is, in fact, not robust, as will be discussed in Section \ref{sec: discussion_loss}.
%\end{enumerate}
%\Zijian{Group DRO, maximin effects, causal learning, federated learning}
% \Zijian{I do not see the key point of the following paragraph.}
% Multi-source domain adaptation methods have been developed to train robust predictors for the target domain using labeled instances from multiple source domains, even when the target domain has limited or no labeled data \citep{sun2015survey,chattopadhyay2012multisource, sun2011two, duan2009domain, duan2012domain}. Many efforts have been put to identify the domain-invariant features shared in the source domains \citep{muandet2013domain, peng2019moment, zhao2020multi, wang2020learning}. Another approach is to utilize pre-learned predictors trained locally on the source domains and the target domain, if available. \citet{schweikert2008empirical, mansour2008domain, sun2013bayesian} proposed different schemes for combining these pre-learned predictors. \citet{mansour2008domain, mansour2012multiple, hoffman2018algorithms} investigated combining pre-learned predictors when the target domain is a mixture of source domains. However, these approaches overlook the uncertainty of the target distribution, and the resulting models lack robustness against a range of target shifts.
% % \citet{schweikert2008empirical} combines prelearned predictors through a simple weight scheme, \citet{mansour2008domain} establishes distribution weighted combining rule to aggregate, while \citet{sun2013bayesian} is based on the Bayesian learning principle.
% \Zijian{What is the key difference from federated learning?}
% Federated learning is another field closely related to our proposed DRL. It enables the training of models across multiple devices or nodes while keeping the data localized \citep{kairouz2021advances, zhang2021survey}. This approach has already found applications in healthcare \citep{xu2021federated}, finance \citep{long2020federated}, and other industries \citep{niknam2020federated} due to data privacy regulations in these sections. Robust federated learning, considering the heterogeneity of source distributions, has become a prominent research topic \citep{mohri2019agnostic, pillutla2022robust, marfoq2021federated, ghosh2019robust, deng2020distributionally}. Our proposed DRL only requires training the machine learning model locally for each source and then aggregating these fitted models. By avoiding the need to share raw data with a central location, the DRL algorithm alleviates data privacy concerns to a considerable extent.
%\Zijian{Zhenyu, add more literatures.}
%\Zhenyu{There are more literature I can add up: (1) prior information of $\HH$, 2-3 papers (2) causal discovery works, which aims to find the causal invariant features for prediction.}\Zijian{Need to add the causal discovery works. Need to mention the prior information $\HH$.}
%%%%%%%%%%%%%%%%%%%%% 1.3 %%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Notations} Define $n = \min_{l\in[L]} n_l$ and $n_\PPall = \sum_{l=1}^L n_l$. For a positive integer $m$, define $[m] = \{1,..., m\}$. For real numbers $a$ and $b$, define $a\vee b = \max\{a,b\}$ and $a\wedge b = \min\{a, b\}$. We use $c$ and $C$ to denote generic positive constants that may vary from place to place. For positive sequence $a(n)$ and $b(n)$, we use $a(n)\lesssim b(n)$ or $a(n) = O(b(n))$ to represent that there exists some universal constant $C > 0$ such that $a(n) \leq C\cdot b(n)$ for all
%the
$n\geq 1$, and denote $a(n) \asymp b(n)$ if $a(n)\lesssim b(n)$ and $b(n)\lesssim a(n)$. We use notations $a(n)\ll b(n)$ or $a(n)=o(b(n))$ if $\limsup_{n\rightarrow\infty} (a(n)/b(n)) = 0$. For a vector $x \in \R^p$, we define its $\ell_q$ norm as $\|x\|_q = (\sum_{i=1}^p |x_i|^q)^{1/q}$ for $q \geq 0$. For a numeric value $a$, we use notation $a_p$ to represent a length $p$ vector whose every entry is $a$. For a matrix $A$, we use $\|A\|_F, \|A\|_2$ and $\|A\|_\infty$ to denote its Frobenius norm, spectral norm, and element-wise maximum norm, respectively. We use $I_p$ to denote the $p$-dimensional identity matrix.
% On the other hand, there is growing literature on invariance causal discover to the multi-source robust learning. They assume the conditional outcome given some predictors or predictors' some representation keeps invariant across each source (referred as ``environment''). Predictions based on a causal model will in general work well under arbitrary perturbations (interventions) on the covariates. However, how many diversified environments are sufficient to exploit the invariance still remains unresolved. But in observational studies, we often have a limited number of environments. Moreover, exact invariance guarantees may be too conservative and yield to a suboptimal performance on observational and moderately shifted data.
% Data heterogeneity is a common challenge in many machine learning applications, where the training data is collected from multiple sources that have different underlying distributions. In many real-world scenarios, it is difficult or impossible to combine multiple sources data due to privacy or regulatory concerns, such as hospitals or commercial companies. Traditional machine learning methods suffer from dealing with such situation, which assume that the training data comes from a single, homogeneous distribution.
% To address this challenge, Distributionally Robust Optimization (DRO) has emerged a promising approach in recent years. DRO provides a way to incorporate distributional uncertainty into the optimization process, which enables the learning algorithms to be more robust to the heterogeneity in the data. Group DRO, a variant of DRO, has been designed to handle the multi-source challenge. It assumes that the group label for each source observation is known, and it seeks to find a model that is robust to the worst-case distribution over groups. However, all these DRO approaches requires learning the algorithms over the entire source data, i.e. each group data should be accessible simultaneously, which may not be possible or desirable in practice.
% It is essential to develop algorithms that can learn a robust model from multiple sources without relying on data sharing. This is particularly relevant in fields such as healthcare, finance, and telecommunications, where data privacy is critical. In this paper, we propose a novel formulation for DRO that is designed for settings where multiple groups of source data are available, and where data cannot be shared across groups. Our formulation leverages ideas from Group DRO but does not require access to all groups of data simultaneously.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%% 2.0 %%%%%%%%%%%%%%%%%%%%%%%%%
\section{Distributionally Robust Prediction Models: Definition and Identification }
\label{sec: formulation}
We start by introducing the multi-source setting considered in this paper. The source data consists of $L$ groups, denoted by $\{(\X{l}, \Y{l})\}_{l\in[L]}$, which may have heterogeneous distributions due to the data being collected from different sources. For the $l$-th group of data $\{(\X{l}_i, \Y{l}_i)\}_{i\in [n_l]}$, let $\PP{l}_X$ represent the distribution of the covariates $\X{l}_i\in\R^p$ and $\PP{l}_{Y|X}$ represent the conditional distribution of the outcome $\Y{l}_i$ given $\X{l}_i$. Equivalently, we write
\begin{equation}
\X{l}_i {\sim} \PP{l}_X
\quad \text{and}\quad
\Y{l}_i \mid \X{l}_i {\sim} \PP{l}_{Y|X}, \quad \text{for} \quad i\in [n_l].
\label{eq: source-pop}
\end{equation}
We model the data heterogeneity by allowing $\{\PP{l}_X\}_{l\in[L]}$ and $\{\PP{l}_{Y|X}\}_{l\in[L]}$ to vary with $l$. For $l\in[L]$, we use $\mathcal{X}_l\subseteq \R^p$ to denote the support of covariates for the $l$-th group
and define the conditional mean function $\f{l}: \mathcal{X}_l \mapsto \R$ as $\f{l}(\X{l}_i) = \E[\Y{l}_i\mid \X{l}_i]$. The data generation in \eqref{eq: source-pop} is the same as that in \eqref{eq: source-models}. %\Zijian{stop here.}
We aim to construct a generalizable prediction model for a target population $\QQ$, where $\QQ_X$ and $\QQ_{Y|X}$ respectively denote the target population's covariate and conditional outcome distributions. When the target population $\QQ=(\QQ_{X}, \QQ_{Y|X})$ has distributional shifts from the source populations $\PP{l}=(\PP{l}_{X}, \PP{l}_{Y|X})$ for $l\in [L]$, the standard ERM estimator based on the pooled source data might lead to a poor prediction model for the test data generated from the target population $\QQ$. We consider a general setting by allowing for the simultaneous existence of the following two types of distributional shifts: (i) covariate shift: $\QQ_X$ might differ from any of $\{\PP{l}_X\}_{l\in[L]}$; (ii) posterior drift: $\QQ_{Y|X}$ might differ from any of $\{\PP{l}_{Y|X}\}_{l\in[L]}$. Since it is typically hard to obtain outcome labels for the target population {such as in the case of electronic health record data analysis \citep{humbert2022strategies} and biological image analysis \citep{zhang2015deep}}, we focus on the unsupervised domain adaptation regime where we only have access to covariate observations $\{\XQ_j\}_{j\in [n_\QQ]}$ but no outcome observations $\{\YQ_j\}_{j\in [n_\QQ]}$.
%%%%%%%%%%%%%%%%%%%%%% 2.1 %%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Group Distributionally Robust Prediction Models}
The identification of $\QQ_{Y|X}$ becomes unfeasible within our framework since the target data is unlabelled, and the target conditional outcome distribution $\QQ_{Y|X}$ is allowed to be different from any of $\{\PP{l}_{Y|X}\}_{l\in[L]}$. Instead of trying to recover the true $\QQ_{Y|X}$, we propose in this subsection a distributionally robust prediction model that guarantees an excellent prediction performance for a wide range of target populations. It is worth noting that $\QQ_X$ can be identified in our setting since we can access covariate observations for the target population.
Since $\QQ_{Y|X}$ is not identifiable, we define in the following a class of target population distributions, %on which the target distribution $\QQ = (\QQ_X, \QQ_{Y|X})$ is assumed to lie:
\begin{equation}
\C(\QQ_X):= \left\{\left.\TT= (\QQ_X, \TT_{Y|X}) \; \right\vert \; \TT_{Y|X} = \sum_{l=1}^L q_l\cdot \PP{l}_{Y|X} \quad \text{with}\quad q\in \Delta^L
\right\}
\label{eq: region-Q full}
\end{equation}
where $\Delta^L= \{q \in \R^L \vert \sum_{l=1}^L q_l=1,\; \min_l q_l \geq 0\}$ denotes the $L$-dimensional simplex. We refer to $\C(\QQ_X)$ as the uncertainty set, representing a class of distributions from which the target population might be drawn. When the true target conditional outcome distribution $\QQ_{Y|X}$ is a mixture of $\{\PP{l}_{Y|X}\}_{l\in[L]}$, the uncertainty set $\C(\QQ_X)$ contains the unidentifiable target population $(\QQ_X, \QQ_{Y|X})$.
For any target distribution $\TT$ and a given prediction model $f$, we define the reward function $\RR_\TT(f)$ as
\begin{equation}
\RR_{\TT}(f) := \E_{(X,Y)\sim\TT}[Y^2 - (Y - f(X))^2]
\label{eq: reward}
\end{equation}
{where $\E_{(X,Y)\sim\TT}$ denotes the expectation taken with respect to the data $(X,Y)$ following the distribution $\TT$. The reward function $\RR_\TT(f)$ compares the prediction accuracy of the model $f(\cdot)$ to the null model, evaluated on the test data following the distribution $\TT$. For centered outcome variable, $\RR_\TT(f)$ represents
the outcome variance explained by model $f$.}
Intuitively, a higher value of $\RR_\TT(f)$ indicates a better prediction performance of the model $f$. %is able to capture a large portion of the variation in the outcome.
Due to the uncertainty of the target population, we evaluate the performance of a prediction model $f$ by the adversarial reward $\min_{\TT \in \C(\QQ_X)} \RR_\TT(f)$ with respect to the uncertainty set $\C(\QQ_X)$ defined in \eqref{eq: region-Q full}.
We define the robust prediction model $f^*$ as the optimizer of this adversarial reward,
\begin{equation}
f^* := \argmax_{f\in \FF}\min_{\TT \in \C(\QQ_X)} \RR_\TT(f)
\label{eq: f_star full}
\end{equation}
where $\FF$ denotes a pre-specified function class. We can define the negative reward function as a loss function, and then \eqref{eq: f_star full} can be formulated as an equivalent minimax optimization problem,
$$
f^* := \argmin_{f\in \FF} \max_{\TT \in \C(\QQ_X)}\E_\TT[\ell(X_i,Y_i; f)]
\quad \text{with}\quad \ell(x,y; f) = (y-f(x))^2 - y^2,$$ which is in a typical form of the group DRO \citep{sagawa2019distributionally, hu2018does}.
\begin{Remark}[Reward/loss function and binary outcome.]\rm
The choice of the reward function or the related loss function is critical in guaranteeing the robustness property of $f^*$. In the current paper, we mainly focus on the reward function introduced in \eqref{eq: reward}, which compares the prediction improvement of $f$ to a null model. In section \ref{sec: discussion_loss}, we discuss the advantage of using $\ell(x,y; f) = (y-f(x))^2 - y^2$ and demonstrate the issues of using the standard squared error loss $(y-f(x))^2$. Even though our proposed distributionally robust prediction model is applicable to the binary outcome regime, it does not contain the popular loss function (e.g., entropy or hinge loss) for binary outcome variables.
\end{Remark}
In the following, we generalize the definition in \eqref{eq: f_star full} by considering a broader class of uncertainty sets.
%and introduce a broader framework for constructing distributional robust models. %To achieve this, we shall consider a collection of shrinking uncertainty setes which are subsets of $\C(\QQ_X)$.
With $\C(\QQ_X)$ defined in \eqref{eq: region-Q full}, the distributionally robust prediction model $f^*$ defined in \eqref{eq: f_star full} optimizes the adversarial reward concerning the entire convex hull of $\{\PP{l}_{Y|X}\}_{l\in[L]}$. Despite its robustness, the model $f^*$ might lead to conservative prediction performance. For example, if the distributions of $\{\PP{l}_{Y|X}\}_{l\in[L]}$ are very different, $f^*$ might become a zero prediction model. Users might be interested in considering a smaller uncertainty set as a subset of $\Delta^{L}$ used in $\C(\QQ_X).$
In practice, there may be situations where we have some prior knowledge regarding $\QQ_{Y|X}$ \citep{hu2018does, mohri2019agnostic, zhang2020coping}. %We can use this prior information to define a narrower uncertainty set than $\C(\QQ_X)$.
As an agnostic default choice, we assume that the target population's conditional outcome model is generated with weights similar to the source data mixture $q^{\rm sou}=(n_1/n_\PPall,\cdots,n_L/n_\PPall)$ and consider the convex subset $\HH = \{q\in \Delta^L \; |\; \mathcal{D}(q,q^{\rm sou})\leq \rho\}$, where $\rho>0$ is a user-specified parameter controlling the size of $\HH$, and $\mathcal{D}(\cdot, \cdot)$ is a dissimilarity measure. For example, we can take $\mathcal{D}(\cdot, \cdot)$ as a Euclidean distance or f-divergence \citep{hu2018does} or Wasserstein-divergence \citep{blanchet2019robust}. %\Zhenyu{Zijian, please take a look at the preceding paragraph.}
With $\HH$ denoting a subset of the $L$-dimensional simplex $\Delta^L$, we define
\begin{equation}
\C(\QQ_X, \HH):= \left\{\left.\TT = (\QQ_X, \TT_{Y|X})\;\right\vert \; \TT_{Y|X} = \sum_{l=1}^L q_l \cdot \PP{l}_{Y|X} \quad \text{with}\quad q\in \HH \right\}.
\label{eq: region-Q H}
\end{equation}
Note that the uncertainty set $\C(\QQ_X, \HH)$ is a subset of $\C(\QQ_X)$. Consequently, with the uncertainty set $\C(\QQ_X, \HH)$, we define the corresponding group distributionally robust prediction model as %\Peter{You often write "model"; to me, a model is a bit different, but it's probably OK.} \Zijian{So is it better to say prediction model? I was also a bit unsure about this. I use "model" to imply that we are looking at the population parameter, but not the usual way of generating model. I am not sure whether we can use algorithm since algorithm is implemented on the finite sample, instead of the population parameter? After we figure out a way to handle this, I can make the change across the whole paper. } \Peter{OK, I agree. I would slightly prefer "prediction model".}
\begin{equation}
\fH := \argmax_{f\in \FF} \min_{\TT \in \C(\QQ_X, \HH)} \RR_\TT (f), \quad \text{with}\quad \HH\subseteq \Delta^{L}.
\label{eq: f_star H}
\end{equation}
Specifying the uncertainty set $\C(\QQ_X, \HH)$ involves a trade-off. On the one hand, if $\C(\QQ_X, \HH)$ contains the target population $\QQ$, a smaller uncertainty set would lead to a more accurate prediction model; on the other hand, a larger class is more likely to contain $\QQ$ and produce a more robust prediction model. The subset $\HH$ is determined by the prior knowledge about the target population. If we set $\HH = \Delta^L$, $f^*$ provides the prediction guarantees as long as the conditional outcome model for the target distribution is nearly a mixture of those for the source populations. However, since the model is guaranteed to optimize the worst-case reward over a large class of possible distributions, it may be too conservative for any given target population. To balance the trade-off between robustness and prediction accuracy, we recommend the users specify a reasonable region $\HH$ based on their domain knowledge (with $\HH=\Delta^L$ as the default choice). {We have explored how different specifications of $\HH$ will affect the model predictiveness and robustness in the numerical studies; see Section \ref{sec: simu-H} for details.}
% Define $\|f\|_{\QQ,2} := \left(\E_{\QQ_X}[f(X)^2]\right)^{1/2}$ as the second norm of $f(X)$ over target distribution $\QQ$.
%With $\HH$ converging to a singleton $\HH=\{q^0\}$, the uncertainty set converges to the single distribution with the corresponding weight $q^0.$
We consider the scenario that $\QQ_{Y|X} = \sum_{l=1}^L q^0_l \cdot \PP{l}_{Y|X}$ for some (unknown) $q^0 \in \Delta^L$ and the corresponding conditional outcome $f^\QQ(x) = \E_{\QQ}[Y|X=x]$ admits the form
$\sum_{l=1}^L q^0_l\cdot \f{l}(x).$
The following proposition establishes that a sufficiently accurate set $\HH$ can lead to the identification of $f^\QQ$.
\begin{Proposition} Suppose that $\QQ_{Y|X} = \sum_{l=1}^L q^0_l \cdot \PP{l}_{Y|X}$ for some $q^0 \in \Delta^L.$
If $\max_{q\in \HH}\|q - q^0\|_2 \leq \rho$, then $
\left(\E_{\QQ_X}[(\fH(X)-f^\QQ(X))^2]\right)^{1/2} \leq \rho \sqrt{L} \max_{l\in [L]}\|\f{l}\|_{\QQ,2} $
with $\fH$ defined in \eqref{eq: identification}. %Consequently, for any $\epsilon_0>0$, with $\rho \rightarrow 0$, $\prob_\QQ(|\fH(X) - f^\QQ(X)|>\epsilon_0) \rightarrow 0$.
% $\fH(X)$ converges to $f^\QQ(X)$ with probability 1 on the target population.
%\Zijian{Can you write a probability statement? The current statement is not in good shape.}\Zhenyu{Zijian: please double check it.}
% with $\rho \rightarrow 0$, $\fH \rightarrow f^\QQ$ \Zijian{Zhenyu, converge with probability 1 on the target population? please double check.}.
\label{prop: shrink H}
\end{Proposition}
The proposition above demonstrates that $\fH$ converges to the true conditional mean model $f^\QQ$ when the subset $\HH$ converges to the true mixing weight $q^0$. If we lack prior information about the potential mixture of the conditional outcome distribution, the most natural and secure approach is to fit a robust prediction model that guarantees all potential mixture weights by setting $\HH = \Delta^L$. However, with richer prior information, we can obtain a {better} predictive model or even recover the true conditional outcome model, but at the risk of reduced robustness if the prior information is incorrect.
A few important remarks on our definition of distributionally robust prediction models in \eqref{eq: f_star H} are in order.
\begin{Remark}[Fairness] \rm Fairness is an important consideration for designing the machine learning algorithm. Considering that the $L$ groups correspond to different subpopulations, we might protect some subpopulations by imposing certain restrictions on the construction of $\HH$, such as imposing lower bounds for the proportions of certain groups. With such constraint on $\HH$, our definition in \eqref{eq: f_star H} is closely related to the fairness notation proposed in \citet{hashimoto2018fairness,mohri2019agnostic,diana2021minimax,martinez2020minimax}.
% For example, \citet{hashimoto2018fairness} gave a lower bound on the smallest group proportion in the latent group structure setting (which defines $\HH$ implicitly), and demonstrated higher minority group performance and lower disparity amplification. \Zijian{What is our main difference?}\Zhenyu{Zijian, take another look.}%\citet{mohri2019agnostic} proposed the good-intent fairness for seeking the agnostic loss within a pre-specified uncertainty set $\HH$.
%\Zhenyu{Zijian, please take a look at this paragraph.}
\end{Remark}
\begin{Remark}[Connection to ERM]\rm Our proposed distributional robust prediction model $\fH$ bridges DRO and ERM by specifying the prior information set $\HH.$ With $n_\PPall=\sum_{l=1}^{L}n_l$, we recall the ERM is defined as
\begin{equation*}
\widehat{f}^{\rm ERM} \coloneqq \argmin_{f\in \FF} \sum_{l=1}^L \frac{n_l}{n_\PPall} \cdot \frac{1}{n_l}\sum_{i=1}^{n_l}\left[Y_i^{(l)} - f(X^{(l)}_i)\right]^2=\argmax_{f\in \FF} \sum_{l=1}^L \frac{n_l}{n_\PPall} \cdot \frac{1}{n_l}\sum_{i=1}^{n_l}\left\{\left[Y_i^{(l)}\right]^2-\left[Y_i^{(l)} - f(X^{(l)}_i)\right]^2\right\}.
%\label{eq: ERM}
\end{equation*}
Then the ERM estimator optimizes a sample version of the population risk $\sum_{l=1}^L q^{\rm sou}_l\cdot \E_{\PP{l}}[Y^2-\left(Y - f(X)\right)^2]$ with $q^{\rm sou}=(n_1/n_\PPall,\cdots,n_L/n_\PPall).$
For the no covariate shift setting, with $\HH = \{q^{\rm sou}\}$, $\fH$ defined in \eqref{eq: f_star H} is optimizing the same population risk as ERM.
% coincides with the RM:
% \begin{align*}
% \fH := \argmin_{f\in \FF} \E_{X\sim \QQ_X} \left[\sum_{l=1}^L q^{\rm sou}_l\cdot \E_{Y\sim \PP{l}_{Y|X}}\left(Y-f(X)\right)^2 \right]= \argmin_{f\in \FF} \sum_{l=1}^L q^{\rm sou}_l\cdot \E_{\PP{l}}\left(Y-f(X)\right)^2 = f^{\rm RM}
% \end{align*}
% For the sample version, we approximate them by the observed source data with the total sample size $n_\PPall$.
% \Zhenyu{Zijian, the population levels if RM and DRL coincides, where the empirical levels they shall differs a little bit $\tfH \neq \widehat{f}^{\rm ERM}$. Should we change the name ``Bridge DRL and ERM'' to ``Bridge DRL and RM''?}
\end{Remark}
% \begin{Remark}[Bridge DRO and ERM] \rm \Zijian{Zhenyu, please rewrite this remark. You need to define $n_p$ and ERM first, then discuss the connection to DRO and ERM.}
% During training, we observed the mixture of $L$ source groups $q^{\rm sou} \in \Delta^L$.
% For the no covariate shift setting with $\PP{l}_X = \QQ_X$ for $l\in [L]$, when $\HH = \{q^{\rm sou}\}$, then $\fH$ defined in \eqref{eq: f_star H} coincides with the ERM approach, which pools over all the source groups data.
% \begin{align*}
% \fH &= \argmax_{f\in \FF} \min_{\TT \in \C(\QQ_X, \HH)} \RR_\TT (f) = \argmin_{f\in \FF} \E_{X\sim \QQ_X} \left[\sum_{l=1}^L q^{\rm sou}_l\cdot \E_{Y\sim \PP{l}_{Y|X}}\left(Y-f(X)\right)^2 \right]\\
% &= \argmin_{f\in \FF} \sum_{l=1}^L q^{\rm sou}_l\cdot \E_{X\sim \QQ_X}\E_{Y\sim \PP{l}_{Y|X}}\left(Y-f(X)\right)^2 \\
% &= \argmin_{f\in \FF} \sum_{l=1}^L q^{\rm sou}_l\cdot \E_{\PP{l}}\left(Y-f(X)\right)^2
% \end{align*}
% The sample version is as follows:
% \[
% \sum_{l=1}^L q^{\rm sou}_l\cdot \E_{\PP{l}}\left(Y-f(X)\right)^2 \approx
% \sum_{l=1}^L \frac{n_l}{n_\PPall} \cdot \frac{1}{n_l}\sum_{i=1}^{n_l}\left[Y_i^{(l)} - f(X^{(l)}_i)\right]^2 = \frac{1}{n_\PPall} \sum_{i=1}^{n_\PPall}\left[Y_i - f(X_i)\right]^2
% \]
% where the last term is exactly the ERM's objective.
% \end{Remark}
%%%%%%%%%%%%%%%%%%%%%% 2.2 %%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Identification for Distributionally robust prediction models}
\label{subsec: identification}
In the following, we establish that the optimizer $\fH$ defined in \eqref{eq: f_star H} attains a closed-form solution as a weighted average of $\left\{\f{l}\right\}_{l\in[L]}$.
\begin{Theorem}
Suppose that the function class $\FF$ is convex with $\f{l} \in \FF$ for $l\in[L]$ and $\HH$ is a convex subset of $\Delta^L$, then $\fH$ defined in \eqref{eq: f_star H} is identified as:
\begin{equation}
\fH = \sum_{l=1}^L q_l^* \cdot \f{l} \quad \textrm{with} \quad
q^* = \argmin_{q\in \HH}q^\intercal \Gamma q
\label{eq: identification}
\end{equation}
% where $\Gamma = \E_{\QQ_X}\left[{\F}(X)\F(X)^\intercal\right]$ and $\F(\cdot)=\left(\f{1}(\cdot),\f{2}(\cdot),...,\f{L}(\cdot)\right)^\intercal.$
where $\Gamma_{k,l} = \E_{\QQ_X}[\f{k}(X)\f{l}(X)]$ for $k,l\in [L]$.
\label{thm: identification}
\end{Theorem}
The above identification result requires the optimization class $\FF$ to be sufficiently large, containing the convex hull of the conditional outcome models $\{\f{l}\}_{l\in [L]}.$ The above identification result might not hold for a small class $\FF$ which does not contain the convex hull of $\{\f{l}\}_{l\in [L]}.$
The above theorem provides a convenient way of computing $\fH$ defined in \eqref{eq: f_star H}.
The quadratic optimization problem in \eqref{eq: identification} yields the optimal aggregation weight $q^*$, which is further used to compute $\fH$ as a weighted average of $\{\f{l}\}_{l\in[L]}$. Crucially, the closed-form of $\fH$ in \eqref{eq: identification} provides the interpretation of the distributionally robust prediction model, where the optimal weight $q^*$ describes the contribution of $\{\f{l}\}_{l\in[L]}$ from different sources.
The objective $q^\intercal \Gamma q$ in \eqref{eq: identification} is equivalent to $\E_{X\sim \QQ_X}\left[(\sum_{l=1}^L q_l\cdot \f{l}(X))^2\right]$, which is a distance measure from the aggregated model $\sum_{l=1}^L q_l \cdot \f{l}$ to the origin. The definition of this distance measure depends on the target covariates distribution $\QQ_X$. Geometrically speaking, $f^*$ represents the point that is closest to the origin in the convex hull of $\{\f{l}\}_{l\in[L]}$ while $\fH$ denotes the point closest to the origin in the $\HH$-constrained set, as shown in the left subfigure of Figure \ref{fig: illus mag}.
%\begin{Remark}[Comparison to Federated and Meta Learning]
%\end{Remark}
\begin{figure}[htp]
\centering
\includegraphics[width=0.9\textwidth]{assets/illus_drl.png}
\caption{Illustration of $f^*$ (the blue point) and $\fH$ (the red point) for $p=2$, $L=6$, and the additive models $\f{l}(x)=\sum_{j=1}^{2}\f{l}_{j}(x_j)$ for ${l\in[L]}$. The left panel: $f^*$ is the point closest to the origin in the convex hull of $\left\{\f{l}\right\}_{l\in[L]}$, and $\fH$ is the point in the $\HH$-constrained set having the smallest distance to the origin; The right panel: consider the setting with shared first component $\f{1}_1 = \f{2}_1 = ... =\f{L}_1 = f_1$ and the second component being scattered around $0$; the distributional robust prediction model $f^*(x)=f_1(x_1)$ (blue point) retains only the shared component and shrinks the sign heterogeneous component to zero. }
\label{fig: illus mag}
\end{figure}
Theorem \ref{thm: identification} generalizes the result obtained in \citet{meinshausen2015maximin,guo2020inference}, mainly focusing on the linear models and $\HH=\Delta^{L}$; see Theorem 1 in \citet{meinshausen2015maximin} and Proposition 1 in \citet{guo2020inference} for exact statements. In the context of linear models, \citet{buhlmann2015magging,meinshausen2015maximin,guo2020inference} have argued that the maximin effects, as a special case of $f^*$ by restricting $\FF$ to the class of linear models, provides a measure of stable associations shared across different sources. As an important generalization, the general $f^*$ is useful in capturing the shared associations of heterogeneous and possibly non-linear relationships across different sources. To illustrate this, we consider the additive models with $\f{l}(x) = \sum_{j=1}^2 \f{l}_j(x_{j})$ for $x\in \R^2$ and $l\in[L]$ where $\f{1}_1 = \f{2}_1 = ... = \f{L}_1=f_1$ and $\{\f{l}_2\}_{l\in[L]}$ are randomly scattered around $0$. Our proposed distributionally robust prediction model admits the form $f^*(x)=f_1(x_1)$, capturing the shared component across different sites.
\begin{Remark}[Different identification strategy from group DRO and agnostic federated learning]\rm
We compare to the existing DRO algorithms and emphasize a few unique advantages of closed-form identification in Theorem \ref{thm: identification}. Firstly, our identification result in Theorem \ref{thm: identification} typically leads to a computationally more efficient algorithm than the existing DRO algorithms.
{The standard approach in DRO literature usually involves two-level gradient-based updates to solve the minimax optimizations \citep{sagawa2019distributionally,mohri2019agnostic,deng2020distributionally,oren2019distributionally}. Specifically, in each iteration, their approach involves fixing the model parameters to update the group weights, followed by updating the model parameters using the updated group weights. This process is repeated until a convergence criterion is met.} In contrast, Theorem \ref{thm: identification} provides a different approach by calculating the closed form of the distributionally robust prediction model. Notably, we only need to train the machine learning models for each source data $\{\f{l}\}_{l\in[L]}$ and compute the optimal aggregation weight. Compared to the interleaving gradient-based updates, the closed-form in Theorem \ref{thm: identification} yields a computationally more efficient algorithm. Secondly, our identification in Theorem \ref{thm: identification} is generally applicable to a wide range of machine learning algorithms, including random forests, boosting, and deep neural networks, while most existing group DRO algorithms are mainly designed for the gradient-based algorithm \citep{sagawa2019distributionally,mohri2019agnostic,deng2020distributionally,oren2019distributionally}, excluding random forests or boosting methods that are based on decision trees.
\end{Remark}
%With Magging, only summary statistics of $\left\{\hf{l}\right\}_{l\in[L]}$ are shared across sites, allowing for privacy-preserving collaboration. In contrast, other DRO approaches struggle with such data privacy issues.
%\Zijian{Stop here.} the implementation of our DRL method In practice, it is not necessary to input the entire source dataset into the estimation algorithm, but only the separate summary statistics $\left\{\hf{l}\right\}_{l\in[L]}$ fitted for each source group. This means that we can fit $L$ models in parallel for each group with a relatively smaller sample size, then identify the weights using Theorem \ref{thm: identification}, and finally aggregate then into one robust model. This approach is in contrast to most DRO methods fitting over the entire source dataset comprising $L$ groups until convergence. \Zhenyu{Want to mention the connection with Federated learning.}
%Theorem \ref{thm: identification} also has significant implications for constructing privacy-preserving distributionally robust algorithms.
%This approach demonstrates that how to construct a group distributionally robust model by determining the optimal weights to aggregate the conditional mean models. A broader class of machine learning models, including random forests and boosting, can be incorporated into out framework. This extends the applicability of the distributionally robustness scheme beyond gradient-based methods, as we do not limit how each $\left\{\f{l}\right\}_{l\in[L]}$ is fitted before aggregation.
%Moreover,
%%%%%%%%%%%%%%%%%%%%%% 2.3 %%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Squared Loss and Other Uncertainty Classes}
\label{sec: discussion_loss}
The loss or reward function and the uncertainty class are key components of defining the distributionally robust prediction model in \eqref{eq: f_star full}. In the following, we discuss the pitfalls with the squared loss function and also highlight the difference between our proposed DRL and the existing works in terms of defining the uncertainty class.
%Pitfalls with the Squared Loss and the Uncertainty Set in Group DRO}
As an alternative to $f^*$ in \eqref{eq: f_star full}, we may define the distributionally robust prediction model based on the standard squared loss function,
\begin{equation}
f^{\rm sq} := \argmin_{f\in \FF} \max_{\TT \in \C(\QQ_X)}\E_\TT[(Y-f(X))^2].
\label{eq: l2 loss def}
\end{equation}
In the following proposition, we consider special settings with $L=2$ and illustrate an interesting and surprising phenomenon {with an explicit solution for}
%that
$f^{\rm sq}$.
%does not provide a robust prediction guarantee.
This phenomenon was informally described at page 6 of \citet{meinshausen2015maximin} for the case of linear regression settings.
%, specifically for
%\Peter{Perhaps add here: "The latter has been informally described in Meinshausen \& PB (2015, page 6)."}
\begin{Proposition}
We consider the model \eqref{eq: source-models} with $L=2$, $\f{1} \neq \f{2}$, $\E[(\eps{1}_i)^2|\X{1}_i] = \sigma^2_1$ and $\E[(\eps{2}_i)^2|\X{2}_i] = \sigma^2_2$, and $\eps{l}_i$ being independent of $\X{l}_i$ for $i\in [n_l]$ and $l\in[L]$. Then $f^{\rm sq}$ admits the following expression,
\[
f^{\rm sq} = q_1 \f{1} + (1-q_1) \f{2} \quad {\rm with}\quad q_1 = 0\vee \left(\frac{1}{2} + \frac{\sigma_1^2 - \sigma_2^2}{\left\|\f{1}-\f{2}\right\|_{\QQ,2}^2}\right) \wedge 1.
\]
Consequently, when $\sigma_1^2 = \sigma_2^2$, $f^{\rm sq} = \frac{1}{2}\f{1} + \frac{1}{2}\f{2}$; when $\sigma_1^2 \gg \sigma_2^2$, $f^{\rm sq} = \f{1}$.
\label{prop: squared loss model}
\end{Proposition}
\begin{figure}[htp]
\centering
\includegraphics[width=0.85\textwidth]{assets/illus_sqloss.png}
\caption{Illustration of our proposed distributionally robust prediction model $f^*$(blue point), squared loss robust prediction model $f^{\rm sq}$(green point) for $p=2$ and $L=2$, and the additive models $\f{l}(x) = \sum_{j=1}^2 \f{l}_j(x_j)$ for $l\in [L]$. The left panel illustrates the homogeneous noise case with $\sigma_1^2 = \sigma_2^2$; the right panel illustrates the heterogeneous noise case with $\sigma_1^2\gg \sigma_2^2$.}
\label{fig: illus squareloss}
\end{figure}
$f^{\rm sq}$ can also be viewed as a convex combination of $f^{(1)}$ and $f^{(2)}$, but it exhibits reduced robustness compared to our proposed $f^*$.
In the scenario of homogeneous noise with $\sigma_1^2=\sigma_2^2$, we observe a surprising phenomenon that $f^{\rm sq}$ equals the simple average (half and half) of $f^{(1)}$ and $f^{(2)}$, which does not rely on the distribution of the target covariates. Moreover, it will be drastically influenced by an outlier group by simple averaging. We compare $f^{\rm sq}$ and $f^*$ in the left panel of Figure \ref{fig: illus squareloss}. In the context of heterogeneous noise settings, it is evident that $f^{\rm sq}$ is highly sensitive to the noise level in each group. The model $f^{\rm sq}$ mainly optimizes the prediction performance for the source data with noisier regression errors. We have illustrated this on the right panel of Figure \ref{fig: illus squareloss}. The main difference for $f^*$ is that it eliminates the effect for heterogeneous noise levels across multiple groups.
% \begin{Remark}[Uncertainty set definition]\rm
In addition to the loss function, it is critical to specify the uncertainty class. We discuss in the following the major difference between the definition in \eqref{eq: f_star H} and the group DRO \citep{sagawa2019distributionally, hu2018does} or agnostic federated learning \citep{mohri2019agnostic,deng2020distributionally} lies how to define the uncertainty set $\C(\QQ_X)$. Particularly, \cite{sagawa2019distributionally, hu2018does,mohri2019agnostic,deng2020distributionally} defined the uncertainty set as $\C_0(\HH):= \left\{\TT | \TT = \sum_{l=1}^L q_l \cdot \PP{l} \; \text{with}\; q\in \HH \subseteq \Delta^L \right\},$ a class containing the mixture of the joint distributions $\{\PP{l}\}_{l\in[L]}$ with the weight belonging to $\HH.$ In contrast, our definition \eqref{eq: region-Q H} only imposes the mixture assumption on the conditional outcome distributions but considers the identifiable target covariate distribution $\QQ_{X}$. When there is no covariate shift between the target and source populations, then the two uncertainty sets coincide with each other. We shall remark that our introduced uncertainty set definition \eqref{eq: region-Q H} is more useful in the presence of covariate shift. Firstly, in our focused unsupervised domain adaption regime, if the target population belongs to $\C_0(\HH)$, then the available covariate observations help us identify the mixture structure of the target population, and there is no need to construct the distributionally robust prediction model. Moreover, in practice, the covariate distributions for the source and target populations may differ; for example, the genetic variants might have a very different distribution when the data are collected for different subpopulations \citep{leek2010tackling,ling2022batch,liu2021integrative,weber2022international}.
% \Zijian{Zhenyu, add some references about this. You need to do some creative searching online.}\Zhenyu{Done}
However, the conditional outcome model across different groups might exhibit more similarity as they are measuring a similar biological process for different subpopulations. The uncertainty set in \eqref{eq: region-Q H} could be more meaningful than $\C_0(\HH)$ for these applications. %\Zhenyu{Zijian, I moved the remark here.}
% \label{rem: DRO comparison}
% \end{Remark}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%% 3.0 %%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Distributionally Robust Machine Learning}
%\Zijian{Peter, for the session title, should we use the term ``federated machine learning" or ``federated learning"? This is also a question for our paper's title.} \Peter{I don't know for sure... I was reading that in Section 3.3 you have "Federated" and thus I thought that we may have it for the Section 3 title as well. It's true that in 3.3 it is "federated learning" which seems to be the more common name than "federated machine learning".}
\label{sec: methodology}
In this section, we leverage the identification result in Theorem \ref{thm: identification} and devise a computationally efficient and privacy-preserving distributional robust learning framework. Our proposal can be seen as distributionally robust federated learning, where the aggregation weight is determined by the sample version of the optimization problem in \eqref{eq: identification}. This is in sharp contrast to the inverse variance weighting or other standard meta or federated learning algorithms.
Throughout the discussion, we use $\hf{l}$ to denote the estimator of $\f{l}$ for $l\in[L]$, where $\hf{l}$ may denote any machine learning estimator, including but not limited to high-dimensional regression, random forests, boosting, and deep neural network.
% To simplify the discussion, we define $\hF = \left(\hf{1}, \hf{2}, ..., \hf{L}\right)$.
We apply Theorem \ref{thm: identification} and construct the following estimator of $\fH$,
\begin{equation}
% \tfH = \sum_{l=1}^L \widetilde{q}_l \cdot \hf{l}\quad \textrm{with} \quad \widetilde{q} = \argmin_{q\in \HH} q^\intercal \tGamma q \quad \textrm{where}\quad \tGamma = \frac{1}{n_\QQ}\sum_{j=1}^{n_\QQ} \hF(\XQ_j)\hF(\XQ_j)^\intercal.
\tfH = \sum_{l=1}^L \widetilde{q}_l \cdot \hf{l}\quad \textrm{with} \quad \widetilde{q} = \argmin_{q\in \HH} q^\intercal \tGamma q \quad \textrm{where}\quad \tGamma_{k,l} = \frac{1}{n_\QQ}\sum_{j=1}^{n_\QQ} \hf{k}(\XQ_j)\hf{l}(\XQ_j) \; \textrm{ for $k,l\in [L]$}.
\label{eq: Magging plugin}
\end{equation}
In the above construction, we estimate $\Gamma$ in \eqref{eq: identification} by its sample version $\tGamma$. The accuracy of the estimator $\tfH$ comes from the estimation errors of $\widehat{f}^{(l)}$ and the weight estimator $\widetilde{q},$ which further depends on the error of the matrix estimator $\tGamma.$ In the following, we propose a novel bias-correction step to improve the estimation accuracy of $\tGamma.$ This bias-correction step will lead to a more accurate estimator than the plug-in estimator $\tfH$ defined in \eqref{eq: Magging plugin}.
Now we explain the main idea of bias correction and then provide the full details immediately. For $k,l \in [L]$, the error of the plug-in estimator $\tGamma_{k,l}$ is decomposed into
{\small
\begin{align*}
&\tGamma_{k,l} - \Gamma_{k,l} = \frac{1}{n_\QQ}\sum_{j=1}^{n_\QQ}\hf{k}(\XQ_j)\hf{l}(\XQ_j) - \E_{\QQ_X}[\f{k}(\XQ)\f{l}(\XQ)] \notag \\
&\quad = \frac{1}{n_\QQ}\sum_{j=1}^{n_\QQ}\left(\hf{k}(\XQ_j)\hf{l}(\XQ_j) - \f{k}(\XQ_j)\f{l}(\XQ_j)\right) + \frac{1}{n_\QQ}\sum_{j=1}^{n_\QQ}
\left(\f{k}(\XQ_j)\f{l}(\XQ_j) - \E_{\QQ_X}[\f{k}(\XQ)\f{l}(\XQ)]\right).
\label{eq: tGamma error decompose}
\end{align*}}
In the above decomposition, the second term is the finite-sample uncertainty even when $\f{l}$ and $\f{k}$ are known, while the first term reflects the estimation error of $\hf{l}$ and $\hf{k}$. We further decompose the first term of the above decomposition as follows, %and highlight its leading bias components, %in $\tGamma_{k,l}$ and can be decomposed further as follows:
{\small
\begin{equation}
\begin{aligned}
&\underbrace{\frac{1}{n_\QQ} \sum_{j=1}^{n_\QQ} \hf{k}(\XQ_j) \left(\hf{l}(\XQ_j) - \f{l}(\XQ_j)\right)}_{\textrm{Bias-kl}}+
\underbrace{\frac{1}{n_\QQ} \sum_{j=1}^{n_\QQ} \hf{l}(\XQ_j) \left( \hf{k}(\XQ_j) - \f{k}(\XQ_j)\right)}_{\textrm{Bias-lk}} \\
& - \underbrace{\frac{1}{n_\QQ} \sum_{j=1}^{n_\QQ} \left(\hf{k}(\XQ_j) - \f{k}(\XQ_j)\right)\left( \hf{l}(\XQ_j) - \f{l}(\XQ_j)\right)}_{\textrm{higher order bias}}.
\end{aligned}
\label{eq: error decomposition}
\end{equation}}
The terms ``Bias-kl'' and ``Bias-lk'' are the dominating terms, while the ``higher-order bias'', as the product of $\hf{k}-\f{k}$ and $\hf{l}-\f{l}$, typically diminishes to zero at a faster rate. %rapidly when both estimation errors converge to $0$ at a sufficiently fast rate.
In the following, we discuss the intuitive idea of estimating the ``Bias-kl'' term, and similar reasoning can be applied to the ``Bias-lk'' term.
%The next step is to correct the bias by estimating the terms ``Bias-kl'' and ``Bias-lk''.
%{The bias correction technique is widely used for accurately estimating quadratic forms and inner products in high-dimensional linear models \citep{verzelen2018adaptive,tony2020semisupervised,guo2020inference}, but is new in terms of machine learning literature. We offer insights
If $\hf{k},\hf{l}$ are independent of $(X, Y)$, then
\begin{equation}
\E_{Y\sim\PP{l}_{Y|X}}\left[\left.\hf{k}(X) \left(\hf{l}(X) - Y\right) \right\vert X = \XQ_j\right]=\hf{k}(\XQ_j) \left(\hf{l}(\XQ_j) - \f{l}(\XQ_j)\right)
\label{eq: bias correct idea}
\end{equation}
where the conditional expectation is taken with respect to $\PP{l}_{Y|X}$ and evaluated at $X = \XQ_j$. We then take the expectation with respect to the covariates and it implies that
\[
\E_{X\sim \QQ_X,\;Y\sim\PP{l}_{Y|X}}\left[\hf{k}(X) \left(\hf{l}(X) - Y\right) \right] = \E_{X\sim \QQ_X}\left[\hf{k}(X) \left(\hf{l}(X) - \f{l}(X)\right)\right].
\]
Therefore, when the data $\{X_i,Y_i\}_{i\in [m]}$ are sampled from $(\QQ_X, \PP{l}_{Y|X})$ and independent of $\hf{k},\hf{l}$, we can estimate the ``Bias-kl'' term in \eqref{eq: error decomposition} based on the following expression
\begin{equation}\frac{1}{m}\sum_{i=1}^{m}\hf{k}(X_i) \left(\hf{l}(X_i) - Y_i\right) \approx \textrm{Bias-kl}.
\label{eq: intuitive expression}
\end{equation}
%\Zijian{Why the above is true? We have integrated out the covariate distribution, right?}
We shall employ the sample splitting to create the independence between the initial estimator and the data used for bias correction as in \eqref{eq: bias correct idea}.
%{Sampling splitting is needed to ensure the independence between the data-dependent estimators $\left\{\hf{l}\right\}_{l\in[L]}$ and the data used for bias correction. }
We randomly partition the source data of the $l$-th group into two disjoint subsets $\A_l$ and $\B_l$. To simplify the discussion, we assume that $\A_l$ contains the first $\lfloor n_l/2\rfloor$ observations, while $\B_l$ contains the remaining ones, i.e., $\A_l = \{1,2,..., \lfloor n_l/2\rfloor \}$ and $\B_l = [n_l] \setminus \A_l$. We then construct the machine learning estimators $\hf{l}_\A$ and $\hf{l}_\B$ using the data $\A_l$ and $\B_l$, respectively, and will further employ the cross-fitting idea \citep{klaassen1987consistent,schick1986asymptotically,chernozhukov2018double} to improve the estimation accuracy of $\Gamma$. We shall emphasize that the sample split estimators $\{\hf{l}_A, \hf{l}_\B\}_{l\in [L]}$ are employed solely for the construction of the bias-corrected matrix estimator $\hGamma$. After obtaining the estimated aggregation weights, we still aggregate the machine learning estimators $\{\hf{l}\}_{l\in [L]}$ that are constructed with the full data.
In the following subsections, we provide the details of constructing the bias-corrected estimator $\hGamma$ of $\Gamma$ for the settings with and without covariate shifts. We construct the bias-corrected DRL estimator as
\begin{equation}
\hfH = \sum_{l=1}^L \widehat{q}_l\cdot \hf{l} \quad \textrm{with} \quad \widehat{q} = \argmin_{q\in \HH} q^\intercal \hGamma q.
\label{eq: Magging cs}
\end{equation}
The main difference from $\tfH$ in \eqref{eq: Magging plugin} is that we use the bias-corrected matrix estimator $\hGamma$ instead of the plug-in estimator $\tGamma$.
%and apply the above bias-correction idea to improve the estimation accuracy of $\Gamma$ in both settings.
%%%%%%%%%%%%%%%%%%%%%% 3.1 %%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Bias-corrected estimator of $\Gamma$ for no covariate shift setting}
\label{sec: no shift}
%%%%%%%%%%%%%%%%%%%%%% 3.1 %%%%%%%%%%%%%%%%%%%%%%%%%%
We first consider the no covariate shift setting with $\PP{l}_X=\QQ_X$ for $l\in[L]$. In this case, the data $(\X{l}_i,\Y{l}_i)\distras{iid} (\QQ_X, \PP{l}_{Y|X})$ for the $l$-th source group. For $1\leq k\leq l\leq L$,
we leverage the expression \eqref{eq: intuitive expression} to construct the following bias-corrected estimator of $\Gamma_{k,l}$,
% \begin{equation}
% \hGamma_{k,l}^\A = \tGamma_{k,l} - \widehat{\mathcal{D}}_{k,l}^\A - \widehat{\mathcal{D}}_{l,k}^\A\quad
% \textrm{where}\quad \widehat{\mathcal{D}}_{k,l}^\A =\frac{1}{|\B_l|}\sum_{i \in \B_l} \hf{k}_\A(\X{l}_i)\left(\hf{l}_\A(\X{l}_i) - \Y{l}_i \right)
% \label{eq: Gamma-debias-NoCS-A}
% \end{equation}
% where $\tGamma_{k,l}$ is defined in \eqref{eq: Magging plugin} and $\widehat{\mathcal{D}}_{k,l}^\A$, $\widehat{\mathcal{D}}_{l,k}^\A$ are used to estimate the ``Bias-kl'', ``Bias-lk'' in \eqref{eq: error decomposition}, respectively.
% \begin{equation}
% \begin{aligned}
% \hGamma_{k,l}^\A =& \frac{1}{n_\QQ} \sum_{j=1}^{n_\QQ} \hf{k}_\A(\XQ_j) \hf{l}_\A(\XQ_j)
% - \frac{1}{|\B_l|}\sum_{i \in \B_l} \hf{k}_\A(\X{l}_i)\left(\hf{l}_\A(\X{l}_i) - \Y{l}_i \right)\\
% &- \frac{1}{|\B_k|}\sum_{i \in \B_k} \hf{l}_\A(\X{k}_i)\left(\hf{k}_\A(\X{k}_i) - \Y{k}_i\right)
% \end{aligned}
% \label{eq: Gamma-debias-NoCS-A}
% \end{equation}
{\small
\begin{equation}
\hGamma_{k,l}^\A = \tGamma_{k,l}^\A - \widehat{\mathcal{D}}_{k,l}^\A - \widehat{\mathcal{D}}_{l,k}^\A\quad
\textrm{with}\quad\tGamma_{k,l}^\A=\frac{1}{n_\QQ}\sum_{j=1}^{n_\QQ} \hf{k}_\A(\XQ_j)\hf{l}_\A(\XQ_j), \quad \widehat{\mathcal{D}}_{k,l}^\A =\frac{1}{|\B_l|}\sum_{i \in \B_l} \hf{k}_\A(\X{l}_i)\left(\hf{l}_\A(\X{l}_i) - \Y{l}_i \right)
\label{eq: Gamma-debias-NoCS-A}
\end{equation}}
where $\tGamma_{k,l}^\A$ denotes the plugin estimator and $\widehat{\mathcal{D}}_{k,l}^\A$, $\widehat{\mathcal{D}}_{l,k}^\A$ are used for estimating ``Bias-kl'', ``Bias-lk'' in \eqref{eq: error decomposition}, respectively. Essentially, we construct $\hGamma^\A_{k,l}$ using the sample split estimators $\hfA{k}$, $\hfA{l}$ which are independent of the data $\B_l$, $\B_k$.
% combine the subsets to form two independent datasets, $\A = \bigcup_{l\in[L]} \A_l$ and $\B = \bigcup_{l\in[L]} \B_l$ and construct the the machine learning estimators $\left\{\hfA{l}\right\}_{l\in[L]}$ and $\left\{\hfB{l}\right\}_{l\in[L]}$ using the data from $\left\{\A_l\right\}_{l\in[L]}$ and $\left\{\B_l\right\}_{l\in[L]}$, respectively.
%Using the estimators $\hfA{k}$, $\hfA{l}$ and data from $\B_l$, $\B_k$, we propose a bias-corrected estimator $\hGamma_{k,l}^\A$ for .
%In the above expression, we leverage the approximation in \eqref{eq: intuitive expression} and
%The expression on the right-hand side of equation \eqref{eq: Gamma-debias-NoCS-A} consists of three terms. The first term is the plugin estimator, the average taken over the target observations $\left\{\XQ_j\right\}_{1\leq j\leq n_\QQ}$. The second and the third terms take the average over the split source data $\B_l$ and $\B_k$. And they are approximations of the ``Bias-kl'' term and ``Bias-lk'' terms in the error decomposition shown in \eqref{eq: error decomposition}, respectively.
%\Zijian{Need to explain why they are approximations.}
%\Zhenyu{To be specific, the second term in \eqref{eq: Gamma-debias-NoCS-A} can be approximated as follows:
%{\small
%\begin{align*}
% \frac{1}{|\B_l|}\sum_{i \in \B_l} \hf{k}_\A(\X{l}_i)\left(\hf{l}_\A(\X{l}_i) - \Y{l}_i \right)
% \approx
% \E_{X\sim \QQ_X, Y\sim\PP{l}_{Y|X}}\left[\hf{k}(X) \left(\hf{l}(X) - Y\right) \right]
%\end{align*}
%}
%as $(\X{l}_i,\Y{l}_i) \sim (\QQ_X, \PP{l}_{Y|X})$ for $i\in \B_l$ under no covariate shift setting. Therefore, it approximates the ``Bias-kl'' term. A similar argument is applied to the third term in \eqref{eq: Gamma-debias-NoCS-A} for approximation of the ``Bias-lk'' term.
Next, we switch the roles of $\{\A_l\}_{l\in[L]}$ and $\{\B_l\}_{l\in[L]}$ and construct the symmetric bias-corrected estimator,
{\small
\begin{equation}
\hGamma_{k,l}^\B = \tGamma_{k,l}^\B - \widehat{\mathcal{D}}_{k,l}^\B - \widehat{\mathcal{D}}_{l,k}^\B\quad
\textrm{with}\quad \tGamma_{k,l}^\B=\frac{1}{n_\QQ}\sum_{j=1}^{n_\QQ} \hf{k}_\B(\XQ_j)\hf{l}_\B(\XQ_j), \quad \widehat{\mathcal{D}}_{k,l}^\B =\frac{1}{|\A_l|}\sum_{i \in \A_l} \hf{k}_\B(\X{l}_i)\left(\hf{l}_\B(\X{l}_i) - \Y{l}_i \right).
\label{eq: Gamma-debias-NoCS-B}
\end{equation}
}
% \begin{equation}
% \begin{aligned}
% \hGamma_{k,l}^\B =& \frac{1}{n_\QQ} \sum_{j=1}^{n_\QQ} \hf{k}_\B(\XQ_j) \hf{l}_\B(\XQ_j)
% - \frac{1}{|\A_l|}\sum_{i \in \A_l} \hf{k}_\B(\X{l}_i)\left(\hf{l}_\B(\X{l}_i) - \Y{l}_i \right)\\
% &- \frac{1}{|\A_k|}\sum_{i \in \A_k} \hf{l}_\B(\X{k}_i)\left(\hf{k}_\B(\X{k}_i) - \Y{k}_i\right).
% \end{aligned}
% \label{eq: Gamma-debias-NoCS-B}
% \end{equation}
We propose the final bias-corrected estimator of $\Gamma$ as
$\hGamma_{k,l} = \frac{1}{2}\left(\hGamma_{k,l}^\A + \hGamma_{k,l}^\B\right)$ {for} $1\leq k \leq l\leq L.$ For $1\leq l<k\leq L,$ we use the symmetry property and set $\hGamma_{k,l} = \hGamma_{l,k}.$
%We then use symmetry to impute the upper triangular part of .
% The above construction ensures that $\hGamma$ is a symmetric matrix.
%\Zijian{Zhenyu, is this really symmetric?} \Zhenyu{Done}
%We construct the bias-corrected DRL estimator as follows:
%\begin{equation}
% \hfH = \sum_{l=1}^L \widehat{q}_l\cdot \hf{l} \quad \textrm{with} \quad \widehat{q} = \argmin_{q\in \HH} q^\intercal \hGamma q.
% \label{eq: Magging nocs}
%\end{equation}
% \begin{algorithm}
% \caption{Bias-corrected Magging Estimator for No Covariate Shift Setting}
% \KwData{Source data $\{(\X{l}_i, \Y{l}_i)\}$ for $1\leq i\leq n_l$ and $l\in[L]$; target data $\XQ_j$ for $1\leq j\leq n_\QQ$; the convex subset $\HH$ (with the default value $\HH=\Delta^{L}$).}
% \KwResult{Magging estimator $\hfH$}
% \For{$l=1,...,L$}{
% Split the $l$-th source data into two disjoint subsets $\A_l$ and $\B_l$
% Using the data $\{(\X{l}_i, \Y{l}_i)\}_{i\in \A_l}$ to construct the model estimator $\hf{l}_\A$
% Using the data $\{(\X{l}_i, \Y{l}_i)\}_{i\in \B_l}$ to construct the model estimator $\hf{l}_\B$
% }
% \For{$k,l = 1,..., L$}{
% % Using models trained on the subdatasets $\A_l$ and $\A_k$ including $\hf{k}_\A$, $\hf{l}_\A$, perform bias correction on the subdatasets $\B_l$ and $\B_k$ to compute $\hGamma^\A_{k,l}$ as \eqref{eq: Gamma-debias-NoCS-A}
% % Using models trained on the subdatasets $\B_l$ and $\B_k$ including $\hf{k}_\B$, $\hf{l}_\B$, perform bias correction on the subdatasets $\A_l$ and $\A_k$ to compute $\hGamma^\B_{k,l}$ as \eqref{eq: Gamma-debias-NoCS-A}
% Compute $\hGamma^\A_{k,l}$ and $\hGamma^\B_{k,l}$ as \eqref{eq: Gamma-debias-NoCS-A} and \eqref{eq: Gamma-debias-NoCS-B}, respectively
% Construct the bias-corrected estimator $\hGamma_{k,l} = \frac{1}{2}\left(\hGamma^\A_{k,l} + \hGamma^\B_{k,l}\right)$
% }
% Construct the data-dependent optimal weight as $\widehat{q} = \argmin_{q\in \mathcal{H}} q^\intercal \hGamma q$
% Return $\widehat{f}_\HH = \sum_{l=1}^L \widehat{q}_l \cdot \hf{l}$
% \label{alg: non-cs mm}
% \end{algorithm}
%%%%%%%%%%%%%%%%%%%%%% 3.2 %%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Bias-corrected estimator of $\Gamma$ for covariate shift setting}
\label{sec: covariate shift}
In this subsection, we consider a more challenging scenario allowing for covariate shifts, that is, $\QQ_X$ may not be identical to any $\{\PP{l}_X\}_{l\in[L]}$, and the covariate distributions $\{\PP{l}_X\}_{l\in[L]}$ may differ from each other.
% \Peter{I would delete the following sentence -- it is not important?}
% The key difference from the no covariate shift setting is that $(\X{l}_i,\Y{l}_i)$ follows $(\PP{l}_X, \PP{l}_{Y|X})$ instead of $(\QQ_X, \PP{l}_{Y|X}).$
We shall adopt the importance weighting strategy, which is commonly used in handling the covariate shift; see \citet{shimodaira2000improving,gretton2009covariate,sugiyama2012density, reddi2015doubly,menon2016linking, tibshirani2019conformal,choi2022density} for examples.
The main idea is to reweight the source data $\{(\X{l}_i,\Y{l}_i)\}_{i\in [n_l]}$ for each source group $l\in [L]$, such that the reweighted source distribution is nearly the same as $(\QQ_X, \PP{l}_{Y|X}).$ Then we can leverage \eqref{eq: intuitive expression} to estimate the bias components with the reweighted source distribution. For each $l\in[L]$, we assume that $\QQ_X$ is absolutely continuous with respect to $\PP{l}_X$ and define the density ratio $\w{l}(x) = d\QQ_X(x)/ d\PP{l}_X(x)$. The key step of implementing reweighting is to estimate the density ratios $\{\w{l}(\cdot)\}_{l\in[L]}$. Notably, our DRL framework is compatible with various existing density ratio estimations; see \citet{sugiyama2012density} for a thorough review.
In the following discussion, we apply the Bayes formula to estimate the density ratios as an illustrative example and then use the density ratio estimators to perform the bias correction.
% we introduce one of the approaches called Probabilistic Classification \citep{sugiyama2012density} to estimate the density ratios and then use the density ratio estimators to perform the bias correction. \Zijian{Can we just say we use the bayes formula to estimate the density ratio?}
\vspace{2mm}
\noindent \textbf{An Illustrative Example: Density Ratio Estimation by Bayes Formula}\\
For each source group $l \in [L]$, we merge its covariates observations $\{{X}^{(l)}_i\}_{i\in [n_l]}$ with the target covariates $\{\XQ_j\}_{j\in [n_\QQ]}$ and denote the merged covariates as $\{\widetilde{X}^{(l)}_i\}_{i\in [n_l + n_\QQ]}$, where the first $n_l$ observations are from the $l$-th source group and the remaining $n_\QQ$ observations are from the target group. We also define an indicator random variable $G^{(l)}_i\in \{0,1\}$ for $i \in [n_l+n_\QQ]$: $\G{l}_i = 0$ if $\widetilde{X}^{(l)}_i$ is drawn from the $l$-th source group and $\G{l}_i = 1$ if $\widetilde{X}^{(l)}_i$ is drawn from the target distribution. The density of $l$-th source covariates and target covariates are written as $d\PP{l}_X(x) = p(x|\G{l}_i=0)$ and $d\QQ_X(x) = p(x|\G{l}_i=1)$. We apply the Bayes formula and obtain the following expression of the density ratio,
\[
\w{l}(x) = \frac{d\QQ_X(x)}{d\PP{l}_X(x)} =
\frac{p(\G{l}_i=0)}{p(\G{l}_i=1)} \frac{p(\G{l}_i = 1|x)}{p(\G{l}_i = 0|x)}.
\]
The ratio ${p(\G{l}_i=0)}/{p(\G{l}_i=1)}$ can be approximated by the ratio of sample size.
%Notably, here we follow the case-control sampling scheme discussed in the \citet{qin1998inferences}, and assume that the marginal density of $\widetilde{X}^{(l)}_i$ exists. \Zhenyu{Zijian, please check the preceding sentence.}
The class posterior probability $p(\G{l}_i|x)$ may be approximated using a probabilistic classifier, like the logistic regression and the random forests classifier. As an example, we consider the binary-class logistic regression with
\begin{equation}
p(\G{l}_i = 1|x) = h(x^\intercal \gamma^{(l)})\quad \textrm{with} \quad h(z) = \exp(z)/(1+\exp(z))
% p(\G{l}_i = 1|x) = {\exp(x^\intercal \gamma^{(l)})}/(1 + \exp(x^\intercal \gamma^{(l)}))
\label{eq: density ratio}
\end{equation}
where $\gamma^{(l)}$ denotes the regression vector.
For implementation, we still employ sample splitting to estimate the density ratios to avoid the potential overfitting issue for machine learning algorithms. As we did previously in Section \ref{sec: methodology}, each source data of the $l$-th group is randomly partitioned into two disjoint subsets $\A_l$ and $\B_l$. We estimate $\gamma^{(l)}$ by $\widehat{\gamma}^{(l)}_\A$ and $\widehat{\gamma}^{(l)}_\B$ based on the data belonging to $A_l$ and $B_l$, respectively. Examples include the MLE estimator in low dimensions and the penalized log-likelihood estimators in high dimensions; see Section \ref{sec: add_methods} in the supplementary material for the exact formulas.
% We may estimate $\gamma^{(l)}$ by the MLE estimator $\widehat{\gamma}^{(l)}$ in the low-dimensional regime while in high dimensions, we propose to estimate $\gamma^{(l)}$ by the penalized log-likelihood estimator defined as
% {\small
% \begin{equation}
% \widehat{\gamma}^{(l)} = \argmin_{\gamma \in \R^p}\left\{-\frac{1}{|\A_l|+n_\QQ} \sum_{i=1}^{|\A_l|+n_\QQ} \left[\log(1+\exp(\widetilde{X}^{(l)\intercal}_{i}\gamma^{(l)})) - \G{l}_i\cdot (\widetilde{X}^{(l)\intercal}_{i}\gamma^{(l)})\right] + \lambda^{(l)} \sum_{j=2}^p \frac{\|\widetilde{X}^{(l)}_{\cdot j}\|_2}{\sqrt{|\A_l|+n_\QQ}} |\gamma_j|\right\}
% \label{eq: penalized MLE}
% \end{equation}}
% where $\widetilde{X}^{(l)} = (X^{(l)\intercal}, X_\QQ^\intercal)^\intercal$, $\widetilde{X}^{(l)}_{\cdot j}$ denotes the $j$-th column of $\widetilde{X}^{(l)}$, the first column of $\widetilde{X}^{(l)}$ is set as the constant 1, and the tuning parameter $\lambda^{(l)}$ is of the order $\sqrt{\log p/(|\A_l|+n_\QQ)}$.
% estimate the ratio $\frac{p(\G{l}_i = 1|x)}{p(\G{l}_i = 0|x)}=\exp(x^\intercal \gamma^{(l)})$ by $\exp(x^\intercal \widehat{\gamma}^{(l)})$.
% With $\widehat{\gamma}^{(l)}$, we construct the class posterior probability $\widehat{p}(\G{l}_i = 1|x) = h(x^\intercal \widehat{\gamma}^{(l)})$
% % $\widehat{p}(\G{l}_i = 1|x) = {\exp(x^\intercal \widehat{\gamma}^{(l)})}/(1 + \exp(x^\intercal \widehat{\gamma}^{(l)}))$
% and $\widehat{p}(\G{l}_i = 0|x) = 1-\widehat{p}(\G{l}_i = 1|x)$ and further estimate the density ratio as follows:
% \begin{equation}
% \hw{l}_\A(x) = \frac{|\A_l|}{n_\QQ} \frac{\widehat{p}(\G{l}_i = 1|x)}{\widehat{p}(\G{l}_i = 0|x)} \quad \text{for}\quad l\in [L].
% \label{eq: density ratio estimation}
% \end{equation}
% The same strategy shall be applied to the disjoint subset $\B_l$, and we can construct the density ratio estimators $\hw{l}_\B(\cdot)$ for each $l\in [L]$ as well.
We then further estimate the density ratio as follows: for $l\in [L]$,
\begin{equation}
\hw{l}_\A(x) = \frac{|\A_l|}{n_\QQ} \frac{h(x^\intercal \widehat{\gamma}^{(l)}_\A)}{1-h(x^\intercal \widehat{\gamma}^{(l)}_\A)}, \quad \textrm{and} \quad
\hw{l}_\B(x) = \frac{|\B_l|}{n_\QQ} \frac{h(x^\intercal \widehat{\gamma}^{(l)}_\B)}{1-h(x^\intercal \widehat{\gamma}^{(l)}_\B)}.
\label{eq: density ratio estimation}
\end{equation}
It is important to highlight that the consistency of our proposed DRL estimator does not rely on the correct specification of the density ratio model, such as \eqref{eq: density ratio}. However, a correctly specified density ratio model leads to a faster rate of convergence; see the details in Section \ref{sec: theory-correct}.
%In addition to the density ratio estimator in \eqref{eq: density ratio estimation}, our DRL framework allows for other approaches for density ratio estimations, such as Kernel Mean Matching \citep{huang2006correcting,gretton2009covariate}.
% our DRL framework is compatible with other density ratio estimators; see \citet{sugiyama2012density} for a thorough review.
%such as Kernel Mean Matching \citep{huang2006correcting,gretton2009covariate}, KL Importance Estimation \citep{sugiyama2007direct}, and Least Squares Importance Fitting \citep{kanamori2009least};
%Once $\left\{\hw{l}(\cdot)\right\}_{l\in [L]}$ is obtained, we split the samples and obtain the model estimators $\hfA{l}(\cdot)$,$\hfB{l}(\cdot)$, on the split data $\A_l$ and $\B_l$, respectively, by following the approach used in the no covariate shift setting, for $l\in[L]$.
\vspace{2mm}
\noindent\textbf{Bias Correction with Estimated Density Ratios}\\
For $1\leq k\leq l\leq L$, we integrate the density ratio estimators $\hw{l}_\A(\cdot)$ and $\hw{k}_\A(\cdot)$ to construct the following bias-corrected estimator of $\Gamma_{k,l}$,
\begin{equation}
\hGamma_{k,l}^\A = \tGamma_{k,l}^\A - \widehat{\mathcal{D}}_{k,l}^\A - \widehat{\mathcal{D}}_{l,k}^\A\quad
\textrm{with}\quad \widehat{\mathcal{D}}_{k,l}^\A =\frac{1}{|\B_l|}\sum_{i \in \B_l} \hw{l}_\A(\X{l}_i)\hf{k}_\A(\X{l}_i)\left(\hf{l}_\A(\X{l}_i) - \Y{l}_i \right)
\label{eq: Gamma-debias-CS-A}
\end{equation}
where $\tGamma_{k,l}^\A$ is defined in \eqref{eq: Gamma-debias-NoCS-A}. With slightly abuse of notations, we still use $\widehat{\mathcal{D}}_{k,l}^\A$, $\widehat{\mathcal{D}}_{l,k}^\A$ for the estimators of terms ``Bias-kl'', ``Bias-lk'' in \eqref{eq: error decomposition}, respectively.
% \begin{equation}
% \begin{aligned}
% \hGamma_{k,l}^\A =& \frac{1}{n_\QQ} \sum_{j=1}^{n_\QQ} \hf{k}_\A(\XQ_j) \hf{l}_\A(\XQ_j)
% - \frac{1}{|\B_l|}\sum_{i \in \B_l} \hw{l}(\X{l}_i)\hf{k}_\A(\X{l}_i)\left(\hf{l}_\A(\X{l}_i) - \Y{l}_i \right)\\
% &- \frac{1}{|\B_k|}\sum_{i \in \B_k} \hw{k}(\X{k}_i)\hf{l}_\A(\X{k}_i)\left(\hf{k}_\A(\X{k}_i) - \Y{k}_i\right).
% \end{aligned}
% \label{eq: Gamma-debias-CS-A}
% \end{equation}
The construction of the above estimator is similar to that in \eqref{eq: Gamma-debias-NoCS-A}, where the main difference is that we reweight the data $\{\X{l}_i\}_{i\in \B_l}$ and $\{\X{k}_i\}_{i\in \B_k}$ by the density ratio functions $\hw{l}_\A(\cdot)$ and $\hw{k}_\A(\cdot)$, respectively. The no covariate shift construction in \eqref{eq: Gamma-debias-NoCS-A} can be viewed as a special case of \eqref{eq: Gamma-debias-CS-A} by taking the density ratio as $1$.
%In an ideal scenario, if the estimated density ratios $\left\{\hw{l}\right\}$ accurately capture the true density ratios, then the covariate distribution of the reweighted source data will be identical to that of the target data.
We then swap the roles of $\left\{\A_l\right\}_{l\in[L]}$ and $\left\{\B_l\right\}_{l\in[L]}$ and construct
\begin{equation}
\hGamma_{k,l}^\B = \tGamma_{k,l}^\B - \widehat{\mathcal{D}}_{k,l}^\B - \widehat{\mathcal{D}}_{l,k}^\B\quad
\textrm{with}\quad \widehat{\mathcal{D}}_{k,l}^\B =\frac{1}{|\A_l|}\sum_{i \in \A_l} \hw{l}_\B(\X{l}_i)\hf{k}_\B(\X{l}_i)\left(\hf{l}_\B(\X{l}_i) - \Y{l}_i \right)
\label{eq: Gamma-debias-CS-B}
\end{equation}
where $\tGamma_{k,l}^\B$ is defined in \eqref{eq: Gamma-debias-NoCS-B}.
% \begin{equation}
% \begin{aligned}
% \hGamma_{k,l}^\B =& \frac{1}{n_\QQ} \sum_{j=1}^{n_\QQ} \hf{k}_\B(\XQ_j) \hf{l}_\B(\XQ_j)
% - \frac{1}{|\A_l|}\sum_{i \in \A_l} \hw{l}(\X{l}_i)\hf{k}_\B(\X{l}_i)\left(\hf{l}_\B(\X{l}_i) - \Y{l}_i \right)\\
% &- \frac{1}{|\A_k|}\sum_{i \in \A_k} \hw{k}(\X{k}_i)\hf{l}_\B(\X{k}_i)\left(\hf{k}_\B(\X{k}_i) - \Y{k}_i\right).
% \end{aligned}
% \label{eq: Gamma-debias-CS-B}
% \end{equation}
Then the final bias-corrected estimator of $\Gamma$ is $\hGamma_{k,l} = \frac{1}{2}\left(\hGamma^\A_{k,l} + \hGamma^\B_{k,l}\right)$ for $1\leq k \leq l\leq L$. For $1\leq l< k\leq L$, we use the symmetry property and set $\hGamma_{k,l} = \hGamma_{l,k}$.
%\Zijian{We need to define $\widehat{f}^{(l)}$. this is a good point.} \Zhenyu{I've highlighted it before the section 3.1.}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Distributionally Robust Federated Learning Algorithm}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
In the following Algorithm \ref{alg: cs mm}, we summarize our proposal for the general covariate-shift setting. We consider that the $L$ source data sets are respectively stored in the $L$ source sites while the data for the target population is stored in the target site. After Algorithm \ref{alg: cs mm}, we provide important remarks on our proposal.
\begin{algorithm}[h]
\DontPrintSemicolon
\SetAlgoLined
\SetNoFillComment
\LinesNotNumbered
\caption{Bias-corrected {\DRL} Estimator}
\KwData{Source data $\{(\X{l}_i, \Y{l}_i)\}_{i\in [n_l]}$ for $l\in[L]$; target covariates $\XQ_j$ for $ j\in [n_\QQ]$; the convex subset $\HH$ (with the default value $\HH=\Delta^{L}$).}
\KwResult{{\DRL} estimator $\hfH$}
\tcc{On source sites}
\For{$l=1,...,L$}{
Construct the initial machine learning estimators $\hf{l}$
Construct the sample splitting density ratio estimator $\hw{l}_\A$ and $\hw{l}_\B$ as in \eqref{eq: density ratio estimation} %with target covariates $\{\XQ_j\}_{j\in [n_\QQ]}$
%Split the $l$-th source data into two disjoint subsets
Construct the sample splitting machine learning estimators $\hf{l}_\A$ and $\hf{l}_\B$ %based on the split data $\A_l$ and $\B_l$, respectively.
}
\tcc{Broadcast $\{\hf{l}_\A, \hw{l}_\A, \hf{l}_\B, \hw{l}_\B\}_{l\in [L]}$ to all source sites}
\For{$k,l=1,...,L$}{
Compute $\widehat{\mathcal{D}}^\A_{k,l}$ and $\widehat{\mathcal{D}}^\B_{k,l}$ as \eqref{eq: Gamma-debias-CS-A} and \eqref{eq: Gamma-debias-CS-B}, respectively
}
\tcc{Transmit $\{\hf{l},\hf{l}_\A,\hf{l}_\B\}_{l\in [L]}$, $\{\widehat{\mathcal{D}}^\A_{k,l},\widehat{\mathcal{D}}^\B_{k,l}\}_{k,l\in [L]}$ to the target site}
\For{$k,l=1,...,L$}{
% Compute the plugin estimators $\tGamma_{k,l}^\A, \tGamma_{k,l}^\B$ as in \eqref{eq: Gamma-debias-NoCS-A} and \eqref{eq: Gamma-debias-NoCS-B}, respectively
% as in \eqref{eq: Magging plugin} but with $\hf{k}_\A,\hf{l}_\A,\hf{k}_\B,\hf{l}_\A$
Compute the sample split estimators $\hGamma^\A_{k,l}, \hGamma^\B_{k,l}$ as in \eqref{eq: Gamma-debias-CS-A} and \eqref{eq: Gamma-debias-CS-B}, respectively
Compute the bias-corrected estimator $\hGamma_{k,l} = \frac{1}{2}\left(\hGamma^\A_{k,l} + \hGamma^\B_{k,l}\right)$
}
% \tcc{Transmit $\{\hf{l}\}_{l\in [L]}$ to the target site}
Construct the data-dependent optimal weight as $\widehat{q} = \argmin_{q\in \mathcal{H}} q^\intercal \hGamma q$
Return $\widehat{f}_\HH = \sum_{l=1}^L \widehat{q}_l \cdot \hf{l}$
\label{alg: cs mm}
\end{algorithm}
%\Zijian{We omit a lot of details here about how to make it private.}\Zhenyu{Modified.}
% \begin{algorithm}
% \caption{Bias-corrected \texttt{DRL} Estimator for Covariate Shift Setting}
% \KwData{Source data $\{(\X{l}_i, \Y{l}_i)\}$ for $i\in [n_l]$ and $l\in[L]$; target data $\XQ_j$ for $ j\in [n_\QQ]$; the convex subset $\HH$ (with the default value $\HH=\Delta^{L}$).}
% \KwResult{\texttt{DRL} estimator $\hfH$}
% \For{$l=1,...,L$}{
% Construct density ratio estimator $\hw{l}$ using the covariates data $\{\X{l}_i\}_{i\in [n_l]}$ and $\{\XQ_j\}_{j\in [n_\QQ]}$
% }
% \For{$l=1,...,L$}{
% Split the $l$-th source data into two disjoint subsets $\A_l$ and $\B_l$
% Construct the model estimator $\hf{l}_\A$ using the data $\{(\X{l}_i, \Y{l}_i)\}_{i\in \A_l}$
% Construct the model estimator $\hf{l}_\B$ using the data $\{(\X{l}_i, \Y{l}_i)\}_{i\in \B_l}$
% }
% \For{$k,l=1,...,L$}{
% Compute $\hGamma^\A_{k,l}$ and $\hGamma^\B_{k,l}$ as \eqref{eq: Gamma-debias-CS-A} and \eqref{eq: Gamma-debias-CS-B}, respectively
% Construct the bias-corrected estimator $\hGamma_{k,l} = \frac{1}{2}\left(\hGamma^\A_{k,l} + \hGamma^\B_{k,l}\right)$
% }
% Construct the data-dependent optimal weight as $\widehat{q} = \argmin_{q\in \mathcal{H}} q^\intercal \hGamma q$
% Return $\widehat{f}_\HH = \sum_{l=1}^L \widehat{q}_l \cdot \hf{l}$
% \label{alg: cs mm}
% \end{algorithm}
%\Zijian{Zhenyu, add a remark on privacy preserving.} \Zijian{My question is that we need to evaluate the estimators on the target data, what should we pass in the implementation? }\Zhenyu{Done}
%\Zijian{Stop here.}
\begin{Remark}[Privacy Preserving Algorithm]\rm
Algorithm \ref{alg: cs mm} does not require the sharing of the $l$-th source site's data $\{(\X{l}_i, \Y{l}_i)\}_{i\in [n_l]}$ to other source sites or the target site. The density ratio estimation step requires transmitting the target covariates to each source group. However, in scenarios where the density ratios are known, such as in the no covariate shift setting, the need for target covariates transmission is eliminated as well. Our proposed DRL algorithm protects data privacy since it mainly requires the trained prediction and the density ratio model from each source data instead of the raw data. It is compiled with the data-privacy regulations such as the HIPPA Privacy Rule \citep{topaloglu2021pursuit} in the United States, which protects patient information and prohibits sharing patient-level data between sites.
%\Zijian{We need to be more explicit about why it is privacy preserving.} \Zijian{I do not think this is a good place to discuss the estimator without bias correction.}\Zhenyu{Modified.}
\end{Remark}
%\Zijian{What is DR random forests?}
%\begin{Remark}[Distributionally robust random forests and neural networks] %\Zijian{Zhenyu, add}\Zhenyu{Modified}
%\rm We consider using random forests to estimate the conditional outcome models, with the random forests estimator $\hf{l}$ averaging the predicted value of $S$ decision trees. Then the distributionally robust estimator $\hfH(x) = \sum_{l=1}^L \widehat{q}_l \cdot \hf{l}(x)$ is still an aggregation of the predicted values given by $L\cdot S$ decision trees, where the trees constructed for different source groups have different aggregation weights (determined by $\widehat{q}_l $), instead of the equal averaging typically used in random forests. Likewise, the insight can be extended to the neural network models. If each $\hf{l}$ is a neural network, $\hfH(x) = \sum_{l=1}^L \widehat{q}_l \cdot \hf{l}(x)$ can be interpreted as adding one more layer at the end to perform a weighted average of the outputs of the $L$ neural networks, respectively. In this way, the final robust model combines the predictions from $L$ neural networks while retaining the form of a neural network. \Peter{This remark is perhaps a bit trivial and not so useful/insightful?}
%For simplicity, we assume that $S^{(l)} = S$ for $l\in [L]$. The growth of each decision tree is determined by the random parameters $\{\theta^{(l)}_1,..., \theta^{(l)}_S\}$, and the conditional outcome model $\f{l}(x)$ is estimated as follows: $\hf{l}(x) = \frac{1}{S}\sum_{s=1}^S \widehat{f}_{\theta^{(l)}_s}(x)$ for $l\in [L]$.
%Following the procedures in Algorithm \ref{alg: cs mm}, we obtain the final robust DRL model with the aggregation weights $\widehat{q} \in \HH$ as follows:
%\[
%\hfH(x) = \sum_{l=1}^L \widehat{q}_l \cdot \hf{l}(x) = \sum_{l=1}^L \sum_{s=1}^S \frac{\widehat{q}_l}{S}\cdot \widehat{f}_{\theta^{(l)}_s}(x)
%\]
%Consequently, $\widehat{f}_{\mathcal{H}}$
% Lastly, we want to highlight the distinction of our proposed {\DRL} Algorithm \ref{alg: cs mm} and directly solving the max-min optimization the empirical version of \eqref{eq: f_star H}. The empirical version can be expressed as follows:
% \[
% \argmax_{f\in \FF} \min_{\TT\in \mathcal{C}(\widehat{\QQ}_X, \HH)} \widehat{\E}_{\TT}[Y^2 - (Y-f(X))^2] = \argmax_{f\in \FF} \min_{q\in \HH} \sum_{l=1}^L q^{(l)}\cdot \widehat{\E}_{}
% \]
% \Zhenyu{Add discussion regarding the difference of directing solving. Issue: there is no direct observations $(X,Y)\sim \QQ$}
%\end{Remark}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%% 4.0 %%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Theoretical Justification}
\label{sec: Theoretical}
In this section, we establish the convergence rate of our proposed DRL estimator to its population counterpart $\fH$. For $l\in[L]$, we use $\Del{l}:= \hf{l}-\f{l}$ to denote the estimation error of $\hf{l}$ and use $n = \min_{l\in[L]} n_l$ to denote the smallest sample size among the $L$ groups of source data.
Before presenting the results, we first state the required conditions.
\begin{Assumption}
{The number of source groups $L$}
%Consider that the group number $L$
is finite. The matrix $\Gamma$ defined in \eqref{eq: identification} is positive definite with
$\eigenGamma > 0$. There exists some positive constant $\sigma^2_{\varepsilon}$ such that
$ \max_{l\in [L]}\max_{i\in [n_l]} \E[(\eps{l}_i)^2|\X{l}_i] \leq \sigma^2_{\varepsilon}.$
% The noise defined in \eqref{eq: source-models} satisfies $\max_{l\in[L]} \max_{1\leq i\leq n_l} \E[(\eps{l}_i)^2|\X{l}_i] \leq C$ for some constant $C$.
\label{ass: eigen}
\end{Assumption}
We assume the constant group number to simplify the presentation while our analysis can be extended to handle a growing $L$. The definition of $\Gamma$ ensures it is a positive semi-definite matrix, and we further require its smallest eigenvalue strictly greater than $0$, which will be fulfilled as long as the models $\{\f{l}\}_{l\in [L]}$ are linearly independent. As for the group noise, we make no assumptions regarding its homogeneity within or across groups but only require its second moment bounded by some constant.
For a function $f$, we define its $\ell_q$ norm measured over the target covariate distribution $\QQ_{X}$ as $\|f\|_{\QQ,q} := \left(\E_{\QQ_X}[f(X)^q]\right)^{1/q}$ for $q\geq 1.$
We define the following norms of $\f{l}$, %measured over the target covariate distribution $\QQ_{X}$, %\Peter{In (23), you write just $\QQ$ and not $\QQ_X$} \Zijian{Peter, I think either notation is fine since $\f{l}$ is only a function of $X$ and even if we integrate with respect to both $X$ and $Y$, it will be integration with respect to $\QQ_{X}$?} \Zijian{Zhenyu, may we replace $\QQ$ by $\QQ_X$? That is, using $\|\f{l}\|_{\QQ_X,2}$?} \Peter{OK, let's leave it as it is!}
\begin{equation}
M\coloneqq\max_{l\in[L]}\|\f{l}\|_{\QQ,2} \quad \text{and}\quad \widetilde{M} \coloneqq \max_{l\in[L]}\|\f{l}\|_{\QQ,4}.
\label{eq: f norm}
\end{equation}
We do not particularly assume $M$ and $\widetilde{M}$ to be bounded but allow them to grow with the sample size $n$ or the dimension $p$. As an example, we consider the high-dimensional sparse regression $\f{l}(X^{(l)}_i)=[X^{(l)}_i]^{\intercal}\beta^{(l)}$ with $X^{(l)}_i\in \R^{p}$ and $\beta^{(l)}\in \R^{p}.$ If we assume $X^{(l)}_i$ to be a sub-Gaussian random vector, we have $M\lesssim \max_{l\in [L]}\|\beta^{(l)}\|_2$ and $\widetilde{M}\lesssim \max_{l\in [L]}\|\beta^{(l)}\|_2$; for Gaussian $X^{(l)}_i$, we have $M\asymp \widetilde{M} \asymp \max_{l\in[L]}\|\beta^{(l)}\|_2.$ In addition, when $\sup_{x\in \R^{p}} |\f{l}(x)|\leq M_0$, we have $\max\{M,\widetilde{M}\}\leq M_0$. This further implies that, for binary outcome variables, $\{\f{l}\}_{l\in[L]}$ are the conditional probability functions, and $M$ and $\widetilde{M}$ are upper bounded by the constant $1$.
In the following, we assume that $\hf{l}$ is a consistent estimator of $\f{l}$ and mainly assume a rate of convergence for the $\ell_2$ norm of estimation error $\Del{l}:=\hf{l}-\f{l}$. Similarly, we assume the same convergence rate for the sample splitting estimators $\Del{l}_\A:=\hf{l}_\A-\f{l}$, and $\Del{l}_\B:=\hf{l}_\B-\f{l}$.
\begin{Assumption} With probability larger than $1-\tau_n$ with $\tau_n\rightarrow 0$, there exists some positive sequence $\delta_n>0$ such that $
\max_{l\in [L]}\max\left\{\|\Del{l}\|_{\QQ,2},\|\Del{l}_\A\|_{\QQ,2},\|\Del{l}_\B\|_{\QQ,2} \right\}\leq \delta_n.
$
%There exists some universal constants $M, \widetilde{M}$ and some term $\delta_n$ (depending on sample size $n$) such that \Zijian{$M, \widetilde{M}$ are not universal constants...}
%where the positive sequences $\tau_n>0$ and $\delta_n>0$ satisfies $\lim_{n\rightarrow \infty}\tau_n=0$. %and $\lim_{n\rightarrow \infty}\delta_n=0.$
% For $l\in[L]$, $\f{l}$ and $\Del{l}$ satisfy:
% \[
% \max_{l\in[L]} \|\f{l}\|_{\QQ,2} \leq M; \quad
% \max_{l\in[L]} \|\f{l}\|_{\QQ,4} \leq \widetilde{M}; \quad
% \max_{l\in[L]} \|\Del{l}\|_{\QQ,2} \leq \delta_n
% \]
\label{ass: plugin}
\end{Assumption}
In the above assumption, $\|\Del{l}\|_{\QQ,2}=\left(\E_{X_i\sim \QQ_X}[\Del{l}(X_i)^2]\right)^{1/2}$ measures the $\ell_2$ norm of $\Del{l}=\hf{l}-\f{l}$ evaluated at the independent target covariates distributions. In the following discussion, we shall require $\hf{l}$ to be a consistent estimator with $\delta_n\rightarrow 0$, though our results will also hold for inconsistent machine learning estimators. Such rate of convergence results have been obtained for various machine learning methods, such as high-dimensional sparse regression \citep[e.g.]{candes2007dantzig, bickel2009simultaneous}, random forests \citep[e.g.]{biau2008consistency,biau2012analysis, scornet2015consistency, meinshausen2006quantile}, and neural networks \citep[e.g.]{schmidt2020nonparametric,farrell2021deep}. In the high-dimensional regression, the Lasso or Dantzig selector meets this assumption \ref{ass: plugin} with $\delta_n= C\sqrt{s_\beta \log p/n}$, where $s_\beta = \max_{l\in [L]}\|\beta^{(l)}\|_0$ denotes the largest sparsity level of $\beta^{(l)}$ for $l\in [L]$. In the context of random forests, \citet{biau2012analysis} established $\delta_n^2 = C\cdot n^{-0.75/ (S\log 2+0.75)}$ under their imposed conditions, where they consider only $S$ out of $p$ features playing a role in the conditional outcome model, see the corresponding discussion after Theorem 2.2 in \citet{biau2012analysis}.
\citet{farrell2021deep} established in their Theorem 1 that certain neural networks satisfy assumption \ref{ass: plugin} with $\delta_n^2 = C\cdot(n^{-\alpha/(\alpha+p)} \log^8 n + \log\log n/n)$, where they assume $\f{l}$ to lie in Sobolev ball $\mathcal{W}^{\alpha, \infty}([-1,1]^p)$ with the smoothness level $\alpha$.
%We assume that for each group $l\in[L]$, $\f{l}$ has bounded second and fourth moments, and $\Del{l}$ has bounded second norm, measured over the target distribution $\QQ$. In practice, $\rm{supp}(\XQ)$, the support of target covariates, is often bounded; then the second and fourth norm of $\f{l}$ is also bounded by some constant. We also ask $\Del{l}$ has bounded second norm by $\delta_n$, which depends on the sample size $n$. Through our discussion, we assume $\delta_n$ decay to $0$ as sample size $n$ grows. It holds if the data-dependent model estimators $\left\{\hf{l}\right\}_{l\in[L]}$ could accurately approximate $\left\{\f{l}\right\}_{l\in[L]}$ with increasing sample size. Extensive well-known machine-learning methods ensure such consistency property, , etc. \Zijian{We need much more references than this and comment that they have rate of control on $\delta_n$.}
%For the convex set $\HH$ of $\Delta^L$, we define its diameter as
The following theorem establishes the rate of convergence for the plugin \texttt{DRL} estimator $\tfH$ defined in \eqref{eq: Magging plugin}.
\begin{Theorem} Suppose that the Assumptions \ref{ass: eigen} and \ref{ass: plugin} hold, then with probability larger than $1-{1}/{t^2}- \tau_n$ for some $t>1$ {and $\tau_n$ as in Assumption \ref{ass: plugin}},
\begin{equation}
\begin{aligned}
\left\|\tfH - \fH\right\|_{\QQ,2} &\lesssim \delta_n + t M \cdot \min\left\{\frac{1}{\eigenGamma}\left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \Err_0\right), \rho_\HH\right\}
\quad \textrm{with} \quad \Err_0 = M\delta_n,
\end{aligned}
\label{eq: error plugin}
\end{equation}
where $M$ and $\widetilde{M}$ are defined in \eqref{eq: f norm}, $\delta_n$ is defined in the Assumption \ref{ass: plugin}, and $
\rho_\HH := \max_{q, q' \in \HH} \|q - q'\|_2$ denotes the diameter of the set $\HH$.
\label{thm: Magging plugin}
\end{Theorem}
The convergence rate of $\left\|\tfH - \fH\right\|_{\QQ,2}$ decomposes of two parts: the first term $\delta_n$ on the right-hand side of \eqref{eq: error plugin} results from the estimation error of $\hf{l}$ while the second term $\min\left\{\frac{1}{\eigenGamma}\left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \Err_0\right), \rho_{\HH}\right\}$ comes from the error of estimating the optimal aggregation rate. The intuition of taking this minimum in the second term is as follows: when the diameter $\rho_{\HH}$ of $\HH$ is relatively large, the error of estimating the optimal weight results from the uncertainty of estimating $\Gamma$; however, when the diameter $\rho_{\HH}$ is smaller than the statistical precision of estimating the aggregation weight, the size of the constraint set $\HH$ determines the weight estimation error. Both the sample sizes of the source data $n$ and unlabelled target data $n_\QQ$ affect the final convergence rate of $\tfH$, where the sample size $n$ of the source data plays a role in affecting the accuracy of $\tfH$ through the term $\delta_n$.
Theorem \ref{thm: Magging plugin} establishes the consistency of $\tfH$ in various scenarios. For example, when $M,\widetilde{M}$ are of constant orders and $\lim_{n\rightarrow \infty}\delta_n=0$, the right-hand side of \eqref{eq: error plugin} converges to $0$ with $n$ and $n_\QQ$ growing to infinity.
We consider the important case $\delta_n\ll M\asymp \widetilde{M}$, essentially requiring $\hf{l}$ to be a consistent estimator and the second and fourth order norm of $\f{l}$ to have the same order. Such a special case is satisfied for various scenarios, including the high-dimensional regression and binary outcome models with consistent initial estimators. Then we can simply {write} the convergence rate as
\begin{equation}
\left\|\tfH - \fH\right\|_{\QQ,2} \lesssim \delta_n+ M\cdot \min\left\{M\delta_n+\frac{M^2}{\sqrt{n_{\QQ}}}, \rho_\HH\right\}.
\label{eq: sim plugin}
\end{equation}
When $M$ is relatively large such that $M\gg 1$, the second term is the dominating term in the above rate of convergence. We will show in the following Theorem \ref{thm: Magging correct} that the second term, especially the term $M^2\delta_n$, can be significantly reduced with our proposed bias-corrected estimator. As one example, we have that $M = \max_{l \in [L]}\|\beta^{(l)}\|_2$ in high-dimensional linear models. It is then conceivable that $M$ grows, and Corollary \ref{cor: high-dim} discusses our distributional robust learning in this context. %\Peter{I do not follow this 100\%. Are you saying: because $M = max_{l} \|\beta^{(l)}\|_2$ in high-dimensional linear models, it typically means that $M$ grows? Or is it more the following:\\
%}%\Zijian{I have implemented}
%\Peter{OK, it is good. But: I do not understand the last sentence. Corollary 2 does not say that $M$ is growing and it does not show it either. Do you mean the following.\\
%As an example, in high-dimensional linear omdels we have that $M = \max_{l \in [L]}\|\beta^{(l)}\|_2$ and Corollary \ref{cor: high-dim} discusses our distributional robust learning in this context.}
%\Peter{But isn't the canonical case where $M$ is of order $O(1)$?}\Zijian{That is true. As one case, $M$ can grow in high dimensions as it is of order $\|\beta\|_2.$ Here is a bit my belief: we can always let $M$ to grow (even in relatively low dimensional setting), which minic a scenario that $\|\f{l}\|_{\QQ,2}$ to be relatively large. This is more about which regime better approximates the finite sample data.} \Peter{Comment: If $\|\beta\|_2$ grows in a high-dim. linear model we are in a regime with increasing signal to noise ratio! I think this is generally the case: if $\|f\|_{Q,2}$ grows we have ncreasing signal to noise ratio -- which is unreasonable for doing asymptotics} \Peter{Perhaps we should just talk about the "second term" instead of the "dominating term". The second terms is about learning the aggregation, correct?} \Zhenyu{Correct, the second term is about the error of $\Gamma$ matrix estimation, and then we derive it to the aggregation weights error.}
%As , . Thus, the plugin \texttt{DRL} estimator $\tfH$ is consistent with $\fH$. The term $\Err_0$ arises from the ``Bias-kl'' and ``Bias-lk'' terms in \eqref{eq: error decomposition}. The potential issue for the plugin \texttt{DRL} estimator $\tfH$ occurs when the second norm scale $M$ is large enough and $\delta_n$ decays at a slow rate. Then $\Err_0$ will dominate the right-hand side of \eqref{eq: error plugin}, and the whole approximation error of $\tfH$ will decline at a rate of $O(M\delta_n)$.
% \Zhenyu{
% In the equation above, the term ${\widetilde{M}^2}/{\sqrt{n_\QQ}}$ arises from the last term in \eqref{eq: tGamma error decompose} that is left uncorrected. The term $\delta_n^2$ corresponds to the ``higher order bias'' term in the error decomposition \eqref{eq: error decomposition}, which typically decays fast. Additionally, $\Err_0$ is generated by the ``Bias-kl'' and ``Bias-lk'' terms in \eqref{eq: error decomposition}.
% We assume that the data-dependent model estimators $\left\{\hf{l}\right\}_{l\in[L]}$ could accurately approximates $\left\{\f{l}\right\}_{l\in[L]}$ as the sample size $n$ grows, such that $\delta_n \rightarrow 0$. The entire upper bound in \eqref{eq: error plugin} decays. As a result, $\tfH$ converges to $\fH$ consistently.
% .
% }
% \Zijian{1. Converge; 2. which term might have a potential issue?}
% {\color{red} 1. Converge as $M$ bounded, $\delta_n$ goes to $0$. 2. $\Err_0$ comes from.}
%\Zijian{It is better to use definitions instead of assumptions...}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Rate improvement with the bias correction}
\label{sec: theory-correct}
In the following, we switch to the bias-corrected {\DRL} estimator $\hfH$ and demonstrate how the bias correction step improves the rate of convergence. To characterize the results, we define the following rates of convergence for estimating $\f{l}$ and the density ratio $\w{l}$.
\begin{Assumption}
% \Zijian{This is not correct. We need a low bound for the density ratio. the absolute value is not useful here as the ratio is always positive?}
There exist constants $C$ such that for any $l\in [L]$ and $x\in \R^p$, $0 < \w{l}(x)\leq C$. With probability larger than $1-\tau_n$ with $\tau_n\rightarrow 0,$ there exist positive sequences $\widetilde{\delta}_n,\eta_{\omega},$ and $\widetilde{\eta}_{\omega}$ such that for any $l\in [L]$, $\max\left\{\|\Del{l}\|_{\QQ,4},\|\Del{l}_\A\|_{\QQ,4},\|\Del{l}_\B\|_{\QQ,4}\right\}\leq \widetilde{\delta}_n$,
\vspace{-4mm}
\[
\max\left\{\|\hw{l}_\A-\w{l} \|_{\QQ, 2}, \|{\hw{l}_\B}-{\w{l}}\|_{\QQ, 2}\right\}\leq \eta_{\omega},\quad \max\left\{\|\hw{l}_\A-\w{l} \|_{\QQ, 4}, \|{\hw{l}_\B}-{\w{l}}\|_{\QQ, 4}\right\}\leq \widetilde{\eta}_{\omega}.
\]
\label{ass: correct}
\end{Assumption}
% {\small
% \begin{equation*}
% \max_{l\in [L]} \max\left\{\|\hw{l}_\A-\w{l} \|_{\QQ, 2}, \|{\hw{l}_\B}-{\w{l}}\|_{\QQ, 2}\right\}\leq \eta_{\omega},\quad \textrm{and}\quad
% \max_{l\in [L]} \max\left\{\|\hw{l}_\A-\w{l} \|_{\QQ, 4}, \|{\hw{l}_\B}-{\w{l}}\|_{\QQ, 4}\right\}\leq \widetilde{\eta}_{\omega}
% \end{equation*}
% }
% As a remark, when $\w{l}(\cdot)$ is bounded such that $\sup_{x\in \R^p}|\w{l}(x)|\leq c$ for some universal constant $c$, it is established that $\left\|\frac{\hw{l}_\A}{\w{l}} - 1\right\|_{\QQ, q}\asymp \left\|\hw{l}_\A - \w{l}\right\|_{\QQ, q}$ and $\left\|\frac{\hw{l}_\B}{\w{l}} - 1\right\|_{\QQ, q}\asymp \left\|\hw{l}_\B - \w{l}\right\|_{\QQ, q}$ for $q\geq 1$.
\vspace{-8mm}
% \Zijian{The tone of the following discussion is not right. We assume finiteness not for simplicity, right? It is essential for our method and theory to work. Wainwright and coauthors have papers to mainly talk about this. We should say this is a strong assumption but commonly used in the literature and then cite some literatures.}
{In the first part of Assumption \ref{ass: correct}, we consider the boundedness condition of $\omega(x)$ for any value of $x$, which is commonly used in the literature. For example, it can be derived from the Positivity/Overlap assumption \citep[e.g.]{rosenbaum1983central,chen2008semiparametric} or
%be
directly assumed,
%in the works of density ratio estimations,
see e.g. Assumption 13.1 and 14.7 in the book \citep{sugiyama2012density}.
In
%the
Section \ref{sec: proof-thm-magging-correct} of the supplementary material, we will present a relaxed version for this assumption without requiring the boundedness.}
% As a remark, we assumed the finiteness condition for density ratios $\w{l}(\cdot)$ for simplicity; see Appendix \Zhenyu{Ref} for a relaxed version for this assumption. In addition, this finiteness property is widely used in literature, which can deduced from the Overlap Assumption in works \citep[e.g.]{chen2008semiparametric,d2021overlap} or the Positivity Assumption in works \Zhenyu{Ref}.
Similarly to
%the
Assumption \ref{ass: plugin}, we focus on the setting where
%that
$\widetilde{\delta}_n\rightarrow 0$.
% \Peter{What does the following half sentence mean... Why do we write it here?} although our result is also applied to the setting with non-vanishing $\widetilde{\delta}_n.$
We still take the high-dimensional sparse regression $\f{l}(x)=x^{\intercal}\beta^{(l)}$ for $l\in [L]$ as the illustrative example. In this setting, we establish in the following Corollary \ref{cor: high-dim} that $\widetilde{\delta}_n\lesssim \sqrt{s_\beta \log p/n}$ for sub-Gaussian covariates $X^{\QQ}_{i}$.
Regarding the density ratio estimators $\hw{l}_\A$ and $\hw{l}_\B$, we do not impose any convergence condition on $\eta_{\omega}$ and $\widetilde{\eta}_{\omega}$. Our theoretical analysis covers consistent and inconsistent density ratio estimators, and we will elaborate on the consequences of our proposal for both scenarios. For the special setting with no covariate shift, we have $\hw{l}_\A(\cdot)=\hw{l}_\B(\cdot)=\w{l}(\cdot)=1$ and then have $\eta_{\omega}=\widetilde{\eta}_{\omega}=0$. In the more general setting allowing for covariate shifts, we can further establish the convergence rate of $\eta_{\omega}$ and $\widetilde{\eta}_{\omega}$ when the density ratio model in \eqref{eq: density ratio} is correctly specified.
For example, we will establish in the following Corollary \ref{cor: high-dim} that
$
\max\{\eta_{\omega},\widetilde{\eta}_{\omega}\}\lesssim \sqrt{{s_{\gamma}\log p}/{(n+n_{\QQ})}}$ for the high-dimensional logistic density ratio model with the maximum sparsity $s_{\gamma} = \max_{l\in [L]}\|\gamma^{(l)}\|_0.$ For the nonparametric density estimation, \citet{nguyen2010estimating} leveraged the non-asymptotic characterization of f-divergence and proposed the optimal density ratio estimator in the minimax sense with the rate $\eta_\omega \asymp O(n^{-1/(2+\nu)})$, where $0<\nu<2$ denotes the complexity of the prespecified function class $\mathcal{G}$ containing the true density ratio; see Theorem 3 in \citet{nguyen2010estimating} for details.
More generally, \citet{sugiyama2012density} thoroughly review the consistency results of different density ratio estimators. %employing different density ratio estimations that have the property of consistency.
% \Zijian{what kind of the density ratio estimator? this discussion still does not read well.}\Zhenyu{Rewrite} \Zijian{Zhenyu, please double check this after you read the covariate balancing literature.}\Zhenyu{I will add covariate balancing rate later.}
% Particularly, for the nonparametric case, \citet{nguyen2010estimating} establishes an optimal estimator in the minimax sense with the rate $\eta_\omega \asymp O(n^{-1/(2+\nu)})$ where $\nu$ determines the smoothness of the function class and $0<\nu < 2$. \Zijian{What is $\nu$ here? it denotes some smoothness?}\Zhenyu{Yes, smoothess of the function class}
%\Zijian{Zhenyu, I think we need to rewrite a bit the following sentence, I do not fully get it.}\Zhenyu{Done}
% various studies employing different density ratio estimations have demonstrated the property of consistency \citep{sugiyama2012density}. This results in both $\eta_\omega$ and $\widetilde{\eta}_\omega$ approaching $0$ as the sample size $n$ and $n_\QQ$ grow. For example, in the \citet{nguyen2010estimating} establishes an optimal estimator in the minimax sense with the rate $\eta_\omega \asymp O(n^{-1/(2+\nu)})$ where $\nu$ determines the complexity of the function class and $0<\nu < 2$.
% \Zijian{Zhenyu, for the following sentence, we need more detailed discussion.}
%In addition to the assumption \ref{ass: plugin}, we further require $\Del{l}$ to have a bounded fourth norm, controlled by the term $\widetilde{\delta}_{n}$, which converges to $0$ in our analysis. For the covariate shift setting when $\QQ_X\neq \PP{l}_X$ for any $l\in[L]$, we assume that the $\hw{l}/\w{l} - 1$ has bounded second and fourth norm, controlled by $\eta_\omega$ and $\widetilde{\eta}_\omega$ respectively.
We now state the convergence rate for the bias-corrected estimator $\hfH$.
\begin{Theorem} Suppose that Assumptions \ref{ass: eigen}, \ref{ass: plugin}, \ref{ass: correct} hold, then with probability larger than $1-{1}/{t^2} - 2\tau_n$ for $t>1$ {and $\tau_n$ as in Assumptions \ref{ass: plugin} and \ref{ass: correct}},
\begin{equation}
\begin{aligned}
\left\|\hfH - \fH\right\|_{\QQ,2} &\lesssim \delta_n + t M\cdot \min\left\{\frac{1}{\eigenGamma}\left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \Err_1 + \Err_2\right), \rho_\HH\right\}\\
\textrm{with}\quad \Err_1 &= \frac{(\widetilde{M}+\widetilde{\delta}_n)\widetilde{\delta}_n}{\sqrt{n \wedge n_\QQ}} + \frac{M+\delta_n}{\sqrt{n}} \quad \text{and} \quad
\Err_2 = (\widetilde{M}+\widetilde{\delta}_n)\widetilde{\delta}_n \eta_\omega + \frac{(\widetilde{M}+\widetilde{\delta}_n)\widetilde{\eta}_\omega}{\sqrt{n}},
\end{aligned}
\label{eq: error correct}
\end{equation}
where $M$ and $\widetilde{M}$ are defined in \eqref{eq: f norm}, $\delta_n$ is defined in the assumption \ref{ass: plugin}, $\widetilde{\delta}_n,\eta_\omega$ and $\widetilde{\eta}_\omega$ are defined in the assumption \ref{ass: correct}, and $
\rho_\HH := \max_{q, q' \in \HH} \|q - q'\|_2$ denotes the diameter of the set $\HH$.
\label{thm: Magging correct}
\end{Theorem}
In comparison to Theorem \ref{thm: Magging plugin}, the above theorem replaces the term $\Err_0$ in \eqref{eq: error plugin} by $\Err_1 + \Err_2$ in \eqref{eq: error correct}, where $\Err_1$ stands for the remaining error after the bias-correction in the regime of known density ratio and $\Err_2$ stands for the estimation error of the density ratio. We consider a few important regimes in the following and illustrate that $\Err_1 + \Err_2$ is typically not larger than $\Err_0.$ To start with, we consider the no covariate shift setting and establish the following corollary by applying Theorem \ref{thm: Magging correct} with $\eta_\omega=0$ and $\widetilde{\eta}_\omega=0.$
\begin{Corollary}
Suppose that Assumptions \ref{ass: eigen}, \ref{ass: plugin}, \ref{ass: correct} hold and $\PP{l}_X = \QQ_X$ for $l\in [L]$. Then with probability larger than $1-{1}/{t^2}-2\tau_n$ for $t>1$ {and $\tau_n$ as in Assumptions \ref{ass: plugin} and \ref{ass: correct}}:
\begin{equation*}
\begin{aligned}
\left\|\hfH - \fH\right\|_{\QQ,2} \lesssim \delta_n + t M\cdot \min\left\{ \frac{1}{\eigenGamma}\left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \Err_1 \right), \rho_\HH\right\}\;
\textrm{ with } \Err_1 = \frac{(\widetilde{M}+\widetilde{\delta}_n)\widetilde{\delta}_n}{\sqrt{n \wedge n_\QQ}} + \frac{M+\delta_n}{\sqrt{n}}.
\end{aligned}
\label{eq: error no shift}
\end{equation*}
\end{Corollary}
Similarly to \eqref{eq: sim plugin}, we consider an important regime $\delta_n\asymp \widetilde{\delta}_n \ll M\asymp \widetilde{M}$ and simplify the above corollary as
\begin{equation}
\left\|\hfH - \fH\right\|_{\QQ,2} \lesssim \delta_n+ M\min\left\{M\delta_n\cdot r_n+\frac{M^2}{\sqrt{n_{\QQ}}}\;, \;\rho_\HH\right\} \quad \text{with}\quad r_n=\frac{1}{\sqrt{n \wedge n_\QQ}}+\frac{\delta_n}{M}+\frac{n^{-1/2}}{\delta_n}.
\label{eq: error no shift}
\end{equation}
By comparing the above rate to \eqref{eq: sim plugin}, we observe that the bias-correction step shrinks the term $M^2\delta_n$ by a factor of $r_n$. The ratio $r_n$ diminishes to zero when the machine learning estimator has a convergence rate slower than the parametric rate $n^{-1/2}$, which is the case for high-dimensional regression, random forests, and deep neural networks. Through comparing the rates in \eqref{eq: sim plugin} and \eqref{eq: error no shift}, we observe %that the plug-in estimator and the bias-corrected estimator have the same convergence rate when $M$ is of a constant scale; however,
that when $M\gg 1$, the bias correction step improves the convergence rate of the plug-in estimator.
We now consider the more general covariate shift setting. To illustrate the main idea, we consider the important regime $\delta_n\asymp \widetilde{\delta}_n \ll M\asymp \widetilde{M}$, $\widetilde{\eta}_\omega\asymp {\eta}_\omega$ and ${n^{-1/2}}\lesssim{\delta_n}$ and simplify \eqref{eq: error correct} as
\begin{equation}
\left\|\hfH - \fH\right\|_{\QQ,2} \lesssim \delta_n+ M\min\left\{M\delta_n\cdot (r_n+\eta_{\omega})+\frac{M^2}{\sqrt{n_{\QQ}}}\;, \;\rho_\HH\right\}.
\label{eq: shift simplified}
\end{equation}
By comparing the above rate to the no covariate shift rate in \eqref{eq: error no shift}, we notice that there is an additional component $\eta_{\omega}$ in the multiplicative factor in front of the term $M^2\delta_n$. This additional piece results from the uncertainty of estimating the density ratio. {Similarly to \eqref{eq: error no shift}, the bias correction mainly improves the rate of the plug-in estimator for the regime with a sufficiently large $M$ such that $M\gg 1.$ For $M\gg 1$,} we consider that the machine learning estimator has a convergence rate slower than $n^{-1/2}$ and then compare the convergence rate in \eqref{eq: shift simplified} to that in \eqref{eq: sim plugin}. We observe an interesting phenomenon: if the density ratio estimator is consistent, the bias-corrected estimator has a faster convergence rate than the plug-in estimator; even if the density ratio estimator is inconsistent but the estimation error $\eta_{\omega}$ is of a constant scale, the convergence rate of the bias-corrected estimator is not worse than that of the plug-in estimator. This demonstrates the robustness of our proposed bias-corrected estimator with respect to
%the
misspecification of the density ratio model. {We explore the numerical performance of the bias-corrected estimator in Section \ref{sec: simu-cs} and demonstrate in the following Figures \ref{fig: mix-L4-cs} and \ref{fig: air-H} that the bias-corrected estimator produces a more accurate estimator than the naive plug-in estimator.}
%\Zijian{Peter, we have addressed your following questions in the above paragraph. Please double check.}
%\Peter{If $M$ is of order $O(1)$, then the rates of the naive and bias-corrected estimator are both $O(\delta_n)$ -- correct?} \Peter{In fact: in the simulations, was the bias-corrected estimator really better than the naive plug-in? At first sight, I could not see this.}\Zhenyu{Through my simulations, I've identified the improvement only for the covariate shift setting, and when the sample size is sufficiently large.} \Peter{Zhenyu: do you mean by sample size $n_Q$? Or the training sample size? Where can we see this?}\Zhenyu{Peter, I mean by training sample size $\{n_l\}_{l\in [L]}$. You can check the Figure \ref{fig: mix-L4-cs} in the simulation section \ref{sec: simu-cs}. The section is covariate shift setting, and we vary training sample size $n_l$ but keep $n_\QQ$ fixed.} \Peter{I see... I made some comments in the figure caption.}
Previously, we have quantified the estimation error of $\hfH$ to the actual $\fH$. In the following theorem, we will quantify their difference in terms of the reward.
\begin{Theorem}
For any target distribution $\QQ$ with the true conditional outcome model $f^\QQ$, we have
\[
\left|\RR_\QQ(\hfH)-\RR_{\QQ}(\fH) \right|\leq 2\|f^\QQ-\fH\|_{\QQ,2} \cdot \|\hfH - \fH\|_{\QQ,2} + \|\hfH - \fH\|_{\QQ,2}^2,
\]
with the reward $\RR_{\QQ}(\fH)$ defined in \eqref{eq: reward}.
\label{prop: reward-diff}
\end{Theorem}
The convergence rate of the reward difference $\left|\RR_\QQ(\hfH)-\RR_{\QQ}(\fH)\right|$ relies on the convergence rate of $\|\hfH - \fH\|_{\QQ,2}$. If $\hfH$ is a consistent estimator of $\fH$ and $\|f^\QQ-\fH\|_{\QQ,2}$ is of a constant scale, the reward of our proposed bias-corrected estimator converges to the reward of its population counterpart $\fH$. Meanwhile, Theorem \ref{prop: reward-diff} also applies to the plugin DRL by replacing $\hfH$ with $\tfH$. As we have discussed previously, the bias-corrected estimator $\hfH$ will improve the rate of convergence in certain circumstances compared to the plugin estimator $\tfH$. Importantly, when the target population is sufficiently different from the {pooled} multi-source data, the reward $\RR_\QQ(\hfH)$ is expected, due to its definition, to outperform the reward given by the ERM. In numerical studies, we demonstrate this reward improvement in Section \ref{sec: simu-varyL} and also explore the effect of $\mathcal{H}$ on $\RR_\QQ(\hfH)$ in Section \ref{sec: simu-H}.
A major difference between the reward convergence in Theorem \ref{prop: reward-diff} and the existing DRO literature results \citep{namkoong2017variance,duchi2021learning} is that we establish the reward convergence for any given target population while they studied the convergence of the adversarial loss (or reward); see for example Corollary 3.2 in \citet{namkoong2017variance} and Corollary 2 in \citet{duchi2021learning}. Different from their analysis, we consider the {case where}
%data with
the group structure {is} known a-priori. In this setting, we leverage the closed-form solution obtained in Theorem \ref{thm: identification} and control the reward convergence rate by the estimation error $\|\hfH - \fH\|_{\QQ,2}$. %\Peter{It's good. I commented out a small part.}
%then we deduce the estimation error in the reward level.
%The existing DRO literature also established the convergence rate of the reward of the distributionally robust estimator t quantify the difference of reward/loss, such as the Corollary 3.2 in \citet{namkoong2017variance} and the Corollary 2 in \citet{duchi2021learning}.
%\Peter{Quite interesting! This is not when using ${\cal H}$, correct? In such a case, the radius $\rho_H$ would come into play as well? (It does come into play via the $\|.\|_{Q,2}$ norm). Should this be a Theorem rather than a Proposition?}\Zijian{I have revised the statement a bit. This theorem basically says that $\hfH$ converges to $\fH$ from the "reward" level. It is the typical result from the literature. Our previous results are a bit more refined in the sense that theorem 4 is more like a corollary. I think the dependence on H should be through $\|\hfH - \fH\|_{\QQ,2}^2$ since the above derivation is almost tight. Any other insights why $\HH$ may play other roles?} \Peter{OK, it's good.}
%\Zijian{Peter, the above theorem still does not answer the following questions: why $\RR_\QQ(\hfH)$ is better than $\RR_\QQ(\hat{f}^{\rm ERM})?$ From Zhenyu's exploration, this question is not addressed too much in theory; most literature results look into numerical comparison. } \Peter{OK. But isn't it clear by definition that $R_Q(f^*_H) \ge R_Q(f^*_{ERM})$ (these are the population versions) if $Q$ is contained in ${\cal H}$? And there must be a gap with a strict $>$ if $Q$ is sufficiently different from the pooled training data (and to capture this, ${\cal H}$ as must have a certain size (it cannot be too small)). I would not make it rigorous... but to me, it sounds plausible that there must be clear gain over ERM in certain settings.}\Zhenyu{Peter: I think ``$R_Q(f^*_H) \ge R_Q(f^*_{ERM})$'' is not correct. Imagine a setting that source mixture is very close to the target $\QQ$, then ERM should be the better one. In the literature, some work gives a concrete example to claim that there exists a distribution that $\fH$ is better than ERM, but they do not provide a result in the general sense.} \Zijian{I have added some of the discussions in the above paragraphs. please double check.} \Peter{Looks all fine now!}
%when the rate of , which causes $r_n$ to vanish, we can observe that as long as $\eta_\omega$ approaches to zero, the bias-correction step improves the rate. Even if the density ratios are poorly estimated with $\eta_\omega \asymp O(1)$, the rate after bias correction remains the same as the plugin rate in \eqref{eq: sim plugin}.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Example: high-dimensional setting}
\label{sec: ex-highd}
In this subsection, we consider the high-dimensional setting as an illustrative example of our main theorems. Particularly, we directly justify the rate of convergences assumed in the assumptions \ref{ass: plugin} and \ref{ass: correct}. We assume that the conditional outcome models in \eqref{eq: source-models} are linear models with $\f{l}(x) = x^\intercal \beta^{(l)}$ for $\beta^{(l)} \in \R^p$ and the density ratio model in \eqref{eq: density ratio} is correctly specified.
% The high-dimensional vectors $\{\beta^{(l)}\}_{l\in [L]}$ and $\{\gamma^{(l)}\}_{l\in [L]}$ are assumed to be sparse with highest sparsity level $s_\beta = \max_{l\in [L]}\|\beta^{(L)}\|_0$, and $s_\gamma = \max_{l\in [L]}\|\gamma^{(l)}\|_0$, respectively.
We briefly review the estimation of $\f{l}$ and the density ratio and then present the technical assumptions. For each group $l\in [L]$, we use $\widehat{\beta}^{(l)}$, $\widehat{\beta}^{(l)}_\A$, and $\widehat{\beta}^{(l)}_\B$ to respectively denote the LASSO estimators denote the estimators fitted with the full data and the subset data belonging to $\A_l$ and $\B_l$; see the exact formulas in Section \ref{sec: add_methods} in the supplementary materials.
% defined as
% \[
% \widehat{\beta}^{(l)} = \argmin_{\beta \in \R^p}\left\{\frac{1}{2n_l} \sum_{i=1}^{n_l} (\Y{l}_i - X^{(l)\intercal}_{i} \beta)^2 +
% A\sqrt{\frac{\log p}{n_l}} \sum_{j=2}^p \frac{\|\X{l}_{\cdot j}\|_2}{\sqrt{n_l}} |\beta_j| \right\},
% \]
% where the first column of $X^{(l)}$ is set as the constant 1, $X^{(l)}_{\cdot j}$ denotes the $j$-th column of $X^{(l)}$, and $A>\sqrt{2}$ is a pre-specified constant.
Under regularity conditions, the LASSO estimators have been shown to satisfy the following rate of convergence: with probability greater than $1-p^{-c}-\exp(-c n_l)$ for some positive constant $c>0$,
\begin{equation}
\max\left\{\|\widehat{\beta}^{(l)} - \beta^{(l)}\|_2, \|\widehat{\beta}^{(l)}_\A - \beta^{(l)}\|_2, \|\widehat{\beta}^{(l)}_\B - \beta^{(l)}\|_2\right\} \lesssim \sqrt{s_\beta \log p/n_l};
\label{eq: beta-rate}
\end{equation}
%the convergence rate for $\|\widehat{\beta}^{(l)}-\beta^{(l)}\|_2$ has been established,
see Theorem 7.2 in \citet{bickel2009simultaneous} and Lemma 6.10 in \citet{buhlmann2011statistics}. %and Corollary 2 in \citet{negahban2012unified}. %Our subsequent discussion follows their results that the initial estimators $\{\widehat{\beta}^{(l)}\}_{l\in [L]}$ satisfy the consistency property with the rate:
% \Zijian{You need to say that which papers show that the lasso estimator satisfy the condition.}Other estimators can also be used including the Dantzig Selector \citep{candes2007dantzig}, the scaled Lasso \citep{sun2012scaled}, square-root Lasso \citep{belloni2011square}, and iterated Lasso \citep{belloni2012sparse}, as long as the initial estimators $\widehat{\beta}^{(l)}$ satisfy the consistency property with the rate:
Similarly, for estimating the density ratio in high dimensions, we follow the literature results and assume that with a high probability greater than $1-p^{-c}-\exp(-c (n_l+n_{\QQ}))$ for some positive constant $c>0$,
the estimators $\widehat{\gamma}^{(l)}_\A, \widehat{\gamma}^{(l)}_\B$ satisfy
% defined in \eqref{eq: penalized MLE} satisfies %the following rate , \Zijian{Zhenyu, please specify what probability? }
% For the high-dimensional generalized linear models (GLM) setting, , and \citet{guo2021inference}'s equation 78, established the rate of convergence for $\widehat{\gamma}^{(l)}$:
\begin{equation}
\max\left\{\|\widehat{\gamma}^{(l)}_\A - \gamma^{(l)}\|_2, \|\widehat{\gamma}^{(l)}_\B - \gamma^{(l)}\|_2\right\} \lesssim \sqrt{s_\gamma \log p / (n+ n_\QQ)};
\label{eq: gamma-rate}
\end{equation}
see Section 4.4 of \citet{negahban2012unified} and Theorem 9 in \citet{huang2012estimation} for detailed statements.
We introduce the following assumptions for the high-dimensional setting. %\Zijian{Peter, please double check the following assumption and its related discussion.} \Peter{It is good.}
\begin{Assumption}
The target covariates observations $\{\XQ_j\}_{j\in [n_\QQ]}$ are i.i.d p-dimensional sub-Gaussian random vectors.
For every group $l\in [L]$, with probability larger than $1-p^{-c_0}$ for some positive constant $c_0>0$, the density ratio model $h(\cdot)$ defined in \eqref{eq: density ratio} satisfies $\min\{h(\widetilde{X}_{i}^{(l)\intercal}\gamma^{(l)}), 1-h(\widetilde{X}_{i}^{(l)\intercal}\gamma^{(l)})\}\geq c_{\rm min}$,
% for $1\leq i \leq n_l+n_{\QQ}$,
with $\widetilde{X}^{(l)} = ((X^{(l)})^{\intercal}, X_\QQ^\intercal)^\intercal$ and $c_{\rm min}\in (0,\frac{1}{2})$ denoting some positive constant.
% \begin{enumerate}[(a)]
% \item The initial estimators $\widehat{\beta}^{(l)}$ and $\widehat{\gamma}^{(l)}$ satisfy the consistency property with the rate as \eqref{eq: beta-rate} and \eqref{eq: gamma-rate}.
% \item
% \end{enumerate}
\label{ass: linear}
\end{Assumption}
%Condition 4 imposes the tail condition on the target covariates observations.
% Condition 4(b) requires the high-dimensional vectors $\{\beta^{(l)}\}_{l\in [L]}$ and $\{\gamma^{(l)}\}_{l\in [L]}$ being sufficiently sparse.
{Assumption} 4 requires the class conditional probability to be bounded away from $0$ and $1$, which is part of the assumption used to establish the consistency of the high-dimensional density ratio model. Such an assumption is commonly assumed for the theoretical analysis of the high-dimensional logistic regression; see condition (iv) of Theorem 3.3 of \citet{van2014asymptotically} and Assumption 6 in \citet{athey2018approximate}. %\Peter{Do we have here binary $Y$?} \Zijian{We do not focus on the binary $Y$ but may use the conditional mean model $\f{l}$ for the binary setting. I think this is a bit a confusion here. The binary logistic regression is not for the outcome model but for the density ratio model. Please double check my revision in Assumption 4.}
%\Zijian{Zhenyu, may you recall where we use such a condition? only for the rate of $\widehat{\gamma}$?}\Zhenyu{Yes, we need it for the assumption 3, $\widetilde{\eta}_\omega$}
% We introduce the following assumptions for the high-dimensional setting.
% \begin{Assumption}
% \begin{enumerate}[(a)] For every group $l\in [L]$:
% \item The $l$-th source covariates observations $\{\X{l}_i\}_{i\in [n_l]}$ are i.i.d $p$-dimensional sub-Gaussian random vectors with $\Sigma^{(l)} = \E(X^{(l)}_{i\cdot} X_{i\cdot}^{(l)\intercal})$ where $\Sigma^{(l)}$ satisfies $c_0\leq \lambda_{\rm min}(\Sigma^{(l)}) \leq \lambda_{\rm max}(\Sigma^{(l)}) \leq C_0$ for some positive constants $C_0\geq c_0>0$. The target covariates observations $\{\XQ_i\}_{i\in [n_\QQ]}$ are i.i.d $p$-dimensional sub-Gaussian random vectors as well.
% \item There exists a positive constant $c>0$ such that $\|\X{l}\|_\infty \lambda^{(l)}_{0,\beta}s_\beta \leq c$ and $\|\widetilde{X}^{(l)}\|_\infty \lambda^{(l)}_{0,\gamma}s_\gamma \leq c$, with $\lambda^{(l)}_{0,\beta} =\|\frac{1}{n_l}\sum_{i=1}^{n_l}\eps{l}_i \X{l}_i\|_\infty$ and $\lambda^{(l)}_{0,\gamma} = \|\frac{1}{n_l+n_\QQ}\sum_{i=1}^{n_l+n_\QQ}\widetilde{\epsilon}_i \widetilde{X}^{(l)}_i\|_\infty$, where $\eps{l}_i = \Y{l}_i - X^{(l)\intercal}_i\beta^{(l)}$ and $\widetilde{\epsilon}^{(l)}_i = G^{(l)}_i - h(\widetilde{X}^{(l)\intercal}_i\gamma^{(l)})$.
% \item With probability larger than $1-p^{-c}$, $\min\{h(X_{i\cdot}^{(l)\intercal}\gamma^{(l)}), 1-h(X_{i\cdot}^{(l)\intercal}\gamma^{(l)})\}\geq c_{\rm min}$ for $i\in [n_l]$, $l\in [L]$, and some small positive constant $c_{\rm min}\in (0,1)$, where $h(\cdot)$ is defined in \eqref{eq: density ratio}.
% \end{enumerate}
% \label{ass: linear}
% \end{Assumption}
% The above conditions are standard in high-dimensional linear and logistic regression analysis \citep{bickel2009simultaneous,van2014asymptotically,guo2021inference}. Condition 4(a) imposes the tail condition on the covariates observations for both source and target data and also requires that the population second-order moment matrices for source covariates are invertible. Condition 4(b) requires the high-dimensional vectors $\{\beta^{(l)}\}_{l\in [L]}$ and $\{\gamma^{(l)}\}_{l\in [L]}$ being sufficiently sparse. Condition 4(c) is imposed such that the class posterior probability is bounded away from $0$ and $1$, which is commonly assumed and convenient for the theoretical analysis of the high-dimensional logistic regression, see condition (iv) of Theorem 3.3 of \citet{van2014asymptotically} and the overlap assumption (Assumption 6) in \citet{athey2018approximate}. \Zhenyu{References added.}
We now compare the convergence rate for $\tfH$ and $\hfH$ in the high-dimensional linear setting.
\begin{Corollary}
Suppose that Assumptions \ref{ass: eigen} and \ref{ass: linear} hold, the sparsity levels $s_{\beta}:= \max_{l\in[L]}\|\beta^{(l)}\|_0$ and $s_{\gamma}:= \max_{l\in[L]}\|\gamma^{(l)}\|_0$ satisfy $\sqrt{s_\beta \log p/n} \longrightarrow 0$ and $\sqrt{s_\gamma \log p / (n+ n_\QQ)} \longrightarrow 0$, and the estimators $\{\widehat{\beta}^{(l)},\widehat{\beta}^{(l)}_\A, \widehat{\beta}^{(l)}_\B\}$ and $\{\widehat{\gamma}^{(l)}_\A, \widehat{\gamma}^{(l)}_\B\}$ satisfy \eqref{eq: beta-rate} and \eqref{eq: gamma-rate}, respectively. Then with probability larger than $1-p^{-c}-\exp(-cn)$ for some positive constant $c>0$,
$\max\{\delta_n, \widetilde{\delta}_n\}\lesssim \sqrt{s_\beta \log p/n}$ and $\max\{\eta_\omega, \widetilde{\eta}_\omega\} \lesssim \sqrt{s_\gamma \log p / (n+n_\QQ)}$.
Additionally, we assume that $M = \max_{l \in [L]}\|\beta^{(l)}\|_2 \gg \sqrt{s_\beta \log p /n}$,
then with probability larger than $1-p^{-c}-\exp(-cn)-{1}/{t^2}$ for some $c>0$ and $t>1$, the bias-corrected DRL estimator satisfies
\begin{equation}
\begin{aligned}
&\left\|\hfH - \fH\right\|_{\QQ,2}
\lesssim
\sqrt{\frac{s_\beta \log p}{n}} + tM\cdot\min\left\{
M\sqrt{\frac{s_\beta \log p}{n}} \cdot r_n + \frac{M^2}{\sqrt{n_\QQ}}\;,\; \rho_\HH\right\} \\
&\quad \textrm{where} \quad r_n =
\frac{1}{\sqrt{s_\beta \log p}} + \frac{1}{\sqrt{n\wedge n_\QQ}} + \frac{1}{M}\sqrt{\frac{s_\beta \log p}{n}} + \sqrt{\frac{s_\gamma \log p}{n+n_\QQ}}.
\end{aligned}
\end{equation}
In contrast, with probability larger than $1-p^{-c}-\exp(-cn)-1/{t^2}$ for some $c>0$ and $t>1$, the plugin DRL estimator satisfies
\begin{equation*}
\left\|\tfH - \fH\right\|_{\QQ,2} \lesssim
\sqrt{\frac{s_\beta \log p}{n}} + tM\cdot\min\left\{M\sqrt{\frac{s_\beta \log p}{n}} +\frac{M^2}{\sqrt{n_\QQ}}\;,\;\rho_\HH\right\}.
\end{equation*}
\label{cor: high-dim}
\end{Corollary}
By comparing the bias-corrected DRL and the plugin DRL, we observe that the bias-correction step shrinks the inner $\sqrt{s_\beta \log p/n}$ by a factor $r_n\rightarrow 0$. {We now point out a regime where the bias-corrected estimator improves the rate of convergence for the high-dimensional setting: the norm $M = \max_{l \in [L]}\|\beta^{(l)}\|_2$ diverges to infinity with $n$ and $p$, the diameter $\rho_{\HH}$ is larger than $M\sqrt{{s_\beta \log p}/{n}}$, and the sample size of the unlabelled data in the target domain is larger than $n M^2/(s_\beta \log p).$}
% We illustrate this improvement with numerical studies in Section \ref{sec: simu-cs}
%\Peter{Again, if $M$ is of order $O(1)$, there is no difference in rates.} \Zijian{Peter, double check this.}\Zhenyu{Zijian: should it be when $\rho_\HH $ larger than ...?}
% \begin{Corollary}
% For the high-dimensional linear models, LASSO or Dantzig selector constructs $\delta_n \asymp \widetilde{\delta}_n \lesssim \sqrt{s_\beta \log p/n}$ for sub-Gaussian covariates $X_i^\QQ$ with $s_\beta = \max_{l\in [L]} \|\beta^{(l)}\|_0$. We also assume that the density ratio convergence rate holds as \eqref{eq: density ratio convergence}, then
% \begin{equation}
% \left\|\hfH - \fH\right\|_{\QQ,2} \lesssim
% \sqrt{\frac{s_\beta \log p}{n}} + M^2\left[\left(\sqrt{s_\beta s_\gamma} + \frac{s_\beta}{M}\right) \frac{\log p}{n} + \frac{s_\beta \log p}{n\wedge n_\QQ} + \frac{M}{\sqrt{n_\QQ}}\right]
% \end{equation}
% In contrast, for the plugin estimator
% \begin{equation}
% \left\|\tfH - \fH\right\|_{\QQ,2} \lesssim
% \sqrt{\frac{s_\beta \log p}{n}} + M^2\left[\sqrt{\frac{s_\beta \log p}{n}} + \frac{M}{\sqrt{n_\QQ}}\right]
% \end{equation}
% \end{Corollary}
%Let's start by discussing the scenario without covariate shift, where for all $l\in [L]$, we have $\PP{l}_X = \QQ_X$. In this case, when we apply the no covariate shift algorithm without estimating the density ratio, $\Err_2$ in \eqref{eq: error correct} becomes zero. Consequently, the bias-corrected estimator $\hfH$ outperforms the plugin estimator $\tfH$, by replacing $\Err_0$ with $\Err_1$. As $\Err_1 \asymp \frac{\widetilde{M}\widetilde{\delta}_n}{\sqrt{n \wedge n_\QQ}} + \frac{M}{\sqrt{n}}$, \Zijian{The following sentence is wrong...} it exhibits an improvement of $O(n^{-1/2})$ compared to $\Err_0=M\delta_n$.
%\Zijian{Zhenyu, I am not going to edit
% the following as it requires your polishing first.} Now let's consider the covariate shift setting. The value of $\Err_2$ depends on the accuracy of the density ratio estimation, represented by $\eta_\omega$ and $\widetilde{\eta}\omega$. Typically, $\eta\omega$ falls between $O(n^{-1/2})$ and $O(1)$. In the best-case scenario, $\eta_\omega$ achieves the optimal rate of $O(n^{-1/2})$, resulting in $\Err_2 \asymp \frac{\widetilde{M}\widetilde{\delta}_n}{\sqrt{n}}$, which is a significant improvement compared to $\Err_0$. However, in the worst-case scenario, if the density ratio is poorly estimated, $\eta_\omega$ and $\widetilde{\eta}_\omega$ become constant. In this case, $\Err_2 \asymp \widetilde{M}\widetilde{\delta}_n$. Hence, $\Err_2$ exhibits the same rate as $\Err_0$.
%In conclusion, the bias-corrected estimator $\hfH$ consistently outperforms the plugin estimator $\tfH$ in terms of the $\ell_2$ norm. Only when the density ratios are poorly estimated with $\eta_\omega$ and $\widetilde{\eta}_\omega$ being constant, the performance of $\hfH$ reverts back to that of $\tfH$.
% {\color{red} For better illustration, $\delta_n$, $\tau_n$ converges to 0.} {\color{red} Is it possible that $\Err_2$ written in $\delta_n$}
% {\color{red} {After the bias correction, the rate is uniformly smaller. In general, $\Err_1$ is smaller than $\Err_0$. $M$,$\tilde{M}$ constant. Maybe we need to replace $n$ by $N$.}}
% {\color{red} Additionally, if $\hf{l}$ is bounded, then $\Err_2$ can be further sharpened into.}
% Comparing the equation above and equation \eqref{eq: error plugin}, the improvement of bias correction lies in replacing the term $\Err_0$ with $\Err_1$ and $\Err_2$. Let us first discuss the $\Err_1$ term. Even if $\tau_n = O(1)$, $\Err_1$ still has a rate of $O(n^{-1/2})$. Thus, the term $\Err_1$ has a rate strictly faster than $\Err_0$. The $\Err_2$ term arises from the density ratio estimation errors. Under no covariate shift setting, the term $\Err_2 = 0$ since there is no need to approximate the density ratios, i.e., $\tau_n = 0$. We now discuss the scenarios under the covariate shift setting. In the \emph{best case}, the data-dependent density ratio estimators $\left\{\hw{l}\right\}_{l\in[L]}$ are well-fitted, such that $\eta_\omega$ converges to $0$ with a rate of $O(n^{-1/2})$. Then the first term in $\Err_2$ matches the rate of the first term in $\Err_1$. The second term in $\Err_2$ is at least $O(n^{-1})$. Therefore, bias correction under the covariate shift setting yields the same rate as in the no covariate shift setting. In the \emph{worst case}, the data-dependent density ratio estimators are poorly fitted with $\eta_\omega = O(1)$. Then the rate of term $\Err_2$ is $O(\widetilde{M} \tau_n)$. Assuming that $\delta_n\asymp \tau_n$, and $M$ is comparable to $\widetilde{M}$, the worst case of bias-corrected $\hfH$ under covariate shift deteriorates to the plugin estimator. To conclude, the covariate shift bias-corrected estimator $\hfH$ has a convergence rate between that of the no covariate shift bias-corrected estimator and that of the plugin estimator. The specific rate depends on how accurate the density ratio estimators are.
% We will now consider the bias-corrected estimators $\hfH$ defined in \eqref{eq: Magging nocs} under no covariate shift setting.
% \begin{Theorem}
% Under no covariate shift setting where $\PP{l}_X = \QQ_X$ for any group $l\in[L]$, suppose that assumptions \ref{ass: eigen}, \ref{ass: plugin}, and \ref{ass: nocs} hold. Then with probability $1-\frac{c}{t^2}$
% \begin{equation}
% \begin{aligned}
% \left\|\hfH - \fH\right\|_{\QQ,2} &\lesssim \delta_n + t\cdot \frac{M}{\eigenGamma}\left(\delta_n^2 + \frac{M^2}{\sqrt{n_\QQ}} + \Err_1\right)\\
% \textrm{with}\quad
% \end{aligned}
% \label{eq: error nocs}
% \end{equation}
% \label{thm: Magging nocs}
% \end{Theorem}
% {\red Err1 is always smaller than Err0}
% The bias-corrected estimator $\hfH$ remains consistent under no covariate shift setting, provided that the fitted models $\left\{\f{l}\right\}$ are consistent and have vanishing $\delta_n$ and $\tau_n$ as $n\rightarrow \infty$. A comparison of the errors in \eqref{eq: error nocs} and \eqref{eq: error plugin} reveals that the improvement lies in the terms $\Err_1$ and $M\delta_n$. \Zijian{Let us go through the following writtings in the meeting. It is not easy for readers to understand.} As discussed previously, when $M$ is large and $\delta_n$ decreases at a slower rate, $M\delta_n$ dominates the error term for $\tfH$. However, after bias correction, $\Err_1$ reduces $\delta_n$ to a rate of $\frac{\delta_n}{\sqrt{n}} \vee \frac{\tau_n}{\sqrt{n\wedge n_\QQ}}$, which would be at least $O(\frac{1}{\sqrt{n}})$ rate provided that $n_\QQ\asymp n$ and $\tau_n$ decays as well.
% In the covariate shift setting, we introduce additional assumptions concerning the estimated density ratios $\left\{\hw{l}\right\}$.
% \Zijian{You can combine theorems 3 and 4 and state theorem 3 as a special case with $\Err_2=0$.}
% Compared to the no covariate shift setting, the approximation error of $\hfH$ in the covariate shift setting includes an additional term $\Err_2$, which arises from the errors in estimating the density ratios $\left\{\hw{l}\right\}$. \emph{In the best scenario}, where the estimated density ratios are well-fitted with a rate of $O(\frac{1}{\sqrt{n}})$, the first component in $\Err_2$ matches with the first component in $\Err_1$. And the second component in $\Err_2$ is dominated by the second component in $\Err_1$, provided that $\delta_n, \tau_n$ have similar convergence rates. Thus, the convergence rate of $\hfH$ is identical in both covariate shift and no covariate shift settings. \emph{In the worst case}, where the estimated density ratios are poorly fitted with a rate of $O(1)$, the rate of $\Err_2$ depends fully on $M\tau_n$, which is similar to the rate of $M\delta_n$ in the plugin Magging error in \eqref{eq: error plugin}. To summarize, if the density ratios are estimated accurately, the approximation error $\hfH$ under covariate shift is the same as that in the no covariate shift setting. However, if the density ratios are estimated poorly, the bias-corrected $\hfH$ under no covariate shift will converge at the same rate as $\tfH$ by the plugin estimators.
% \Zijian{I think the above discussion is enough with a polished discussion.}
% \Zhenyu{Need to mention doubly robust.} We will use the following Table to illustrate the approximation error of $\hfH$ in the covariate shift setting and its relation to the results obtained in the no covariate shift setting, as well as the performance of $\tfH$ without bias correction.
% To simplify discussion, we consider the case where $n_\QQ \asymp n$, $\tau_n \asymp \delta_n$ and the scale $M$ is sufficiently large. We analyze the component in the apprenthesis in these upper bounds for estimators. The discussion is built upon the assumption that $\delta_n$ converges to $0$.
% \begin{table}[htp]
% \centering
% \begin{tabular}{c|c|c}
% Estimators & the main term & \\
% Plugin & $M \delta_n$
% \end{tabular}
% \caption{Caption}
% \label{tab:my_label}
% \end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%% 5.0 %%%%%%%%%%%%%%%%%%%%%%%%%%
% \section{Robust Wasserstein Magging}
% \label{sec: rowmag}
% \Zijian{Should we put it to a separate paper?}
% \Zhenyu{This section deals with a more demanding scenario where both the target covariate observations and the target outcome observations are missing. Consequently, we cannot directly identify $\QQ_X$, which is required in the joint distribution class (uncertainty set for the target distribution $\QQ = (\QQ_X, \QQ_{Y|X})$) defined in \eqref{eq: region-Q H} and its associated DRL model in \eqref{eq: f_star H}. To tackle this challenge, we introduce a new robust model called WasDRL (Wasserstein DRL), which leverages the identification Theorem \ref{thm: identification} for DRL model.}
% %%%%%%%%%%%%%%%%%%%%%% 5.1 %%%%%%%%%%%%%%%%%%%%%%%%%%
% \subsection{Robustness to the target covariate distribution}
% Let $\PPall_X$ represent the covariate distribution of the entire source data, which is a combination of $L$ groups of covariates $\PPall = \sum_{l=1}^L p_l \cdot \PP{l}_X$ with mixture weight $p \in \Delta^L$. Here $p$ is a nuisance parameter that can be interpreted as the probability of drawing an observation from each group.
% If the target covariate distribution $\QQ_X$ cannot be identified, we assume that it is in some neighborhood around the source covariate distribution $\PPall_X$, such that
% \[
% \QQ_X \in \left\{\left. \TT_X \;\right\vert \; \mathcal{D}(\PPall_X, \TT_X)\leq \rho \right\}
% \]
% where $\mathcal{D}(\cdot,\cdot)$ is a metric measuring the distance between two probability distributions, and $\rho > 0$ is a hyperparameter that controls the degree of uncertainty. In our framework, we use the Wasserstein distance to measure the distance between two distributions, defined as
% \[
% \mathcal{D}_c(\TT_X, \PPall_X) := \inf_{\gamma \in \Gamma(\PPall_X, \TT_X)} \E_{\gamma} c(X, X^\prime)
% \]
% where $\Gamma(\PPall_X, \TT_X)$ is the set of all couplings of $\PPall_X$ and $\TT_X$, \Zijian{We need to define what is the meaning of all couplings. We need formulas!} and $c(\cdot, \cdot)$ is a pre-specified cost function.
% Rather than minimizing the second moment over $\QQ_X$, provided it is identifiable in the identification Theorem \ref{thm: identification}, we study the \emph{distributionally robust} problem:
% \begin{equation}
% \frow = \sum_{l=1}^L \qrow_l\cdot \f{l} \quad
% \textrm{with}\quad
% \qrow = \argmin_{q\in \HH} \sup_{\TT_X: \mathcal{D}_c(\PPall_X, \TT_X)\leq \rho} q^\intercal \Gamma^{\rm w} q
% \label{eq: f rowmag}
% \end{equation}
% where $\Gamma^{\rm w} = \E_{\TT_X}[\F(X)\F(X)^\intercal].$ \Zijian{Stop here for this subsection.}
% Let us revisit the identification Theorem \ref{thm: identification} which states that given a target covariate distribution $\QQ_X$, Magging $\fH$ is a weighted average of $L$ conditional mean models $\left\{\f{l}\right\}$, where the mixture weight $q^*$ minimize the second moment measuring over $\QQ_X$, i.e., $q^* = \argmin_{q\in \HH}\E_{\QQ_X}\left[(\sum_{l=1}^L q_l \cdot \f{l}(X))^2\right]$. In the RoWMagging model $\frow$ defined in \eqref{eq: f rowmag}, the linearly weighted average of $\left\{\f{l}\right\}_{l\in[L]}$ is still maintained, ensuring its interpretability, computation efficiency and other advantages as the Magging model. However the mixture weight $\qrow$ is chosen to minimize the worst second moment over the neighborhood around $\PPall_X$, which further strengthens the model's robustness against potential covariate shift between the target distribution and the source distribution.
% In our motivation scenarios, RoWMagging $\frow$ is robust to not only the conditional outcome shift between $\QQ_{Y|X}$ and $\left\{\PP{l}_{Y|X}\right\}_{l\in[L]}$ but also the covariate shift between $\QQ_X$ and $\PPall_X$, as long as the target covariate distribution $\QQ_X$ is within $\rho$ distance of the source covariate distribution $\PPall_X$.
% %%%%%%%%%%%%%%%%%%%%%% 5.2 %%%%%%%%%%%%%%%%%%%%%%%%%%
% \subsection{Identification}
% In the following, we present the identification of $\frow$ defined in \eqref{eq: f rowmag}
% \begin{Theorem}
% With the cost function $c(X, X^\prime) = \|\F(X)-\F(X^\prime)\|_2^2$, $\frow$ defined in \eqref{eq: f rowmag} is identified as
% \begin{equation}
% \begin{gathered}
% \frow = \sum_{l=1}^L \qrow_l \cdot \f{l} \quad \textrm{with} \quad
% \qrow = \argmin_{q\in \HH} \sqrt{q^\intercal \Gamma_\PPall q} + \sqrt{\rho}\|q\|_2\\
% \textrm{for }\; \Gamma_\PPall = \E_{\PPall_X}[\F(X)\F(X)^\intercal]
% \end{gathered}
% \label{eq: identification rowmag}
% \end{equation}
% \label{thm: identification rowmag}
% \end{Theorem}
% We make use of the dual property of the Wasserstein distance to convert the minimax objective in \eqref{eq: f rowmag} into a simpler convex optimization problem \eqref{eq: identification rowmag}. This objective consists of two terms: the first term is equivalent to the square root of the second moment measured over the source covariate distribution $\PPall_X$, i.e., $\left(\E_{\PPall_X}[(\sum_{l=1}^L q_l \cdot \f{l})^2]\right)^{1/2}$, and the second term is a penalty on the weights $q$, that is dependent on the hyperparameter $\rho$. When there is no covariate shift between the source and target distributions ($\rho = 0$), the RoWMagging model reduces to the Magging model, i.e., $\frow = \fH$. As $\rho$ increases, we allow for a more extensive range of potential target covariate distributions in the neighborhood around $\PPall_X$ to be accounted for since the shift level is greater. The second term gradually becomes dominant in the objective, and the RoWMagging model tends to make a more equally weighted average of the models $\{\f{l}\}_{l\in[L]}$. As $\rho$ approaches infinity, we must account for any arbitrary $\QQ_X$, and the best robust model is to average all of the models $\left\{\f{l}\right\}_{l\in[L]}$.
% It should be noted that in Theorem \ref{thm: identification rowmag}, the cost function $c(X, X^\prime)$ is designed to evaluate the difference between the two points after undergoing the $\F(\cdot)$ transformation, which is computed as the difference between $\F(X)$ and $\F(X^\prime)$, for \emph{generic functions} $\F$. This method of setting up the cost function in Wasserstein distance is subtle and not commonly used. However, in the case of linear models $\F$, we can define the cost function directly as $c(X,X^\prime) = \|X - X^\prime\|_2^2$, without applying the $\F(\cdot)$ transformation.
% If the conditional outcome models are linear such that for $l\in[L]$, $\f{l}(X) = X^\intercal \coef{l}$, then $\F(X) = \BB^\intercal X$, with $\BB = \left(\coef{1}, \coef{2}, ..., \coef{L}\right)$. We will now present a proposition that provides identification of the RoWMagging model $\frow$ for linear models.
% \begin{Proposition}
% If $\F$ are linear models and $c(X, X^\prime) = \|X - X^\prime\|_2^2$, $\frow$ defined in \eqref{eq: f rowmag} is identified as
% \begin{equation}
% \begin{gathered}
% \frow(X) = \sum_{l=1}^L \qrow\cdot X^\intercal \coef{l}\quad \textrm{with} \quad
% \qrow = \argmin_{q\in \HH} \sqrt{q^\intercal \BB^\intercal \Sigma_{\PPall}\BB q} + \sqrt{\rho}\sqrt{q^\intercal \BB^\intercal \BB q} \\
% \textrm{for }\; \Sigma_\PPall = \E_{\PPall_X}X X^\intercal
% \end{gathered}
% \label{eq: identification rowmag linear}
% \end{equation}
% \label{thm: identification rowmag linear}
% \end{Proposition}
% The only difference between the formulation of \eqref{eq: identification rowmag linear} in linear models and \eqref{eq: identification rowmag} in generic models is in the penalty term (last component). In \eqref{eq: identification rowmag}, the cost function is penalized directly on the weight $q$ to make the RoWMagging $\frow$ more equally weighted average of $\left\{\f{l}\right\}$, since the cost function undergoes the $\F$ transformation. On the other hand, in \eqref{eq: identification rowmag linear}, since the cost function is without the $\F$ transformation, the penalty term penalizes over the model level $\BB q$ (to ensure that the coefficients of the final model have a smaller second norm), which has similar effects to the Ridge penalty.
% %%%%%%%%%%%%%%%%%%%%%% 5.3 %%%%%%%%%%%%%%%%%%%%%%%%%%
% \subsection{Theoretical Justification}
% We provide a detailed explanation of notation used. The second and fourth norm of $f$ over the source distribution $\PPall$ are defined as $\|f\|_{\PPall,2} := \left(\E_{\PPall_X}[f(X)^2]\right)^{1/2}$ and $\|f\|_{\PPall,4} := \left(\E_{\PPall_X}[f(X)^4]\right)^{1/4}$, respectively. The sample size of source data is denoted as $n_{\PPall}$ with $n_{\PPall} = \sum_{l=1}^L n_l$.
% \Zijian{It is better that we just keep a generic model or linear model. This is already not the focus. }
% \subsubsection{Generic models}
% We start with the generic models $\left\{\f{l}\right\}$. The population version of the objective function given in Theorem \ref{thm: identification rowmag} is denoted as $\La(q)$ such that:
% \[
% \La(q) = \sqrt{q^\intercal \Gamma_{\PPall} q} + \sqrt{\rho}\|q\|_2
% \]
% with $\Gamma_{\PPall}$ defined in \eqref{eq: identification rowmag}.
% And for the sample version, we define the optimal aggregation weight $\qrowt$ given by
% \begin{equation}
% \qrowt = \argmin_{q\in \HH}\sqrt{q \tGamma_{\PPall} q} + \sqrt{\rho}\|q\|_2 \; \textrm{ for }\; \tGamma_{\PPall} = \frac{1}{n_{\PPall}} \sum_{i=1}^{n_{\PPall}}\hF(X_i)\hF(X_i)^\intercal
% \label{eq: weight rowmag plugin}
% \end{equation}
% With slightly abuse of notation, we modified assumption \ref{ass: plugin} originally formulated for the target distribution $\QQ$ to the source distribution $\PPall$.
% \begin{assumptionp}{\ref*{ass: plugin}$'$}
% Over the source distribution $\PPall$, $\left\{\f{l}\right\}$ have bounded fourth norm and $\left\{\Del{l}\right\}$ have bounded second norm such that
% \[
% \max_{l\in[L]}\|\f{l}\|_{\PPall, 4}\leq M;
% \quad
% \max_{l\in[L]}\|\Del{l}\|_{\PPall, 2}\leq \delta_n
% \]
% \label{ass: plugin rowmag}
% \end{assumptionp}
% \begin{Proposition}
% Under assumption \ref{ass: plugin rowmag}, with probability $1-\frac{c}{t^2}$:
% \begin{equation}
% \left[\La(\qrowt) - \La(\qrow)\right]^2 \lesssim t\cdot \left(\delta_n^2 + \frac{M^2}{\sqrt{n_{\PPall}}} + M\delta_n\right)
% \end{equation}
% where $\qrowt$ is defined in \eqref{eq: weight rowmag plugin} and $\qrow$ is defined in \eqref{eq: identification rowmag}.
% \end{Proposition}
% The above proposition aims to quantify the difference between the sample version weight $\qrowt$ and the population version weight $\qrow$. It reveals that as long as $\delta_n$ vanishes, $\qrowt$ is consistent on the objective level such that $\left|\La(\qrowt) - \La(\qrow)\right|$ decays to $0$.
% \subsubsection{Linear models}
% In the following analysis, we consider linear models where $\f{l}(X) = X^\intercal \coef{l}$ with coefficients $\coef{l}$, for $l\in[L]$. The estimators of these coefficients are denoted as $\left\{\hcoef{l}\right\}$, and $\hBB = \left(\hcoef{1}, \hcoef{2}, ..., \hcoef{L}\right)$.
% \begin{Assumption}
% The bias of estimated coefficients is bounded:
% \[
% \max_{l\in[L]} \left\|\hcoef{l} - \coef{l}\right\|_2 \leq \kappa_n
% \]
% \label{ass: coefs}
% \end{Assumption}
% The objective in Proposition \ref{thm: identification rowmag linear} is denoted as $\Lb(q)$ with
% \[
% \Lb(q) = \sqrt{q^\intercal \BB^\intercal \Sigma_{\PPall} \BB q} + \sqrt{\rho} \sqrt{q^\intercal \BB^\intercal \BB q}
% \]
% with $\Sigma_\PPall$ defined in \eqref{eq: identification rowmag linear}. The sample versioned weight $\qrowt$ is defined as
% \begin{equation}
% \qrowt = \argmin_{q\in \HH}\sqrt{q \hBB^\intercal \widetilde{\Sigma}_\PPall \hBB q} + \sqrt{\rho}\sqrt{q^\intercal \hBB^\intercal \hBB q}\; \textrm{ for }\; \widetilde{\Sigma}_{\PPall} = \frac{1}{n_\PPall}\sum_{i=1}^{n_\PPall} X_i X_i^\intercal
% \label{eq: weight rowmag linear plugin}
% \end{equation}
% \begin{Proposition}
% Under assumption \ref{ass: plugin rowmag} and \ref{ass: coefs}, with probability $1-\frac{c}{t^2}$:
% \begin{equation}
% \left[\Lb(\qrowt) - \Lb(\qrow)\right]^2 \lesssim t\cdot \left(\delta_n^2 + \frac{M^2}{\sqrt{n_{\PPall}}} + M\delta_n + \rho \kappa_n^2\right)
% \end{equation}
% where $\qrowt$ is defined in \eqref{eq: weight rowmag linear plugin} and $\qrow$ is defined in \eqref{eq: identification rowmag linear}.
% \end{Proposition}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%% 6.0 %%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Simulations}
In this section, we demonstrate the effectiveness of the proposed DRL via simulated data. The code for implementation of DRL can be accessed via \url{https://github.com/zywang0701/DRL-simus}. We consider the simulated data mixed with $L$ groups, each following the data generation mechanism described in \eqref{eq: source-models}. For the $l$-th group, we set the covariates distribution $\PP{l}_X = \N(0_5, I_5)$ and generate the noise $\eps{l}_i \distras{iid} \N(0,1)$ being independent of the covariates; unless otherwise stated, we set the conditional outcome model $\f{l}$ as
$\f{l}(x) = \sum_{1\leq j\leq 5} \alpha_j^{(l)} x_j + \sum_{1\leq j\leq k\leq 5} \beta_{jk}^{(l)} \left(x_j x_k - \E[x_j x_k]\right)$ where $\alpha_j^{(l)}$ and $\beta_{jk}^{(l)}$ represent the coefficients of covariates and interaction terms, respectively. We randomly generate the coefficients across the $L$ groups to approximate the heterogeneous conditional outcome models. Particularly, for $1\leq l\leq L$ and $1\leq j\leq k\leq 5$, the coefficients $\alpha_j^{(l)}$ and $\beta_{jk}^{(l)}$ are independently generated random variables, taking values in $\{0.4, 0.2, 0\}$ with probability 0.3, 0.4, and 0.3, respectively.
We use the notation $n_\PPall$ to represent the sample size of the entire source data, which is mixed from $L$ groups with the known mixture weight $q^{\rm sou} \in \Delta^L$. We generate the test data, with sample size $n_{\QQ}$, as a mixture of $L$ source populations with the mixture weight $q^{\rm tar} \in \Delta^L$. When $q^{\rm tar}$ is very different from $q^{\rm sou}$, the {\ERM} model fitted over the pooled source data sets might lead to poor prediction performance for the target data.
In the following, unless otherwise stated, we consider the no covariate shift setting and compare our proposed distributionally robust prediction models with the standard {\ERM},
\begin{itemize}
%\item ERM: fitting the model on the pooled source data.
\item {\DRLn}: the proposed \texttt{DRL} using Algorithm \ref{alg: cs mm} with default $\HH = \Delta^L$;
\item {\DRL}: the proposed \texttt{DRL} with a specified $\HH$, which will be detailed in the specific settings.
\end{itemize}
To illustrate the main idea, we use random forests as our major machine learning algorithms, implemented using the R package \texttt{Ranger} \citep{wright2015ranger}, with hyperparameters tuned through out-of-bag error. We also demonstrate our general DRL framework using deep neural networks and high-dimensional sparse regression models as the baseline machine learning algorithms.
In the reported numerical evaluations, we primarily focus on the empirical version of the reward defined in \eqref{eq: reward}, denoted as $\sum_{i=1}^{n_{\QQ}}\left[Y_i^2 - (Y_i - \widehat{f}(X_i))^2\right]/n_{\QQ}$, where $\widehat{f}(\cdot)$ represents the estimated conditional outcome model. A higher value of the evaluated reward indicates a better prediction performance. It is important to note that during the training step, we only have access to target covariate observations but not their outcomes. The target outcome observations are generated only for evaluation purposes but not for training prediction models. To provide a comprehensive assessment, we conduct 500 rounds of simulations to summarize the numerical performance.
%\Zijian{Zhenyu, we should not put the following paragraph here as we have not discussed how to generate thetesco target population.}
%%%%%%%%%%%%%%%%%%%%%% 6.1 %%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Robustness of {\DRLn} over {\ERM}}
\label{sec: simu-varyL}
We begin by examining the scenario with a varying number of groups, denoted by $L\in \{2,3,..., 9, 10\}$. The sample size for the source data is set as $n_\PPall = 2000 \cdot L$, while the sample size for the target data is $n_\QQ = 10,000$.
The target distribution $\QQ$ can be a mixture of any $L$ groups, that is $\QQ \in \C(\QQ_X)$ with $\C(\QQ_X)$ defined in \eqref{eq: region-Q full}. We evaluate the performance of different prediction models by calculating the worst-case reward across all possible target populations belonging to $\C(\QQ_X).$ In this case, the worst-case reward over $\C(\QQ_X)$ corresponds to the reward of the worst-performing group.
Our proposed Algorithm \ref{alg: cs mm} involves sample splitting. However, we believe sample splitting is mainly introduced to facilitate the theoretical justification. In Figure \ref{fig: mix-varyL}, we have implemented Algorithm \ref{alg: cs mm} and the no sample splitting version of Algorithm \ref{alg: cs mm}, which corresponds to Algorithm \ref{alg: cs mm} by replacing $\hf{l}_\A, \hf{l}_\B$ with $\hf{l}.$ In Figure \ref{fig: mix-varyL}, we observe that Algorithm \ref{alg: cs mm} and its no sample splitting version lead to a similar predictive performance. In the remaining of numerical studies, we focus on implementing Algorithm \ref{alg: cs mm} with no sample splitting.
%splitting performs even better in the experiments, see Figure . For the non-splitting method, we replace $\hf{l}_\A, \hf{l}_\B$ (based on the split data $\A_l,\B_l$) with $\hf{l}$ (based on the full data) for each $l\in [L]$. From this point onward, the terms {\DRLn} or {\DRL} will denote the methods without sample splitting.
% The comparison of versions with/without sample splitting is presented in Appendix \ref{sec: simu-varyL-split}
\begin{figure}[htp!]
\centering
\includegraphics[width=0.8\textwidth]{assets/plot-mix-varyL-split.pdf}
\caption{
Comparison of worst group reward for \texttt{DRL0-nosplit}, \texttt{DRL0-split}, and \texttt{ERM}, with the number of groups $L$ varied across $\{2,3,...,10\}$. Both \texttt{DRL0-split} and \texttt{DRL0-nosplit} are implemented by Algorithm \ref{alg: cs mm} with default $\HH=\Delta^L$ assuming no covariate shift. \texttt{DRL0-split} and \texttt{DRL0-nosplit} stand for Algorithm \ref{alg: cs mm} with and without sample splitting, respectively. {\ERM} stands for the empirical risk minimizer fitted over the pooled source data sets.
}
\label{fig: mix-varyL}
\end{figure}
In Figure \ref{fig: mix-varyL}, we compare {\DRLn} (with and without sample splitting) and {\ERM} for two different $q^{\rm sou}$: (1) even mixture with $q^{\rm sou} = (\frac{1}{L}, ..., \frac{1}{L})$; and (2) uneven mixture with $q^{\rm sou} = (0.55, \frac{0.45}{L-1}, ..., \frac{0.45}{L-1})$. As expected, our proposed {\DRLn} consistently outperforms {\ERM} in both scenarios. This happens because when the target data has a distributional shift from the source data, the {\ERM} does not generalize well while {\DRLn} ensures a more robust prediction model.
Instead of investigating the worst-group reward, in Figure \ref{fig: mix-L2}, we further look into the case with $L=2$ and report the reward for a collection of target populations (with various target mixture $q^{\rm tar}$). For comparison, we also include the approach \texttt{DRO-sq}, which is the sample version of $f^{\rm sq}$ defined in \eqref{eq: l2 loss def}. %As shown in Proposition \ref{prop: squared loss model}, \texttt{DRO-sq} is also a weighted average of the prediction models fitted for the two source groups. In Figure \ref{fig: mix-L2}, we consider three different source mixture weights:
We compare the approaches for three different $q^{\rm sou}$: (1) $q^{\rm sou} = (0.2, 0.8)$, (2) $q^{\rm sou} = (0.5, 0.5)$, and (3) $q^{\rm sou} = (0.8, 0.2)$. For each source mixture $q^{\rm sou}$, we vary $q^{\rm tar}$ and evaluate the prediction performance on a collection of test data sets generated with different $q^{\rm tar}$. Particularly, we vary the first weight $q^{\rm tar}_1$ across $\{0, 0.05, 0.1, ..., 0.95, 1\}$.
%In Figure \ref{fig: mix-L2}, the worst-case rewards for each method, among all values of $q^{\rm tar}_1$, are .
Our proposed {\DRLn} consistently outperforms the other two methods regarding the worst reward (represented by a triangle). The vertical black dashed lines correspond to the case $q^{\rm tar} = q^{\rm sou}$. As expected, {\ERM} performs the best when $q^{\rm tar}$ is close to $q^{\rm sou}$.
However, as $q^{\rm tar}$ gets more different from $q^{\rm sou}$, {\DRLn} starts to demoninate {\ERM} in terms of the reward.
%Lastly, when $q^{\rm sou} = (0.5, 0.5)$, the performance of \texttt{DRO-sq} is similar to {\ERM}, since the source data mixture aligns with the simple average generated by \texttt{DRO-sq}.
\begin{figure}[ht]
\centering
\includegraphics[width=0.9\textwidth]{assets/plot-mix-L2.pdf}
\caption{ Comparison of rewards for {\DRLn}, {\ERM} and \texttt{DRO-sq} for $L=2$ over a collection of target populations with mixture weight $q^{\rm tar}_1$ varied across $\{0, 0.05, 0.1, ..., 0.95, 1\}$. {\DRLn} is implemented by Algorithm \ref{alg: cs mm} with the default $\HH = \Delta^L$ and no sample splitting, assuming no covariate shift; \texttt{DRO-sq} stands for the sample version of $f^{\rm sq}$ defined in \eqref{eq: l2 loss def}; {\ERM} stands for the empirical risk minimizer fitted over the pooled source data sets. The triangles correspond to the worst reward for each method. The vertical black dashed lines represent the cases when $q^{\rm tar} = q^{\rm sou}$.
}
\label{fig: mix-L2}
\end{figure}
%For each specific value of $q^{\rm tar}$, we generate the target data, on which we evaluate the reward of each method. \Zijian{Stop correcting the grammar after here.}
%\Zijian{shall we use square for the worst case?}\Zhenyu{Done. Use triangles here.} in Figure \ref{fig: mix-L2}. In the left sub-figure, {\DRLn} (blue), {\ERM} (red), and \texttt{DRO-sq} (green) attain their worst cases when $q^{\rm tar}_1 = 0$, $q^{\rm tar}_1 = 1$, and $q^{\rm tar}_1 = 0$, respectively. All methods achieve their worst cases in the middle and right sub-figures when $q^{\rm tar}_1 = 0$. It is evident that the
%However, when $q^{\rm sou} = (0.2,0.8)$ or $(0.8, 0.2)$, there is no specific pattern for DRO-sq, as it can perform worse or better than ERM and DRL0 depending on the target mixture.
\subsection{{\DRL}: balancing robustness and predictiveness}
\label{sec: simu-H}
%\Zijian{For Section 5, I stop here.}
The construction of a distributionally robust prediction model might lead to an overly conservative prediction model \citep{zhang2020coping, duchi2023distributionally,hu2018does}.
%If we solely focus on the worst group reward during the training step, we may end up with an overly conservative prediction model . %As shown in Figure \ref{fig: mix-L2}, when the target distribution mixture $q^{\rm tar}$ is close to source mixture $q^{\rm sou}$, {\DRLn} might perform less favorably than {\ERM}, as it guarantees the worst case over an overly broad uncertainty region for the target distribution.
However, when we have the prior information about the target distribution, we only consider a convex subset $\HH\subset \Delta^{L}$ and implement our proposed {\DRL} in Algorithm \ref{alg: cs mm}. We shall demonstrate that our proposed {\DRL} algorithm bridges the methods {\DRLn} and {\ERM} and guarantees the prediction model's predictiveness and robustness simultaneously. We consider the source data mixed by $L=4$ groups with mixture weights $q^{\rm sou} = (0.55, 0.15, 0.15, 0.15)$ and implement Algorithm \ref{alg: cs mm} by specifying $\HH = \{q\in \Delta^L \; |\; \|q - q^{\rm sou} \|_2\leq \rho \sqrt{L}\}$ with $\rho$ controlling the size of the uncertainty region. %The parameter $\delta$ in $\HH$ controls the size over which our proposed {\DRL} guarantees robust prediction performance.
For $\rho=1$, {\DRL} guarantees a robust prediction model for any target distribution $\QQ$ belonging to $\mathcal{C}(\QQ_{X})$ and in this case {\DRL} is the same as {\DRLn}. On the other hand, when $\rho=0$, {\DRL} trains a prediction model for the target data generated as a mixture of the source populations with the weight $q^{\rm sou}$. We evaluate the prediction performance by calculating the worst-case reward over the target distribution $\QQ$ generated from the class $\C(\QQ_X, \mathcal{E})$ with $\mathcal{E} = \{q \in \Delta^L \; | \; \|q - q^{\rm sou}\|_2 \leq e\sqrt{L}\}$. Here, the parameter $e$ controls the size of the target region and we vary the size $e$ across $\{0.4, 0.39, ..., 0.01, 0\}$. %For instance, when $e=1$, $\mathcal{E} = \Delta^L$, implying that $\QQ$ can be any mixture of the source groups. Conversely, when $e = 0$, $\mathcal{E}= \{q^{\rm sou}\}$ indicates that the target mixture is exactly the same as the source mixture.
\begin{figure}[!h]
\centering
\includegraphics[width=0.95\textwidth]{assets/plot-mix-L4-worstQ.pdf}
\caption{
Comparison of {\DRL}, {\DRLn} and {\ERM} in terms of the worst-case reward evaluated over the target distribution $\QQ \in \C(\QQ_X, \mathcal{E})$, with $\mathcal{E} = \{q \in \Delta^L \; | \; \|q - q^{\rm sou}\|_2 \leq e\sqrt{L}\}$ and $e$ varied across $\{0.4,0.39,...,0.01,0\}$. {\DRL} is implemented in Algorithm \ref{alg: cs mm} with the prior information $\HH = \{q\in \Delta^L \; |\; \|q - q^{\rm sou} \|_2\leq \rho \sqrt{L}\}$, where the left, middle, and right panels correspond to $\rho=0.25, 0.15, 0.02$, respectively. {\DRLn} is implemented by Algorithm \ref{alg: cs mm} with the default $\HH=\Delta^L$. Both {\DRL} and {\DRLn} are implemented without sample splitting and assuming no covariate shift. {\ERM} stands for the empirical risk minimizer fitted over the pooled source data.}
\label{fig: mix-L4}
\end{figure}
\begin{figure}[!h]
\centering
\includegraphics[width=0.95\textwidth]{assets/plot-mix-L4-fixQ.pdf}
\caption{ Comparison of {\DRL}, {\DRLn}, and {\ERM} in terms of the reward evaluated over three different target settings. {\DRL} is implemented in Algorithm \ref{alg: cs mm} with prior information $\HH = \{q\in \Delta^L \; |\; \|q - q^{\rm sou} \|_2\leq \rho \sqrt{L}\}$, with $\rho$ varied across $\{0.30,0.29,...,0.01,0\}$. {\DRLn} is implemented by Algorithm \ref{alg: cs mm} with the default $\HH=\Delta^L$. Both {\DRL} and {\DRLn} are implemented without sample splitting and assuming no covariate shift. {\ERM} stands for the empirical risk minimizer fitted over the pooled source data. The left, middle, and right panels correspond to three different target distribution settings with mixture weights $q^{\rm tar}$ being set as $(0, 0.1, 0.9, 0)$, $(0.25, 0.25, 0.25, 0.25)$, and $(0.5, 0.15, 0.2, 0.15)$, respectively.
}
\label{fig: mix-L4-fixQ}
\end{figure}
In Figure \ref{fig: mix-L4}, we consider three different values of $\rho$ that control the size of the prior information $\HH$. From the leftmost to the rightmost panel in Figure \ref{fig: mix-L4}, the prior information of $\HH$ shrinks. Since $\HH$ only affects the performance of $\DRL$, the worst-case rewards of {\DRLn} and {\ERM} remain the same for different specifications of $\HH$. On the leftmost panel, the relatively large value of $\rho$ provides little prior information and {\DRL} performs similarly to {\DRLn}; on the rightmost panel, $\rho=0.02$ provides strong prior information and {\DRL} performs similarly to {\ERM}. More interestingly, as illustrated in the middle panel, with a proper specification of $\rho$ (and the prior set $\HH$), {\DRL} combines the strength of {\DRLn} and {\ERM} and strikes a balance between the prediction model's robustness and predictiveness. Particularly, {\DRL} achieves robustness when we consider the target population being far away from the source data (with a relatively large $e$), and {\DRL} ensures predictiveness when all target populations are close to the source data (with a relatively small $e$).
In Figure \ref{fig: mix-L4-fixQ}, we evaluate the prediction accuracy over three representative target distributions instead of computing the worst-case reward over a class of target populations. We generate the three specific target populations, where the target settings 1, 2, and 3 correspond to the mixture weights $q^{\rm tar}$ being set as $(0, 0.1, 0.9, 0)$, $(0.25, 0.25, 0.25, 0.25)$, and $(0.5, 0.15, 0.2, 0.15)$, respectively. We vary the prior information $\HH$ by considering a wide range of $\rho$ values with $\rho \in \{0.30,0.29,...,0.01,0\}$. %\Zhenyu{Zijian, check the caption.}
Across all three panels of Figure \ref{fig: mix-L4-fixQ}, we observe that {\DRL} bridges {\DRLn} and {\ERM}: {\DRL} with a relatively large $\rho$ (little prior information) leads to a similar prediction reward to that of {\DRLn}, while {\DRL} with a relatively small $\rho$ (strong prior information) leads to similar prediction reward to that of {\ERM}. In Figure \ref{fig: mix-L4-fixQ}, from the leftmost to the rightmost panels, we consider settings with a decreasing difference between the source and target distributions.
In the leftmost, with $\|q^{\rm tar} - q^{\rm sou}\|_2/\sqrt{L} = 0.472$, the target mixture $q^{\rm tar} $ is considerably distinguished from the source mixture $q^{\rm sou}$. In this case, a larger $\rho$ value is necessary to guarantee a robust prediction performance for the target population. Particularly, with a small $\rho\in (0,0.3)$, {\DRL} cannot generalize well with this misspecified prior information $\HH$. Therefore, {\DRLn} or {\DRL} (with $\rho\geq 0.3$) outperforms {\ERM} in terms of prediction accuracy for target setting-1. In the middle panel of Figure \ref{fig: mix-L4-fixQ}, we consider the target setting-2 with $\|q^{\rm tar} - q^{\rm sou}\|_2/\sqrt{L} = 0.173$, representing that the target mixture is mildly different from the source mixture. The reward of $\DRL$ firstly increases since a smaller uncertainty set $\HH$ guarantees a more accurate prediction model. However, when the uncertainty set $\HH$ becomes too narrow, the reward evaluated at this specific target distribution drops due to the fact that uncertainty class $\HH$ no longer guarantees a generalizable prediction model for the particular target distribution. In the rightmost panel of Figure \ref{fig: mix-L4-fixQ}, the target mixture $q^{\rm tar}$ is nearly identical to the source mixture $q^{\rm sou}$ with $\|q^{\rm tar} - q^{\rm sou}\|_2/\sqrt{L} = 0.035$. In this case, as expected, {\ERM} outperforms {\DRLn}. With a smaller $\rho$, {\DRL} performs similarly to {\ERM} and outperforms {\DRLn}.
%The parameter $\delta$ for $\HH$ decreases from $0.3$ to $0$ on the x-axis. Initially, when $\delta=0.3$, the broader uncertainty region of $\HH$ (which includes $q^{\rm tar}$) results in a conservative prediction model. Thus, as $\delta$ decreases, the scope of $\HH$ becomes narrower, leading to a more accurate performance by {\DRL}, until it eventually fails to encompass $q^{\rm tar}$. Beyond this point, as $\delta$ continues to decrease, $\HH$ becomes increasingly misspecified, resulting in a decline in the generalization performance of ${\DRL}$. This pattern accounts for the observed reward curve that initially increases and then decreases as $\delta$ decreases.
% As $\delta$ decreases from $0.3$ to $0$, the prior information $\HH$ initially contains the target mixture $q^{\rm tar}$ but then fails covering it, leading to the observed increasing-then-decreasing reward curve. Specifically, when $\delta$ is relatively large (approximately $\geq 0.17$), a smaller $\delta$ implies a narrower uncertainty region, resulting in a more accurate prediction model. However, when $\delta$ becomes small (approximately $< 0.17$), the prior information $\HH$ is misspecified and fails to cover $q^{\rm tar}$ effectively, leading to a decrease in generalization performance for {\DRL}.
%Additionally, as the provided $\HH$ already contains the target mixture $q^{\rm tar}$ for any $\delta \in (0.3,0.01)$, selecting a smaller $\delta$ value brings {\DRL} closer to the performance of {\ERM}, resulting in a more accurate model.
% \Zijian{Zhenyu, revise}In the mid dle sub-figure, for the target setting-2 with $\|q^{\rm tar} - q^{\rm sou}\|_2/\sqrt{L} = 0.173$, the initial larger $\delta$ for the prior $\HH$ correctly include $q^{\rm tar}$, leading to better accuracy and higher rewards for {\DRL} as $\delta$ values decreases. However, with excessively small $\delta$ values, $\HH$ fails to cover $q^{\rm tar}$ correctly, resulting in a decrease in robustness and poorer rewards with $\delta$ keep decreasing. This explains the observed increasing-then-decreasing reward curve as $\delta$ decreases.
% \Zijian{Zhenyu, revise} In the right sub-figure, for the target setting-3 with $\|q^{\rm tar} - q^{\rm sou}\|_2/\sqrt{L} = 0.035$, the source and target distributions are almost identical, favoring {\ERM} over {\DRLn}, as expected. Given that the provided $\HH$ correctly contains the target mixture, smaller $\delta$ values yield higher rewards for {\DRL}, significantly improving its performance.
% \subsection{Accommodations to Various Density Ratio Estimators For Covariate Shift Setting}
% \subsection{Accommodations to Various Machine Learning Models}
\subsection{Distributionally robust boosting and neural networks}
In our previous simulations, we mainly utilized random forests for our proposed framework. The following simulation shows that our proposed {\DRL} framework can accommodate various machine learning models, including neural networks and boosting methods.
% The parameters of the neural networks and boosting will be delayed in Appendix \ref{sec: exp-setup}.
\begin{figure}[ht]
\centering
\includegraphics[width=0.95\textwidth]{assets/plot-mix-L4-boost.pdf}
\caption{Comparison of {\DRL} with various machine learning models. \texttt{DRL-rf},\texttt{DRL-xgb} and \texttt{DRL-nn} are implemented by Algorithm \ref{alg: cs mm} (with no sample splitting and assuming no covariate shift) with random forests, xgboost, and neural networks as the base machine learning algorithms, respectively. {Top panel:} the settings correspond to the Figure \ref{fig: mix-L4}; Bottom panel: the settings correspond to the Figure \ref{fig: mix-L4-fixQ}.}
\label{fig: mix-L4-nn}
\end{figure}
Figure \ref{fig: mix-L4-nn} presents a comparison of {\DRL} with various machine learning algorithms, including random forests, boosting, and neural networks. The top panel corresponds to the settings in Figure \ref{fig: mix-L4}, while the bottom one corresponds to the settings in Figure \ref{fig: mix-L4-fixQ}. In both panels, we observe that \texttt{DRL-rf},\texttt{DRL-xgb} and \texttt{DRL-nn} have similar patterns and \texttt{DRL-nn} has uniformly better prediction performance than \texttt{DRL-rf} and \texttt{DRL-xgb}, suggesting that neural networks provide more accurate estimators $\{\hf{l}\}_{l\in [L]}$ compared to random forests and xgboost in our current simulation settings.
\subsection{Effectiveness of bias correction}
\label{sec: simu-cs}
We explore the effect of estimating the density ratio, as specified in Algorithm \ref{alg: cs mm}. Particularly, we assess the performance of the bias-corrected {\DRLn} with various density ratio estimators and compare it with the plugin {\DRLn} estimator to highlight the effectiveness of bias correction.
With the number of groups $L=5$, we set the conditional outcome model $\f{l}$ as
$\f{l}(x) = \sum_{1\leq j\leq 4} \alpha_j^{(l)} \ind_{[x_j>0]} + \sum_{1\leq j\leq k\leq 4} \beta_{jk}^{(l)}\ind_{[x_j>0]}\ind_{[x_k>0]} - \sum_{1\leq j\leq k\leq 4}\gamma_{jk}^{(l)}\ind_{[x_j<2]}\ind_{[x_k>-2]}$, for $x\in \R^4$. We randomly generate the coefficients across the $L$ groups to approximate the heterogeneous conditional outcome models. Particularly, for $1\leq l\leq L$ and $1\leq j\leq k\leq 4$, the coefficients $\alpha_j^{(l)}$,$\beta_{jk}^{(l)}$, and $\gamma_{jk}^{(l)}$ are independently generated random variables, taking values in $\{8,1.6,0,-4\}$ with equal probability, respectively. The sample size for each source group is varied across $\{200,400,600\}$, while the sample size for the target covariates is fixed at $n_\QQ = 1000$.
% To facilitate this simulation, we generate the target covariate data according to the following distribution:
% \[
% \{\XQ_j\}_{j\in [n_\QQ]} \distras{iid} \N(\mu_0, \Sigma_0)\quad \textrm{with}\; \mu_0 = (1/4,1/4,1/4,1/4,1/4)^\intercal \;\textrm{and}\; \Sigma_0 = \textrm{diag}(1.5, 1.2, 1, 0.8, 1.5)
% \]
% We use the same source data setting as before with an even mixture of $L=4$ groups and vary the sample size per group across $\{2000,4000,8000\}$.
In Figure \ref{fig: mix-L4-cs}, we compare the following versions of $\DRLn$. %(with default $\HH = \Delta^L$) . Each version of $\DRLn$ will generate the aggregation weights and the robust model estimator, denoted as $\widehat{q}$ and $\widehat{f}$.
\begin{itemize}
\item \texttt{DRL0-plug}: the plug-in DRL estimator in \eqref{eq: Magging plugin} with $\HH = \Delta^L$.
\item \texttt{DRL0-log}: Algorithm \ref{alg: cs mm} with $\HH = \Delta^L$ and density ratios estimated by the logistic regression as in \eqref{eq: density ratio}.
% and \eqref{eq: density ratio estimation}. %This approach utilizes logistic regression for binary classification in the density ratio estimation step.
\item \texttt{DRL0-rf}: Algorithm \ref{alg: cs mm} with $\HH = \Delta^L$ and density ratios estimated by random forests. %random forests are used for binary classification in the density ratio estimation step.
\end{itemize}
% \Zijian{Zhenyu, we are a bit repeating now.}The aim of this investigation is to verify the approximation error mentioned in Theorem Theorem \ref{thm: Magging correct}.
\begin{figure}[!h]
\centering
\includegraphics[width=0.9\textwidth]{assets/plot-improve-rf-nocs.pdf}
\caption{ Comparison of {\DRLn} estimators for the covariate shift setting with $L=5$. \texttt{DRL0-plug} represents the plug-in estimator in \eqref{eq: Magging plugin} with $\HH = \Delta^L$.
\texttt{DRL0-log} and \texttt{DRL0-rf} are implemented by Algorithm \ref{alg: cs mm} with $\HH = \Delta^L$ and the density ratios estimated by the logistic regression and random forests, respectively. The sample size per source group is varied across $\{200,400,600\}$. The sample size for target data is fixed at $n_\QQ = 1000$. Top panel: the distance of the data-dependent aggregation weight to the true aggregation weight, measured by $\|\widehat{q} - q^*\|_2^2$. Bottom panel: the distance of the estimator of {\DRLn} to the truth, measured by $\|\widehat{f} - f^*\|_{\QQ,2}^2$.}
\label{fig: mix-L4-cs}
\end{figure}
In the top panel of Figure \ref{fig: mix-L4-cs}, we calculate the aggregation weight $\widehat{q}$ for each method and present $\|\widehat{q} - q^*\|_2^2$ with $q^*$ denoting the true aggregation weight. The bottom panel displays the estimation error $\|\widehat{f} - f^*\|_2^2$. %The true aggregation weights $q^*$ and the actual robust model $f^*$ for {\DRLn} are determined by directly incorporating the true conditional outcome models $\{\f{l}\}_{l\in [L]}$ into the optimization process outlined in \eqref{eq: Magging plugin}.
%For better illustration, we remove outliers of each method.
In the top panel of Figure \ref{fig: mix-L4-cs}, the aggregation weights computed by \texttt{DRL0-log} and \texttt{DRL0-rf} are more accurate than \texttt{DRL0-plug}, and the improvements are more significant with a larger source data sample size. In the bottom panel of Figure \ref{fig: mix-L4-cs}, we still observe that the bias correction leads to improvement in estimating $f^*$, even though the improvement is not as huge as that for the aggregation weights. This happens since the estimation error $\|\widehat{f} - f^*\|_{\QQ,2}^2$ decomposes of two parts: the error of estimating $\hf{l}$ and estimating the aggregation weight and the bias correction procedure mainly improves the rate of estimating the aggregation weight. %, see \eqref{eq: shift simplified}.
%\Zhenyu{Peter, Zijian: please check the preceding paragraph.}
% However, for the method \texttt{DRL0-rf}, when $n=1000$, it is not as effective as the other two methods due to the poor estimation of density ratio estimations with random forests. Instead, as the sample size per group increases to $n=2000$, \texttt{DRL0-rf} starts dominating \texttt{DRL0-plug}, and when $n=4000$, \texttt{DRL0-rf} performs equally well as \texttt{DRL0-log}. Hence, from the above figures, we can conclude that both \texttt{DRL0-log} and \texttt{DRL0-rf} work well for this covariate shift setting, especially when the sample size per group is sufficiently large.
\subsection{Capturing shared associations in high-dimensional setting}
Section \ref{subsec: identification} emphasizes that {\DRL} can capture the common factor across different sites.
We now demonstrate this phenomenon using high-dimensional regressions. We set the group number $L=5$, the dimension $p=200$, and set each group to have the same number of observations, denoted as $n_l = n$ for all $l\in [L]$. We define the base model $f_{\rm base}: \R^p \mapsto \R$ as $f_{\rm base}(x) = \sum_{j=1}^{10} 0.5 x_j$, where only the first ten of $p=200$ variables have non-zero coefficients here.
For $l \in [L]$, we generate the $l$-th group source data following \eqref{eq: source-models} with $\PP{l}_X = \N(0_p, I_p)$, $\f{l}(x) = f_{\rm base}(x) + \sum_{j=11}^{13} \gamma^{(l)}_{j} x_{j}$, and $\eps{l}_i \distras{iid} \N(0,1)$.
The coefficients $\gamma^{(l)}_j$ are independently drawn from $\N(0,1)$ for $11\leq j\leq 13$. Consequently, the first $10$ coefficients remain the same across each source group, while the $11$-th to $13$-th coefficients are randomly generated to be scattered around $0$.
We consider the no covariate shift setting with $\QQ_X = \N(0_p, I_p)$. For each round of simulation, the target conditional outcome model is generated as: $f^\QQ(x) = f_{\rm base}(x) + \sum_{j=11}^{13} \gamma^\QQ_{j} x_{j}$, where $\gamma^\QQ_j \sim \N(0,1)$ for $11\leq j\leq 13$. As a result, the target data shares the same base model with source groups, but the $11$-th to $13$-th coefficients are randomly generated, not necessarily seen in the source groups. The target data is generated with $n_\QQ = 10000$ observations. We utilize the LASSO algorithm with the hyperparameters tuned by cross-validation for fitting models $\f{l}$ for $l\in [L]$.
In the left sub-figure of Figure \ref{fig: mix-highd-fig1}, we can observe that {\DRLn} consistently outperforms {\ERM} with a higher reward, indicating that {\DRLn} consistently produces a better prediction model compared to {\ERM}. In the right sub-figure, we compare the squared $\ell_2$ distance of the computed coefficients of {\DRLn} and {\ERM} to the coefficients of the base model. With an increasing sample size per group, the distance for {\DRLn} steadily diminishes, approaching zero. This trend indicates that the initially randomly generated coefficients (11th to 13th) are gradually nullified, leading to a successful recovery of the base model. In contrast, the distance associated with {\ERM} remains non-negligible and tends to stabilize at approximately 0.4 as the sample size grows. This analysis confirms that {\DRLn} successfully captures the shared factors across the groups while effectively discarding the heterogeneous elements. On the contrary, {\ERM} retains those randomly generated coefficients, resulting in a failure to recover the underlying base model.
The proposed {\DRLn} provides aggregation weights for each source group. In the left sub-figure of Figure \ref{fig: mix-highd-fig2}, we display the computed weights for the five groups with varying sample sizes per group, ranging from $\{100,200,300,400,800,1200\}$. In the right sub-figure of Figure \ref{fig: mix-highd-fig2}, we present the 11th to 13th coefficients (which are randomly generated) of the five groups. The groups G1, G3, and G5 have leading weights in determining the final \texttt{DRL} model. Note that the 11th to 13th coefficients for group G5 are of the opposite sign to those for group G3 and are mostly of the same sign as those for group G1. The weights determined by \texttt{DRL} lead to the cancellation of the heterogeneous effects with opposite signs.
\begin{figure}[!h]
\centering
\includegraphics[width=0.8\textwidth]{assets/plot-highd-L5-fig1.pdf}
\caption{Comparison of {\DRLn} and {\ERM} in high-dimensional setting with sample size per group varied across $\{100,150,200,250,300,400,800,1200\}$. \texttt{DRL0} is implemented by Algorithm \ref{alg: cs mm} with default $\HH=\Delta^L$, no sample splitting, and assuming no covariate shift. {\ERM} stands for the empirical risk minimizer fitted over the pooled source data sets. {Left panel:} the reward evaluated over the target data; { Right panel:} the squared $\ell_2$ distance between the computed coefficients to the base model's coefficients. {\DRLn} is implemented by Algorithm \ref{alg: cs mm} without sample splitting.}
\label{fig: mix-highd-fig1}
\end{figure}
\begin{figure}[!h]
\centering
\includegraphics[width=0.8\textwidth]{assets/plot-highd-L5-fig2.pdf}
\caption{{Left panel:} the computed aggregation weights for the five source groups with sample size per group varied across $\{100,200,400,800,1200\}$; {Right panel:} the actual randomly generated coefficients (11th to 13th) for five groups.}
\label{fig: mix-highd-fig2}
\end{figure}
%and the dominating the only one with a positive signal for the 11th coefficient, while group G3 exhibits the sole positive signal for the 12th coefficient. To effectively cancel out these signals, we can deduce that groups G3 and G5 play a crucial role in the final aggregation.
% For the 11th coefficient, group G5 is the only one with a positive signal. For the 12th coefficient, group G3 is the only group with a negative signal. Additionally, the positive effect of group G1, in conjunction with group G3, on the 13th coefficient results in the overall cancellation of the signal on the 13th coefficient as well.
% It is evident that groups G2 and G4 have vanishing weights with increasing sample size, while the remaining groups G1,G3,G5 contribute significantly to the final model.
% In the middle and right sub-figures of Figure \ref{fig: mix-highd-fig2}, we perform experiments where we remove each group and train a {\DRL} model with the remaining four source groups using the default input $\HH$. In the middle sub-figure, the rewards after removing G3 (green points) and after removing G5 (orange points) are the lowest, which signifies a significant decrease in reward after removing these two groups. In the right sub-figure, we evaluate the $\ell_2$ distance of the {\DRL} model after removing each group, and once again, we observe that G5 and G3 have the highest distances, indicating the failure to recover the base model after removing these two groups. Thus, it is evident that groups G3 and G5 play a crucial role in contributing to the final robust model, which aligns with their higher aggregation weights.
%\subsection{Distributionally Robust Neural Networks}
%\Zijian{Zhenyu, fill in this section.}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%% 7.0 %%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Real Data}
We validate our proposed DRL approach using the Beijing PM2.5 Air Pollution data, initially analyzed by \citet{zhang2017cautionary} and publicly available at \url{https://archive.ics.uci.edu/dataset/501/beijing+multi+site+air+quality+data}. The dataset consists of PM2.5 concentration measurements from 2013 to 2016 at 12 nationally controlled air-quality monitoring sites in Beijing, along with meteorological covariates including temperature (TEMP), pressure (PRES), dew point temperature (DEWP), precipitation (RAIN), wind direction (WD), and wind speed (WSPM). \footnote{The wind direction variable (WD) is categorized into five groups: northwest (NW), northeast (NE), calm and variable (CV), southwest (SW), and southeast (SE).} For a given PM2.5 observational site (denoted as $s$), we represent the centralized logarithm transformed PM2.5 concentration at a particular hour $t$ in season $j$ of a specific year $l$ as $Y^{(l)}_{t}(j,s)$. The seasons are indexed as follows: 1 for spring (March to May), 2 for summer (June to August), 3 for autumn (September to November), and 4 for winter (December to February of the following year). We use $X^{(l)}_{t}(j,s)$ to denote the corresponding measurements of the six meteorological variables.
\begin{figure}[!h]
\centering
\includegraphics[width=0.95\textwidth]{assets/plot-air-fig1.pdf}
\caption{PM2.5 concentration analysis on sites ``Dingling'', ``Huairou'', ``Aotizhongxin'' and ``Dongsi'', denoted as sites A, B, C, and D, respectively. The {\DRLn} method is implemented by Algorithm \ref{alg: cs mm} (using default $\HH=\Delta^L$) with no sample splitting and with the density ratios estimated by logistic regression. {Top Panel}: comparison of reward (prediction accuracy) for approaches {\DRLn} and {\ERM} evaluated on Year 2016. We also include the model fitted solely on Year 2015 for comparison. {Bottom Panel}: the computed weights for the {\DRLn} model. }
\label{fig: air}
\end{figure}
We first fix the monitoring site $s$ and season $j$ and train our proposed $\DRL$ algorithm and $\ERM$ using the data collected for years $l=2013, 2014,$ and $2015$, where the year denotes the group label $l$ in our analysis. We evaluate the performance of $\DRL$ and $\ERM$ using the data collected in 2016. We train the base prediction models using random forests with hyperparameters tuned by out-of-bag error. In our analysis, we consider all four seasons and specifically focus on four monitoring sites: ``Dingling''(located in the northwest), ``Huairou''(located in the northeast), ``Aotizhongxin'' and ``Dongsi''(both located centrally). For example, we train the prediction models using 2013-2015's spring data for the site ``Dingling" and evaluate their performance using the spring data of 2016 in ``Dingling". %These four sites are denoted as site A,B,C, and D respectively, as shown in Figure \ref{fig: air}.
In Figure \ref{fig: air}, we compare the {\DRLn} and {\ERM} approaches in the top panel, which displays the rewards evaluated for each season in the year 2016. We also include the model fitted solely on the year 2015 since our computed {\DRLn} weights indicate that the year 2015 has the highest weight among the three years and also it is the latest year to the target year 2016. Our proposed {\DRLn} achieves the best prediction accuracy among these three prediction models. For sites A (``Dingling'') and B (``Huairou''), {\DRLn} significantly outperforms {\ERM} in terms of prediction accuracy across all four seasons. This happens since the conditional outcome model for the year 2016 may have a shift from those for the years 2013-2015. For sites C (``Aotizhongxin'') and D (``Dongsi''), {\DRLn} performs similarly to {\ERM} during autumn and winter while {\DRLn} outperforms {\ERM} during spring and summer. The bottom panel of Figure \ref{fig: air} presents the aggregation weights for the {\DRLn} model. We observe that the year 2015 contributes the most to the final robust prediction model, particularly during spring. For the years 2013 and 2014, in spring and autumn, 2014 is relatively more important, while 2013 has a greater contribution during summer and winter.
\begin{figure}[h]
\centering
\includegraphics[width=0.95\textwidth]{assets/plot-air-fig2.pdf}
\caption{Comparison of rewards for approaches {\ERM}, {\DRL}, and \texttt{DRL-plug} under two different prior information $\HH$ on the site ``Dingling''. The input prior info $\HH$ is in the form as $\HH = \{q\in \Delta^L \;|\; \|q - v_0\|_2\leq \rho\sqrt{L}\}$, where each entry of $v_0 \in \R^3$ corresponds to the weight for the group Year 2013, 2014, and 2015 respectively. {\DRL} is implemented by Algorithm \ref{alg: cs mm} with no sample splitting and with the density ratios estimated by logistic regression. \texttt{DRL-plug} is implemented as in \eqref{eq: Magging plugin}. {Top Panel:} $v_0 = (0.15, 0.35, 0.5)$ for the approaches {\DRL} and \texttt{DRL-plug}. {Bottom Panel:} $v_0 = (0.35, 0.15, 0.5)$ for the approaches {\DRL} and \texttt{DRL-plug}.}
\label{fig: air-H}
\end{figure}
Figure \ref{fig: air-H} illustrates the rewards obtained using two different prior information for the {\DRL} approach with $\HH = \{q \;|\; \|q - v_0\|_2 \leq \rho \sqrt{L}\}$, where each entry of $v_0\in \R^3$ corresponds to the prior weight of 2013, 2014, and 2015, respectively. The top panel and bottom panel correspond to $v_0 = (0.15, 0.35, 0.5)$ and $v_0 = (0.35, 0.15, 0.5)$, respectively. For both $\HH$, we compare the approaches {\DRL} (Algorithm \ref{alg: cs mm}), \texttt{DRL-plug} (without bias-correction), and {\ERM}. It is evident that {\DRL} consistently outperforms \texttt{DRL-plugin}, highlighting the advantage of bias correction for this real data application. Additionally, both {\DRL} and \texttt{DRL-plugin} outperform {\ERM}.
In Figure \ref{fig: air-H}, the difference between the top and bottom panels lies in whether we prioritize the Year 2013 or 2014. For the spring and autumn seasons with both specifications of $v_0$ and the winter season with $v_0 = (0.35, 0.15, 0.5)$, the reward given by different $\rho$ (or different prior information $\HH$) either remain nearly the same or only slightly increase. However, for the summer season and the winter season with $v_0 = (0.15, 0.35, 0.5)$, we observe a sharp decline in rewards for {\DRL} as $\rho$ decreases. This suggests that the set $\HH$ is likely to be far from the true mixture weight in the Year 2016 (if it exists). In this case, the misspecification of $\HH$ does not guarantee {\DRL} to have a nice generalizable predictive performance for the Year 2016. With a smaller value of $\rho$ (or $\HH$), the predictive performance of {\DRL} gets worse. %with a smaller $\delta$ and $\HH$. %In contrast, in the bottom panel with the center weight $v_0 = (0.35, 0.15, 0.5)$, the predictive performance decline for {\DRL} (with a smaller $\delta$ value) is mitigated in seasons of summer and winter. By comparing the top and bottom panels, we can infer that the conditional outcome model at site ``Dingling'' is relatively more aligned with the aggregated models using weights $(0.35, 0.15, 0.5)$ than with weights $(0.15, 0.35, 0.5)$.
% Figure \ref{fig: air} shows the PM2.5 concentration analysis for the ``Dingling'',``Huairou'',``Dongsi'' sites. The top panel compares the rewards obtained by the DRL and ERM approaches in the year 2016. We also include the model fitted solely on Year 2015, since our DRL weights indicate that the Year 2015 has the highest weight among the three years. For the ``Dingling'' and ``Huairou'' sites, the DRL approach consistently outperforms the ERM and Year2015 approaches. Notably, in some seasons, like 2016-2 in ``Dingling'' site, the ERM approach yields a negative reward, indicating worse performance than a null model that predicts a constant value of 0.
% Regarding the "Dongsi" site, the DRL is comparable with ERM as these two approaches perform similarly in all four seasons.
% In the bottom panel of Figure \ref{fig: air}, we present the computed weights for each group (year) in the DRL model. These weights indicate the contribution of each year to the DRL model's predictions. For instance, at the ``Dingling'' site across four seasons, the Year2015 accounts for nearly 50\% in the final aggregated DRL model. In the other two sites, ``Huairou'' and ``Dongsi'', we can see that Year2015 contributes the most as well. This phenomenon confirms our expectation that the Year2015 helps more to predict the latest Year2016 than the previous Year2013 and Year2014. In addition, combined with the top panel, we can see it is not enough to use Year2015 solely to predict Year2016, as the green lines are the worst compared with blue lines DRL and red lines ERM.
% For $i\in \{2013,2014,2015,2016\}$ and $j\in \{1,2,3,4\}$, we use the non-parametric regression model for each site $s$:
% \[
% Y^{(s)}_{ij}(t) = f^{(s)}_{ij}(X^{(s)}_{ij}(t)) + \epsilon^{(s)}_{ij}(t)
% \]
% where $\epsilon^{(s)}_{ij}(t)$ is the residual, and $f^{(s)}_{ij}(x) = \E[Y^{(s)}_{ij}(t) | X^{(s)}_{ij}(t) = x]$.
% We selected five monitoring sites, namely ``Aotizhongxin'', ``Dongsi'', ``Guanyang'', ``Gucheng'', and ``Huairou'', as our source sites. Our analysis focuses on two aspects: (1) generalization to the mixture of these five sites, as measured by the worst site mean squared error (MSE); and (2) generalization to an unseen site, specifically the MSE on a new site called ``Wanliu''. It is important to note that, for the purpose of clear presentation and illustration, we employed the MSE metric instead of the Reward metric in our real-world application. However, the same conclusions apply when using the Reward metric. In Figure \ref{fig: real-air}, we compare approaches including
% \begin{itemize}
% \item {\bf mag}: bias-corrected Magging in \ref{alg: cs mm} without specification of $\HH$
% \item {\bf mag-H}: bias-corrected Magging with $\HH = \{q\; |\; \|q - q_{\rm so}\|_2 \leq 0.1 \sqrt{L},\; q\in \Delta^L\}$, where $q_{\rm so} = (\frac{1}{5},\frac{1}{5},\frac{1}{5},\frac{1}{5},\frac{1}{5})^\intercal$ and $L=5$.
% \item ERM: pool over 5 source sites data and fit the regression over the pooled data.
% \end{itemize}
% All the models are fitted with random forests, with hyperparameters tuned by cross-validation.
% \begin{figure}[ht]
% \centering
% \includegraphics[width=0.8\textwidth]{assets/plot-air.pdf}
% \caption{Comparison of Approaches in terms of generalization to sites mixture and the unseen site. {\bf Top Panel}: At each year-season, for every site (among the 5 sites), we pick 70\% observations as source to fit models, and use the remaining 30\% observations as the target to evaluate the worst group MSE. {\bf Bottom Panel}: At each year-season, for every site (among the 5 sites), we use all the observations to fit models, and evaluate the MSE on the unseen site.}
% \label{fig: real-air}
% \end{figure}
% In the top panel of Figure \ref{fig: real-air}, we compare the generalization performance of each approach to the mixture of the five source sites. We measure the performance by evaluating the worst site mean squared error (MSE), where a lower value indicates greater robustness to the potential mixture. Across different years and seasons, we observe that the ERM approach (represented by red lines) consistently has a larger worst site MSE compared to {\bf mag} and {\bf mag-H}. This aligns with our expectations, as DRL is specifically designed to handle the mixture of conditional outcomes from multiple sites. When comparing the two DRL algorithms, {\bf mag-H} consistently outperforms {\bf mag}, indicating that, for this dataset, the prior information provided by $\HH$ is correctly specified, resulting in more accurate predictions.
% In the bottom panel of Figure \ref{fig: real-air}, we illustrate the generalization of DRL to an unseen site ``Wanliu'' , assuming that the conditional outcome of this unseen site closely resembles a mixture of the source sites. We first compare the ERM and {\bf mag} approaches. In most cases, we observe significant advantages of DRL, particularly from 2014-4 to 2016-1. However, there are instances, such as the period from 2013-4 to 2014-3, where {\bf mag} does not perform as well as ERM, suggesting that {\bf mag} may be too conservative as it explores the entire simplex $\Delta^L$. Next, we examine the {\bf mag-H} approach. For nearly all time points, {\bf mag-H} achieves a lower MSE than ERM, indicating better generalization. Importantly, for time points where {\bf mag} falls behind ERM, such as the period from 2013-4 to 2014-3, {\bf mag-H} substantially improves the MSE, demonstrating that the conditional outcome for the unseen site ``Wanliu'' is accurately captured in our specified $\HH$. Therefore, {\bf mag-H} achieves both robustness and accuracy simultaneously. However, there is a period from 2014-4 to 2015-4 where {\bf mag-H} does not perform as well as {\bf mag}. From this, we can conclude that the specification of $\HH$ may be incorrect, although {\bf mag-H} still outperforms ERM.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%% bib %%%%%%%%%%%%%%%%%%%%%%%%%%
%\section*{Acknowledgement}
\bibliographystyle{abbrvnat}
\bibliography{ref}
\newpage
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\appendix
\section{Additional Methods and Theories}
\label{sec: add_methods}
In this section, we will give detailed formulas for the high-dimensional regime
as in Section \ref{sec: ex-highd}. We firstly introduce the estimators $\{\widehat{\beta}^{(l)}, \widehat{\beta}^{(l)}_\A, \widehat{\beta}^{(l)}_\B\}$ in \eqref{eq: beta-rate} for conditional outcome models; then move forward to estimators $\{\widehat{\gamma}^{(l)}_\A, \widehat{\gamma}^{(l)}_\B\}$ in \eqref{eq: gamma-rate} for density ratio estimation.
For each source group $l\in [L]$, we use LASSO to estimate $\beta^{(l)}$ and obtain the estimators $\widehat{\beta}^{(l)}$, $\widehat{\beta}^{(l)}_\A$, $\widehat{\beta}^{(l)}_\B$. They are fitted on the full data and subset data belonging to $\A_l$ and $\B_l$, respectively. $\widehat{\beta}^{(l)}$ (with full data) is defined as follows:
{\small
\[
\widehat{\beta}^{(l)} = \argmin_{\beta \in \R^p}\left\{\frac{1}{2n_l} \sum_{i=1}^{n_l} (\Y{l}_i - X^{(l)\intercal}_{i} \beta)^2 +
A\sqrt{\frac{\log p}{n_l}} \sum_{j=2}^p \frac{\|\X{l}_{\cdot j}\|_2}{\sqrt{n_l}} |\beta_j| \right\},
\]
} where the first column of $X^{(l)}$ is set as the constant 1, $X^{(l)}_{\cdot j}$ denotes the $j$-th column of $X^{(l)}$, and $A>\sqrt{2}$ is a pre-specified constant. Similarly, we can define $\widehat{\beta}^{(l)}_\A$ and $\widehat{\beta}^{(l)}_\B$ on the split data subset $\A_l$ and $\B_l$, respectively.
\[
\begin{aligned}
\widehat{\beta}^{(l)}_\A &= \argmin_{\beta \in \R^p}\left\{\frac{1}{2|\A_l|} \sum_{i=1}^{|\A_l|} (\Y{l}_{\A,i} - X^{(l)\intercal}_{\A, i} \beta)^2 +
A\sqrt{\frac{\log p}{|\A_l|}} \sum_{j=2}^p \frac{\|\X{l}_{\A, \cdot j}\|_2}{\sqrt{|\A_l|}} |\beta_j| \right\}, \\
\widehat{\beta}^{(l)}_\B &= \argmin_{\beta \in \R^p}\left\{\frac{1}{2|\B_l|} \sum_{i=1}^{|\B_l|} (\Y{l}_{\B,i} - X^{(l)\intercal}_{\B, i} \beta)^2 +
A\sqrt{\frac{\log p}{|\B_l|}} \sum_{j=2}^p \frac{\|\X{l}_{\B, \cdot j}\|_2}{\sqrt{|\B_l|}} |\beta_j| \right\},
\end{aligned}
\]
where $\{\X{l}_\A, \Y{l}_\A\}$ denotes data belonging to the subset $\A_l$ in $\{\X{l}, \Y{l}\}$, and similar representation holds for $\{\X{l}_\B, \Y{l}_\B\}$.
For the density ratio estimation part, for each source group $l\in [L]$, we propose to estimate $\gamma^{(l)}$ by the penalized log-likelihood estimators $\widehat{\gamma}^{(l)}_\A$ and $\widehat{\gamma}^{(l)}_\B$, which are fitted with the data subset $\A_l$ and $\B_l$, respectively. We start with the definition of $\widehat{\gamma}^{(l)}_\A$, and then we define $\widehat{\gamma}^{(l)}_\B$ following the similar argument.
{\small
\begin{equation}
\widehat{\gamma}^{(l)}_\A = \argmin_{\gamma \in \R^p}\left\{-\frac{1}{|\A_l|+n_\QQ} \sum_{i=1}^{|\A_l|+n_\QQ} \left[\log(1+\exp(\widetilde{X}^{(l)\intercal}_{\A, i}\gamma^{(l)})) - \G{l}_{\A,i}\cdot (\widetilde{X}^{(l)\intercal}_{\A, i}\gamma^{(l)})\right] + \lambda^{(l)}_\A \sum_{j=2}^p \frac{\|\widetilde{X}^{(l)}_{\A, \cdot j}\|_2}{\sqrt{|\A_l|+n_\QQ}} |\gamma_j|\right\},
\label{eq: penalized MLE}
\end{equation}}
where $\widetilde{X}^{(l)}_\A = (X^{(l)\intercal}_\A, X_\QQ^\intercal)^\intercal$, $\widetilde{X}^{(l)}_{\A,\cdot j}$ denotes the $j$-th column of $\widetilde{X}^{(l)}_\A$, the first column of $\widetilde{X}^{(l)}_\A$ is set as the constant 1, and the tuning parameter $\lambda^{(l)}_\A$ is of the order $\sqrt{\log p/(|\A_l|+n_\QQ)}$. The indicator random variable $G^{(l)}_{\A,i}$ is set as $G^{(l)}_{\A,i}=0$ if $\widetilde{X}^{(l)}_{\A,i}$ is drawn from the $l$-th source group's subset $\A_l$ and is set as $G^{(l)}_{\A,i}=1$ if $\widetilde{X}^{(l)}_{\A,i}$ is drawn from the target distribution. By replacing $\A_l$ with $\B_l$ in \eqref{eq: penalized MLE}, we can obtain the estimator of $\widehat{\gamma}^{(l)}_\B $, that can be expressed as follows:
{\small
\begin{equation*}
\widehat{\gamma}^{(l)}_\B = \argmin_{\gamma \in \R^p}\left\{-\frac{1}{|\B_l|+n_\QQ} \sum_{i=1}^{|\B_l|+n_\QQ} \left[\log(1+\exp(\widetilde{X}^{(l)\intercal}_{\B, i}\gamma^{(l)})) - \G{l}_{\B,i}\cdot (\widetilde{X}^{(l)\intercal}_{\B, i}\gamma^{(l)})\right] + \lambda^{(l)}_\B \sum_{j=2}^p \frac{\|\widetilde{X}^{(l)}_{\B, \cdot j}\|_2}{\sqrt{|\B_l|+n_\QQ}} |\gamma_j|\right\},
\label{eq: penalized MLE-2}
\end{equation*}}
where $\widetilde{X}^{(l)}_\B = (X^{(l)\intercal}_\B, X_\QQ^\intercal)^\intercal$, $\widetilde{X}^{(l)}_{\B,\cdot j}$ denotes the $j$-th column of $\widetilde{X}^{(l)}_\B$, the first column of $\widetilde{X}^{(l)}_\B$ is set as the constant 1, and the tuning parameter $\lambda^{(l)}_\B$ is of the order $\sqrt{\log p/(|\B_l|+n_\QQ)}$. The indicator random variable $G^{(l)}_{\B,i}$ is set as $G^{(l)}_{\B,i}=0$ if $\widetilde{X}^{(l)}_{\B,i}$ is drawn from the $l$-th source group's subset $\B_l$ and is set as $G^{(l)}_{\B,i}=1$ if $\widetilde{X}^{(l)}_{\B,i}$ is drawn from the target distribution.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Proofs}
Here, we start with the proofs for our main results. In subsections \ref{sec: proof-thm-identification}, \ref{sec: proof-thm-magging-plugin}, and \ref{sec: proof-thm-magging-correct}, we will subsequently present the identification theorem in Theorem \ref{thm: identification}, the approximation errors for plugin DRL and bias-corrected DRL estimators in Theorem \ref{thm: Magging plugin} and Theorem \ref{thm: Magging correct}, respectively. Then in the remaining subsections, we will present the proofs for Corollary \ref{cor: high-dim}, Theorem \ref{prop: reward-diff}, Propositions \ref{prop: shrink H}, \ref{prop: squared loss model}. The proofs of technical lemmas will be delayed to Section \ref{sec: proof-lemmas}.
\subsection{Proof of Theorem \ref{thm: identification}}
\label{sec: proof-thm-identification}
We fix $q \in \mathcal{H} \subseteq \Delta^L$ and consider the joint distribution $\TT = \left(\QQ_X, \TT_{Y|X} \right)$ with $\TT_{Y|X} = \sum_{l=1}^L q_l \cdot \PP{l}_{Y|X}$.
Then the reward for the function $f(\cdot)$ with respect to $\TT$ is defined as:
\begin{align*}
\RR_\TT(f) &:= \E_{(X_i, Y_i) \sim \TT} \left[Y_i^2 - (Y_i - f(X_i))^2 \right]
= \E_{(X_i, Y_i) \sim \TT} \left[ 2Y_i f(X_i) - f(X_i)^2 \right].
\end{align*}
By definition of $\TT$, $\TT_{Y|X}$ is a weighted average of $\PP{l}_{Y|X}$,
\begin{align*}
\RR_\TT(f) &= \E_{X_i \sim \QQ_X}\left[\sum_{l=1}^L q_l \E_{Y_i|X_i \sim \PP{l}_{Y|X}}\left[ 2Y_i f(X_i) - f(X_i)^2 \right] \right] \\
&= \sum_{l=1}^L q_l \E_{X_i \sim \QQ_X} \E_{Y_i|X_i \sim \PP{l}_{Y|X}}\left[ 2Y_i f(X_i) - f(X_i)^2 \right].
\end{align*}
We can further derive the following equation of $\RR_\TT(f)$
\[
\RR_\TT(f) = \sum_{l=1}^L q_l \E_{X_i \sim \QQ_X} \left[ 2\f{l}(X_i)f(X_i) - f(X_i)^2\right],
\]
where the equality follows from the model for each source group in \eqref{eq: source-models}.
Then the objective in \eqref{eq: f_star H} is simplified as:
\begin{align}
&\max_{f \in \FF} \min_{\TT \in \mathcal{C}(\QQ_X, \mathcal{H})}
\R_{\TT}(f)
= \max_{f \in \FF} \min_{q \in \HH}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[\f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\}.
\label{eq: thm1-obj1}
\end{align}
Suppose we can swap the order of max and min, which would be justified later, the new objective after swapping is,
\begin{equation}
\min_{q\in \HH} \max_{f \in \FF}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[\f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[f(X_i)^2 \right]\right\}.
\label{eq: thm1-obj2}
\end{equation}
Fix any $q\in \HH$, denote $f^q$ as the solution to the inner maximizer, i.e.,
\begin{equation}
f^q(\cdot) = \argmax_{f \in \FF} \left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[\f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[f(X_i)^2\right]\right\}.
\label{eq: thm1-fq-objective}
\end{equation}
For any $h(\cdot) \in \FF$ and $t\in (0,1)$, since $\FF$ is a convex set, it holds that $f^q(\cdot) + t\left( h(\cdot) - f^q(\cdot) \right) \in \FF$. And by definition of $f^q(\cdot)$, we have
\begin{multline*}
2\sum_{l=1}^L q_l \E_{\QQ_X}\left[ \f{l}(X_i)\left(f^q(X_i) + t (h(X_i) - f^q(X_i))\right) \right] - \E_{\QQ_X}\left[f^q(X_i) + t(h(X_i) - f^q(X_i))\right]^2 \\
\leq 2\sum_{l=1}^L q_l \E_{\QQ_X}\left[\f{l}(X_i)f^q(X_i)\right] - \E_{\QQ_X}\left[f^q(X_i)\right]^2.
\end{multline*}
Hence for any $t \in (0,1)$, we have
$$
2t\E_{\QQ_X}\left[\left(\sum_{l=1}^L q_l \f{l}(X_i) - f^q(X_i)\right)\left(h(X_i) - f^q(X_i)\right)\right] - t^2 \E_{\QQ_X}\left[h(X_i) - f^q(X_i)\right]^2 \leq 0
$$
By taking $t \rightarrow 0_+$, for any $h(\cdot) \in \FF$, we have:
$$
\E_{\QQ_X}\left[\left(\sum_{l=1}^L q_l \f{l}(X_i) - f^q(X_i)\right)\left(h(X_i) - f^q(X_i)\right)\right] \leq 0.
$$
We take $h(\cdot) = \sum_{l=1}^L q_l \f{l}(\cdot)$. Since $q \in \HH$ and $\f{l}(\cdot)\in \FF$ for each $l$, we have $h(\cdot) = \sum_{l=1}^L q_l \f{l}(\cdot) \in \FF$. Then we obtain
$$
\E_{\QQ_X}\left[\sum_{l=1}^L q_l \f{l}(X_i) - f^q(X_i)\right]^2 = 0,
$$
which implies
\begin{equation}
f^q(\cdot)= \sum_{l=1}^L q_l \f{l}(\cdot) \quad \textrm{almost surely}.
\label{eq: thm1-fq}
\end{equation}
Thus the optimizers for the objective \eqref{eq: thm1-obj2} is:
\begin{equation}
q^* = \argmin_{q\in \mathcal{H}} \E_{\QQ_X} \left[\sum_{l=1}^L q_l \cdot \f{l}(X_i) \right]^2 \quad \textrm{and} \quad f^*(\cdot) = \sum_{l=1}^L q_l^* \cdot \f{l}(\cdot).
\label{eq: thm1-fstar}
\end{equation}
Next we justify that the interchange of max and min by showing that $f^*(\cdot) = \sum_{l=1}^L q_l^* \cdot \f{l}(\cdot)$ is also the solution to the original objective \eqref{eq: thm1-obj1}, then the proof is completed. For any $t\in [0,1]$ and $\lambda \in \mathcal{H}$, we use the fact that $q^* + t(\lambda - q^*) \in \mathcal{H}$ and by the definition of $q^*$, we obtain that:
$$
\E_{\QQ_X} \left[ \sum_{l=1}^L (q^*_l + t (\lambda_l - q^*_l) ) \cdot \f{l}(X_i) \right]^2 \geq \E_{\QQ_X} \left[ \sum_{l=1}^L q^*_l \cdot \f{l} (X_i) \right]^2.
$$
The above inequality further implies
$$
2 t \E_{\QQ_X} \left[\left( \sum_{l=1}^L q_l^* \f{l}(X_i)\right) \left(\sum_{l=1}^L (\lambda_l - q^*_l) \f{l}(X_i) \right)\right] + t^2 \E_{\QQ_X} \left[\sum_{l=1}^L (\lambda_l - q^*_l) \f{l}(X_i)\right]^2 \geq 0.
$$
% The above inequality holds for any $t \in [0,1]$, only if for any $\lambda \in \mathcal{H}$ the following holds:
By taking $t\rightarrow 0_+$, we obtain
$$
\E_{\QQ_X} \left[\left( \sum_{l=1}^L q_l^* \f{l}(X_i)\right) \left(\sum_{l=1}^L (\lambda_l - q^*_l) \f{l}(X_i) \right)\right] \geq 0,
$$
which is equivalent as, for any $\lambda\in \HH$,
\begin{equation}
2\E_{\QQ_X} \left[ \left(\sum_{l=1}^L q_l^* \f{l}(X_i)\right) \left(\sum_{l=1}^L \lambda_l \f{l}(X_i) \right) \right] - \E_{\QQ_X} \left[\sum_{l=1}^L q_l^* \cdot \f{l}(X_i) \right]^2 \geq \E_{\QQ_X} \left[\sum_{l=1}^L q_l^* \cdot \f{l}(X_i) \right]^2.
\label{eq: thm1-interm1}
\end{equation}
For the original objective \eqref{eq: thm1-obj1}, the following inequality holds:
{\small
\begin{equation}
\begin{aligned}
&\max_{f \in \FF} \min_{q \in \HH}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\} \\
&\quad \geq \min_{q \in \mathcal{H}}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f^*(X_i) \right] - \E_{\QQ_X}\left[ f^*(X_i)^2 \right]\right\}.
\end{aligned}
\label{eq: thm1-interm2}
\end{equation}
}
Since $f^*(X_i)=\sum_{l=1}^L q_l^* \f{l}(X_i)$, we can rewrite the right hand side of \eqref{eq: thm1-interm2} as follows:
{\small
\begin{equation}
\begin{aligned}
&2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f^*(X_i) \right] - \E_{\QQ_X}\left[ f^*(X_i)^2 \right]\\
&\quad =2\E_{\QQ_X} \left[\left(\sum_{l=1}^L q_l \f{l}(X_i)\right) \left(\sum_{l=1}^L q_l^* \f{l}(X_i) \right) \right] - \E_{\QQ_X} \left[ \sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2.
\end{aligned}
\label{eq: thm1-interm3}
\end{equation}}
Combining \eqref{eq: thm1-interm2} and \eqref{eq: thm1-interm3}, we obtain:
{\small
\begin{equation*}
\begin{aligned}
& \max_{f \in \FF} \min_{q \in \HH}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\} \\
&\quad \geq \min_{q\in \mathcal{H}} \left\{2\E_{\QQ_X} \left[\left(\sum_{l=1}^L q_l \f{l}(X_i)\right) \left(\sum_{l=1}^L q_l^* \f{l}(X_i) \right) \right] - \E_{\QQ_X} \left[ \sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2\right\}.
\end{aligned}
\end{equation*}}
As \eqref{eq: thm1-interm1} holds for any $\lambda \in \HH$, we apply it to the above inequality and establish that:
\begin{equation}
\max_{f \in \FF} \min_{q \in \HH}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\} \geq \E_{\QQ_X} \left[\sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2.
\label{eq: thm1-direction-1}
\end{equation}
% \begin{align*}
% &\max_{f \in \FF} \min_{q \in \HH}
% \left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\} \\
% &\geq \min_{q \in \mathcal{H}}
% \left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f^*(X_i) \right] - \E_{\QQ_X}\left[ f^*(X_i)^2 \right]\right\} \\
% &= \min_{q\in \mathcal{H}} \left\{2\E_{\QQ_X} \left[\left(\sum_{l=1}^L q_l \f{l}(X_i)\right) \left(\sum_{l=1}^L q_l^* \f{l}(X_i) \right) \right] - \E_{\QQ_X} \left[ \sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2\right\} \\
% &\geq \E_{\QQ_X} \left[\sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2
% \end{align*}
For the other direction, since $q^* \in \HH$, it holds that
{\small
\[
\begin{aligned}
&\min_{q \in \mathcal{H}}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\}\\
&\quad \leq 2\sum_{l=1}^L q_l^* \E_{\QQ_X}\left[\f{l}(X_i) f(X_i) \right] - \E_{\QQ_X} \left[f(X_i)^2\right] \leq \max_{f\in \mathcal{F}} \left\{2\sum_{l=1}^L q_l^* \E_{\QQ_X}\left[\f{l}(X_i) f(X_i) \right] - \E_{\QQ_X} \left[f(X_i)^2\right]\right\}.
\end{aligned}
\]}
Since the above inequality holds for any $f\in \FF$, we then have
{\small
\begin{equation}
\begin{aligned}
&\max_{f \in \mathcal{F}} \min_{q \in \mathcal{H}}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\} \\
&\quad \leq \max_{f\in \mathcal{F}} \left\{2\sum_{l=1}^L q_l^* \E_{\QQ_X}\left[\f{l}(X_i) f(X_i) \right] - \E_{\QQ_X} \left[f(X_i)^2\right]\right\}.
\end{aligned}
\label{eq: thm1-interm4}
\end{equation}}
Recall the definition of $f^q(\cdot)$ in \eqref{eq: thm1-fq-objective}, the maximizer for the right hand side of the above inequality will be exactly $f^{*}$, that is defined in \eqref{eq: thm1-fq}. And by the result in \eqref{eq: thm1-fstar}, we obtain that
\begin{equation}
\max_{f\in \mathcal{F}} \left\{2\sum_{l=1}^L q_l^* \E_{\QQ_X}\left[\f{l}(X_i) f(X_i) \right] - \E_{\QQ_X} \left[f(X_i)^2\right]\right\} = \E_{\QQ_X} \left[\sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2.
\label{eq: thm1-interm5}
\end{equation}
We combine \eqref{eq: thm1-interm4} and \eqref{eq: thm1-interm5} to establish
\begin{equation}
\max_{f \in \mathcal{F}} \min_{q \in \mathcal{H}}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\} \leq \E_{\QQ_X} \left[\sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2.
\label{eq: thm1-direction-2}
\end{equation}
% \begin{align*}
% &\max_{f \in \mathcal{F}} \min_{q \in \mathcal{H}}
% \left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\} \\
% &\leq \max_{f\in \mathcal{F}} \left\{2\sum_{l=1}^L q_l^* \E_{\QQ_X}\left[\f{l}(X_i) f(X_i) \right] - \E_{\QQ_X} \left[f(X_i)^2\right]\right\}\\
% &= \E_{\QQ_X} \left[\sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2
% \end{align*}
By matching the above two directions \eqref{eq: thm1-direction-1} and \eqref{eq: thm1-direction-2}, we have
\[
\max_{f \in \mathcal{F}} \min_{q \in \mathcal{H}}
\left\{2\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[ \f{l}(X_i) f(X_i) \right] - \E_{\QQ_X}\left[ f(X_i)^2 \right]\right\} = \E_{\QQ_X} \left[\sum_{l=1}^L q_l^* \f{l}(X_i) \right]^2.
\]
Since $f^*(\cdot) = \sum_{l=1}^L q_l^* \f{l}(\cdot) \in \FF$ and $q^*\in \HH$, we can draw the conclusion that $f^*(\cdot)$ is the optimizer to the original objective.
\subsection{Proof of Theorem \ref{thm: Magging plugin}}
\label{sec: proof-thm-magging-plugin}
With $\tfH$ and $\fH$ defined in \eqref{eq: Magging plugin} and \eqref{eq: identification}, respectively, it holds that
\[
\left(\tfH(x) - \fH(x)\right)^2 = \left(\sum_{l=1}^L \widetilde{q}_l \hf{l}(x) - \sum_{l=1}^L q_l^* \f{l}(x)\right)^2 = \left(\sum_{l=1}^L \widetilde{q}_l \Del{l}(X_i) + \sum_{l=1}^L (\widetilde{q}_l - q_l^*) \f{l}(x)\right)^2.
\]
where $\Del{l}(\cdot) := \hf{l}(\cdot) - \f{l}(\cdot)$. It follows that
\[
\left(\sum_{l=1}^L \widetilde{q}_l \Del{l}(X_i) + \sum_{l=1}^L (\widetilde{q}_l - q_l^*) \f{l}(x)\right)^2 \leq 2\left(\left(\sum_{l=1}^L \widetilde{q}_l \Del{l}(x)\right)^2 + \left(\sum_{l=1}^L (\widetilde{q}_l - q_l^*) \f{l}(x)\right)^2 \right).
\]
We then apply the Cauchy-Schwarz inequality to both terms on the right-hand side of the above inequality and establish that
\[
\left(\sum_{l=1}^L \widetilde{q}_l \Del{l}(x)\right)^2\leq \sum_{l=1}^L \widetilde{q}_l^2 \sum_{l=1}^L \Del{l}(x)^2 \leq \sum_{l=1}^L \Del{l}(x)^2,
\]
where the last inequality is due to the fact that $\widetilde{q}\in \Delta^L$, and
\[
\left(\sum_{l=1}^L (\widetilde{q}_l - q_l^*) \f{l}(x)\right)^2
\leq \|\widetilde{q} - q^*\|_2^2 \sum_{l=1}^L \f{l}(x)^2.
\]
Therefore, it follows that
\[
\left(\tfH(x) - \fH(x)\right)^2 \leq 2\left(\sum_{l=1}^L \Del{l}(x)^2 + \|\widetilde{q} - q^*\|_2^2 \sum_{l=1}^L \f{l}(x)^2 \right).
\]
Then we can quantify the error of $\tfH$:
\begin{equation}
\begin{aligned}
\left\|\tfH - \fH\right\|_{\QQ,2}
&= \sqrt{\E_\QQ \left(\tfH(X) - \fH(X)\right)^2} \leq \sqrt{2\left(\sum_{l=1}^L \E_\QQ \left[\Del{l}(X)^2\right] + \|\widetilde{q} - q^*\|_2^2 \sum_{l=1}^L \E_\QQ \left[\f{l}(x)^2\right]\right)}\\
&\leq \sqrt{2L \left(\delta_n^2 + \|\widetilde{q} - q^*\|_2^2M^2\right)} \leq \sqrt{2L}\left(\delta_n + \|\widetilde{q} - q^*\|_2 M\right)
\end{aligned}
\label{eq: approx1-f_diff}
\end{equation}
To upper bound the term $\|\widetilde{q} - q^*\|_2$, we will use the following Lemma, which essentially follows the Lemma 6 in \citet{guo2020inference}. We delay its proof to Appendix \ref{sec: proof-lemma-approx}.
\begin{Lemma}
For the convex subset $\HH$ of $\Delta^L$, its diameter $\rho_{\HH}$ is defined as $\rho_\HH = \max_{q,q'\in \HH}\|q - q'\|_2$. For the semi-positive definite matrix $\Gamma$, $\widehat{\Gamma}$, we define
$q^* = \argmin_{q\in \HH} q^\intercal\Gamma q$ and $\widehat{q} = \argmin_{q\in \HH }q^\intercal \hGamma q $.
If $\lambda_{\rm min}(\Gamma) > 0$, then
\[
\|\widehat{q} - q^*\|_2 \leq \frac{L\|\widehat{\Gamma} - \Gamma\|_\infty}{\lambda_{\rm min}(\Gamma)} \wedge \rho_\HH.
\]
\label{lemma: approx}
\end{Lemma}
\noindent We apply Lemma \ref{lemma: approx} with $\hGamma=\tGamma$, and obtain that
\begin{equation}
\|\widetilde{q} - q^*\|_2 \leq \frac{L\|\tGamma - \Gamma\|_\infty}{\lambda_{\rm min}(\Gamma)} \wedge \rho_\HH.
\label{eq: approx1-qdiff}
\end{equation}
Combining \eqref{eq: approx1-f_diff} and \eqref{eq: approx1-qdiff} and \eqref{eq: approx1-Gamma_diff}, we establish that,
\begin{equation}
\left\|\tfH - \fH\right\|_{\QQ,2} \leq \sqrt{2L}\left[\delta_n + \left(\frac{L\|\tGamma - \Gamma\|_\infty}{\lambda_{\rm min}(\Gamma)} \wedge \rho_\HH\right)M\right].
\label{eq: approx1-basic}
\end{equation}
In the following discussions, we will continue the analysis of the term $\|\tGamma - \Gamma\|_\infty$.
For fixed $k,l\in [L]$, we decompose the error of $\widetilde{\Gamma}$ as follows:
\begin{equation}
\begin{aligned}
&\left({\tGamma}- \Gamma\right)_{k,l} = \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\hf{k}(X_i)\hf{l}(X_i) - \E_\QQ \f{k}(X)\f{l}(X) \\
&= \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \hf{k}(X_i)\hf{l}(X_i) - \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{k}(X_i)\f{l}(X_i)}_{(\RN{1})_{k,l}}
+ \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{k}(X_i)\f{l}(X_i) - \E_\QQ \f{k}(X)\f{l}(X)}_{(\RN{2})_{k,l}}.
\end{aligned}
\label{eq: proof-thm-plug-error-decompose}
\end{equation}
We will upper bound the error parts $(\RN{1})_{k,l}$ and $(\RN{2})_{k,l}$ in the followings.
\subsubsection{\underline{\textbf{Part $(\RN{1})_{k,l}$}}}
% \noindent{\underline{\textbf{Part $(\RN{1})_{k,l}$:}}}
We define $\Del{l}(\cdot) = \hf{l}(\cdot) - \f{l}(\cdot)$, then $(\RN{1})_{k,l}$ can be re-expressed as follows:
\begin{align*}
\left|(\RN{1})_{k,l}\right| &= \left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{k}(X_i)\f{l}(X_i) + \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{l}(X_i)\f{k}(X_i) + \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{k}(X_i)\Del{l}(X_i)\right| \\
&\leq \left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{k}(X_i)\f{l}(X_i)\right| + \left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{l}(X_i)\f{k}(X_i)\right| +\left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{k}(X_i)\Del{l}(X_i)\right| \\
&\leq \sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}(\Del{k}(X_i))^2}\sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}(\f{l}(X_i))^2} + \sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}(\Del{l}(X_i))^2}\sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}(\f{k}(X_i))^2} \\
&\quad \quad +\sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}(\Del{k}(X_i))^2}\sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}(\Del{l}(X_i))^2}.
\end{align*}
where the first inequality is due to the triangle inequality, and the second inequality is derived from the Cauchy-Schwarz inequality. Next, we apply Markov's inequality to further upper bound the terms in the above inequality. For any $t>1$, the following inequalities hold separately:
\begin{equation}
\prob\left\{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{l}(X_i)^2 \geq 2 t L \delta_n^2\right\} \leq \frac{\E_\QQ \left[\Del{l}(X)^2\right]}{2 t L \delta_n^2} \leq \frac{1}{2tL}.
\label{eq: approx1-markov1}
\end{equation}
\begin{equation}
\prob\left\{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{l}(X_i)^2 \geq 2 t L M^2 \right\} \leq \frac{\E_\QQ \left[\f{l}(X)^2\right]}{2tLM^2}\leq \frac{1}{2tL}.
\label{eq: approx1-markov2}
\end{equation}
Define the event $\Omega_1$ as
\[
\Omega_1 = \left\{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{l}(X_i)^2 \geq 2 t L \delta_n^2,\; \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{l}(X_i)^2 \geq 2 t L M^2 \quad \textrm{for some $l=1,..., L$} \right\}.
\]
Then by union bound, we obtain that $\prob(\Omega_1) \leq 1/t$. It follows that, with probability at least $1 - \frac{1}{t}$,
\begin{equation}
\max_{k,l}\left|(\RN{1})_{k,l}\right| \leq 2tL( 2 M \delta_n + \delta_n^2).
\label{eq: approx1-term1}
\end{equation}
% Next, we turn to study the part $(\RN{2})_{k,l}$, where we will apply Chebyshev's inequality to upper bound the term. $\forall a > 0$:
\subsubsection{\underline{\textbf{Part $(\RN{2})_{k,l}$}}}
We apply Chebyshev's inequality to upper bound the term.
{\small
\[
\prob\left\{\left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{k}(X_i)\f{l}(X_i) - \E_\QQ\f{k}(X)\f{l}(X)\right| \geq a \right\} \leq \frac{\textrm{Var}_\QQ(\f{k}(X)\f{l}(X))}{n_\QQ a^2} \leq \frac{\E_\QQ \left[\left(\f{k}(X)\f{l}(X)\right)^2\right]}{n_\QQ a^2},
\]}
where $\textrm{Var}_\QQ(\f{k}(X)\f{l}(X))$ is the variance for $\f{k}(X)\f{l}(X)$ evaluated on the target covariates population. With the Cauchy-Schwarz inequality, we have
\[
\E_\QQ \left[\left(\f{k}(X)\f{l}(X)\right)^2\right] \leq \sqrt{\E_\QQ\left[\f{k}(X)^4\right]} \sqrt{\E_\QQ\left[\f{l}(X)^4\right]} \leq \widetilde{M}^4.
\]
Therefore, it leads to
\[
\prob\left\{\left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{k}(X_i)\f{l}(X_i) - \E_\QQ\f{k}(X)\f{l}(X)\right| \geq a \right\} \leq \frac{\widetilde{M}^4}{n_\QQ a^2}.
\]
For $t>1$, we take $a = t L \widetilde{M}^2/ \sqrt{n_\QQ}$, then the above inequality becomes:
\begin{equation}
\prob\left\{\left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{k}(X_i)\f{l}(X_i) - \E_\QQ\f{k}(X)\f{l}(X)\right| \geq t L \frac{\widetilde{M}^2}{\sqrt{n_\QQ}}\right\} \leq \frac{1}{t^2 L^2}.
\label{eq: approx1-chebyshev1}
\end{equation}
% \begin{align*}
% &\prob\left\{\left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{k}(X_i)\f{l}(X_i) - \E_\QQ\f{k}(X)\f{l}(X)\right| \geq a \right\} \leq \frac{\textrm{Var}_\QQ(\f{k}(X)\f{l}(X))}{n_\QQ a^2}\\
% &\quad \leq \frac{\E_\QQ \left[\left(\f{k}(X)\f{l}(X)\right)^2\right]}{n_\QQ a^2}
% \leq \frac{\sqrt{\E_\QQ\left[\f{k}(X)^4\right]} \sqrt{\E_\QQ\left[\f{l}(X)^4\right]}}{n_\QQ a^2}
% \leq \frac{\widetilde{M}^4}{n_\QQ a^2}
% \end{align*} \Zijian{Have you defined $\textrm{Var}_\QQ$ and $\E_\QQ $?}
% Take $a = t \frac{\widetilde{M}^2}{\sqrt{n_\QQ}}$, we obtain that
% \begin{equation}
% \prob\left\{\left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{k}(X_i)\f{l}(X_i) - \E_\QQ\f{k}(X)\f{l}(X)\right| \geq t \frac{\widetilde{M}^2}{\sqrt{n_\QQ}}\right\} \leq \frac{1}{t^2}
% \label{eq: approx1-chebyshev1}
% \end{equation}
Define the event $\Omega_2$ as
\[
\Omega_2 = \left\{ \left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \f{k}(X_i)\f{l}(X_i) - \E_\QQ\f{k}(X)\f{l}(X)\right| \geq t L \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} \quad \textrm{for some $k,l\in [L]$}\right\}.
\]
it holds that $\prob(\Omega_2)\leq \frac{1}{t^2}$ by union bound. It follows that with probability at least $1 - \frac{1}{t^2}$,
\begin{equation}
\max_{k,l}\left|(\RN{2})_{k,l}\right| \leq t L \frac{\widetilde{M}^2}{\sqrt{n_\QQ}}.
\label{eq: approx1-term2}
\end{equation}
\subsubsection{\underline{\textbf{Proof of \eqref{eq: error plugin}}}}
Combining the inequalities in \eqref{eq: approx1-term1} and \eqref{eq: approx1-term2}, we establish that, with probability at least $1 - \frac{1}{t} - \frac{1}{t^2}$
\begin{equation}
\left\|\widetilde{\Gamma} - \Gamma\right\|_\infty \leq \max_{k,l}\left|(\RN{1})_{k,l}\right| + \max_{k,l}\left|(\RN{2})_{k,l}\right| \leq t L \cdot \left(2 \delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + 4M \delta_n\right).
\label{eq: approx1-Gamma_diff}
\end{equation}
Combining \eqref{eq: approx1-basic} and \eqref{eq: approx1-Gamma_diff}, we establish that with probability at least $1-\frac{1}{t}-\frac{1}{t^2}$:
{\small
\begin{equation}
\left\|\tfH - \fH\right\|_{\QQ,2} \leq \sqrt{2L}\left[\delta_n + \left(\frac{L\|\tGamma - \Gamma\|_\infty}{\lambda_{\rm min}(\Gamma)} \wedge \rho_\HH\right)M\right] \lesssim \delta_n + t\cdot \max\left\{\frac{M}{\lambda_{\rm min}(\Gamma)}\left(\delta_n^2+ \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + M\delta_n\right), M\rho_\HH\right\}.
\label{proof-drl-plugin-f-diff}
\end{equation}}
% \begin{align*}
% \left(\widetilde{f}(x) - f^*(x)\right)^2 &= \left(\sum_{l=1}^L \widetilde{q}_l \hf{l}(x) - \sum_{l=1}^L q_l^* \f{l}(x)\right)^2 \\
% &= \left(\sum_{l=1}^L \widetilde{q}_l \left(\hf{l}(x)-\f{l}(x)\right) + \sum_{l=1}^L (\widetilde{q}_l - q_l^*) \f{l}(x)\right)^2 \\
% &\leq 2\left(\left(\sum_{l=1}^L \widetilde{q}_l \Del{l}(x)\right)^2 + \left(\sum_{l=1}^L (\widetilde{q}_l - q_l^*) \f{l}(x)\right)^2 \right)\\
% &\leq 2\left(\sum_{l=1}^L \widetilde{q}_l^2 \sum_{l=1}^L \Del{l}(x)^2 + \|\widetilde{q} - q^*\|_2^2 \sum_{l=1}^L \f{l}(x)^2 \right)\\
% &\leq 2\left(\sum_{l=1}^L \Del{l}(x)^2 + \|\widetilde{q} - q^*\|_2^2 \sum_{l=1}^L \f{l}(x)^2 \right)
% \end{align*}
% Its expectation over $\QQ$ is upper bounded by
% \begin{equation}
% \begin{aligned}
% \left\|\widetilde{f} - f^*\right\|_{\QQ}
% &= \sqrt{\E_\QQ \left(\widetilde{f}(x) - f^*(x)\right)^2} \\
% &\leq \sqrt{2\left(\sum_{l=1}^L \E_\QQ \left[\Del{l}(X)^2\right] + \|\widetilde{q} - q^*\|_2^2 \sum_{l=1}^L \E_\QQ \left[\f{l}(x)^2\right]^2\right)}\\
% &\leq \sqrt{2L \left(\delta_n^2 + \|\widetilde{q} - q^*\|_2^2M^2\right)}\\
% &\leq \sqrt{2L}\left(\delta_n + \|\widetilde{q} - q^*\|_2 M\right)\\
% &\leq \sqrt{2}L^{1/2}\delta_n + \frac{\sqrt{2}L^{3/2}M}{\lambda_{\rm min}(\Gamma)}\left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}}+ 2M_1\delta_n\right) \cdot t
% \end{aligned}
% \label{eq: approx1-f_diff}
% \end{equation}
% \Zijian{What is $M_1$ in the above equation?}
% where we apply lemma \ref{lemma: approx} in the last inequality above.
\subsection{Proof of Theorem \ref{thm: Magging correct}}
\label{sec: proof-thm-magging-correct}
We present the relaxed version of Assumption \ref{ass: correct} without requiring the boundeness of $\w{l}(x)$ for any $x$.
\begin{Assumption}
For any $l\in [L]$ and $x\in \R^p$, $\w{l}(x)>0$. With probability larger than $1-\tau_n$ with $\tau_n\rightarrow 0,$ there exist positive sequences $\widetilde{\delta}_n,\eta_{\omega},$ and $\widetilde{\eta}_{\omega}$ such that for any $l\in [L]$, $\max\left\{\|\Del{l}\|_{\QQ,4},\|\Del{l}_\A\|_{\QQ,4},\|\Del{l}_\B\|_{\QQ,4}\right\}\leq \widetilde{\delta}_n$,
{\small
\begin{equation*}
\max\left\{\left\|\frac{\hw{l}_\A}{\w{l}} - 1\right\|_{\QQ, 2}, \left\|\frac{\hw{l}_\B}{\w{l}} - 1\right\|_{\QQ, 2}\right\}\leq \eta_{\omega},\quad
\max\left\{\left\|\frac{\hw{l}_\A}{\w{l}} - 1\right\|_{\QQ, 4},\left\|\frac{\hw{l}_\B}{\w{l}} - 1\right\|_{\QQ, 4}, \left\|\hw{l}_\A - \w{l}\right\|_{\QQ, 4}, \left\|\hw{l}_\B - \w{l}\right\|_{\QQ, 4}\right\}\leq \widetilde{\eta}_{\omega}
\end{equation*}}
\label{ass: correct-relax}
\end{Assumption}
We follow the same argument as \eqref{eq: approx1-basic}, which is derived for the plugin DRL estimator, but replace $\tGamma$ with $\hGamma$. Here, we present the analysis for the term $\left\|\hGamma - \Gamma\right\|_\infty$.
In this proof, we use subscript $i$ as the index for the observations $\{\XQ_i\}_{i\in [n_\QQ]}$ in the target data, and use subscript $j$ for the observations $\{(\X{l}_j, \Y{l}_j)\}_{j\in [n_l], l\in [L]}$ in the source data. We mainly study the error of $\hGamma_\A$ that is defined in \eqref{eq: Gamma-debias-CS-A}. With the same strategy, we can control the error of $\hGamma_\B$.
For any fixed $k,l\in [L]$, we combine \eqref{eq: error decomposition} and \eqref{eq: Gamma-debias-CS-A} and decompose $(\hGamma_\A-\Gamma)_{k,l}$ as follows:
\begin{align*}
(\widehat{\Gamma}_\mathcal{A} - \Gamma)_{k,l}
&= \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\f{k} (X_i)\f{l}(X_i)-\E_\QQ \f{k}(X)\f{l}(X)}_{(\RN{1})_{k,l}} \\
&+ \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\hf{k}_\A(X_i)\left(\hf{l}_\A(X_i) - \f{l}(X_i)\right) - \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l} \hw{l}_\A(X_j)\hf{k}_\A(X_j)\left(\hf{l}_\A(X_j) - Y_j\right)}_{(\RN{2})_{k,l}} \\
& + \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\hf{l}_\A(X_i)\left(\hf{k}_\A(X_i) - \f{k}(X_i)\right) - \frac{1}{|\mathcal{B}_k|}\sum_{j\in \mathcal{B}_k} \hw{k}_\A(X_j)\hf{l}_\A(X_j)\left(\hf{k}_\A(X_j) - Y_j\right)}_{(\RN{3})_{k,l}} \\
&- \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \left(\hf{k}_\A(X_i) - \f{k}(X_i)\right)\left(\hf{l}_\A(X_i) - \f{l}(X_i)\right)}_{(\RN{4})_{k,l}}.
\end{align*}
Subsequently, we will study the terms $(\RN{1})_{k,l}$, $(\RN{4})_{k,l}$, $(\RN{2})_{k,l}$, and $(\RN{3})_{k,l}$ one by one.
\subsubsection{\underline{\textbf{Term $(\RN{1})_{k,l}$}}}
Following from \eqref{eq: approx1-chebyshev1} and \eqref{eq: approx1-term2}, we show that with probability at least $1-\frac{1}{t^2}$, for $t>1$,
\begin{equation}
\max_{k,l}|(\RN{1})_{k,l}|\leq tL \frac{\widetilde{M}^2}{\sqrt{n_\QQ}}.
\label{proof-drl-correct-part1-final}
\end{equation}
\subsubsection{\underline{\textbf{Term $(\RN{4})_{k,l}$}}}
It follows from the Cauchy-schwarz inequality that
\[
|(\RN{4})_{k,l}| = \left|\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{k}_\A(X_i)\Del{l}_\A(X_i)\right| \leq \sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} (\Del{k}_\A(X_i))^2} \sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} (\Del{l}_\A(X_i))^2}.
\]
where we denote $\Del{l}_\A(X_i) = \hf{l}_\A(X_i) - \f{l}(X_i)$ for each $l\in [L]$.
We leverage Markov's inequality as \eqref{eq: approx1-markov1}, then by union bound, it holds that: with probability at least $1-\frac{1}{t}$,
\begin{equation}
\max_{k,l}|(\RN{4})_{k,l}|\leq tL\delta_n^2.
\label{proof-drl-correct-part4-final}
\end{equation}
\subsubsection{\underline{\textbf{Term $(\RN{2})_{k,l}$}}}
% We continue studying the part $(\RN{2})_{k,l}$.
For the observation $j\in \B_l$, we have $\Y{l}_j = \f{l}(X_j) + \eps{l}_j$, and express $(\RN{2})_{k,l}$ as follows:
\begin{equation}
(\RN{2})_{k,l} = \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \hf{k}_\A(X_i)\Del{l}_\A(X_i) - \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\hw{l}_\A(X_j)\hf{k}_\A(X_j)\Del{l}_\A(X_j) + \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\hw{l}_\A(X_j)\hf{k}_\A(X_j)\varepsilon_j^{(l)}.
\label{proof-drl-correct-part2-decompose-1}
\end{equation}
After investigating the above decomposed parts on the right-hand side, the first and the second components have the same population mean $\E_{\QQ}\hf{k}_\A(X)\Del{l}_\A(X)$, if we replace $\hw{l}_\A$ with $\w{l}$. Therefore, we can decompose the first two components in \eqref{proof-drl-correct-part2-decompose-1} as follows:
{\small
\begin{equation}
\begin{aligned}
&\underbrace{\left(\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \hf{k}_\A(X_i)\Del{l}_\A(X_i) - \E_{\QQ}\hf{k}_\A(X)\Del{l}_\A(X)\right)}_{(\RN{2}-1)_{k,l}} - \underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\w{l}(X_j)\hf{k}_\A(X_j)\Del{l}_\A(X_j) - \E_{\QQ}\hf{k}_\A(X)\Del{l}_\A(X)\right)}_{(\RN{2}-2)_{k,l}}\\
&\quad - \underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\left(\hw{l}_\A(X_j)-\w{l}(X_j)\right)\hf{k}_\A(X_j)\Del{l}_\A(X_j)\right)}_{(\RN{2}-3)_{k,l}}.
\end{aligned}
\label{proof-drl-correct-part2-decompose-2}
\end{equation}
}
The third component in \eqref{proof-drl-correct-part2-decompose-1} can be further decomposed as
{\small
\begin{equation}
\underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\w{l}(X_j)\hf{k}_\A(X_j)\varepsilon_j^{(l)}\right)}_{(\RN{2}-4)_{k,l}} +
\underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\left(\hw{l}_\A(X_j)-\w{l}(X_j)\right)\hf{k}_\A(X_j)\varepsilon_j^{(l)}\right)}_{(\RN{2}-5)_{k,l}}.
\label{proof-drl-correct-part2-decompose-3}
\end{equation}}
% We continue decomposing $(\RN{2})_{k,l}$ as follows:
% \begin{equation}
% \begin{aligned}
% (\RN{2})_{k,l} &= \underbrace{\left(\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \hf{k}_\A(X_i)\Del{l}_\A(X_i) - \E_{\QQ}\hf{k}_\A(X)\Del{l}_\A(X)\right)}_{(\RN{2}-1)_{k,l}} \\
% &\quad - \underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\w{l}(X_j)\hf{k}_\A(X_j)\Del{l}_\A(X_j) - \E_{\QQ}\hf{k}_\A(X)\Del{l}_\A(X)\right)}_{(\RN{2}-2)_{k,l}}\\
% &\quad - \underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\left(\hw{l}(X_j)-\w{l}(X_j)\right)\hf{k}_\A(X_j)\Del{l}_\A(X_j)\right)}_{(\RN{2}-3)_{k,l}}\\
% & \quad + \underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\w{l}(X_j)\hf{k}_\A(X_j)\varepsilon_j^{(l)}\right)}_{(\RN{2}-4)_{k,l}} +
% \underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\left(\hw{l}(X_j)-\w{l}(X_j)\right)\hf{k}_\A(X_j)\varepsilon_j^{(l)}\right)}_{(\RN{2}-5)_{k,l}}.
% \end{aligned}
% \label{proof-drl-correct-part2-decompose-2}
% \end{equation}
It is noteworthy that for terms $(\RN{2}-2)_{k,l}$ and $(\RN{2}-4)_{k,l}$, each observation $j$ is originally drawn from the $l$-th source population $\PP{l}$, however with the true density ratio $\w{l}$, it can be viewed as being drawn from target population $\QQ$ instead.
Due to sample splitting, both estimators $\hf{l}_\A$ and $\hw{l}_\A$ are independent of $\{(X_j, \eps{l}_j)\}_{j\in \B_l}$, thus the term $(\RN{2}-2)_{k,l}$, $(\RN{2}-4)_{k,l}$, and $(\RN{2}-5)_{k,l}$ is mean 0 by construction. In addition, the term $(\RN{2}-1)_{k,l}$ is mean 0 as well.
In the following, we will construct the upper bounds for the $(\RN{2}-1)_{k,l}$, $(\RN{2}-2)_{k,l}$, and $(\RN{2}-4)_{k,l}$, subsequently. After that, we will continue studying the components $(\RN{2}-3)_{k,l}$, and $(\RN{2}-5)_{k,l}$.
% \vspace{1mm}
\begin{itemize}
\item \underline{{\textbf{Component $(\RN{2}-1)_{k,l}$}:}}
\end{itemize}
% \noindent{\underline{{\textit{$\cdot$ Term $(\RN{2}-1)_{k,l}$}:}}}\\
The following inequality holds due to Chebyshev's inequality that for any $a_1>0$,
\[
\prob\left\{ \left|(\RN{2}-1)_{k,l}\right| \geq a_1 \right\} \leq \frac{\Var_\QQ\left(\hf{k}_\A(X)\Del{l}_\A(X)\right)}{n_\QQ a_1^2}.
\]
To upper bound $\Var_\QQ\left(\hf{k}_\A(X) \Del{l}_\A(X)\right)$, the followings hold that
\begin{equation}
\begin{aligned}
\Var_\QQ\left(\hf{k}_\A(X)\Del{l}_\A(X)\right) \leq \E_\QQ \left[\left(\hf{k}_\A(X)\Del{l}_\A(X)\right)^2\right] = \E_\QQ\left[\left(\f{k}(X)+\Del{l}_\A(X)\right)^2 \left(\Del{l}_\A(X)\right)^2\right].
% \\
% &\leq \sqrt{\E_\QQ\left[\left(\f{k}(X)+\Del{l}_\A(X)\right)^4 \right]} \sqrt{\E_\QQ\left[ \left(\Del{l}_\A(X)\right)^4 \right]} \\
% &\leq \left[\left(\E_\QQ\left[\f{k}(X)^4\right]\right)^{1/4} + \left(\E_\QQ\left[\Del{l}(X)^4\right]\right)^{1/4} \right]^2 \left[\left(\E_\QQ\left[\Del{l}(X)^4\right]\right)^{1/4} \right]^2\\
% &\leq (\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2
\end{aligned}
\label{proof-drl-correct-part2-VAR1-interm1}
\end{equation}
It follows from the Cauchy-Schwarz inequality that:
\begin{equation}
\E_\QQ\left[\left(\f{k}(X)+\Del{l}_\A(X)\right)^2 \left(\Del{l}_\A(X)\right)^2\right] \leq \sqrt{\E_\QQ\left[\left(\f{k}(X)+\Del{l}_\A(X)\right)^4 \right]} \sqrt{\E_\QQ\left[ \left(\Del{l}_\A(X)\right)^4 \right]}.
\label{proof-drl-correct-part2-VAR1-interm2}
\end{equation}
With the fourth-moment Minkowski's inequality, we have
\begin{equation}
\left(\E_\QQ\left[\left(\f{k}(X)+\Del{l}_\A(X)\right)^4\right]\right)^{1/4} \leq \left(\E_\QQ\left[\f{k}(X)^4\right]\right)^{1/4} + \left(\E_\QQ\left[\Del{l}_\A(X)^4\right]\right)^{1/4}.
\label{proof-drl-correct-part2-VAR1-interm3}
\end{equation}
Combining \eqref{proof-drl-correct-part2-VAR1-interm1},\eqref{proof-drl-correct-part2-VAR1-interm2}, and \eqref{proof-drl-correct-part2-VAR1-interm3}, we establish that
\begin{equation}
\begin{aligned}
\Var_\QQ\left(\hf{k}_\A(X)\Del{l}_\A(X)\right)
&\leq \E_\QQ \left[\left(\hf{k}_\A(X)\Del{l}_\A(X)\right)^2\right]\\
&\leq \left[\left(\E_\QQ\left[\f{k}(X)^4\right]\right)^{1/4} + \left(\E_\QQ\left[\Del{l}_\A(X)^4\right]\right)^{1/4} \right]^2 \left[\left(\E_\QQ\left[\Del{l}_\A(X)^4\right]\right)^{1/4} \right]^2 \\
&\leq (\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2.
\end{aligned}
\label{proof-drl-correct-part2-VAR1}
\end{equation}
With Chebyshev's inequality, we apply the above inequality in \eqref{proof-drl-correct-part2-VAR1} and establish that for any $a_1>0$,
\[
\prob\left\{ \left|(\RN{2}-1)_{k,l}\right| \geq a_1 \right\} \leq \frac{\Var_\QQ\left(\hf{k}_\A(X)\Del{l}_\A(X)\right)}{n_\QQ a_1^2} \leq \frac{(\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2}{n_\QQ a_1^2}.
\]
We set $a_1 = \frac{t L}{\sqrt{n_\QQ}}(\widetilde{M} + \widetilde{\delta}_n) \widetilde{\delta}_n$, then by union bound with probability $1-\frac{1}{t^2}$,
\begin{equation}
\max_{k,l}\left|(\RN{2}-1)_{k,l}\right|\leq \frac{t L}{\sqrt{n_\QQ}}(\widetilde{M} + \widetilde{\delta}_n) \widetilde{\delta}_n.
\label{proof-drl-correct-part2-1}
\end{equation}
% \vspace{1mm}
\begin{itemize}
\item \underline{{\textbf{Component $(\RN{2}-2)_{k,l}$}:}}
\end{itemize}
% \noindent{\underline{{\textit{Term $(\RN{2}-2)_{k,l}$}:}}}\\
Following the same argument in the analysis of $(\RN{2}-1)_{k,l}$, we can construct the similar arguments for the term $(\RN{2}-2)_{k,l}$. With Chebyshev's inequality, for any $a_2>0$,
\[
\prob\left\{ \left|(\RN{2}-2)_{k,l}\right| \geq a_2 \right\} \leq \frac{\Var_\QQ\left(\hf{k}_\A(X)\Del{l}_\A(X)\right)}{|\mathcal{B}_l| a_2^2} \leq \frac{(\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2}{|\mathcal{B}_l| a_2^2}.
\]
We set $a_2 = \frac{t L}{\sqrt{n/2}}(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n$, then by union bound with probability $1 - \frac{1}{t^2}$,
\begin{equation}
\max_{k,l}\left|(\RN{2}-2)_{k,l}\right| \leq \frac{tL}{\sqrt{n/2}}(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n .
\label{proof-drl-correct-part2-2}
\end{equation}
% \vspace{1mm}
\begin{itemize}
\item \underline{\textbf{Component $(\RN{2}-4)_{k,l}$}:}
\end{itemize}
% \noindent{\underline{{\textit{Term $(\RN{2}-4)_{k,l}$}:}}}\\
We apply Chebyshev's inequality, the following inequality holds for any $a_3>0$
\[
\prob\left\{ \left|(\RN{2}-4)_{k,l}\right| \geq a_3\right\} \leq \frac{\Var_\QQ\left(\hf{k}_\A(X)\varepsilon^{(l)}\right)}{|\mathcal{B}_l| a_3^2}.
\]
Now we construct the upper bound for $\Var_\QQ\left(\hf{k}_\A(X)\eps{l}\right)$:
\begin{equation*}
\Var_\QQ\left(\hf{k}_\A(X)\eps{l}\right) \leq \E_\QQ\left[(\hf{k}_\A(X)\eps{l})^2\right] = \E_\QQ\left[(\hf{k}_\A(X))^2 \E[(\eps{l})^2|X]\right] =\E_\QQ\left[(\hf{k}_\A(X))^2 \E[(\eps{l})^2]\right].
\end{equation*}
where the last inequality is due to the independence of $\eps{l}$ and target covariates.
With the assumption \ref{ass: eigen} requiring the bounded moments for $\eps{l}$, it follows that
\begin{equation}
\Var_\QQ\left(\hf{k}_\A(X)\eps{l}\right)
\leq \sigma_\varepsilon^2 \E_\QQ\left[(\hf{k}_\A(X))^2\right]
= \sigma_\varepsilon^2 \E_\QQ\left[\left(\f{k}(X) + \Del{k}_\A(X)\right)^2\right].
\label{proof-drl-correct-part2-VAR2-interm1}
\end{equation}
With the second-moment Minkowski's inequality, it holds that
\begin{equation}
\left(\E_\QQ\left[\left(\f{k}(X) + \Del{k}_\A(X)\right)^2\right]\right)^{1/2} \leq
\left(\E_\QQ\left[\f{k}(X)^2\right]\right)^{1/2} + \left(\E_\QQ\left[\Del{l}_\A(X)^2\right]\right)^{1/2}.
\label{proof-drl-correct-part2-VAR2-interm2}
\end{equation}
Combining \eqref{proof-drl-correct-part2-VAR2-interm1} and \eqref{proof-drl-correct-part2-VAR2-interm2}, we construct the following upper bound,
\begin{equation}
\Var_\QQ\left(\hf{k}_\A(X)\eps{l}\right) \leq
\sigma_\varepsilon^2\left[\left(\E_\QQ\left[\f{k}(X)^2\right]\right)^{1/2} + \left(\E_\QQ\left[\Del{l}_\A(X)^2\right]\right)^{1/2}\right]^{2} \leq \sigma_\varepsilon^2 (M + \delta_n)^2.
\label{proof-drl-correct-part2-VAR2}
\end{equation}
Now with the upper bound result in \eqref{proof-drl-correct-part2-VAR2}, we construct for any $a_3>0$,
\[
\prob\left\{ \left|(\RN{2}-4)_{k,l}\right| \geq a_3\right\} \leq \frac{\Var_\QQ\left(\hf{k}_\A(X)\varepsilon^{(l)}\right)}{|\mathcal{B}_l| a_3^2} \leq \frac{\sigma_{\varepsilon}^2(M + \delta_n)^2}{|\mathcal{B}_l| a_3^2}.
\]
Set $a_3 = \frac{tL }{\sqrt{n/2}} \sigma_\varepsilon(M + \delta_n)$. Then by union bound with probability $1 - \frac{1}{t^2}$,
\begin{equation}
\max_{k,l}\left|(\RN{2}-4)_{k,l}\right|\leq \frac{t L}{\sqrt{n/2}} \sigma_\varepsilon (M + \delta_n).
\label{proof-drl-correct-part2-4}
\end{equation}
Now in the decomposition of $(\RN{2})_{k,l}$ in \eqref{proof-drl-correct-part2-decompose-2} and \eqref{proof-drl-correct-part2-decompose-3}, the only terms left are terms $(\RN{2}-3)_{k,l}$ and $(\RN{2}-5)_{k,l}$. In the following, we will leverage the approximation error of the density ratio estimator $\hw{l}_\A$ to quantify these two terms further.
% \vspace{1mm}
\begin{itemize}
\item \underline{\textbf{Component $(\RN{2}-3)_{k,l}$:}}
\end{itemize}
% \noindent{\underline{{\textit{Term $(\RN{2}-3)_{k,l}$}:}}}\\
We re-express $(\RN{2}-3)_{k,l}$ as follows:
\[
(\RN{2}-3)_{k,l} = \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\left(\frac{\hw{l}_\A(X_j)}{\w{l}(X_j)}-1\right)\w{l}(X_j)\hf{k}_\A(X_j)\Del{l}_\A(X_j).
\]
It follows from the Cauchy-Schwarz inequality that
\begin{equation}
|(\RN{2}-3)_{k,l}| \leq \sqrt{\frac{1}{|\mathcal{B}_l|} \sum_{j\in \mathcal{B}_l} \w{l}(X_j) \left(\frac{\hw{l}_\A(X_j)}{\w{l}(X_j)}-1\right)^2} \sqrt{\frac{1}{|\mathcal{B}_l|} \sum_{j\in \mathcal{B}_l} \w{l}(X_j)\left(\hf{k}_\A(X_j)\Del{l}(X_j)\right)^2}.
\label{proof-drl-correct-part2-3-interm1}
\end{equation}
Since the density ratio estimator $\hw{l}_\A(\cdot)$ is independent of the data in $\B_l$, we have
\[
\E_{\PP{l}}\left[\frac{1}{|\mathcal{B}_l|} \sum_{j\in \mathcal{B}_l} \w{l}(X_j) \left(\frac{\hw{l}_\A(X_j)}{\w{l}(X_j)}-1\right)^2\right] = \E_{\QQ} \left(\frac{\hw{l}_\A(X_j)}{\w{l}(X_j)}-1\right)^2.
\]
Therefore, the upper bound for the first component on the right-hand side of \eqref{proof-drl-correct-part2-3-interm1} can be established by Markov's inequality. For any $a_1>0$,
\[
\prob\left\{
\frac{1}{|\mathcal{B}_l|} \sum_{j\in \mathcal{B}_l} \w{l}(X_j)\left(\frac{\hw{l}_\A(X_j)}{\w{l}(X_j)}-1\right)^2
\geq a_1
\right\} \leq \frac{\E_{\QQ} \left(\frac{\hw{l}_\A(X_j)}{\w{l}(X_j)}-1\right)^2}{a_1}
\leq \frac{\eta_\omega^2}{a_1}.
\]
The similar strategy using Markov's inequality can be applied to the second component on the right-hand side of \eqref{proof-drl-correct-part2-3-interm1}, such that for any $a_2>0$:
\[
\prob\left\{
\frac{1}{|\B_l|} \sum_{j\in \B_l} \w{l}(X_j) \left(\hf{k}_\A(X_j)\Del{l}_\A(X_j)\right)^2
\geq a_2
\right\} \leq \frac{\E_{\QQ} \left(\hf{k}_\A(X)\Del{l}_\A(X)\right)^2}{a_2}\leq \frac{(\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2}{a_2}.
\]
where the last inequality follows the result in \eqref{proof-drl-correct-part2-VAR1}. We then take $a_1 = tL \eta_\omega^2$ and $a_2 = t L(\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2$. By the union bound, with probability at least $1-\frac{2}{t}$, the following inequalities hold
\[
\max_{l\in [L]}\left|\frac{1}{|\mathcal{B}_l|} \sum_{j\in \mathcal{B}_l} \w{l}(X_j)\left(\frac{\hw{l}_\A(X_j)}{\w{l}(X_j)}-1\right)^2 \right| \leq t L \eta_\omega^2,
\]
and
\[
\max_{l\in [L]}\left|\frac{1}{|\mathcal{B}_l|} \sum_{j\in \mathcal{B}_l} \w{l}(X_j) \left(\hf{k}_\A(X_j)\Del{l}_\A(X_j)\right)^2\right| \leq t L (\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2.
\]
Combining the above two inequalities together with \eqref{proof-drl-correct-part2-3-interm1}, we establish that with probability greater than $1-\frac{2}{t}$,
\begin{equation}
\max_{k,l} |(\RN{2}-3)_{k,l}| \leq t L (\widetilde{M} + \widetilde{\delta}_n) \widetilde{\delta}_n \eta_\omega.
\label{proof-drl-correct-part2-3}
\end{equation}
% \vspace{1mm}
\begin{itemize}
\item \underline{\textbf{Component $(\RN{2}-5)_{k,l}$:}}
\end{itemize}
% \noindent{\underline{{\textit{Term $(\RN{2}-5)_{k,l}$}:}}}\\
For the term $\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\left(\hw{l}_\A(X_j)-\w{l}(X_j)\right)\hf{k}_\A(X_j)\varepsilon_j^{(l)}$, its expectation satisfies mean 0 as follows:
{\small
\[
\E_{\PP{l}}\left[\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\left(\hw{l}_\A(X_j)-\w{l}(X_j)\right)\hf{k}_\A(X_j)\varepsilon_j^{(l)}\right] = \E\left[\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\left(\hw{l}_\A(X_j)-\w{l}(X_j)\right)\hf{k}_\A(X_j)\left[\E[\varepsilon_j^{(l)}|X_j\right]\right] = 0.
\]
}
We apply Chebyshev's inequality to upper bound the term, such that $\forall a>0$:
\begin{equation*}
\prob\left\{\left|(\RN{2}-5)_{k,l}) - 0\right|\geq a\right\} \leq \frac{\Var_{\PP{l}}\left((\hw{l}_\A(X)-\w{l}(X))\hf{k}_\A(X)\varepsilon^{(l)}\right)}{|\mathcal{B}_l|a^2}.
\end{equation*}
In the following, we will establish the upper bound for the variance term. With the assumption \ref{ass: eigen}, it holds that
\begin{equation}
\Var_{\PP{l}}\left((\hw{l}_\A(X)-\w{l}(X))\hf{k}_\A(X)\varepsilon^{(l)}\right) \leq \sigma_\varepsilon^2\E_{\PP{l}}\left[\left((\hw{l}_\A(X)-\w{l}(X))\hf{k}_\A(X)\right)^2\right].
\label{proof-drl-correct-part2-5-1}
\end{equation}
The right-hand side of the above inequality can be rewritten as follows:
\[
\E_{\PP{l}}\left[\left((\hw{l}_\A(X)-\w{l}(X))\hf{k}_\A(X)\right)^2\right] = \E_{\PP{l}}\left[\w{l}(X)^2\left(\frac{\hw{l}_\A(X)}{\w{l}(X)}-1\right)^2\hf{k}_\A(X)^2\right].
\]
It follows from the Cauchy-Schwarz inequality that
{\small
\begin{equation}
\E_{\PP{l}}\left[\w{l}(X)^2\left(\frac{\hw{l}_\A(X)}{\w{l}(X)}-1\right)^2\hf{k}_\A(X)^2\right] \leq
\sqrt{\E_{\PP{l}}\left[\w{l}(X)^3\left(\frac{\hw{l}_\A(X)}{\w{l}(X)}-1\right)^4\right] \E_{\PP{l}}\left[\w{l}(X)\left(\hf{k}_\A(X)\right)^4\right]}.
\label{proof-drl-correct-part2-5-2}
\end{equation}}
We further upper bound both terms on the right hand side in \eqref{proof-drl-correct-part2-5-2}. For the first term, we have
\begin{equation}
\E_{\PP{l}}\left[\w{l}(X)^3\left(\frac{\hw{l}_\A(X)}{\w{l}(X)}-1\right)^4\right] = \E_\QQ\left[\w{l}(X)^2\left(\frac{\hw{l}_\A(X)}{\w{l}(X)}-1\right)^4\right].
\label{proof-drl-correct-part2-5-3}
\end{equation}
Since $1 + \w{l}(X)^4\geq 2\w{l}(X)^2$, we have
\begin{equation}
\E_{\QQ}\left[\w{l}(X)^2\left(\frac{\hw{l}_\A(X)}{\w{l}(X)}-1\right)^4\right]\leq \frac{1}{2}\E_\QQ\left[\left(\frac{\hw{l}_\A(X)}{\w{l}(X)}-1\right)^4 + \left(\hw{l}_\A(X)-\w{l}(X)\right)^4\right] \leq \widetilde{\eta}_\omega^4.
\label{proof-drl-correct-part2-5-3}
\end{equation}
For the other term, it holds that
\begin{equation}
\E_{\PP{l}}\left[\w{l}(X)\left(\hf{k}_\A(X)\right)^4\right] = \E_{\QQ}\left[\left(\hf{k}_\A(X)\right)^4\right] = \E_{\QQ}\left[\left(\f{k}(X) + \Del{k}_\A(X)\right)^4\right].
\label{proof-drl-correct-part2-5-4}
\end{equation}
It follows from the fourth moment Minowski's inequality that
\begin{equation}
\left(\E_{\QQ}\left[\left(\f{k}(X) + \Del{k}_\A(X)\right)^4\right]\right)^{1/4} \leq \left(\E_\QQ[\f{k}(X)^4]\right)^{1/4} + \left(\E_\QQ[\Del{k}_\A(X)^4]\right)^{1/4}\leq \widetilde{M} + \widetilde{\delta}_n.
\label{proof-drl-correct-part2-5-5}
\end{equation}
Combining \eqref{proof-drl-correct-part2-5-1},\eqref{proof-drl-correct-part2-5-2},\eqref{proof-drl-correct-part2-5-3},\eqref{proof-drl-correct-part2-5-4}, and \eqref{proof-drl-correct-part2-5-5}, we establish the upper bound for the variance:
\begin{equation*}
\Var_{\PP{l}}\left((\hw{l}_\A(X)-\w{l}(X))\hf{k}_\A(X)\varepsilon^{(l)}\right) \leq \sigma_\varepsilon^2 (\widetilde{M} + \widetilde{\delta}_n)^2\widetilde{\eta}_\omega^2.
\end{equation*}
Therefore, with Chebyshev's inequality, $\forall a>0$:
\begin{equation*}
\prob\left\{\left|(\RN{2}-5)_{k,l})\right|\geq a\right\} \leq \frac{\Var_{\PP{l}}\left((\hw{l}_\A(X)-\w{l}(X))\hf{k}_\A(X)\varepsilon^{(l)}\right)}{|\mathcal{B}_l|a^2}
\leq \frac{\sigma_\varepsilon^2 (\widetilde{M} + \widetilde{\delta}_n)^2\widetilde{\eta}_\omega^2}{|\mathcal{B}_l|a^2} .
\end{equation*}
Let $a = t L \frac{\sigma_\varepsilon (\widetilde{M} + \widetilde{\delta}_n)\widetilde{\eta}_\omega}{\sqrt{n/2}} $, and by union bound, with probability at least $1-\frac{1}{t^2}$, the following holds:
\begin{equation}
\max_{k,l}\left|(\RN{2}-5)_{k,l}\right| \leq t L \frac{\sigma_\varepsilon (\widetilde{M} + \widetilde{\delta}_n)\widetilde{\eta}_\omega}{\sqrt{n/2}}.
\label{proof-drl-correct-part2-5}
\end{equation}
Combining all components for $(\RN{2})_{k,l}$ in \eqref{proof-drl-correct-part2-1},\eqref{proof-drl-correct-part2-2},\eqref{proof-drl-correct-part2-4},\eqref{proof-drl-correct-part2-3}, and \eqref{proof-drl-correct-part2-5}, we have with probability at least $1-\frac{1}{t}-\frac{1}{t^2}$:
{\small
\begin{equation}
\begin{aligned}
\max_{k,l}|(\RN{2})_{k,l}|&\leq \max_{k,l}|(\RN{2}-1)_{k,l}|+\max_{k,l}|(\RN{2}-2)_{k,l}|+\max_{k,l}|(\RN{2}-3)_{k,l}|+\max_{k,l}|(\RN{2}-4)_{k,l}|+\max_{k,l}|(\RN{2}-5)_{k,l}|\\
&\lesssim \frac{M+\delta_n}{\sqrt{n}} + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n}{\sqrt{n\wedge n_\QQ}} + (\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n\eta_\omega + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\eta}_\omega}{\sqrt{n}}.
\end{aligned}
\label{proof-drl-correct-part2-final}
\end{equation}}
\subsubsection{\underline{\textbf{Term $(\RN{3})_{k,l}$}}}
For part $(\RN{3})_{k,l}$, we just switch the roles indices $k$ and $l$, the same arguments holds just as ones for $(\RN{2})_{k,l}$, that is, with probability at least $1-\frac{1}{t}-\frac{1}{t^2}$:
\begin{equation}
\max_{k,l}|(\RN{3})_{k,l}|\lesssim \frac{M+\delta_n}{\sqrt{n}} + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n}{\sqrt{n\wedge n_\QQ}} + (\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n\eta_\omega + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\eta}_\omega}{\sqrt{n}}.
\label{proof-drl-correct-part3-final}
\end{equation}
\subsubsection{\underline{\textbf{Proof of \eqref{eq: error correct}}}}
Till now, we've finished analyzing all parts in the error decomposition of $\hGamma^\A_{k,l}$. We combine the results for terms $(\RN{1})_{k,l}$, $(\RN{2})_{k,l}$, $(\RN{3})_{k,l}$, and $(\RN{4})_{k,l}$ in \eqref{proof-drl-correct-part1-final}, \eqref{proof-drl-correct-part2-final}, \eqref{proof-drl-correct-part3-final}, and \eqref{proof-drl-correct-part4-final}, respectively. We construct the following: with the probability at least $1-\frac{1}{t}-\frac{1}{t^2}$:
\begin{equation*}
\begin{aligned}
\left\|\hGamma^\A - \Gamma\right\|_\infty &\leq \max_{k,l}|(\RN{1})_{k,l}|+\max_{k,l}|(\RN{2})_{k,l}|+\max_{k,l}|(\RN{3})_{k,l}|+\max_{k,l}|(\RN{4})_{k,l}| \\
&\lesssim t \cdot \left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \frac{M+\delta_n}{\sqrt{n}} + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n}{\sqrt{n\wedge n_\QQ}} + (\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n\eta_\omega + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\eta}_\omega}{\sqrt{n}}\right).
\end{aligned}
\end{equation*}
Next, we switch the roles of sample split dataset $\{\A_l\}_{l\in [L]}$ and $\{\B_l\}_{l\in [L]}$, and construct the same upper bound for $\left\|\hGamma^\B - \Gamma\right\|_\infty$. By construction, $\hGamma = \frac{1}{2}(\hGamma^\A + \hGamma^\B)$. We apply the triangle inequality and construct the upper bound for the error of $\hGamma$, that is, with the probability at least $1-\frac{1}{t}-\frac{1}{t^2}$:
\begin{equation}
\left\|\hGamma - \Gamma\right\|_\infty \lesssim t \cdot \left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \frac{M+\delta_n}{\sqrt{n}} + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n}{\sqrt{n\wedge n_\QQ}} + (\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n\eta_\omega + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\eta}_\omega}{\sqrt{n}}\right)
\end{equation}
With the error of $\hGamma$ bounded, we move forward bounding the error of the bias-corrected DRL estimator $\hfH$. We follow the same argument as \eqref{eq: approx1-basic} but replace $\tGamma$ with $\hGamma$. Then with probability at least $1-\frac{1}{t}-\frac{1}{t^2}$, we have
{\small
\begin{equation}
\begin{aligned}
&\left\|\hfH - \fH\right\|_{\QQ,2} \lesssim \left[\delta_n + \left(\frac{L\|\hGamma - \Gamma\|_\infty}{\lambda_{\rm min}(\Gamma)} \wedge \rho_\HH\right)M\right] \\
&\lesssim \delta_n + t\cdot \max\left\{\frac{M}{\lambda_{\rm min}(\Gamma)}\left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \frac{M+\delta_n}{\sqrt{n}} + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n}{\sqrt{n\wedge n_\QQ}} + (\widetilde{M}+\widetilde{\delta_n})\widetilde{\delta}_n\eta_\omega + \frac{(\widetilde{M}+\widetilde{\delta_n})\widetilde{\eta}_\omega}{\sqrt{n}}\right), M\rho_\HH\right\}.
\end{aligned}
\label{proof-drl-correct-f-diff}
\end{equation}
}
% \subsection{No CS:Proof for the bias-corrected DRL in Theorem \ref{thm: Magging correct}}
% In the proof, we use subscript $i$ for the observations in \Zijian{$\QQ$ is not a data.} target data $\QQ$, and use subscript $j$ for the observations in \Zijian{$\PP{l}$ is not a data.}source data $\{\PP{l}\}_{1\leq l\leq L}$. For the $k,l$-th entry in the matrix $\widehat{\Gamma}_\mathcal{A} - \Gamma$, \Zijian{check the following sentence.} we omit the subscript $\mathcal{A}$ in the $\{\widehat{f}_\mathcal{A}^{(l)}\}_{1\leq l\leq L}$ for simplicity: \Zijian{You do not omit $\mathcal{A}$ in the following.}
% \begin{align*}
% &(\widehat{\Gamma}_\mathcal{A} - \Gamma)_{k,l} =
% \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \hf{k}(X_i)\hf{l}(X_i) -
% \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l} \hf{k}(X_j)\left(\hf{l}(X_j) - Y_j\right) \\
% &\quad \quad -\frac{1}{|\mathcal{B}_k|}\sum_{j\in \mathcal{B}_k} \hf{l}(X_j)\left(\hf{k}(X_j) - Y_j\right) - \E_{\QQ}\f{k}(X)\f{l}(X) \\
% &= \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\f{k}(X_i)\f{l}(X_i)-\E_\QQ \f{k}(X)\f{l}(X)}_{(\RN{1})_{k,l}} \\
% &+ \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\hf{k}(X_i)\left(\hf{l}(X_i) - \f{l}(X_i)\right) - \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l} \hf{k}(X_j)\left(\hf{l}(X_j) - Y_j\right)}_{(\RN{2})_{k,l}} \\
% & + \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\hf{l}(X_i)\left(\hf{k}(X_i) - \f{k}(X_i)\right) - \frac{1}{|\mathcal{B}_k|}\sum_{j\in \mathcal{B}_k} \hf{l}(X_j)\left(\hf{k}(X_j) - Y_j\right)}_{(\RN{3})_{k,l}} \\
% &- \underbrace{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \left(\hf{k}(X_i) - \f{k}(X_i)\right)\left(\hf{l}(X_i) - \f{l}(X_i)\right)}_{(\RN{4})_{k,l}}
% \end{align*}
% For part $(\RN{1})_{k,l}$, it is controlled as in \eqref{eq: approx1-chebyshev1} such that with probability at least $1 - \frac{L(L+1)}{2t^2}$:
% \begin{equation}
% \max_{k,l}\left| (\RN{1})_{k,l}\right| \leq \frac{t}{\sqrt{n_\QQ}}\widetilde{M}^2
% \label{eq: approx2-part1}
% \end{equation}
% For part $(\RN{4})_{k,l}$, we use Cauchy-schwartz inequality and Markov inequality in \eqref{eq: approx1-markov1} and obtain that with probability at least $1 - \frac{1}{t}$:
% \[
% (\RN{4})_{k,l} = \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} \Del{k}(X_i)\Del{l}(X_i) \leq \sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} (\Del{k}(X_i))^2} \sqrt{\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ} (\Del{l}(X_i))^2} \leq t\delta_n^2
% \]
% then by union bound, it holds that: with probability at least $1 - \frac{L}{t}$
% \begin{equation}
% \max_{k,l}\left|(\RN{4})_{k,l}\right| \leq t\delta_n^2
% \label{eq: approx2-part4}
% \end{equation}
% \Zijian{Similar to what?}
% The remaining parts $(\RN{2})_{k,l}$, $(\RN{3})_{k,l}$ follow the similar arguments. We use the term $(\RN{2})_{k,l}$ for illustration, which can be further decomposed into:
% \begin{align*}
% (\RN{2})_{k,l}
% &= \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\hf{k}(X_i)\left(\hf{l}(X_i) - \f{l}(X_i)\right) - \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l} \hf{k}(X_j)\left(\hf{l}(X_j) - \f{l}(X_j) - \varepsilon_j^{(l)}\right)\\
% &= \frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\hf{k}(X_i)\Del{l}(X_i) - \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\hf{k}(X_j)\Del{l}(X_j)
% + \frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\hf{k}(X_j)\varepsilon_j^{(l)}\\
% &= \underbrace{\left(\frac{1}{n_\QQ}\sum_{i=1}^{n_\QQ}\hf{k}(X_i)\Del{l}(X_i) - \E_\QQ \hf{k}(X)\Del{l}(X)\right)}_{(\RN{2}-1)_{k,l}} -
% \underbrace{\left(\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\hf{k}(X_j)\Del{l}(X_j) - \E_\QQ \hf{k}(X)\Del{l}(X) \right)}_{(\RN{2}-2)_{k,l}} \\
% & \quad \quad + \underbrace{\frac{1}{|\mathcal{B}_l|}\sum_{j\in \mathcal{B}_l}\hf{k}(X_j)\varepsilon_j^{(l)}}_{(\RN{2}-3)_{k,l}}
% \end{align*}
% \Zijian{Explain a bit why independence is useful here.}
% \Zijian{In the following, we establish the upper bounds for $\Var\left(\hf{k}(X)\Del{l}(X)\right)$ and $\Var\left(\hf{k}(X)\varepsilon^{(l)}\right)$, which further provides the upper bound for the above items $(\RN{2}-1)_{k,l}$, $(\RN{2}-2)_{k,l}$, and $(\RN{2}-3)_{k,l}$.} It can be shown that $\Var\left(\hf{k}(X)\Del{l}(X)\right)$ and $\Var\left(\hf{k}(X)\varepsilon^{(l)}\right)$ are bounded, then the above items $(\RN{2}-1)_{k,l}$, $(\RN{2}-2)_{k,l}$, and $(\RN{2}-3)_{k,l}$ can be separately bounded by \Zijian{the} Chebyshev's inequality.
% \Zijian{We start with ... and move to...}
% \begin{equation}
% \begin{aligned}
% &\Var_\QQ\left(\hf{k}(X)\Del{l}(X)\right)
% \leq \E_\QQ \left[\left(\hf{k}(X)\Del{l}(X)\right)^2\right] \\
% &= \E_\QQ\left[\left(\f{k}(X)+\Del{l}(X)\right)^2 \left(\Del{l}(X)\right)^2\right] \\
% &\leq \sqrt{\E_\QQ\left[\left(\f{k}(X)+\Del{l}(X)\right)^4 \right]} \sqrt{\E_\QQ\left[ \left(\Del{l}(X)\right)^4 \right]} \\
% &\leq \left[\left(\E_\QQ\left[\f{k}(X)^4\right]\right)^{1/4} + \left(\E_\QQ\left[\Del{l}(X)^4\right]\right)^{1/4} \right]^2 \left[\left(\E_\QQ\left[\Del{l}(X)^4\right]\right)^{1/4} \right]^2\\
% &\leq (\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2
% \end{aligned}
% \label{eq: approx2-var1}
% \end{equation}
% where we apply fourth-moment Minkowski inequality in the last second inequality above.
% \Zijian{We do not have the following assumption...}By assumption that second moment of noise is bounded by constant $C$, we have \Zijian{The following argument is more than necessary.}
% \begin{equation}
% \begin{aligned}
% &\Var_\QQ\left(\hf{k}(X)\varepsilon^{(l)}\right)\\
% &= \E_\QQ\left[\Var_\QQ\left(\hf{k}(X)\varepsilon^{(l)}\mid X\right)\right] + \Var_\QQ\left[\E_\QQ\left(\hf{k}(X)\varepsilon^{(l)}\mid X\right)\right] \quad \textrm{by Eve's law} \\
% &= \E_\QQ\left[\left(\hf{k}(X)\right)^2 \Var_\QQ\left(\varepsilon^{(l)}\mid X\right)\right] + \Var_\QQ\left[\hf{k}(X)\E_\QQ\left(\varepsilon^{(l)}\mid X\right)\right] \\
% &= \E_\QQ\left[\left(\hf{k}(X)\right)^2 \Var_\QQ\left(\varepsilon^{(l)}\mid X\right)\right] \quad \textrm{by model assumption for $\varepsilon^{(l)}$}\\
% &\leq C \E_\QQ\left[\hf{k}(X)^2\right]
% = C \E_\QQ\left[\left(\f{k}(X) + \Del{k}(X)\right)^2\right] \\
% &\leq C \left[\left(\E_\QQ\left[\f{k}(X)^2\right]\right)^{1/2} + \left(\E_\QQ\left[\Del{l}(X)^2\right]\right)^{1/2}\right]^{2} \\
% &\leq C(M + \delta_n)^2
% \end{aligned}
% \label{eq: approx2-var2}
% \end{equation}
% where we apply second-moment Minkowski inequality in the last second inequality above. Now we go back to the part $(\RN{2})_{k,l}$. Apply chebyshev's inequality and \eqref{eq: approx2-var1}, \eqref{eq: approx2-var2}, $\forall a_1, a_2, a_3 > 0$
% \[
% \prob\left\{ \left|(\RN{2}-1)_{k,l}\right| \geq a_1 \right\} \leq \frac{\Var_\QQ\left(\hf{k}(X)\Del{l}(X)\right)}{n_\QQ a_1^2} \leq \frac{(\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2}{n_\QQ a_1^2}
% \]
% \[
% \prob\left\{ \left|(\RN{2}-2)_{k,l}\right| \geq a_2 \right\} \leq \frac{\Var_\QQ\left(\hf{k}(X)\Del{l}(X)\right)}{|\mathcal{B}_l| a_2^2} \leq \frac{(\widetilde{M} + \widetilde{\delta}_n)^2 \widetilde{\delta}_n^2}{|\mathcal{B}_l| a_2^2}
% \]
% \[
% \prob\left\{ \left|(\RN{2}-3)_{k,l}\right| \geq a_3\right\} \leq \frac{\Var_\QQ\left(\hf{k}(X)\varepsilon^{(l)}\right)}{|\mathcal{B}_l| a_3^2} \leq \frac{C(M + \delta_n)^2}{|\mathcal{B}_l| a_3^2}
% \]
% Set $a_1 = \frac{t}{\sqrt{n_\QQ}}(\widetilde{M} + \widetilde{\delta}_n) \widetilde{\delta}_n$, $a_2 = \frac{t}{\sqrt{|\mathcal{B}_l|}}(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n$, and $a_3 = \frac{t}{\sqrt{|\mathcal{B}_l|}} \sqrt{C}(M + \delta_n)$ in the above three inequalities, respectively. Then with probabililty $1 - \frac{1}{t^2}$ separately, \Zijian{the following is not rigorous, please write $\lesssim$}
% \begin{equation}
% \left|(\RN{2}-1)_{k,l}\right|\leq \frac{t}{\sqrt{n_\QQ}}(\widetilde{M} + \widetilde{\delta}_n) \widetilde{\delta}_n;
% \quad
% \left|(\RN{2}-2)_{k,l}\right| \leq \frac{t}{\sqrt{|\mathcal{B}_l|}}(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n;
% \quad
% \left|(\RN{2}-3)_{k,l}\right|\leq \frac{t}{\sqrt{|\mathcal{B}_l|}} \sqrt{C} (M + \delta_n)
% \label{eq: approx2-part2-sep}
% \end{equation}
% thus under probability at least $1 - \frac{3L(L+1)}{2t^2}$:
% \[
% \begin{aligned}
% \max_{k,l}|(\RN{2})_{k,l}| &\leq \max_{k,l}|(\RN{2}-1)_{k,l}| + \max_{k,l}|(\RN{2}-2)_{k,l}| + \max_{k,l}|(\RN{2}-3)_{k,l}|\\
% &\leq \frac{t}{\sqrt{n_\QQ}}(\widetilde{M} + \widetilde{\delta}_n) \widetilde{\delta}_n + \frac{t}{\sqrt{|\mathcal{B}_l|}}(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n + \frac{t}{\sqrt{|\mathcal{B}_l|}} \sqrt{C} (M + \delta_n) \\
% &\leq \sqrt{2}t\cdot\left(\frac{(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n}{\sqrt{\min(n_\QQ, n)}} + \frac{\sqrt{C} (M + \delta_n)}{\sqrt{n}} \right)
% \end{aligned}
% \]
% where $n := \min_{1\leq l\leq L} n_l$ and $|\mathcal{B}_l| = \frac{n_l}{2}$. Similary, for the part $(\RN{3})_{k,l}$
% \[
% \begin{aligned}
% \max_{k,l}|(\RN{3})_{k,l}|
% &\leq \frac{t}{\sqrt{n_\QQ}}(\widetilde{M} + \widetilde{\delta}_n) \widetilde{\delta}_n + \frac{t}{\sqrt{|\mathcal{B}_k|}}(\widetilde{M} + \widetilde{\delta}_n) \widetilde{\delta}_n + \frac{t}{\sqrt{|\mathcal{B}_k|}} \sqrt{C}(M + \delta_n) \\
% &\leq \sqrt{2}t\cdot\left(\frac{(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n}{\sqrt{\min(n_\QQ, n)}} + \frac{\sqrt{C} (M + \delta_n)}{\sqrt{n}} \right)
% \end{aligned}
% \]
% Till now, we can upper bound the approximation error $\|\widehat{\Gamma}_\mathcal{A} - \Gamma\|_\infty$:
% \[
% \begin{aligned}
% \|\widehat{\Gamma}_\mathcal{A} - \Gamma\|_\infty
% &\leq
% \max_{k,l}|(\RN{1})_{k,l}| + \max_{k,l}|(\RN{2})_{k,l}| + \max_{k,l}|(\RN{3})_{k,l}| + \max_{k,l}|(\RN{4})_{k,l}|\\
% &\lesssim
% t\cdot \left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \frac{(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n}{\sqrt{\min(n_\QQ, n)}} + \frac{M + \delta_n}{\sqrt{n}}\right)
% \end{aligned}
% \]
% similarly, we can bound $\|\widehat{\Gamma}_\mathcal{B} - \Gamma\|_\infty$. Then our bias-corrected estimator $\widehat{\Gamma} = \frac{1}{2}\left(\widehat{\Gamma}_\mathcal{A} + \widehat{\Gamma}_\mathcal{B}\right)$ has approximation error:
% \[
% \begin{aligned}
% \left\|\widehat{\Gamma} - \Gamma\right\|_\infty &= \left\|\frac{1}{2}\left(\widehat{\Gamma}_\mathcal{A} - \Gamma\right) + \frac{1}{2}\left(\widehat{\Gamma}_\mathcal{B} - \Gamma\right)\right\|_\infty \leq
% \left\|\frac{1}{2}\left(\widehat{\Gamma}_\mathcal{A} - \Gamma\right)\right\|_\infty + \left\|\frac{1}{2}\left(\widehat{\Gamma}_\mathcal{B} - \Gamma\right)\right\|_\infty \\
% &\lesssim
% t\cdot \left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \frac{(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n}{\sqrt{\min(n_\QQ, n)}} + \frac{M + \delta_n}{\sqrt{n}}\right)
% \end{aligned}
% \]
% with probability at least $1 - \frac{c_1}{t} - \frac{c_2}{t^2}$. And following the approach similar with \eqref{eq: approx1-f_diff}:
% \[
% \left\|\widehat{f} - f^*\right\|_{\QQ,2} \lesssim \delta_n + \frac{M}{\lambda_{\rm min}(\Gamma)}\left(\delta_n^2 + \frac{\widetilde{M}^2}{\sqrt{n_\QQ}} + \frac{(\widetilde{M} + \widetilde{\delta}_n)\widetilde{\delta}_n}{\sqrt{\min(n_\QQ, n)}} + \frac{M + \delta_n}{\sqrt{n}}\right) \cdot t
% \]
\subsection{Proof of Corollary \ref{cor: high-dim}}
\label{sec: appendix-density}
We start with the analysis of $\delta_n$ and $\widetilde{\delta}_n$. To simplify the discussion, we only present the analysis of $\widehat{\beta}^{(l)}$ fitted over the full data, and the estimators $\widehat{\beta}^{(l)}_\A$ and $\widehat{\beta}^{(l)}_\B$ fitted on the subset data $\A_l$ and $\B_l$ shall follow the same argument. and Known that $\XQ\sim \QQ_X$ is sub-Gaussian random vector, as in the assumption \ref{ass: linear}, and $\XQ$ is independent of $\widehat{\beta}^{(l)} - \beta^{(l)}$, we establish that
$(\XQ)^\intercal \frac{\widehat{\beta}^{(l)} - \beta^{(l)}}{\|\widehat{\beta}^{(l)} - \beta^{(l)}\|_2} $
is a sub-Gaussian random variable, following the sub-Gaussian random vector definition (Def 5.22) in the book \citet{vershynin2010introduction}. Then for all $q\geq 1$:
{\small
\[
\left(\E_{\QQ_X}\left[\left((\XQ)^\intercal \frac{\widehat{\beta}^{(l)} - \beta^{(l)}}{\|\widehat{\beta}^{(l)} - \beta^{(l)}\|_2}\right)^q\right]\right)^{1/q}\leq C\sqrt{q}.
\]}
Since $\widehat{\beta}^{(l)}$ satisfies \eqref{eq: beta-rate}, the above inequality implies that
\[
\|(\XQ)^\intercal(\widehat{\beta}^{(l)} - \beta^{(l)})\|_{\QQ,2}\leq \sqrt{2}C \|\widehat{\beta}^{(l)} - \beta^{(l)}\|_2 \lesssim \sqrt{\frac{s_\beta \log p}{n_l}},
\]
and
\[
\|(\XQ)^\intercal(\widehat{\beta}^{(l)} - \beta^{(l)})\|_{\QQ,4}\leq 2C \|\widehat{\beta}^{(l)} - \beta^{(l)}\|_2 \lesssim \sqrt{\frac{s_\beta \log p}{n_l}}.
\]
Therefore,
\[
\max\{\delta_n\;, \; \widetilde{\delta}_n\} \lesssim \sqrt{s_\beta \log p/n}.
\]
\noindent Next, we move forward to the analysis of $\eta_\omega$ and $\widetilde{\eta}_\omega$. For simiplicity, we only analyze the density ratio estimator $\widehat{\omega}^{(l)}_\A$ fitted on the split sample subset $\A_l$. We replace the notations $\widehat{\gamma}^{(l)}_\A, \widehat{\omega}^{(l)}_\A$ with $\widehat{\gamma}^{(l)},\widehat{\omega}^{(l)}$ for the use of convenience.
Assuming that the density ratio models in \eqref{eq: density ratio} are correctly specified, we follow the regime in the \citet{qin1998inferences}, and denote $\pi(x)$ as the marginal density for the random vector $\widetilde{X}^{(l)}_i$. We define $p_l$ to be the marginal probability of $\{G^{(l)}_i = 0\}$ in the population, i.e., $p_l = p(G^{(l)}_i = 0) = \int p(G^{(l)}_i = 0|\widetilde{X}^{(l)}_i = x) \pi(x) dx$. Due to the Assumption \ref{ass: linear}, $p_l \in [c_{\rm min}, \frac{1}{2})$ by definition. In the following discussion, we treat $p_l$ as an unknown constant satisfying $p_l \in (0,\frac{1}{2})$ without reaching the boundaries.
% We then denote the conditional density functions of $\widetilde{X}^{(l)}_i$ given $G^{(l)}_i$ by
% \[
% p(\widetilde{X}^{(l)}_i = x| G^{(l)}_i = 0) = d F_0(x)/dx, \quad p(\widetilde{X}^{(l)}_i = x| G^{(l)}_i = 1) = d F_1(x)/dx
% \]
% By Bayes' rule, we have
% \[
% p(x| G^{(l)}_i = 0) = \frac{1}{p_l(1+\exp(x^\intercal \gamma^{(l)}))} f(x),\quad
% p(x| G^{(l)}_i = 1)= \frac{\exp(x^\intercal \gamma^{(l)})}{(1-p_l)(1+\exp(x^\intercal \gamma^{(l)}))} f(x)
% \]
\noindent We can see that the density ratio is:
\[
\frac{p(x| G^{(l)}_i = 1)}{p(x| G^{(l)}_i = 0)} = \frac{p_l}{1-p_l} \exp(x^\intercal \gamma^{(l)}).
\]
The random variable $G_i^{(l)}$ for $i\in [n_G]$ with $n_G = |\A_l| + n_\QQ$ is bernoulli distributed with $p(G_i^{(l)}=0) = p_l$ and $p(G_i^{(l)}=1) = 1-p_l$. We consider the independent $U_i \sim {\rm Unif}(0,1)$ for $i\in [n_G]$, then $n_l = \sum_{i=1}^{n_G}\ind_{[U_i\leq p_l]}$.
\noindent Again, known that $\XQ\sim \QQ_X$ is sub-Gaussian random vector,
it holds that
$(\XQ)^\intercal \frac{\widehat{\gamma}^{(l)} - \gamma^{(l)}}{\|\widehat{\gamma}^{(l)} - \gamma^{(l)}\|_2} $
is a sub-Gaussian random variable. Thus for all $q\geq1$:
{\small
\[
\left(\E_{\QQ_X}\left[\left((\XQ)^\intercal \frac{\widehat{\gamma}^{(l)} - \gamma^{(l)}}{\|\widehat{\gamma}^{(l)} - \gamma^{(l)}\|_2}\right)^q\right]\right)^{1/q}\leq C\sqrt{q}.
\]}
The above inequality implies that
\begin{equation*}
\|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,q}\leq \sqrt{q}C \|\widehat{\gamma}^{(l)} - \gamma^{(l)}\|_2.
\end{equation*}
Together with $\widehat{\gamma}^{(l)}$ satisfying \eqref{eq: gamma-rate}, we know that for all $q\geq 1$,
\begin{equation}
\|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,q}\leq \sqrt{q}C \|\widehat{\gamma}^{(l)} - \gamma^{(l)}\|_2 \lesssim \sqrt{\frac{s_\gamma \log p}{n_G}} \longrightarrow 0,
\label{proof-coro-linear-5}
\end{equation}
as $n_G\rightarrow \infty$.
\noindent For the term $\frac{\hw{l}_\A(x)}{\w{l}(x)} - 1$, we analyze it as follows:
\begin{align*}
\frac{\hw{l}(x)}{\w{l}(x)} - 1 &= \frac{\frac{|\A_l|}{n_\QQ}\exp(x^\intercal \widehat{\gamma}^{(l)})}{\frac{p_l}{1-p_l}\exp(x^\intercal \gamma^{(l)})} - 1
= \frac{\frac{|\A_l|}{n_\QQ}\exp(x^\intercal \widehat{\gamma}^{(l)}) - \frac{p_l}{1-p_l}\exp(x^\intercal \gamma^{(l)})}{\frac{p_l}{1-p_l}\exp(x^\intercal \gamma^{(l)})} \\
&= \frac{\frac{|\A_l|}{n_\QQ}\left(\exp(x^\intercal \widehat{\gamma}^{(l)}) - \exp(x^\intercal\gamma^{(l)})\right) + \left(\frac{|\A_l|}{n_\QQ} - \frac{p_l}{1-p_l}\right)\exp(x^\intercal \gamma^{(l)})}{\frac{p_l}{1-p_l}\exp(x^\intercal \gamma^{(l)})} \\
&= \frac{|\A_l|/n_\QQ}{p_l/(1-p_l)}\left[\exp\left(x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right) - 1\right] + \frac{|\A_l|/n_\QQ - p_l/(1-p_l)}{p_l/(1-p_l)} .
\end{align*}
It follows from the triangle inequality that
\begin{equation}
\left|\frac{\hw{l}(x)}{\w{l}(x)} - 1\right|\leq \left|\frac{|\A_l|/n_\QQ}{p_l/(1-p_l)}\right|\left|\exp\left(x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right) - 1\right| + \left|\frac{|\A_l|/n_\QQ - p_l/(1-p_l)}{p_l/(1-p_l)}\right|.
\label{proof-coro-linear-6}
\end{equation}
\begin{Lemma}
With probability greater than $1-2c_1$, for $c_1\in (0,\frac{1}{2})$, if $n_G > \frac{\log(1/c_1)}{\min\{(1-p_l)^2, p_l^2\}}$, it holds that
\[
\left|\frac{|\A_l|}{n_\QQ} - \frac{p_l}{1-p_l}\right| \leq C_1 \sqrt{\frac{1}{n_G}},
\]
where $C_1 = \frac{\sqrt{\log(1/c_1)}}{(\sqrt{2}-1)(1-p_l)^2}$.
\label{lemma: proof-coro-linear}
\end{Lemma}
\noindent Applying the above Lemma \ref{lemma: proof-coro-linear} together with the inequality in \eqref{proof-coro-linear-6}, we establish that
{\small
\begin{equation}
\left|\frac{\hw{l}(x)}{\w{l}(x)} - 1\right|\leq \left(1+C_1\sqrt{\frac{1}{n_G}}\right)\left|\exp\left(x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right) - 1\right| + \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}}.
\label{proof-coro-linear-3}
\end{equation}}
We continue upper-bounding the term $\left|\exp\left(x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right) - 1\right|$ in the above inequality \eqref{proof-coro-linear-3}.
Due to the Mean Value Theorem,
\[
\exp\left(x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right) - 1 = \theta \cdot x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)}),
\]
where $\theta$ takes some value between $0$ and $\exp(x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)}))$, and then
\begin{equation}
\left|\exp\left(x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right) - 1\right| = \left|\theta\right| \cdot \left|x^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right|.
\label{proof-coro-linear-7}
\end{equation}
{
Since $(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})$ is subGaussian, we obtain that for all $t>0$,
\[
\prob\left\{\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right| > t\right\} \leq \exp(1- t^2/K_1^2) .
\]
and $K_1\lesssim \sqrt{\frac{s_\gamma \log p}{n_G}}$ by the equivalence of sub-Gaussian property and the result in \eqref{proof-coro-linear-5}.
Taking $t = \log 2$, we establish that with probability greater than $1-\exp(1-(\log 2)^2/K_1^2)$, the following event holds.
\[
\Omega_1 = \left\{\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right| \leq \log 2\right\}.
\]
i.e., $\prob(\Omega_1)\geq 1-\exp(1-(\log 2)^2/K_1^2) = 1-p^{-c} - \exp(-cn)$.
\underline{Conditional on event $\Omega_1$}, we have
\[
|\theta| \leq \exp(\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right|)\leq 2.
\]
Then, it follows from \eqref{proof-coro-linear-3} and \eqref{proof-coro-linear-7} that
{\small
\begin{equation}
\left|\frac{\hw{l}(\XQ)}{\w{l}(\XQ)} - 1\right|\leq 2\left(1+C_1\sqrt{\frac{1}{n_G}}\right)\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right| + \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}}.
\label{eq: proof-coro-linear-8}
\end{equation}}
Therefore, its $\ell_2$ and $\ell_4$ norms satisfy:
\begin{equation}
\begin{aligned}
&\left\|\frac{\hw{l}}{\w{l}} - 1\right\|_{\QQ,2} \leq 2\left(1+C_1\sqrt{\frac{1}{n_G}}\right)
\left(\|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,2}\right) + \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}}\\
&\left\|\frac{\hw{l}}{\w{l}} - 1\right\|_{\QQ,4} \leq 2\left(1+C_1\sqrt{\frac{1}{n_G}}\right)
\left(\|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,4}\right) + \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}}
\end{aligned}
\end{equation}
Combining \eqref{proof-coro-linear-5}, we establish that
\begin{equation}
\max\left\{\left\|\frac{\hw{l}}{\w{l}} - 1\right\|_{\QQ,2}\;,\;\left\|\frac{\hw{l}}{\w{l}} - 1\right\|_{\QQ,4}\right\}\lesssim \sqrt{\frac{s_\gamma \log p}{n_G}} \lesssim \sqrt{\frac{s_\gamma \log p}{n + n_\QQ}}.
\end{equation}
Next, we upper bound the term $\left\|\hw{l}-\w{l}\right\|_{\QQ,4}$. By assumption \ref{ass: linear}, with probability larger than $1-p^{-c_0}$, the following holds:
\[
\w{l}(x) = \frac{p_l}{1-p_l} \exp(x^\top \gamma^{(l)}) \leq \frac{p_l}{1-p_l} \frac{1-c_{\rm min}}{c_{\rm min}}.
\]
It implies that
{\small
\begin{equation}
\left\|\hw{l}-\w{l}\right\|_{\QQ,4}= \left\|\w{l}\left(\frac{\hw{l}}{\w{l}}-1\right)\right\|_{\QQ,4} \leq
\frac{p_l}{1-p_l} \frac{1-c_{\rm min}}{c_{\rm min}} \left\|\frac{\hw{l}}{\w{l}} - 1\right\|_{\QQ,4} \lesssim \sqrt{\frac{s_\gamma \log p}{n_G}}\lesssim \sqrt{\frac{s_\gamma \log p}{n + n_\QQ}}.
\label{proof-coro-linear-density-2}
\end{equation}}
Therefore, we show that
\[
\max\{\eta_\omega\;,\; \widetilde{\eta}_\omega\}\lesssim \sqrt{s_\gamma \log p/(n+n_\QQ)}.
\]
}
% Therefore,
% {\small
% \begin{equation}
% \begin{aligned}
% \left\|\left.\frac{\hw{l}_\A(\XQ)}{\w{l}(\XQ)} - 1\right\vert \Omega_1 \right\|_{\QQ,2}&\leq c\log(n_G)\left(1+C_1\sqrt{\frac{1}{n_G}}\right)\|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,2}+ \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}}\\
% &\lesssim c\log(n_G)\left(1+C_1\sqrt{\frac{1}{n_G}}\right)\sqrt{\frac{s_\gamma\log p}{n_G}} + \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}} \\
% &\lesssim \log(n_G)\sqrt{\frac{s_\gamma \log p}{n_G}}.
% \end{aligned}
% \end{equation}
% }
% Conditioning on event $\Omega_1^c$, since $\theta$
% \Zhenyu{\color{red} the following will be deleted.}
% It follows from Markov's inequality that
% \[
% \prob\left\{\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right| > \log 2\right\} \leq
% \frac{\E_{\QQ}\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right|}{\log 2}
% \]
% By taking $q=1$ in \eqref{proof-coro-linear-5}, we obtain that
% \[
% \E_{\QQ}\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right|:= \|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,1} \leq C_2 \sqrt{\frac{s_\gamma \log p}{n_G}}.
% \]
% for some constant $C_2$.
% If $n_G > ({s_\gamma \log p})^2(C_2/{c})^2$,
% combining the above two inequalities, we establish that
% \[
% \prob\left\{\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right| > \log 2\right\}\leq \frac{C_2}{\log 2}\sqrt{\frac{s_\gamma \log p}{n_G}}.
% \]
% with probability at least $1 - c$, the following event holds
% \begin{equation*}
% \left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right| \leq 2
% \end{equation*}
% Since $|\theta| \leq \left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right|$, we combine the above inequality and \eqref{proof-coro-linear-7} to establish
% \begin{equation}
% \left|\exp\left((\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right) - 1\right|\leq 2\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right|.
% \label{proof-coro-linear-4}
% \end{equation}
% % Considering \eqref{proof-coro-linear-5}, together with Markov's inequality,
% % \[ \prob\left((\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)}) > 2\right) \leq \frac{\|(\XQ)^\intercal (\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,1}}{2}\longrightarrow 0.
% % \]
% % Therefore, with a high probability, as $n_G$ is sufficiently large, $(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)}) \leq 2$.
% Combining \eqref{proof-coro-linear-3} and \eqref{proof-coro-linear-4}, we obtain that, with high probability greater than $1 - c$,
% \[
% \left|\frac{\hw{l}_\A(\XQ)}{\w{l}(\XQ)} - 1\right|\leq
% 2\left(1+C_1\sqrt{\frac{1}{n_G}}\right)\left(\left|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\right|\right) + \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}} .
% \]
% Next, we derive the upper bound for the $\ell_2$ and $\ell_4$ norm of the term $\left|\frac{\hw{l}_\A}{\w{l}} - 1\right|$ over the target distribution $\QQ_X$.
% \begin{align*}
% \left\|\frac{\hw{l}}{\w{l}} - 1\right\|_{\QQ,2} \leq 2\left(1+C_1\sqrt{\frac{1}{n_G}}\right)
% \left(\|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,2}\right) + \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}},
% \end{align*}
% and
% \begin{align*}
% \left\|\frac{\hw{l}_\A}{\w{l}} - 1\right\|_{\QQ,4} \leq 2\left(1+C_1\sqrt{\frac{1}{n_G}}\right)
% \left(\|(\XQ)^\intercal(\widehat{\gamma}^{(l)} - \gamma^{(l)})\|_{\QQ,4}\right) + \frac{C_1}{p_l/(1-p_l)}\sqrt{\frac{1}{n_G}} .
% \end{align*}
% Combining \eqref{proof-coro-linear-5}, we establish that
% \begin{equation}
% \max\left\{\left\|\frac{\hw{l}_\A}{\w{l}} - 1\right\|_{\QQ,2}\;,\;\left\|\frac{\hw{l}_\A}{\w{l}} - 1\right\|_{\QQ,4}\right\}\lesssim \sqrt{\frac{s_\gamma \log p}{n_G}} \lesssim \sqrt{\frac{s_\gamma \log p}{n + n_\QQ}}.
% \label{proof-coro-linear-density-1}
% \end{equation}
% Next, we upper bound the term $\left\|\hw{l}_\A-\w{l}\right\|_{\QQ,4}$. By assumption \ref{ass: linear}, with probability larger than $1-p^{-c_0}$, the following holds:
% \[
% \w{l}(x) = \frac{p_l}{1-p_l} \exp(x^\top \gamma^{(l)}) \leq \frac{p_l}{1-p_l} \frac{1-c_{\rm min}}{c_{\rm min}}.
% \]
% It implies that
% {\small
% \begin{equation}
% \left\|\hw{l}_\A-\w{l}\right\|_{\QQ,4}= \left\|\w{l}\left(\frac{\hw{l}_\A}{\w{l}}-1\right)\right\|_{\QQ,4} \leq
% \frac{p_l}{1-p_l} \frac{1-c_{\rm min}}{c_{\rm min}} \left\|\frac{\hw{l}_\A}{\w{l}} - 1\right\|_{\QQ,4} \lesssim \sqrt{\frac{s_\gamma \log p}{n_G}}\lesssim \sqrt{\frac{s_\gamma \log p}{n + n_\QQ}}.
% \label{proof-coro-linear-density-2}
% \end{equation}}
% The same arguments can be established for $\hw{l}_\B$ as $\hw{l}_\A$ in \eqref{proof-coro-linear-density-1} and \eqref{proof-coro-linear-density-2}.
% Therefore, we show that
% \[
% \max\{\eta_\omega\;,\; \widetilde{\eta}_\omega\}\lesssim \sqrt{s_\gamma \log p/(n+n_\QQ)}.
% \]
\subsection{Proof of Theorem \ref{prop: reward-diff}}
We care about the reward difference of the population version DRL $\fH$ and empirical version DRL $\hfH$. For any fixed target distribution $\QQ$:
\[
\begin{aligned}
\RR_\QQ(\fH) - \RR_\QQ(\hfH) &= \E_\QQ(Y - \fH(X))^2 - \E_\QQ(Y - \hfH(X))^2\\
&= \E_\QQ(Y - \fH(X) + \fH(X) - \hfH(X))^2 - \E_\QQ(Y-\fH(X))^2 \\
&= 2\E_\QQ\left[(Y-\fH(X))(\fH(X) - \hfH(X))\right] + \E_\QQ(\fH(X) - \hfH(X))^2.
\end{aligned}
\]
We rewrite the target outcome as follows: $Y = f^\QQ(X) + \epsilon^\QQ$ with $f_\QQ(X) = \E_{\QQ}[Y|X]$ and $\E_\QQ[\epsilon^\QQ|X] = 0$.
Thus
\[
\E_\QQ\left[(Y-\fH(X))(\fH(X) - \hfH(X))\right] =
\E_\QQ\left[(f^\QQ(X)-\fH(X))(\fH(X) - \hfH(X))\right].
\]
By the Cauchy-Schwarz inequality, it holds that
\[
\E_\QQ\left[(f^\QQ(X)-\fH(X))(\fH(X) - \hfH(X))\right] \leq \sqrt{\E_\QQ[(f^\QQ(X)-\fH(X))^2]}\sqrt{\E_\QQ[(\fH(X) - \hfH(X))^2]}.
\]
It implies that
\[
\begin{aligned}
\left|\RR_\QQ(\fH) - \RR_\QQ(\hfH)\right|&\leq 2\sqrt{\E_\QQ[(f^\QQ(X)-\fH(X))^2]}\|\hfH - \fH\|_{\QQ,2} + \|\hfH - \fH\|_{\QQ,2}^2\\
&= 2\|f^\QQ - \fH\|_{\QQ,2}\cdot \|\hfH - \fH\|_{\QQ,2} + \|\hfH - \fH\|_{\QQ,2}^2.
\end{aligned}
\]
% We further study the term $\E_\QQ[(f^\QQ(X)-\fH(X))^2]$.
% As $f^\QQ(x) = \sum_{l=1}^L q_l^0 \f{l}(x)$ and $\fH(x) = \sum_{l=1}^L q_l^* \f{l}(x)$, we can rewrite $ \E_\QQ[(f^\QQ(X) - \fH(X))^2]$ as follows:
% \[
% \begin{aligned}
% \E_\QQ[(f^\QQ(X) - \fH(X))^2] = \E_\QQ\left[\left(\sum_{l=1}^L (q_l^0 - q_l^*)\f{l}(X)\right)^2\right].
% \end{aligned}
% \]
% \Zijian{This is very loose. We should not continue like this.}We further have
% % \[
% % \E_\QQ\left[\left(\sum_{l=1}^L (q_l^0 - q_l^*)\f{l}(X)\right)^2\right] \leq \sum_{l=1}^L (q_l^0 - q_l^*)^2 \sum_{l=1}^L \E_{\QQ}[\f{l}(X)]^2 \leq L \sum_{l=1}^L (q_l^0 - q_l^*)^2 M^2 \leq 2L M^2
% % \]
% \[
% \E_\QQ\left[\left(\sum_{l=1}^L (q_l^0 - q_l^*)\f{l}(X)\right)^2\right] \leq \E_\QQ\left[L\left(\sum_{l=1}^L (q_l^0 - q_l^*)^2 (\f{l}(X))^2\right)\right]\leq L\|q^0 - q^*\|_2^2 M^2.
% \]
% Therefore, we conclude that
% \[
% \begin{aligned}
% \left|\RR_\QQ(\fH) - \RR_\QQ(\hfH)\right|\leq 2\sqrt{L}M\|q^0 - q^*\|_2\|\hfH - \fH\|_{\QQ,2} + \|\hfH - \fH\|_{\QQ,2}^2.
% \end{aligned}
% \]
\subsection{Proof of Proposition \ref{prop: shrink H}}
Since $\QQ_{Y|X} = \sum_{l=1}^L q_l^0 \cdot \PP{l}_{Y|X}$, the conditional outcome model for target distribution satisfies:
\begin{equation}
f^\QQ(x) = \E_{\QQ}[Y|X=x] = \sum_{l=1}^L q_l^0 \cdot \E_{\PP{l}}[Y|X=x] = \sum_{l=1}^L q_l \cdot \f{l}(x)
\label{eq: prop-shinkH-1}
\end{equation}
It follows from the identification theorem \ref{thm: identification} that
% \begin{align*}
% \|\fH - f^\QQ\|_{\QQ,2}^2 = \E_{\QQ_X}\left[\sum_{l=1}^L (q_l^* - q_l^0)\cdot \f{l}(X)\right]^2 \leq \sum_{l=1}^L (q_l^* - q_l^0)^2 \sum_{l=1}^L \E_{\QQ_X}[\f{l}(X)]^2
% \end{align*}
\begin{align*}
\|\fH - f^\QQ\|_{\QQ,2}^2 = \E_{\QQ_X}\left[\sum_{l=1}^L (q_l^* - q_l^0)\cdot \f{l}(X)\right]^2 \leq
\E_{\QQ_X}\left[L \sum_{l=1}^L (q_l^* - q_l^0)^2 (\f{l}(X))^2 \right] = L\sum_{l=1}^L(q_l^* - q_l^0)^2 \|\f{l}\|_{\QQ,2}^2
\end{align*}
If $\max_{q\in \HH}\|q - q^0\|_2 \leq \rho$, since in the above inequality $q^*\in \HH$, it follows that
\[
\|\fH - f^\QQ\|_{\QQ,2}^2 \leq L\rho^2\max_{l\in [L]}\|\f{l}\|_{\QQ,2}^2.
\]
which implies that $\|\fH - f^\QQ\|_{\QQ,2} \leq \rho\sqrt{L}\max_{l\in [L]}\|\f{l}\|_{\QQ,2}$.
\subsection{Proof of Proposition \ref{prop: squared loss model}}
The squared loss for the function $f(\cdot)$ with respect to the distribution $\TT$ is defined as:
\begin{align*}
\E_\TT[(Y_i - f(X_i))^2] &= \E_{X_i\sim \QQ_X}\E_{Y_i\sim \TT_{Y|X}}[(Y_i - f(X_i))^2] .
\end{align*}
As the distribution $\TT$ falls into the uncertainty class $\C(\QQ_X, \HH)$, the preceding can be expressed as
\[
\E_{X_i \sim \QQ_X}\left[\sum_{l=1}^L q_l \E_{Y_i\sim \PP{l}_{Y|X}}[(Y_i - f(X_i))^2]\right].
\]
Considering the data generation mechanism and the noise level assumption in Proposition \ref{prop: squared loss model}, it can be further expressed as follows:
\begin{align*}
\E_\TT[(Y_i - f(X_i))^2]
&= \E_{X_i \sim \QQ_X}\left[\sum_{l=1}^L q_l \cdot \E_{\eps{l}_i|X_i}\left[(\f{l}(X_i) - f(X_i))^2 + (\eps{l}_i)^2\right]\right] \\
&= \E_{X_i \sim \QQ_X}\left[\sum_{l=1}^L q_l \cdot (\f{l}(X_i) - f(X_i))^2 + \sum_{l=1}^L q_l \cdot \sigma_l^2\right] \\
&= \sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[\f{l}(X_i) - f(X_i)\right]^2 + \sum_{l=1}^L q_l \cdot \sigma_l^2.
\end{align*}
The above inequality leads to the following equivalence in terms of the optimization problem,
\[
\min_{f\in \FF}\max_{\TT \in \mathcal{C}(\QQ_X, \HH)} \E_\TT[(Y_i - f(X_i))^2] =
\min_{f\in \FF}\max_{q\in \HH} \left\{\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[\f{l}(X_i) - f(X_i)\right]^2 + \sum_{l=1}^L q_l \cdot \sigma_l^2\right\}.
\]
By minimax theorem \citep{sion1958general,komiya1988elementary} with the convex concave objective, we can swap the order of min and max, the objective after swapping is:
\begin{equation}
\max_{q\in \HH} \min_{f\in \FF} \left\{\sum_{l=1}^L q_l \cdot \E_{\QQ_X} \left[\f{l}(X_i) - f(X_i)\right]^2 + \sum_{l=1}^L q_l \cdot \sigma_l^2\right\}.
\label{proof-prop-sqloss-obj-swap}
\end{equation}
For any fixed $q\in \HH$, we define $f^q$ as the solution for the inner minimization problem in \eqref{proof-prop-sqloss-obj-swap}.
\begin{align*}
f^q &:= \argmin_{f\in \FF}\left\{\sum_{l=1}^L q_l \E_{\QQ_X} \left[\f{l}(X_i) - f(X_i)\right]^2 + \sum_{l=1}^L q_l \sigma_l^2\right\} \\
&= \argmin_{f\in \FF}\left\{\sum_{l=1}^L q_l \E_{\QQ_X} \left[\f{l}(X_i) - f(X_i)\right]^2\right\},
\end{align*}
where we remove the term $\sum_{l=1}^L q_l \sigma_l^2$ as it is irrelevant to the $f$. Then we can further express $f^q$ as follows:
\begin{equation}
\begin{aligned}
f^q &= \argmin_{f\in \FF}\left\{\sum_{l=1}^L q_l \E_{\QQ_X}\left[(\f{l}(X_i))^2 - 2\f{l}(X_i)f(X_i) + (f(X_i))^2\right]\right\} \\
&= \argmin_{f\in \FF}\left\{-2\sum_{l=1}^L q_l \cdot \E_{\QQ_X}\left[\f{l}(X_i)f(X_i)\right] + \E_{\QQ_X}(f(X_i))^2\right\},
\end{aligned}
\end{equation}
where we remove term $\sum_{l=1}^L q_l \E_{\QQ_X}[\f{l}(X_i)^2]$ as it is irrelevant to the $f$ again. With this new expression for $f^q$, we note that $f^q$ is defined exactly the same way as \eqref{eq: thm1-fq-objective} in Theorem \ref{thm: identification}. Therefore by \eqref{eq: thm1-fq}, we can obtain that
\[
f^q(\cdot) = \sum_{l=1}^L q_l \f{l}(\cdot) \quad \textrm{almost surely}.
\]
Hence the objective after swapping in \eqref{proof-prop-sqloss-obj-swap} can be expressed as
\[
\max_{q\in \HH} \left\{\sum_{l=1}^L q_l \E_{\QQ_X}\left[\f{l}(X_i) - f^q(X_i)\right]^2 + \sum_{l=1}^L q_l \sigma_l^2\right\}
\]
Next we define $q^*$ as the solution for the above objective as follows,
\begin{gather}
\begin{aligned}
q^* &= \argmax_{q\in \HH} \left\{\sum_{l=1}^L q_l \E_{\QQ_X}\left[\f{l}(X_i) - f^q(X_i)\right]^2 + \sum_{l=1}^L q_l \sigma_l^2\right\} \\
&= \argmax_{q\in \HH} \left\{\sum_{l=1}^L q_l \E_{\QQ_X}\left[\f{l}(X_i)^2\right] - \E_{\QQ_X}\left(\sum_{l=1}^L q_l \f{l}(X_i)\right)^2 + \sum_{l=1}^L q_l \sigma_l^2\right\}.
\end{aligned}
\label{eq: prop-squaredloss-1}
\end{gather}
Considering $\HH = \Delta^L$ and $L=2$,
with Lagrangian multiplier $\lambda$, we define
\[
\mathcal{L}(q_1, q_2, \lambda) = q_1\E(\f{1})^2 + q_2\E(\f{2})^2 - \E(q_1 \f{1} + q_2 \f{2})^2 + q_1 \sigma_1^2 + q_2 \sigma_2^2 + \lambda (1 - q_1 - q_2),
\]
where we use the simplified notations for convenience. Then the following three equalities hold simultaneously,
\[
\frac{\partial \mathcal{L}}{\partial q_1} = \E(\f{1})^2 - 2q_1\E(\f{1})^2 - 2q_2 \E(\f{1}\f{2}) + \sigma_1^2 - \lambda = 0,
\]
\[
\frac{\partial \mathcal{L}}{\partial q_2} = \E(\f{2})^2 - 2q_2\E(\f{2})^2 - 2q_1 \E(\f{1}\f{2}) + \sigma_2^2 - \lambda = 0,
\]
\[
\frac{\partial \mathcal{L}}{\partial \lambda} = 1 - q_1 - q_2 = 0.
\]
After simple organization of the preceding equalities, we establish that
\[
\E(\f{1})^2 - 2q_1\E(\f{1})^2 - 2(1-q_1) \E(\f{1}\f{2}) + \sigma_1^2 = \E(\f{2})^2 - 2(1-q_1)\E(\f{2})^2 - 2q_1 \E(\f{1}\f{2}) + \sigma_2^2,
\]
which implies that
\[
q_1\E \left[2(\f{1})^2 + 2(\f{2})^2 - 4\f{1}\f{2}\right] = \E \left[(\f{1})^2 + (\f{2})^2 - 2(\f{1}\f{2}) \right] + \sigma_1^2 - \sigma_2^2.
\]
Therefore, we obtain that
\[
2q_1 \E(\f{1}- \f{2})^2 = \E(\f{1}- \f{2})^2 + \sigma_1^2 - \sigma_2^2
\]
By computation, with boundary constraints that $0\leq q_1\leq 1$, we can obtain that
\[
q_1^* = 0 \vee \left(\frac{1}{2}+\frac{\sigma_1^2 - \sigma_2^2}{\left\|\f{1} - \f{2})^2\right\|_{\QQ,2}^2}\right)\wedge 1; \quad
q_2^* = 1-q_1^*
\]
\section{Proofs of Lemmas}
\label{sec: proof-lemmas}
\subsection{Proof for Lemma \ref{lemma: approx}}
\label{sec: proof-lemma-approx}
The proof essentially follows the proof of Lemma 6 in \citet{guo2020inference}. By definition of $q^*$, for any $t\in (0,1)$, we have
\[
(q^*)^\intercal \Gamma q^* \leq [q^* + t(\widehat{q} - q^*)]^\intercal \Gamma [q^* + t(\widehat{q} - q^*)].
\]
hence
\[
0\leq 2 t(q^*)^\intercal \Gamma (\widehat{q} - q^*) + t^2 (\widehat{q} - q^*)^\intercal \Gamma (\widehat{q} - q^*).
\]
By taking $t \rightarrow 0_+$, we have
\begin{equation}
(q^*)^\intercal \Gamma (\widehat{q} - q^*) \geq 0
\label{eq: lemma-approx-1}.
\end{equation}
Similarly, by definition of $\widehat{q}$, for any $t\in (0,1)$, we have
\[
\widehat{q}\widehat{\Gamma}\widehat{q} \leq \left[\widehat{q} + t(q^* - \widehat{q})\right]^\intercal \widehat{\Gamma}\left[\widehat{q} + t(q^* - \widehat{q})\right].
\]
that results in
\[
2(q^*)^\intercal \widehat{\Gamma}(q^* - \widehat{q}) + (t-2)(q^* - \widehat{q})\widehat{\Gamma}(q^* - \widehat{q}) \geq 0.
\]
Since $2-t>0$, it holds that
\begin{equation}
\left(q^* - \widehat{q}\right)^\intercal \widehat{\Gamma}\left(q^* - \widehat{q}\right) \leq \frac{2}{2-t}(q^*)^\intercal \widehat{\Gamma}(q^* - \widehat{q}).
\label{eq: lemma-approx-2}
\end{equation}
It follows from \eqref{eq: lemma-approx-1} that
\[
(q^*)^\intercal \widehat{\Gamma}(q^* - \widehat{q}) = (q^*)^\intercal {\Gamma}(q^* - \widehat{q}) + (q^*)^\intercal \left(\widehat{\Gamma}-\Gamma\right)(q^* - \widehat{q}) \leq (q^*)^\intercal \left(\widehat{\Gamma}-\Gamma\right)(q^* - \widehat{q}).
\]
Combined with \eqref{eq: lemma-approx-2}, we have
\begin{equation}
\left(q^* - \widehat{q}\right)^\intercal \widehat{\Gamma}\left(q^* - \widehat{q}\right) \leq \frac{2}{2-t}(q^*)^\intercal \left(\widehat{\Gamma}-\Gamma\right)(q^* - \widehat{q}) \leq \frac{2\|q^*\|_2}{2-t} \|\widehat{\Gamma} - \Gamma\|_2 \left\|q^* - \widehat{q}\right\|_2.
\label{eq: lemma-approx-3}
\end{equation}
Since the $\widehat{q}$ and $q$ are symmetric, we switch the roles of $\{\widehat{\Gamma},\widehat{q}\}$, $\{\Gamma, q\}$ and establish:
\begin{equation}
\left(q^* - \widehat{q}\right)^\intercal \Gamma \left(q^* - \widehat{q}\right) \leq \frac{2\|\widehat{q}\|_2}{2-t} \|\widehat{\Gamma} - \Gamma\|_2 \left\|q^* - \widehat{q}\right\|_2.
\end{equation}
If $\lambda_{\rm min}(\Gamma) > 0$, it follows that
\[
\lambda_{\rm min}(\Gamma)\left\|q^* - \widehat{q}\right\|_2^2 \leq \frac{2\|\widehat{q}\|_2}{2-t} \|\widehat{\Gamma} - \Gamma\|_2 \left\|q^* - \widehat{q}\right\|_2.
\]
Take $t\rightarrow 0_+$, we establish that
\[
\left\|q^* - \widehat{q}\right\|_2 \leq \frac{1}{\lambda_{\min}(\Gamma)}\left\|\widehat{\Gamma} - \Gamma\right\|_2 \|\widehat{q}\|_2 \leq \frac{\|\widehat{\Gamma} - \Gamma\|_F}{\lambda_{\min}(\Gamma)} \leq \frac{L\|\widehat{\Gamma} - \Gamma\|_\infty}{\lambda_{\min}(\Gamma)}.
\]
Lastly, considering that both $q^*\in \HH$ and $\widehat{q}\in \HH$, so
\[
\left\|q^* - \widehat{q}\right\|_2\leq \frac{L\|\widehat{\Gamma} - \Gamma\|_\infty}{\lambda_{\min}(\Gamma)}\wedge \rho_\HH.
\]
\subsection{Proof for Lemma \ref{lemma: proof-coro-linear}}
Let $|\A_l| = \sum_{i=1}^{n_G}\ind_{[U_i\leq p_l]}$, by Hoeffeding's inequality, it holds that
\[
\prob\left(\left|\frac{1}{n_G}\sum_{i=1}^{n_G}(\ind_{[U_i\leq p_l]} - p_l)\right|\geq t\right)\leq 2\exp(-2n_G t^2).
\]
Let $t = \sqrt{\frac{\log(1/c_1)}{2n_G}}$. Then with probability greater than $1-2c$, the following event holds
\begin{equation}
\left|\frac{|\A_l|}{n_G} - p_l\right|\leq \sqrt{\frac{\log(1/c_1)}{2n_G}}.
\label{proof-coro-linear-2}
\end{equation}
In the following, we will discuss how to identify such constant $C_1$ to make $\left||\A_l|/n_\QQ - p_l/(1-p_l)\right|\leq C_1 \sqrt{\frac{1}{n_G}}$.
\noindent We firstly discuss the case that $|\A_l| /n_\QQ - p_l/(1-p_l) \leq C_1 \sqrt{\frac{1}{n_G}}$.
Considering the case that $n_G$ is sufficiently large such that $n_G > \frac{\log(1/c_1)}{(1-p_l)^2}$,
\begin{align*}
\frac{|\A_l|}{n_\QQ} = \frac{|\A_l|/n_G}{n_\QQ/n_G} = \frac{|\A_l|/n_G - p_l + p_l}{1 - p_l + p_l - |\A_l|/n_G} \leq \frac{\sqrt{\frac{\log(1/c_1)}{2n_G}}+p_l}{1-p_l - \sqrt{\frac{\log(1/\epsilon)}{2n_G}}},
\end{align*}
where the last inequality is due to \eqref{proof-coro-linear-2}.
To show that $|\A_l| /n_\QQ - p_l/(1-p_l) \leq C_1\frac{1}{\sqrt{n_G}}$, it suffices to show that
\[
\frac{\sqrt{\frac{\log(1/c_1)}{2n_G}}+p_l}{1-p_l - \sqrt{\frac{\log(1/c_1)}{2n_G}}} - \frac{p_l}{1-p_l} \leq C_1\frac{1}{\sqrt{n_G}}.
\]
The preceding inequality holds as long as
\[
C_1 \geq \frac{\frac{1}{1-p_l}\sqrt{\log(1/c_1)}}{\sqrt{2}(1-p_l) - \sqrt{\frac{\log(1/c_1)}{n_G}}} \quad \textrm{for any $n_G>\frac{\log(1/c_1)}{(1-p_l)^2}$}.
\]
Thus we make $C_1 \geq \frac{\sqrt{\log(1/c_1)}}{(\sqrt{2}-1)(1-p_l)^2}$.
\noindent For the other case that $p_l/(1-p_l) - |\A_l|/n_\QQ \leq C_1\sqrt{\frac{1}{n_G}}$. Consider the case that $n_G$ is sufficiently large such that $n_G > \frac{\log(1/c_1)}{p_l^2}$,
\[
\frac{|\A_l|}{n_\QQ} = \frac{|\A_l|/n_G}{n_\QQ/n_G} = \frac{|\A_l|/n_G - p_l + p_l}{1 - p_l + p_l - |\A_l|/n_G} \geq \frac{p_l - \sqrt{\frac{\log(1/c_1)}{2n_G}}}{1-p_l + \sqrt{\frac{\log(1/c_1)}{2n_G}}},
\]
where the last inequality is due to \eqref{proof-coro-linear-2}. To show that $p_l/(1-p_l) - |\A_l|/n_\QQ \leq C_1\sqrt{\frac{1}{n_G}}$, it suffices to show that
\[
\frac{p_l}{1-p_l} - \frac{p_l-\sqrt{\frac{\log(1/c_1)}{2n_G}}}{1-p_l+ \sqrt{\frac{\log(1/c_1)}{2n_G}}} \leq C_1\sqrt{\frac{1}{n_G}}.
\]
The preceding inequality holds as long as
\[
C_1 \geq \frac{\frac{1}{1-p_l}\sqrt{\log(1/c_1)}}{\sqrt{2}(1-p_l) + \sqrt{\frac{\log(1/c_1)}{n_G}}} \quad \textrm{for any $n_G>\frac{\log(1/c_1)}{p_l^2}$}.
\]
Thus we make $C_1 = \frac{\sqrt{\log(1/c_1)}}{\sqrt{2}(1-p_l)^2}$.
\noindent To summarize, when $n_G$ is sufficiently large such that $n_G > \frac{\log(1/c_1)}{\min\{(1-p_l)^2, p_l^2\}}$, with probability at least $1-2c$, it holds that
\[
\left|\frac{|\A_l|}{n_\QQ} - \frac{p_l}{1-p_l}\right|\leq C_1\sqrt{\frac{1}{n_G}},
\]
where $C_1 = \frac{\sqrt{\log(1/c_1)}}{(\sqrt{2}-1)(1-p_l)^2}$.
\section{Appendix-Experiments}
\subsection{Experimental Setups}
\label{sec: exp-setup}
For the neural network, we use fully connected MLP with 3 hidden layers and Tanh activation function. The number of hidden units are 5, 7, and 5 respectively. Among the whole data, 95\% are randomly selected as training data, leaving the remaining 5\% for hyperparameters tuning. For experiments on the simulated data, we train the neural network using Adam with a learning rate of 0.005 and weight decay of 0.001. For XGBoost, we use 5-fold cross-validation to choose the hyperparameters.
\end{document}
\begin{thebibliography}{81}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Athey et~al.(2018)Athey, Imbens, and Wager]{athey2018approximate}
S.~Athey, G.~W. Imbens, and S.~Wager.
\newblock Approximate residual balancing: debiased inference of average
treatment effects in high dimensions.
\newblock \emph{Journal of the Royal Statistical Society Series B: Statistical
Methodology}, 80\penalty0 (4):\penalty0 597--623, 2018.
\bibitem[Ben-Tal et~al.(2013)Ben-Tal, Den~Hertog, De~Waegenaere, Melenberg, and
Rennen]{ben2013robust}
A.~Ben-Tal, D.~Den~Hertog, A.~De~Waegenaere, B.~Melenberg, and G.~Rennen.
\newblock Robust solutions of optimization problems affected by uncertain
probabilities.
\newblock \emph{Management Science}, 59\penalty0 (2):\penalty0 341--357, 2013.
\bibitem[Bertsimas et~al.(2018)Bertsimas, Gupta, and Kallus]{bertsimas2018data}
D.~Bertsimas, V.~Gupta, and N.~Kallus.
\newblock Data-driven robust optimization.
\newblock \emph{Mathematical Programming}, 167:\penalty0 235--292, 2018.
\bibitem[Biau(2012)]{biau2012analysis}
G.~Biau.
\newblock Analysis of a random forests model.
\newblock \emph{The Journal of Machine Learning Research}, 13\penalty0
(1):\penalty0 1063--1095, 2012.
\bibitem[Biau et~al.(2008)Biau, Devroye, and Lugosi]{biau2008consistency}
G.~Biau, L.~Devroye, and G.~Lugosi.
\newblock Consistency of random forests and other averaging classifiers.
\newblock \emph{Journal of Machine Learning Research}, 9\penalty0 (9), 2008.
\bibitem[Bickel et~al.(2009)Bickel, Ritov, and
Tsybakov]{bickel2009simultaneous}
P.~J. Bickel, Y.~Ritov, and A.~B. Tsybakov.
\newblock Simultaneous analysis of lasso and dantzig selector.
\newblock \emph{The Annals of Statistics}, 37\penalty0 (4):\penalty0
1705--1732, 2009.
\bibitem[Blanchet and Murthy(2019)]{blanchet2019quantifying}
J.~Blanchet and K.~Murthy.
\newblock Quantifying distributional model risk via optimal transport.
\newblock \emph{Mathematics of Operations Research}, 44\penalty0 (2):\penalty0
565--600, 2019.
\bibitem[Blanchet et~al.(2019{\natexlab{a}})Blanchet, Kang, and
Murthy]{blanchet2019robust}
J.~Blanchet, Y.~Kang, and K.~Murthy.
\newblock Robust wasserstein profile inference and applications to machine
learning.
\newblock \emph{Journal of Applied Probability}, 56\penalty0 (3):\penalty0
830--857, 2019{\natexlab{a}}.
\bibitem[Blanchet et~al.(2019{\natexlab{b}})Blanchet, Kang, Murthy, and
Zhang]{blanchet2019data}
J.~Blanchet, Y.~Kang, K.~Murthy, and F.~Zhang.
\newblock Data-driven optimal transport cost selection for distributionally
robust optimization.
\newblock In \emph{2019 winter simulation conference (WSC)}, pages 3740--3751.
IEEE, 2019{\natexlab{b}}.
\bibitem[Blodgett et~al.(2016)Blodgett, Green, and
O'Connor]{blodgett2016demographic}
S.~L. Blodgett, L.~Green, and B.~O'Connor.
\newblock Demographic dialectal variation in social media: A case study of
african-american english.
\newblock \emph{arXiv preprint arXiv:1608.08868}, 2016.
\bibitem[B{\"u}hlmann and Meinshausen(2015)]{buhlmann2015magging}
P.~B{\"u}hlmann and N.~Meinshausen.
\newblock Magging: maximin aggregation for inhomogeneous large-scale data.
\newblock \emph{Proceedings of the IEEE}, 104\penalty0 (1):\penalty0 126--135,
2015.
\bibitem[B{\"u}hlmann and van~de Geer(2011)]{buhlmann2011statistics}
P.~B{\"u}hlmann and S.~van~de Geer.
\newblock \emph{Statistics for high-dimensional data: methods, theory and
applications}.
\newblock Springer Science \& Business Media, 2011.
\bibitem[Candes and Tao(2007)]{candes2007dantzig}
E.~Candes and T.~Tao.
\newblock The dantzig selector: Statistical estimation when p is much larger
than n.
\newblock \emph{Annals of statistics}, 35\penalty0 (6):\penalty0 2313--2404,
2007.
\bibitem[Chen et~al.(2019)Chen, Bentley, Mori, Misawa, Fujiwara, and
Rueckert]{chen2019self}
L.~Chen, P.~Bentley, K.~Mori, K.~Misawa, M.~Fujiwara, and D.~Rueckert.
\newblock Self-supervised learning for medical image analysis using image
context restoration.
\newblock \emph{Medical image analysis}, 58:\penalty0 101539, 2019.
\bibitem[Chen et~al.(2008)Chen, Hong, and Tarozzi]{chen2008semiparametric}
X.~Chen, H.~Hong, and A.~Tarozzi.
\newblock Semiparametric efficiency in gmm models in auxiliary data.
\newblock \emph{The Annals of Statistics}, 36\penalty0 (2):\penalty0 808--843,
2008.
\bibitem[Chernozhukov et~al.(2018)Chernozhukov, Chetverikov, Demirer, Duflo,
Hansen, Newey, and Robins]{chernozhukov2018double}
V.~Chernozhukov, D.~Chetverikov, M.~Demirer, E.~Duflo, C.~Hansen, W.~Newey, and
J.~Robins.
\newblock Double/debiased machine learning for treatment and structural
parameters: Double/debiased machine learning, 2018.
\bibitem[Choi et~al.(2022)Choi, Meng, Song, and Ermon]{choi2022density}
K.~Choi, C.~Meng, Y.~Song, and S.~Ermon.
\newblock Density ratio estimation via infinitesimal classification.
\newblock In \emph{International Conference on Artificial Intelligence and
Statistics}, pages 2552--2573. PMLR, 2022.
\bibitem[Deng et~al.(2020)Deng, Kamani, and Mahdavi]{deng2020distributionally}
Y.~Deng, M.~M. Kamani, and M.~Mahdavi.
\newblock Distributionally robust federated averaging.
\newblock \emph{Advances in neural information processing systems},
33:\penalty0 15111--15122, 2020.
\bibitem[Diana et~al.(2021)Diana, Gill, Kearns, Kenthapadi, and
Roth]{diana2021minimax}
E.~Diana, W.~Gill, M.~Kearns, K.~Kenthapadi, and A.~Roth.
\newblock Minimax group fairness: Algorithms and experiments.
\newblock In \emph{Proceedings of the 2021 AAAI/ACM Conference on AI, Ethics,
and Society}, pages 66--76, 2021.
\bibitem[Du et~al.(2021)Du, Xu, Wu, and Tong]{du2021fairness}
W.~Du, D.~Xu, X.~Wu, and H.~Tong.
\newblock Fairness-aware agnostic federated learning.
\newblock In \emph{Proceedings of the 2021 SIAM International Conference on
Data Mining (SDM)}, pages 181--189. SIAM, 2021.
\bibitem[Duchi et~al.(2023)Duchi, Hashimoto, and
Namkoong]{duchi2023distributionally}
J.~Duchi, T.~Hashimoto, and H.~Namkoong.
\newblock Distributionally robust losses for latent covariate mixtures.
\newblock \emph{Operations Research}, 71\penalty0 (2):\penalty0 649--664, 2023.
\bibitem[Duchi and Namkoong(2021)]{duchi2021learning}
J.~C. Duchi and H.~Namkoong.
\newblock Learning models with uniform performance via distributionally robust
optimization.
\newblock \emph{The Annals of Statistics}, 49\penalty0 (3):\penalty0
1378--1406, 2021.
\bibitem[Farrell et~al.(2021)Farrell, Liang, and Misra]{farrell2021deep}
M.~H. Farrell, T.~Liang, and S.~Misra.
\newblock Deep neural networks for estimation and inference.
\newblock \emph{Econometrica}, 89\penalty0 (1):\penalty0 181--213, 2021.
\bibitem[Gao and Kleywegt(2023)]{gao2023distributionally}
R.~Gao and A.~Kleywegt.
\newblock Distributionally robust stochastic optimization with wasserstein
distance.
\newblock \emph{Mathematics of Operations Research}, 48\penalty0 (2):\penalty0
603--655, 2023.
\bibitem[Gao et~al.(2022)Gao, Chen, and Kleywegt]{gao2022wasserstein}
R.~Gao, X.~Chen, and A.~J. Kleywegt.
\newblock Wasserstein distributionally robust optimization and variation
regularization.
\newblock \emph{Operations Research}, 2022.
\bibitem[Gretton et~al.(2009)Gretton, Smola, Huang, Schmittfull, Borgwardt, and
Sch{\"o}llkopf]{gretton2009covariate}
A.~Gretton, A.~J. Smola, J.~Huang, M.~Schmittfull, K.~M. Borgwardt, and
B.~Sch{\"o}llkopf.
\newblock Covariate shift by kernel mean matching.
\newblock pages 131--160, 2009.
\bibitem[Grother and Phillips(2011)]{grother2011report}
P.~J. Grother and P.~J. Phillips.
\newblock \emph{Report on the evaluation of 2D still-image face recognition
algorithms}.
\newblock 2011.
\bibitem[Guo(2023)]{guo2020inference}
Z.~Guo.
\newblock Statistical inference for maximin effects: Identifying stable
associations across multiple studies.
\newblock \emph{Journal of the American Statistical Association}, pages 1--32,
2023.
\bibitem[Guo et~al.(2023)Guo, Li, Han, and Cai]{guo2023robust}
Z.~Guo, X.~Li, L.~Han, and T.~Cai.
\newblock Robust inference for federated meta-learning.
\newblock \emph{arXiv preprint arXiv:2301.00718}, 2023.
\bibitem[Hashimoto et~al.(2018)Hashimoto, Srivastava, Namkoong, and
Liang]{hashimoto2018fairness}
T.~Hashimoto, M.~Srivastava, H.~Namkoong, and P.~Liang.
\newblock Fairness without demographics in repeated loss minimization.
\newblock In \emph{International Conference on Machine Learning}, pages
1929--1938. PMLR, 2018.
\bibitem[Hu et~al.(2018)Hu, Niu, Sato, and Sugiyama]{hu2018does}
W.~Hu, G.~Niu, I.~Sato, and M.~Sugiyama.
\newblock Does distributionally robust supervised learning give robust
classifiers?
\newblock In \emph{International Conference on Machine Learning}, pages
2029--2037. PMLR, 2018.
\bibitem[Huang and Zhang(2012)]{huang2012estimation}
J.~Huang and C.-H. Zhang.
\newblock Estimation and selection via absolute penalized convex minimization
and its multistage adaptive applications.
\newblock \emph{The Journal of Machine Learning Research}, 13\penalty0
(1):\penalty0 1839--1864, 2012.
\bibitem[Humbert-Droz et~al.(2022)Humbert-Droz, Mukherjee, Gevaert,
et~al.]{humbert2022strategies}
M.~Humbert-Droz, P.~Mukherjee, O.~Gevaert, et~al.
\newblock Strategies to address the lack of labeled data for supervised machine
learning training with electronic health records: Case study for the
extraction of symptoms from clinical notes.
\newblock \emph{JMIR Medical Informatics}, 10\penalty0 (3):\penalty0 e32903,
2022.
\bibitem[Jiao et~al.(2022)Jiao, Zhang, Ding, Cai, and Zhang]{jiao2022learning}
R.~Jiao, Y.~Zhang, L.~Ding, R.~Cai, and J.~Zhang.
\newblock Learning with limited annotations: a survey on deep semi-supervised
learning for medical image segmentation.
\newblock \emph{arXiv preprint arXiv:2207.14191}, 2022.
\bibitem[Klaassen(1987)]{klaassen1987consistent}
C.~A. Klaassen.
\newblock Consistent estimation of the influence function of locally
asymptotically linear estimators.
\newblock \emph{The Annals of Statistics}, 15\penalty0 (4):\penalty0
1548--1562, 1987.
\bibitem[Koh et~al.(2021)Koh, Sagawa, Marklund, Xie, Zhang, Balsubramani, Hu,
Yasunaga, Phillips, Gao, et~al.]{koh2021wilds}
P.~W. Koh, S.~Sagawa, H.~Marklund, S.~M. Xie, M.~Zhang, A.~Balsubramani, W.~Hu,
M.~Yasunaga, R.~L. Phillips, I.~Gao, et~al.
\newblock Wilds: A benchmark of in-the-wild distribution shifts.
\newblock In \emph{International Conference on Machine Learning}, pages
5637--5664. PMLR, 2021.
\bibitem[Komiya(1988)]{komiya1988elementary}
H.~Komiya.
\newblock Elementary proof for sion's minimax theorem.
\newblock \emph{Kodai mathematical journal}, 11\penalty0 (1):\penalty0 5--7,
1988.
\bibitem[Kuhn et~al.(2019)Kuhn, Esfahani, Nguyen, and
Shafieezadeh-Abadeh]{kuhn2019wasserstein}
D.~Kuhn, P.~M. Esfahani, V.~A. Nguyen, and S.~Shafieezadeh-Abadeh.
\newblock Wasserstein distributionally robust optimization: Theory and
applications in machine learning.
\newblock In \emph{Operations research \& management science in the age of
analytics}, pages 130--166. Informs, 2019.
\bibitem[Leek et~al.(2010)Leek, Scharpf, Bravo, Simcha, Langmead, Johnson,
Geman, Baggerly, and Irizarry]{leek2010tackling}
J.~T. Leek, R.~B. Scharpf, H.~C. Bravo, D.~Simcha, B.~Langmead, W.~E. Johnson,
D.~Geman, K.~Baggerly, and R.~A. Irizarry.
\newblock Tackling the widespread and critical impact of batch effects in
high-throughput data.
\newblock \emph{Nature Reviews Genetics}, 11\penalty0 (10):\penalty0 733--739,
2010.
\bibitem[Ling et~al.(2022)Ling, Lu, Zhao, Lulla, Plantinga, Fu, Zhang, Liu,
Song, Li, et~al.]{ling2022batch}
W.~Ling, J.~Lu, N.~Zhao, A.~Lulla, A.~M. Plantinga, W.~Fu, A.~Zhang, H.~Liu,
H.~Song, Z.~Li, et~al.
\newblock Batch effects removal for microbiome data via conditional quantile
regression.
\newblock \emph{Nature communications}, 13\penalty0 (1):\penalty0 5418, 2022.
\bibitem[Liu et~al.(2021)Liu, Xia, Cho, and Cai]{liu2021integrative}
M.~Liu, Y.~Xia, K.~Cho, and T.~Cai.
\newblock Integrative high dimensional multiple testing with heterogeneity
under data sharing constraints.
\newblock \emph{The Journal of Machine Learning Research}, 22\penalty0
(1):\penalty0 5607--5632, 2021.
\bibitem[Maity et~al.(2022)Maity, Sun, and Banerjee]{maity2022meta}
S.~Maity, Y.~Sun, and M.~Banerjee.
\newblock Meta-analysis of heterogeneous data: integrative sparse regression in
high-dimensions.
\newblock \emph{The Journal of Machine Learning Research}, 23\penalty0
(1):\penalty0 8975--9024, 2022.
\bibitem[Malinin et~al.(2021)Malinin, Band, Chesnokov, Gal, Gales, Noskov,
Ploskonosov, Prokhorenkova, Provilkov, Raina, et~al.]{malinin2021shifts}
A.~Malinin, N.~Band, G.~Chesnokov, Y.~Gal, M.~J. Gales, A.~Noskov,
A.~Ploskonosov, L.~Prokhorenkova, I.~Provilkov, V.~Raina, et~al.
\newblock Shifts: A dataset of real distributional shift across multiple
large-scale tasks.
\newblock \emph{arXiv preprint arXiv:2107.07455}, 2021.
\bibitem[Martinez et~al.(2020)Martinez, Bertran, and
Sapiro]{martinez2020minimax}
N.~Martinez, M.~Bertran, and G.~Sapiro.
\newblock Minimax pareto fairness: A multi objective perspective.
\newblock In \emph{International Conference on Machine Learning}, pages
6755--6764. PMLR, 2020.
\bibitem[Meinshausen and B{\"u}hlmann(2015)]{meinshausen2015maximin}
N.~Meinshausen and P.~B{\"u}hlmann.
\newblock Maximin effects in inhomogeneous large-scale data.
\newblock \emph{The Annals of Statistics}, 43\penalty0 (4):\penalty0
1801--1830, 2015.
\bibitem[Meinshausen and Ridgeway(2006)]{meinshausen2006quantile}
N.~Meinshausen and G.~Ridgeway.
\newblock Quantile regression forests.
\newblock \emph{Journal of machine learning research}, 7\penalty0 (6), 2006.
\bibitem[Menon and Ong(2016)]{menon2016linking}
A.~Menon and C.~S. Ong.
\newblock Linking losses for density ratio and class-probability estimation.
\newblock In \emph{International Conference on Machine Learning}, pages
304--313. PMLR, 2016.
\bibitem[Mohri et~al.(2019)Mohri, Sivek, and Suresh]{mohri2019agnostic}
M.~Mohri, G.~Sivek, and A.~T. Suresh.
\newblock Agnostic federated learning.
\newblock In \emph{International Conference on Machine Learning}, pages
4615--4625. PMLR, 2019.
\bibitem[Nado et~al.(2021)Nado, Band, Collier, Djolonga, Dusenberry, Farquhar,
Feng, Filos, Havasi, Jenatton, et~al.]{nado2021uncertainty}
Z.~Nado, N.~Band, M.~Collier, J.~Djolonga, M.~W. Dusenberry, S.~Farquhar,
Q.~Feng, A.~Filos, M.~Havasi, R.~Jenatton, et~al.
\newblock Uncertainty baselines: Benchmarks for uncertainty \& robustness in
deep learning.
\newblock \emph{arXiv preprint arXiv:2106.04015}, 2021.
\bibitem[Namkoong and Duchi(2016)]{namkoong2016stochastic}
H.~Namkoong and J.~C. Duchi.
\newblock Stochastic gradient methods for distributionally robust optimization
with f-divergences.
\newblock \emph{Advances in neural information processing systems}, 29, 2016.
\bibitem[Namkoong and Duchi(2017)]{namkoong2017variance}
H.~Namkoong and J.~C. Duchi.
\newblock Variance-based regularization with convex objectives.
\newblock \emph{Advances in neural information processing systems}, 30, 2017.
\bibitem[Negahban et~al.(2012)Negahban, Ravikumar, Wainwright, and
Yu]{negahban2012unified}
S.~N. Negahban, P.~Ravikumar, M.~J. Wainwright, and B.~Yu.
\newblock A unified framework for high-dimensional analysis of m-estimators
with decomposable regularizers.
\newblock 2012.
\bibitem[Nguyen et~al.(2010)Nguyen, Wainwright, and
Jordan]{nguyen2010estimating}
X.~Nguyen, M.~J. Wainwright, and M.~I. Jordan.
\newblock Estimating divergence functionals and the likelihood ratio by convex
risk minimization.
\newblock \emph{IEEE Transactions on Information Theory}, 56\penalty0
(11):\penalty0 5847--5861, 2010.
\bibitem[Oren et~al.(2019)Oren, Sagawa, Hashimoto, and
Liang]{oren2019distributionally}
Y.~Oren, S.~Sagawa, T.~B. Hashimoto, and P.~Liang.
\newblock Distributionally robust language modeling.
\newblock \emph{arXiv preprint arXiv:1909.02060}, 2019.
\bibitem[Ouali et~al.(2020)Ouali, Hudelot, and Tami]{ouali2020overview}
Y.~Ouali, C.~Hudelot, and M.~Tami.
\newblock An overview of deep semi-supervised learning.
\newblock \emph{arXiv preprint arXiv:2006.05278}, 2020.
\bibitem[Qin(1998)]{qin1998inferences}
J.~Qin.
\newblock Inferences for case-control and semiparametric two-sample density
ratio models.
\newblock \emph{Biometrika}, 85\penalty0 (3):\penalty0 619--630, 1998.
\bibitem[Quinonero-Candela et~al.(2008)Quinonero-Candela, Sugiyama,
Schwaighofer, and Lawrence]{quinonero2008dataset}
J.~Quinonero-Candela, M.~Sugiyama, A.~Schwaighofer, and N.~D. Lawrence.
\newblock \emph{Dataset shift in machine learning}.
\newblock Mit Press, 2008.
\bibitem[Rahimian and Mehrotra(2019)]{rahimian2019distributionally}
H.~Rahimian and S.~Mehrotra.
\newblock Distributionally robust optimization: A review.
\newblock \emph{arXiv preprint arXiv:1908.05659}, 2019.
\bibitem[Rasmy et~al.(2018)Rasmy, Wu, Wang, Geng, Zheng, Wang, Wu, Xu, and
Zhi]{rasmy2018study}
L.~Rasmy, Y.~Wu, N.~Wang, X.~Geng, W.~J. Zheng, F.~Wang, H.~Wu, H.~Xu, and
D.~Zhi.
\newblock A study of generalizability of recurrent neural network-based
predictive models for heart failure onset risk using a large and
heterogeneous ehr data set.
\newblock \emph{Journal of biomedical informatics}, 84:\penalty0 11--16, 2018.
\bibitem[Reddi et~al.(2015)Reddi, Poczos, and Smola]{reddi2015doubly}
S.~Reddi, B.~Poczos, and A.~Smola.
\newblock Doubly robust covariate shift correction.
\newblock In \emph{Proceedings of the AAAI Conference on Artificial
Intelligence}, volume~29, 2015.
\bibitem[Reisizadeh et~al.(2020)Reisizadeh, Farnia, Pedarsani, and
Jadbabaie]{reisizadeh2020robust}
A.~Reisizadeh, F.~Farnia, R.~Pedarsani, and A.~Jadbabaie.
\newblock Robust federated learning: The case of affine distribution shifts.
\newblock \emph{Advances in neural information processing systems},
33:\penalty0 21554--21565, 2020.
\bibitem[Rosenbaum and Rubin(1983)]{rosenbaum1983central}
P.~R. Rosenbaum and D.~B. Rubin.
\newblock The central role of the propensity score in observational studies for
causal effects.
\newblock \emph{Biometrika}, 70\penalty0 (1):\penalty0 41--55, 1983.
\bibitem[Rothenh{\"a}usler et~al.(2016)Rothenh{\"a}usler, Meinshausen, and
B{\"u}hlmann]{rothenhausler2016confidence}
D.~Rothenh{\"a}usler, N.~Meinshausen, and P.~B{\"u}hlmann.
\newblock Confidence intervals for maximin effects in inhomogeneous large-scale
data.
\newblock In \emph{Statistical Analysis for High-Dimensional Data: The Abel
Symposium 2014}, pages 255--277. Springer, 2016.
\bibitem[Sagawa et~al.(2019)Sagawa, Koh, Hashimoto, and
Liang]{sagawa2019distributionally}
S.~Sagawa, P.~W. Koh, T.~B. Hashimoto, and P.~Liang.
\newblock Distributionally robust neural networks for group shifts: On the
importance of regularization for worst-case generalization.
\newblock \emph{arXiv preprint arXiv:1911.08731}, 2019.
\bibitem[Schick(1986)]{schick1986asymptotically}
A.~Schick.
\newblock On asymptotically efficient estimation in semiparametric models.
\newblock \emph{The Annals of Statistics}, pages 1139--1151, 1986.
\bibitem[Schmidt-Hieber(2020)]{schmidt2020nonparametric}
J.~Schmidt-Hieber.
\newblock Nonparametric regression using deep neural networks with relu
activation function.
\newblock \emph{The Annals of Statistics}, 48\penalty0 (4):\penalty0
1875--1897, 2020.
\bibitem[Scornet et~al.(2015)Scornet, Biau, and Vert]{scornet2015consistency}
E.~Scornet, G.~Biau, and J.-P. Vert.
\newblock Consistency of random forests.
\newblock \emph{Annals of Statistics}, 43\penalty0 (4):\penalty0 1716--1741,
2015.
\bibitem[Shimodaira(2000)]{shimodaira2000improving}
H.~Shimodaira.
\newblock Improving predictive inference under covariate shift by weighting the
log-likelihood function.
\newblock \emph{Journal of statistical planning and inference}, 90\penalty0
(2):\penalty0 227--244, 2000.
\bibitem[Singh et~al.(2022)Singh, Mhasawade, and
Chunara]{singh2022generalizability}
H.~Singh, V.~Mhasawade, and R.~Chunara.
\newblock Generalizability challenges of mortality risk prediction models: A
retrospective analysis on a multi-center database.
\newblock \emph{PLOS Digital Health}, 1\penalty0 (4):\penalty0 e0000023, 2022.
\bibitem[Sinha et~al.(2017)Sinha, Namkoong, Volpi, and
Duchi]{sinha2017certifying}
A.~Sinha, H.~Namkoong, R.~Volpi, and J.~Duchi.
\newblock Certifying some distributional robustness with principled adversarial
training.
\newblock \emph{arXiv preprint arXiv:1710.10571}, 2017.
\bibitem[Sion(1958)]{sion1958general}
M.~Sion.
\newblock On general minimax theorems.
\newblock 1958.
\bibitem[Sugiyama et~al.(2012)Sugiyama, Suzuki, and
Kanamori]{sugiyama2012density}
M.~Sugiyama, T.~Suzuki, and T.~Kanamori.
\newblock \emph{Density ratio estimation in machine learning}.
\newblock Cambridge University Press, 2012.
\bibitem[Tibshirani et~al.(2019)Tibshirani, Foygel~Barber, Candes, and
Ramdas]{tibshirani2019conformal}
R.~J. Tibshirani, R.~Foygel~Barber, E.~Candes, and A.~Ramdas.
\newblock Conformal prediction under covariate shift.
\newblock \emph{Advances in neural information processing systems}, 32, 2019.
\bibitem[Topaloglu et~al.(2021)Topaloglu, Morrell, Rajendran, and
Topaloglu]{topaloglu2021pursuit}
M.~Y. Topaloglu, E.~M. Morrell, S.~Rajendran, and U.~Topaloglu.
\newblock In the pursuit of privacy: the promises and predicaments of federated
learning in healthcare.
\newblock \emph{Frontiers in Artificial Intelligence}, page 147, 2021.
\bibitem[van~de Geer et~al.(2014)van~de Geer, B{\"u}hlmann, Ritov, and
Dezeure]{van2014asymptotically}
S.~van~de Geer, P.~B{\"u}hlmann, Y.~Ritov, and R.~Dezeure.
\newblock On asymptotically optimal confidence regions and tests for
high-dimensional models.
\newblock \emph{The Annals of Statistics}, 42\penalty0 (3):\penalty0
1166--1202, 2014.
\bibitem[Vershynin(2010)]{vershynin2010introduction}
R.~Vershynin.
\newblock Introduction to the non-asymptotic analysis of random matrices.
\newblock \emph{arXiv preprint arXiv:1011.3027}, 2010.
\bibitem[Weber et~al.(2022)Weber, Hong, Xia, Palmer, Avillach, L’yi, Keller,
Murphy, Guti{\'e}rrez-Sacrist{\'a}n, Bonzel, et~al.]{weber2022international}
G.~M. Weber, C.~Hong, Z.~Xia, N.~P. Palmer, P.~Avillach, S.~L’yi, M.~S.
Keller, S.~N. Murphy, A.~Guti{\'e}rrez-Sacrist{\'a}n, C.-L. Bonzel, et~al.
\newblock International comparisons of laboratory values from the 4ce
collaborative to predict covid-19 mortality.
\newblock \emph{NPJ digital medicine}, 5\penalty0 (1):\penalty0 74, 2022.
\bibitem[Wright and Ziegler(2015)]{wright2015ranger}
M.~N. Wright and A.~Ziegler.
\newblock ranger: A fast implementation of random forests for high dimensional
data in c++ and r.
\newblock \emph{arXiv preprint arXiv:1508.04409}, 2015.
\bibitem[Zhang et~al.(2020{\natexlab{a}})Zhang, Menon, Veit, Bhojanapalli,
Kumar, and Sra]{zhang2020coping}
J.~Zhang, A.~Menon, A.~Veit, S.~Bhojanapalli, S.~Kumar, and S.~Sra.
\newblock Coping with label shift via distributionally robust optimisation.
\newblock \emph{arXiv preprint arXiv:2010.12230}, 2020{\natexlab{a}}.
\bibitem[Zhang et~al.(2017)Zhang, Guo, Dong, He, Xu, and
Chen]{zhang2017cautionary}
S.~Zhang, B.~Guo, A.~Dong, J.~He, Z.~Xu, and S.~X. Chen.
\newblock Cautionary tales on air-quality improvement in beijing.
\newblock \emph{Proceedings of the Royal Society A: Mathematical, Physical and
Engineering Sciences}, 473\penalty0 (2205):\penalty0 20170457, 2017.
\bibitem[Zhang et~al.(2020{\natexlab{b}})Zhang, Li, Zeng, Sun, Kumar, Ye, and
Ji]{zhang2015deep}
W.~Zhang, R.~Li, T.~Zeng, Q.~Sun, S.~Kumar, J.~Ye, and S.~Ji.
\newblock Deep model based transfer and multi-task learning for biological
image analysis.
\newblock volume~6, pages 322--333, 2020{\natexlab{b}}.
\end{thebibliography}