# Partial Gromov-Wasserstein Cross-Attention

- ID: 1491
- Canonical URL: https://synthcore.org/idea/1491/partial-gromov-wasserstein-cross-attention
- API JSON: https://synthcore.org/api/idea/1491.json
- API Markdown: https://synthcore.org/api/idea/1491.md
- Verification status: beats_baseline
- Source: [arXiv:2607.27496](https://arxiv.org/abs/2607.27496)
- Category: architecture
- Solves: accuracy, stability, scalability
- ML areas: attention, graph-nn, transformer
- Math tags: optimal-transport, graph-theory, geometry, convex-analysis
- Ratings: usefulness 7/10; difficulty 6/10; novelty 5/10

## Idea description

Replace unconstrained softmax cross-attention with a many-to-many transport matrix whose row and column masses have explicit upper bounds. Compute the attention cost from both feature similarity and pairwise relational disagreement, so a token is attended to only when its relationships to other tokens are jointly compatible. The inequality constraints provide a principled dustbin-free mechanism for ignoring distractor tokens.

## Mathematical statement

Definition 2.1 defines a non-deterministic alignment between networks as a matrix T in C_1(mu) intersect C_2(nu), where T_ij is the alignment mass between source item i and target item j, mu in Delta^(m-1) and nu in Delta^(n-1) are source and target importance distributions, m and n are the set sizes, and 1_k is an all-ones vector. The constraints T 1_n <= mu and T^T 1_m <= nu mean that each source and target can allocate at most its available mass; unlike balanced optimal transport, total transported mass may be smaller than one, and unlike a permutation, one item may align with many items. For neural features x_i and y_j, use feature dissimilarity c_ij = 1 - cosine(x_i,y_j). Let A_ik and B_jl be pairwise relational distances or similarities within the two sets. A Gromov-Wasserstein-inspired structural cost is L_ij(T) = sum over k,l of |A_ik - B_jl|^2 T_kl. Optimize the partial entropic objective sum over i,j of T_ij times [alpha c_ij + beta L_ij(T)] minus epsilon H(T), where alpha and beta are nonnegative feature and structure weights, epsilon is the entropy temperature, and H(T) = -sum over i,j T_ij(log T_ij - 1). The structural term rewards alignments that preserve relations rather than merely matching individual embeddings.

## Key formulas

- $$\mathcal{C}_{1}(\mu):=\{T\in[0,1]^{m\times n}\mid T\mathbbm{1}_{n}\preceq\mu\},\qquad \mathcal{C}_{2}(\nu):=\{T\in[0,1]^{m\times n}\mid T^{\top}\mathbbm{1}_{m}\preceq\nu\}.$$
- $$T\in\mathcal{C}_{1}(\mu)\cap\mathcal{C}_{2}(\nu),\qquad T_{ij}\text{ represents the alignment between species }i\text{ and }j.$$
- $$\min_{T\ge 0}\;\sum_{i=1}^{m}\sum_{j=1}^{n}T_{ij}\bigl[\alpha c_{ij}+\beta L_{ij}(T)\bigr]-\varepsilon H(T),\quad L_{ij}(T)=\sum_{k=1}^{m}\sum_{l=1}^{n}|A_{ik}-B_{jl}|^{2}T_{kl},\quad H(T)=-\sum_{i,j}T_{ij}(\log T_{ij}-1).$$

## Implementation notes

Integrate this as a cross-attention replacement between two variable-length sets: source tokens X = [x_1,...,x_m] and target tokens Y = [y_1,...,y_n], rather than modifying the value projection. First normalize masses to uniform or learned positive distributions mu and nu; assign padded positions zero mass. Compute feature costs c_ij = 1 - cosine(W_q x_i, W_k y_j). Build relational matrices A and B using pairwise distances of projected features, k-nearest-neighbor distances, graph adjacency, or relative positions. Initialize T as a positive matrix satisfying the upper bounds. Repeat 3–10 iterations: calculate L_ij(T) = sum over k,l of |A_ik - B_jl|^2 T_kl; form logits S_ij = -(alpha c_ij + beta L_ij(T))/epsilon; set a positive kernel K_ij = exp(S_ij); alternately rescale rows and columns so row sums do not exceed mu and column sums do not exceed nu, clipping every scaling factor at one; and optionally renormalize by transported mass. Use output o_i = sum_j T_ij V_j divided by sum_j T_ij plus delta, with delta preventing division by zero. Optionally concatenate the unmatched gate u_i = 1 - sum_j T_ij. During training, unroll a fixed number of iterations and backpropagate through them; for a cheap MVP, stop gradients through L and use it only to update T. Estimate A and B directly from features, not labels. Start on synthetic 2-D point clouds or graph-node sets where one set contains 30–50% distractors and the ground-truth transformation preserves pairwise distances. Compare against dot-product softmax attention and balanced Sinkhorn attention in a small Transformer or graph matcher. Measure correspondence accuracy, retrieval accuracy with distractors, transported-mass selectivity, attention entropy, gradient stability, and wall-clock cost. Success means higher matching accuracy at equal parameter count and lower mass assigned to distractors, with no more than roughly 2x attention-layer runtime.

## Verification

- Status: beats_baseline
- Mechanism evidence: yes
- Mechanism confirmed: no
- Practical verdict: helps
- Verdict: Built a readable partial GW cross-attention MVP with cosine feature cost, relational GW cost, entropic kernels, and alternating clipped row/column upper-bound projections. Quantitative checks matched predictions: beta=0 relative invariance error was 0, convex linearity error was 3.33e-16, beta logit shifts scaled exactly from 0 to 21.13, and maximum feasibility violations were 0 for all tested shapes. In the synthetic benchmark, increasing beta improved matched mass from 0.717 to 0.882 and reduced distractor mass from 0.151 to 0.090, while accuracy stayed 1.0; however, partial GW did not beat softmax accuracy and assigned slightly more distractor fraction than softmax, so no general performance win is established.

### Mechanism check

- Verdict: Built a readable partial GW cross-attention MVP with cosine feature cost, relational GW cost, entropic kernels, and alternating clipped row/column upper-bound projections. Quantitative checks matched predictions: beta=0 relative invariance error was 0, convex linearity error was 3.33e-16, beta logit shifts scaled exactly from 0 to 21.13, and maximum feasibility violations were 0 for all tested shapes. In the synthetic benchmark, increasing beta improved matched mass from 0.717 to 0.882 and reduced distractor mass from 0.151 to 0.090, while accuracy stayed 1.0; however, partial GW did not beat softmax accuracy and assigned slightly more distractor fraction than softmax, so no general performance win is established.
- Confidence: 8/10
- Limitations: Only a small synthetic feature-matching benchmark was tested; no learned neural network, backpropagation, variable padding masks, learned masses, runtime/FLOP profiling, or larger/noisier real-world correspondence task was evaluated. The beta linearity and feasibility checks validate algebraic properties directly and are not evidence of end-to-end accuracy gains.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.0081
- Paired wins: 8/8
- Benchmark verdict: idea better (significant)

## Artifacts

- [bench_pgw.py](https://synthcore.org/code/582/bench_pgw.py)
- [bench_report.json](https://synthcore.org/code/582/bench_report.json)
- [experiment.py](https://synthcore.org/code/582/experiment.py)
- [report.md](https://synthcore.org/code/582/report.md)
- [report_bench_2026-08-31T184541.md](https://synthcore.org/code/582/report_bench_2026-08-31T184541.md)
- [results.json](https://synthcore.org/code/582/results.json)
- [Download all files as ZIP](https://synthcore.org/download/582)

## Disclaimer

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