Add Attract-Repel link prediction model #10105

tommyly201 · 2025-03-09T03:34:08Z

Add Attract-Repel Embeddings for Link Prediction

This PR adds an implementation of Attract-Repel embeddings from the paper "Pseudo-Euclidean Attract-Repel Embeddings for Undirected Graphs" by Peysakhovich et al.

Overview

Attract-Repel embeddings address a fundamental limitation in traditional graph embeddings: their inability to effectively represent non-transitive relationships. The method splits node representations into two components:

Attract component: nodes with similar attract vectors are likely to connect
Repel component: nodes with similar repel vectors are unlikely to connect

Link prediction scores are calculated as: attract_dot_product - repel_dot_product

Implementation

Added ARLinkPredictor class that implements the core functionality
Added thorough tests to validate the implementation
The implementation is minimal and focused on the core concept

Benefits

Improves link prediction performance on graphs with non-transitive relationships
Requires minimal architectural changes to existing models
Provides useful information about graph structure through R-fraction analysis
Example PR #10107

for more information, see https://pre-commit.ci

puririshi98

LGTM just get CI green (the precommit needs fix from you, the other one isnt related to you so hopefully it just goes away when you fix the precommit one)

for more information, see https://pre-commit.ci

tommyly201 · 2025-04-01T05:25:24Z

Thank you, Rishi

It passed the precommit now! looking forward to seeing this in the repo 😄 . so excited for this!

puririshi98 · 2025-04-01T18:37:48Z

need fix for https://github.com/pyg-team/pytorch_geometric/actions/runs/14187814330/job/39794651601?pr=10105:
Warning, treated as error:
/home/runner/work/pytorch_geometric/pytorch_geometric/torch_geometric/nn/models/attract_repel.py:docstring of torch_geometric.nn.models.attract_repel.ARLinkPredictor:16:Field list ends without a blank line; unexpected unindent.

…/tommyly201/pytorch_geometric into feature/attract-repel-embeddings

for more information, see https://pre-commit.ci

…/tommyly201/pytorch_geometric into feature/attract-repel-embeddings

for more information, see https://pre-commit.ci

puririshi98

LGTM, the CI failure remaining does not seem related, hopefully gets fixed soon

puririshi98 · 2025-04-01T21:27:34Z

@tommyly201 just keep updating the branch as updates come to master and it should merge itself once its green, then ill merge the example PR

…epel-embeddings

…/tommyly201/pytorch_geometric into feature/attract-repel-embeddings

…epel-embeddings

…/tommyly201/pytorch_geometric into feature/attract-repel-embeddings

# Attract-Repel Link Prediction Example This PR adds an example implementation showcasing the attract-repel embeddings (from PR #10105) on link prediction tasks. The example demonstrates how attract-repel embeddings can significantly improve performance on citation networks. ## Overview The example: - Compares traditional vs attract-repel embeddings on link prediction - Works with standard Planetoid datasets (Cora, CiteSeer, PubMed) - Shows performance improvements across different networks - Calculates and reports the R-fraction for analysis ## Results | Dataset | Method | Test AUC | |---------|--------|----------| | Cora | Traditional | 0.6624 | | Cora | Attract-Repel | 0.8945 | | PubMed | Traditional | 0.8977 | | PubMed | Attract-Repel | 0.9607 | | CitiSeer | Traditional | 0.8067 | | CitiSeer | Attract-Repel | 0.8206 | The attract-repel approach shows consistent improvements across datasets: - **Cora**: +0.2321 AUC (R-fraction: 0.4075) - **PubMed**: +0.0630 AUC (R-fraction: 0.5062) - **CitiSeer**: +0.0139 AUC (R-fraction: 0.4869) The varying R-fractions indicate that the model adapts to different graph structures, with PubMed using more of the repel component than Cora. ## Usage Traditional approach: python examples/ar_link_pred.py --dataset Cora Attract-Repel approach python examples/ar_link_pred.py --dataset Cora --use_ar This example complements the attract-repel embeddings implementation in PR #10105. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Rishi Puri <[email protected]> Co-authored-by: Rishi Puri <[email protected]>

# Add Attract-Repel Embeddings for Link Prediction This PR adds an implementation of Attract-Repel embeddings from the paper "Pseudo-Euclidean Attract-Repel Embeddings for Undirected Graphs" by Peysakhovich et al. ## Overview Attract-Repel embeddings address a fundamental limitation in traditional graph embeddings: their inability to effectively represent non-transitive relationships. The method splits node representations into two components: - Attract component: nodes with similar attract vectors are likely to connect - Repel component: nodes with similar repel vectors are unlikely to connect Link prediction scores are calculated as: attract_dot_product - repel_dot_product ## Implementation - Added `ARLinkPredictor` class that implements the core functionality - Added thorough tests to validate the implementation - The implementation is minimal and focused on the core concept ## Benefits - Improves link prediction performance on graphs with non-transitive relationships - Requires minimal architectural changes to existing models - Provides useful information about graph structure through R-fraction analysis - Example PR [pyg-team#10107 ](pyg-team#10107) --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Rishi Puri <[email protected]> Co-authored-by: Rishi Puri <[email protected]>

