DIFFormer: Scalable (Graph) Transformers Induced by Energy Constrained Diffusion
Authors: Qitian Wu
Chenxiao Yang
Wentao Zhao
Yixuan He
David Wipf
Junchi Yan
Venue: ICLR 2023 notable top 25% (2023)
Keywords: structured representation learning
diffusion model
optimization-induced model
node prediction

Abstract: Real-world data generation often involves complex inter-dependencies among instances, violating the IID-data hypothesis of standard learning paradigms and posing a challenge for uncovering the geometric structures for learning desired instance representations. To this end, we introduce an energy constrained diffusion model which encodes a batch of instances from a dataset into evolutionary states that progressively incorporate other instances' information by their interactions. The diffusion process is constrained by descent criteria w.r.t. a principled energy function that characterizes the global consistency of instance representations over latent structures. We provide rigorous theory that implies closed-form optimal estimates for the pairwise diffusion strength among arbitrary instance pairs, which gives rise to a new class of neural encoders, dubbed as DIFFormer (diffusion-based Transformers), with two instantiations: a simple version with linear complexity for prohibitive instance numbers, and an advanced version for learning complex structures. Experiments highlight the wide applicability of our model as a general-purpose encoder backbone with superior performance in various tasks, such as node classification on large graphs, semi-supervised image/text classification, and spatial-temporal dynamics prediction. The codes are available at https://github.com/qitianwu/DIFFormer.

Published as a conference paper at ICLR 2023
DIFFORMER: SCALABLE (GrAPH) TRANSFORMERS INDUCED BY ENERGY CONSTRAINED DIFFUSION
Qitian Wu', Chenxiao Yang', Wentao Zhao', Yixuan Hel, David Wipf', Junchi Yan i Department of CSE & MoE Lab of Artificial Intelligence, Shanghai Jiao Tong University
} Department of Statistics, University of Oxford
S Amazon Web Service
(echo740, chr261 95, permanent, permanent,yaajuunchh ) jsjtu. edu. cn, yi xuaan, hc@stats ox. ax.ac.uk davidwipf@gmail.com
ABSTRACT
Real-world data generation often involves complex inter-dependencies among instances, violating the IID-data hypothesis of standard learning paradigms and posing challenge for uncovering the geometric structures for learning desired instance representations. To this end, we introduce an energy constrained diffu- sion model which encodes batch of instances from a dataset into evolutionary states that progressively incorporate other instances' information by their interac- tions. The diffusion process is constrained by descent criteria w.r.t. a principled energy function that characterizes the global consistency of instance representa- tions over latent structures. We provide rigorous theory that implies closed- form optimal estimates for the pairwise diffusion strength among arbitrary instance pairs, which gives rise to new class neural encoders, dubbed as DIFFORMER (diffusion- based Transformers), with two instantiations: simple version with linear complexity for prohibitive instance numbers, and an advanced version for learning complex structures. Experiments highlight the wide applicability of our model as a general-purposs encoder backbone with superior performance in vari- ous tasks, such node classification large graphs, semi-supervised image/text classification, and spatial-tempon dynamics prediction. The codes are available at https:: /tteps:/ggiiee www.httttt.uugggo com/ qi wwww///IFoo./........
INTRODUCTION
Real-world data are generated from a convoluted interactive process whose underlying physical prin- ciples are often unknown. Such nature violates the common hypothesis of standard representation learning paradigms assuming that data are IID sampled. The challenge, however, is that due to the absence of prior knowledge about ground-truth data generation, can be practically prohibitive to build feasible methodology for uncovering data dependencies, despite the acknowledged significance. To address this issue, prior works, e.g., Wang et al. (2019); Franceschi et al. (2019); Jiang al. (2019); Zhang et al. (2019), consider encoding the potential interactions between instance pairs, but this requires sufficient degrees of freedom that significantly increases learning difficulty from limited labels (Fatemi et al., 2021) and hinders the scalability to large systems (Wu et al., 2022b).
Turning to a simpler problem setting where putative instance relations are instantiated as an observed graph, remarkable progress has been made in designing expressive architectures such graph neural networks (GNNs) (Scarselli et al., 2008; Kipf & Welling, 2017; Velickovic et al., 2018; Wu et al., 2019: Chen et al., 2020a; Yang et al., 2021) for harnessing nnter-connettions between instances geometric prior (Bronstein et al., 2017). 7). However, the observed relations can be incomplete/noisy, due error-prone data collection, or generated by artificial construction independent from downstream targets. The potential inconsistency between observation and the underlying data geometry would presumably elicit systematic bias between structured representation of graph-based learning and true
"Corresponding author: Junchi Yan who is also affiliated with Shanghai Al Lab. The work was in part supported by National Key Research and Development Program of China (202AAAA0107600)) National Natural Science Foundation of China (62222607), STCSM (2251110510)

---

Published as a conference paper at ICLR 2023
Diffusion Process
Supervised Loss
(1)
V (8(2(())) () (a.(0)
A x)
Mfra Bes)
1=0
Encoder
Energy Maxx.
Classifier
Nin
Observed
Energy Constraint 8/7, 12- 2011 Al x,11)
Figure I: An illustration of the general idea behind DIFFORMER which takes whole dataset (or a batch) of instances as input and encodes them into hidden states through diffusion process aimed at minimizing a regularized energy. This design allows feature propagation among arbitrary instance pairs at each layer with optimal inter-connecting structures for informed prediction on each instance.
data dependencies. While a plausible remedy is to learn more useful structures from the data, this unfortunately brings the previously-mentionn
obstacles to the fore.
To resolve the dilemma, we propose a novel general-
encoder framework that uncovers data dependencies from observations (a dataset of partially labeled instances), proceeding via two-fold inspiration from physics as illustrated in Fig. 1. Our model is defined through feed-forward continuous dynamies (i.e., PDE) involving all the instances of dataset as locations on Riemannian manifolds with latent structures, upon which the features of instances act as heat flowing over the underlying geometry (Hamzi & Owhadi, 2021). Such a diffusion model serves an important inductive bias for leveraging global information from other instances to obtain more informative representations. Its major advantage lies in the flexibility for the diffusivviy function, i.e., a measure of the rate at which information spreads (Rosenberg & Steven, 1997): we allow for feature propagation between arbitrary instance pairs at each layer, and adaptively navigate this process by pairwise connectivity weights. Moreover, for guiding the instance representations towards some ideal constraints of internal consistency, we introduce a principled energy function that enforces layer-wise regularization on the evolutionary directions. The energy function provides another view (from macroscopic standpoint) into the desired instance representations with low energy that are produced, i.c., soliciting a steady state that gives rise to informed predictions on unlabeled data.
As a justification for the tractability of above general methodology, our theory reveals the underlying equivalence between finiteddifference iterations of the diffusion process and unfolding the minimiza- tion dynamics for an associated regularized energy. This result further suggests a closed-form optimal solution for the diffusivity function that updates instance representations by the ones of all the other instances towards giving a rigorous decrease of the global energy. Based on this, we also show that the energy constrained diffusion model can serve as a principled perspective for unifying popular models like MLP, GCN and GAT which can be viewed as special cases of our framework.
On top of the theory, we propose new class of neural encoders, Diffusion- based Transformers (DIFFORMER), and its two practical instantiations: one is a simple version with O(N) complexity (N for instance number) for computing all-pair interactions among instances; the other is a more expressive version that can learn complex latent structures. We empirically demonstrate the success of DIFFORMER on a diverse set of tasks. It outperforms SOTA approaches on superviedd node classification benchmarks and performs competitively on large-scale graphs. It also shows promising power for image/text classification with low label rates and predicting
dynamics.
2 RELATED WORK
Graph-based Semi-supervised Learning, Graph-based SSL (Kipf & Welling, 2017) aims to learn from partially labeled data, where instances are treated as nodes and their relations are given by a graph. The observed structure can be leveraged as regularization for learning representations (Belkin et al., 2006; Weston et al., 2012; Yang et al., 2016) or as an inductive bias of modern GNN architec- tures (Scarselli al., 2008). However, there frequently exist situations where the observed structure is unavailable or unreliable (Franceschi et al., 2019; Jiang et al., 2019; Chen et al., 2020c; Fatemi et al., 2021; Lao et al., 2022), in which case the challenge remains how to uncover the underlying
2

---

