Representation-Aware DRO (READ)

• The paper introduces READ, a framework that refines WDRO by using stable, transferable representations to protect informative feature directions during distribution shifts.
• It leverages an alignment-sensitive transport cost, replacing isotropic penalties with differential protection that yields exact convex reformulations and adaptive regularization.
• READ calibrates ambiguity radii and constructs anisotropic confidence regions, thereby reducing bias and improving out-of-sample performance in supervised estimation.

Representation-Aware Distributionally Robust Estimation (READ) is a Wasserstein distributionally robust learning framework for supervised estimation under feature distribution shifts when a low-dimensional predictive representation is believed to remain stable across environments. Its defining modification to classical Wasserstein distributionally robust optimization (WDRO) is to replace representation-agnostic transport penalties with a transport geometry that differentially discourages perturbations along representation directions associated with informative or transferable structure. In the formulation introduced in “Representation-Aware Distributionally Robust Optimization: A Knowledge Transfer Framework,” READ is developed for linear regression and linear binary classification, with exact convex reformulations, robust Wasserstein profile inference (RWPI) for radius calibration, representation-aware confidence regions, and a data-adaptive procedure for selecting alignment parameters (Wang et al., 11 Sep 2025).

1. Problem setting and conceptual motivation

READ is formulated for covariates $X \in \mathbb{R}^d$, response $Y$, unknown data-generating distribution $\mathbb{P}^*$, and empirical measure $\mathbb{P}_N^*$ built from $\{(x_i,y_i)\}_{i=1}^N$. The population target is

$$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$

Classical WDRO replaces this with

$$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$

where the ambiguity set is a Wasserstein ball

$$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$

and

$$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$

The central critique of standard WDRO is that transport costs such as

$c(x,u)=\|x-u\|_q^k$

treat all feature perturbation directions as equally plausible and equally harmful. READ instead assumes that predictive signal may be concentrated in a representation subspace that should be protected more strongly than unstable or weakly informative directions. The framework encodes this idea through a representation matrix $Y$0, whose columns may come from external source coefficients, pre-trained predictors, or learned shared task representations. In a low-rank latent-factor example,

$Y$1

so $Y$2 is treated as a more stable and informative predictor than raw $Y$3 (Wang et al., 11 Sep 2025).

A recurring interpretation is causal-invariance motivated: shifts may occur primarily in directions orthogonal to the stable representation, while predictive structure remains concentrated in the representation-aligned component. READ operationalizes robustness by preserving $Y$4 more strongly than the orthogonal complement. This suggests a notion of robustness tied not merely to perturbation magnitude, but to perturbation geometry relative to informative structure.

2. Representation-aware transport geometry

The basic READ feature cost is

$Y$5

where $Y$6, $Y$7, and $Y$8. In the paper’s notation, the intended action is projection of $Y$9 onto the representation directions. For regression, the induced transport cost on $\mathbb{P}^*$0 is

$\mathbb{P}^*$1

so only feature perturbations are admissible (Wang et al., 11 Sep 2025).

The alignment parameter $\mathbb{P}^*$2 is a multidimensional nonnegative diagonal matrix, with entries allowed in $\mathbb{P}^*$3. It is not a generic dense matrix in the main formulation. Its role is component-specific protection of representation coordinates. Increasing $\mathbb{P}^*$4 makes perturbations along $\mathbb{P}^*$5 more expensive. If $\mathbb{P}^*$6, that direction receives no extra protection. If $\mathbb{P}^*$7, perturbations changing that coordinate are effectively prohibited, yielding a hard invariance constraint.

This construction induces a continuous family of models between two limiting cases. At $\mathbb{P}^*$8, READ reduces to standard WDRO with no representation guidance. At $\mathbb{P}^*$9, it coincides with the hard-transfer regime identified in the paper with knowledge-guided DRO (KG-DRO). The intermediate regime is the distinctive one: alignment is finite, component-specific, and data-dependent.

Geometrically, READ replaces isotropic adversarial transport with anisotropic transport aligned to representation structure. In the special case $\mathbb{P}_N^*$0, the induced metric simplifies to

$\mathbb{P}_N^*$1

and the corresponding regularizer becomes $\mathbb{P}_N^*$2. The paper interprets this as contraction along trusted representation directions and relative expansion in orthogonal directions (Wang et al., 11 Sep 2025).

3. Exact reformulations and induced regularization

A core theoretical contribution of READ is an exact duality between WDRO with extended-norm transport geometry and seminorm regularization. The starting point is an extended norm $\mathbb{P}_N^*$3. If $\mathbb{P}_N^*$4, then $\mathbb{P}_N^*$5 is a linear subspace, and the dual norm

$\mathbb{P}_N^*$6

is a seminorm whose null space is $\mathbb{P}_N^*$7. Lemma 1 gives

$\mathbb{P}_N^*$8

which is the algebraic basis for the robust reformulations (Wang et al., 11 Sep 2025).

For linear regression with squared loss $\mathbb{P}_N^*$9, if

$\{(x_i,y_i)\}_{i=1}^N$0

