1 Introduction↩︎

Machine learning models are often trained on data from a source domain and deployed in a target domain where the data distribution may differ. This discrepancy between training and deployment distributions can degrade model performance. This work considers a scenario where both features and labels are available in the source domain, but only unlabeled features are observed in the target domain. This mirrors applications such as medical diagnosis or credit scoring, where predictions involve future outcomes and immediate access to labels during deployment is not feasible.

When the target distribution can differ arbitrarily from the source distribution, estimating target performance or adapting the model becomes an ill-posed problem. To address this, the literature often imposes invariance assumptions about the nature of the distribution shift [1]. Common assumptions include changes in the marginal feature distribution \(P_{\mathbf{X}}\) while the conditional label distribution \(P_{Y|\mathbf{X}}\) remains unchanged, known as covariate shift [2]; shifts in the marginal label distribution \(P_Y\) while \(P_{\mathbf{X}|Y}\) remains constant, referred to as label shift [3], [4]; or changes in the joint distribution \(P_{\mathbf{X}, Y}\) under the condition that there exists a subset of features \(S\) such that \(P_{\mathbf{X}_{\bar{S}} | \mathbf{X}_S, Y}\) is invariant, known as sparse joint shift [1], [5], [6].

Instead of committing to these specific shift assumptions, we adopt the perspective of domain adaptation theory [7], [8], which assumes only that there exists a hypothesis that performs reasonably well across both domains. In addition, we focus on machine learning models that process semantically meaningful features, such as those in tabular data, rather than unstructured data like images. Within this context, we introduce a unified framework that: estimates a model’s performance on the unlabeled target domain; explains performance changes by identifying the features responsible under sparse joint shifts; and improves the model to enhance its target-domain performance.

Our method transforms the source distribution by solving a constrained optimization problem: it matches key statistics of the target domain while minimizing the KL divergence from the source. This formulation yields a unique closed-form solution for the importance weights, which guarantees efficient computation and reduced estimation variance. The central theoretical insight is that accurate performance estimation and adaptation don’t require recovering the full distribution ratio, but matching carefully chosen moments suffices. Our framework is supported by strong theoretical foundations and extensive empirical validation: across several distribution shifts, EPA achieves \(30-70\%\) better target error estimation, more precise feature attribution, and superior model adaptation over state-of-the-art methods, all while running an order of magnitude faster.

Related Works. The work most closely related to ours is SEES [5], which also generates importance weights, estimates model performance, and identifies features driving sparse joint shifts. However, our approach diverges fundamentally in its objective. SEES attempts to directly learn the density ratio \(w_S(x, y) = \frac{dQ}{dP}(\boldsymbol{x}_S, y)\) by using a class of basis functions. In contrast, we circumvent the notoriously difficult task of density ratio estimation. Instead, we find the distribution closest to the source – thereby reducing weight variance – while satisfying selected statistics of the target domain, which ensures accuracy for downstream estimation and adaptation. This formulation is not only more robust but also yields a closed-form solution, making our method significantly faster. Furthermore, unlike SEES, our approach is not limited to classification tasks.

Our approach also shares conceptual similarities with kernel mean matching (KMM) [9], [10], which aligns feature means in an RKHS. KMM, however, requires solving a quadratic program whose complexity scales with the number of samples \(n\), making it inefficient for large datasets. Our entropic projection approach scales with the feature dimension \(k \ll n\). Moreover, the KL minimisation in our objective is equivalent to entropy maximisation, which provides implicit regularisation and prevents the unstable, high-variance weights that KMM can sometimes produce.

Finally, to our knowledge, no prior work has integrated all three aspects, estimating, explaining, and improving, into a single framework; most focus on only some. Refer to 20 for further discussion of prior works.

2 Entropic Projection Alignment↩︎

Notations. Let \(\small \mathcal{X}\) and \(\small\mathcal{Y}\) denote the input and label spaces. We consider a predictive model \(\small f : \mathcal{X} \to \mathcal{Y}\) from a hypothesis class \(\mathcal{H}\) and bounded loss function \(\small \mathcal{L} : \mathcal{Y}^2 \to [0, 1]\). We are given a labeled dataset , drawn i.i.d from a source distribution \(\small P = P_X \times P_{Y|X}\), and an unlabeled dataset drawn from a target distribution \(\small Q = Q_X \times Q_{Y|X}\). We define the risk of a hypothesis \(f \in \mathcal{H}\) under distribution \(P\) with respect to its true labeling function \(c_P\) as \(\epsilon_P(f) := \mathbb{E}_{x \sim P_X}\big[\ell(f(x), c_P(x))\big]\) and analogously for \(Q\). We use the hat notation \(\small\hat{\cdot}\) to represent empirical distributions.

EPA casts alignment as an entropic projection [11] – a classical optimisation problem that seeks the distribution closest in \(\mathrm{KL}\) divergence to a reference distribution while satisfying linear constraints. This framework was recently used by [12] for sensitivity analysis, where it addresses counterfactual questions such as: how would a model’s predictions change if the data distribution were projected to increase the mean or correlations of the features?

In contrast, we use entropic projection for a different goal: domain alignment. Our EPA method imposes constraints that align the source with the target distribution. Concretely, EPA seeks a distribution \(P^\star\) that is as close as possible to \(P\) while matching moment conditions derived from the target domain \(Q_{\boldsymbol{X}}\).

If \(\widehat\mu_Q\) lies outside \(\operatorname{conv}\{\Phi(X_i)\}_{i=1}^n\), the equality-constrained problem is infeasible. In that case, we project \(\widehat\mu_Q\) onto the source convex hull and run EPA with the projected moment; §35 shows that this is equivalent to imposing a hard proximity constraint around the target moment.

In our setting, both estimation and explanation revolve around finding the alignment (via the matching function \(\Phi\)) that most accurately estimates and explains the model’s performance. Then, we leverage the distribution \(P^\star\), drawn from the best estimate, to improve the model.

As noted by [12], this result follows as a corollary of the theorems in [11]. In §22.1, we provide a simpler direct proof that avoids relying on those general results. Moreover, in  §30 we establish new finite-sample concentration bounds comparing the empirical EPA in  §1 with its population analogue.

The EPA approach offers several advantages:

1. Uniqueness and implicit regularisation. The solution is unique and corresponds to the maximum-entropy set of weights satisfying the constraints, since \(\mathrm{KL}(\tilde{P} \,\Vert\, \hat{P}) = \log n - H(\lambda)\). This encourages weights close to uniform, acting as a strong regulariser that penalises concentration and reduces the variance of importance-weighted estimators. A formal variance analysis is provided in §22.2.

2. Scalability. The optimisation involves only \(k = \dim(\Phi)\) parameters, whereas KMM requires solving a quadratic program scaling with \(n\), which is prohibitive for large datasets.

3. Distribution-free formulation. EPA requires no explicit assumptions on the underlying data distribution to approximate KL divergence, unlike prior approaches such as SEES.

3 Estimating Model Performance↩︎

The previous section introduced the EPA framework, which produces a reweighted source distribution \(P^\star\) that matches the target moments of \(\Phi\) under \(Q\). The ultimate goal is to ensure that the risk on the reweighted source, \(\epsilon_{P^\star}(f)\), is a reliable estimator for the target risk \(\epsilon_Q(f)\). This raises the central question:

To address this, we draw on the theory of domain adaptation, which provides formal tools for reasoning about performance under distribution shift.

This measure captures the maximum disagreement between the two distributions over the worst-case pair of hypotheses from \(\mathcal{H}\). A standard result [7], [8], shown below, bounds the estimation gap using this discrepancy.

See derivations in 23. This bound clarifies our objective: to make the estimation gap small, we must choose our feature map \(\Phi\) such that the resulting \(P^\star\) has small discrepancy \(\mathrm{disc}_\ell(P^\star_X, Q_X)\). The following example shows that naively matching simple moments is not enough in general.

3.1 Alignment↩︎

