TAG-DLM: Diffusion Language Models for Text-Attributed Graph Learning

arXiv cs.CL Papers

Summary

TAG-DLM unifies textual reasoning and graph message passing within a masked diffusion language model, enabling joint reasoning over text and graph topology for node classification and link prediction tasks.

arXiv:2606.31166v1 Announce Type: new Abstract: Text-attributed graphs (TAGs), where each node carries a natural language description, require models to jointly reason over text and graph topology. Existing approaches often handle the two modalities separately: graph neural networks operate on shallow text features, while hybrids of LLMs and graphs use the language model mainly as a text encoder and delegate structure learning to a separate graph module. We propose method that unifies textual reasoning and graph message passing within a masked diffusion language model, a language model with bidirectional attention and generative decoding. For each graph instance, method linearises a sampled local neighbourhood into a token sequence and injects graph structure through a topology attention mask, which realises message passing over the graph. Because the diffusion language model can both interpret and generate text, the method adapts to different tasks simply by changing the prompt, supporting node classification, link prediction, and cross-dataset transfer with no target-specific fine-tuning. Experiments show that method outperforms graph neural networks, graph transformers, and LLM-based baselines on all three TAG benchmarks across two tasks, improving over the strongest baseline by up to 3.9 points.
Original Article
View Cached Full Text

Cached at: 07/01/26, 05:33 AM

