ICML Workshop (2026)
Yong-Min Shin, Dae-Woong Jeong, Sehui Han
Abstract
We investigate the internal mechanisms of a tokenized graph transformer — a TokenGT-style encoder trained on graphs represented as sequences of node and edge tokens — to understand if and how a transformer can understand and solve graph tasks. Using manual attention map inspections, activation patching, and linear probing of the residual stream, we identify circuits for three fundamental graph tasks: node degree classification, ring-node detection, and shortest-path distance prediction. Our analysis reveals task-specific attention head specialization that emerges from permutation-augmented training: heads that are weakly tuned under fixed node orderings sharpen into functionally distinct roles under permutation invariance constraints. For shortest-path prediction, causal evidence shows signs of a multi-layer pipeline in which early heads construct adjacency routing structure and a single mid-network head performs what resembles a breadth-first search tree selection. For node degree and ring-node tasks, the residual stream becomes linearly decodable at shallow layers, with per-layer probe accuracy tracking the causal contributions identified by ablation. Together, these findings suggest that the model develops structured, localized computational circuits rather than diffuse representations, and that training with permutation augmentation is a key driver of this specialization. We discuss what the current evidence does and does not support, as well as limitations and potential direction of research.