The failure of the previous example demonstrates that the feature map \(\Phi\) must be chosen in a way that is sensitive to the hypothesis class \(\mathcal{H}\) and loss function \(\ell\). We formalise this intuition with the concept of alignment.

Intuitively, alignment means that \(\Phi\) is expressive enough to exactly represent any disagreement pattern between hypotheses. When alignment holds, moment matching eliminates discrepancy entirely.

The remaining term \(\lambda_{P^\star,Q}\) captures the shared-good-hypothesis assumption. Under standard covariate or label shift, it is zero whenever the hypothesis class can represent the common labelling rule. If source and target are genuinely incompatible, then \(\lambda_{P^\star,Q}\) is large, and neither EPA nor any other reweighting method can be expected to provide accurate target risk estimation or reliable adaptation.

To illustrate how this abstract condition manifests in practice, we now present several concrete examples.

See proofs and additional examples in 24. Our focus is on tabular data, where tree-based models are state-of-the-art and thus the models we use in our experiments. Because these models can be represented as piecewise-constant functions, they also admit alignment, as shown below.

See proof in 25. This result shows that, in principle, exact alignment can be achieved for decision trees: the indicator features \(\mathbf{1}\{x \in R\}\) span the full disagreement space. In practice, however, such a construction is infeasible as \(\mathcal{R}_\mathcal{H}\) can be extremely large, and most regions will be empty.

To make the approach tractable, we turn to approximations. Instead of explicitly representing all tree regions, we approximate the disagreement space using quantile binning of the feature space. Concretely, motivated by the examples and our empirical findings, we use the following general-purpose matching function: \[\Phi(\boldsymbol{X})=\big(\Psi_1(X^{(1)}),\ldots,\Psi_d(X^{(d)}),\,\Psi_f(f(\boldsymbol{X}))\big),\] where each \(\Psi_j\) is a low-cardinality quantile binning. We use \(10\) bins in all our experiments. This approach has proven empirically effective for most common tabular models and datasets. See 34 for further discussion of alternative binning strategies for decision trees that may reduce approximation gap.

4 Explaining Model Performance↩︎

Understanding differences between source and target domains is crucial for providing guidelines to monitor similar changes or to anticipate further shifts. This is often approached by learning a constrained map \(T\) that transforms the source data to the target data \(T(\hat{P}_{\boldsymbol{X}}) = \hat{Q}_{\boldsymbol{X}}.\), e.g., via optimal transport [13].

A challenge with this problem is identifiability: without assumptions of invariance between the source and target data or knowledge about the causal structure, there are infinitely many plausible mappings \(T\), making it impossible to identify which one generated the given target data. Therefore, we need to specify certain invariance assumptions to make this problem identifiable. Unlike the estimating and improving components, the explaining aspect of our methodology focuses specifically on explaining sparse shifts, which assumes there exists a subset \(S \subset \{1, \dots, d\}\) such that \(P(\mathbf{X}_{\bar{S}} | \mathbf{X}_S, Y) = Q(\mathbf{X}_{\bar{S}} | \mathbf{X}_S, Y).\) Notably, the sparse joint shift generalizes both classical label shift and sparse covariate shift. See [1], [6] for a comprehensive analysis of these shifts. In this context, explaining the causes involves identifying the subset of features responsible for generating the shift [5].

See derivation in 28. Our approach is to operationalise this idea using EPA. We iteratively test candidate feature subsets. For each subset, we use EPA to perform a histogram matching, reweighting the source data to align the distribution of those specific features with the target’s. We then select the subset that yields the best overall approximation of the target data, as measured by an estimated KL distance. Specifically, among all subsets within \(p\%\) (default \(p=5\)) of the minimum, we then select the subset with the smallest size. We can also combine feature matchings with the target variable to capture shifts involving both inputs and labels. Refer to §6 and §4 for a step-by-step description of histogram matching and details on EPA’s explanations.

5 Improving Model Performance↩︎

To enhance the model’s performance under the target distribution \(Q\), we propose to leverage the reweighted data generated by the matching \(\Phi\) used for target error estimation, i.e., matching the binned version of all features and model prediction.

The use of reweighting to improve model performance has been extensively studied. [2] demonstrated that, for parametric models, reweighting can be beneficial when the model is misspecified (e.g., using a linear model when the true relationship is quadratic), but offers no advantage when the model is well-specified. However, recent works such as [14] and [15] report negative results, indicating that reweighting does not outperform classical training with uniform weights for high-capacity models. In contrast, [16] argue that these negative findings are due to evaluations conducted on well-specified cases, and they show that reweighting can be effective even for high-capacity models if the model is misspecified.

Additionally, it is well-known that directly using the exact ratio of source and target densities as weights may not be optimal in practice. [2] discusses this issue in detail, highlighting that while such weights are asymptotically unbiased, they can suffer from high variance in practice. For example, these weights may disproportionately emphasize a small subset of the data, potentially leading to poor model. This phenomenon reflects the classic bias-variance trade-off: accepting a small bias can lead to more stable and effective learning. Our reweighting technique addresses this trade-off by ensuring that the reweighted distribution remains close to the original distribution, thereby preventing the weights from concentrating on few points while still matching key statistics of the target data useful for estimation and adaptation.

When learning with reweighted data, most prior work focuses on training a completely new model [14]–[16]. In contrast, we propose adding a corrective term to the existing model \(f\). We find this approach more effective than retraining a new model from scratch under reweighting, as it builds upon the valuable information already captured by \(f\). Specifically, we employ a boosting procedure to minimize the error under the reweighted distribution \(P^\star\).

Starting with \(f_0 = f\), we iteratively update the model as \(f_t = f_{t-1} + \alpha_t h_t,\) where at each iteration \(t\), \(h_t \in \mathcal{H}\) is a base learner chosen from a class of functions \(\mathcal{H}\) to minimize the weighted loss \(h_t = \arg\min_{h \in \mathcal{H}} \sum_{i=1}^n \lambda_i \, \mathcal{L}\big( f_{t-1}(\mathbf{X}_i) + h(\mathbf{X}_i), \, Y_i \big),\) and \(\alpha_t\) is a learning rate determined by: \(\alpha_t = \arg\min_{\alpha} \sum_{i=1}^n \lambda_i \, \mathcal{L}\big( f_{t-1}(\mathbf{X}_i) + \alpha h_t(\mathbf{X}_i), \, Y_i \big).\) After \(M\) iterations, we obtain the improved model \(\tilde{f} = f_M\), which is expected to perform better under \(P^\star\) and, consequently, in the target domain \(Q\). This procedure can be efficiently implemented using gradient boosting libraries like XGBoost [17], which natively support instance weights.

The following theorem provides a formal guarantee: if the performance gain from boosting on the reweighted source data exceeds a certain threshold, the adapted model is provably better than the baseline on the target distribution.

Our matching \(\Phi\) is designed to make the discrepancy term \(\mathrm{disc}_\ell(P^\star_X, Q_X)\) small, while boosting decreases the reweighted error \(\epsilon_{\hat{P}^\star}(\tilde{f})\). Theorem 9 then ensures these two effects combine into a provable improvement guarantee over the baseline model \(f\). See proof in 29.

6 Experimental Setup↩︎

In this section, we detail the experimental setup used to determine which method – relying only on source data and target domain features – yields the best weights across three core objectives: estimating the target domain performance, identifying the features responsible for performance changes in sparse shift scenarios, and adapting the model to the target data.

We use three shift settings: sparse covariate, sparse joint and natural shifts. For each case, we have source and target datasets. We split the source data into a training set, used to train the initial model \(f\), and a calibration set, which is made accessible to each method for weight estimation and model adaptation. Similarly, we split the target data into an evaluation and calibration sets.

Sparse Covariate Shifts: We use the Adult [18], HELOC [19], and NHANES I [20] datasets. To create different source and target domains, we split the data based on the values of one feature at a time. For each continuous feature, we split the data at its median value. For each categorical feature, we exclude data belonging to one category. In either case, splitting results in two subsets, \(A\) and \(B\). We then randomly select 10% of \(B\) and add it to \(A\) to form our source data, while the remaining 90% of \(B\) constitutes our target data.

