Skip to content

Add ppr_relation_features token input to the PPR Graph Transformer path - #767

Draft
mkolodner-sc wants to merge 1 commit into
mainfrom
mkolodner-sc/gt-ppr-relation-features
Draft

Add ppr_relation_features token input to the PPR Graph Transformer path#767
mkolodner-sc wants to merge 1 commit into
mainfrom
mkolodner-sc/gt-ppr-relation-features

Conversation

@mkolodner-sc

Copy link
Copy Markdown
Collaborator

What

Adds a reserved ppr_relation_features anchor-relative token input to the PPR
Graph Transformer path, and consolidates the reserved-feature validation shared
by the transform and the encoder.

Details

  • transforms/graph_transformer.py
    • New PPR_RELATION_FEATURES_NAME token input: the columns of the PPR
      edge_attr beyond the scalar weight (col 0) are gathered per (anchor,
      neighbor) into a (batch, max_seq_len, k) feature tensor in
      _build_sequence_layout_from_ppr_edges and exposed through the anchor
      feature composers.
    • Neighbors stay ordered by descending PPR weight; relation-feature rows are
      reordered by the same permutation so they remain aligned with their edges.
    • Relaxes _validate_ppr_sequence_input so the batch may also carry the
      original relation edge types (which feed the relation-message channel); only
      the "ppr" edges drive the sequence.
  • New _validate_reserved_anchor_feature_usage helper centralizes the
    placement/method checks for ppr_weight and ppr_relation_features, replacing
    the duplicated blocks in heterodata_to_graph_transformer_input and
    GraphTransformerEncoder.

Testing

  • ruff check / ruff format clean; module compiles.
  • Unit-level: validator rejects misuse (pairwise/scalar bias, non-ppr method);
    PPR sequences are weight-ordered with relation features aligned to their edges;
    truncation to max_seq_len keeps the top-weight neighbors.

Introduce a reserved `ppr_relation_features` anchor-relative token input for
sequence_construction_method='ppr'. The columns of the PPR edge_attr beyond the
scalar weight are laid out per (anchor, neighbor) into a per-token feature
tensor, so relation-aware downstream channels can consume the original edge
features carried on PPR edges. Neighbors stay ordered by descending PPR weight
and the relation-feature rows are permuted in lockstep to stay aligned.

Factor the reserved-feature validation (placement of ppr_weight /
ppr_relation_features and the method requirement) into a single
_validate_reserved_anchor_feature_usage helper shared by
heterodata_to_graph_transformer_input and GraphTransformerEncoder, replacing the
duplicated check blocks in the transform and the encoder.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant