Graph Machine: Exploring Edge Mechanisms as an Inductive Bias

# Exploring Edge Mechanisms as an Inductive Bias
Source: [https://arxiv.org/html/2608.06834](https://arxiv.org/html/2608.06834)
###### Abstract

Transformers provide a powerful architecture for global content\-based matching, but reasoning problems may benefit from a stronger inductive bias toward iterative traversal of latent relations\. We introduce Graph Machine, an architecture with two explicit edge\-based mechanisms: Edge\-augmented attention, in which edges modulate attention between nodes, and edge\-centric referral, in which nodes exchange addresses to update their edges\. Conceptually, this enables the model to dynamically and differentiably construct and revise relational graphs across layers\. We study this inductive bias using Sudoku under controlled settings and find that Graph Machine outperforms Transformer baselines, with ablation studies and mechanistic analysis attributing the gains to the edge mechanisms\. Surprisingly, we found that the model discovers a compact edge\-based construction for Sudoku geometry\. Our results support explicit edge mechanisms as a promising architectural design, motivating broader evaluation\.

![[Uncaptioned image]](https://arxiv.org/html/2608.06834v1/figures/overview.png)## 1Introduction

Most ML architectures, including Transformers[1](https://arxiv.org/html/2608.06834#bib.bib1)and graph neural networks \(GNNs\)[2](https://arxiv.org/html/2608.06834#bib.bib2),[3](https://arxiv.org/html/2608.06834#bib.bib3), can be viewed as operating over nodes that aggregate information from other nodes before applying node\-wise operations\. This view raises a question: from which other nodes should a node absorb information? Modern architectures commonly answer this question through some combination of global content\-based matching and iterative relational traversal\. The former is exemplified by dot\-product self\-attention, which selects nodes based on their features\. The latter, by contrast, relies on edges, whether explicit or implicit, to propagate information through a relational structure\. A directed edge can be understood as carrying the addresses of its target nodes, and potentially also features that describe the relation between the source and the target\.

We hypothesize that the relative ease of global content\-based matching versus iterative relational traversal constitutes an important inductive bias in architectural design\. Systems like large language models \(LLMs\)[4](https://arxiv.org/html/2608.06834#bib.bib4)are known to favor surface heuristics over latent constraints[5](https://arxiv.org/html/2608.06834#bib.bib5)\. Among many examples of what is described as shortcut learning, one demonstrates it with particular clarity: “I want to wash my car\. The car wash is 50 meters away\. Should I walk or drive?” An LLM often answers “walk” by treating the problem as a familiar distance\-focused question and attending to the salient distance cue\. The correct answer, however, is to drive, because the car itself must be brought to the car wash[6](https://arxiv.org/html/2608.06834#bib.bib6)\. Global content\-based matching gives a model broad reach and rich representational capacity, but it also makes shallow shortcuts immediately available\. Iterative relational traversal, by contrast, may be better aligned with reasoning over latent relations and may regularize against content\-based shortcuts\. If this prior holds, then making the latter easy becomes a central desideratum\.

However, we argue that architectures like Transformers and GNNs often do the opposite\. One important reason is the lack of an edge\-centric referral mechanism\. Such a mechanism would encourage the model to pass not only messages between nodes, but also addresses\. These addresses could then be used to construct new edges for the next information aggregation step\. Intuitively, this allows nodes to fetch a neighbor’s neighbor by letting the neighbor provide the referral\.

GNNs commonly model edges directly by providing each node with a set of neighbors, which the node then uses for message passing\. However, the topology is typically fixed, or dynamic only through a procedure that is external to the model’s own computation\. A GNN may be forced to operate over a graph structure that is not amenable to the computation the model wishes to perform, and therefore suffers from linearly depth\-limited receptive fields, 1\-WL expressivity[7](https://arxiv.org/html/2608.06834#bib.bib7)bound, and oversquashing[8](https://arxiv.org/html/2608.06834#bib.bib8)\.

By contrast, dot\-product self\-attention in Transformers directly supports global content\-based matching\. Although self\-attention may have enough capacity to implement iterative relational traversal and referral, it makes these operations difficult in three ways: limited addressing precision, noisy address channels, and competition with global content\-based matching\.

First, dot\-product attention retrieves nodes through dot products in adkd\_\{k\}\-dimensional key/query space\. Viewing exact node addresses as one\-hot vectors in annn\-dimensional address space, this can be interpreted as preserving address distinguishability in adkd\_\{k\}\-dimensional geometry, suggesting a trade\-off between dimension and distortion analogous to Johnson–Lindenstrauss\-type bounds[9](https://arxiv.org/html/2608.06834#bib.bib9)\.

Second, fetched addresses must be carried through the value stream and subsequent feed\-forward networks\. These transformations are useful for representation learning, but noisy for exact bookkeeping: an address should ideally be forwarded without modification\.

Third, global content\-based matching and iterative relational traversal/referral share the same infrastructure, forcing them to compete for dimensions, heads, and optimization pressure\. Thus, the balance between the two modes of computation is determined by implicit trade\-offs, leaving no explicit control to bias the model toward the latter\.

This work introduces Graph Machine \(GM\), an architecture that directly supports edge\-centric referral while generalizing Transformers in expressivity\. Graph Machines update node features through edge\-augmented attention and construct new edges through edge\-centric referral\.

In controlled Sudoku experiments, Graph Machines outperform same\-scale and prior\-advantaged Transformer baselines, while remaining comparable to substantially enlarged and prior\-advantaged Transformer baselines\. Mechanistic analysis suggests that this advantage reflects the intended inductive bias\. Starting from only local adjacency edges, GM learns to construct higher\-level Sudoku relations, including box, row, and column regions, through a compact edge\-based construction that resembles a11\-22\-44expansion process\. This provides evidence that the model is not merely fitting Sudoku solutions, but using referral to build task\-relevant latent graph structure\.

These results suggest that explicit edge mechanisms can provide a useful inductive bias for structured reasoning tasks\.

## 2Architecture

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/representations.png)Figure 1:GM representations\.
![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/layers.png)Figure 2:GM layers\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/edge-sublayer.png)Figure 3:GM edge sublayer\.
![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/node-sublayer.png)Figure 4:GM node sublayer\.

### 2\.1Architecture overview

We conceptualize a Graph Machine as operating on a graph withnnnodes, where each node maintainskkoutgoing edge slots \(Figure[2](https://arxiv.org/html/2608.06834#S2.F2)\)\. In addition to node features of shapen×dnn\\times d\_\{n\}, the model maintains edge features of shapen×k×den\\times k\\times d\_\{e\}and edge addresses of shapen×k×nn\\times k\\times n, wherednd\_\{n\}andded\_\{e\}denote the node and edge hidden sizes, respectively\. The edge addresses represent target mass: each edge assigns weights over thenntarget nodes\.

The initial edge features and edge addresses are provided by the input layer, either by converting task\-specific input graphs or through priors such askkpast neighbors in the case of sequence modeling \(Figure[2](https://arxiv.org/html/2608.06834#S2.F2)\)\. Each intermediate layer of a Graph Machine consists of one or more edge sublayers followed by one or more node sublayers, allowing many\-to\-one or one\-to\-many arrangements\. Edge sublayers update edge representations using edge\-centric referral, followed by a position\-wise feed\-forward network that updates edge features \(Figure[4](https://arxiv.org/html/2608.06834#S2.F4)\)\. Node sublayers update node features using edge\-augmented attention followed by their own position\-wise feed\-forward network \(Figure[4](https://arxiv.org/html/2608.06834#S2.F4)\)\.

Viewed through a programming analogy, node features act like an object’s general state, while each edge slot resembles a relational field that stores both relation\-specific state and a soft pointer\. Each object can traverse its pointers to reach target objects, fetch their states and pointers, and use the retrieved information to revise its own pointers while updating its general and relational states\.

Another way to view Graph Machines is through relational composition\. For an ordinary graph with adjacency matrixPP, the composition operator∘\\circis defined by matrix multiplication:P∘P=P2P\\circ P=P^\{2\}\. Graph Machines generalize relational representation from connectivity to multiple feature\-bearing soft edges\. The edge representation is a pair\(Xe,Pe\)\(X\_\{e\},P\_\{e\}\), whereXnX\_\{n\}andXeX\_\{e\}stores node and node\-relation\-specific features andPeP\_\{e\}is edge addresses, equivalentlykkrow\-normalized soft adjacency matrices\. Edge\-centric referral then implicitly defines a learned differentiable composition operator⋄θ\\diamond\_\{\\theta\}over this representation:\(Xe′,Pe′\)=⋄θ\(Xn,Xe,Pe\)\(X^\{\\prime\}\_\{e\},P^\{\\prime\}\_\{e\}\)=\\diamond\_\{\\theta\}\(X\_\{n\},X\_\{e\},P\_\{e\}\)\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/attention.png)Figure 5:Edge\-augmented attention tensor operations \(11head\)\.
### 2\.2Edge\-augmented attention

Standard multi\-head dot\-product self\-attention computes attention logits from queries and keys projected from node features \(XnX\_\{n\}\)\. Graph Machines retain this term as a*node factor*\(FnodeF\_\{\\mathrm\{node\}\}\), but augment it with an*edge factor*\(FedgeF\_\{\\mathrm\{edge\}\}\) derived from edge features \(XeX\_\{e\}\) and edge addresses \(PeP\_\{e\}\)\. The final attention weights are a product\-of\-experts[10](https://arxiv.org/html/2608.06834#bib.bib10)combination between the factors, so that node\-based content affinity and edge\-based relational evidence jointly shapes the target map\.

For simplicity, we henceforth omit the head dimension and denote a node\-edge\-node\-edge\-node chain by\(n​1,e​1,n​2,e​2,n​3\)\(n1,e1,n2,e2,n3\), wheren​1n1is the source node,e​1e1is an edge of the source node,n​2n2is a target node, ande​2e2is an edge of the target node\. Queries \(QQ\) are projected fromn​1n1features, while the standard attention keys are projected fromn​2n2features; under the chain notation above, we refer to them as*n​2n2keys*\(Kn​2K\_\{n2\}\)\. To produce the edge factor, the samen​1n1queries attend over*e​1e1keys*\(Ke​1K\_\{e1\}\), which are projected from thekkedge features associated withe​1e1\. The corresponding edge addresses are used as values, yielding a mixture of edge\-address distributions over target nodes \(Ae​1A\_\{e1\}\)\. To convert this mixture into logit space, we then take the logarithm of this target mass after clipping away from0with a smallε\\varepsilon\. Concretely,

Ae​1​\(n​1,e​1\)=\\displaystyle A\_\{e1\}\(n1,e1\)=\{\}softmaxe​1⁡\(Q​\(n​1\)​Ke​1​\(n​1,e​1\)⊤dak\),\\displaystyle\\operatorname\{softmax\}\_\{e1\}\\\!\\Big\(\\frac\{Q\(n1\)\\,K\_\{e1\}\(n1,e1\)^\{\\top\}\}\{\\sqrt\{d\_\{\\mathrm\{ak\}\}\}\}\\Big\),Medge​\(n​1,n​2\)=\\displaystyle M\_\{\\mathrm\{edge\}\}\(n1,n2\)=\{\}∑e​1Ae​1​\(n​1,e​1\)​Pe​\(n​1,e​1,n​2\),\\displaystyle\\sum\\nolimits\_\{e1\}A\_\{e1\}\(n1,e1\)\\,P\_\{e\}\(n1,e1,n2\),Fedge​\(n​1,n​2\)=\\displaystyle F\_\{\\mathrm\{edge\}\}\(n1,n2\)=\{\}log⁡\(max⁡\(Medge​\(n​1,n​2\),ε\)\),\\displaystyle\\log\\\!\\Big\(\\\!\\max\\\!\\big\(M\_\{\\mathrm\{edge\}\}\(n1,n2\),\\,\\varepsilon\\,\\big\)\\Big\),Fnode​\(n​1,n​2\)=\\displaystyle F\_\{\\mathrm\{node\}\}\(n1,n2\)=\{\}Q​\(n​1\)​Kn​2​\(n​2\)⊤dak\.\\displaystyle\\frac\{Q\(n1\)\\,K\_\{n2\}\(n2\)^\{\\top\}\}\{\\sqrt\{d\_\{\\mathrm\{ak\}\}\}\}\.
The node and edge factors are multiplied by learned temperature scalars projected from node features and parameterized by a positive function such asexp\\exporsoftplus\\operatorname\{softplus\}\(tnodet\_\{\\mathrm\{node\}\}andtedget\_\{\\mathrm\{edge\}\}\)\. This allows the model to modulate the relative strength of the two experts\. Summing the22factors and applying a softmax over the target\-node dimension gives the final attention weights\. The mechanism is at least as expressive as standard Transformer attention: vanilla self\-attention is recovered by setting the edge\-factor temperature to0, so that the edge expert contributes a uniform distribution\. Concretely,

A​\(n​1,n​2\)=\\displaystyle A\(n1,n2\)=\{\}softmaxn​2⁡\(tnode​\(n​1\)​Fnode​\(n​1,n​2\)\+tedge​\(n​1\)​Fedge​\(n​1,n​2\)\),\\displaystyle\\operatorname\{softmax\}\_\{n2\}\\\!\\Big\(t\_\{\\mathrm\{node\}\}\(n1\)\\,F\_\{\\mathrm\{node\}\}\(n1,n2\)\+t\_\{\\mathrm\{edge\}\}\(n1\)\\,F\_\{\\mathrm\{edge\}\}\(n1,n2\)\\Big\),Z​\(n​1\)=\\displaystyle Z\(n1\)=\{\}∑n​2A​\(n​1,n​2\)​V​\(n​2\)\.\\displaystyle\\sum\\nolimits\_\{n2\}A\(n1,n2\)\\,V\(n2\)\.
We visualize the tensor operations during one head of the edge\-augmented attention mechanism in Figure[5](https://arxiv.org/html/2608.06834#S2.F5)and provide simplified code in Section[A\.1](https://arxiv.org/html/2608.06834#A1.SS1)\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/referral.png)Figure 6:Edge\-centric referral tensor operations \(11head\)\.
### 2\.3Edge\-centric referral

Edge\-centric referral constructs a new set of edges by composing two\-hopn​1→n​2→n​3n1\\to n2\\to n3chains into single\-hopn​1→n​3n1\\to n3edges with new features and addresses\.

During edge\-centric referral, thekknew edges play a dimensional role analogous to attention heads\. Because self\-connections are representable in Graph Machines, the mechanism does not preclude preserving existing edges\.

Referral shares its early computational pattern with edge\-augmented attention\. The queries produce node and edge factors over possible intermediate nodesn​2n2, denoted the*n​2n2node factor*\(Fn​2,nodeF\_\{n2,\\mathrm\{node\}\}\) and*n​2n2edge factor*\(Fn​2,edgeF\_\{n2,\\mathrm\{edge\}\}\), while the attention weights over the originale​1e1edge slots are retained as*e​1e1weights*\(Ae​1A\_\{e1\}\)\. The same queries also attend over the outgoinge​2e2edge slots of eachn​2n2to produce an*e​2e2factor*\(Fe​2F\_\{e2\}\)\. Combining then​2n2node factor,n​2n2edge factor, ande​2e2factor with learned temperatures yields logits over referred edges; applying a softmax gives*e​2e2weights*\(Ae​2A\_\{e2\}\)\. These weights define a distribution over candidatee​2e2edges\. The corresponding*n​2n2weights*\(An​2A\_\{n2\}\) are obtained by summing thee​2e2weights over thee​2e2edge dimension\. By contrast, a different variant would first computen​2n2weights and then multiply them by softmax\-normalizede​2e2factors, thereby modeling a conditional distribution overe​2e2givenn​2n2\. Concretely, after early analogous operations,

Fe​2​\(n​1,n​2,e​2\)=\\displaystyle F\_\{e2\}\(n1,n2,e2\)=\{\}Q​\(n​1\)​Ke​2​\(n​2,e​2\)⊤drk,\\displaystyle\\frac\{Q\(n1\)\\,K\_\{e2\}\(n2,e2\)^\{\\top\}\}\{\\sqrt\{d\_\{\\mathrm\{rk\}\}\}\},Ae​2​\(n​1,n​2,e​2\)=\\displaystyle A\_\{e2\}\(n1,n2,e2\)=\{\}softmaxn​2,e​2\(tn​2,node\(n1\)Fn​2,node\(n1,n2\)\\displaystyle\\operatorname\{softmax\}\_\{n2,\\,e2\}\\\!\\Big\(t\_\{n2,\\mathrm\{node\}\}\(n1\)\\,F\_\{n2,\\mathrm\{node\}\}\(n1,n2\)\+tn​2,edge​\(n​1\)​Fn​2,edge​\(n​1,n​2\)\\displaystyle\\phantom\{\\operatorname\{softmax\}\_\{n2,\\,e2\}\\\!\\Big\(\}\+t\_\{n2,\\mathrm\{edge\}\}\(n1\)\\,F\_\{n2,\\mathrm\{edge\}\}\(n1,n2\)\+te​2\(n1\)Fe​2\(n1,n2,e2\)\),\\displaystyle\\phantom\{\\operatorname\{softmax\}\_\{n2,\\,e2\}\\\!\\Big\(\}\+t\_\{e2\}\(n1\)\\,F\_\{e2\}\(n1,n2,e2\)\\Big\),An​2​\(n​1,n​2\)=\\displaystyle A\_\{n2\}\(n1,n2\)=\{\}∑e​2Ae​2​\(n​1,n​2,e​2\)\.\\displaystyle\\sum\\nolimits\_\{e2\}A\_\{e2\}\(n1,n2,e2\)\.
The new edge features are obtained by taking weighted sums of the values associated withe​1e1,n​2n2, ande​2e2\(Ve​1V\_\{e1\},Vn​2V\_\{n2\}, andVe​2V\_\{e2\}\), using their respective weights, then adding the results and projecting back to the edge hidden dimension\. The new edge addresses are computed by weighting the addresses of thee​2e2edges with thee​2e2weights\. Consequently, edge addresses would remain convex combinations of previous addresses absent the sharpening procedure\. Concretely,

Zx​\(n​1\)=\\displaystyle Z\_\{x\}\(n1\)=\{\}∑e​1Ae​1​\(n​1,e​1\)​Ve​1​\(n​1,e​1\)\\displaystyle\\sum\\nolimits\_\{e1\}A\_\{e1\}\(n1,e1\)\\,V\_\{e1\}\(n1,e1\)\+∑n​2An​2​\(n​1,n​2\)​Vn​2​\(n​2\)\\displaystyle\+\\sum\\nolimits\_\{n2\}A\_\{n2\}\(n1,n2\)\\,V\_\{n2\}\(n2\)\+∑n​2,e​2Ae​2​\(n​1,n​2,e​2\)​Ve​2​\(n​2,e​2\),\\displaystyle\+\\sum\\nolimits\_\{n2,\\,e2\}A\_\{e2\}\(n1,n2,e2\)\\,V\_\{e2\}\(n2,e2\),Zp​\(n​1,n​3\)=\\displaystyle Z\_\{p\}\(n1,n3\)=\{\}∑n​2,e​2Ae​2​\(n​1,n​2,e​2\)​Pe​\(n​2,e​2,n​3\)\.\\displaystyle\\sum\\nolimits\_\{n2,\\,e2\}A\_\{e2\}\(n1,n2,e2\)\\,P\_\{e\}\(n2,e2,n3\)\.
We visualize the tensor operations during one head of the edge\-centric referral mechanism in Figure[6](https://arxiv.org/html/2608.06834#S2.F6)and provide simplified code in Section[A\.2](https://arxiv.org/html/2608.06834#A1.SS2)\.

### 2\.4Sharpening

Because Shannon entropyH​\(⋅\)H\(\\cdot\)is concave in the edge address distributionPeP\_\{e\}, Jensen’s inequality gives

H​\(∑n​2,e​2Ae​2​\(n​1,n​2,e​2\)​Pe​\(n​2,e​2,⋅\)\)≥∑n​2,e​2Ae​2​\(n​1,n​2,e​2\)​H​\(Pe​\(n​2,e​2,⋅\)\),H\\Big\(\\sum\\nolimits\_\{n2,\\,e2\}A\_\{e2\}\(n1,n2,e2\)\\,P\_\{e\}\(n2,e2,\\cdot\)\\Big\)\\geq\\sum\\nolimits\_\{n2,\\,e2\}A\_\{e2\}\(n1,n2,e2\)\\,H\\big\(P\_\{e\}\(n2,e2,\\cdot\)\\big\),whereAe​2​\(n​1,n​2,e​2\)≥0A\_\{e2\}\(n1,n2,e2\)\\geq 0and∑n​2,e​2Ae​2​\(n​1,n​2,e​2\)=1\\sum\\nolimits\_\{n2,\\,e2\}A\_\{e2\}\(n1,n2,e2\)=1\. Thus, repeatedly forming weighted averages of edge addresses during referral can dilute the target mass of the resulting edges\.

To counteract this effect, Graph Machines apply an edge\-sharpening operation at the start of edge and node sublayers\. This operation performs learned temperature scaling on edge addresses: each edge address is mapped to logit space by taking the logarithm, scaled by a learned temperature \(tsharpenert\_\{\\mathrm\{sharpener\}\}\) projected from edge features and parameterized by a positive temperature function, and renormalized with a softmax\. Concretely,

Le​\(n​1,e​1,n​2\)=\\displaystyle L\_\{e\}\(n1,e1,n2\)=\{\}log⁡\(max⁡\(Pe​\(n​1,e​1,n​2\),ε\)\),\\displaystyle\\log\\\!\\Big\(\\\!\\max\\\!\\big\(P\_\{e\}\(n1,e1,n2\),\\,\\varepsilon\\,\\big\)\\Big\),Pe′​\(n​1,e​1,n​2\)=\\displaystyle P^\{\\prime\}\_\{e\}\(n1,e1,n2\)=\{\}softmaxn​2⁡\(tsharpener​\(n​1,e​1\)​Le​\(n​1,e​1,n​2\)\)\.\\displaystyle\\operatorname\{softmax\}\_\{n2\}\\\!\\Big\(t\_\{\\mathrm\{sharpener\}\}\(n1,e1\)\\,L\_\{e\}\(n1,e1,n2\)\\Big\)\.
In edge sublayers, one could use separate sharpener temperatures for the two roles played by edge addresses: one version for producingn​2n2edge factors and another as thee​2e2address being collected\. For simplicity, we use a single sharpener temperature for both roles\.

### 2\.5Computational optimizations

A direct implementation of the conceptual algorithm incurs anO​\(n2\)O\(n^\{2\}\)memory footprint and anO​\(n3\)O\(n^\{3\}\)compute cost in the number of nodesnn\. These costs arise from materializing ann×h×nn\\times h\\times nweight tensor for each batch element in attention and ann×k×n×kn\\times k\\times n\\times kweight tensor in referral, and performing matrix multiplication between the weight tensor and ann×k×nn\\times k\\times naddress tensor during referral\. Two natural optimization strategies reduce this cost by reducing the footprint of the edge addresses\.

The first direction is to compress edge addresses into a lower\-dimensional representation and recover them with a linear decoder during attention and referral\. Because a vanilla compression\-decompression process does not respect the simplex constraint, it is more natural to represent the edge address in logit space instead of weight space \(we present the results on edge address space without compression in Section[C](https://arxiv.org/html/2608.06834#A3.SS0.SSS0.Px5)\)\. This requires interpreting edge\-address aggregation steps as also a product\-of\-experts operation rather than a mixture\. In this setting, the edge factors can be computed, together with the node factors, through dot products between queries concatenated with compressed target edge addresses and keys concatenated with the decoder matrix\. This makes the mechanism compatible with FlashAttention\-style kernels[11](https://arxiv.org/html/2608.06834#bib.bib11)\.

The second direction is to sparsify edge addresses to contain only the top\-ssentries, storing them in a sparse coordinate format and utilizing sparse operations \(`scatter`/`gather`\) during referral and attention\. This allows the memory footprint to beO​\(n\)O\(n\)and computational cost to beO​\(n2\)O\(n^\{2\}\), as long as edge factor temperatures are enforced away from0so that attention/referral weights inherit the sparsity of edge addresses\. Each edge\-address aggregation step requires coalescing, and to maintain sparsity, subsequent selection, either via hard top\-ssor probabilistic Gumbel\-top\-ss\. Additionally,sscould potentially differ between various operations\. In contrast to vanilla Transformers with dynamic dense interactions and sparse Transformers with fixed routing, a sparse Graph Machine would provide sparse but dynamic attention\. We present preliminary simulated sparsification results in Section[C](https://arxiv.org/html/2608.06834#A3.SS0.SSS0.Px6)\.

One limiting case is where edges become hard and referral amounts to assigning weights to at mostk2k^\{2\}neighbor candidates and selectingkkof them\. With such hardened edges, optimization becomes more difficult: referral is no longer fully differentiable, and would need reinforcement\-learning\-style estimators such as policy gradients\. In this sense, GM can be viewed as a soft generalization of hard referral, allowing varying degrees of edge softness so that gradients can flow through the referral process\.

A further optimization is to avoid requiring each of thekkheads to attend separately to each of thekke​2e2edges of ann​2n2node during referral\. Instead, the model can first aggregate edge addresses for each head using a procedure analogous toe​1e1aggregation, reducing then×k×n×kn\\times k\\times n\\times kmemory cost ton×k×nn\\times k\\times n\.

## 3Experiments

### 3\.1Experimental setup

Our experimental design follows the principle of controlled comparison: we isolate the causal effects of the variables of interest by keeping all other factors fixed\. All GM and Transformer conditions are therefore considered instances of a generalized model class and share the same model and training implementation, with configurations kept at their default except for those relevant to the experiments\. We use a deliberately simple setup and avoid many straightforward optimizations; for example, increasing model size or training steps would improve performance but offer little additional insight\.

We use the Sudoku\-3M dataset from Kaggle[12](https://arxiv.org/html/2608.06834#bib.bib12),[13](https://arxiv.org/html/2608.06834#bib.bib13)as our benchmark, which contains33M Sudoku puzzles with2323to2626clues and varying levels of difficulty\. Controlling the task prior given to the model is crucial, since the prior can strongly influence both the task the model faces and its performance\. At one extreme, Sudoku can be solved entirely by human\-engineered symbolic algorithms, without using a neural network\. As we intend to use Sudoku to test a model’s capability to use input relations to construct relations amenable to reasoning, relations such as the Sudoku constraint regions should not be made readily available\. Therefore, our standard task prior is defined as local information: for each cell, knowledge of the cell itself and of its four adjacent neighbors\.

We use a single model pass for both training and testing\. Each of the8181cells is treated as a node \(n=81n=81\), with88edges per node \(k=8k=8\)\. In conditions with edges, the initial edges point to the cell itself and its four adjacent neighbors; the remaining edge slots, up to55for corner cells, are assigned to empty edges\. Edge features are initialized from learned embeddings of the edge category: self, up, down, left, right, or empty\. Initial edge addresses place all mass on the target cell for non\-empty edges and distribute mass uniformly over all cells for empty edges\. Node features are initialized by adding two learned embeddings: one for the cell content, which is either blank or one of1,…,91,\\ldots,9, and one for the cell position, with each of the8181positions treated as a separate category\. Prediction logits are produced from the final node embeddings\. Cross\-entropy loss is applied only to logits at blank positions, and then averaged over puzzles and subsequently batches\. We report board accuracy as our main metric, where a board is considered correct if and only if each blank cell is predicted correctly\.

We use a tiny model that is relatively narrow and deep, motivated by Sudoku’s requirement for multi\-step reasoning\. For our default configuration, the node hidden sizednd\_\{n\}and edge hidden sizeded\_\{e\}are6464and88, respectively; the number of layersllis3232, with each layer containing11edge sublayer and11node sublayer\. Both the edge degree, equivalently the number of referral heads,kkand the number of attention headshhare88, and the referral key sizedrkd\_\{\\mathrm\{rk\}\}and the attention key and value sizesdakd\_\{\\mathrm\{ak\}\}anddavd\_\{\\mathrm\{av\}\}are all88\. The detailed default configuration for all conditions is provided in Section[B](https://arxiv.org/html/2608.06834#A2)\. In general, we follow current best practices for Transformers\.

Each run uses a single NVIDIA GeForce RTX 4090 GPU with2424GB of VRAM and generally completes in under1515hours\. We use integers starting from0as random seeds, randomizing the model initialization, the train\-evaluation\-test split, and the epoch shuffling\. Conditions in Section[3\.2](https://arxiv.org/html/2608.06834#S3.SS2)use1010seeds; all other conditions use33seeds\. Parenthesized conditions in later sections are equivalent to one of the main conditions in Section[3\.2](https://arxiv.org/html/2608.06834#S3.SS2)\(up to numerical error\), so their values are taken from the corresponding1010\-seed averages\.

### 3\.2Comparison against baselines

Table 1:Comparison against baselines\.We compare Graph Machines against several Transformer\-based baselines\.

#### Transformer

We include a standard Transformer as one of our baselines\. To maintain symmetry with GM conditions, we retain factor temperature and auxiliary entropy losses even in baseline conditions, where attention contains only the node factor\. Ablations suggest that these choices do not materially affect Transformer performances; see Section[C](https://arxiv.org/html/2608.06834#A3.SS0.SSS0.Px2)and[C](https://arxiv.org/html/2608.06834#A3.SS0.SSS0.Px3)\.

#### Transformer\-static

For the Transformer\-static condition, we disable edge referral by removing edge sublayers\. The given input edge embeddings and edge addresses are therefore not updated, but are still used by edge\-augmented attention\. This condition remains a generalization of Transformer in expressivity: since empty edges with even mass distribution are supplied, the model can construct a uniform edge expert, effectively adding a constant term to attention scores, therefore not limiting message passing to the given topology\.

#### Transformer\-sin\-PE

For the Transformer\-sin\-PE condition, we provide the Transformer with22D positional information through sinusoidal positional embeddings[1](https://arxiv.org/html/2608.06834#bib.bib1)with base10,00010\{,\}000at the input layer\. This is the best\-performing positional embeddings option for Transformer among sinusoidal, rotary[14](https://arxiv.org/html/2608.06834#bib.bib14), and row\-column factorized positional embeddings schemes of various bases \(Section[C](https://arxiv.org/html/2608.06834#A3.SS0.SSS0.Px4)\)\.

#### Transformer\-sin\-PE\-2x

We include a scaled\-up Transformer with22D sinusoidal position embeddings as a stronger comparison, using models that are2×2\\timeswide in hidden size and head dimension sizes and2×2\\timesdeep in number of layers\.

These conditions provide useful isolation of the relevant factors\. The comparison between GM and Transformer isolates the overall effect of edge representations and mechanisms, while the comparison between GM and Transformer\-static isolates the edge\-updating referral mechanism\. Transformer\-sin\-PE examines the effect of task priors by making0to88\-hop positional information directly accessible at input, from which the Sudoku constraint regions can be easily derived, even though this arguably removes much of the challenge of constructing useful relations from given ones, which was our intention when choosing Sudoku as the benchmark and local information as our task priors\. Together with an experiment where we project input edges into the node feature space \(Section[C](https://arxiv.org/html/2608.06834#A3.SS0.SSS0.Px1)\), they remove or reverse the effect of the0to11\-hop positional task priors given to GM\. Finally, the scaled\-up Transformer\-sin\-PE\-2x provides a generous baseline with advantaged task priors and approximately8×8\\timesthe parameter count\.

Results in Table[1](https://arxiv.org/html/2608.06834#S3.T1)show that GM outperforms Transformer, Transformer\-static, and Transformer\-sin\-PE, while remaining competitive with Transformer\-sin\-PE\-2x\. The Transformer\-static condition yields an average per\-cell accuracy of77\.3%77\.3\\%, moderately lower than Transformer’s96\.3%96\.3\\%, with this discrepancy amplified under the per\-board accuracy metric\.

Transformer\-static’s poor performance, together with the first three ablation studies, presents an interesting picture\. Several conditions are equally or more expressive than the models they underperform: Transformer\-static relative to Transformer, GM with RoPE relative to GM, and 5\-edge\-degree GM and 4\-edge\-sublayer GM relative to Transformer\. The natural interpretation is that our architectural changes introduce inductive biases that can materially move models into different regimes, and that these regimes yield gains only when they are both sufficiently supported and advantageous\. The similar accuracy of GM and the strongest Transformer baseline may therefore reflect two distinct regimes being pushed near the limit under our task and setup, rather than a lack of meaningful effect from inductive bias\.

### 3\.3Ablation studies

#### Positional encoding

We add the best performing sinusoidal and rotary positional embeddings schemes \(evaluated on Transformer, Section[C](https://arxiv.org/html/2608.06834#A3.SS0.SSS0.Px4)\) to GM\. Results \(Table[3](https://arxiv.org/html/2608.06834#S3.T3)\) show that sinusoidal PE appears to slightly improve GM performance to be on par with Transformer\-sin\-PE\-2x\. On the other hand, consistent with our unreported preliminary experiments, RoPE appears to stably decrease GM’s performance to Transformer\-RoPE level shown in Section[C](https://arxiv.org/html/2608.06834#A3.SS0.SSS0.Px4)\. We theorize that RoPE may compete to supply positional information during attention, resulting in a regime that underperforms relative to when edges are used in that role by the model\.

Table 2:Positional encoding\.
Table 3:Edge degree\.

#### Edge degree

The edge degreekkcontrols the sparsity of the edge representation and affects referral compute at ratek2k^\{2\}\. In addition to an edge degree of0conceptually represented by the Transformer baseline, we test an edge degree of55, the smallest value that can accommodate the five provided non\-empty edges \(up, down, left, right, self\) in the input graph\. Interestingly, lower\-degree GM appears to underperform relative to both endpoints \(Table[3](https://arxiv.org/html/2608.06834#S3.T3)\)\.

#### Number of edge sublayers

The ratio of edge to node sublayers controls the relative amount of computation allocated to message passing and address passing\. The Transformer\-static condition can be viewed as the limiting case with no edge sublayers, so edge representations are never updated\. We test the effect of varying the number of edge sublayers while fixing the number of node sublayers to3232\. Performance decreases with fewer edge sublayers, with a trend of diminishing returns toward the upper range \(Table[5](https://arxiv.org/html/2608.06834#S3.T5)\)\. The average for the1616\-edge\-sublayer condition is skewed by seed11’s outlier performance of51\.5%51\.5\\%, compared with79\.1%79\.1\\%and85\.9%85\.9\\%for seeds0and22, respectively\.

Table 4:Number of edge sublayers\.
Table 5:Attention experts\.

#### Attention experts

Standard attention in Transformer uses only the node expert, whereas edge\-augmented attention in GM combines the node and edge experts\. To complete the picture, we also test a variant that uses only the edge expert\. The edge\-only variant underperforms the node\-only variant, while both single\-expert conditions underperform the GM double\-expert condition, suggesting that node and edge experts provide complementary utility \(Table[5](https://arxiv.org/html/2608.06834#S3.T5)\)\.

#### Sharpener temperature

We compare several parameterizations of the sharpener temperature: a fixed value shared across layers and edges, learned values per edge and layer, and projections from edge features at each layer\. Performance decreases under the alternative parameterizations, suggesting that feature\-dependent sharpening is better in our setting \(Table[6](https://arxiv.org/html/2608.06834#S3.T6)\)\.

Table 6:Sharpener temperature\.
#### Factor temperatures

Factor temperatures allow the model to modulate the relative contribution of node and edge experts\. Turning off factor temperatures reduces performance, suggesting that adaptive weighting between the two factors is beneficial \(Table[8](https://arxiv.org/html/2608.06834#S3.T8)\)\.

Table 7:Factor temperatures\.
Table 8:Entropy losses\.

#### Entropy losses

Preliminary experiments suggest that small entropy losses on the edge and noden2n\_\{2\}factors of both edge and node sublayers, each with weight starting at10−310^\{\-3\}and annealed linearly to0, reduce performance variability early in training\. We also observe a slight but non\-significant improvement in final performance attributed to the entropy losses \(Table[8](https://arxiv.org/html/2608.06834#S3.T8)\)\.

### 3\.4Mechanistic analysis

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-puzzle.png)Figure 7:Seed0, test batch0, sample0: puzzle, solution, and predictions \(GM\-0and T\-0\)\.The explicit edge representations and mechanisms lend GM to self\-interpretability\. Various variables in each sublayer are weights over the target nodes, which in our experiments correspond to the8181Sudoku cells\. These variables include edge addresses, attention/referral edge/node factors, and attention/referral weights\. To visualize the representations and mechanisms, we record these variables during the testing of the final models and plot them as heatmaps\.

We restrict this analysis to the GM and Transformer conditions from Section[3\.2](https://arxiv.org/html/2608.06834#S3.SS2), using seed0, test batch0, and samples0–1515\. We focus primarily on one pre\-specified source cell,r​2​c​2r2c2, located at the center of the top\-left box, and extend the analysis to two additional cells,r​2​c​5r2c5andr​5​c​2r5c2, where relevant\. Throughout this section, we use0\-based indexing and refer to the two examined models as GM\-0and T\-0\. We treat each consecutive pair of an edge sublayer and a node sublayer as a layer\. We use the term referral head to denote the new\-edge\-producing dimension during referral;ss,ll, andeeto denote sample, layer, and edge;aaandrrto denote attention and referral; and address\-in to denote edge addresses after sharpening in eitheraaorrr\. For example, referral heads​0​l​0​r​0s0l0r0produces an edge address that becomes two address\-in variables after sharpening:s​0​l​0​e​1​rs0l0e1rin the edge sublayer ands​0​l​0​e​1​as0l0e1ain the node sublayer\.

Our analysis focuses mainly on sample0, which both GM\-0and T\-0solve correctly \(Figure[7](https://arxiv.org/html/2608.06834#S3.F7)\)\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-construct-constraints.png)Figure 8:GM\-0constructs edges representing constraint regions in early layers\.In early layers, GM\-0builds Sudoku geometry from the input edges through a process resembling how convolutional neural networks compose elementary filters into more complex shapes[15](https://arxiv.org/html/2608.06834#bib.bib15)\(Figure[8](https://arxiv.org/html/2608.06834#S3.F8)\)\. For the examined noder​2​c​2r2c2, referral head22in layer0uses the up, self, and down edges to construct an edge factor corresponding to a vertical3×13\\times 1rectangle containingr​1​c​2r1c2,r​2​c​2r2c2, andr​3​c​2r3c2\. It then uses this factor to fetch the left and right edges of the up and down neighbors\. The fetched corner cells complete the box region representing the box constraint ofr​2​c​2r2c2\. Across all1616samples, the model consistently uses referral headl​0​r​2l0r2and the resultingl​1​e​2l1e2edges for the box constraint\.

Similarly, the model obtains the1×91\\times 9rectangle corresponding to the row constraint region through thel​3​e​2l3e2edge, and the9×19\\times 1rectangle corresponding to the column constraint region through thel​3​e​4l3e4edge\. It does so by fetching edges from cells near the middle of the row or column, namelyr​2​c​5r2c5orr​5​c​2r5c2\(Figure[15](https://arxiv.org/html/2608.06834#A4.F15)\)\. These middle\-of\-row/column cells appear to serve as providers of the row and column regions for cells aligned with them \(Figure[16](https://arxiv.org/html/2608.06834#A4.F16)\)\. This is natural given their centrality: they are the only cells that can reach both ends of a row or column, and therefore provide the complete region by the third layer using edge traversal alone\.

They can do so by recursively reaching for the farthest available cells in each direction, in coordination with other nodes\. At layer0, the11\-distance neighbors are provided as input\. At layer11, the22\-distance neighbors are fetched as the11\-distance neighbors’11\-distance neighbors\. At layer22, the44\-distance neighbors are fetched as the22\-distance neighbors’22\-distance neighbors\. Although the model could in principle memorize each full row or column region, the observed pattern is consistent with a simpler and more general11\-22\-44construction process shared across such nodes\. This exemplifies that, unlike standard GNNs, GMs have exponential rather than linear reach on the input graph given layer, and that explicit edge mechanisms may support stronger generalization under relational tasks\.

Two additional phenomena merit discussion\. First, the partial edges formed during this process, namely the horizontall​2​e​5l2e5edge and verticall​2​e​1l2e1edge, also appear in source cells that are not middle\-of\-row/column cells \(Figure[16](https://arxiv.org/html/2608.06834#A4.F16)\)\. Thus, they appear to be byproducts of the11\-22\-44construction applied broadly across nodes\. They may be fully useful only at middle\-of\-row/column cells, while remaining harmless or weakly beneficial elsewhere\.

Second, GM\-0appears to locate the middle\-of\-row/column target cells through different mechanisms for rows and columns\. For the row region, it seems to rely on thee​2e2factor; for the column region, it relies on a joint contribution from then​2n2edge and node factors \(Figure[15](https://arxiv.org/html/2608.06834#A4.F15)\)\. This holds even when the target cell is the source cell itself, as in the case ofr​2​c​5r2c5fetching its own edge for the row region\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-compose-regions.png)Figure 9:GM\-0composes existing regions through self\-referral to produce more complex regions\.The row and column regions are further composed into a cross region, represented by thel​4​e​0l4e0edge, which later combines with the box region to form a fused cross\-box region \(Figure[9](https://arxiv.org/html/2608.06834#S3.F9)\)\. As permitted by the inclusion of self\-edges, these regions are constructed and maintained through self\-referral: the model retrieves its own previous edges using a referral edge factor concentrated on the source cell itself\.

Unexpectedly, rather than preserving a self\-edge across layers as the basis for self\-referral, the model often adopts a more efficient strategy\. It composes existing non\-self edges and applies a high edge\-factor temperature to obtain a factor sharply concentrated on the source cell, thereby using one fewer edge slot while preserving the functional effect\. Across the1616samples, GM\-0generally does not maintain explicit self\-edges beyond the second layer\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-adapt-edges-part.png)Figure 10:GM\-0’s edges are input\-adaptive \(excerpt\)\.Apart from these fixed constructions, many of GM\-0’s edges are input\-adaptive, suggesting that the model also represents abstract relations \(Figure[10](https://arxiv.org/html/2608.06834#S3.F10); additional examples are shown in Figures[17](https://arxiv.org/html/2608.06834#A4.F17)and[18](https://arxiv.org/html/2608.06834#A4.F18)\)\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-divide-labor-part.png)Figure 11:GM\-0’s edge and node factors specialize in relational and content\-based roles, respectively \(excerpt\)\.While GM\-0’s edge factors seem to focus on relations, such as constraints, without encoding content preference, such as given or inferred digits, its node factors seem to match content without encoding relations\. This produces scattered patterns with no apparent geometric structure \(Figure[11](https://arxiv.org/html/2608.06834#S3.F11); additional examples are shown in Figure[19](https://arxiv.org/html/2608.06834#A4.F19)\)\. Often, the pattern of a node factor closely matches cells associated with a particular digit\.

This suggests that, as intended, the model uses a division of labor between the two experts, which are then combined through a PoE to specify the attention or referral target\. In most cases, the resulting attention/referral weights appear to receive substantial contributions from both factors, although either expert can sometimes dominate\. Overall, across GM models, edge factors contribute more strongly to target specification: their normalized entropies average0\.7410\.741, compared with0\.8910\.891for node factors, during attention, and0\.4020\.402, compared with0\.9450\.945, during referral \(Table[17](https://arxiv.org/html/2608.06834#A3.T17)\)\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-t-struggle-part.png)Figure 12:T\-0struggles to combine relational and content\-based roles in early\-to\-middle layers \(excerpt\)\.In contrast, T\-0’s node factors, which are the sole contributors of attention weights, appear to struggle to fulfill both roles \(Figure[12](https://arxiv.org/html/2608.06834#S3.F12); additional examples are shown in Figure[20](https://arxiv.org/html/2608.06834#A4.F20)\)\. In early layers, although their node factors can attend based on both content and simple relational information, such as row, column, and box constraint regions, they appear to struggle with relations beyond these simple groups, as well as with combining relational and content\-based information\. In the final layers, T\-0appears to combine relational and content\-based information in a primitive manner, although its attention remains less global and less interpretable than that of GM\-0\(Figure[21](https://arxiv.org/html/2608.06834#A4.F21)\)\.

Overall, attention in the Transformer models tends to be less sharp than in the GM models, with an average normalized entropy of0\.7360\.736, compared with0\.6520\.652for GM \(Table[17](https://arxiv.org/html/2608.06834#A3.T17)\)\. Since our Sudoku instances can usually be solved using rudimentary “Singles” techniques[16](https://arxiv.org/html/2608.06834#bib.bib16), these weaker attention patterns may still suffice for many examples, narrowing the performance gap with GM\. This highlights the need to test on tasks with a more complex relational structure\.

## 4Related works

Many Transformer and attention variants, such as Graphormer[17](https://arxiv.org/html/2608.06834#bib.bib17), incorporate notions of edges, relations, or graphs\. However, much of this work treats these as auxiliary to the node representation and mechanisms\. Graph Machine is instead situated among architectures that make edges a first\-class component of the model through explicit and dynamic edge states\. Some representative examples include Edge Transformers[18](https://arxiv.org/html/2608.06834#bib.bib18), Edge\-augmented Graph Transformers \(EGT\)[19](https://arxiv.org/html/2608.06834#bib.bib19), and relational attention[20](https://arxiv.org/html/2608.06834#bib.bib20)\. Among this category of architectures, GM differs along four dimensions\.

First, GM preserves the node\-centric Transformer pathway\. For example, our edge\-augmented attention allows joint contribution from both node factors and edge factors\. This makes GM a more direct expressivity generalization of Transformers and enables it to apply naturally beyond graph\-specific settings\.

Second, GM uses a factorized edge state\. Many edge\-state architectures maintain representations of shapen×n×den\\times n\\times d\_\{e\}, assigning a dense vector to every ordered pair of nodes\. GM instead represents edges using edge features of shapen×k×den\\times k\\times d\_\{e\}and edge addresses of shapen×k×nn\\times k\\times n\. This can be viewed as a generalization of the dense pairwise edge state: a densen×n×den\\times n\\times d\_\{e\}representation corresponds tonnedge slots per node with one\-hot addresses over destinations\. This separates what an edge contains from where it points, makes edge addresses directly interpretable as target mass, exposes edge degree as another controllable architectural inductive bias, and clarifies ways to reduce the cost of edge referral through address compression or sparsification\.

Third, GM differs in how edges shape information aggregation between nodes\. Rather than allowing edge representations to enter global content\-based matching in an unconstrained manner, GM uses node and edge features to aggregate edge addresses, which then directly modulate attention\. This biases the edge pathway toward relational traversal, while leaving global content\-based matching primarily to the node pathway\.

Fourth, GM differs in how edge states are updated\. Alternative mechanisms do not explicitly support edge referral or regularize against rich contributions from the node pathway\. Our referral mixes edge addresses across node and edge\-slot dimensions, such that new edge addresses are weighted combinations of previous edge addresses prior to sharpening\.

Sudoku has also been used as a benchmark for relational and iterative reasoning\. Recurrent Relational Networks[21](https://arxiv.org/html/2608.06834#bib.bib21)inject a strong task prior by constructing a graph in which each cell is connected to cells in the same row, column, and box\. More recent recurrent reasoning architectures, including the Hierarchical Reasoning Model[22](https://arxiv.org/html/2608.06834#bib.bib22)and Tiny Recursive Model[23](https://arxiv.org/html/2608.06834#bib.bib23), evaluate on harder Sudoku variants using Transformer modules applied recurrently for many steps\. Our use of Sudoku is different: we do not provide a full hand\-specified Sudoku constraint graph and only use a simple pass of the model\.

## 5Conclusion and limitations

We introduced Graph Machine, a neural architecture that represents graphs through both node features and dynamic edge representations\. By combining edge\-augmented attention and edge\-centric referral, the model can use soft graph structure during node updates while also revising that structure across layers\. In controlled Sudoku experiments, GM improves over standard\-scale Transformer baselines and remains competitive with enlarged Transformer variants, suggesting that explicit edge representations and mechanisms can improve relational reasoning\.

The main limitation of the current implementation is computational cost\. Dense and uncompressed edge addresses yield cubic time complexity and quadratic memory footprint\. Future work should aim for more efficient referral mechanisms, including with sparse edge addresses and operations\. Because GM factorizes edge features and edge addresses, it provides a natural interface for such approximations\.

Our evaluation is also limited to Sudoku\. Sudoku is useful as a controlled relational benchmark, but its latent graph structure is superficial, concrete, fixed, and regular\. This can favor Transformer baselines, since the latent structure can be relatively easily memorized via node embeddings\. GM may be more advantageous in settings where useful relations are latent, abstract, dynamic, and complex\.

Additionally, Sudoku places unusually high demands on dense working memory: humans often require help from verbal rehearsal and pencil marks to solve Sudoku puzzles\. Thus, although edge address distributions are broad in our experiment, it remains worth examining whether this pattern persists in tasks with more typical working\-memory demands\.

More broadly, our results motivate continued exploration of edge mechanisms as architectural inductive biases\. We hypothesize that such mechanisms may help improve abstraction, reasoning, and generalization, capabilities that remain important challenges for current machine learning systems\. Graph Machine offers an initial step in this direction\.

## References

- 1Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N\. Gomez, Lukasz Kaiser, and Illia Polosukhin\.Attention Is All You Need, August 2023\.
- 2Thomas N\. Kipf and Max Welling\.Semi\-Supervised Classification with Graph Convolutional Networks, February 2017\.
- 3Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio\.Graph Attention Networks, February 2018\.
- 4Tom B\. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert\-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M\. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei\.Language Models are Few\-Shot Learners, July 2020\.
- 5Yubo Li, Lu Zhang, Tianchong Jiang, Ramayya Krishnan, and Rema Padman\.The Model Says Walk: How Surface Heuristics Override Implicit Constraints in LLM Reasoning, April 2026\.
- 6Kévin \(@knowmadd@mastodon\.world\)\.https://mastodon\.world/@knowmadd/116072773118828295, February 2026\.
- 7Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka\.How Powerful are Graph Neural Networks?, February 2019\.
- 8Uri Alon and Eran Yahav\.On the Bottleneck of Graph Neural Networks and its Practical Implications, March 2021\.
- 9William B\. Johnson and Joram Lindenstrauss\.Extensions of Lipschitz mappings into a Hilbert space\.In Richard Beals, Anatole Beck, Alexandra Bellow, and Arshag Hajian, editors,Contemporary Mathematics, volume 26, pages 189–206\. American Mathematical Society, Providence, Rhode Island, 1984\.
- 10Geoffrey E\. Hinton\.Training Products of Experts by Minimizing Contrastive Divergence\.Neural Computation, 14\(8\):1771–1800, August 2002\.
- 11Tri Dao, Daniel Y\. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré\.FlashAttention: Fast and Memory\-Efficient Exact Attention with IO\-Awareness, June 2022\.
- 123 million Sudoku puzzles with ratings\.https://www\.kaggle\.com/datasets/radcliffe/3\-million\-sudoku\-puzzles\-with\-ratings\.
- 13Blagovest Dachev\.Dachev/sudoku, March 2026\.
- 14Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu\.RoFormer: Enhanced transformer with rotary position embedding\.CoRR, abs/2104\.09864, 2021\.
- 15Matthew D\. Zeiler and Rob Fergus\.Visualizing and Understanding Convolutional Networks\.In David Fleet, Tomas Pajdla, Bernt Schiele, and Tinne Tuytelaars, editors,Computer Vision – ECCV 2014, pages 818–833, Cham, 2014\. Springer International Publishing\.
- 16HoDoKu: Solving Techniques \- Singles \(Hidden Single, Naked Single, Full House\)\.https://hodoku\.sourceforge\.net/en/tech\_singles\.php\.
- 17Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie\-Yan Liu\.Do Transformers Really Perform Bad for Graph Representation?, November 2021\.
- 18Leon Bergen, Timothy J\. O’Donnell, and Dzmitry Bahdanau\.Systematic Generalization with Edge Transformers\.InAdvances in Neural Information Processing Systems, November 2021\.
- 19Md Shamim Hussain, Mohammed J\. Zaki, and Dharmashankar Subramanian\.Global Self\-Attention as a Replacement for Graph Convolution\.InProceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 655–665, August 2022\.
- 20Cameron Diao and Ricky Loynd\.Relational Attention: Generalizing Transformers for Graph\-Structured Tasks, March 2023\.
- 21Rasmus Berg Palm, Ulrich Paquet, and Ole Winther\.Recurrent Relational Networks, November 2018\.
- 22Guan Wang, Jin Li, Yuhao Sun, Xing Chen, Changling Liu, Yue Wu, Meng Lu, Sen Song, and Yasin Abbasi Yadkori\.Hierarchical Reasoning Model, August 2025\.
- 23Alexia Jolicoeur\-Martineau\.Less is More: Recursive Reasoning with Tiny Networks, October 2025\.
- 24Biao Zhang and Rico Sennrich\.Root Mean Square Layer Normalization, October 2019\.
- 25Noam Shazeer\.GLU Variants Improve Transformer, February 2020\.
- 26Diederik P\. Kingma and Jimmy Ba\.Adam: A Method for Stochastic Optimization, January 2017\.

## Appendix ASimplified codes

### A\.1Simplified code for edge\-augmented attention

defedge\_augmented\_attention\(n1\_queries,e1\_keys,e1\_addresses,n2\_keys,n2\_values\):

"""

Args:

n1\_queries\(b,h,n,d\_ak\):queriesprojectedfromn1’snodefeatures\.

edge\_factor\_temps\(b,h,n\):tempsprojectedfromn1’snodefeatures\.

node\_factor\_temps\(b,h,n\):tempsprojectedfromn1’snodefeatures\.

e1\_keys\(b,h,n,K,d\_ak\):keysprojectedfrome1’sedgefeatures\.

e1\_addresses\(b,n,K,n\):e1’sedgeaddresses\.

n2\_keys\(b,h,n,d\_ak\):keysprojectedfromn2’snodefeatures\.

n2\_values\(b,h,n,d\_av\):valuesprojectedfromn2’snodefeatures\.

Rets:

outs\(b,h,n,d\_av\):attentionoutputsforn1’snodefeatures\.

Einsumsymbols:

\[b,h,n,k,m,a/d\]

"""

e1\_scores=einsum\("bhnd,bhnkd\-\>bhnk",n1\_queries,e1\_keys\)/d\_ak\*\*0\.5

e1\_weights=softmax\(e1\_scores,dim=\-1\)

edge\_raw\_factors=log\(einsum\("bhnk,bnkm\-\>bhnm",e1\_weights,e1\_addresses\)\)

node\_raw\_factors=einsum\("bhnd,bhmd\-\>bhnm",n1\_queries,n2\_keys\)/d\_ak\*\*0\.5

edge\_factors=edge\_raw\_factors\*edge\_factor\_temps\.unsqueeze\(\-1\)

node\_factors=node\_raw\_factors\*node\_factor\_temps\.unsqueeze\(\-1\)

logits=edge\_factors\+node\_factors

weights=softmax\(logits,dim=\-1\)

outs=einsum\("bhnm,bhmd\-\>bhnd",weights,n2\_values\)

returnouts

Figure 13:Simplified code for edge\-augmented attention\.
### A\.2Simplified code for edge\-centric referral

defedge\_centric\_referral\(n1\_queries,\.\.\.,e2\_addresses\):

"""

Args:

n1\_queries\(b,k,n,d\_rk\):queriesprojectedfromn1’snodefeatures\.

n2\_edge\_factor\_temps\(b,k,n\):tempsprojectedfromn1’snodefeatures\.

n2\_node\_factor\_temps\(b,k,n\):tempsprojectedfromn1’snodefeatures\.

e2\_factor\_temps\(b,k,n\):tempsprojectedfromn1’snodefeatures\.

e1\_keys\(b,k,n,k,d\_rk\):keysprojectedfrome1’sedgefeatures\.

e1\_values\(b,k,n,k,d\_e\):valuesprojectedfrome1’sedgefeatures\.

e1\_addresses\(b,n,k,n\):e1’sedgeaddresses\.

n2\_keys\(b,k,n,d\_rk\):keysprojectedfromn2’snodefeatures\.

n2\_values\(b,k,n,d\_e\):valuesprojectedfromn2’snodefeatures\.

e2\_keys\(b,k,n,k,d\_rk\):keysprojectedfrome2’sedgefeatures\.

e2\_values\(b,k,n,k,d\_e\):valuesprojectedfrome2’sedgefeatures\.

e2\_addresses\(b,n,k,n\):e2’sedgeaddresses\.

Rets:

feature\_outs\(b,k,n,d\_e\):referraloutputsfore1’sedgefeatures\.

address\_outs\(b,k,n,n\):referraloutputsfore1’sedgeaddresses\.

Einsumsymbols:

\[b,h,n,k,m,l,a/d\]

"""

\*\_,d\_rk=n1\_queries\.shape

e1\_scores=einsum\("bhnd,bhnkd\-\>bhnk",n1\_queries,e1\_keys\)/d\_rk\*\*0\.5

e1\_weights=softmax\(e1\_scores,dim=\-1\)

n2\_edge\_raw\_factors=log\(einsum\("bhnk,bnka\-\>bhna",e1\_weights,e1\_addresses\)\)

n2\_node\_raw\_factors=einsum\("bhnd,bhmd\-\>bhnm",n1\_queries,n2\_keys\)/d\_rk\*\*0\.5

e2\_raw\_factors=einsum\("bhnd,bhmld\-\>bhnml",n1\_queries,e2\_keys\)/d\_rk\*\*0\.5

n2\_edge\_factors=n2\_edge\_raw\_factors\*n2\_edge\_factor\_temps\.unsqueeze\(\-1\)

n2\_node\_factors=n2\_node\_raw\_factors\*n2\_node\_factor\_temps\.unsqueeze\(\-1\)

e2\_factors=e2\_raw\_factors\*e2\_factor\_temps\.unsqueeze\(\-1\)\.unsqueeze\(\-1\)

n2\_factors=n2\_edge\_factors\+n2\_node\_factors

e2\_logits=n2\_factors\.unsqueeze\(\-1\)\+e2\_factors

e2\_weights=softmax\(dim=\(\-1,\-2\)\)

n2\_weights=e2\_weights\.sum\(dim=\-1\)

e1\_feature\_outs=einsum\("bhnk,bhnkd\-\>bhnd",e1\_weights,e1\_values\)

n2\_feature\_outs=einsum\("bhnm,bhmd\-\>bhnd",n2\_weights,n2\_values\)

e2\_feature\_outs=einsum\("bhnml,bhmld\-\>bhnd",e2\_weights,e2\_values\)

feature\_outs=e1\_feature\_outs\+n2\_feature\_outs\+e2\_feature\_outs

address\_outs=einsum\("bhnml,bmla\-\>bhna",e2\_weights,e2\_addresses\)

returnfeature\_outs,address\_outs

Figure 14:Simplified code for edge\-centric referral\.

## Appendix BDefault configurations

For our default configurations, we follow common Transformer design choices, using RMSNorm[24](https://arxiv.org/html/2608.06834#bib.bib24), no bias, and SwiGLU[25](https://arxiv.org/html/2608.06834#bib.bib25)activation function\. Both the feed\-forward networks in the edge and node sublayers have an expansion factor of44\. We use as our temperature function a horizontally shifted softplus that maps0to11\.

Training uses PyTorch automatic mixed precision viatorch\.autocast\. Adam optimizer[26](https://arxiv.org/html/2608.06834#bib.bib26)is used with\(β1,β2\)\(\\beta\_\{1\},\\beta\_\{2\}\)of\(0\.9,0\.95\)\(0\.9,0\.95\), no weight decay, a1%1\\%warmup, a peak learning rate of10−310^\{\-3\}, and a cosine schedule down to10%10\\%\. Gradient clipping with max grad norm of1\.01\.0is applied\.

Train\-eval\-test split on the dataset is redrawn for each run with proportions0\.90\.9\-0\.050\.05\-0\.050\.05\.100​k100ksteps are performed with batch size of6464each, corresponding to2\.372\.37epochs over the train set; each epoch reshuffles the dataset\. Testing is done with the entirety of the test split at the end of each run\.

## Appendix CAdditional experiments and statistics

#### Transformer \- project input edges

We test a Transformer condition where we project the initial edge embeddings and edge addresses to the node feature space in the input layer\. We do so by

1. 1\.perform weighted sum of the learned node embeddings using the initial edge addresses as weights for each edge,
2. 2\.concatenate with the edge embeddings for each edge,
3. 3\.pass through an MLP for each edge,
4. 4\.and add to the node embeddings\.

Results \(Table[9](https://arxiv.org/html/2608.06834#A3.T9)\) suggest this procedure is not advantageous in our case\.

Table 9:Transformer \- project input edges\.
#### Transformer\-sin\-PE\-2x \- factor temperatures

On top of the Transformer\-sin\-PE\-2x condition, we test the effect of removing factor temperatures\. Results \(Table[10](https://arxiv.org/html/2608.06834#A3.T10)\) suggest that this seems to slightly improve performance\.

Table 10:Transformer\-sin\-PE\-2x \- factor temperatures\.
#### Transformer\-sin\-PE\-2x \- entropy losses

On top of the Transformer\-sin\-PE\-2x condition, we test the effect of removing entropy losses\. Results \(Table[10](https://arxiv.org/html/2608.06834#A3.T10)\) suggest that this does not seem to materially alter performance\.

Table 11:Transformer\-sin\-PE\-2x \- entropy losses\.
#### Transformer \- positional encodings

We test a range of22D positional encoding options, including RoPE of various bases, sinusoidal PE of various bases, as well as a learned row\-column factorized PE, where a node inherits the sum of the learned embeddings of its row and column\. Our experiments suggest that sinusoidal PE with a base of1000010000performs the best among our conditions \(Table[10](https://arxiv.org/html/2608.06834#A3.T10)\)\. Among RoPE conditions, base1010is best, which we label as Transformer\-RoPE to offer comparison with GM\-RoPE in Section[3\.3](https://arxiv.org/html/2608.06834#S3.SS3)\.

Table 12:Transformer \- positional encodings\.
#### GM \- edge address space

We test a variant of GM where we let edge addresses represent logits directly, by removing the conversion to and from weight space when logits are needed\. During input, to convert the input edge\-address weights to logits, we simply multiply the weights by a chosen factor of55\. As a result, the uniform weights of the empty edges convert to uniform logits and thus distribution, and edge addresses with a single\-cell support convert to a distribution where the mass on the intended cell ise5≈148\.41e^\{5\}\\approx 148\.41larger than that of any other cell\. Results \(Table[13](https://arxiv.org/html/2608.06834#A3.T13)\) indicate that this reduces the performance to around Transformer level\.

Table 13:Edge address space\.
#### GM \- address sparsity

We conduct preliminary experiments on address sparsity by varying sparsityss, simulated using our dense algorithm and representation by masking non\-top\-ssentries to0along the last dimension of edge addresses\. We optionally use a Gumbel\-top\-sstemperatureτ\\taufor our selection by adding Gumbel noise to our address logits before selecting the topssof the original logits and renormalizing\. In our implementation, we do not further sparsify the edge factors or mask the attention/referral weights to maintain the top\-sssparsity\. We find that sparsification yields promising results, where a sparsity of1616yeilds similar performance to Transformer; however, using Gumbel\-top\-sswithout annealing seems to yield no benefit \(Table[14](https://arxiv.org/html/2608.06834#A3.T14)\)\. What remains future work includes enforcing sparsity on attention/referral weights for potentially significant efficiency benefits, and where we annealssand/orτ\\taufor potentially improved optimization\.

Table 14:GM \- address sparsity\.Base configurationssτ\\tauAccuracy \(%\\%\)GM0\(Transformer\)0\.00\.071\.771\.7110\.00\.036\.236\.2440\.00\.050\.950\.916160\.00\.070\.370\.38181\(GM\)0\.00\.086\.086\.0440\.010\.0127\.527\.5440\.10\.130\.430\.4441\.01\.053\.653\.6
#### GM and Transformer \- some statistics

We offer some statistics on the GM and Transformer conditions\. Factor entropies are logged after applying the factor temperatures\. In\-addresses are the input addresses after sharpening, while out\-addresses are the output addresses after the referral process\.

Table 15:Mean sharpener temperatures\.Table 16:Mean factor temperatures\.Table 17:Factor mean normalized entropies\.Table 18:Address mean normalized entropies\.

## Appendix DAdditional mechanistic analysis figures

This section collects some additional mechanistic analysis figures\.

![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-target-middle.png)Figure 15:GM\-0uses different mechanisms to targetr​2​c​5r2c5andr​5​c​2r5c2, consistently across source cells\.![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-perform-algorithm.png)Figure 16:GM\-0applies the11\-22\-44construction broadly, including in cells that are not middle\-of\-row/column\.![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-adapt-edges-early.png)Figure 17:GM\-0’s edges are input\-adaptive \(early\-layer\)\.![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-adapt-edges-late.png)Figure 18:GM\-0’s edges are input\-adaptive \(late\-layer\)\.![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-divide-labor.png)Figure 19:GM\-0’s edge and node factors specialize in relational and content\-based roles, respectively\.![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-t-struggle.png)Figure 20:T\-0struggles to combine relational and content\-based roles in early\-to\-middle layers\.![Refer to caption](https://arxiv.org/html/2608.06834v1/figures/mechanistic-gm-t-differ.png)Figure 21:GM\-0and T\-0differ in late\-layer attention weights: GM\-0’s attention appears more global and more interpretable\.