For each dataset, the number of distribution shifts studied equals the number of features times the number of splits: two for continuous features (since we split at the median) and one for each category of discrete features. This setup results in 292 different distribution shifts. A similar methodology was used in prior work [5] on a smaller scale with selected features, whereas we evaluate all possible feature combinations, aligning with recent tabular shift benchmarks [21].

Sparse Joint Shifts: In this setting, both features and labels shift. We apply the same subset-splitting strategy from the covariate-shift setup but additionally modify the label distribution by randomly undersampling the source label mean to \(0.2\) and oversampling the target label mean to \(0.5\).

Natural Shifts: We use the Folktables dataset [22], which comprises U.S.Census data (50 states plus Puerto Rico) from multiple years (2014-2018), capturing both geographic and temporal shifts. The predictive task is to determine whether income exceeds $50,000.

As with the previous datasets, we define subsets \(A\) and \(B\) as different state-year combinations, thereby leveraging the inherent distribution shifts across time and location.

For further details on the datasets, such as size and feature count, see §19.2.

Given our focus on tabular data, we utilize XGBoost [17], a well-established and high-performing model for this type of data. We perform extensive fine-tuning to ensure optimal performance at every step involving training. Details of the tuning process can be found in §18.

In 21, we give a detailed algorithmic description of all components of EPA. We use the following baselines for comparison:

• SEES-C and SEES-D: The two versions of the method introduced by [5]. SEES-C uses continuous features, and SEES-D uses a discretized version of the features to estimate the sample weights. We use the original implementation available on Github.

• KMM [9], [23] using a linear kernel to match the mean of all features.

• Domain Classifier Weighting (COV-CLF): A commonly used method [10] for covariate shift that trains a classifier to distinguish between source and target domains. The predicted probabilities \(\hat{p}(\mathbf{X})\) are used to derive sample weights as \(\hat{p}(\mathbf{X}) / (1 - \hat{p}(\mathbf{X}))\). We use XGBoost as classifier.

7 Experiments: Estimating↩︎

1.5pt

Bin ranges: Sparse Covariate: B1=(-0.149,-0.075], B2=(-0.075,0], B3=(0,0.106], B4=(0.106,0.211]; Sparse Joint: B1=(0.082,0.146], B2=(0.146,0.209], B3=(0.209,0.272], B4=(0.272,0.336]; Natural: B1=(0.042,0.141], B2=(0.141,0.24], B3=(0.24,0.338], B4=(0.338,0.437]. Notes: Estimation shows inaccuracy (lower is better); Explanation shows Acc/Corr pairs; Improvement shows average improvement (higher is better). Power/Prop column shows Power for Estimation and I-prop for Improvement. FP column shows false positives for Estimation only. Bold indicates best performance and underline second best performance within each task-shift. – indicates unavailable data. Note that SRC-UNI and TAR-UNI are not baselines but serve as control group

We evaluate each method’s ability to accurately estimate model performance on target domains and detect harmful distribution shifts using the following metrics:

1. Estimation Inaccuracy: Mean absolute error between estimated and true target errors (lower is better; \(\downarrow\)). True errors are defined as the absolute difference between predicted probabilities and true labels. The results are organized into increasing bins (B1-B4), which reflect the difference between the model error in the target data and the source data. A positive difference indicates a harmful shift; otherwise, the shift is considered benign.

2. Detection Capability: Assesses the method’s ability to detect when a harmful shift occurs—that is, when the target error exceeds the source error. Power (proportion of harmful shifts correctly detected; higher is better \(\uparrow\)) and False Positive rate (FP; lower is better \(\downarrow\)).

Takeaway : As shown in Table [results], EPA significantly outperforms the baselines in both estimation inaccuracy and detection capability across all shift scenarios, achieving 30–70% better performance, particularly in the more harmful bins (B3, B4). It is followed by SEES-C, SEES-D, and then COV-CLF. KMM is the poorest-performing method. In the natural shift setting, which reflects real-world data variability, EPA maintains strong accuracy and the highest power, while other methods degrade more noticeably. Due to scalability issues with larger datasets, KMM is omitted under natural shifts.

8 Experiments: Explaining↩︎

Next, we evaluate each method’s explanation performance, focusing on sparse covariate and sparse joint shifts. In contrast to the natural shift scenario, these cases have well-defined ground truth explanations, as they correspond to the splitting features used to generate the shifts, allowing straightforward evaluation. Among the methods considered, only our approach (EPA) and the two baselines (SEES-C and SEES-D) provide feature-level explanations.

We measure the average detection rate of the features responsible for the shifts (Acc). To account for potential biases arising from correlated features, we also compute the correlation (Corr), which quantifies the relationship between the selected features and the ground truth features across the dataset (see details in §19).

Takeaway . As shown in Table [results], EPA outperforms the baselines across most bins and shift types, with up to 60% higher accuracy under sparse joint shifts. SEES-C generally outperforms SEES-D.

9 Experiments: Improving↩︎

We now turn to adaptation effectiveness, which measures how effectively the reweighted source data from each method can be used to adapt the model and improve its performance on the target domain. We report two relative improvement metrics to isolate the contribution of the reweighting scheme from the strenght of the starting model:

• Average Improvement: The relative percentage difference in performance between the adapted model and the original model on the target data (higher is better; \(\uparrow\)). Results are averaged across bins representing different performance shift intensities.

• Improvement Proportion (I-prop): The proportion of shifts in which the adapted model outperforms the original model ranges from 0 to 1 (higher is better; \(\uparrow\)).

For model adaptation comparison, we include two additional baselines as control methods:

• Source Calibration with Uniform Weights (SRC-UNI): Uses the source calibration set with uniform weights as a naive approach without reweighting.

• Target Calibration with Uniform Weights (TAR-UNI): Uses the target calibration set with uniform weights as an idealized or strong baseline as it leverages ground truth labels from the target domain. Note that the target calibration set is not used for evaluation but comes from the same distribution as the one used for evaluation. It provides a reference for what can be achived with the same calibration budget when target labels are available.

Table [results] summarizes the results across three shift types: Sparse Covariate, Sparse Joint, and Natural shifts. A horizontal line separates the control baselines (SRC-UNI, TAR-UNI) from the actual baselines.

9.1 Sparse Covariate Shifts↩︎

Benign Shifts (Negative Bins) When distribution shifts benefit performance, most methods perform similarly, with differences in average improvement typically under \(0.5\%\) for bin B1.

Harmful Shifts (Positive Bins) When shifts become harmful, TAR-UNI (which leverages labeled target data) achieves the highest improvements. Among methods using only unlabeled target data, EPA stands out in the moderate bin \(B_3=(0, 0.106]\) with improvements close to TAR-UNI. In the highest bin \(B_4=(0.106, 0.211]\), all methods show comparable results, with COV-CLF having a slight edge in average improvement.

Interestingly, SRC-UNI (the naive approach) also performs well, supporting the findings from [16] that uniform weights may often suffice when the model is well specified. Next, we examine performance when the model is misspecified, where naive training provides less benefit, allowing a clearer comparison of each method’s reweighting effectiveness.

9.2 Sparse Joint and Natural Shifts↩︎

Under sparse joint shifts, both the features and the labels change, creating more pronounced and inherently harmful distribution shifts even for a strong model like XGBoost. As expected, TAR-UNI achieves the highest performance gains by leveraging labeled target data. Among methods that only use unlabeled target data, EPA achieves the largest overall improvements, followed closely by SEES-D, SEES-C, and COV-CLF.

For natural shifts, TAR-UNI again demonstrates superior performance, especially under severe shifts. EPA maintains its position as the strongest method among those not using target labels across all bins, while SEES-D, SEES-C, perform similarly, with COV-CLF showing slightly better results.