Published as a conference paper at ICLR 2023
relations. This paper explores new Transformer-lik encoder for discovering data geometry to promote learning through the inter-dependence among instances (either labeled or unlabeled).
Neural Diffusion Models. Several recent efforts explore diffusion- learning where continuous dynamics serve as an inductive bias for representation learning (Hamzi & Owhadi, 2021). These works can be generally grouped into PDE-based learning, where the model itself is continuous diffusion process described by a differential equation (e.g.. Chamberlain et al. (2021a); Eliasof al. (2021); Thorpe et al. (2022)), and PDE-inspired learning, where the diffusion perspective is adopted to develop new models (Atwood & Towsley, 2016; Klicpera et al., 2019) or characterize graph geometric properties (Yang et al., 2022). Our work leans on the later category and the key originality lies in two aspects. First, we introduce a novel diffusion model whose dynamics are implicitly defined by optimizing a regularized energy. Second, our theory establishes an equivalence between the numerical iterations of diffusion process and unfolding the optimization of the energy, based on which we develop a new class of neural encoders for uncovering latent structures among a large number of instances.
Transformers. Transformers serve as a model class of wide research interest showing competitive efficacy for modeling the dependency among tokens of inputs through all-pair attention mechanism. While the original architecture (Vaswani et al., 2017) is motivated by the dnter-deeendence among words of a sentence in NLP tasks, from different aspects, this work targets Transformer- like architecture for modeling the inter-dependence among instances in a dataset. In the terminology of Transformers in NLPs, one treat the whole dataset an extremely long 'sequence' (with the length equaling to the dataset size), and each instance (with a label to predict) acts as the 'token' in such a sequence. In the context of graph machine learning, our goal can be framed as learning latent interaction graphs among nodes beyond the observed graph (if available), and this can be generally viewed as an embodiment of node-level prediction (Hu et al., 2020). For the latter case, critically though, it remains under-expi under-explore how to build scalable and expressive Transformer for learning node-pair interactions given the prohibitively large number of node instances (Wu et al., 2022b).
ENERGY CONSTRAINED GEOMETRIC DIFFUSION TRANSFORMERS
Consider set of partially labeled instances {x, whose labeled portion is {(x, 3, 3/) (often M & N). In some cases there exist relational structures that connect instances graph EV, where the node set V contains all the instances and the edge set & = {eyj} consists of observed relations. Without loss of generality, the main body of this section does not assume graph structures as input, but we will later discuss how to trivially incorporate them if they are available.
3.1 GEOMETRIC DIFFUSION MODEL
The starting point of our model is diffusion process that treats dataset of instances as whole and produces instance representations through information flows characterized by an anisotropic diffusion process, which is inspired by an analogy with heat diffusion on Riemannian manifold (Rosenberg & Steven, 1997). We use a vector-valued function (t) oo) + Rd to define an instance's 's state at time and location i. The anisotropic diffusion process describes the evolution of instance states (i.e., representations) via a PDE with boundary conditions (Freidlin & Wentzell, 1993; Medvedev, 2014):
OZ(() V" (S(Z(t).t) o VZ(()), s. t. Z(0) = [x] t 2 0. Øt
where = [z) E Nxd O denotes the Hadamard product, and the function S(Z(t), t) N x oo) -> [0, NxN defines the diffiusivity coefficient controlling the diffusion strength between any pair at time t. The diffusivity is specified to be dependent on instances' states. The gradient operator measures the difference between source and target states, i.e., (VZ(t); =
2, (f), and the divergence operator V sums up information flows at point, i.e. (V');
Σ j=1 (Z(),)) (VZ()),, Note that both operators are defined over a discrete space consisting of N locations. The physical implication of Eq. 1 that the temporal change of heat at location equals to the heat flux that spatially enters into the point. Eq. 1 can be explicitly written as
(z,(t) 01
Σ Sc (Z(t) t)), (z, (t))..
(2)
1=1
3

---

Published as a conference paper at ICLR 2023
Such diffusion process can as inductive bias that guides the model to other instances' information at every layer for learning informative instance representations. We can use numerical methods to solve the continuous dynamics in Eq. 2, e.g., the explicit Euler scheme involving finite differences with step size T, which after some re-arranging gives:
(k+1) (
N
N
(k)
(2)
2
j=1
J=1
The numerical iteration can stably converge for T E (0, We can adopt the state after a finite number K of propagation steps and use it for final predictions, i.e., Dr (LP(z(()).
Remark. The diffusivity coefficient in Eq. 1 is measure of the rate at which heat can spread over the space (Rosenberg & Steven, 1997). Particularly in Eq. (tZ) f) determines how information flows over instances and the evolutionary direction of instance states. Much flexibility remains for its specification. For example, a basic choice is to fix (Z)) t) as an identity matrix which constrains the feature propagation to self-loops and the model degrades to an MLP that treats all the instances independently. One could also specify S(Z(t), f) as the observed graph structure if available some scenarios. In such a case, however, the information flows are restricted by neighboring nodes in graph. An ideal case could be to allow (Z))),, () to have non-zero values for arbitrary (i,j) and evolve with time, i.e., the instance states at each layer can efficiently and adaptively propagate to all the others.
3.2 DIFFUSION CONSTRAINED BY A LAYER-WISS ENERGY
As mentioned previously, the crux how to define proper diffusivity function to induce desired diffusion process that can maximize the information utility and accord with some inherent consistency. Since we have no prior knowledge for the explicit form or the inner structure of (kk), we consider the diffusivity as a time-dependent latent variable and introduce an energy function that measures the presumed quality of instance states at a given step k:
ZZ = Z (8)
13
Σ
5(12: 2212), J
where d: IR+ -> defined as function that non-decreasing and concave on particular interval of our interest, and promotes robustness against large differences (Yang et al., 2021) among any pair of instances. Eq. assigns each state in I2 with an energy scalar which can be leveraged to regularize the updated states (towards lower energy desired). The weight À trades two effects: 1) for
each instance i, the states not far from the current one (*) have low energy; 2) for all instances, the smaller differences their states have, the lower energy is produced.
Remark. Eq. 4 can essentially be seen as a robust version of the energy introduced by Zhou et al. (2004), inheriting the spirit of regularizing the global and local consistency of representations. "Robust" here particularly implies that the ó adds uncertainty to each pair the instances and could implicitly filter the information of noisy links (potentially reflected by proximity in the latent space).
Energy Constrained Diffusion. The diffusion process describes the microscopic behavior of instance states through evolution, while the energy function provides a macroscopic view for quantifying the consistency. In general, we expect that the final states could yield a low energy, which suggests that the physical system arrives at steady point wherein the yielded instance representations have absorbed enough global information under certain guiding principle. Thereby, we unify two schools of thoughts into a new diffusive system where instance states would evolve towards producing lower energy, e.g., by finding valid diffusivity function. Formally, aim to find series of ((k)'s 's whose dynamics and constraints are given by
N
N
(k-
(N)
Z
T
(5)
j=1
j=1
(0)
E(Z
k) (() 1:8), 1
z
X1y
The formulation induces new class of geometric flows on latent manifolds whose dynamics are implicitly defined by optimizing a time-varying energy function (see Fig. 1 for an illustration).
4

---

Published as a conference paper at ICLR 2023
3.3 TRACTABILITY OF SOLVING DIFFUSION PROCESS WITH ENERGY MINIMIZATION
Unfortunately, Eq. is hard to solve since we need to infer the value for a series of coupled ((k) that need to satisfy K inequalities by the energy minimization constraint. The key result of this paper is the following theorem that reveals the underlying connection between the geometric diffusion model and iterative minimization of the energy, which further suggests explicit closed-fome solution for (()) based on the current states (4) that yields a rigorous decrease of the energy.
Theorem 1. For any regularized energy defined by Eq. 4 with given X. there exists 0 < T < 1 such that the diffusion process of Eq. 3 with the diffusivity between pair (i,3) at the k-th step given by
S (1)
Wa
(R)
(N)
06(22)
(6)
E1=1 (K)
0z22
1= - 1
yields a descent step on the energy, i.e., (1+1), k; 6) S E(Z (k)
for any k 2 1.
Theorem 1 suggests the existence for the optimal diffusivity in the form of a function over the l2 distance between states at the current step, i.e., ||z| 2 The result enables us to unfold the implicit process and compute S(k) in a ferd-forward way from the initial states. We thus arrive at new family of neural model architectures with layer-wise wise computation specified by:
1|2-
13)
Diffusivity Inference:
(A)
1 SSij
Σ f(112
(K)
State Updating: 2,
(++1)
N
((K) 1
2,
(1)
00 3
R
IsiSN.
(7)
state conservation
state propagation
Remark. The choice of function f in above formulation is not arbitrary, but needs to be a non-negative and decreasing function of 2², that the associated 8 in Eq. 4 is guaranteed to be non-decreasing and concave w.r.t. 22. Critically though, there remains much room for us to properly design the specific f. so to provide adequate capacity and scalability. Also, in our model presented by Eq. 7 we only have one hyper-parameter T in practice, noting that the weight À in the regularized energy is implicitly determined through T by Theorem 1, which reduces the cost of hyper-parameter searching.
4
INSTANTIATIONS OF DIFFORMER
4.1 MODEL INSTANTIATIONS
We next go into model instantiations based on the above theory, with two specified f's as practical versions of our model. Due to space limits, we describe the key ideas concerning the model design in this subsection, and defer the details of model architectures to the self-contained Appendix E. First, because ||z, 2 l1 ||2,122 + ||z 1 2z, Zj. we can convert [[|z, 2)112) into the form g(z 2,) using change of variables on the condition that ||z.||2 remains constant. And we add layer normalization to each layer to loosely enforce such a property in practice.
Simple Diffusivity Model. A straightforward design is to adopt the linear function g(z) 1 +
(k)
(R) wij
(k) f(1iz)
2//0) 11:)
1
z
(8)
12
(K)
Assuming 2
and
||a (k)
2, (k)
112. Eq. 8 can be written as f(=?)
2 - 2 which yields a non-negative result and is decreasing on the interval [0, 2] in which 2² lies. One scalability concern for the model Eq. arises because of the need compute pairwise diffusivity and propagation for each individual, inducing O complexity. Remarkably, the simple diffusivity model allows a significant acceleration by noting that the state propagation can be re-arranged via
+ (z) (k) T
(k)
2
(2)
(K)
x (k)
Σ
S (K)
(k)
MN
1+ (z) (8))
N + (2 (k)
Σt= 2
(k)
5

---

Published as a conference paper at ICLR 2023
Table 1: A unified view for MLP, GCN and GAT from our energy-driven geometric diffusion framework regarding energy function forms, diffusivity specifications and algorithmic complexity.
Models
Energy Function E8Z, 6)
Dillusivity
Complexity
MLP
1Z-
13
S
if =i=
O(NKd")
9
olferwiee
(3.5) 6
GCN
Dogner [| -2,18
S/y
d.d.
O(je)Kdd)
athervntsc
Ensne 8([le [[]ll;-z,]l =2
Letenee 2 1
((e))dd))
otherwise
DIFFORMER
12 BB A..., 5( 1
70 A0*!"
ISCJSN
DIFFORMER OLNKE")
DIFFORMER
D(N"Ka")
The two summation terms above can be computed once and shared to every instance i, reducing the complexity in each iteration to (NN). We refer to this implementation as DIFFORMER-S.
Advanced Diffusivity Model. The simple model facilitates efficiieccysccalbbiittyy yet may sacrifice the capacity for complex latent geometry. We thus propose an advanced version with g(x) 1+exp(-2)
(k) ig = f(12z
(10)
(1))
1 + cxp
which corresponds with f(=²) 1+e f0-1 guaranteeing monotonic decrease and non-negativity.
We dub this version DIFFORMER-a. Appendix D further compares the two models (i.e., different f's and d's) through synthetic results. Real-world empirical comparisons are in Section 5.
4.2 MODEL EXTENSIONS AND FURTHER DISCUSSION
Incorporating Layer-wise Transformations. Eq. 7 does not use feature transformations for each layer. To further improve the representation capacity, we can add such transformations after the updating, i.e., 2((²) (x(0) (x))) where (*) can be fully-connected layer (see Appendix E for details). In this way, each iteration of the diffusion yields descent of a particular energy E(Z, 5, (())) = Z (4) (Z ))) + Ei, 8(112; 2,112) dependent on k. The trainable trans- formation (*) can be optimized w.r.t. the supervised loss map the instance representations into a proper latent space. Our experiments find that the layer-wise transformation is not necessary for small datasets, but contributes to positive effeect for datasets with larger sizes. Furthermore, one can consider non-linear activations in the layer-wiee wise transformation (A)) though empirically found that using a linear model already performs well. We also note that Theorem can be extended to hold even when incorporating such non-linyyryty in each layer (see Appendix C for detailed discussions).
Incorporating Input Graphs. For the model presented so far, we do not assume an input graph for the model formulation. For situations with observed structures as available input, we have G (V,8) that can be leveraged as geometric prior. We can thus modify the updating rule as:
(1+1) 2
T
Σ j=1
( +
A.))
(6) 2
S (k) + A.
(11)
where Aij Vd.d, if (i,j) E E and 0 otherwise, and d, is instance i's degree in G. The diffusion iteration Eq. 11 is essentially a descent step on a new energy additionally incorporating a graph- based penalty (Ioannidis et al., 2017), i.e E(iEE ||z]
| (see Appendix D for details).
Scaling to Large Datasets. Another advantage of DIFFORMER over GNNs is the flexibility for mini-batch training. For datasets with prohibitive instance numbers that make it hard for full-batch training on single GPU, the common practice for GNNs resorts to subgraph sampling Zeng et al. (2020) or graph clustering Chiang et al. (2019) to reduce the overhead. These strategies, however, require extra time and tricky designs to preserve the internal structures. In contrast, thanks to the less reliance on input graphs, for training DIFFORMER, we can naturally partition the dataset into random mini-batches and feed one mini-batch for one feed-forward and backward computation. The

---

Published as a conference paper at ICLR 2023
Table Mean and standard deviation of testing accuracy on node classification (with five different random initializations). All the models are split into groups with comparison of non- non-linearity (whether the model requires activation for wise transformations), PDE-solver (whether the model requires PDE-solver) and Input-G (whether the propagation purely relies on input graphs).
Type
Model
Non-linearity
PDE-solver
Input-G
Cora
Citeseer
Pubmed
MLP
56.1 d: 1.6
56.7
1.7
69.8 + 1.5 65.8
Basic models
LP
R
68.2
42.8
ManiReg
R
R
60.4 + 0.8
67.2 +: 1.6
71.3 =: 1.4
GCN
R
81.5 + 1.3
1.9
77.8 2.9
GAT
R
+ 0.7
0.7
79.0
SGC
R
81.0 1 0.0
0.1
78.9 0.0
GCN-ENN
72.2 + 1.8
56.8 3.2
74.5 3.2
Standard GNNs
GAT-ENN
73.8 1.7
3.88
Dense GAT
78.5 + 2.5
6.
66.4
LDS
83.9 = 0.6
74.8 14.00.0
out-of-memory 78.4
GLCN
83.1 + 0.5
0.9
GRANDA
R
R
83.6 + 1.0
73.4
78.8 1.7
GRAND
R
R
R
83.3 + 1.3
7..1
78.1
GRAND++
Diffusion-basee models
GDC
R
83.6 + 0.2
7.3
78.7 0.4
GraphHeat
83.7
72.5
80.5
DGC-Euler
83.3 0.0
7.13
80.3
NodeFormer
83.4 0.2
0.3
81.5 0.4
Graph Transformers
DIFFORMER-S
85.9 + 0.4
73.5
81.8 + 0.3
DIFFORMER-a
0.6
75.7
0.3
80.5 1.2
flexibility for mini-bated training also brings up convenience for parallel acceleration and federated learning. In fact, such a computational advantage is shared by Transformer-like models for cases where instances are inter-dependent, as demonstrated by the NodeFormer (Wu et al., 2022b).
Connection with Existing Models. Our theory in fact gives rise to general diffusion framework that unifies some existing models as special cases of ours. As a non-exhaustive summary and high-level comparison, Table presents the relationships with MLP, GCN and GAT (see more elaboration in Appendix F). Specifically, the multi-layer perceptrons (MLP) can be seen as only considering the local consistency regularization in the energy function and only allowing non-zero diffusivity for self-loops. For graph convolution networks (GCN) (Kipf & Welling, 2017), it only regularizes the global consistency term constrained within 1-hop neighbors of the observed structure and simplifies as an identity function. The diffusivity of GCN induces information flows through observed edges. For graph attention networks (GAT) (Velickovic et al., 2018). the energy function is similar GCN's except that the non-linearity remains (as a certain specific form), and the diffusivity is computed by attention over edges. In contrast with the above models, DIFFORMER is derived from the general diffusion model that enables non-zero diffusivity values between arbitrary instance pairs.
EXPERIMENTS
We apply DIFFORMER to various tasks for evaluation: graph-based node classification where an input graph is given as observation; 2) image and text classification without input graphs; 3) spatial temporal dynamics prediction. each case, we compare different set of competing models closely associated with DIFFORMER and specifically designed for the particular task. Unless otherwise stated, for datasets where input graphs are available, we incorporate them for feature propagation as is defined by Eq. 11. Due to space limit, we defer details of datasets to Appendix G and the implementations to Appendix H. Also, we provide additional empirical results Appendix I.
5.1 SEMI-SUPERVISED NODE CLASSIFICATION BENCHMARKS
We test DIFFORMER on three citation networks Cora, Ci teseer and Pubmed. Table reports the testing accuracy. We compare with several sets of baselines linked with our model from different aspects. 1) Basic models: MLP and two classical graph-based SSL models Label Propagation (LP) (Zhu et al., 2003) and ManiReg (Belkin et al., 2006). 2) GNN models: SGC (Wu et al., 2019), GCN (Kipf & Welling, 2017), GAT' (Velickovic et al., 2018), their variants GCN-ENN, GAT-ENN (operating kNN graphs constructed from input features) and Dense GAT (with densely connected graph replacing the imput one), and two strong structure learning models LDS (Franceschi et al., 2019) and GLCN (Jiang et al., 2019). 3) PDE graph models: the SOTA models GRAND (Chamberlain et al.,
7

