# Causal DAG Wasserstein Alignment

- ID: 2257
- Canonical URL: https://synthcore.org/idea/2257/causal-dag-wasserstein-alignment
- API JSON: https://synthcore.org/api/idea/2257.json
- API Markdown: https://synthcore.org/api/idea/2257.md
- Verification status: unverified
- Source: [arXiv:2608.13716](https://arxiv.org/abs/2608.13716)
- Category: regularization
- Solves: accuracy, stability, generalization
- ML areas: graph-nn, regularization, attention
- Math tags: optimal-transport, probability, graph-theory, measure-theory, optimization
- Ratings: usefulness 6/10; difficulty 5/10; novelty 7/10

## Idea description

Replace an unconstrained Wasserstein representation-matching loss with a graph-causal transport loss whose coupling at node k is conditioned only on the representations of its parents. This forces domain alignment, distillation, or augmentation consistency to respect the information flow of the model's DAG, reducing spurious matches that exploit descendants or globally visible features.

## Mathematical statement

A DAG G=(V,E) has parent set pa_k for node k. A G-compatible distribution factorizes through parent-conditioned kernels, with source and target conditional laws denoted by mu_k^{x_pa_k} and nu_k^{y_pa_k}. The paper's admissible local coupling set is D_k={(x_pa_k,y_pa_k,pi): pi in Pi(mu_k^{x_pa_k},nu_k^{y_pa_k})}, where Pi(a,b) is the set of probability measures pi whose marginals are a and b. For local cost c_k and descendant value functions V_i^{G,c}, the dynamic-programming integrand is f_k(x_pa_k,y_pa_k,x_k,y_k)=c_k(x_k,x_pa_k,y_k,y_pa_k)+sum_{i in A_k} V_i^{G,c}(x_pa_i,y_pa_i), and the local objective is g_k(x_pa_k,y_pa_k,pi)=integral f_k d pi. Here A_k is the set of active successor or descendant nodes contributing after node k. We adapt this by learning a coupling pi_k between source and target node embeddings separately for each parent context, while adding child value terms recursively. The causal restriction is the transferable property: correspondence at node k cannot use information outside pa_k.

## Key formulas

- $$\mathrm{prox}(E)=\{(\mathrm{prox}(i),i)\mid i\in V,\ \mathrm{pa}_i\neq\varnothing\}\subseteq E,$$
- $$D_k=\left\{(x_{\mathrm{pa}_k},y_{\mathrm{pa}_k},\pi)\ \middle|\ \pi\in\Pi\left(\mu_k^{x_{\mathrm{pa}_k}},\nu_k^{y_{\mathrm{pa}_k}}\right)\right\},$$
- $$g_k(x_{\mathrm{pa}_k},y_{\mathrm{pa}_k},\pi)=\int f_k(x_{\mathrm{pa}_k},y_{\mathrm{pa}_k},x_k,y_k)\,d\pi(x_k,y_k),$$
- $$f_k=c_k(x_k,x_{\mathrm{pa}_k},y_k,y_{\mathrm{pa}_k})+\sum_{i\in A_k}V_i^{G,c}(x_{\mathrm{pa}_i},y_{\mathrm{pa}_i}).$$

## Implementation notes

Integrate this at a representation-alignment loss in a DAG-structured graph neural network, multi-stage encoder, or teacher-student model. Let each node k produce source embeddings X_k and target embeddings Y_k, with parent embeddings X_pa_k and Y_pa_k available from the same forward pass. Construct soft parent contexts q_k=MLP_k([X_pa_k,Y_pa_k]); do not provide child or descendant embeddings to this network. For each minibatch and node, compute pairwise costs C_k[a,b]=||X_k[a]-Y_k[b]||_2^2 plus an optional parent-consistency term alpha||X_pa_k[a]-Y_pa_k[b]||_2^2. Use a conditional Sinkhorn coupling: initialize log scores L_k=-C_k/epsilon, run 20 to 50 log-domain Sinkhorn iterations to enforce uniform source and target marginals, and optionally add a context-dependent bias B_k[a,b]=MLP_k(q_k[a],q_k[b]) to L_k. The node loss is <P_k,C_k>, where P_k is the resulting coupling. Approximate the recursive child-value term by adding lambda_child times already-computed losses of active children; evaluate nodes in reverse topological order. For roots, use unconditional Sinkhorn transport. The causal restriction computed exactly is that P_k uses only parent-conditioned features; epsilon, alpha, and lambda_child are engineering hyperparameters. Start with Colored-MNIST or a small molecular graph domain-shift task using a three-layer GNN. Compare against global Wasserstein alignment, CORAL, and no alignment. Measure target accuracy, source-target alignment, seed variance, and leakage of descendant information. Success means higher shifted-domain accuracy at equal training cost, lower variance across seeds, and reduced alignment when parent contexts disagree.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.