Takeaway : When using boosting approaches, even simple uniform weighting on a new independent source dataset can yield modest improvements as the algorithm naturally adapts. However, as distributional shifts become more pronounced relative to the model class, our results demonstrate that reweighting becomes particularly beneficial, aligning with findings from [16]. Among the methods using only unlabelled target data, EPA performs best overall across diverse types of shifts, especially when the shift becomes harmful. Another observation from these experiments is that even suboptimal weights for estimating model performance can still lead to substantial model improvement.

10 Ablation Studies↩︎

Appendices 13-17 present additional ablation experiments. These experiments demonstrate that the correction or boosting strategies (outlined in Section 5) consistently outperform learning from scratch under the studied shifts. We analyze which method generates reweighted data that aligns most closely with the target data. We conduct simulations demonstrating EPA’s stability advantage over KMM in importance weight estimation. Further experiments are included on linear models and regression problems.

11 Complexity Analysis and Running Time↩︎

We compare the computational costs of computing the sample weights for our method (EPA) against the direct baselines: SEES-C and SEES-D. In terms of complexity, EPA requires estimating a vector proportional to the number of features by minimizing the convex loss in ?? . In contrast, the baselines aim to learn basis functions to approximate the density ratio \(w_S(x, y) = \frac{dQ}{dP}(x_S, y)\). This approach involves a number of parameters dependent on the size of the basis class function, which can often exceed the dimensionality of the feature space.

Table 2 provides the observed runtimes for each method during sparse-shift experiments. These results show that EPA is significantly faster than its closest baselines, particularly the best baseline SEES-C, which requires an order of magnitude more time. Notably, COV-CLF, which simply trains a domain classifier to distinguish between source and target labels and converts the predictions into sample weights, is the quickest, requiring only 2 seconds at most.

8pt

12 Conclusion↩︎

We introduced EPA, an entropic projection-based framework to address three core challenges of distribution shift: estimating target domain performance, identifying the features driving the shift, and adapting models to improve target domain performance, all using only unlabeled target data and the source data. Empirical evaluations across various shifts demonstrate that EPA outperforms state-of-the-art baselines while providing improved computational efficiency.

A key element of our framework is the choice of a feature map \(\Phi\) that is well aligned with the model disagreement class under study. Future work could extend this framework beyond the tabular setting to unstructured domains, such as images, and to richer model classes, such as deep neural networks, where alignment may be achieved through learned representations. This direction is especially relevant for cases where explicit mappings are difficult or intractable, such as tree-based models. Another promising avenue is to employ hypothesis-specific discrepancy bounds rather than worst-case bounds over the entire function class, which could lead to tighter guarantees and improved adaptability.

Disclaimer↩︎

This paper was prepared for informational purposes by the Artificial Intelligence Research group of JPMorgan Chase & Co., and its affiliates (“JP Morgan”) and is not a product of the Research Department of JP Morgan. JP Morgan makes no representation, and warranty whatsoever, and disclaims all liability, for the completeness, accuracy or reliability of the information contained herein. This document is not intended as investment research or investment advice, or a recommendation, offer or solicitation for the purchase or sale of any security, financial instrument, financial product or service, or to be used in any way for evaluating the merits of participating in any transaction, and shall not constitute a solicitation under any jurisdiction or to any person, if such solicitation under such jurisdiction or to such person would be unlawful.

© 2026 JP Morgan Chase & Co. All rights reserved.

Checklist↩︎

1. For all models and algorithms presented, check if you include:

   1. A clear description of the mathematical setting, assumptions, algorithm, and/or model. [Yes. Proofs and detailed statements of the theorems, including their assumptions, are provided in the Appendix. In Section 21, we give a detailed algorithmic description of all components of EPA.]

   2. An analysis of the properties and complexity (time, space, sample size) of any algorithm. [Yes]

   3. (Optional) Anonymized source code, with specification of all dependencies, including external libraries. [Code will be released after acceptance]

2. For any theoretical claim, check if you include:

   1. Statements of the full set of assumptions of all theoretical results. [Yes]

   2. Complete proofs of all theoretical results. [Yes]

   3. Clear explanations of any assumptions. [Yes]

3. For all figures and tables that present empirical results, check if you include:

   1. The code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL). [Code to reproduce all experiments will be released after acceptance]

   2. All the training details (e.g., data splits, hyperparameters, how they were chosen). [Yes. Additional details can be found in Appendix, and code to reproduce all experiments will be released after acceptance]

   3. A clear definition of the specific measure or statistics and error bars (e.g., with respect to the random seed after running experiments multiple times). [Yes. We report mean, standard deviation, min, and max for runtime statistics related to weight computation (Table 2). However, due to the extensive hyperparameter tuning detailed in Appendix 18 and the large number of shift scenarios evaluated (292 for sparse covariate shifts alone), computing comprehensive error bars (e.g., via multiple runs with full tuning) for the primary performance metrics in Tables 1 was computationally prohibitive: each full evaluation (with hyperparameter tuning) took approximately one week per method. However, we ensured robustness by evaluating a wide range of settings and reporting aggregated results across multiples shifts and several datasets.]

   4. A description of the computing infrastructure used. (e.g., type of GPUs, internal cluster, or cloud provider). [Yes. The weight computation experiments were conducted on a MacBook Pro M2, as noted in the caption of Table 2. This figure provides mean, std, min, and max runtimes. Other experiments were run on AWS c5.4xlarge (16 vCPU, 32 GB RAM).]

4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets, check if you include:

   1. Citations of the creator If your work uses existing assets. [Yes]

   2. The license information of the assets, if applicable. [Not Applicable]

   3. New assets either in the supplemental material or as a URL, if applicable. [Not Applicable]

   4. Information about consent from data providers/curators. [Not Applicable]

   5. Discussion of sensible content if applicable, e.g., personally identifiable information or offensive content. [Not Applicable]

5. If you used crowdsourcing or conducted research with human subjects, check if you include:

   1. The full text of instructions given to participants and screenshots. [Not Applicable]

   2. Descriptions of potential participant risks, with links to Institutional Review Board (IRB) approvals if applicable. [Not Applicable]

   3. The estimated hourly wage paid to participants and the total amount spent on participant compensation. [Not Applicable]

Appendix: Table of Contents↩︎

13 Stability: EPA vs KMM↩︎

In this section, we present a simulation to demonstrate the advantages of the EPA framework over KMM, complementing our earlier experiments on real-world datasets. As discussed, the regularization in EPA reduces variance in importance weight estimation. To validate this, we generate datasets multiple times under controlled conditions.

We define two distributions, \(P \sim \mathcal{N}(0_p, 5)\) and \(Q \sim \mathcal{N}(2\times I_p, 1)\), and sample 1000 points from each. The target output is modeled as \(Y = f(X) = \sum_{i=1}^{p} X_i\), with \(p = 100\).

Both KMM and EPA are used to estimate importance weights, and we fit a linear model using the reweighted samples. Note that the optimal model remains \(f\), even after the distribution shift. Performance is evaluated by the Relative Absolute Difference with respect to the optimal model by \[\text{Relative Absolute Difference} = \frac{\text{MSE}(\text{linear\_reweighted\_method}) - \text{MSE}(\text{true\_model})}{\text{MSE}(\text{true\_model})}.\]

The experiment is repeated 500 times, and the results, illustrated in figure 1, highlight that EPA achieves zero variance, consistently recovering the true model, while KMM exhibits high variance. Moreover, in figure [fig:entropy] we show the distribution of the KL distance between the sample weights of each method and the uniform weights. It shows that, as expected, the EPA weights are closer to the uniform distribution than those estimated by KMM.

14 Learning from Scratch vs. Correcting the Existing Model↩︎

We evaluate the effectiveness of training a new model from scratch compared to our proposed method of adding corrective terms to an existing model (as described in Section 5). Table 3 presents the results. When comparing these findings with Table [results] (where correction was applied), we observe that training from scratch is significantly less effective across all experiments.