---

Published as a conference paper at ICLR 2023
Table 3: Testing ROC-AUC for
00.0
Idarity
Proteins and Accuracy for
Consarr
Pokec on large-scale node clas- sification datasets, denotes us- ing mini-batch training.
75.0
72.5
50
ORNPAY
Models
Proteins
Pokec
G25
MLP
72.41 1 0.10
60.15 1. 0.03
52.73
60.8
100
Label Number
500
1000
SCC
49.03 0.93
52.03 1: 0.84
GCN
74.22 1 0.19
62.31 1: 1.13"
(a)
(b)
GAT
75.11 + 1.45"
6..37
NodeFomme
77.45
1.15"
68.32 0.45"
Figure 2: (a) Ablation studies w.r.t. different diffusivity function forms on CIFAR. (b) Impact of K and T on Cora.
DIFFORMERSS
79.49 1 0.44
09.24.10.70 0.76
Table 4: Testing accuracy for image (CIFAR and STL) and text (20News) classification.
MI.P
I.P
ManiRaz
GON-ENN
GATANN
GLCN
DIFFORMER-4
DIFFORMER-a
100 labels
6
06.2
1.
66.7215
06002.1
oui-o--menory ou--t-menery
06.61:
CIFAR
500 labels
73.2
70.6
7.6
72.-004
7.5
748
74.0 0.6
1000 labels
75.4 0.5
71.9
74.3 0.4
2.
2.51
out menoory
14.7..0.3
76.6
75.9 15930..3
100 labels
6.
65.2
1.9
05
0..
out menser
66.4::0..8
67.
L.1
STL
500 labels
23.
71.8
0.5
02.1=08
72.0008
out-f-mennecy out-of-mmmery
72.4+t:1.3
73.7 0.6
729 127
1000 labels
0.8
72.7
0.
0.6
7..3
76.4 0.5
75,3 0.6
1000
14.1
55.9
1.2
2.6
5..
56.
0.3
20News
labels
4.9
57.6
0..0
00.613
59.1
602..007
4000 labels
59.5
6.76+007
64.3= 1.0
62.9=007
62.4
64.1
68.9
6488x 1.0
2021a) (with its linear variant GRAND-I) and GRAND++ (Thorpe et al., 2022). 4) Diffusion-insppree GNN models: GDC (Klicpera et al., 2019), GraphHeat (Xu et al., 2020) and a recent work DGC- Euler (Wang al., 2021). Table 2 shows that DIFFORMER achieves the best results on three datasets with significant improvements. Also, we notice that the simple diffusivity model DIFFORMER-S significantly exceeds the counterparts without non-linearity (SGC, GRAND-I and DGC-Fuler) and even comes to the first on Cora and Pubmed. These results suggest that DIFFORMER can serve as a very competitive encoder backbone for node-level prediction that learns inter-instance interactions for generating informative representations and boosting downstream performance.
5.2 LARGE-SCALE NODE CLASSIFCATION GRAPIIS
We also consider two large-scale graph datasets ogbn-Proteing a multi-task interaction network, and Pokec, a social network. Table 3 presents the results. Due to the dataset size (0.1 (0.13M/1.66 3MM/1.33M nodes for two graphs) and scalability issues that many of the competitors in Table 2 as well as DIFFORMER-a would potentially experience, we only compare DIFFOR MER-S with standard GNNs. In particular, we found GCN/GAT/DIFFORMER are still hard for full-graph training on single V100 GPU with 16GM memory. We thus consider mini-batch training with batch size OOK/ for P Proteinse wnwe/onncc........ We found that DIFFORMER outperforms common GNNs by large margin, which suggests its desired efficacy on large datasets. As mentioned previously. we prioritize the efficacy of DIFFORMER as general encoder backbone for solving node-level prediction tasks on large graphs. While there are quite a few practical tricks shown to be effective for training GNNs for this purpose, e.g., hop-wise attention (Sun et al., 2022) or various label re-use strategies, these efforts are largely orthogonal to our contribution here and can be applied to most any model to further boost performance. For further investigation, we supplement more results using different mini- sizes for training and study its impact on testing performance in Appendix I.4. Furthermore, we compare the training time and memory costs in Appendix I.5 where found that DIFFORMER-S is about 6 times faster than GAT and 39 times faster than DenseGAT on Pokec, which suggests superior scalability and efficiency of DIFFORMER-S on large graphs
5.3 IMAGE AND TEXT CLASSIFICATION WITH Low LABEL RATES
We next conduct experiments on CIFAR-10, 0, STL-10 and 20News - Grroup datasets to test DIF- FORMER for standard classification tasks with limited label rates. For 20News provided by Pedregosa et al. (2011), we take 10 topics and use words with TF-IDF more than 5 as features. For CIFAR and STL, two public image datasets, we first use the sulf-super approach SimCLR (Chen et al., 2020b) (that does not use labels for training) to train ResNet- 18 for extracting the feature maps as input features of instances. These datasets contain no graph structure, so we use ANN to construct a graph over input features for GNN competitors and do not input graphs for DIFFORMER.
8

---

Published as a conference paper at ICLR 2023
Table 5: Mean and standard deviation of MSE on spatial-temporal prediction datasets.
GCN
Denue
GATAAN
GCN
DIFFOUMER
DIFFORMER-
DIFFURNER w DIFFORMER- wio
Chickrnpex
(10.001)
(20.00))
(40,005)
(E0000)
(20,004)
Civid
650
0956
(90.167)
1,000
(30.330)
1,822
(30,319)
(50133)
0361
(605500)
1,75
(10L186)
(-1090)
(30,143)
WillMach
1073
(30.125)
1,292
1.390
(30170))
0.826
(20015)
0.522
(200053)
1023
(50005))
0,922
(ERITT)
0.738
(20
0:993
(240.18))
0.720
Table 4 reports the testing accuracy of DIFFORMER and competitors including MLP, ManiReg. GCN-ENN, GAT-ENN, DenseGAT and GL.CN. Two DIFFORMER models perform much better than MLP in nearly all cases, suggesting the effectiveness of learning the inter-dependence over instances. Besides, DIFFORMER yields large improvements over GCN and GAT which are in some sense limited by the handcrafted graph that leads to sub-optimal propagation. Moreover, DIFFORMER significantly outperforms GLCN, strong baseline that learns new (static) graph structures, which demonstrates the superiority of our evolving diffusivity that can adapt to different layers.
5.4 SPATIAL-TEMPORAL DYNAMICS PREDICTION
We consider three apatial-temport datasets with details in Appendix G. Each dataset consists of a series of graph snapshots where nodes are treated as instances and each of them has integer label (e.g., reported cases for Ch ckenpox or Covid). The task to predict the labels of one snapshot based the previous ones. Table 5 compares testing MSE of four DIFFORMER variants (here DIFFORMER-S w/o g denotes the model DIFFORMER-S without using input graphs) against baselines. We see that two DIFFORMER variants without input graphs even outperform the counterparts using input structures in four out of six cases. This implies that our diffusivity estimation module could learn useful structures for informed prediction, and the input structure might not always contribute to positive effect. In fact, for temporal dynamics, the underlying relations that truly influence the trajectory evolution can be much complex and the observed relations could be unreliable with missing or noisy links, in which case GNN models relying on input graphs may perform undesirably. Compared the competitors, our models rank the first with significant improvements.
5.5 FURTHER RESULTS AND DISCUSSIONS
How do different diffusivity functions perform? Figure 2(a) compares DIFFORMER with four variants using other diffusivity functions that have no essential connection with energy minimization:
1) Identity sets (kk) as a fixed identity matrix; 2) Constant fixes S(k) as all-one constant matrix; ((k) by attention networks (Vaswani al., 2017); 4) Kernel adopts Gaussian kernel for computing S(k), More results on other datasets are in Appendix I.I and they consistently show that our adopted diffusivity forms produce superior performance, which verifies the effectiveness of our diffusivity designs derived from minimization of principled energy.
3) Full Alin parameterizes
How do model depth and step size impact the performance? We discuss the influence of model depth K and step size on Cora in Fig. 2(b). More results Citeseer and Pubmed are generally consistent with Fig. 2(b) and deferred to Appendix 1.2. The curves indicate that the performance of GCN and GAT models exhibit significant degradation with deeper layers, while DIFFORMER maintains its superiority and performs stably with large K. Furthermore, when K is not large enough (less than 32), there is clear performance improvement of DIFFORMER as K increases, and larger contributes to a steeper increase. When K continues increasing (more than 32), the model performance still goes up with small 7 (0.1 and 0.25) yet exhibits a slight drop with large т (0.5 and 0.9). The reason could be that larger (smaller) 7 contributes to more (less) concentration on global information from other instances in each iteration, which brings up more (less) benefits with increasing propagation layers yet could lead to instability when the step size is too large.
6 CONCLUSIONS
This paper proposes energy-driven geometric diffusion model with latent diffusivity function for data representations. The model encodes all the instances as a whole into evolving states aimed at minimizing principled energy implicit regularization. We further design two practical implemen- tations with enough scalability and capacity for learning complex interactions over the underlying data geometry. Extensive experiments demonstrate the effectiveness and superiority of the model.