then

$\{(x_i,y_i)\}_{i=1}^N$1

where $\{(x_i,y_i)\}_{i=1}^N$2. For the standard choice $\{(x_i,y_i)\}_{i=1}^N$3, this recovers

$\{(x_i,y_i)\}_{i=1}^N$4

For binary classification with logistic or hinge loss, if

$\{(x_i,y_i)\}_{i=1}^N$5

then

$\{(x_i,y_i)\}_{i=1}^N$6

READ is obtained by taking

$\{(x_i,y_i)\}_{i=1}^N$7

Proposition 4 gives the conjugate

$\{(x_i,y_i)\}_{i=1}^N$8

with $\{(x_i,y_i)\}_{i=1}^N$9, $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$0, and $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$1. The READ seminorm regularizer is therefore

$$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$2

The regularizer shrinks $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$3 toward $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$4, while allowing differential shrinkage across representation coordinates. Large $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$5 implies a small penalty on the corresponding $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$6, so the $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$7th representation direction is more freely used in approximating $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$8. At $$\beta^* \coloneqq \arg\min_{\beta \in \mathbb{R}^d} \mathbb{E}_{\mathbb{P}^*}\!\left[\ell(X,Y;\beta)\right].$$9, that coordinate is freely available in the projection; at $$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$0, the convention forces $$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$1, so that source direction is ignored.

The corollaries specialized to READ yield the exact convex reformulations

$$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$2

and, for logistic or hinge loss,

$$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$3

These results place READ within a broader WDRO tradition in which the transport geometry determines the induced regularizer. Earlier work had already shown that learning the ground cost in OT-DRO can produce adaptive regularization, including Mahalanobis-norm penalties, through data-driven geometry selection (Blanchet et al., 2017). READ generalizes that principle from metric learning on observed features to a representation-aware cost built around $$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$4 and $$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$5.

4. Radius calibration, confidence regions, and asymptotic behavior

READ does not treat the Wasserstein radius $$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$6 solely as a generic tuning parameter. For fixed $$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$7, it adopts RWPI to calibrate a radius that gives a target coverage statement for the population optimum. The set of parameters optimal for some distribution in the ambiguity ball is

$$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$8

and the calibrated radius is

$$\beta_{\rm DRO}\in\arg\min_{\beta} \sup_{\mathbb{P}\in \mathcal{P}_{\rm amb}} \mathbb{E}_{\mathbb{P}}\!\left[\ell(X,Y;\beta)\right],$$9

The corresponding robust Wasserstein profile function is

$$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$0

where $$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$1. The key equivalence is

$$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$2

Hence $$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$3 is the upper $$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$4-quantile of $$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$5 (Wang et al., 11 Sep 2025).

Under the paper’s regularity, optimality, and positive-definiteness assumptions, the asymptotic theory yields

$$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$6

where $$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$7 is the $$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$8-quantile of $$\mathcal{B}_\delta (\mathbb{P}^*_N;c) \coloneqq \left\{ \mathbb{P} \in \mathcal{P}(\mathbb{R}^{d+1}) \,\middle|\, \mathcal{D}_c(\mathbb{P},\mathbb{P}^*_N) \leq \delta \right\},$$9, with $$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$0. The resulting asymptotically valid confidence region is

$$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$1

where

$$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$2

The geometry of these confidence regions is representation-dependent. In the $$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$3 linear regression case, the region is an ellipsoid. The paper emphasizes that increasing alignment contracts the region in the trusted representation direction and expands it in orthogonal directions. It also provides a support-function description: $$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$4 which permits polyhedral outer approximations from random support directions.

READ also studies the asymptotic distribution of the estimator itself. Writing the robust empirical objective as

$$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$5

and the empirical sensitivity as

$$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$6

Proposition 5 states that when $$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$7,

$$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$8

uniformly on compact sets. The robustification therefore adds a first-order variation penalty tied to representation-aware feature sensitivity.

If

$$\mathcal{D}_c(\mathbb{P},\mathbb{Q})\coloneqq\inf_{\pi \in \Pi(\mathbb{P},\mathbb{Q})} \mathbb{E}_\pi[c(U,V)].$$9

then Theorem 2 gives

$c(x,u)=\|x-u\|_q^k$0

The limiting law contains the usual stochastic fluctuation plus a deterministic robustification bias term. That decomposition is central to the paper’s later alignment-selection criterion.

5. Alignment selection and computational procedures

For regression, the READ reformulation is convex and can be implemented through conic programming. In the $c(x,u)=\|x-u\|_q^k$1-th order cone formulation, the paper introduces auxiliary variables $c(x,u)=\|x-u\|_q^k$2 with constraints

$c(x,u)=\|x-u\|_q^k$3

This yields a finite-dimensional convex program amenable to standard conic solvers (Wang et al., 11 Sep 2025).

The harder problem is selecting $c(x,u)=\|x-u\|_q^k$4. The paper frames this in terms of a “solution surface”

$c(x,u)=\|x-u\|_q^k$5