In the “learning from scratch” setting, most methods fail to improve performance and even degrade performance in most cases. The best-performing approach in this case is calibration with uniform weights (SRC-UNI), which produces results similar to the original model. This outcome is expected, as the calibration data follows the same distribution as the data used to train the original model.

This pattern holds across all other experiments. These findings underscore the limitations of training a new model on reweighted data in sparse shifts. In contrast, our correction approach effectively utilizes reweighted data to achieve better adaptation and performance improvements.

2pt

15 Experiments: Estimation and Improvement using a Linear base↩︎

This section evaluates the performance of the methods in terms of estimation and improvement under sparse covariate shifts, sparse joint shifts, and natural shifts using a linear base model. This setup highlights scenarios where the model class is underperforming, and reweighting is known to help [16] in mitigating the performance degradation.

15.1 Sparse Covariate Shift: Linear Model↩︎

15.1.0.1 Key Observations:

• Estimation Inaccuracy (Table 4):

  • EPA achieves the best estimation accuracy across all bins, with significantly lower errors compared to other methods. For detection power, EPA significantly outperforms baselines, achieving 0.848 compared to approximately 0.14 for others.

  • Baselines SEES-C, SEES-D, and COV-CLF struggle, particularly in cases of large performance shifts.

• Model Improvement (Table 5):

  • While most methods perform similarly in lower shift bins, EPA demonstrates the highest average improvement in bins with the largest shifts.

  • Uniform calibration (SRC-UNI) fails to deliver meaningful improvements, often resulting in negligible or negative performance gains.

2pt

2pt

15.2 Sparse Joint Shift: Linear Model↩︎

15.2.0.1 Key Observations:

• Estimation Inaccuracy (Table 6):

  • EPA achieves near-perfect accuracy with negligible errors across all bins.

  • Baselines show large errors and low detection power (near \(0\)).

• Model Improvement (Table 7):

  • All methods demonstrate similar improvement levels, but EPA lags slightly in improvement proportion (I-prop) compared to the baselines.

  • SRC-UNI continues to underperform, offering limited or negligible improvements.

2pt

2pt

15.3 Natural Shift: Linear model↩︎

15.3.0.1 Key Observations:

• Estimation Inaccuracy (Table 8):

  • EPA demonstrates consistently strong performance, achieving low estimation errors across all bins.

  • Baselines such as SEES-C, SEES-D, and COV-CLF exhibit significantly higher errors.

• Model Improvement (Table 9):

  • All methods show comparable levels of improvement.

  • TAR-UNI and SRC-UNI fail to provide meaningful enhancements in lower bins. TAR-UNI gives the highest improvement on the most harmful bin.

2pt

2pt

15.3.0.2 Takeaway:

In terms of estimation, EPA is by far the most accurate method, significantly outperforming all baselines. Regarding improvement, all methods perform similarly overall, highlighting that reweighting can still drive performance gains even when error estimation is suboptimal—showing the learning algorithm’s ability to benefit from reweighted data itself. Additionally, when the model is not high-performing, classical training with uniform weighting (SRC-UNI) fails to improve and can even degrade performance. Interestingly, revealing labels (TAR-UNI) can also underperform in this context, emphasizing the challenges of adaptation when the model class and data distribution are misaligned.

16 Experiments on Regression Model↩︎

In this section, we extend the evaluation of our method to regression problems using the California Housing Pricing dataset [18] following the previous section by using a Linear base. Among the baselines, only COV-CLF is applicable, as SEES-C and SEES-D are not designed for regression settings.

16.0.0.1 Key Observations:

• Estimation Inaccuracy (Table 8):

  • EPA achieves the lowest estimation errors across all bins, significantly outperforming the COV-CLF baseline.

  • The detection power of EPA is high (0.889), while COV-CLF fails to detect shifts at all (Power (\(\uparrow\)) = 0).

• Model Improvement (Table 5):

  • EPA demonstrates consistent improvements across all bins, slightly outperforming the COV-CLF baseline.

2pt

2pt

16.0.0.2 Takeaway:

Once again, EPA significantly outperforms the baseline in terms of error estimation. In terms of model improvement, EPA shows a slight advantage. This example also demonstrates that it is possible to have weights that are not necessarily accurate for estimation but can still contribute to improving model performance.

17 Target Data Estimation↩︎

In this section, we evaluate how well each method reweights the source data to approximate the target data. Two metrics are used for evaluation:

• Mean Matching for All Features: Measures the discrepancy in feature means between the reweighted source data and the target data.

• Histogram Matching for All Features: Measures the discrepancy in feature distributions (histogram binning) computed marginally for each features independently and averaged. The number of bins is set at 10.

17.1 Sparse Covariate Shift↩︎

Table 12 summarizes the results for sparse covariate shifts. EPA achieves the lowest values for both metrics, indicating superior alignment with the target data. SEES-C performs moderately well, followed by COV-CLF and SEES-D, which show larger discrepancies.

8pt

17.2 Sparse Joint Shift↩︎

Table 13 presents the results for sparse joint shifts. Similar to the covariate shift setting, EPA demonstrates the best performance, achieving the smallest mean approximation. Regarding histogram differences SEES-C is slightly better, while COV-CLF and SEES-D perform worse, with higher values for both metrics.

8pt

17.2.0.1 Takeaway:

These findings reinforce the effectiveness of EPA in approximating target distributions, making it a reliable choice for reweighting source data under distribution shifts.

18 Hyperparameter Optimization for XGBoost Model↩︎

To fine-tune the XGBoost model, a hyperparameter optimization process was conducted using 1000 iterations. The optimization strategy varied depending on the booster type (gbtree vs gblinear). The following details the process used to suggest and train the optimal hyperparameters. The hyperparameters were dynamically suggested during training based on the booster type:

• For gbtree boosters, the following hyperparameters were optimized:

  • eta (learning rate): Log-uniformly sampled in the range [0.01, 0.3].

  • max_depth: Integer value between 3 and 10.

  • subsample: Float value between 0.6 and 1.0.

  • colsample_bytree: Float value between 0.6 and 1.0.

  • lambda (L2 regularization): Log-uniformly sampled in the range [1e-3, 10.0].

  • alpha (L1 regularization): Log-uniformly sampled in the range [1e-3, 10.0].

  • min_child_weight: Integer value between 1 and 10.

• For the gblinear booster, the following hyperparameters were optimized:

  • lambda (L2 regularization): Log-uniformly sampled in the range [1e-3, 10.0].

  • alpha (L1 regularization): Log-uniformly sampled in the range [1e-3, 10.0].

The trained model was optimized for performance on the validation dataset, leveraging early stopping and dynamically selected hyperparameters to achieve the best results using the package Optuna [24].

19 Additional experiments details↩︎

19.1 Correlation metric for the explainability part↩︎

Given two features, \(A\) and \(B\), we determine their types and apply an appropriate correlation metric:

• Both Continuous: We use the Pearson correlation coefficient, which measures the linear relationship between two continuous variables: \[\rho(A, B) = \frac{\mathrm{cov}(A, B)}{\sigma_A \sigma_B},\] where \(\mathrm{cov}(A, B)\) is the covariance of \(A\) and \(B\), and \(\sigma_A\) and \(\sigma_B\) are their respective standard deviations.

• Both Categorical: We use Cramér’s V, which is derived from the chi-squared (\(\chi^2\)) statistic computed from the contingency table: \[V = \sqrt{\frac{\chi^2/n}{\min(k_A - 1, k_B - 1)}},\] where \(\chi^2\) is the chi-squared statistic, \(n\) is the total number of observations, and \(k_A\) and \(k_B\) are the number of categories in \(A\) and \(B\), respectively.

• One Continuous, One Binary: We use the Point-Biserial correlation coefficient, which measures the correlation between a continuous variable and a binary categorical variable: \[r_{pb} = \frac{\overline{X}_1 - \overline{X}_0}{s} \sqrt{\frac{n_1 n_0}{n(n - 1)}},\] where \(\overline{X}_1\) and \(\overline{X}_0\) are the means of the continuous variable for the two binary groups, \(s\) is the standard deviation of the continuous variable, and \(n_1\), \(n_0\), and \(n\) are the sizes of the respective groups and the total sample.