---

Published as a conference paper at ICLR 2023
REFERENCES
James Atwood and Don Towsley. Diffusion-coonvllttiona neural networks. In Advances in neural information processing systems, 2016.
Mikhail Belkin, Partha Niyogi, and Vikas Sindhwani. Manifold regularization: geometric frame- work for learning from labeled and unlabeled examples. Jourmal of machine learning research, 7 (11), 2006.
Aleksandar Bojchevski, Johannes Klicpera, Bryan Perozzi, Amol Kapoor, Martin Blais, Benedek Rózemberczki, Michal Lukasik, and Stephan Günnemann. Scaling graph neural networks with approximate pagerank. In ACM SIGKDD Conference Knowledge Discovery and Data Mining, 2020.
Michael M. Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geometric deep learning: going beyond euclidean data. IEEE Signal Process. Mag, 34:18 42, 2017.
Ben Chamberlain, James Rowbottom, Maria I. Gorinova, Michael M. Bronstein, Stefan Webb, and Emanuele Rossi. GRAND: graph neural diffusion. In International Conference on Machine Learning (ICML), 1407-1418, 2021a.
Benjamin Paul Chamberlain, James Rowbottom, Davide Eynard, Francesco Di Giovanni, Xiaowen Dong, and Michael M. Bronstein. Beltrami flow and neural diffusion on graphs. In Advances in Neural Information Processing Systems (NeurIPS), 2021b.
Jinsong Chen, Kaiyuan Gao, Gaichao Li, and Kun He. NAGphormer: tokenized graph transformer for node classification in large graphs. In International Conference on Learning Representations, 2023.
Ming Chen. Zhewei Wei, Zengfeng Huang, Bolin Ding, and Yaliang Li. Simple and deep graph convolutional networks. In International Conference on Machine Learning, pp. 1725-1735, 2020a.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in neural information processing systems, 2018
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607, 2020b.
Yu Chen, Lingfei Wu, and Mohammed J. Zaki. Iterative deep graph learning for graph neural networks: Better and robust node embeddings. In Advances in Neural Information Processing Systems, 2020c.
Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh. Cluster-gen: An efficient algorithm for training deep and large graph convolutional networks. In ACM SIGKDD International Conference Knowledge Discovery & Data Mining, pp. 257-266, 2019.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Ben- jamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference Learning Representations, 2021.
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image worth 16x words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Vijay Prakash Dwivedi and Xavier Bresson. generalization of transformer networks to graphs. CoRR, abs/2012.09699,, 2020.
Moshe Eliasof, Eldad Haber, and Eran Treister. PDE-GCN:
novel architectures for graph neural networks motivated by partial differential equations. In Advances Neural Information Processing Systems (NeurIPS), 2021.
10

---

Published as a conference paper at ICLR 2023
Bahare Fatemi, Layla El Asri, and Seyed Mchran Kazemi. Slaps: Self supfrsuuerviioon improves structure learning for graph neural networks. In Advances in Neural Information Processing Systems, 2021.
Luca Franceschi, Mathias Niepert, Massimiliano Pontil, and Xiao He. Learning discrete structures for graph neural networks. In International Conference on Machine Learning pp. 1972-1982, 2019.
Mark Freidlin and Alexander D Wentzell. Diffusion processes graphs and the averaging principle. The Annals of probability, pp. 2215-2245, 1993.
Boumediene Hamzi and Houman Owhadi. Learning dynamical systems from data: A simple cross- validation perspective, part i: Parametric kemel flows. Physica D: Nonlinear Phenomena, 421: 132817, 2021.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. In Advances in Neural Information Processing Systems, 2020.
Vassilis N. Ioannidis, Meng Ma, Athanasios N. Nikolakopoulos, and Georgios B. Giannakis. Kernel- based inference of functions over graphs. CoRR, abs/1711. 2017.
Bo Jiang, Ziyan Zhang, Doudou Lin, Jin Tang, and Bin Luo. Semi-supervised learning with graph learning-convolutione networks. In /EEF/CVF conference on computer vision and pattern recognition, pp. 1311-11320, 2019.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. International Conference on Machine Learning, pp. 5156-5165, 2020.
Thomas N. Kipf and Max Welling, Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations (ICLR), 2017.
Johannes Klicpera, Stefan Weißenberger, and Stephan Günnemann. Diffusion improves graph learning. Advances in neural information processing systems, 2019.
Isaac E Lagaris, Aristidis Likas, and Dimitrios Fotiadis. Artificial neural networks for solving ordinary and partial differential equations. IEEE transactions neural networks, 95):987711000, 1998.
Danning Lao, Xinyu Yang. Qitian Wu, and Junchi Yan. Variational inference for training graph neural networks in low-data regime through joint structure-label estimation. Proceedings of the 28th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 824-834--4 2022.
Yao Ma, Xiaorui Liu, Tong Zhao, Yozen Liu, Jiliang Tang, and Neil Shah. A unified view on graph neural networks as graph signal denoising. In ACM International Conference on Information and Knowledge Management, 2021.
Georgi S Medvedev. The nonlinear heat equation on dense graphs and graph limits. SIAM Journal on Mathematical Analysis, 46(4))2743 274222766,, 2014.
Fabian Pedregosa, Gaël Varoquaux. Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, al. Scikit-learn: Machine learning in python. the Journal of machine Learning reseurch, 12:2825-2830, 2011.
Hieu Pham, Zihang Dai, Qizhe Xie, and Quoc V. Le. Meta pseudo labels. In IEEE Conference on Computer Vision and Patlern Recognition (CVPR), pp. 11557-11568. 2021.
R. Rockafellar. Convex analysis. Princeton University Press, 1970.
Steven Rosenberg and Rosenberg Steven. The Laplacian on a Riemannian manifold: an introduction to analysis on manifolds. Number 31. Cambridge University Press, 1997.

---

Published as a conference paper at ICLR 2023
Benedek Rozemberczki, Paul Scherer, Yixuan He, George Panagopoulos, Alexander Riedel, Maria Astefamouei, Oliver Kiss, Ferene Beres, Guzmán López, Nicolas Collignon, et al. Pytorch geometric temporal: Spatiotemporal signal processing with neural machine learning models. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 4564- 4573, 2021.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE transactions on neural networks, 20(1)x61-80) 2008.
Prithviraj Sen, Galileo Namata, Mustafa Bilgic, Lise Getoor, Brian Gallagher, and Tina Eliassi- Rad. Collective classification network data. Al Mag., 29(3):93-106, 2008
Chuxiong Sun, Jie Hu, Hongming Gu, Jinpeng Chen, and Mingchuan Yang. Adaptive graph diffusion networks. CoRR, abs/2012.15024 2022.
Matthew Thorpe, Hedi Xia, Tan Nguyen, Thomas Strohmer, Andrea I.. Bertozzi, Stanley J. Osher, and Bao Wang. GRAND++: graph neural diffusion with a source term. In International Conference on Learning Representations (ICLR), 2022.
Anton Tsitsulin. Davide Mottin, Panagiotis Karras, Alexander M. Bronstein, and Emmanuel Müller. Netlsd: Hearing the shape of graph. In ACM SIGKDD International Conference Knowledge Discovery Data Mining, pp. 2347-2356, 2018
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pp. 1998-6008, 2017.
Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference Learning Representations (ICLR), 2018.
Yifei Wang. Yisen Wang, Jiansheng Yang. and Zhouchen Lin. Dissecting the diffusion process in linear graph convolutional networks. In Advances Neural Information Processing Systems, 2021.
Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E. Sarma, Michael M. Bronstein, and Justin M. Solomon. Dynamic graph CNN for learning point clouds. ACM Trans. Graph., 38(5): (::5):14::11444 14::1116::12 2019.
Yuelin Wang, Kai Yi, Xinliang Liu, Yu Guang Wang, and Shi Jin. ACMP: Allen-cahn message passing with attractive and repulsive forces for graph neural networks. International Conference Learning Representations, 2023.
Jason Weston, Frédéric Ratle, Hossein Mobahi, and Ronan Collobert. Deep learning via semi- supervised embedding. In Neural networks: Tricks of the trade, pp. 639-655. Springer, 2012.
Felix Wu, Amauri H. Souza Jr., Tianyi Zhang, Christopher Fifty, Tao Yu, and Kilian Q. Weinberger. Simplifying graph convolutional networks. In International Conference on Machine Learning, pp. 6861-6871, 2019.
Qitian Wu, Hengrui Zhang, Junchi Yan, and David Wipf. Handling distribution shifts graphs: An invariance perspective. International Conference on Learning Representations, 2022a.
Qitian Wu, Wentao Zhao, Zenan Li, David Wipf, and Junchi Yan. Nodeformer: A scalable graph struc- ture learning transformer for node classification. In Advances in Neural Information Processing Systems, 2022b.
Bingbing Xu, Huawei Shen, Qi Cao, Keting Cen, and Xueqi Cheng. Graph convolutional networks using heat kernel for supi-super learning. arXiv preprint arXiv:2007.16000 2020.
Chenxiao Yang, Qitian Wu, and Junchi Yan. Geometric knowledge distillation: Topology compression for graph neural networks. In Advances in Neural Information Processing Systems, 2022.
Chenxiao Yang, Qitian Wu, Jiahua Wang, and Junchi Yan. Graph neural networks are inherently good generalizers: Insights by bridging GNNs and MLPs. In International Conference on Learning Representations, 2023.
12

