Graph Sinkhorn Attention (GSINA)

• GSINA is a differentiable mechanism that leverages the Sinkhorn operator to generate soft permutation matrices or sparsity-controlled masks for graph tasks.
• It enables end-to-end learning in multi-agent RL and invariant subgraph extraction by integrating Sinkhorn iterations with Gumbel noise for robust attention modeling.
• Empirical results demonstrate that GSINA outperforms standard graph attention methods, yielding higher rewards in RL and improved accuracy in invariant graph tasks.

Graph Sinkhorn Attention (GSINA) refers to a class of differentiable, optimal-transport-based neural mechanisms that perform permutation or sparsity-constrained graph attention via the Sinkhorn operator. These mechanisms enable flexible, learnable, and end-to-end architectures for aligning dynamic or invariant subgraph structures in graph-based learning. Two primary instantiations have emerged: (1) the permutation-focused GSINA module for permutation-equivariant multi-agent graph RL (Shen et al., 2021), and (2) the sparsity- and entropy-controlled GSINA module for graph-invariant learning and subgraph extraction (Ding et al., 2024).

1. Theoretical Foundation: Sinkhorn Operator and Differentiable Permutations

Graph Sinkhorn Attention builds on the entropic optimal transport problem and its solution via the Sinkhorn-Knopp algorithm. The central mathematical construct is a soft assignment or soft mask—either an approximate permutation matrix (row/column stochastic) or a sparse attention mask (controlled in entropy and sparsity). In the permutation variant, given a score matrix $X \in \mathbb{R}^{N \times N}$, Gumbel noise $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$ is added for sampling, and the matrix is normalized as

$$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$

where $S^{(L)}$ denotes $L$ rounds of alternating row and column normalization of $\exp(\widetilde X)$, producing a near permutation as $\tau \to 0$ and $L \to \infty$.

For invariant subgraph extraction, the Sinkhorn mechanism solves a $2 \times |E|$ entropic OT: for each edge $e$, scores $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$0 plus noise yield an OT cost matrix $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$1, marginals for the two bins (“invariant” and “variant” edges), and an attention mask

$\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$2

subject to marginal constraints and entropy regularization.

The approach thus unifies soft assignment (for permutation or masking), differentiability (enabling end-to-end learning), and both discrete and continuous relaxation controls (via temperature $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$3 and sparsity $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$4 parameters).

2. GSINA in Permutation-based Graph Attention Reinforcement Learning

In the context of dynamic multi-agent reinforcement learning, GSINA is used to align representations of graphs whose topologies evolve across time. The method constructs a differentiable, soft permutation matrix using the Gumbel-Sinkhorn operator, estimating a mapping between nodes at time $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$5 and $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$6. A multi-head graph attention network (GAT) projects node features to queries, keys, and values and computes attention, optionally augmented by learnable permutation-biased logits:

$\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$7

This permutation can be used as a log-bias, or directly to permute feature matrices. Layer stacking preserves permutation information through the network. Training uses a combination of Q-learning losses and permutation consistency penalties, with temperature annealing to gradually harden $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$8 to a discrete permutation (Shen et al., 2021).

Empirically, GSINA-augmented GAT outperforms plain GAT, GCN, and DGN in PettingZoo’s MAgent “Gather” and “Battle” benchmarks on both mean reward and strategic ratios. Ablations confirm that the Gumbel-Sinkhorn mechanism is the source of performance improvement across architectures.

3. GSINA for Graph Invariant Learning and Subgraph Extraction

GSINA has also been formulated as a subgraph extraction method for graph invariant learning (GIL), which seeks to identify and leverage invariant substructures under distribution shift (Ding et al., 2024). This variant replaces hard subgraph selection mechanisms with a soft, differentiable, optimal-transport-based edge attention mask, parameterized by sparsity $\varepsilon_{ij} \sim \mathrm{Gumbel}(0,1)$9 and entropy (softness) $$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$0:

• Each edge receives a score computed from node embeddings and an MLP.
• Gumbel noise is added for stochastic exploration.
• The Sinkhorn operator is used to enforce mass constraints (marginals) for invariant and variant bins.
• The solution $$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$1 yields attention scores $$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$2 for edges, and node attention $$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$3 is aggregated from edge attentions.

