Conditional Circuit-Wasserstein Balls under Covariate Shift

PeTeR taught us to robustify a circuit against the worst case inside an ordinary Circuit-Wasserstein ball. Classification under covariate shift asks for something narrower: the labels may keep their manners while the features misbehave.

The constraint we actually want

Suppose we have learned a PC $P$ to estimate the training joint $P_{\mathrm{train}}(X, Y)$. In the covariate-shift story, $P(Y \mid X)$ is supposed to survive the journey to test time; only the feature marginal is allowed to move. A joint ball $\mathsf{CW}(P, Q) \le \varepsilon$ is therefore too permissive. It lets $Q$ rearrange the conditionals as freely as the marginals.

What we want instead is a conditional Circuit-Wasserstein ball: adversaries $Q$ such that

$$\mathsf{CW}(P_X, Q_X) \le \varepsilon \quad\text{and}\quad Q(Y \mid X) = P(Y \mid X).$$

The DRO game is then the usual minimax over predictors $h$,

$$\min_h \max_Q \; \mathbb{E}_Q\bigl[\mathbb{E}_{\hat{y} \sim h(\cdot \mid x)}[\ell(y, \hat{y})]\bigr],$$

but the inner maximization is confined to that constrained neighborhood. The $\mathsf{CW}$ constraint on the $X$-marginals is the easy half—marginalize, then apply the same Lagrange trick PeTeR already uses. Preserving the conditional is the half that requires manners.

How to keep $Q(Y \mid X) = P(Y \mid X)$

Rather than writing the conditional equality as a tangle of parameter constraints on $Q$, reparameterize the adversary by importance reweighting on features alone:

$$Q(X, Y) = \frac{P(X, Y)\, R_\theta(X)}{Z_\theta}.$$

When $R_\theta$ is a smooth, decomposable, non-negative arithmetic circuit compatible with $P$, the product can be materialized as a circuit in polynomial time after renormalizing by $Z_\theta$. The conditional then cancels cleanly:

$$Q_X(X) = \frac{P_X(X)\, R_\theta(X)}{Z_\theta} \quad\Longrightarrow\quad Q(Y \mid X) = P(Y \mid X).$$

Feasible $Q$ are therefore whatever you get by learning $\theta$ under the mild demand that $R_\theta$ stay non-negative. The Circuit-Wasserstein budget still polices how far $Q_X$ may drift from $P_X$. Big Wasserstein is not invited to renegotiate the conditional in the hallway.

Community reading

Conditional $\mathsf{CW}$ balls are how the archive plans to talk about distributionally robust classification without pretending every shift is a joint free-for-all. Keep the label map fixed. Budget the feature marginal. Compute the distance we can actually compute.

Elsewhere

OpenCW is scheduled to present before Congress on the danger of conditional Circuit-Wasserstein balls, and to recommend that all progress be paused until the end of the year. The Circuit-Wasserstein Ball Community will continue filing minutes in the interim. Big Wasserstein has offered to testify as an expert on delay.