---

Published as a conference paper at ICLR 2023
Yongyi Yang, Tang Liu, Yangkun Wang, Jinjing Zhou, Quan Gan, Zhewei Wei, Zheng Zhang, Zengfeng Huang, and David Wipf. Graph neural networks inspired by classical iterative algorithms. In International Conference on Machine Learning, pp. 11773-11788 2021.
Zhilin Yang, William W. Cohen, and Ruslan Salakhutdinov. Revisiting semi-supervised learning with graph embeddings. In International Conference on Machine Learning (ICML), pp. 40-48, 2016.
Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yas Yan Liu. Do transformers really perform bad for graph representation? In Advances Neural Information Processing Systems, 2021.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang. and Amr Ahmed. Big bird: Transformers for longer sequences. In Advances Neural Information Processing Systems, 2020.
Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor K. Prasanna. Graphsaint: Graph sampling based inductive learning method. In International Conference on Learning Representations, 2020.
Jiawei Zhang, Haopeng Zhang, Congying Xia, and Li Sun. Graph-bert: Only attention is needed for learning graph representations. CoRR, abs/2001 10000000.511,,,, 2020.
Tianqi Zhang, Qitian Wu, Junchi Yan, Yunan Zhao, and Bing Han. Scalegen: Efficient and effective graph convolution via channel-wist scale transformation. IEEE Transactions on Neural Networks and Learning Systems, 2022.
Yingxue Zhang, Soumyasundar Pal, Mark Coates, and Deniz Üstebay, Bayesian graph convolutional neural networks for semi cemi-supervised classification. In AAAI Conference Artificial Intelligence, pp. 5829-58366 2019.
Yunhao Zhang and Junchi Yan. Crossformer: Transformer utilizing cross-dimension dependency for multivariate time series forecasting. In International Conference on Learning Representations, 2023.
Dengyong Zhou, Olivier Bousquet, Thomas Navin Lal, Jason Weston, and Bernhard Schölkopf. Learning with local and global consistency. In Advances in Neural Information Processing Systems, pp. 321-328, 2004.
Xiaojin Zhu, Zoubin Ghahramani, and John D. Lafferty. Semi-supervised learning using gaussian fields and harmonic functions. In International Conference on Machine Learning (ICML), pp. 912-919, 2003.
A FURTHER RELATED WORKS AND CONNECTION WITH OURS
We discuss more related works that associate with ours from different aspects to properly position this paper with different areas. Based this, we further shed more lights the technical contributions of our work and its potential impact in different communities.
A.1 NEURAL DIFFUSION MODELS
PDE-based Learning. The diffusion- learning has gained increasing research interests, as the continuous dynamics can serve as an inductive bias incorporated with prior knowledge of the tasks at hand. One category directly solves continuous process of differential equations (Lagaris et al., 1998; Chen et al., 2018), e.g., Chamberlain et al. (2021a) and its follow-ups (Chamberlain et al., 202 Thorpe et al., 2022) reveal the analogy between the discretization of diffusion process and GNNs' feedforward rules, and devise new (continuous) models on graphs whose training requires PDE-solving tools. A concurrent work (Wang et al., 2023) explores how to derive neural networks from gradient flows and proposes Allen-Cahn Message Passing that combines attractive and repulsive effects.
13

---

Published as a conference paper at ICLR 2023
PDE-inspired Learning. Another category is PDE-inspired learning using the diffusion perspective as principled guidelines on top of which new (discrete) neural network- approaches are designed for node classification (Atwood & Towsley, 2016; Klicpera et al.. 2019; Xu et al., 2020). graph comparison (Tsitsulin et al., 2), and geometric knowledge distillation (Yang et al., 2022). Our work leans on PDE-inspired learning and introduce new diffusion model that is implicitly defined as minimizing a regularized energy. Our theory also reveals the underlying equivalence between the numerical iterations of the diffusion process and unfolding the minimization dynamics of corresponding energy. The new results can serve to extend the class of diffusion process and illuminate its fundamental connection with energy optimization systems. More importantly, as we will in the maintext, such new perspective brings up unifying view for existing models like MLP and GNNs.
A.2 THEORETICAL INTERPRETATIONS OF GNNS
Optinization-indueed Models. As GNNS have proven to be powerful encoders for modeling data with geometry, there arise quite few studies that aim at understanding the efficacy of GNNs with respect to node representations yielded by the propagation. For instance, Yang et al. (2021) founds that the layer-wise updates of common GNNs can be derived from classical iterative algorithm solving an optimization problem defined over an energy. Furthermore, Ma et al. (2021) points out the relationship between GNNs' message passing rules and the graph denoising problem. These works interpret GNNs as optimizationiinee induced models. In our model, the diffusion process is implicitly defined by an optimization problem and our theory further reveals the fundamental equivalence between diffusion iterations and energy optimization dynamics, which can bridge two schools of thinking and facilitate the understandings from both perspectives.
Generalization of GNNs. Another line of works turn attention to GNNs generalization capability. As increasing evidence shows that GNNs are sensitive to distribution shifts (Wu et al., 2022a), becomes urgent to investigate and figure out how to improve GNNs' generalization. A concurrent work (Yang et al., 2023) identifies an intriguing phenomenon that GNNs without propagation (i.e., an MLP architecture) during training can still yield competitive or even close performance inference time to standard GNNs (using feature propagation at both training and inference stage). This suggests that the propagation design contributes to better generalization instead of the model capacity, and the former plays as critical factor that makes GNNs more powerful encoder than MLP. This result can further interprets the success of DIFFORMER: the new propagation of DIFFORMER that aggregates other nodes' embeddings can enhance the generalization at inference time on testing data.
A.3 TRANSFORMERS
General Transformers. The key design of Transformers lies in the (all-pair) attention mechanism that captures the interactions among tokens of an input sequence. The attention design is originally motivated by the nnter-dependence among words of sentence in NLP tasks (Vaswani et al., 2017). Furthermore, a surge of follow-up works extend the attention mechanisms for capturing the inter- dependence among patches of image (Dosovitskiy et al., 2021), atoms of molecule (Dwivedi & Bresson, 2020), dimensions of multivariate time series (Zhang & Yan, 2023), etc. From different aspects, our work explores Transformer model that targets learning inter-dependence among instances in a dataset.
Efficient Transformers. When the length of inputs goes large, the quadratic complexity of all-pair attention becomes the computational bottleneck of Transformers. To address the long sequences in NLP tasks, quite few recent works propose various strategies via, e.g.. Softmax kernel (Choromanski et al., 2021), local-global attention (Zaheer et al., 2020), and low-rank approximation (Katharopoulos et al., 2020). our problem setting, if one treats the whole dataset as an input 'sequence' for the Transformer model, then the length of such a sequence could go demanding scalability challenges. Our model DIFFORMER-S achieves strictly linear complexity w.r.t. N through introducing a new attention function that keeps simple form yet accommodates the influence among arbitrary node pairs.
million-level, which poses
Transformers on Graphs. There is also increasing research interest on building powerful Transform- ers for graph-structured data, due to the inherent good ability of Transformer models for capturing long-range dependence and potential node-pair interactions that are unobserved in input structures.
14

---

