Exact First-Passage Structural Encodings for Graph Transformers

Reference implementation and reproducibility package for the two exact first-passage structural encodings (FPSE) — FPSE-P (first-passage probability) and FPSE-T (first-passage time) — evaluated as drop-in relative structural encodings inside the GRIT graph transformer on ZINC and Peptides-struct.

The code is a fork of GRIT (Ma et al., 2023; see Attribution below). All FPSE additions live in grit/transform/rrwp.py and are selected through a single config knob, posenc_RRWP.se_type, so that none / RRWP / FPSE-P / FPSE-T run through an identical model, optimizer, encoder, and attention-bias pipeline. Only the relative structural signal changes between arms.

---

1. What FPSE is

GRIT injects graph structure through Relative Random-Walk Probabilities (RRWP): for the row-stochastic transition matrix $P = D^{-1}A$, the relative encoding of the ordered pair $(i,j)$ is the length-$K$ stack

$$\mathrm{RRWP}_k(i,j) = (P^k)_{ij}, \qquad k = 0,1,\dots,K-1 ,$$

i.e. the probability that a $k$-step walk started at $i$ occupies $j$.

FPSE replaces this occupation signal with an exact first-passage signal, computed in closed form by a short matrix recurrence (no sampling). Let $\tau_j = \min{t \ge 1 : X_t = j}$ be the first-passage step of the walk to node $j$.

FPSE-P — first-passage probability CDF. Off the diagonal the encoding stores the probability that the walk from $i$ first reaches $j$ within $k$ steps,

$$\mathrm{FPSE\text{-}P}_k(i,j) = \Pr{}_i[\tau_j \le k],$$

produced by the iterate $\psi_p^{(k)} = T_1!\left(P,\psi_p^{(k-1)}\right)$ with $\psi_p^{(0)} = I$, where $T_1$ clamps the diagonal to $1$ (an absorbed walk stays absorbed). The diagonal carries the first-return CDF and is routed to the absolute node channel.

FPSE-T — truncated first-passage time. The encoding stores the truncated first moment of the first-passage step,

$$\mathrm{FPSE\text{-}T}_k(i,j) = \mathbb{E}_i!\left[\tau_j \cdot \mathbf{1}{\tau_j \le k}\right],$$

produced by a two-iterate recurrence that advances $\psi_p$ and $\psi_t$ in lockstep ($T_0$ clamps the return-time diagonal to $0$). Because $\psi_t$ is an unbounded truncated moment, the time channels are passed through a fixed log1p + per-graph per-channel z-score normaliser.

none — parameter-matched control. Identical encoder / attention-bias machinery and tensor shape, but a constant, structurally uninformative relative signal. This isolates the information content of the structural encoding from the extra parameters its encoders introduce.

All arms share $P = D^{-1}A$, the same walk length $K$ (posenc_RRWP.ksteps), add_identity: True, and the locked off-diagonal-only relative policy (mask_rel_diag: True), so the diagonal handling is never a confound between encodings.

---

2. Repository layout

.
├── README.md                       # this file
├── LICENSE                         # MIT (upstream GRIT, © 2023 Liheng Ma)
├── NOTICE                          # attribution + list of FPSE modifications
├── environment.yml                 # conda environment (exact tested versions)
├── requirements.txt                # pip fallback
├── main.py                         # GRIT / GraphGym entry point
├── setup.py
├── grit/                           # GRIT framework
│   ├── transform/rrwp.py           # ← FPSE-P / FPSE-T implementation
│   └── config/posenc_config.py     # ← se_type / mask_rel_diag config knobs
├── configs/GRIT/                   # the 8 paper configs (4 arms × 2 datasets)
│   ├── zinc-GRIT-{none,rrwp,fphit_p,fphit_t}.yaml
│   └── peptides-struct-GRIT-{none,rrwp,fphit_p,fphit_t}.yaml
├── scripts/
│   ├── run_zinc_4seed.sh           # ZINC 4-arm × 4-seed driver
│   └── run_peptides_4seed.sh       # Peptides-struct 4-arm × 4-seed driver
├── results/
│   ├── four_seed_summary.md        # verified paper-protocol result table
│   └── four_seed_results.json      # per-seed raw values + provenance
├── tests/  unittests/              # upstream GRIT tests
└── docs/UPSTREAM_GRIT_README.md    # original GRIT README

The config name suffix maps to posenc_RRWP.se_type as: none → none, rrwp → rrwp, fphit_p → fphit (FPSE-P), fphit_t → fphit_t (FPSE-T).

