LG AI Research Applies Mechanistic Interpretability to Tokenized Graph Transformers

LG AI Research presented new findings at the ICML 2026 Mechanistic Interpretability Workshop on applying mechanistic interpretability to TokenGT, a graph transformer model. While most interpretability research focuses on LLMs, this work addresses the underexplored domain of graph transformers used in chemistry and biology. The team analyzed how vanilla transformers process graphs represented as sequences of node and edge tokens, aiming to reverse-engineer the internal computational paths.

Shared Computational Mechanisms

The analysis of models trained on degree calculation, cycle detection, and shortest-path distance revealed a shared initial mechanism. In the first transformer layer, node tokens attend strongly to edge tokens with matching endpoint identifiers. This ID-matching behavior allows the model to recover incident edges without built-in graph operations. Linear probes confirmed that node degree information is accurately identified immediately after this first attention operation, with the magnitude of the write vector in the residual stream proportional to the ground truth degree.

Task-Specific Adaptations

While the early local-structure computation is shared, the models diverge in how they compose these features for specific tasks:

• Ring Membership: The model does not implement an explicit cycle-tracing algorithm. Instead, it initially classifies nodes as part of a ring and progressively gathers evidence to flip the prediction if the node is not part of a ring. Ablation studies showed that removing the second layer primarily damaged non-ring predictions.
• Shortest-Path Distance: The model employs a more sophisticated mechanism where the third attention head in the third layer concentrates on adjacent nodes that match a parent in a breadth-first-search (BFS) tree. Replacing this attention target with the corresponding BFS parent largely preserved performance, indicating the model views the input graph as a rooted BFS tree.

Future Directions in AI4Science

This research provides an initial mechanistic account of how transformers perform graph computation without architectural biases. LG AI Research plans to extend this work to realistic molecular graphs to understand how models learn chemical facts and rules. The goal is to bridge the gap between AI models and domain scientists, facilitating the extraction of scientific knowledge through interpretability analysis.