Finally, we take the absolute value of each metric to ensure that all correlation values fall within the range \([0,1]\). We compute the values using all the dataset.

19.2 Dataset Details↩︎

Our experiments leverage a diverse array of publicly available, real-world datasets. These datasets vary in terms of sample size, feature dimensionality, and application domain, facilitating a comprehensive evaluation of our proposed framework under different conditions. Specific details for each dataset are provided below. For all datasets, categorical features were one-hot encoded, and the resulting feature dimensions are reported.

• Adult: This dataset, commonly used for income prediction (binary classification of whether income exceeds $50K/yr), is sourced from the UCI Machine Learning Repository [18]. It originally contains 14 features (a mix of continuous and categorical) for 48,842 samples. After one-hot encoding of categorical variables, the feature space expands to 109 dimensions.

• HELOC (Home Equity Line of Credit): Provided by [19], this dataset is used for credit risk assessment. It comprises 9,871 samples and 24 features. The task is to predict whether a homeowner will default on their HELOC.

• NHANES I (National Health and Nutrition Examination Survey): This dataset consists of health-related data from the first National Health and Nutrition Examination Survey, made available by the [20]. In our experiments, we use a version with 8,593 samples and 18 features, focusing on a mortality prediction task.

• California Housing Prices: Also from the UCI Machine Learning Repository [18], this dataset contains 20,640 samples and 9 features. It is a standard benchmark for regression tasks, where the objective is to predict the median house value in California districts.

Folktables (ACS-based Tasks)↩︎

The Folktables suite [22] provides benchmark tasks derived from the U.S. Census Bureau’s American Community Survey (ACS). These datasets allow for the study of model performance under various demographic shifts. Table 14 summarizes the characteristics of several tasks available within this suite.

For our experiments involving natural distribution shifts, we primarily utilize the ACSIncome task from Folktables. While the original ACSIncome task includes 10 features, after one-hot encoding of its categorical variables, the dimensionality increases to approximately 284 features. The full ACSIncome dataset across all available years and states contains over 1.6 million records. In our experimental setup, we generate source and target domains by selecting data for specific U.S. states and years. This typically results in experimental datasets ranging from 30,000 to 100,000 samples, though some state-year combinations may yield larger datasets.

20 Additional Related Works Discussion↩︎

One important related body of literature is representation learning. Methods within this field have progressed in managing unstructured data such as images by using mean-matching techniques to align learned embedding features for adapting to distributional changes, as highlighted in works like [25], [25]–[29] or sub-domain alignment methods [30]. However, the EPA approach differs from these works. These approaches primarily address domain adaptation alone. In contrast, EPA is specifically designed to tackle three interrelated problems simultaneously: (1) performance evaluation, (2) explanation of shifts, and (3) model adaptation. Thus, these earlier methods do not fully compare across all three aspects. In addition, direct empirical comparisons with these methods are challenging, as many are tailored to specific model architectures, such as deep neural networks with unique loss functions and feature alignment strategies—whereas EPA is not restricted to a particular model type. Rather than replacing existing models, EPA integrates a corrective component using estimated importance weights, setting it apart in its goals and operation. Putting differently, EPA is a post-hoc method. However, extending EPA to fields like computer vision or textual data would require an additional phase for identifying meaningful feature representations where mean-matching could be effectively applied. Developing this extension is complex and separate from EPA’s core function and thus is not covered in this paper. As stated in our conclusion, we have left this extension for future exploration.

Given EPA’s core mechanism of deriving importance weights, its contributions are particularly relevant to the literature on reweighting techniques for distribution shift. Within this domain, SEES [5] is our main competitor and best among the reweighting techniques as demonstrated in their paper. Our work therefore emphasizes comparison with such state-of-the-art reweighting approaches to clearly benchmark EPA’s advancements in robust weight estimation, along several additional dimensions such as performance prediction, explanation of shifts, and model adaptation.

Lastly, EPA can be contrasted with distributionally robust optimization (DRO) strategies [31]. DRO methods typically aim to optimize model performance under worst-case distributional shifts, often by reweighting data from predefined groups. While effective for enhancing worst-case generalization, DRO approaches require these group definitions. Specifying such groups comprehensively can be challenging, may not capture all relevant types of shifts, or the necessary groupings might not be apparent in all domains (for instance, while image datasets sometimes allow for grouping by attributes like gender or hair color, this approach has its limitations and is not universally applicable). Furthermore, DRO techniques typically do not offer the direct target performance estimation or the explanatory insights inherent to the EPA framework.

21 Algorithmic Description of EPA↩︎

This section provides a detailed algorithmic description of the proposed Entropic Projection Alignment (EPA) framework. The following algorithms outline the core components and applications of EPA:

• Algorithm 2 describes the core EPA procedure for computing importance weights and the reweighted source distribution.

• Algorithm 3 details how EPA is used to estimate model performance on an unlabeled target domain.

• Algorithm 4 outlines the process for identifying feature subsets responsible for distribution shifts, utilizing Algorithm 6 for discrepancy measurement.

• Algorithm 5 shows the methodology for adapting an existing model to improve its performance in the target domain using EPA-derived weights.

• Algorithm 6 provides the method for computing the histogram-based symmetric KL divergence used in the shift explanation process.

• Algorithm 7 provide the method for computing the target estimation with label prevalence.

22 Theoretical analysis of EPA↩︎

22.1 Proof of Proposition 1↩︎

To prove Proposition 1, we begin with the fundamental theorem governing the minimization of KL divergence relative to a prior measure under linear moment constraints. This result corresponds to Theorem A.1 in [12], stated there without proof as a consequence of the theorems of [11].

22.2 EPA’s weights minimises variance upper-bound by construction↩︎

Interestingly, EPA’s definition in 1 can be rewritten as the following optimisation problem:

Let \(S=(X_{1:n},\tilde{X}_{1:m})\) denote the (unlabeled) features from source and target. EPA produces weights \[\lambda(S)\in\Delta_n,\qquad \sum_{i=1}^n \lambda_i\,\Phi(X_i)=\widehat\mu_Q,\qquad \lambda_i\propto \exp\{\langle\xi,\Phi(X_i)\rangle\},\] which depend only on \(S\) (not on labels). For any \(f\in\mathcal{H}\), define the centered label-noise term \[\Delta_f\;:=\;\sum_{i=1}^n \lambda_i\Big(\ell(f(X_i),Y_i)-\mathbb{E}[\ell(f(X_i),Y_i)\mid X_i]\Big).\] Write \(u\) for the uniform weights \(u_i=1/n\), and set \(N_{\mathrm{eff}}(\lambda):=1/\sum_{i=1}^n\lambda_i^2\).

23 Proof of Proposition 2↩︎

24 Proofs of the results on Alignment↩︎

24.1 Proof of Proposition 3↩︎

24.2 Alignment Under Covariate Shift↩︎

24.3 Alignment of prediction histogram↩︎

24.3.0.1 Setup.

Let \(p_f(X)\in[0,1]\) and partition \([0,1]\) into bins \(B_1,\ldots,B_K\). Define model-derived features \(\psi_j(x):=\mathbb{1}\{p_f(x)\in B_j\}\) and set \(\Phi(x)=(\psi_1(x),\ldots,\psi_K(x))\). Let \(\alpha_{D,j}:=\Pr_D(p_f(X)\in B_j)\) so that \(\sum_j \alpha_{D,j}=1\). Suppose the conditional risk is well-approximated as piecewise-constant on the bins: \(r_j\approx \mathbb{E}[\ell(f(X),Y)\mid p_f(X)\in B_j]\). Then \[\epsilon_D(f)\approx \sum_{j=1}^K \alpha_{D,j}\, r_j.\]

24.3.0.2 EPA with prediction-histogram features.

