Rotary position encodings improved to capture graph routes
Analytic-Walk Rotary Positional Encodings for Graphs
Machine Learning
Summary
When computers analyze graphs, they often need to know how nodes (points) connect to each other. Previous methods used angles assigned to each node but couldn't tell the difference between different paths connecting the same points. The authors created a new way to assign angles to edges and combine information across all paths between nodes. This helps the computer understand the specific routes in a graph, improving accuracy on tests and real datasets.
What this means in practice
- •For graph neural network engineers: Improve models by distinguishing different connection paths in graphs, enhancing accuracy in tasks like molecule and superpixel classification.
- •For bioinformatics developers: Enhance peptide structure analysis tools by better encoding relational graph information through edge-based positional information.
Authors
Jiaqing Xie, Yuxin Wang, Xipeng Qiu
Abstract
Rotary position encodings make attention sensitive to relative position, but extending them to graphs requires choosing how graph structure enters the rotation. Previous works assign each node a rotation from spectral coordinates, so the rotary factor between two nodes depends only on their endpoints and cannot distinguish the routes connecting them. We introduce \textit{Analytic-Walk Rotary Positional Encodings} (AW-RoPE), which place the rotations on edges and sum the transported features over all walks, so contributions along different routes can reinforce or cancel. An exact variant evaluates the complete sum by a differentiable linear solve, and a sparse variant truncates it at a finite depth. We prove forward and parameter-derivative truncation bounds at fixed inputs and parameters. Both variants act on projected queries and keys, and the sparse recurrence also augments message-passing networks. Across five synthetic tasks both variants reduce nRMSE by $15$--$58\%$ relative to the strongest baseline, and on real superpixel, peptide and OGB benchmarks the sparse recurrence attains the best mean on every dataset with Performer kernels and on twelve of thirteen datasets with GIN. Analysis shows that AW-RoPE can distinguish routes whose only cue is how two endpoints are connected, while node-wise rotary encodings cannot.