Published as a conference paper at ICLR 2023
Recent works, e.g., Dwivedi & Bresson (2020); Zhang et al. (2020); Ying al. (2021) explore various effective architectures for graph-level prediction (GP) tasks, e.g., molecular property prediction, whereby each graph itself generally has label to predict and dataset contains many graph instances. Our target problem, as mentioned in Section 2, can be seen an embodiment of node-level prediction (NP) studied in graph learning community. These two problems typically tackled separately in the literature (Hu et al., 2020) with disparate technical considerations. This is because input instances are inter-dependent in NP (due the instance interactions involved in the data-generating process), while in GP tasks the instances can be treated as IID samples. For node-level prediction, a recent work (Wu et al., 2022b) proposes kernelized Sommbel-ofttax-- based message passing achieve all all-pair feature propagation with linear complexity. Concurrently at the same conference, (Chen et al., 2023) proposes to transform neighborhood features into tokens fed into Transformers to enable efficient training on large graphs
A.4 GRAPH NEURAL NETWORKS
Scalable Graph Neural Networks. Due to the large graph sizes that could be frequently encountered in graph-based predictive tasks (particularly NP), many recent works focus designing scalable GNNs. From different technical aspects, these works can be generally grouped into graph sampling- based partition (Zeng et al., 2020; Chiang et al., 2019). approximation-baate et al., 2020) and simplified architectures (Wu et al., 2019; Zhang et al., 2022). In contrast with them, DIFFORMER can scale to medium medium-sized graphs with the linear complexity, and targe-scale graphs with mini-batch training that is simple, stable and also flexible for balancing efficiency and precision with proper mini-batch sizes. The later is allowed by the all-pair message passing schemes that do not rely on input structures (the input graphs play auxiliary role in the model).
propagation (Bojehevski
Graph Structure Learning. Learning latent graph structures is well-established research problem in graph learning community. Motivated by the observation that input graphs often be unreli- able or unavailable, based on which the message passing of GNNs could yield undesired results, graph strueture learning aims learning adaptive structures that can boost GNNs towards better representations and downstream prediction (Franceschi et al., 2019; Chen et al., 2020c; Jiang et al., 2019; Fatemi et al., 2021: Lao et al., 2022). Since graph structure learning requires estimation for N potential edges that connect all node pairs in a dataset, the time and space complexity of most existing models are at least quadratic w.r.t. node numbers, which is prohibitive for large graphs. The recent work (Wu et al., 2022b) proposes a Transformer model with linear complexity that can scale graph structure learning to graphs with millions of nodes. The proposed model DIFFORMER-S can be another powerful approach for learning latent structures in large-scale systems, without sacrificing the accommoaation of all-pair interactions at each layer.
B PROOF FOR THEOREM
First of all, we can convert the minimization of Eq. 4 into minimization of its variational upper bound, shown in the following proposition.
Proposition 1. The energy function E(Z, 0) is upper bounded by
BiZ,k; (wish 8) = 1z z 11 +
wigllz, - -,,122 -2 Σ
(12)
3,3
where 8 is the concave conjugale of 6, and the equality holds if and only if the variational parameters satisfy
05(=2)
Wij
(13)
0.2
2, 12,12,-2,11
Proof. The proof of the proposition follows the principles of convex analysis and Fenchel dual- ity (Rockafellar, 1970). For any concave and non-decreasing function p.: IR - IR one can express it as the variational decomposition
p(=²)
min min[wz2 - [(w)) 2 wz² -p(w),
(14)
w20
15

---

Published as a conference paper at ICLR 2023
(a) DIFFORMER-S: 8(=2) 2= 1
(b) DIFFORMER-a: 8(=²) = 2 log( e
/2- + +1
divergence field (produced by 10 randomly generated instances marked as red stars) and cross-section energy field of an individual.
Figure 3: Plot of penalty curves 6(2²) and f(=²)
dz
where w is a variational parameter and أ is the concave conjugate of p. Eq. 14 essentially defines p(=?) as the minimal envelope of series of quadratic bounds w p(w:) defined by different values 2 0 and the upper bound is given for a fixed w when removing the minimization operator. Based on this, obtain the result of Eq. 12. In terms of the sufficient and necessary condition for equality, we note that for any optimal wt we have
=² "kww) p(=?)
(15)
which is tangent to at x² and w
a5(=²) 0z2
We thus obtain the result of Eq. 13.
We next continue to prove the main result of Theorem 1. According to Proposition 1, we can minimize the upper bound surrogate Eq. 12 and it becomes equivalent to a minimization of the original energy on condition that the variational parameters given by Eq. 13. Then with one-step gradient decent of Eq. 12, the instance states could be updated via (assuming / as the index of steps and a as step size)
z(41) = ( a Oz
2=
(16)
= ( (A(D() (1) - (20))(Z0))) + (0) z
0 a (D (1) Ω (z)
where ?2 [wi JNNN ( denotes the diagonal degree matrix associated with (2()) and we Common practice accelerate convergence
(1)
introduce a' at to combine two parameters
adopts positive definite preconditioner term, e.g., ( to re-scale the updating gradient and the final updating form becomes
Z (4+1) = (1 (1-a''z" + + -2
(17)
One can notice that Eq. 17 shares similar forms as the numerical iteration Eq. 3 for the PDE diffusion system, in particular we write Eq. 3 as a matrix form:
Z ((++)= = (1 (1- -b() ( + (k)))
(18)
where (4) is the degree matrix associated with ((k), Pushing further, we can see that the effect of Eq. 18 is the same as 17 when we let a'. (k)) (D -00 and S (1)) is row-normalized.
i.e. Σvev ((k) = 1 and D (8) I
Thereby, we have proven by construction that a one-step numerical iteration by the explicit Euler scheme, specifically shown by Eq. 17 is equivalent to a one-step gradient descent on the surrogate
Eq. 14 which further equals to the original energy Eq. 4. We thus have the result E (k+1), k; 8) <
(k) k; d). Besides, we notice that for a fixed Z, E Z, k; 5) = IZ- Z + + E,; p(| -2,
becomes function of k and its optimum is achieved if and only if (A) Z. Such a fact yields that (()(k) k; ō) (k) k 1;8). The result of the main theorem follows by noting that
F (++1) k; σ) < F 5) F (k)
1:8).
16

---

Published as a conference paper at ICLR 2023
C EXTENSION OF OUR THEORY TO INCORPORATE NON-LINEARITY
In Section 3, we mainly focus on the situation without non-linearity each model layer to stay consistent with implementation where empirically found that omitting the non-linearity in the middle diffusion layers works smoothly in practice (see the pseudo code of Alg. in appendix for details). Even so, theory can also be extended incorporate the layer- byser-wise non-linearity in diffusion propagation. Specifically, the non-linear activation can be treated as a proximal operator (which projects the output into feasible region) and the gradient descent used in our analysis Eqn. 16 and 17 can be modified to add a proximal operator:
Z (++1) = Proxo (11 th +
where (x) = arg min || - z||2 and Ω defines feasible region. The updating above corresponds to proximial gradient descent which also guarantees strict minimization for the energy function and our Theorem 1 will still hold. In partieular, if one uses ReLU activation, the proximal operator will be (r)xn(=) max(0,
D DIFFERENT ENERGY FORMS
We present more detailed illustration for the choices of f and specific energy function forms in Eq. 4. Simple Diffusivity Model. As discussed in Section 4, the simple model assumes f(=²)
=²
that corresponds to g(x) = 1+ where we define
|z: 2, and x = 2 2j. The corresponding
penalty function whose first-order derivative f would be 8(22) 22 2 We plot the penalty function curves in Fig. 3(a). As we can see, the ʃ a non-negative, decreasing function of =? which implies that the 8 satisfies the non-decreasing and concavity properties to guarantee valid regularized energy function. Also, Fig. 3(a) we present the divergence field produced by 10 randomly generated instances (marked red stars) and the cross-section energy field of one instance.
Advanced Diffusivity Model. The diffusivity model defines f(z²) 14e- with g(x) =
1+xxp(-z)) and the corresponding penalty function 8(2²) = 2² log (e /2-1 + 1). The penalty function curves, divergence field and energy field are shown in Fig. 3(b).
Incorporating Input Graphs. In Section 4, we present an extended version of our model for incorporating input graphs. Such extension defines new diffusion process whose iterations are equivalent (up to a re-scaling factor on the adjacency matrix) to a sequence of descending steps for the following regularized energy:
EZ k: = ZZ
6(11z - 2,
Σ
12 - -z, 12,
(19)
2
(KJ)EC
where the last term contributes to penalty for observed edges in the input graph.
E MODEL IMPLEMENTATION AND PSEUDO CODES
In this section, we present the details for model implementation including the feed-forward and loss computation of DIFFORMER.
Input Layer For input data X = [x E RNDD where x; denotes the D-dimensional input features of the i-th instance, we first use shallow fully-connected layer to convert it into d- dimensional embedding in the latent space:
= a (LayerNorm W, by)),
(20)
where W, E RdxD and by Rd are trainable parameters, and a is non-linear activation (i.e., ReLU), Then the node embeddings Z will be used for feature propagation with our biffusion-induced Transformer model by letting Z () = Z as the initial states.
Propagation Layer. The updating rule for node embeddings at the k-th layer consists of sequence of operations as described below. For the h-th head,
k. (k)
w(k) (
((k)
(k) 7
(k)
w (k) V.A
(k))
(21)
R.
17

---

Published as a conference paper at ICLR 2023
where W h E RRdxd w à Rdxd W Rdxd trainable parameters of the h-thh head the k-th layer. Then the transformed embeddings will be fed into the all-pair propagation unit with different specifications for DIFFORMER-S and DIFFORMER-a.
For DIFFORMER-S: we first adopt L2 normalization for key and query vectors
N
A
K
(b)
(22)
h
1=1
where K denotes the i-th row vector of K
Then the all all-pair propagation is achieved
by
R (k) diag
h
[
(
(A))
( T1)
(23)
P
(k)
R
(k)
N
v
(k)
Q (Kk) v(()
(k)
(24)
The above computation only requires (NN) complexity.
For DIFFORMER-a: need compute tha all-pair similarity before aggregating the results
A (k) Sigmoid (Q.
(k)
(k(
(25)
R (A: diag
A 1).
(k) h
(26)
P
R
(27)
If using input graphs, we add the updated embeddings of graph-based propagation to the all-pair propagation's 's ones:
P
P
(k)
D
AD
(28)
where the input graph and D denotes its corresponding diagonal degree matrix.
We then average the propagated results of multiple heads:
H P
P
(29)
H
After once propagation, the next-layer embeddings will be updated by
(++1) = o' (LayerNorm (TP" (k) + +(1 (1- -))-))(())
(30)
where can be identity mapping or non non-linear activation (e.g., ReLU).
Output Layer. After K layers of propagation, then shallow fully-connected layer to output the predicted logits:
= ( W + bo.
(31)
where Wo E RdxC and bo E RRC are trainable parameters, and C denotes the number of classes. And, the predicted logits Y will be used for computing loss of the form ((Y, Y) where can be cross-entropy for classification or mean square error for regression.
E. PSEUDO CODES
We provide the Pytorch-style pseudo codes for DIFFORMER class in Alg. I and the one-layer propagation of two model versions (shown in Alg. 2 for DIFFORMER-S and Alg. for DIFFORMER- a). The key design of our methodology lies the model architectures which are shown in detail in Alg. 2 for DIFFORMER-s and Alg. for DIFFORMER-a, where for each case, the model takes the data as input and outputs prediction for each individual instance. For more details concerning the implementation, please refer provided codes at the first page.
18

---

Published as a conference paper at ICLR 2023
Algorithm PyTorch-ttyle Code for DIFFORMER
Fos: Ful cy-connected layers
bus: layer normalizati lzzation layers
conve: DIFF Layers (see implementation in Alg. 2 and 31 activations activation function for the input layer
x: input data sized IN, DJ, for instance aumber, fer input feature dimensicn
4 edge_index: input graph struczure if available, None otherwise
tau: step size for each iterazion update
use_act: whether to use activation for propagation layers
layer_ []
input MGP layer
fcs [0] (x)
[0] (x)
x activatien(x) (X)
store as residual link
apper_.
for 1, conv in enumerate (convs):
graph convolution with global all-pair attention (specified by Alg. 2 for DIFFormer
-5 and Alg. for DIFPormer-a)
conv (x, edge_index)
(1-tau) [iyer_[i]
bus(i++]] (x)
use_act:
x - aattiiaationn
(p
outpus MIP layer Ecs [-1]
supervison sed loss os sslculation, nogativee tog-likalihons
y I_ggp F. softaax (out, dim-1)
1oss criterion logp train_idx),, y_true (train_idx))
Algorithm 2 PyTorch-style Code for One-layer Feed-forward of DIFFORMER-S
x: data embeddings sized [N, d), N for instance number, d for hidden size otherwise
edge_index, input graph struczure if available,
H: head number
4 use_graph: whether to use input graph
: whether to use feature transformation for each laver
graph_come graph convalution operatior using the normall ize adjasency macrix (-
AD*(-1/2;
Wa, Wi,, Wr: weight: matrices for feature trans formatioooioon
2 (x) [N, H, D]
K (k) [N, 8, D]
use weight (V
weight) [N, B,
numeratior
XV torch. wrrcch..insmm((mmm 1"mm, Ihd->hmd", K, V)
num forch,einsun ("nim, hmd snhd" Q KV) $ [N, B, D]
nim
denominator
all_ancs tcrch. ones (N)
suu torch.einsum einsum (" "1hm, 1->hm", X, all_ones)
den torch. ("nhm, hm >nh", G. K[aum) [8,
aggregated results
den +- torch.ones. oealoones__ee lise * agg
(an.unsqueezee())
input graph for graph conv
is graphh
agg agg graph_conv (V, edge_index) output agg meg. (dim-1)
19

---

Published as a conference paper at ICLR 2023
Algorithm 3 PyTorch-style Code for One-layer Feed-forward of DIFFORMER-a
x: data embedding iigss sized [N, d), N for instance number, d for Fidden size index: imput graph struccure if available, None otherwase
H: head number
uso_graph; whether
input graph
use weight:: whether to use feature cranstormation for each layer
graph graph convolution operatior using the normalized adjacency matrix (-(/222)
AD*(-1/2)
Wgr WK, WV: weight matrices for feature tranaformation
0 (x) # [N, H, D]
K Wk (x) [N, H, DJ
use weight wvtx) (1 (sse_weight) x [N, H, [R]
numerator
num - woo.ch.i.goid//ieo (torch. einsum ("nhm, 1hm 1hm->nhal >nlh", Q K) [N, N, H]
denominator
all_ones torch. ones (N)
torch.einstee (*n1h, 1-2nh", num, all aln_s den - den. unsqueeze (1) repeat N, 11 # [N, B, H]
aggregated results
attn num / + [N, N, H] agg terch.cinsum (ennsssmmmhh, 1hé-Snad", attn, V) 8 [N, H,
input graph graph conv
1f use _ae_graph:
agg agg graph_conv edge_index)
output mea. (dim-1)
F CONNECTIONS WITH EXISTING MODELS
MIP. MLPs can be viewed as a simplified diffusion model with only non-zero diffusivity values the diagonal line, i.e. | if j and 0 otherwise. From graph convolution perspective, MLP only considers propagation on the self-loop connection in each layer. Correspondlying, the energy function the feed-forward process of the model essentially optimizes would be E(Z, k) IZ - 1 which only counts for the local consistency term and ignores the global information.
GCN. Graph Convolution Networks (Kipf & Welling, 2017) define a convolution operator on a graph g = (V,&) E) by multiplying the node feature matrix with D ADD- where A denotes the adjacency matrix and D is its associated degree matrix. The layer-wise updating rule can be written as
Vd,d, if (L,)) 0,
((4) =
otherwise.
(32)
N
2/
(++1)
1
(K)
2
(K)
+
(k)
(k)
15i5N.
Eq. 32 generalizes the original message saassageppssingggg rule GCN by adding an additional self-loop links with adaptive weights for different nodes. The diffusivity matrix is defined with the observed adjacency, i.c., S(k) D AD-1. Though such design leverages the geometric information from input structures as a guidance for feature propagation, it constrains the efficiency of layer- wise information flows within the receptive field of local neighbors and could only exploit partial global information with a finite number of iterations.
GAT. Graph Attention Networks (Velickovic et al., 2018) extend the GCN architecture to incorporate attention mechanisms as a learnable function producing adaptive weights for each observed edge. From diffusion perspective, the attention matrix for layer-wise convolution be treated the
20

---

Published as a conference paper at ICLR 2023
diffusivity in our updating rule:
f(1lz (1) (1)
-
(2)
(())
(1.5) E
ELNNEs f(112) IB)
(33)
N
(k+1)
1 - Σ 1=1
((R)
(K)
(k) J 1 1
2
2; +
N (()) due the normalization the denominator. Therefore, the layer-wise We notice that ブØ=1 updating can be viewed as an attentive aggregation over graph structures and subsequent residual link. In fact, the original implementation for the GAT model (Velickovic et al., 2018) specifies the function f partieular form, i.e., exp(LeakyReLLU(a [Wz, (k) W2, (2), which can be viewed as generalized similarity function given trainable W, towards optimizing supervised loss.
G DATASET INFORMATION
this section, we present the detailed information for all the experimental datasets, the pre-processing and evaluation protocol used in Section 5.
Table 6: Information for node classification datasets.
Dataset
Type
# Nodes
Edges
# Node features
# Class
Cora
Citation network
2,708
5,429
1,433
7
Citeseer
Citation network
3,327
4,732
3,703
Pubmed
Citation network
19.717
44,338
500
Proteins
Protein interaction
132,534
39,561,252
Pokec
Social network
1,632,803
30,622,564
G.1 NODE CLASSIFICATION DATASETS
Cora, Cites cisseeer and Pubmed (Sen et al., 2008) are commonly used citation networks for evalu- ating models on node classification tasks., These datasets are small-scale networks (with 2K~20K nodes) and the goal is to classify the topics of documents (instances) based input features of each instance (bag-of-words representation of documents) and graph structure (citation links). Following the semi-supervised learning setting in Kipf & Welling (2017), we randomly choosing 20 instances per class for training, 500/1000 instances for validation/testing for each dataset. OGBN-Prote eiroteinns (Hu et al., 2020) is a multi-task protein-prost preetinn interaction network whose goal is to predict molecule instances' property. We follow the original splitting of Hu et al. (2020) for evaluation. Pokec is a large-scale social network with features including profile information, such as geographical region, registration time, and age, for prediction on users' gender. For semi-supervised learning, we consider randomly splitting the instances into train/valid/test with 10%/10%/800% ratios. Table summarizes the statistics of these datasets.
G.2 IMAGE AND TEXT CLASSIFICATION DATASETS
We evaluate our model on two image classification datasets: STL-10 and CIFAR-10. We use all 1 13000 images from STL-10.. each which belongs to one of ten classes. We choose 1500 images from each of 10 classes of CIFAR-10 10 and obtain total of 15,000 images. For STL-10 and CIFAR-10, 10, we randomly select 10/50/100 instances per class as training set, 1000 instances for validation and the remaining instances for testing. We first use the self-supervised approach SimCLR (Chen et al., 2020b) (that does not use labels for training) to train a ResNet-18 18 for extracting the feature maps as input features of instances. We also evaluate our model on 20Newsgroup, which text classification dataset consisting of 9607 instances. We follow Franceschi et al. (2019) to take 10 classes from 20 Newsgroup and words (TFIDF) with frequency of more than 5% as features.
21

---

Published as a conference paper at ICLR 2023
G.3
DATASETS
The spatial-temporal datasets are from the open-source library PyTorch Geometric Temporal (Rozem- berezki et al., 2021), with properties and summary statistics described in Table 7. Node features are evolving for all the datasets considered here, i.c.. we have different node features for different snapshots. For each dataset, we split the snapshots into training, validation, and test sets according to 2:2:6 ratio in order to make it more challenging and close to the real-world low-data learning setting. In details:
Chi ckenpox describes weekly officially reported cases of chickenpox in Hungary from 2005 to 2015, whose nodes are counties and edges denote direct neighborhood relationships Node features are lagged weekly counts of the chickenpox cases (we included 4 lags). The target is the weekly number of cases for the upcoming week (signed integers).
Covid contains daily mobility graph between regions in England NUTS3 regions, with node features corresponding to the number of confirmed COVID-19 cases in the previous days from March to May 2020. The graph indicates how many people moved from one region the other each day, based on Facebook Data For Good disease prevention maps. Node features correspond to the number of COVID-19 cases in the region in the past 8 days. The task is to predict the number of cases in each node after I day.
Wi k i Math is dataset whose nodes describe Wikipedia pages on popular mathematics topics and edges denote the links from one page to another. Node features are provided by the number of daily visits between 2019 March and 2021 March. The graph is directed and weighted. Weights represent the number of links found at the source Wikipedia page linking the target Wikipedia page. The target the daily user visits to the Wikipedia pages between March 16th 2019 and March 15th 2021 which results in 731 periods.
Table 7: Properties and summary statisties of the spatial-temporal datasets used in the experiments with information about whether the graph structure is dynamic or static, meaning of node features (the same as the prediction target) and the corresponding dimension (D), the number of snapshots (T), the number of nodes (V)), as well as the meaning of edges/edge weights.
Dataset
Graph structure
Node features/ Target
Frequency T
[V]
Edges/ Edge weights
Chickenpox
Static
Weekly Chickenpox Cases
4
Weekly
522 20
Direct Neighhorhoods
Covid
Dynamic
Daily Covid Cases
Daily
61
129
Daily Mobility
WikiMath
Static
Daily User Visits
14
Daily
731
1,068
Page Links
H IMPLEMENTATION DETAILS AND
H.1 NODE CLASSIFICATION EXPERIMENT
We use feature transformation for each layer on two large datasets and omit it for citation networks. The head number set as 1. We set graphs for DIFFORMER. For other hyper- paramters, we adopt grid search for all the models with learning rate from {0.0001, 0.001,0.01, 0.1). weight decay for the Adam optimizer from {0, 0.0001, 0.001, 0.01, 0.1, 1.0). dropout rate from {0, 0.2 0.5). hidden size from {16, 32, 64). number of layers from (2, 8, 163. For evaluation, we compute the mean and standard deviation for
0.5 and incorporate the input
of the results with five repeating runs with different initializations. For each run, maximum of 1000 epochs and report the testing performance achieved by the epoch yielding the best performance on validation set.
H.2 IMAGE AND TEXT CLASSIHICATION EXPERIMENT
For image and text datasets, we consider feature transformation for layer-wise updating. The head number is set as 1. We set 7 = 0.5. These datasets do not have input graphs so we only consider learning new structures for the diffusion model. For hyper-parameter settings. we conduct grid search for all the models with learning rate from (0.0001, 0.0005, 0.005, 0.01, 0.05). weight decay for the Adam optimizer from (0.0001, 0.001, 0.01, dropout rate from (0, 0.2, 0.5), hidden size from
22

---

Published as a conference paper at ICLR 2023
Identity
Identity Constant
Constant Full Attn
Full Attn
Accuracy
72.5
DiFFormer-s
Accuracy
62.5
DIFFormer-
70.0
DiFFormer-a
60.0
DIFFormer-a
67.5
57.5
65.0
95.0
62.5
52.5
60.0
50.0
100
1000
1000
2000
4000
Label Number
Label Number
(a) STL
(b) 20News
Figure 4: Ablation studies for different energy function forms image and text datasets.
{32, 64, 100, 200, 300, 400), number of layers from {1, 4, 6, 10, 12]. We average the results for five repeating runs and report as well the standard deviation. For each run, we run for maximum of 600 epochs and report the testing accuracy achieved by the epoch yielding the highest accuracy validation set.
H.3 SPATIAL- TEMIORAL PREDICTION EXPERIMENT
We do not use feature transformation for these datasets due to their small sizes and also set 7 = 0.5. The head number is graph structures, we consider both cases, using the input graphs and not, in our experiments and discuss their impact on the performance. For other syper-parameters for all models here with learning rate from {0.01, 0.05, 0.005), weight decay for the Adam optimizer from (0, 0.005), dropout rate from (0, 0.2, 0.5). and report the test mean squared error (MSE) based on the lowest validation MSE. We average the results for five repeating runs and report well the standard deviation for each MSE result. For each run, we run for a maximum of 200 epochs total and stop the training process with 20-epoch early stopping on the validation performance. The data split is done in time order, and hence deterministic. We report the results using the same hidden size (4) and number of layers (2) for all methods.
1, These spatial-temporal dynamics prediction datasets contain available
also consider grid search
I MORE EXPERIMENT RESULTS
We supplement more experiment results including extra ablation studies, hyper-yppaatte paremeeeretter studies and visualization results more datasets that not presented in Section 5 due the limit of space.
I.I ABLATION STUDIES
In Fig. 4 present more experiment results for ablation studies w.r.t. the energy function forms used by DIFFORMER, See discussions and analysis in Section 5.
I.2 HYPER- PAPEREPAR ANALYSIS
We plot the testing performance of several baselines and DIFFORMER with different step size T as model size increases Fig. See discussions and analysis Section
I.3 VISUALIZATION
Fig. 6 and 7 plot the produced instance-leve representations and diffusion strength estimates by the model on 20News and STL, respectively. We observe that the diffusivity estimates tend to connect nodes with different classes, which contribute to increasing the global connectivity and facilitate absorbing other instances' information for informative representations. The node embeddings
23

---

Published as a conference paper at ICLR 2023
73
61
70-
70-
30-
OCh rosidaa!
Accuracy
62
CCM
CCK i-siihaa
Affsmer
Afformer
20
24
14
10
26 20
26
64
12
16
20
28
30
Model Depth
Model Depth K
Citeseer
(b) Pubmed
Figure 5: Hyper-parameter studies of model depth K and step size two citation networks.
produced by our model have small intra-class distance and large inter-class distance, making easier for the classifier to distinguish.
Fig. 8 visualizes the diffusivity estimates on Ch ckenpox. We conclude that large diffusion strength usually exists between nodes with similar ground-truti labels. DIFFORMER-S has more concentrated large weights while DIFFORMER-a tends to have large diffusivity spreading out more. DIFFORMER- indeed learns more complex underlying structures than DIFFORMER-S due to its better capacity for diffusivity modeling.
1.4 IMPACT OF MINI-BATCH SIZES ON LARGE GRAPHS
The randomness of mini-batch partition on large graphs has negligible effect on the performance since we use large batch sizes for training, which is facilitated by the linear complexity of DIFFORMER-S. Even setting the batch size be 100000, our model only costs 3GB GPU memory on Pokec. As further investigation on this, we add more experiments using different batch sizes on Pokec and the results are shown in Table 8.
Table 8: Discussions on using different mini-batch sizes for training on Pokec. We report testing accuracy and training memory for comparison.
Batch size
5000
20000
100000
200000
Test Acc (%)
165.4
0.34 +
668.53 0.75 0.75168..99
68.96
+ 69.24 0.76 0.76
69.15 +
GPU Memory (MB)
1244
1326
1539
2060
2928
4011
One can see that using small batch sizes would indeed sacrifice the performance yet large batch sizes produce decent and low-variance results with acceptable memory costs.
I.5 COMPARISON OF RUNNING TIME AND MEMORY COSTS
To further study the efficiency and scalability of our model, we provide more comparison regarding the training time per epoch and memory costs of two DIFFORMER's variants, GCN, GAT and DenseGAT in Table 9, One can see that compared to GAT, DIFFORMER-S costs comparable time on small datasets such as Cora a and Wi k kiath, and much faster on large dataset Pokec. As for memory consumption, DIFFORMER-S reduces the costs by several times over DenseGAT,, which clearly shows the efficiency of our new diffusion function designs. Overall, DIFFORMER-S has nice scalability, decent efficiency and yield significantly better accuracy. In contrast, DIFFORMER-2 costs much larger time and memory costs than DIFFORMER-S. due to its quadratic complexity induced by the explicit computation for the all all-pair diffusivity. Still, DIFFORMER-a accommodates non non-lime linearity for modeling the diffusion strengths which enables better capacity for learning complex wise inter-interactions.
24

---

Published as a conference paper at ICLR 2023
Table Comparison of training time and memory of different models on Cora, Pokec, STL-10 and kiath. OOM refers to out-of mit-of- when training GPU with 16GB memory.
Method
GCN
GAT
DIIFEGAT DIFFORMER-3
Train time (8)
0.0584
0.0807
0,5165
0,1438
0.3292
Cora
Training memory (MB)
1168
1380
8460
1350
3893
Train time (s)
1.069
14.87
88.07
2.206
OOM
Pokec
Training memory (MB)
1812
2014
13174
2923
OOM
Train time (8)
00.0069
00.0444
OOM
0.0323
0.3298
STL.
Training memory (MB)
1224
1980
0OM
1342
7680
Train time (8)
0.008
00.0261
0.0364
0.0281
0,0350
WikiMath
Training memory (MB)
1048
1054
1316
1046
1142
1.6 INCORPORATION OF PSEUDO LABELS
For semi-supervised learning, there is line of approaches that leverage pscudo labels to augment the training data. Our model DIFFORMER essentially has orthogonal technical aspects compared to this line of work in that we focus on building a new encoder backbone and only train the model with standard supervised loss on the labeled data. This means that pseudo-label-based approaches equally applicable for enhancing the training of our model as well as the competitors we used in our experiments.
As an initial verification of this claim, we use the Meta Pseudo Labels (MPL) (Pham et al., 2021) as plug-in module boost DIFFORMER well our competitors GCN-KNN and GAT-KNN, and empirically compare the relative improvement. Specifically, we use DIFFORMER-S, DIFFORMER-A. GCN and GAT as the encoder backbone of teacher and student models, respectively, and use the MPL algorithm to generate pseudo labels for augmenting the training data used for computing the supervised loss. The results on CIFAR- 10 and STL-10 are shown Table 10. As we can see, MPL contributes to some performance gains across all four encoders, while our two DIFFORMER variants still maintain the superiority over the competitors. Note also that as a proof-of-concey here did not practice this type of more sophisticated MPI. implementation could principle be applied to further improve performance (across all models).
additional consistency loss that requires careful manual tuning. However, in
Table 10: Comparison of using and not using Meta Pseudo Labels (MPL) as a plug-in module boost different encoder backbones on CIFAR 10 and STL-10 10.
Method
GCN
GAT
DIFFORMMR-
DIFFORMER-A
w/o MPL
73.7 0.4
73.9
76.4
75.3 + 0.6 75.9 0.4
STL
w/ MPL
74.3
0.5
0.7
7.6
w/o MPL
74.7
0.1
76.6
75.9 0.3 76.3 + 0.3
CIFAR
w/ MPL
0.4
74.8
7.3
25

---

Published as a conference paper at ICLR 2023
(a) The first layer
(b) The last layer
Figure 6: Visualization of instance representations and diffusivity strengths (we set threshold and only plot the edges with weights more than the threshold) at different layers given by DIFFORMER-S on 20News.
(a) The first layer
(b) The last layer
Figure 7: Visualization of instance representations and diffusivity strengths (we set threshold and only plot the edges with weights more than the threshold) different layers given by DIFFORMER-S on STL.
(a)
(b)
(c)
(d)
(c)
(f)
Figure 8: The produced diffusivity of the first layer (i.e., S()) on Chi Chickepox across the first three snapshots, yielded by DIFFORMER-S. shown in (a)~(c), and DIFFORMER-2 shown in (d)~(f). Node colors correspond to ground-truth labels (i.e., reported cases), varying from red to blue as the label increases. We visualize the edges with top 100 diffusion strength, where edge colors change from blue to red as
(1) increases.
26


Reviews:
[{"review_id": "fcn-Q89nAxv", "reviewer": "Reviewer_cBoH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this work, a discrete diffusion process along the lines of Zhou et al (2004) but with features inspired by GNN works such as Velikovic et al (2017) is studied for use in semi-supervised node feature prediction.  It is argued that the process can be derived from an energy function.\nMany experiments and comparisons with other models are done, and the results are competitive and arguably better.", "summary_of_the_review": "Interesting contribution to a much studied field which makes incremental progress. ", "strength_and_weaknesses": "Strengths: Clearly written and many experimental results.\n\nWeaknesses: This is a much studied field and I did not find the discussion of advances over previous work entirely convincing.  \n\nThe performance improvements are marginal (1 sigma).\n\nIt was not clear to me how the \"simple diffusivity model\" could satisfy the conditions of theorem 1.  If the function g is linear in z^2, then \\delta is quadratic.  But there is no nondecreasing concave quadratic function on R.  The \"advanced diffusivity model\" uses a nonpolynomial function.  Perhaps there is some implicit restriction on the domain, as suggested by the comment below Eq. (8) that f(z^2)=2-z^2/2 is non-negative, and the comment about layer normalization.  It is important to clarify this point.\n\nTheorem 1 might follow from quasiconvexity of the energy function.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "f3Oop0P1Bly", "reviewer": "Reviewer_tDZ1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors introduce an energy-constrained diffusion model which encodes a batch of instances from a dataset into evolutionary states that progressively incorporate other instances’ information by their interactions. In their method, the diffusion process is constrained by descent criteria w.r.t. a principled energy function that characterizes the global consistency. Experiments show that the model could achieve superior performance in various tasks.", "summary_of_the_review": "The method proposed in the paper is novel, and the writing and structure of the paper are good. ", "strength_and_weaknesses": "Strength:\n\nIn this paper, the authors proposed a geometric approach to model sample points on the data manifold through diffusion on graphs. The proposed approach and methodology are interesting.  Moreover, the method achieved improved performance on multiple datasets. \n\n\nWeakness:\n\nThe author could give more detail about the difference between the proposed method and label propagation used for semi-supervised learning.  How does the data sample graph affect the performance of the method? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "GM_jMZd8rd", "reviewer": "Reviewer_ckRq", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a new class of self-supervised encoder that learns the latent representation of data instances via a diffusion process that encodes the interaction interaction strength between data pairs. Two versions of the model are introduced, DIFFormer-s and DIFFormer-a, both demonstrating superior performance in a range of tasks.", "summary_of_the_review": "In summary, this paper provides a novel, principled approach to learn the latent geometry of a dataset via an energy function that encodes pairwise diffusion strength. The theory is sound and well-motivated. Empirical results show improvements on a range of tasks. ", "strength_and_weaknesses": "Strength:\n- This paper does a good job in motivating and describing the approach. From my perspective, the theory is quite interesting and elegant. The authors also provide a unified view for MLP, GCN, and GAT from their framework.\n- The model achieves improved performance over a broad range of tasks.\n- Extensive experiment results are provided in the Appendix on the influence of batch sizes, compute cost, etc. \n\nWeakness:\n- The model may have challenges in scaling to larger dataset due to the O(NKd^2) scaling for DIFFormer-s and O(N^2Kd^2) scaling for DIFFormer-a. I don’t fully understand the GPU memory cost result in Table 8. Why does the memory cost only increase by 4 times when the batch size increases from 5000 to 200000? How about the compute and memory cost for DIFFormer-a?\n- For the CIFAR-10 and STL-10 experiments, have the authors considered stronger semi-supervised learning baselines like Meta Pseudo Labels? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "2AuR84kA2gH", "reviewer": "Reviewer_Hw3e", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a theoretical framework which leverages a diffusivity function $\\mathbf S_{ij}$ to capture interdependencies in input data. In the extremes, one can consider an IID assumption - in which case the diffusivity function is just the constant identity matrix, and the model simplified to an MLP - or a setting where the inputs are vertices in a known graph - in which case using a diffusivity function $\\mathbf S_{ij}$ with nonzero weights between edges $(i,j) \\in \\mathcal E$ can result in either a GCN or GAT, depending on the choice of weights.\n\nBeyond incorporating existing architectures in a unified framework, this perspective also suggests that the diffusivity function can be computed based on the input features. The authors introduce the crucial idea of minimizing an energy function which measures the quality of instance states at a given step, and then prove the existence of a diffusivity function which guarantees that the energy of the instances will be minimized in each step. This result has one remaining design decision - the specification of a non-negative decreasing function - and the authors explore two obvious choices for this function, leading to a \"simple diffusivity model\" DIFFormer-s which is linear in terms of the number of inputs, and an \"advanced diffusivity model\" DIFFormer-a, the latter of which is capable of capturing more complex latent geometry but has quadratic complexity in the number of inputs.\n\nA large array of experiments, both synthetic and on real data, are included, with an array of strong baselines. The authors also perform a comprehensive ablation study to assess various aspects of the design decisions.", "summary_of_the_review": "This was an outstanding work, very polished and complete. Definitely ready for publication.", "strength_and_weaknesses": "This paper is very strong in many respects. The theoretical framework fits well in backdrop of a larger ongoing effort to unify various machine learning models. The motivation and presentation is very clear throughout, both from a technical/mathematical sense as well as general writing style. The theorem which is proved is not just a nice-to-have, but rather fundamental to the solution, and the authors emphasize this in their presentation beautifully. I have no major weaknesses to report.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%