For each $c(x,u)=\|x-u\|_q^k$6, RWPI calibration yields an estimator with the same nominal robustness level, but different alignment profiles move the estimator across this surface. The stated goal is to estimate the projection of the global optimum onto this surface under predictive $c(x,u)=\|x-u\|_q^k$7 risk.

Using the asymptotic expansion above, the excess prediction risk is approximated by a term involving

$c(x,u)=\|x-u\|_q^k$8

which leads to the asymptotic $c(x,u)=\|x-u\|_q^k$9-selection target

$Y$00

The proposed algorithm does not optimize over a free matrix. Instead, it first computes an elastic-net projection of $Y$01 onto $Y$02,

$Y$03

then parameterizes the diagonal alignment profile as

$Y$04

componentwise. A grid search over $Y$05 minimizes

$Y$06

and a final cross-validation step selects $Y$07. The returned values are

$Y$08

The paper characterizes this as a practical stability-oriented procedure rather than a globally solved continuous optimization over $Y$09. It does not provide a formal global convergence guarantee for the $Y$10-selection algorithm. It does, however, state that confidence regions remain asymptotically valid after tuning $Y$11 in this way because the RWPI theory is uniform over $Y$12.

A special “perfect alignment” regime is also analyzed. If $Y$13, where $Y$14 is the submatrix corresponding to coordinates with $Y$15, positive-definiteness assumptions may fail and the limiting object becomes a constrained estimator. Under appropriate rates for $Y$16, $Y$17, and $Y$18, the resulting constrained estimator is asymptotically more efficient than the unconstrained one when $Y$19 (Wang et al., 11 Sep 2025).

6. Empirical behavior, antecedents, and limitations

The empirical study compares READ against transfer-learning and DRO baselines: Dist-Based, Angle-Based, LowRep-Based, KG-DRO, and standard DRO. The simulation metric is the percentage reduction in bias norm relative to ERM,

$Y$20

In the low-source, non-sparse regime, READ is reported to improve as the informativeness parameter $Y$21 increases, with gains up to about $Y$22. As $Y$23 increases, all methods improve, but READ’s advantage over the second-best method widens from roughly $Y$24 to $Y$25. In the high-source, sparse setting, READ again can outperform competitors by up to $Y$26. When the number of source directions rises from $Y$27 to $Y$28, KG-DRO drops by about $Y$29, whereas READ is reported not to suffer this degradation. In a one-source ablation with $Y$30, correlation $Y$31, and $Y$32, the average bias norm is minimized at a finite $Y$33, and $Y$34 achieves an $Y$35 reduction in bias relative to $Y$36. Under the paper’s approximately $Y$37-invariant test shift

$Y$38

READ gives the best out-of-sample MSE improvement across 20 test environments (Wang et al., 11 Sep 2025).

These experiments support two recurring claims. First, the best estimator often lies between standard DRO and hard alignment. Second, adaptive finite alignment is particularly important when source information is informative but contaminated. A plausible implication is that the principal modeling choice is not merely the radius of the ambiguity set, but the geometry of admissible perturbations.

READ sits in a specific lineage within Wasserstein DRO. The monograph “Distributionally Robust Learning” develops the broader WDRO machinery on which READ relies: ambiguity sets around the empirical distribution, strong duality, geometry-induced regularization, and RWPI-style calibration. Its recurring lesson is that the ground metric is the modeling object that determines what robustness means, but it does not formulate representation-aware transfer through a multidimensional alignment parameter (Chen et al., 2021). The earlier paper “Data-driven Optimal Cost Selection for Distributionally Robust Optimization” is a more direct precursor to READ in that it learns OT geometry from data via metric learning, replacing isotropic transport with a Mahalanobis cost and thereby inducing adaptive regularization. That framework is representation-aware in a modular sense, but it does not jointly formulate representation-transfer geometry, RWPI-calibrated confidence regions, or the solution-surface analysis developed in READ (Blanchet et al., 2017).

Several misconceptions are therefore excluded by the formulation itself. READ is not standard WDRO with a differently tuned radius; the decisive modification is the alignment-weighted transport cost. It is not equivalent to hard transfer unless some entries of $Y$39 are sent to $Y$40. It is not an end-to-end deep latent-space DRO framework in its current form; the development is linear, with a linear representation matrix $Y$41, linear regression, and linear binary classification. The paper also states clear limitations: the asymptotic theory requires differentiability and local regularity assumptions; hard alignment can be harmful under severe misspecification of $Y$42; the $Y$43-selection procedure is heuristic in its parameterization and grid search; and radius calibration requires solving many variational programs in Monte Carlo, even though the variational representation substantially reduces the burden relative to nested conic formulations (Wang et al., 11 Sep 2025).

In that sense, READ is best understood as a transfer-aware, representation-parameterized refinement of WDRO. Its defining contribution is to formalize how prior or learned representation structure can be embedded directly into the transport cost, producing exact seminorm regularizers, calibrated ambiguity radii, anisotropic confidence regions, and a principled bias-geometry tradeoff under distribution shift.