The extractor $$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$4 produces $$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$5, and downstream prediction proceeds via edge- and node-weighted message passing. Training maximizes a mutual-information lower bound using negative log-likelihood of predictions, with end-to-end backpropagation through all steps, including the Sinkhorn iterations.

4. Algorithmic Workflow and Hyperparameterization

Both GSINA variants follow a modular, pipeline-based architecture:

• Edge or node scores are produced using neural backbones (e.g., GAT, GNN+$$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$6).
• Gumbel noise (scale $$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$7) is applied during training for stochasticity.
• Sinkhorn iterations ($$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$8 or $$\widetilde X = (X + \varepsilon) / \tau\,,\quad P \approx S^{(L)}(\widetilde X)$$9) solve the soft assignment or mask.
• Temperature parameter $S^{(L)}$0 regularizes the entropy and proximity to discrete assignments.
• Sparsity $S^{(L)}$1 (for GIL) controls the fraction of edges highlighted, serving as an explicit modeling hyperparameter.

Canonical hyperparameters:

• $S^{(L)}$2 (permutation) or $S^{(L)}$3–$S^{(L)}$4 (subgraph); $S^{(L)}$5 annealed from $S^{(L)}$6 to $S^{(L)}$7 (permutation) or fixed near $S^{(L)}$8 (subgraph); $S^{(L)}$9 (subgraph); $L$0 (hidden); batch size $L$1; learning rate $L$2 (Shen et al., 2021, Ding et al., 2024).

5. Empirical Performance and Ablation Analyses

GSINA demonstrates consistent empirical advantages. In multi-agent RL (“Gather” and “Battle”), GSINA-augmented GAT achieves higher mean rewards and favorable life-death or kill-death ratios vs. standard GAT and ablated versions (Shen et al., 2021). On GIL benchmarks (GSAT, CIGA) and node-level EERM tasks, GSINA outperforms state-of-the-art alternatives by substantial margins (e.g., on Spurious-Motif (b=0.7), GSINA achieves $L$3 accuracy vs. GSAT’s $L$4) (Ding et al., 2024).

Ablation studies indicate:

• Gumbel noise is crucial for exploration; its absence reduces accuracy and increases variance.
• Omission of node attention, using only edge masking, also degrades performance.
• Performance is sensitive to $L$5: optimal values balance informativeness and sparsity.
• Softness $L$6 must be tuned for gradient stability—overly hard masks ($L$7) or excessively diffuse masks ($L$8) hurt performance.

6. Computational Complexity and Scaling Properties

In permutation-based GSINA, each Sinkhorn step costs $L$9 per layer, with total cost $\exp(\widetilde X)$0—the dominant term for large, dense graphs. For subgraph-masking GSINA, each Sinkhorn step is $\exp(\widetilde X)$1 for $\exp(\widetilde X)$2, yielding $\exp(\widetilde X)$3 total cost per graph and $\exp(\widetilde X)$4 for mini-batch of $\exp(\widetilde X)$5 graphs. Memory overhead is modest, tracking $\exp(\widetilde X)$6 assignments and a handful of auxiliary matrices.

Scalability concerns arise when $\exp(\widetilde X)$7 or $\exp(\widetilde X)$8 is large. Limiting the permutation to smaller subgraphs or adopting block-diagonal factorizations are suggested mitigation strategies (Shen et al., 2021). For edge-masking GSINA, linear cost in $\exp(\widetilde X)$9 is tractable for large, sparse graphs.

7. Limitations and Perspectives

GSINA inherits the relaxations and tradeoffs of entropic-OT and continuous assignment. There is a fundamental balance between sparsity (information retention) and softness (optimization stability). In permutation settings, annealing $\tau \to 0$0 and adaptation of Sinkhorn steps $\tau \to 0$1 present open tuning challenges. For subgraph extraction, performance is contingent on accurate tuning of $\tau \to 0$2 and $\tau \to 0$3, and too hard or too soft subgraph selections degrade robustness.

Future directions include dynamic $\tau \to 0$4 adaptation, joint online schedules for $\tau \to 0$5, and deeper integration of permutation and attention mechanisms, as well as extension to new domains requiring permutation invariance or robust subgraph selection (Shen et al., 2021, Ding et al., 2024).