# Attract-Repel Link Prediction Example This PR adds an example implementation showcasing the attract-repel embeddings (from PR pyg-team#10105) on link prediction tasks. The example demonstrates how attract-repel embeddings can significantly improve performance on citation networks. ## Overview The example: - Compares traditional vs attract-repel embeddings on link prediction - Works with standard Planetoid datasets (Cora, CiteSeer, PubMed) - Shows performance improvements across different networks - Calculates and reports the R-fraction for analysis ## Results | Dataset | Method | Test AUC | |---------|--------|----------| | Cora | Traditional | 0.6624 | | Cora | Attract-Repel | 0.8945 | | PubMed | Traditional | 0.8977 | | PubMed | Attract-Repel | 0.9607 | | CitiSeer | Traditional | 0.8067 | | CitiSeer | Attract-Repel | 0.8206 | The attract-repel approach shows consistent improvements across datasets: - **Cora**: +0.2321 AUC (R-fraction: 0.4075) - **PubMed**: +0.0630 AUC (R-fraction: 0.5062) - **CitiSeer**: +0.0139 AUC (R-fraction: 0.4869) The varying R-fractions indicate that the model adapts to different graph structures, with PubMed using more of the repel component than Cora. ## Usage Traditional approach: python examples/ar_link_pred.py --dataset Cora Attract-Repel approach python examples/ar_link_pred.py --dataset Cora --use_ar This example complements the attract-repel embeddings implementation in PR pyg-team#10105. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Rishi Puri <[email protected]> Co-authored-by: Rishi Puri <[email protected]>

Add Attract-Repel link prediction model

bb306d0

tommyly201 requested review from EdisonLeeeee and wsad1 as code owners March 9, 2025 03:34

pre-commit-ci bot and others added 4 commits March 9, 2025 03:35

[pre-commit.ci] auto fixes from pre-commit.com hooks

0d4bca8

for more information, see https://pre-commit.ci

Update CHANGELOG.md for ARLinkPredictor

0ce1e05

Update CHANGELOG.md for ARLinkPredictor

e9474b4

[pre-commit.ci] auto fixes from pre-commit.com hooks

4ec3497

for more information, see https://pre-commit.ci

This was referenced Mar 9, 2025

[Feature Proposal] Attract-Repel Embeddings for PyTorch Geometric #10084

Open

Add attract-repel link prediction example #10107

Merged

Feature/attract repel wrapper #10117

Closed

Attract Repel wrapper around existing GNN models (GCN, GAT, GraphSage) #10120

Open

puririshi98 and others added 2 commits March 20, 2025 13:55

Merge branch 'master' into feature/attract-repel-embeddings

a41cbe0

Merge branch 'master' into feature/attract-repel-embeddings

523a951

puririshi98 requested changes Apr 1, 2025

View reviewed changes

tommyly201 and others added 5 commits April 1, 2025 01:02

Fix CI issues: flake8 line length limits and missing docstring params

81b41c9

[pre-commit.ci] auto fixes from pre-commit.com hooks

bc2f033

for more information, see https://pre-commit.ci

Fix CI issues: flake8 line length limits

fa9e9df

fix CI and flak8

5666dbc

[pre-commit.ci] auto fixes from pre-commit.com hooks

258a45d

for more information, see https://pre-commit.ci

tommyly201 and others added 7 commits April 1, 2025 15:18

fix docstring for AR Link predictor

3378758

Merge branch 'feature/attract-repel-embeddings' of https://github.com…

a67c350

…/tommyly201/pytorch_geometric into feature/attract-repel-embeddings

[pre-commit.ci] auto fixes from pre-commit.com hooks

d73c9a8

for more information, see https://pre-commit.ci

docstring for encode method

7f605e3

Merge branch 'feature/attract-repel-embeddings' of https://github.com…

c6fb7a8

…/tommyly201/pytorch_geometric into feature/attract-repel-embeddings

[pre-commit.ci] auto fixes from pre-commit.com hooks

9d4afae

for more information, see https://pre-commit.ci

fix indenttation for encode method

1e7534f

puririshi98 approved these changes Apr 1, 2025

View reviewed changes

puririshi98 enabled auto-merge (squash) April 1, 2025 21:27

puririshi98 and others added 6 commits April 1, 2025 18:09

Merge branch 'master' into feature/attract-repel-embeddings

7e2a531

Merge remote-tracking branch 'upstream/master' into feature/attract-r…

f10ab96

…epel-embeddings

Merge branch 'feature/attract-repel-embeddings' of https://github.com…

dada105

…/tommyly201/pytorch_geometric into feature/attract-repel-embeddings

no change, just trigger new CI

60c35f7

Merge remote-tracking branch 'upstream/master' into feature/attract-r…

144809a

…epel-embeddings

Merge branch 'feature/attract-repel-embeddings' of https://github.com…

eb160a8

…/tommyly201/pytorch_geometric into feature/attract-repel-embeddings

auto-merge was automatically disabled April 3, 2025 22:04
Invalid email address

puririshi98 merged commit 81d562e into pyg-team:master Apr 3, 2025
16 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add Attract-Repel link prediction model #10105

Add Attract-Repel link prediction model #10105

Uh oh!

tommyly201 commented Mar 9, 2025 •

edited

Loading

Uh oh!

puririshi98 left a comment

Uh oh!

tommyly201 commented Apr 1, 2025

Uh oh!

puririshi98 commented Apr 1, 2025

Uh oh!

puririshi98 left a comment

Uh oh!

puririshi98 commented Apr 1, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Add Attract-Repel link prediction model #10105

Add Attract-Repel link prediction model #10105

Uh oh!

Conversation

tommyly201 commented Mar 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Add Attract-Repel Embeddings for Link Prediction

Overview

Implementation

Benefits

Uh oh!

puririshi98 left a comment

Choose a reason for hiding this comment

Uh oh!

tommyly201 commented Apr 1, 2025

Uh oh!

puririshi98 commented Apr 1, 2025

Uh oh!

puririshi98 left a comment

Choose a reason for hiding this comment

Uh oh!

puririshi98 commented Apr 1, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tommyly201 commented Mar 9, 2025 •

edited

Loading