EPA enforces \(\mathbb{E}_{P^\star}[\psi_j]=\mathbb{E}_Q[\psi_j]\) for all \(j\), i.e. it matches the entire prediction histogram: \(\alpha_{P^\star,j}=\alpha_{Q,j}\). Hence \(\epsilon_{P^\star}(f)\approx \epsilon_Q(f)\), and exactly equals it if the piecewise-constant assumption holds.

25 Proof of Theorem 5 - Alignment for Decision Trees↩︎

25.1 The Hypothesis Class of Decision Trees↩︎

Let the input space be \(\mathcal{X}\subseteq \mathbb{R}^d\). A decision tree classifier partitions the input space into a set of disjoint regions and assigns a class label to each region.

Let \(\mathcal{H}\) be the class of all decision trees on \(\mathcal{X}\) up to a maximum depth \(D\). For our analysis, we consider the binary classification setting where \(\mathcal{Y}=\{0,1\}\) and the loss function is the 0–1 loss, \(\ell(y_1, y_2) = \mathbf{1}\{y_1 \neq y_2\}\).

26 Characterizing the Disagreement Space↩︎

To construct an aligned feature map, we must first understand the structure of the disagreement functions \(g(x) = \mathbf{1}\{f(x) \neq h(x)\}\) where \(f, h \in \mathcal{H}\).

26.1 Constructing the Aligned Feature Map↩︎

The previous proposition shows that any disagreement function is an indicator function of a set composed of elementary hyper-rectangles. This suggests that the basis for our feature map should be the indicator functions of these elementary regions.

26.2 Main Result and Proof of Alignment↩︎

We now prove that this construction of \(\Phi\) provides perfect alignment.

27 Approximate Alignment and the Residual↩︎

Let \(\mathcal{A}_\Phi=\{x\mapsto \beta_0+\beta^\top\Phi(x):\beta_0\in\mathbb{R},\;\beta\in\mathbb{R}^k\}\) be the affine hull of \(\Phi\). For \(g\in\mathcal{G}_{\mathcal{H},\ell}\), define the residual relative to \((\beta_0,\beta)\): \(r_g(x)=g(x)-\beta_0-\beta^\top\Phi(x).\)

When moments are matched only approximately, let \(\Delta=\mathbb{E}_{Q_X}[\Phi]-\mathbb{E}_{P^\star_X}[\Phi]\) and fix any norm pair \((\|\cdot\|,\|\cdot\|_\ast)\) on \(\mathbb{R}^k\).

27.1 How to Control the Residual: Examples↩︎

Takeaway: Coarse bin features turn interval-type disagreements into affine functions of \(\Phi\), with residuals controlled by grid width.

28 Proof of proposition 8↩︎

29 Proof of Theorem 9↩︎

Recall the weight construction:

These weights define a reweighted source distribution: \[\hat{P}^\star = \sum_{i=1}^n \lambda_i \delta_{(X_i, Y_i)}\]

The first part of the theorem is proved in the following lemma.

Now, we control the population discrepancy by combining exact empirical moment matching with the approximate-alignment bound from §27.

30 Proof analysing the empirical version of EPA with its population version.↩︎

Note that this section complements [12], which provides an asymptotic analysis of the entropic variable projection, by instead offering a non-asymptotic analysis of the solution.

30.1 Problem Formulation↩︎

We recall the population and empirical problems.

• Population Problem: Given a source distribution \(P_X\) and target moment \(\mu_Q = \mathbb{E}_{Q_X}[\Phi(X)]\), find \(P^* = \mathop{\mathrm{argmin}}_{P \ll P_X} \mathrm{KL}(P \| P_X)\) subject to \(\mathbb{E}_{P}[\Phi(X)] = \mu_Q\). The solution is given by the density \(\frac{dP^*}{dP_X}(x) = \exp(\langle \xi^*, \Phi(x) \rangle - H(\xi^*))\), where \(\xi^*\) is the true dual parameter and \(H(\xi)\) is the log-partition function.

• Empirical Problem: Given source samples \(\{X_i\}_{i=1}^n \sim P_X\) and target samples \(\{Z_j\}_{j=1}^m \sim Q_X\), we compute the empirical measure \(\hat{P}_n = \frac{1}{n}\sum_{i=1}^n \delta_{X_i}\) and estimate the target moment as \(\hat{\mu}_{Q,m} = \frac{1}{m}\sum_{j=1}^m \Phi(Z_j)\). We then find \(\hat{P}_{n,w}^* = \mathop{\mathrm{argmin}}_{\tilde{P} \ll \hat{P}_n} \mathrm{KL}(\tilde{P} \| \hat{P}_n)\) subject to \(\mathbb{E}_{\tilde{P}}[\Phi(X)] = \hat{\mu}_{Q,m}\). The solution is a discrete distribution on \(\{X_i\}_{i=1}^n\) with weights \[\lambda_i^* \;=\; \frac{\exp\{\langle \hat{\xi}_{n,m}^*, \Phi(X_i) \rangle\}}{\sum_{\ell=1}^n \exp\{\langle \hat{\xi}_{n,m}^*, \Phi(X_\ell) \rangle\}},\] where \(\hat{\xi}_{n,m}^*\) is the empirical dual parameter.

For simplicity, we will denote \(\hat{\xi}_{n,m}^*\) as \(\hat{\xi}^*\).

30.2 Assumptions↩︎

30.3 Main Theorem↩︎

The following theorem provides finite-sample bounds for the estimated continuous model.

The remainder of this document is dedicated to proving this theorem and connecting it to the discrete empirical solution.

31 Proof of the Bound on the Dual Parameter↩︎

The proof relies on establishing a basic inequality that links the output error \(\|\hat{\xi}^* - \xi^*\|\) to input errors, which are then bounded using concentration inequalities.

31.1 Step 1: A Basic Inequality↩︎

Let \(H(\xi)=\log \mathbb{E}_{P_X}[e^{\langle \xi,\Phi(X)\rangle}]\) and \(H_{\hat{P}_n}(\xi)=\log \Big(\frac{1}{n}\sum_{i=1}^n e^{\langle \xi,\Phi(X_i)\rangle}\Big)\). Set \(G=\nabla H\) and \(\hat{G}_n=\nabla H_{\hat{P}_n}\). The optimality conditions are \(G(\xi^*)=\mu_Q\) and \(\hat{G}_n(\hat{\xi}^*)=\hat{\mu}_{Q,m}\).

By strong convexity of \(H\) on \(\Xi\) (Assumption 3), the gradient map \(G\) is \(\sigma\)-strongly monotone: \(\langle G(y)-G(x), y-x\rangle \ge \sigma\|y-x\|_2^2\). Hence, by Cauchy–Schwarz, \(\|G(y)-G(x)\|_2 \ge \sigma\|y-x\|_2\). Thus \[\sigma \|\hat{\xi}^* - \xi^*\|_2 \le \|G(\hat{\xi}^*) - G(\xi^*)\|_2.\] Substituting \(G(\xi^*) = \mu_Q\) and using the empirical optimality condition, \[\begin{align} \sigma \|\hat{\xi}^* - \xi^*\|_2 &\le \|G(\hat{\xi}^*) - \mu_Q\|_2 \\ &= \|G(\hat{\xi}^*) - \hat{G}_n(\hat{\xi}^*) + \hat{\mu}_{Q,m} - \mu_Q\|_2 \\ &\le \|G(\hat{\xi}^*) - \hat{G}_n(\hat{\xi}^*)\|_2 + \|\hat{\mu}_{Q,m} - \mu_Q\|_2. \end{align}\] This is our basic inequality.

31.2 Step 2: Bounding the Error Terms via Concentration↩︎

31.3 Step 3: Combining the Bounds↩︎

By a union bound applied to Lemmas 3 and Assumptions 4, with probability at least \(1-2\delta\), \[\sigma \|\hat{\xi}^* - \xi^*\|_2 \le \alpha_n(\delta) + \beta_m(\delta),\] which yields Item (1) of Theorem 24 after dividing by \(\sigma\).

32 Proof of the Bounds on the Distribution Error↩︎