# Diffusion Language Models for Text-Attributed Graph Learning
Source: [https://arxiv.org/html/2606.31166](https://arxiv.org/html/2606.31166)
Lingjie Chen UIUC lingjie7@illinois\.edu &Yuanchen Bei∗ UIUC bei4@illinois\.edu &Haobo Xu∗ UIUC haoboxu@illinois\.edu Yanjun Zhao UIUC yanjunzh@illinois\.edu &Yuzhong Chen VISA yuzchen@visa\.edu &Hanghang Tong UIUC htong@illinois\.edu

###### Abstract

Text\-attributed graphs \(TAGs\), where each node carries a natural language description, require models to jointly reason over text and graph topology\. Existing approaches often handle the two modalities separately: graph neural networks operate on shallow text features, while hybrids of LLMs and graphs use the language model mainly as a text encoder and delegate structure learning to a separate graph module\. We proposeTAG\-DLM, which unifies textual reasoning and graph message passing within a masked diffusion language model, a language model with bidirectional attention and generative decoding\. For each graph instance,TAG\-DLMlinearises a sampled local neighbourhood into a token sequence and injects graph structure through atopology attention mask, which realises message passing over the graph\. Because the diffusion language model can both interpret and generate text,TAG\-DLMadapts to different tasks simply by changing the prompt, supporting node classification, link prediction, and cross\-dataset transfer with no target\-specific fine\-tuning\. Experiments show thatTAG\-DLMoutperforms graph neural networks, graph transformers, and LLM\-based baselines on all three TAG benchmarks across two tasks, improving over the strongest baseline by up to 3\.9 points\.

TAG\-DLM: Diffusion Language Models for Text\-Attributed Graph Learning

Lingjie Chen††thanks:Equal contribution\.UIUClingjie7@illinois\.eduYuanchen Bei∗UIUCbei4@illinois\.eduHaobo Xu∗UIUChaoboxu@illinois\.edu

Yanjun ZhaoUIUCyanjunzh@illinois\.eduYuzhong ChenVISAyuzchen@visa\.eduHanghang TongUIUChtong@illinois\.edu

## 1Introduction

Graphs whose nodes carry natural language descriptions arise across a wide range of real\-world systems, such as citation networks in which each paper is represented by its title and abstract, biomedical literature graphs whose nodes describe research topics or articles, and knowledge bases in which entities are described by free\-form passagesYanet al\.\([2023](https://arxiv.org/html/2606.31166#bib.bib23)\); Wanget al\.\([2025](https://arxiv.org/html/2606.31166#bib.bib24)\)\. A graph of this kind is atext\-attributed graph\(TAG\): every node is annotated with a piece of text, and the edges represent relations between nodes\. TAG learning asks a model to make predictions over such graphs from both the node text and the graph topologyWanget al\.\([2024](https://arxiv.org/html/2606.31166#bib.bib25)\)\. We study two standard tasks: in node classification \(NC\), the model assigns a label to a target node; in link prediction \(LP\), it decides whether an edge exists between a candidate pair\. In both cases the answer is specified by the prompt and must be inferred from two complementary sources of evidence, the textual content of the local neighbourhood and its connectivity\. Solving these tasks well therefore requires a model that can simultaneously understand natural language and reason over graph structure\.

![Refer to caption](https://arxiv.org/html/2606.31166v1/x1.png)Figure 1:Comparison betweenTAG\-DLMwith existing text\-attributed graph learning pipelines\.Existing approaches usually place language understanding and graph reasoning in separate components, which limits how the two sources of evidence interact\. Graph Neural Networks \(GNNs\), such as GCN\(Kipf and Welling,[2017](https://arxiv.org/html/2606.31166#bib.bib3)\)and GAT\(Veličkovićet al\.,[2018](https://arxiv.org/html/2606.31166#bib.bib5)\)model graph structure effectively but represent node text as shallow bag of words or truncated embeddings, discarding most of the semantic content\. More recent large language model \(LLM\) hybrids\(Chenet al\.,[2024](https://arxiv.org/html/2606.31166#bib.bib12); Tanget al\.,[2024](https://arxiv.org/html/2606.31166#bib.bib26)\)treat the language model as a frozen text encoder: node texts are embedded independently, and a separate GNN is then attached to propagate the resulting embeddings over the graph\. This pipeline keeps the two modalities structurally separate\. The language model never conditions on graph topology, and the graph module operates on these precomputed text embeddings rather than the raw node text\. As a result, the contextual understanding of the LLM cannot interact with graph\-level reasoning: each side sees only a degraded view of the other\. A further limitation is task specificity: because the task is read out by a fixed output head, these methods must be retrained whenever the task changes\.

First, MDLMs use fully bidirectional self\-attention, so information flows in all directions within a single forward pass\. This suits graph prompts, where the linearised order of neighbour nodes is arbitrary and the target and its neighbours should exchange evidence regardless of that order, rather than along a fixed left\-to\-right direction\. Second, MDLMs are generative: they fill in masked spans in a sequence, so the answer space can be specified by the prompt rather than by an output head specific to a task\. These two properties motivate the key insight behind our work:*graph neighbourhood relationships can be encoded directly as an attention mask*\. By restricting which tokens may attend to which, we embed graph structure into the denoising process without a separate graph module or any change to the backbone\. And because the answer is simply a filled\-in prompt position, the same formulation expresses both NC and LP through prompt variation, with no task\-specific graph head\.

Building on this, we introduceTAG\-DLM\(Text\-AttributedGraphDiffusionLanguageModel\), a framework that casts TAG learning as a masked infill problem\. For a target graph instance, such as a node to classify or a candidate edge to score,TAG\-DLMlinearises the sampled local context into a single token sequence, fills a masked answer position with an option defined by the prompt, and attaches a binarytopology attention mask\. The mask enforces structured attention over the local graph context, so a single forward pass through the MDLM backbone simultaneously reads textual evidence and propagates information along the graph, with no separate GNN component required\.

Our contributions are summarized as follows:

- •Topology attention mask as graph operator\.Graph structure is injected into an MDLM purely through a binary attention mask for each sample, eliminating the need for a separate GNN encoder or any architectural modification to the backbone\.
- •MDLMs for TAG learning\.TAG\-DLMadapts the masked diffusion paradigm to graph\-structured data, combining bidirectional language understanding with graph topology in a single model\.
- •Theoretical grounding\.We prove that a single denoising forward pass with the topology attention mask performs attention\-based message passing on the graph induced by the mask\.
- •Unified graph learning evaluation\.SinceTAG\-DLMis generative, node classification, link prediction, and cross\-dataset transfer can be expressed through answer options defined by the prompt rather than output heads specific to a task\.

## 2Preliminaries

### 2\.1Text Attributed Graphs

A*text attributed graph*\(TAG\) is a triple𝒢=\(𝒱,ℰ,𝒳\)\\mathcal\{G\}=\(\\mathcal\{V\},\\mathcal\{E\},\\mathcal\{X\}\)in which𝒱\\mathcal\{V\}is the node set,ℰ⊆𝒱×𝒱\\mathcal\{E\}\\subseteq\\mathcal\{V\}\\times\\mathcal\{V\}is the edge set, and𝒳=\{xv\}v∈𝒱\\mathcal\{X\}=\\\{x\_\{v\}\\\}\_\{v\\in\\mathcal\{V\}\}assigns to every nodevva piece of natural language textxvx\_\{v\}\. We consider graph prediction tasks defined by prompts over such graphs\. In*node classification*\(NC\), each node carries a discrete labelyv∈𝒴y\_\{v\}\\in\\mathcal\{Y\}from a finite, language expressible label set𝒴\\mathcal\{Y\}\(e\.g\. “Neural Networks”, “cs\.CV”\), the labels of a subset𝒱train⊆𝒱\\mathcal\{V\}\_\{\\mathrm\{train\}\}\\subseteq\\mathcal\{V\}are observed, and the goal is to predictyvy\_\{v\}for held out nodes\. In*link prediction*\(LP\), the model receives a candidate pair\(u,v\)\(u,v\)and predicts whether the edge belongs to the target graph under a binary answer set\. Both tasks are cast as selecting an answer from a candidate set specified in the prompt\. The ego graph of a target nodevvwithinkkhops is denoted𝒩k​\(v\)⊆𝒱\\mathcal\{N\}\_\{k\}\(v\)\\subseteq\\mathcal\{V\}and containsvvtogether with the nodes reachable fromvvin at mostkkedge steps; for pairwise tasks, we use the sampled local context around the candidate endpoints\.

### 2\.2Masked Diffusion Language Models

TAG\-DLMuses a*masked diffusion language model*\(MDLM\) as its backbone\(Sahooet al\.,[2024](https://arxiv.org/html/2606.31166#bib.bib13); Nieet al\.,[2025](https://arxiv.org/html/2606.31166#bib.bib14)\)\. LetVVdenote the model vocabulary augmented with a special mask token\[M\]and letx=\(x1,…,xL\)∈VLx=\(x\_\{1\},\\ldots,x\_\{L\}\)\\in V^\{L\}be a clean token sequence\. The*forward*process corruptsxxover a continuous timet∈\[0,1\]t\\in\[0,1\]by independently replacing each token with\[M\]with probabilityβ​\(t\)\\beta\(t\), whereβ\\betais a monotone schedule withβ​\(0\)=0\\beta\(0\)=0andβ​\(1\)=1\\beta\(1\)=1\. Writingx~t\\widetilde\{x\}\_\{t\}for the sequence at timett, the forward kernel factorises across positions:

q​\(x~t∣x\)=∏i=1Lq​\(x~t,i∣xi\)\.q\(\\widetilde\{x\}\_\{t\}\\mid x\)=\\prod\_\{i=1\}^\{L\}q\\bigl\(\\widetilde\{x\}\_\{t,i\}\\mid x\_\{i\}\\bigr\)\.\(1\)
A denoising networkfθf\_\{\\theta\}predicts each clean token from a corrupted sequence using*bidirectional*self\-attention over all positions\. Training maximises a Rao Blackwellised evidence lower boundLiuet al\.\([2019a](https://arxiv.org/html/2606.31166#bib.bib27)\)that reduces, up to a reweightingw​\(t\)w\(t\)that depends on time, to a masked cross entropy on the corrupted positions:

ℒMDLM​\(θ\)=𝔼t,x,x~t​\[w​\(t\)​∑i:x~t,i=\[M\]−log⁡fθ​\(xi∣x~t,t\)\]\.\\mathcal\{L\}\_\{\\mathrm\{MDLM\}\}\(\\theta\)=\\mathbb\{E\}\_\{t,x,\\widetilde\{x\}\_\{t\}\}\\left\[w\(t\)\\sum\_\{i:\\widetilde\{x\}\_\{t,i\}=\\texttt\{\[M\]\}\}\-\\log f\_\{\\theta\}\(x\_\{i\}\\mid\\widetilde\{x\}\_\{t\},t\)\\right\]\.

\(2\)At inference time, generation proceeds by an iterative denoising loop that starts from a fully or partially masked sequence and progressively unmasks tokens by sampling fromfθf\_\{\\theta\}\. Crucially for our analysis in Section[4](https://arxiv.org/html/2606.31166#S4),fθf\_\{\\theta\}uses attention that is not causal, so the only structural restriction on information flow within a layer is whatever attention mask is supplied externally\.

### 2\.3Attention\-Based Message Passing

Many graph layers can be written as message passing with learned attention over a node’s neighbourhood\. Given node representationshuℓh\_\{u\}^\{\\ell\}at layerℓ\\ell, an attention\-based message passing layer has the form

muℓ\\displaystyle m\_\{u\}^\{\\ell\}=∑w∈𝒩​\(u\)αu​wℓ​ϕℓ​\(hwℓ\),\\displaystyle=\\sum\_\{w\\in\\mathcal\{N\}\(u\)\}\\alpha\_\{uw\}^\{\\ell\}\\,\\phi\_\{\\ell\}\(h\_\{w\}^\{\\ell\}\),\(3\)αu​wℓ\\displaystyle\\alpha\_\{uw\}^\{\\ell\}=exp⁡sℓ​\(huℓ,hwℓ\)∑r∈𝒩​\(u\)exp⁡sℓ​\(huℓ,hrℓ\),\\displaystyle=\\frac\{\\exp s\_\{\\ell\}\(h\_\{u\}^\{\\ell\},h\_\{w\}^\{\\ell\}\)\}\{\\sum\_\{r\\in\\mathcal\{N\}\(u\)\}\\exp s\_\{\\ell\}\(h\_\{u\}^\{\\ell\},h\_\{r\}^\{\\ell\}\)\},\(4\)huℓ\+1\\displaystyle h\_\{u\}^\{\\ell\+1\}=Uℓ​\(huℓ,muℓ\),\\displaystyle=U\_\{\\ell\}\(h\_\{u\}^\{\\ell\},m\_\{u\}^\{\\ell\}\),\(5\)wheresℓs\_\{\\ell\}is an attention score,ϕℓ\\phi\_\{\\ell\}transforms neighbour features, andUℓU\_\{\\ell\}updates the node state\. GAT\(Veličkovićet al\.,[2018](https://arxiv.org/html/2606.31166#bib.bib5)\)is a standard instance of this family: it uses an additive attention scoresℓ​\(hu,hw\)=LeakyReLU⁡\(a⊤​\[W​hu∥W​hw\]\)s\_\{\\ell\}\(h\_\{u\},h\_\{w\}\)=\\operatorname\{LeakyReLU\}\(a^\{\\top\}\[Wh\_\{u\}\\,\\\|\\,Wh\_\{w\}\]\)and a shared linear value transform\. Section[4](https://arxiv.org/html/2606.31166#S4)shows that topology masked denoising realizes this broader attention\-based message passing form, with GAT recovered as a special case\.

## 3Method

### 3\.1Overall Framework

TAG\-DLMcasts graph prediction as masked label infilling over a linearised sequence that contains the target instance and its sampled local context\. Graph structure is injected into self\-attention through a topology attention mask rather than through a separate GNN encoder\. Whereas existing hybrids of LLMs and graphs treat the language model as a frozen text encoder attached to a graph module,TAG\-DLMmakes the language model itself the graph operator: the only graph aware input is a binary attention mask for each sample, and the only trained parameters are LoRA adapters on the linear projections\. This design gives two useful properties\. First, the underlying MDLM weights are reused without architectural changes or additional graph modules\. Second, masked self\-attention performs message passing on the graph induced by the mask \(Section[4](https://arxiv.org/html/2606.31166#S4)\), so a single bidirectional forward pass through the MDLM reads textual evidence and propagates information along the graph\. As illustrated in Figure[2](https://arxiv.org/html/2606.31166#S3.F2),TAG\-DLMhas two core components: prompt construction from a sampled local graph context \(§[3\.2](https://arxiv.org/html/2606.31166#S3.SS2)\) and a topology attention mask that constrains information flow inside the MDLM \(§[3\.3](https://arxiv.org/html/2606.31166#S3.SS3)\)\.

![Refer to caption](https://arxiv.org/html/2606.31166v1/x2.png)Figure 2:Overview ofTAG\-DLM\. The ego graph of target nodevvwithinkkhops is linearised into a token sequenceSvS\_\{v\}with the answer position replaced by\[M\]\. A binary topology attention maskMvM\_\{v\}enforces a star\-shaped attention pattern:vvattends to all tokens; each neighbour attends only to itself andvv\. Both are fed into LLaDA\-8B with LoRA tuning, which predicts the answer defined by the prompt from the answer position logits\.
### 3\.2Prompt Construction

##### Subgraph Linearisation\.

For each target nodev∈𝒱v\\in\\mathcal\{V\}in a TAG𝒢=\(𝒱,ℰ,𝒳\)\\mathcal\{G\}=\(\\mathcal\{V\},\\mathcal\{E\},\\mathcal\{X\}\), we sample an ego graph withinkkhops,𝒩k​\(v\)=\{v\}∪\{u1,…,unv\}\\mathcal\{N\}\_\{k\}\(v\)=\\\{v\\\}\\cup\\\{u\_\{1\},\\ldots,u\_\{n\_\{v\}\}\\\}, by uniformly selecting at mostKKneighbours per hop \(defaultk=2k\{=\}2,K=10K\{=\}10\)\. The linearised sequence is

Sv=T​\(v\)​‖T​\(u1\)‖​⋯∥T​\(unv\),S\_\{v\}\\;=\\;T\(v\)\\;\\\|\\;T\(u\_\{1\}\)\\;\\\|\\;\\cdots\\;\\\|\\;T\(u\_\{n\_\{v\}\}\),\(6\)whereT​\(⋅\)T\(\\cdot\)denotes the template for each node and∥\\\|is sequence concatenation\. The target node always occupies the first segment so that its token span can be recovered directly\. For pairwise prediction, we use the same construction around the candidate endpoints and encode the candidate relation in the prompt\.

##### Multiple Choice Prompt\.

We use the*multiple choice digit*prompt format for both node classification and link prediction\. The answer position is a single\[M\]token, and each candidate answer is represented by a digit token\. For node classification, the digit options index the dataset label names\. For link prediction, the same scoring rule is used with binary options indicating whether the candidate edge exists\. Neighbour templates may include observed training labels inside square brackets, while neighbours without observed labels drop the bracket prefix\. We provide a concrete prompt example in Appendix[C](https://arxiv.org/html/2606.31166#A3)\.

### 3\.3Topology Attention Mask

A central design choice inTAG\-DLMis to expose graph adjacency only through attention\. The base MDLM weights remain frozen, and graph structure enters as an input\-dependent attention constraint rather than as a separate graph encoder\. For the tokenised sequenceSvS\_\{v\}, let\|Sv\|\|S\_\{v\}\|denote its number of tokens\. We attach to every sample a binary maskMv∈\{0,1\}\|Sv\|×\|Sv\|M\_\{v\}\\in\\\{0,1\\\}^\{\|S\_\{v\}\|\\times\|S\_\{v\}\|\}whose rows and columns correspond to token positions inSvS\_\{v\}\. Conceptually,MvM\_\{v\}is a token span level graph operator: each row specifies which target or context span may read from which other span during self\-attention\.

The default mask follows a star\-shaped attention pattern\. Tokens in the target span attend to every token inSvS\_\{v\}, so the target sees context from one tokkhops\. Tokens in a neighbour span attend to themselves and to the target span, but not to other neighbours, yielding star\-shaped message flow with the target as the root\. Padding positions are masked in both directions, so they neither send nor receive information\. This star mask is an intentional local context abstraction for target\-centred prediction\. The target instance acts as the unique aggregation point and reads evidence from all sampled nodes, which keeps the same prompt and mask interface for node classification and link prediction\. Suppressing direct communication among neighbours also reduces sensitivity to dataset\-specific ego graph wiring, so transfer across datasets depends less on idiosyncratic local topology\. Richer masks could expose more of the original ego graph structure, but we use the star mask as the default because it is simple, agnostic to the task, and stable across graph datasets\.

## 4Theoretical Analysis

We now show that, at any fixed diffusion time, one self\-attention layer of the topology masked MDLM denoiser implements an attention\-based message passing update over the neighbourhoods induced by the topology attention mask\. The argument explains why injecting the graph through attention alone, without a separate GNN encoder, is expressive enough to subsume standard graph attention, and makes explicit which node spans exchange information\. The full derivation, multihead extension, and the corollary recovering GAT appear in Appendix[B](https://arxiv.org/html/2606.31166#A2)\.

### 4\.1Fixed Time Denoising Operator

TheTAG\-DLMdenoising networkfθ​\(X~t,t,Mv\)f\_\{\\theta\}\(\\widetilde\{X\}\_\{t\},t,M\_\{v\}\)takes the corrupted token sequenceX~t\\widetilde\{X\}\_\{t\}at diffusion timett, the topology attention maskMvM\_\{v\}for the target, and predicts the clean tokens*in parallel*through stacked self\-attention layers\. Throughout this section, we fixttand analyse a single layer, separating the temporal denoising process from graph\-structured communication\. The reverse diffusion process evolves the corrupted sequence across timesteps, whereas graph information flow is induced within each denoising call by the topology attention mask over token spans\. Our claim is therefore local to a fixed time denoising operator: a masked self\-attention layer realizes spatial message passing on the graph induced by the mask, independently of the noising and denoising dynamics overtt\.

The crucial property of the backbone is that, unlike a causal language model whose attention mask forbids tokens from attending to later positions, the MDLM denoiser uses*bidirectional*self\-attention\. After the topology attention maskMvM\_\{v\}is applied, the*only*structural constraint on information flow within a layer isMvM\_\{v\}itself, rather than token order from left to right\. This is what allows the masked denoiser to be read as graph message passing\.

### 4\.2Neighbourhoods Induced by the Mask

Using the linearised sequence from Section[3\.2](https://arxiv.org/html/2606.31166#S3.SS2), letIuI\_\{u\}denote the contiguous token span of nodeuu\. The topology attention mask determines which node spans can exchange information\. We write𝒩M​\(u\)\\mathcal\{N\}\_\{M\}\(u\)for the set of nodes whose spans may be attended to by tokens inIuI\_\{u\}underMvM\_\{v\}\. For the star mask used inTAG\-DLM, the target node attends to all sampled nodes, while each node other than the target attends only to itself and the target\. Thus, the message passing structure analysed below is induced by the mask, rather than by the original TAG edges\.

### 4\.3Masked Denoising as Message Passing

Lethuℓ∈ℝdh\_\{u\}^\{\\ell\}\\in\\mathbb\{R\}^\{d\}denote the node level representation obtained by pooling the token hidden states over the spanIuI\_\{u\}at layerℓ\\ell, and letWQℓ,WKℓ,WVℓW\_\{Q\}^\{\\ell\},W\_\{K\}^\{\\ell\},W\_\{V\}^\{\\ell\}be the layer\-ℓ\\ellquery, key, and value projections\.

###### Proposition 1\(Topology masked denoising as message passing\)\.

At any fixed diffusion timett, one bidirectional self\-attention layer offθf\_\{\\theta\}restricted by the topology attention maskMvM\_\{v\}admits the form

huℓ\+1=Uℓ​\(huℓ,∑w∈𝒩M​\(u\)αu​wℓ​WVℓ​hwℓ\),h\_\{u\}^\{\\ell\+1\}=U\_\{\\ell\}\\\!\\left\(h\_\{u\}^\{\\ell\},\\;\\sum\_\{w\\in\\mathcal\{N\}\_\{M\}\(u\)\}\\alpha\_\{uw\}^\{\\ell\}\\,W\_\{V\}^\{\\ell\}h\_\{w\}^\{\\ell\}\\right\),\(7\)where the attention weight is

αu​wℓ=exp⁡\(sℓ​\(huℓ,hwℓ,t\)\)∑r∈𝒩M​\(u\)exp⁡\(sℓ​\(huℓ,hrℓ,t\)\),\\alpha\_\{uw\}^\{\\ell\}=\\frac\{\\exp\\\!\\bigl\(s\_\{\\ell\}\(h\_\{u\}^\{\\ell\},h\_\{w\}^\{\\ell\},t\)\\bigr\)\}\{\\sum\_\{r\\in\\mathcal\{N\}\_\{M\}\(u\)\}\\exp\\\!\\bigl\(s\_\{\\ell\}\(h\_\{u\}^\{\\ell\},h\_\{r\}^\{\\ell\},t\)\\bigr\)\},\(8\)sℓ​\(⋅\)s\_\{\\ell\}\(\\cdot\)is the attention score conditioned on the timestep, andUℓU\_\{\\ell\}is the layer update formed by the residual connection, LayerNorm, and the FFN\. Equation \([7](https://arxiv.org/html/2606.31166#S4.E7)\) is the canonical attention\-based message passing layer over the neighbourhoods induced by the mask\.

##### Proof Sketch\.

Masked self\-attention can be matched term by term to an attention\-based MPNN\. The topology attention mask zeros outαu​wℓ\\alpha\_\{uw\}^\{\\ell\}forw∉𝒩M​\(u\)w\\notin\\mathcal\{N\}\_\{M\}\(u\), so the softmax in Eq\. \([8](https://arxiv.org/html/2606.31166#S4.E8)\) reduces to a normalisation over the induced neighbourhood\. The value projectionWVℓ​hwℓW\_\{V\}^\{\\ell\}h\_\{w\}^\{\\ell\}defines the message content, the normalised weightαu​wℓ\\alpha\_\{uw\}^\{\\ell\}is the message coefficient, and the masked attention sum is the aggregation operator\. The residual connection together with the output projection, LayerNorm, and feed\-forward block constitutes the update functionUℓU\_\{\\ell\}\. Replacing token\-level attention with span\-pooled attention under the compatibility assumptions in Appendix[B](https://arxiv.org/html/2606.31166#A2)yields Eqs\. \([7](https://arxiv.org/html/2606.31166#S4.E7)\)\-\([8](https://arxiv.org/html/2606.31166#S4.E8)\)\. ∎

### 4\.4Implications

The proposition above has two consequences\. First, choosingsℓs\_\{\\ell\}to be additive attention recovers the GAT update exactly, so GAT is a special case ofTAG\-DLMrather than a separate design; the dot\-product score used byTAG\-DLMgeneralises this to a broader family of attention MPNNs\. Second, the full reverse diffusion procedure composes timestep conditioned message passing operators over the same neighbourhoods induced by the mask, which justifies treating the iterative infill mode ofTAG\-DLMas repeated graph reasoning rather than a separate inference mechanism\.

## 5Experiments

Table 1:Main results on node classification and link prediction across Cora, PubMed, and ogbn\-arxiv\. Values are test accuracy \(%\);TAG\-DLMachieves the best accuracy in all reported task\-dataset settings\. Best results are inbold, second\-best results areunderlined, and N/A indicates no comparable recorded result\.### 5\.1Experimental Setup

##### Datasets\.

We evaluateTAG\-DLMon three widely used TAG learning benchmarks: Cora and PubMed\(Yanget al\.,[2016](https://arxiv.org/html/2606.31166#bib.bib1)\), and ogbn\-arxiv\(Huet al\.,[2020](https://arxiv.org/html/2606.31166#bib.bib2)\)\. The benchmarks span two orders of magnitude in graph size, from a small citation graph to a much larger arXiv citation network, and differ in domain, label granularity, and label space size\. We use the preprocessing aligned with LLaGA\(Chenet al\.,[2024](https://arxiv.org/html/2606.31166#bib.bib12)\)for all three datasets and the public train/validation/test splits\. Full statistics and split sizes appear in Appendix[A](https://arxiv.org/html/2606.31166#A1)\.

##### Training\.

We trainTAG\-DLMon each dataset and task configuration by fine\-tuning the LLaDA\-8B\-Instruct backbone with LoRA adapters\(Huet al\.,[2022](https://arxiv.org/html/2606.31166#bib.bib17)\)\(rankr=64r\{=\}64, scalingα=64\\alpha\{=\}64, target modulesall\-linear\); the base MDLM weights stay frozen\. We linearise each target instance’s local graph context with at most1010uniformly sampled neighbours per hop and use the*multiple\-choice digit*prompt template \(Section[3\.2](https://arxiv.org/html/2606.31166#S3.SS2)\)\. The maximum sequence length is20482048tokens for Cora and PubMed and40964096tokens for ogbn\-arxiv\. We optimise with AdamW\(Loshchilov and Hutter,[2019](https://arxiv.org/html/2606.31166#bib.bib18)\)at learning rate5×10−55\\times 10^\{\-5\}for2020epochs, with batch size44per device and gradient accumulation44steps \(effective batch size128128\)\. Training runs on a single node of eight NVIDIA A100 80GB GPUs with DeepSpeed ZeRO\-2\(Rasleyet al\.,[2020](https://arxiv.org/html/2606.31166#bib.bib19)\)\. The full hyperparameter list is in Appendix[D](https://arxiv.org/html/2606.31166#A4)\.

##### Baselines\.

We compareTAG\-DLMagainst three families of baselines\. The first family is classical GNNs trained from scratch on shallow text features: GCN\(Kipf and Welling,[2017](https://arxiv.org/html/2606.31166#bib.bib3)\), GraphSAGE\(Hamiltonet al\.,[2017](https://arxiv.org/html/2606.31166#bib.bib4)\), GAT\(Veličkovićet al\.,[2018](https://arxiv.org/html/2606.31166#bib.bib5)\), GATv2\(Brodyet al\.,[2022](https://arxiv.org/html/2606.31166#bib.bib6)\), GIN, and MixHop\. The second family is graph transformers that fuse attention with graph structure but, like the GNNs, do not condition on raw text: GraphTransformer\(Dwivedi and Bresson,[2021](https://arxiv.org/html/2606.31166#bib.bib7)\), NodeFormer\(Wuet al\.,[2022](https://arxiv.org/html/2606.31166#bib.bib8)\), DIFFormer\(Wuet al\.,[2023a](https://arxiv.org/html/2606.31166#bib.bib9)\), and SGFormer\(Wuet al\.,[2023b](https://arxiv.org/html/2606.31166#bib.bib10)\)\. The third family is text\-aware language\-model baselines: RoBERTa\(Liuet al\.,[2019b](https://arxiv.org/html/2606.31166#bib.bib11)\), a text\-only encoder fine\-tuned on the linearised ego graph; LLaGA\(Chenet al\.,[2024](https://arxiv.org/html/2606.31166#bib.bib12)\), the current state of the art hybrid of LLMs and graphs; and a*frozen*LLaDA\-8B\-Instruct\(Nieet al\.,[2025](https://arxiv.org/html/2606.31166#bib.bib14)\), which shares the backbone ofTAG\-DLMbut receives no fine\-tuning and isolates the gain attributable to our training recipe\. For GNN and graph transformer baselines, we tune learning rate, hidden dimension, number of layers, and dropout on the validation split and report test accuracy from the best validation configuration\. For the language\-model baselines, we use the same neighbourhood sampling and prompt formatting recipe asTAG\-DLM\.

##### Evaluation protocol\.

We report accuracy on test splits for both NC and LP\. For NC, all neighbour labels from validation and test nodes are masked during training and evaluation, so label information cannot leak through the neighbourhood sampler\. For LP, we follow the same held\-out edge split and negative sampling protocol used by the corresponding graph baselines\. To prevent edge leakage, any positive candidate edge being scored is removed from the graph before constructing the local context, and held\-out validation/test edges are not exposed through neighbourhood sampling\. We score each candidate pair with binary prompt options\.TAG\-DLMuses a single forward pass scoring rule across tasks: each candidate answer is represented by a digit option, scored by the log\-likelihood of that digit token at the masked answer position, and the highest\-scoring candidate is selected\.

### 5\.2Main Results

Table[1](https://arxiv.org/html/2606.31166#S5.T1)reports accuracy on NC and LP across the three TAG benchmarks with comparable recorded results\. On NC,TAG\-DLMachieves the best accuracy on all three datasets, with the clearest gains on Cora and PubMed and a smaller but consistent improvement on ogbn\-arxiv\. The result is notable because these benchmarks differ substantially in scale and label structure: Cora is a small citation graph, PubMed is larger and comes from a biomedical citation domain, and ogbn\-arxiv is much larger with a finer\-grained label space\. Thus, the same topology masked MDLM remains competitive across both small and large graphs, and across label spaces\.

The advantage is more pronounced on LP, whereTAG\-DLMleads on every benchmark by a wider margin\. This supports our intuition that the topology\-aware denoising formulation is especially effective when the task directly depends on relational evidence\. Unlike methods that encode text first and then pass fixed vectors via a separate graph module,TAG\-DLMlets textual evidence and graph structure interact inside the same bidirectional denoising network, which is useful for both tasks\.

Table 2:Cross\-dataset transfer accuracy \(%\)\. Rows indicate the source dataset used for training, and columns indicate the target dataset used for evaluation\. Diagonal entries are in\-domain evaluations, while off\-diagonal entries measure direct transfer without target\-specific fine\-tuning\. Parentheses report absolute gains over Frozen LLaDA\-8B on the same target dataset\.
### 5\.3Cross Dataset Transfer

Table[2](https://arxiv.org/html/2606.31166#S5.T2)evaluates whether a topology\-aware checkpoint trained on one source graph transfers directly to unseen target graphs\. The diagonal entries provide the in\-domain reference point, while the entries outside the diagonal test a stricter setting: the target answer space is supplied only through the prompt, with no fine\-tuning specific to the target\. Among the benchmarked methods in Table[1](https://arxiv.org/html/2606.31166#S5.T1),TAG\-DLMis the only one that naturally supports this direct transfer setting defined by prompts, since it does not rely on an output head or graph decoder specific to a dataset\. Most GNN and graph transformer baselines in Table[1](https://arxiv.org/html/2606.31166#S5.T1)learn a classifier or decoder tied to the source dataset, so applying them to a target graph with different label names, class granularity, or edge distribution would require replacing or retraining that component\. That would turn the comparison into target adaptation rather than direct transfer\. Frozen LLaDA\-8B is prompt compatible and therefore serves as the reference in parentheses in Table[2](https://arxiv.org/html/2606.31166#S5.T2), but it does not include the graph adaptation learned byTAG\-DLM\.

The results show useful cross dataset generalisation, with the expected task dependence\. NC transfer is strongest between datasets with closer label semantics, such as Cora and PubMed, and is weaker when transferring into ogbn\-arxiv, whose finer\-grained label space and label distribution differ more substantially\. This makes the cross dataset setting challenging: the source and target graphs may vary in size by orders of magnitude, the class name inventories are not the same, and the citation domains differ between datasets\. Despite these shifts,TAG\-DLMstill produces meaningful predictions outside the source domain because the target answer space is expressed in language through the prompt rather than encoded in a fixed output head\.

LP is more stable across datasets, indicating that pairwise relational evidence transfers more robustly than dataset\-specific label semantics\. The gains in parentheses further separate prompting from adaptation: Frozen LLaDA\-8B is far below the trained model across both tasks, so the transfer behavior is not coming from prompt wording alone\. It depends on training the MDLM to use the linearised graph context and topology mask as task\-relevant evidence\. Overall,TAG\-DLMtransfers best when source and target prompts share label meaning, but the broader result is that the model can be applied across datasets without changing the architecture or adding a new dataset\-specific head\.

### 5\.4Ablation Studies

#### 5\.4\.1Effect of Topology Attention Mask

Table 3:Topology attention mask ablation accuracy \(%\)\. For each task and dataset, we compareTAG\-DLMwith and without the topology attention mask while keeping the same prompt format and backbone\.Table[3](https://arxiv.org/html/2606.31166#S5.T3)reports the topology\-mask ablation\.Removing the topology attention mask consistently weakens performance on both NC and LP, while the full topology\-aware model retains the best accuracy across all three benchmarks\. The effect is especially visible on LP, where the prediction depends directly on pairwise relational evidence\.

#### 5\.4\.2Effect of Attention Mask Type

Table 4:Attention mask type ablation accuracy \(%\)\. We compare the default star mask with an ego\-graph mask that exposes observed neighbour\-neighbour edges in the sampled local context\.Table[4](https://arxiv.org/html/2606.31166#S5.T4)compares the default star mask with an ego\-graph mask\. The ego\-graph mask allows attention along the observed edges in the sampled k\-hop ego graph, so it exposes the real local graph relationships among context nodes rather than using a target\-centred abstraction\.The star mask performs better on both NC and LP tasks, with a particularly large gap on LP\. The result supports our design choice: because each prompt asks for a target node label or a candidate edge decision, the target instance should remain the primary aggregation point\. Allowing neighbour nodes to communicate through observed ego\-graph edges introduces extra context\-node interactions that are not always relevant to the target prediction and can amplify dataset\-specific local wiring\. The star mask instead keeps information flow simple, target\-centred, and consistent across tasks\.

#### 5\.4\.3Effect of Neighbour Number

Table 5:Neighbour number ablation accuracy \(%\)\.KKdenotes the maximum number of sampled neighbours per hop, and each row reports one task\-dataset setting\.Table[5](https://arxiv.org/html/2606.31166#S5.T5)studies how the sampled neighbourhood size affectsTAG\-DLM\. VaryingKKchanges the maximum number of neighbours sampled per hop while keeping the prompt format, topology attention mask type, and scoring rule fixed, so the ablation isolates the effect of local context size\. The completed rows show that using the defaultK=10K\{=\}10is generally strongest, while very small neighbourhoods lack enough relational evidence and larger neighbourhoods introduce less relevant context\.

## 6Related Work

##### GNNs and graph transformers for TAG learning\.

GNNs propagate node representations along graph edges through message passing, such as GCN\(Kipf and Welling,[2017](https://arxiv.org/html/2606.31166#bib.bib3)\), GraphSAGE\(Hamiltonet al\.,[2017](https://arxiv.org/html/2606.31166#bib.bib4)\), GAT\(Veličkovićet al\.,[2018](https://arxiv.org/html/2606.31166#bib.bib5)\), and GATv2\(Brodyet al\.,[2022](https://arxiv.org/html/2606.31166#bib.bib6)\)\. Graph transformers further combine attention with structural biases\(Dwivedi and Bresson,[2021](https://arxiv.org/html/2606.31166#bib.bib7); Wuet al\.,[2022](https://arxiv.org/html/2606.31166#bib.bib8),[2023a](https://arxiv.org/html/2606.31166#bib.bib9),[2023b](https://arxiv.org/html/2606.31166#bib.bib10)\)\. These methods model topology explicitly, but typically operate on fixed node features rather than letting a language model read raw node text while performing graph reasoning\.

##### Language models for TAG learning\.

Language\-model approaches strengthen the textual side of TAG learning: RoBERTa\(Liuet al\.,[2019b](https://arxiv.org/html/2606.31166#bib.bib11)\)can be fine tuned over linearised graph contexts, while LLaGA\(Chenet al\.,[2024](https://arxiv.org/html/2606.31166#bib.bib12)\)combines LLM representations with graph\-side aggregation\. They confirm the value of textual semantics, but still separate language understanding from graph propagation\.TAG\-DLMinstead places graph structure inside the language model’s attention pattern, so the denoising network itself reasons over topology\.

##### MDLMs for graph learning\.

MDLMs\(Sahooet al\.,[2024](https://arxiv.org/html/2606.31166#bib.bib13); Nieet al\.,[2025](https://arxiv.org/html/2606.31166#bib.bib14)\)generate text by iteratively denoising masked tokens with bidirectional self\-attention\. This non\-causal attention suits graph contexts, where evidence should flow across target and neighbour text rather than left to right\. To our knowledge,TAG\-DLMis the first topology\-aware graph learner built on an MDLM with an input\-dependent attention mask\.

## 7Conclusion

We presentedTAG\-DLM, an MDLM defined by prompts for TAG learning\. By linearising local graph context and injecting topology through an attention mask,TAG\-DLMlets textual evidence and graph structure interact inside the denoising model without a separate GNN encoder or graph head specific to a task\. The theoretical analysis connects topology masked denoising to attention based message passing, and the experiments show strong performance on NC, LP, and cross dataset transfer\. Future work should explore richer topology attention masks, larger graph contexts, and broader graph learning settings beyond the studied benchmarks\.

## Limitations

WhileTAG\-DLMprovides a unified framework for text\-attributed graph learning with topology\-masked diffusion language models, several limitations remain\. First, the current framework mainly focuses on static homogeneous graphs and does not explicitly model evolving graph structures\. In real\-world settings such as citation streams and social networks, nodes, edges, and textual attributes may change over time\. ExtendingTAG\-DLMto dynamic graphs would require time\-aware linearization and temporal topology masks\. Second, our current topology mask does not distinguish different node or edge types\. This limits its applicability to heterogeneous graphs, where relation types often carry different semantics\. Future work could introduce type\-aware masks or relation\-conditioned attention biases to support heterogeneous message passing inside the denoising backbone\.

## Ethics Statement

This work uses public TAG benchmarks and does not involve human subjects, private user data, or personally identifiable information\. Potential risks mainly arise if graph prediction models are deployed in sensitive domains, where inferred labels or links may affect individuals or communities\. Such use requires additional fairness, privacy, and domain\-specific auditing beyond our offline benchmark setting\.

## References

- S\. Brody, U\. Alon, and E\. Yahav \(2022\)How attentive are graph attention networks?\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px1.p1.1)\.
- R\. Chen, T\. Zhao, A\. K\. Jaiswal, N\. Shah, and Z\. Wang \(2024\)LLaGA: large language and graph assistant\.InProceedings of the 41st International Conference on Machine Learning \(ICML\),pp\. 7809–7823\.Cited by:[§1](https://arxiv.org/html/2606.31166#S1.p2.1),[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px1.p1.1),[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px2.p1.1)\.
- V\. P\. Dwivedi and X\. Bresson \(2021\)A generalization of transformer networks to graphs\.AAAI Workshop on Deep Learning on Graphs \(DLG\-AAAI\)\.Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px1.p1.1)\.
- W\. L\. Hamilton, R\. Ying, and J\. Leskovec \(2017\)Inductive representation learning on large graphs\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px1.p1.1)\.
- E\. J\. Hu, Y\. Shen, P\. Wallis, Z\. Allen\-Zhu, Y\. Li, S\. Wang, L\. Wang, and W\. Chen \(2022\)LoRA: low\-rank adaptation of large language models\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px2.p1.10)\.
- W\. Hu, M\. Fey, M\. Zitnik, Y\. Dong, H\. Ren, B\. Liu, M\. Catasta, and J\. Leskovec \(2020\)Open graph benchmark: datasets for machine learning on graphs\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px1.p1.1)\.
- T\. N\. Kipf and M\. Welling \(2017\)Semi\-supervised classification with graph convolutional networks\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2606.31166#S1.p2.1),[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px1.p1.1)\.
- R\. Liu, J\. Regier, N\. Tripuraneni, M\. Jordan, and J\. Mcauliffe \(2019a\)Rao\-blackwellized stochastic gradients for discrete distributions\.InInternational Conference on Machine Learning,pp\. 4023–4031\.Cited by:[§2\.2](https://arxiv.org/html/2606.31166#S2.SS2.p2.2)\.
- Y\. Liu, M\. Ott, N\. Goyal, J\. Du, M\. Joshi, D\. Chen, O\. Levy, M\. Lewis, L\. Zettlemoyer, and V\. Stoyanov \(2019b\)RoBERTa: a robustly optimized BERT pretraining approach\.arXiv preprint arXiv:1907\.11692\.Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px2.p1.1)\.
- I\. Loshchilov and F\. Hutter \(2019\)Decoupled weight decay regularization\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px2.p1.10)\.
- S\. Nie, F\. Zhu, Z\. You, X\. Zhang, J\. Ou, J\. Hu, J\. Zhou, Y\. Lin, J\. Wen, and C\. Li \(2025\)Large language diffusion models\.Advances in Neural Information Processing Systems38,pp\. 50608–50646\.Cited by:[§2\.2](https://arxiv.org/html/2606.31166#S2.SS2.p1.10),[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px3.p1.1)\.
- J\. Rasley, S\. Rajbhandari, O\. Ruwase, and Y\. He \(2020\)DeepSpeed: system optimizations enable training deep learning models with over 100 billion parameters\.InProceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining \(KDD\),pp\. 3505–3506\.Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px2.p1.10)\.
- S\. S\. Sahoo, M\. Arriola, Y\. Schiff, A\. Gokaslan, E\. M\. Marroquin, J\. T\. Chiu, A\. M\. Rush, and V\. Kuleshov \(2024\)Simple and effective masked diffusion language models\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§2\.2](https://arxiv.org/html/2606.31166#S2.SS2.p1.10),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px3.p1.1)\.
- J\. Tang, Y\. Yang, W\. Wei, L\. Shi, L\. Su, S\. Cheng, D\. Yin, and C\. Huang \(2024\)GraphGPT: graph instruction tuning for large language models\.InProceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval,pp\. 491–500\.Cited by:[§1](https://arxiv.org/html/2606.31166#S1.p2.1)\.
- P\. Veličković, G\. Cucurull, A\. Casanova, A\. Romero, P\. Liò, and Y\. Bengio \(2018\)Graph attention networks\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§B\.5](https://arxiv.org/html/2606.31166#A2.SS5.p1.3),[§1](https://arxiv.org/html/2606.31166#S1.p2.1),[§2\.3](https://arxiv.org/html/2606.31166#S2.SS3.p1.6),[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px1.p1.1)\.
- Y\. Wang, Y\. Zhu, W\. Zhang, Y\. Zhuang, Y\. Li, and S\. Tang \(2024\)Bridging local details and global context in text\-attributed graphs\.InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing,pp\. 14830–14841\.Cited by:[§1](https://arxiv.org/html/2606.31166#S1.p1.1)\.
- Z\. Wang, S\. Liu, Z\. Zhang, T\. Ma, C\. Zhang, and Y\. Ye \(2025\)Can LLMs convert graphs to text\-attributed graphs?\.InProceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies \(Volume 1: Long Papers\),pp\. 1412–1432\.Cited by:[§1](https://arxiv.org/html/2606.31166#S1.p1.1)\.
- Q\. Wu, C\. Yang, W\. Zhao, Y\. He, D\. P\. Wipf, and J\. Yan \(2023a\)DIFFormer: scalable \(graph\) transformers induced by energy constrained diffusion\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px1.p1.1)\.
- Q\. Wu, W\. Zhao, Z\. Li, D\. P\. Wipf, and J\. Yan \(2022\)NodeFormer: a scalable graph structure learning transformer for node classification\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px1.p1.1)\.
- Q\. Wu, W\. Zhao, C\. Yang, H\. Zhang, F\. Nie, H\. Jiang, Y\. Bian, and J\. Yan \(2023b\)SGFormer: simplifying and empowering transformers for large\-graph representations\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2606.31166#S6.SS0.SSS0.Px1.p1.1)\.
- H\. Yan, C\. Li, R\. Long, C\. Yan, J\. Zhao, W\. Zhuang, J\. Yin, P\. Zhang, W\. Han, H\. Sun,et al\.\(2023\)A comprehensive study on text\-attributed graphs: benchmarking and rethinking\.Advances in Neural Information Processing Systems36,pp\. 17238–17264\.Cited by:[§1](https://arxiv.org/html/2606.31166#S1.p1.1)\.
- Z\. Yang, W\. W\. Cohen, and R\. Salakhutdinov \(2016\)Revisiting semi\-supervised learning with graph embeddings\.InProceedings of the 33rd International Conference on Machine Learning \(ICML\),pp\. 40–48\.Cited by:[§5\.1](https://arxiv.org/html/2606.31166#S5.SS1.SSS0.Px1.p1.1)\.

## Appendix ADataset Statistics

Table[6](https://arxiv.org/html/2606.31166#A1.T6)summarises the three TAG benchmarks used in our experiments\. All datasets are loaded from the processed version aligned with the LLaGA cache so that node text, edges, and label space match the inputs seen by the language model baselines\. Edge counts are reported as undirected edges\.

Table 6:Dataset statistics for the TAG benchmarks; edge counts are undirected\.##### Splits\.

For Cora and PubMed we follow the splits aligned with LLaGA used in our training pipeline; the test split contains542542nodes for Cora and3,9443\{,\}944nodes for PubMed\. For ogbn\-arxiv, we use the official OGB split:90,94190\{,\}941/29,79929\{,\}799/48,60348\{,\}603nodes for train/validation/test\. Validation labels and test node labels are masked when those nodes appear as neighbours of a training target, preventing label leakage through the neighbourhood sampler\.

##### Label semantics\.

Cora and PubMed labels are research topics \(e\.g\. “Neural Networks”, “Diabetes Mellitus, Type 1”\), and ogbn\-arxiv labels are arXiv computer science subject areas \(e\.g\. “cs\.CV”, “cs\.LG”\)\. In themc\_digitprompt, these surface forms are listed as options indexed by digit options, and the model predicts the digit corresponding to the target label\.

## Appendix BProof of Proposition[1](https://arxiv.org/html/2606.31166#Thmproposition1)and Relation to GAT

This appendix gives the formal assumptions, the full derivation of Proposition[1](https://arxiv.org/html/2606.31166#Thmproposition1), the multihead and reverse diffusion extensions, and the corollary that recovers GAT as a special case\.

### B\.1Notation and Assumptions

Fix a target nodevvwith sampled ego graphVv=\{v,u1,…,unv\}V\_\{v\}=\\\{v,u\_\{1\},\\ldots,u\_\{n\_\{v\}\}\\\}and linearised token sequence of lengthNN\. For each nodeu∈Vvu\\in V\_\{v\}letIu⊆\{1,…,N\}I\_\{u\}\\subseteq\\\{1,\\ldots,N\\\}be its contiguous token span; the spans partition the node token positions\. We use the following assumptions throughout\.

##### \(A1\) Span pooling\.

The node level representation at layerℓ\\ellis

huℓ=Pool⁡\(\{Hiℓ:i∈Iu\}\),h\_\{u\}^\{\\ell\}=\\operatorname\{Pool\}\\bigl\(\\\{H\_\{i\}^\{\\ell\}:i\\in I\_\{u\}\\\}\\bigr\),\(9\)whereHiℓH\_\{i\}^\{\\ell\}is the layer\-ℓ\\elltoken hidden state andPool\\operatorname\{Pool\}is mean pooling\.

##### \(A2\) Neighbourhood induced by the mask\.

For eachu∈Vvu\\in V\_\{v\},

𝒩M​\(u\)=\{w∈Vv:Iu​may attend to​Iw​under​Mv\}\.\\mathcal\{N\}\_\{M\}\(u\)=\\\{w\\in V\_\{v\}:I\_\{u\}\\text\{ may attend to \}I\_\{w\}\\text\{ under \}M\_\{v\}\\\}\.\(10\)For the topology attention mask ofTAG\-DLM,𝒩M​\(v\)=Vv\\mathcal\{N\}\_\{M\}\(v\)=V\_\{v\}and𝒩M​\(u\)=\{u,v\}\\mathcal\{N\}\_\{M\}\(u\)=\\\{u,v\\\}foru≠vu\\neq v\.

##### \(A3\) Fixed diffusion time\.

The proposition is established for one denoising layer at a fixed diffusion timett\. The timestep embedding may enter the attention score as an additional argumentsℓ​\(⋅,⋅,t\)s\_\{\\ell\}\(\\cdot,\\cdot,t\)but does not modify the spatial structure\.

##### \(A4\) Bidirectional denoising attention\.

The denoising network is not causal: in the absence ofMvM\_\{v\}every token may attend to every other\. Hence, the only structural restriction on information flow within a layer is the topology attention mask itself, rather than token order from left to right\.

##### \(A5\) Attention compatible with pooling\.

The token\-level attention weights inside a span are approximately constant across the span, so that pooling commutes with the attention sum up to the modelling abstraction in Eq\. \([9](https://arxiv.org/html/2606.31166#A2.E9)\)\. This is the only abstraction that converts token\-level to node\-level message passing\.

### B\.2Proof of Proposition[1](https://arxiv.org/html/2606.31166#Thmproposition1)

LetWQℓ,WKℓ,WVℓW\_\{Q\}^\{\\ell\},W\_\{K\}^\{\\ell\},W\_\{V\}^\{\\ell\}be the layer\-ℓ\\ellquery, key, and value projections\. The token\-level self\-attention output at positioniiis

Ziℓ=∑j=1Nαi​jℓ​WVℓ​Hjℓ,αi​jℓ=ei​jℓ∑kei​kℓ,Z\_\{i\}^\{\\ell\}=\\sum\_\{j=1\}^\{N\}\\alpha\_\{ij\}^\{\\ell\}\\,W\_\{V\}^\{\\ell\}H\_\{j\}^\{\\ell\},\\quad\\alpha\_\{ij\}^\{\\ell\}=\\frac\{e\_\{ij\}^\{\\ell\}\}\{\\sum\_\{k\}e\_\{ik\}^\{\\ell\}\},\(11\)with raw scoresei​jℓ=exp⁡\(\(WQℓ​Hiℓ\)⊤​WKℓ​Hjℓ/d\)e\_\{ij\}^\{\\ell\}=\\exp\\\!\\bigl\(\(W\_\{Q\}^\{\\ell\}H\_\{i\}^\{\\ell\}\)^\{\\top\}W\_\{K\}^\{\\ell\}H\_\{j\}^\{\\ell\}/\\sqrt\{d\}\\bigr\)\. Applying the topology attention maskMvM\_\{v\}setsei​jℓ=0e\_\{ij\}^\{\\ell\}=0whenever the node containingiiis not allowed to attend to the node containingjj, which by \(A2\) means

αi​jℓ=0wheneverj∈Iw,w∉𝒩M​\(u\),i∈Iu\.\\alpha\_\{ij\}^\{\\ell\}=0\\quad\\text\{whenever\}\\quad j\\in I\_\{w\},\\ w\\notin\\mathcal\{N\}\_\{M\}\(u\),\\ i\\in I\_\{u\}\.\(12\)
Aggregate token positions by their owning node\. Fori∈Iui\\in I\_\{u\},

Ziℓ=∑w∈𝒩M​\(u\)∑j∈Iwαi​jℓ​WVℓ​Hjℓ\.Z\_\{i\}^\{\\ell\}=\\sum\_\{w\\in\\mathcal\{N\}\_\{M\}\(u\)\}\\sum\_\{j\\in I\_\{w\}\}\\alpha\_\{ij\}^\{\\ell\}\\,W\_\{V\}^\{\\ell\}H\_\{j\}^\{\\ell\}\.\(13\)Pooling overi∈Iui\\in I\_\{u\}and using \(A5\) to commute pooling with the inner sum yields

muℓ:=Pooli∈Iu⁡\(Ziℓ\)=∑w∈𝒩M​\(u\)αu​wℓ​WVℓ​hwℓ,m\_\{u\}^\{\\ell\}:=\\operatorname\{Pool\}\_\{i\\in I\_\{u\}\}\(Z\_\{i\}^\{\\ell\}\)=\\sum\_\{w\\in\\mathcal\{N\}\_\{M\}\(u\)\}\\alpha\_\{uw\}^\{\\ell\}\\,W\_\{V\}^\{\\ell\}h\_\{w\}^\{\\ell\},\(14\)with the node\-level weight

αu​wℓ=exp⁡sℓ​\(huℓ,hwℓ,t\)∑r∈𝒩M​\(u\)exp⁡sℓ​\(huℓ,hrℓ,t\),\\alpha\_\{uw\}^\{\\ell\}=\\frac\{\\exp s\_\{\\ell\}\(h\_\{u\}^\{\\ell\},h\_\{w\}^\{\\ell\},t\)\}\{\\sum\_\{r\\in\\mathcal\{N\}\_\{M\}\(u\)\}\\exp s\_\{\\ell\}\(h\_\{u\}^\{\\ell\},h\_\{r\}^\{\\ell\},t\)\},\(15\)wheresℓ​\(hu,hw,t\)=\(WQℓ​hu\)⊤​WKℓ​hw/ds\_\{\\ell\}\(h\_\{u\},h\_\{w\},t\)=\(W\_\{Q\}^\{\\ell\}h\_\{u\}\)^\{\\top\}W\_\{K\}^\{\\ell\}h\_\{w\}/\\sqrt\{d\}plus any timestep bias\. The transformer’s residual connection, LayerNorm, output projectionWOℓW\_\{O\}^\{\\ell\}, and feed\-forward block together define the update

Uℓ​\(huℓ,muℓ\)=FFNℓ⁡\(LN⁡\(huℓ\+WOℓ​muℓ\)\)\.U\_\{\\ell\}\(h\_\{u\}^\{\\ell\},m\_\{u\}^\{\\ell\}\)=\\operatorname\{FFN\}\_\{\\ell\}\\\!\\left\(\\operatorname\{LN\}\\\!\\bigl\(h\_\{u\}^\{\\ell\}\+W\_\{O\}^\{\\ell\}m\_\{u\}^\{\\ell\}\\bigr\)\\right\)\.\(16\)Combining Eqs\. \([14](https://arxiv.org/html/2606.31166#A2.E14)\) and \([16](https://arxiv.org/html/2606.31166#A2.E16)\) gives

huℓ\+1=Uℓ​\(huℓ,∑w∈𝒩M​\(u\)αu​wℓ​WVℓ​hwℓ\),h\_\{u\}^\{\\ell\+1\}=U\_\{\\ell\}\\\!\\left\(h\_\{u\}^\{\\ell\},\\;\\sum\_\{w\\in\\mathcal\{N\}\_\{M\}\(u\)\}\\alpha\_\{uw\}^\{\\ell\}\\,W\_\{V\}^\{\\ell\}h\_\{w\}^\{\\ell\}\\right\),\(17\)which is the canonical attention\-based message passing layer on𝒢M\\mathcal\{G\}\_\{M\}\. ∎

### B\.3Multihead Extension

For anRR\-head attention layer, repeating the argument per head yields a message for each head

mu,rℓ=∑w∈𝒩M​\(u\)αu​w,rℓ​WV,rℓ​hwℓ,m\_\{u,r\}^\{\\ell\}=\\sum\_\{w\\in\\mathcal\{N\}\_\{M\}\(u\)\}\\alpha\_\{uw,r\}^\{\\ell\}\\,W\_\{V,r\}^\{\\ell\}h\_\{w\}^\{\\ell\},\(18\)and the node level update concatenates the heads,muℓ=Concatr=1R⁡mu,rℓm\_\{u\}^\{\\ell\}=\\operatorname\{Concat\}\_\{r=1\}^\{R\}\\,m\_\{u,r\}^\{\\ell\}, before applyingWOℓW\_\{O\}^\{\\ell\}\. Multihead attention is therefore a multichannel attention\-based message passing layer on the same𝒢M\\mathcal\{G\}\_\{M\}\.

### B\.4Reverse Diffusion as Composed Message Passing

The reverse diffusion procedure with timestepst1\>⋯\>tTt\_\{1\}\>\\cdots\>t\_\{T\}applies the same denoiser at decreasing noise levels:

H\(k−1\)=fθ​\(H\(k\),tk,Mv\)\.H^\{\(k\-1\)\}=f\_\{\\theta\}\\bigl\(H^\{\(k\)\},t\_\{k\},M\_\{v\}\\bigr\)\.\(19\)By Proposition[1](https://arxiv.org/html/2606.31166#Thmproposition1), each callfθ​\(⋅,tk,Mv\)f\_\{\\theta\}\(\\cdot,t\_\{k\},M\_\{v\}\)is a stack of attention\-based message passing layers on𝒢M\\mathcal\{G\}\_\{M\}whose attention scores are conditioned on the timestep\. The composition

H\(0\)=fθ​\(⋅,t1,Mv\)∘⋯∘fθ​\(⋅,tT,Mv\)​\(H\(T\)\)H^\{\(0\)\}=f\_\{\\theta\}\(\\cdot,t\_\{1\},M\_\{v\}\)\\circ\\cdots\\circ f\_\{\\theta\}\(\\cdot,t\_\{T\},M\_\{v\}\)\\bigl\(H^\{\(T\)\}\\bigr\)\(20\)is therefore a sequence of message passing operators conditioned on the timestep on the same graph induced by the mask\.

### B\.5Corollary: GAT as a Special Case

Leta∈ℝ2​d′a\\in\\mathbb\{R\}^\{2d^\{\\prime\}\}be a learnable vector and choose the attention score

sℓ​\(hu,hw,t\)=LeakyReLU⁡\(a⊤​\[W​hu∥W​hw\]\),s\_\{\\ell\}\(h\_\{u\},h\_\{w\},t\)=\\operatorname\{LeakyReLU\}\\\!\\bigl\(a^\{\\top\}\[Wh\_\{u\}\\,\\\|\\,Wh\_\{w\}\]\\bigr\),\(21\)i\.e\. the additive form ofVeličkovićet al\.\([2018](https://arxiv.org/html/2606.31166#bib.bib5)\)withWWa shared linear projection and the timestep argument dropped\. Substituting this choice into Eqs\. \([14](https://arxiv.org/html/2606.31166#A2.E14)\) and \([15](https://arxiv.org/html/2606.31166#A2.E15)\) and identifyingWVℓ≡WW\_\{V\}^\{\\ell\}\\equiv Wrecovers the standard GAT update exactly\. Thus GAT is the special case of Proposition[1](https://arxiv.org/html/2606.31166#Thmproposition1)obtained by restricting the score function to additive attention with a single shared projection\. The dot\-product score used byTAG\-DLMgeneralises this to a broader class of attention\-based MPNNs\.

### B\.6Limits of the Equivalence

We close with the limits of the equivalence\. \(i\) The result is to*attention\-based*message passing, not to arbitrary MPNNs whose aggregation is not a softmax weighted sum\. \(ii\) The message passing graph is the graph induced by the mask𝒢M\\mathcal\{G\}\_\{M\}, which for our default mask is a star centred on the target and therefore blocks direct interaction among neighbours; equivalence to a layer over the original ego graph requires the mask to expose those edges\. \(iii\) Diffusion time is not graph diffusion time: the composition over timesteps in Section[B\.4](https://arxiv.org/html/2606.31166#A2.SS4)is a sequence of distinct message passing operators conditioned on the timestep rather than the integration of any continuous graph diffusion process\. \(iv\) Eq\. \([14](https://arxiv.org/html/2606.31166#A2.E14)\) relies on the span pooling abstraction \(A5\); more fine\-grained, position\-dependent attention inside a span yields a strictly more expressive operator that contains attention\-based message passing as a coarse\-grained quotient\.

## Appendix CPrompt Format Example

This appendix gives a concrete example of howTAG\-DLMconverts a sampled ego graph into a token sequence\. The example is illustrative: node texts are shortened for readability, but the segment order, options indexed by digits, neighbour label brackets, and topology attention mask match themc\_digitformat used by the model\.

##### Linearised prompt\.

Suppose the target nodevvis a Cora paper whose label is hidden, and the sampled neighbourhood contains two papersu1u\_\{1\}andu2u\_\{2\}\. The first neighbour has an observed training label; the second is treated as unlabelled and therefore drops the bracket prefix\.

> Target paper: Graph neural networks for citation data\. Options: 0=Case Based; 1=Genetic Algorithms; 2=Neural Networks; \.\.\. Answer: \[M\] Neighbor 1 \[2\]: Semisupervised classification with graph convolutions\. Neighbor 2: Learning representations for linked documents\.

The single\[M\]token is the answer position\. At scoring time,TAG\-DLMcompares the logits of the valid digit tokens at this position; in this example, choosing digit2corresponds to predicting the label “Neural Networks\.” For LP, the same template is used with task specific options such as0=Noand1=Yes\.

##### Topology attention mask\.

After tokenisation, all tokens belonging to the same node text are grouped into a segment\. The segment level mask below shows the star shaped attention pattern used in the running example; each entry indicates whether the row segment can attend to the column segment\. Padding tokens are never attended to\.

Table 7:Segment level topology attention mask for the example prompt; the token level mask expands each entry to all token pairs in the corresponding segments\.The target segment can aggregate evidence from every sampled neighbour\. Each neighbour can attend to itself and the target, but not directly to other neighbours\. This keeps the prompt bidirectional within allowed segments while making structured graph information flow explicit in the attention mask\.

## Appendix DTraining Details and Hyperparameters

This appendix lists the full set of hyperparameters used to trainTAG\-DLMon each dataset and task configuration\. The same recipe is used for all three TAG benchmarks unless explicitly noted; the only adjustment by dataset is the maximum sequence length, which is increased on ogbn\-arxiv to fit longer abstracts\. Table[8](https://arxiv.org/html/2606.31166#A4.T8)summarises the configuration\.

Table 8:Default training hyperparameters forTAG\-DLM\.##### Loss\.

The training objective is the standard MDLM evidence lower bound restricted to the answer position: only the masked digit token contributes to the cross entropy loss, while the textual context \(target instance text, neighbour text, and observed neighbour label digits when available\) is treated as conditioning\. For NC, the neighbour label masking option masks neighbour labels with the same probability as the answer position during training, preventing the model from learning a trivial copy shortcut from a matching neighbour label to the target\.

##### LP protocol\.

For LP, each example is a candidate node pair with binary options0=Noand1=Yes\. Positive examples are held out edges from the corresponding split, and negative examples are sampled from node pairs without edges following the same protocol as the graph baselines\. Before sampling the local context for a positive candidate pair, we remove that candidate edge from the graph and construct neighbourhoods using only the permitted training graph\. This prevents the model from observing the answer through the sampled topology\. Accuracy is computed over the resulting positive and negative candidate pairs\.

##### Reproducibility\.

When repeated runs are available, dataset/configuration pairs use seeds\{0,1,2\}\\\{0,1,2\\\}; the current tables report the selected checkpoints recorded in the experiment logs\. The training script, environment specification, and exact launch commands appear underexamples/tmdlm/in the released code\. The full run specific configuration, including the resolved LoRA target module list, prompt template strings, and tokenisation parameters, is dumped to aconfig\.jsonfile alongside the checkpoint to enable exact replays\.

## Appendix EChecklist

##### Artifacts and licenses\.

This work uses public research artifacts\. Cora and PubMed are used through the Planetoid/LLaGA\-aligned benchmark distribution; the Planetoid repository is MIT\-licensed and the LLaGA codebase is Apache\-2\.0 licensed\. ogbn\-arxiv is distributed by OGB under the ODC\-BY license\. The LLaDA\-8B\-Instruct backbone is released under the MIT license\. We release only our code, scripts, and configuration files under the MIT license, and do not redistribute raw datasets or processed dataset caches\.

##### Intended use and privacy\.

The datasets are public citation\-style research benchmarks, and our experiments are limited to offline academic evaluation\. We do not collect new human\-subject data, private user data, or personally identifying information\. The model is intended for research on text\-attributed graph learning and should not be directly deployed for sensitive decisions involving individuals or communities without additional privacy, fairness, and domain\-specific auditing\.

##### Computational budget\.

Each dataset\-task configuration is trained separately on a single node with eight NVIDIA A100 80GB GPUs\. One training run takes approximately 4 hours for Cora, 12 hours for PubMed, and 50 hours for ogbn\-arxiv, corresponding to 32, 96, and 400 A100 GPU\-hours per task run\. Since NC and LP are trained separately, the main in\-domain training runs require approximately 1056 A100 GPU\-hours in total\. Cross\-dataset transfer evaluations reuse trained checkpoints and require no additional fine\-tuning\.

##### Result reporting\.

The reported table entries are from single runs using the selected checkpoints recorded in the experiment logs, rather than means over multiple random seeds\. We provide the training setup, seeds, prompt format, and checkpoint/configuration logging details in Appendix[D](https://arxiv.org/html/2606.31166#A4)to make this reporting convention explicit\. Appendix[D](https://arxiv.org/html/2606.31166#A4)also reports the package, implementation, and parameter settings used for both baselines and our method\.

##### Use of AI assistants\.

AI assistants were used to review code and to review manuscript writing for clarity, consistency, and formatting\. The authors verified all technical claims, experimental results, code changes, and final manuscript text\.

Similar Articles

TextLDM: Language Modeling with Continuous Latent Diffusion

Hugging Face Daily Papers

This paper introduces TextLDM, a method that adapts visual latent diffusion transformers for language modeling by mapping discrete tokens to continuous latents. It demonstrates that this approach, enhanced by representation alignment, matches GPT-2 performance and unifies visual and text generation architectures.

Continuous Latent Diffusion Language Model

Hugging Face Daily Papers

Cola DLM is a hierarchical latent diffusion language model that uses text-to-latent mapping and conditional decoding to achieve efficient, non-autoregressive text generation.

PerceptionDLM: Parallel Region Perception with Multimodal Diffusion Language Models

Hugging Face Daily Papers

PerceptionDLM introduces a multimodal diffusion language model that enables parallel region perception via structured attention masking and efficient prompting, achieving faster inference without sacrificing caption quality. Experiments show competitive performance with substantial speed improvements for multi-region perception tasks.