Partial Gromov-Wasserstein Cross-Attention

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.