We now translate the bound on the dual parameter error into bounds on statistical distances between the continuous distributions \(P^*\) and \(\hat{P}_n^* = P_{\hat{\xi}^*}\).

32.1 Step 1: Bounding the KL Divergence↩︎

The KL divergence between two distributions from an exponential family is equal to the Bregman divergence of their natural parameters with respect to the log-partition function.

Combining Propositions 25 and 26 proves Item (2) of Theorem 24.

32.2 Step 2: Bounding the Total Variation Distance↩︎

Applying Lemma 4 together with Item (2) of Theorem 24, we obtain Item (3): \[d_{\text{TV}}(P^*, \hat{P}_n^*)^2 \le \tfrac{1}{2}\mathrm{KL}(P^* \| \hat{P}_n^*) \le \tfrac{L}{4}\,\|\hat{\xi}^* - \xi^*\|_2^2.\]

33 Connection to the Discrete Empirical Solution↩︎

We now relate the continuous tilted model \(P_{\hat{\xi}^*}\) to the practical discrete optimizer \(\hat{P}_{n,w}^*\). Since these measures need not be mutually absolutely continuous, Wasserstein distance is the appropriate metric.

34 Discussion of The Approximation Gap of Axis-Aligned Binning↩︎

Matching moments with \(\Phi_{hist}\) enforces that the marginal probability mass in each bin is the same for \(P^\star_X\) and \(Q_X\), but it does not constrain the joint distribution of features. This lack of constraint on feature interactions is the source of the gap.

To close the gap, we suggest exploring feature maps that can capture more complex relations, which may be worth investigating in future work.

34.1 Method 1: Feature Crosses↩︎

We can explicitly model interactions by binning the joint distribution of feature pairs (or higher-order tuples).

34.2 Method 2: Tree-Based Binning↩︎

A more adaptive approach is to use a decision tree to define the bins, as its structure naturally captures interactions relevant to the data.

By construction, the regions \(\{L_i\}\) are sensitive to the parts of the feature space where \(P_X\) and \(Q_X\) differ. This makes \(\Phi_{tree}\) a far better approximation of the disagreement space \(\mathcal{G}_{\mathcal{H},\ell}\) than \(\Phi_{hist}\).

Another solution is to use a tighter bound of 23 that is specific to a given hypothesis \(h\) for the decision tree in question.

The original discrepancy term can be replaced by one that measures the divergence only between the specific hypothesis \(h\) and the ideal joint hypothesis \(h_{\text{joint}}\).

This bound is tighter than the original, as the hypothesis-specific term is necessarily less than or equal to the supremum taken over all pairs in \(\mathcal{H}\): \[\big|d_{Q_X}(h, h_{\text{joint}}) - d_{P_X}(h, h_{\text{joint}})\big| \le \sup_{g,g'\in\mathcal{H}}\big|d_{Q_X}(g,g')-d_{P_X}(g,g')\big|\]

Using the tighter, hypothesis-specific bound for your decision tree provides a more precise and practical target for your feature mapping.

Instead of needing a feature map that works for all possible trees, you only need one that works for your specific tree (according to its leaves). This is a less demanding and more realistic goal.

35 EPA when the target mean lies outside the convex hull↩︎

Let \(\Phi_n := \operatorname{conv}\{\Phi(X_1),\dots,\Phi(X_n)\}\subset\mathbb{R}^k\) be the convex hull of source features and let \(t\in\mathbb{R}^k\) denote the (target) mean we wish to match. If \(t\notin\Phi_n\), the equality-constrained EPA problem is infeasible. A principled remedy is to project \(t\) onto \(\Phi_n\) (with respect to a chosen norm) and then impose equality to this projection.

35.0.0.1 Step 1: metric projection onto \(\Phi_n\).

Fix a norm \(\|\cdot\|\) on \(\mathbb{R}^k\). Define the projection \[s^\star \;\in\; \mathop{\mathrm{argmin}}_{s\in \Phi_n}\;\|s - t\| \qquad\text{and its distance}\qquad \delta_n(t)\;:=\;\|s^\star - t\|.\] Equivalently, in the weights \(\lambda\in\Delta_n\), \[\label{eq:proj-program} s^\star \in \mathop{\mathrm{argmin}}_{\lambda\in\Delta_n}\;\Big\|\sum_{i=1}^n \lambda_i \Phi(X_i) - t\Big\|.\tag{6}\] For \(\|\cdot\|_2\) this is a small QP; for \(\ell_1/\ell_\infty\) it is an LP.

35.0.0.2 Step 2: EPA at the projected target.

Run EPA with the equality constraint set to \(s^\star\): \[\label{eq:epa-projected} \lambda^\star \in \mathop{\mathrm{argmin}}_{\lambda\in\Delta_n}\;\mathrm{KL}(\lambda\Vert u) \quad\text{s.t.}\quad \sum_{i=1}^n \lambda_i\,\Phi(X_i) = s^\star, \qquad u_i=\tfrac1n.\tag{7}\] By Theorem 12, the unique solution has the exponential-tilt form \[\lambda_i^\star \;=\; \frac{\exp\{\langle \xi^\star,\Phi(X_i)\rangle\}}{\sum_{j=1}^n \exp\{\langle \xi^\star,\Phi(X_j)\rangle\}}, \qquad \sum_{i=1}^n \lambda_i^\star\,\Phi(X_i)=s^\star,\] where \(\xi^\star\) is the (unique) minimizer of \(H(\xi)-\langle \xi,s^\star\rangle\) with \(H(\xi)=\log\!\big(\frac{1}{n}\sum_j e^{\langle \xi,\Phi(X_j)\rangle}\big)\).

35.0.0.3 Equivalence to hard proximity EPA.

The projection approach is equivalent to solving EPA with a hard proximity tolerance equal to the projection distance: \[\label{eq:epa-ball} \min_{\lambda\in\Delta_n}\;\mathrm{KL}(\lambda\Vert u) \quad\text{s.t.}\quad \Big\|\sum_{i=1}^n \lambda_i\,\Phi(X_i) - t\Big\| \le \delta_n(t).\tag{8}\] Indeed, let \(B:=\{s:\|s-t\|\le \delta_n(t)\}\). By definition of \(s^\star\), \(B\cap \Phi_n=\{s^\star\}\), so the feasible set of 8 is exactly the affine face \(\{\lambda\in\Delta_n:\sum_i \lambda_i\Phi(X_i)=s^\star\}\), and 8 reduces to 7 . Consequently the unique optimizer is the same \(\lambda^\star\).

35.0.0.4 Dual viewpoint and computation.

Let \(\|\cdot\|_*\) denote the dual norm. The Lagrange dual of 8 is the strictly convex problem \[\min_{\xi\in\mathbb{R}^k}\;\; H(\xi) - \langle \xi, t\rangle - \delta_n(t)\,\|\xi\|_*,\] whose minimizer \(\xi^\star\) yields \(\lambda^\star\) via the exponential tilt above. When \(t\notin\Phi_n\) the proximity constraint is active and \(\|\xi^\star\|_*>0\). Practically, one may either (i) solve the projection 6 and then the equality EPA 7 , or (ii) solve the single smooth/nonsmooth dual above; both give the same \(\lambda^\star\).

35.0.0.5 Remarks.

(i) The choice of norm encodes the geometry you trust: Mahalanobis (with target covariance) is often statistically natural; Euclidean is a robust default. (ii) If \(t\in\Phi_n\), then \(\delta_n(t)=0\) and we recover the usual EPA. (iii) Geometrically, \(s^\star\) lies on a face of \(\Phi_n\) exposed by some \(\xi^\star\); the EPA weights are the maximum-entropy distribution realizing that face-average.

References↩︎

---

1. A quick proof: for any random variable \(W\in[a,b]\), \(\mathrm{Var}(W)=\min_{t\in\mathbb{R}}\mathbb{E}[(W-t)^2]\le \mathbb{E}[(W-\tfrac{a+b}{2})^2]\le (b-a)^2/4\).↩︎