When Do Transformers Learn Heuristics for Graph Connectivity?

TL;DR

It is proved that an $L$-layer model has capacity to solve for graphs with diameters up to exactly $3^L$, implementing an algorithm equivalent to computing powers of the adjacency matrix.

Abstract

Transformers often fail to learn generalizable algorithms, instead relying on brittle heuristics. Using graph connectivity as a testbed, we explain this phenomenon both theoretically and empirically. We consider a simplified Transformer architecture, the disentangled Transformer, and prove that an $L$-layer model has capacity to solve for graphs with diameters up to exactly $3^L$, implementing an algorithm equivalent to computing powers of the adjacency matrix. We analyze the training-dynamics, and show that the learned strategy hinges on whether most training instances are within this model capacity. Within-capacity graphs (diameter $\leq 3^L$) drive the learning of a correct algorithmic solution while beyond-capacity graphs drive the learning of a simple heuristic based on node degrees. Finally, we empirically demonstrate that restricting training data within a model's capacity leads to both standard and disentangled transformers learning the exact algorithm rather than the degree-based heuristic.

Figures (11)

• Figure 1: We train 2-layer Transformers models on Erdős-Rényi graphs. (Left) Visualization of input, target and model prediction of a sample graph. (Right) Although trained models are able to predict perfectly on every edge within distribution, they failed to generalize to out-of-distribution graphs such as graphs with two isolated chains or cliques.
• Figure 2: Capacity of Disentangled Transformers. We train 2-layer (left) and 3-layer (right) disentangled transformers on $\mathsf{ER}(n=24)$ and $\mathsf{ER}(n=64)$ graphs respectively. When evaluated on hold-out sets, both models can only make reliable predictions ($\geq 99\%$ accuracy) on node pairs $u, v$ if and only if $d_G(u,v) \leq 3^L$. It resonates with our theoretical observations in \ref{['thm:capacity']}.
• Figure 3: Training Dynamics of Disentangled Transformers. We train an 1-layer disentangled transformer on graphs from $\mathsf{ER}(n=8, p=0.2)$ distribution. Weight $W$ will approximately approach to $A \otimes I_n + B \otimes J_n$ form. (Left) There are two major phases during training, where during Phase 1, model focuses on learning the equivariant parameterizations so both $I$ and $J$ channel's share of energy in $W$ increases, and during Phase 2, the algorithmic $I$-channel is promoted and the heuristic $J$-channel is suppressed. (Right) Visualization of the learned weights during training and its projection to the closest $\hat{W} = \hat{A} \otimes I_n + \hat{B} \otimes J_n$ form.
• Figure 4: Following insights from \ref{['thm:training-suppresses-J', 'thm:dynamics']}, we repeat the same experiment setup as in \ref{['fig:1-layer-dynamics']} but only training on within-capacity graphs (see \ref{['def:capacity_dichotomy']} and the left figure). As shown in the solid lines in the right figure, restricting training samples by capacity pushes the energy share of the algorithmic mechanism (the $A \otimes I_n$ channel) further to nearly 100% in the weight $W$. It simultaneously prevents the growth of the heuristic portion (the $B\otimes J_n$ channel).
• Figure 5: With 1-layer disentangled transformers with capacity $\mathsf{Cap} = 3$ following \ref{['thm:capacity']}, we vary the $d$ such that we restrict our training graphs to have $\operatorname{diam}(G) \leq d$. We also vary the edge probability of our training distribution $\mathsf{ER}(n=8, p = \cdot)$ for generality. We test on $\mathsf{2Chain}(n=8, k=\cdot)$ graphs with $k=2$ or $3$ and show the exact match accuracy on configurations where the accuracy is non-zero for readability. We find if the training $d \leq \mathsf{Cap}$, models still learns the algorithmic solution up to problem size $d$ (see $d=2, k=2$ case on the left in orange) but fails to length generalize (see $d=2, k=3$ in orange on the right). On the other hand, if the training $d > \mathsf{Cap}$, model struggles to learn the algorithmic solution (see $d=4$ cases in red on both $k = 2$ or $3$). The best case overall is when setting $d = \mathsf{Cap}$, i.e., preventing the model from seeing beyond-capacity samples but still preserving at-capacity samples for better generalization. As shown in the green lines, with $d=3$, model achieves balanced testing accuracy on both $k=2$ and $3$.

Theorems & Definitions (41)

• Definition 1: Self-loop-augmented adjacency matrix
• Definition 2: Connectivity
• Definition 3: Transformers for Graph Connectivity
• Definition 4: Disentangled Transformer for Graphs
• Definition 5: Graph distances and diameter
• Theorem 4.1: Expressivity
• proof : Sketch of proof
• Theorem 4.2: Capacity
• proof : Sketch of proof
• Definition 6: Within-capacity and beyond-capacity pairs at depth $L$