---

3. Installation

Tested environment (the exact versions used to produce the reported numbers):

• Python 3.10, CUDA 11.8, one NVIDIA GPU
• PyTorch 2.1.0, PyTorch Geometric 2.2.0 (+ torch-scatter, torch-sparse)
• ogb, numpy, scipy, scikit-learn, yacs, networkx, einops

conda env create -f environment.yml
conda activate fpse

The torch-scatter / torch-sparse wheels must match your Torch/CUDA build; if conda env create cannot resolve them on your platform, install PyTorch + PyG first and then pip install -r requirements.txt. See environment.yml for the exact pinned versions.

4. Datasets

Both datasets download automatically on first run into --dataset.dir (default datasets/); nothing needs to be fetched by hand.

• ZINC (subset, 12k molecular regression) via torch_geometric.datasets.ZINC.
• Peptides-struct (LRGB, 11-D structural regression) via the OGB / PyG LRGB loader (dataset.format: OGB, name: peptides-structural).

---

5. Reproducing the paper (four-seed protocol)

Every reported number is the validation-selected test MAE (metric_best: mae, metric_agg: argmin on the validation MAE) of a single training run, averaged over the four seeds 0, 1, 2, 3. Lower is better.

One arm, one seed

ZINC (2000 epochs, ksteps = 21):

python main.py --cfg configs/GRIT/zinc-GRIT-fphit_p.yaml \
    seed 0 accelerator cuda:0 dataset.dir datasets out_dir results/zinc-fphit_p-s0

Peptides-struct (200 epochs, ksteps = 24; on-the-fly PE + batch size 8 keep host RAM and GPU memory bounded on the ~150-node graphs):

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
python main.py --cfg configs/GRIT/peptides-struct-GRIT-fphit_t.yaml \
    seed 0 accelerator cuda:0 dataset.dir datasets out_dir results/pept-fphit_t-s0

Swap the config file to select the arm: *-none, *-rrwp, *-fphit_p (FPSE-P), *-fphit_t (FPSE-T).

Full matrix (4 arms × 4 seeds)

bash scripts/run_zinc_4seed.sh          # 16 ZINC runs
bash scripts/run_peptides_4seed.sh      # 16 Peptides-struct runs

Both scripts run strictly one seed at a time and honour DATA_DIR / OUT_DIR environment variables (defaults: datasets/, results/). Set MAX_EPOCH to shorten a run for a quick smoke check, e.g. MAX_EPOCH=5 bash scripts/run_zinc_4seed.sh.

Reading the result

Each run prints, and logs to OUT_DIR/.../logging.log, a line of the form

> Epoch NNNN: ... | Best so far: epoch E  ... val_mae: V  test_mae: T

The final test_mae on that "Best so far" line (equivalently best/test_mae) is the validation-selected test MAE for that seed. Average the four seeds to get the reported cell.

---

6. Verified results (this repository's runs)

Four seeds (0,1,2,3), validation-selected test MAE ± sample std. Full per-seed values and provenance are in results/four_seed_summary.md and results/four_seed_results.json.

Structural encoding	ZINC (MAE ↓)	Peptides-struct (MAE ↓)
none (control)	0.1387 ± 0.0113	0.3804 ± 0.0608
RRWP (baseline)	0.0633 ± 0.0004	0.2634 ± 0.0129
FPSE-P	0.0577 ± 0.0046	0.2640 ± 0.0138
FPSE-T	0.0623 ± 0.0061	0.2465 ± 0.0023

Parameter counts are matched within each dataset (ZINC ≈ 473k, Peptides-struct ≈ 450k across all four arms). Peptides-struct here uses batch size 8 and mask_rel_diag under an identical pipeline for all arms; the comparison is controlled and internal (absolute MAE is not directly comparable to literature GRIT numbers trained with a different batch/PE budget).

---

7. Attribution and license

This project is a derivative of GRIT — Graph Inductive Bias Transformer — by Liheng Ma et al., "Graph Inductive Biases in Transformers without Message Passing," ICML 2023 (https://github.com/LiamMa/GRIT). GRIT is released under the MIT License, retained here in LICENSE (© 2023 Liheng Ma). The FPSE additions are released under the same MIT License. Modifications are enumerated in NOTICE.

8. Citation

@misc{fpse_graph_transformers,
  title  = {Exact First-Passage Structural Encodings for Graph Transformers},
  note   = {Manuscript under review},
  year   = {2026}
}