% Title: Discovering Data Geometry: \\

% ArXiv ID: 2301.09474




\maketitle

\begin{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 \emph{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 \model (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 \url{https://github.com/qitianwu/DIFFormer}.
\end{abstract}

\vspace{-4pt}
\section{Introduction}\label{sec:intro}
\vspace{-4pt}

Real-world data are generated from a convoluted interactive process whose underlying physical principles are often unknown. %The observational data used for model development or experimental purposes can be viewed as a subset from a hypothetical pool composed of an exhaustive enumeration of all concepts or objects. 
Such a 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, it can be practically prohibitive to build feasible methodology for uncovering data dependencies, despite the acknowledged significance. To address this issue, prior works, e.g., \cite{pointcloud-19,LDS-icml19,jiang2019glcn,Bayesstruct-aaai19}, consider encoding the potential interactions between instance pairs, but this requires sufficient degrees of freedom that significantly increases learning difficulty from limited labels~\citep{fatemi2021slaps} and hinders the scalability to large systems~\citep{wunodeformer}. 



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 as graph neural networks (GNNs)~\citep{scarselli2008gnnearly,GCN-vallina,GAT,SGC-icml19,gcnii-icml20,RWLS-icml21} for harnessing inter-connections between instances as a geometric prior~\citep{geometriclearning-2017}. 
However, the observed relations can be incomplete/noisy, due to error-prone data collection, or generated by an 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 data dependencies. While a plausible remedy is to learn more useful structures from the data, this unfortunately brings the previously-mentioned obstacles to the fore.%embodied by emerging evidence showing GNNs' deficiency for heterophily~\cite{}, long-range dependence~\cite{} and noisy/incomplete edges~\cite{}. A plausible remedy is to learn more useful structure from data~\cite{LDS-icml19,jiang2019glcn,gong2020geometrically}, which, however, brings the previous obstacles to the fore.



To resolve the dilemma, we propose a novel general-purpose 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.~\ref{fig:model}. Our model is defined through feed-forward continuous dynamics (i.e., a PDE) involving all the instances of a dataset as locations on Riemannian manifolds with \emph{latent} structures, upon which the features of instances act as heat flowing over the underlying geometry~\citep{hamzi2021learning}. Such a diffusion model serves an important \emph{inductive bias} for leveraging global information from other instances to obtain more informative representations. Its major advantage lies in the flexibility for the \emph{diffusivity} function, i.e., a measure of the rate at which information spreads~\citep{rosenberg1997laplacian}: we allow for feature propagation between arbitrary instance pairs at each layer, and adaptively navigate this process by pairwise connectivity weights.  %The diffusion process plays an inductive bias for representing instances into evolving states that progressively integrate information flows from other particles in the system. %and compose a continuous dynamics (i.e., a PDE) with instantaneous diffusion strength, a.k.a. \emph{diffusivity}, as \emph{latent} variables. The diffusion process encodes all the instances into a series of evolutionary states that progressively integrate information flows from other particles in the system. 
Moreover, for guiding the instance representations towards some ideal constraints of internal consistency, we introduce a principled energy function that enforces layer-wise \emph{regularization} on the evolutionary directions. The energy function provides another view (from a macroscopic standpoint) into the desired instance representations with low energy that are produced, i.e., soliciting a steady state that gives rise to informed predictions on unlabeled data.%We call our model Energy Constrained Neural Diffusion (\model) whose dynamics are \emph{implicitly} constrained by descent criteria w.r.t. the regularized energy, i.e., soliciting a final state that gives rise to informed prediction on unlabeled data.}

\begin{figure}[tb!]
			\centering
			\hspace{5pt}
			\includegraphics[width=0.85\textwidth]{figure/model.pdf}
	    \vspace{-5pt}
		\caption{\looseness=-1 An illustration of the general idea behind \model which takes a whole dataset (or a batch) of instances as input and encodes them into hidden states through a 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.}
		\label{fig:model}
	\vspace{-20pt}
\end{figure}

As a justification for the tractability of above general methodology, our theory reveals the underlying equivalence between finite-difference iterations of the diffusion process and unfolding the minimization 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 a new class of neural encoders, Diffusion-based Transformers (\model), and its two practical instantiations: one is a simple version with $\mathcal 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 \model on a diverse set of tasks. It outperforms SOTA approaches on semi-supervised 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 spatial-temporal dynamics.







    
    

\vspace{-4pt}
\section{Related Work}\label{sec:related}
\vspace{-4pt}

\textbf{Graph-based Semi-supervised Learning.}
Graph-based SSL~\citep{GCN-vallina} 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~\citep{belkin2006manireg,weston2012semiemb,Panetoid-icml19} or as an inductive bias of modern GNN architectures~\citep{scarselli2008gnnearly}. However, there frequently exist situations where the observed structure is unavailable or unreliable~\citep{LDS-icml19,jiang2019glcn,IDGL-neurips20,fatemi2021slaps,lao2022variational}, in which case the challenge remains how to uncover the underlying relations. This paper explores a new Transformer-like encoder for discovering data geometry to promote learning through the inter-dependence among instances (either labeled or unlabeled).

\textbf{Neural Diffusion Models.} Several recent efforts explore diffusion-based learning where continuous dynamics serve as an inductive bias for representation learning~\citep{hamzi2021learning}. These works can be generally grouped into PDE-based learning, where the model itself is a continuous diffusion process described by a differential equation (e.g., \cite{grand,PDE-GCN,GRAND++}), and PDE-inspired learning, where the diffusion perspective is adopted to develop new models~\citep{atwood2016diffusion,klicpera2019diffusion} or characterize graph geometric properties~\citep{yang2022geometric}. 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.

\modify{\textbf{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~\citep{transformer} is motivated by the inter-dependence among words of a sentence in NLP tasks, from different aspects, this work targets a Transformer-like architecture for modeling the inter-dependence among instances in a dataset. In the terminology of Transformers in NLPs, one can treat the whole dataset as 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~\citep{ogb-nips20}. For the latter case, critically though, it remains under-explored how to build a scalable and expressive Transformer for learning node-pair interactions given the prohibitively large number of node instances~\citep{wunodeformer}.}


\vspace{-4pt}

\section{Energy Constrained Geometric Diffusion Transformers}\label{sec:model}
\vspace{-4pt}



Consider a set of partially labeled instances $\{\mathbf x_i\}_{i=1}^N$, whose labeled portion is $\{(\mathbf x_j, y_j)\}_{j=1}^M$ (often $M \ll N$). In some cases there exist relational structures that connect
instances as a graph $\mathcal G = (\mathcal V, \mathcal E)$, where the node set $\mathcal V$ contains all the instances and the edge set $\mathcal E=\{e_{ij}\}$ consists of observed relations. Without loss of generality, the main body of this section does \emph{not} assume graph structures as input, but we will later discuss how to trivially incorporate them if they are available. 





    


\vspace{-4pt}

\subsection{Geometric Diffusion Model}
\vspace{-4pt}








The starting point of our model is a diffusion process that treats a dataset of instances as a whole and produces instance representations through information flows characterized by an anisotropic diffusion process, which is inspired by an analogy with heat diffusion on a Riemannian manifold~\citep{rosenberg1997laplacian}. We use a vector-valued function $\mathbf z_i(t): [0, \infty) \rightarrow \mathbb R^d$ to define an instance's state at time $t$ and location $i$. The anisotropic diffusion process describes the evolution of instance states (i.e., representations) via a PDE with boundary conditions~\citep{freidlin1993diffusion,medvedev2014nonlinear}:
\begin{equation}\label{eqn-diffuse}
    \frac{\partial \mathbf Z(t)}{\partial t} = \nabla^*\left(\mathbf S(\mathbf{Z}(t), t) \odot \nabla \mathbf Z(t)\right), ~~~ \mbox{s. t.} ~~ \mathbf Z(0) = [\mathbf x_i]_{i=1}^N, ~~ t\geq 0, %+ f(\mathbf z_u^{(0)}; W_F),
\end{equation}
where $\mathbf Z(t) = [\mathbf z_i(t)]_{i=1}^N \in \mathbb R^{N\times d}$, $\odot$ denotes the Hadamard product, and the function $\mathbf S(\mathbf Z(t), t): \mathbb R^{N\times d} \times [0, \infty) \rightarrow [0, 1]^{N\times N}$ defines the \emph{diffusivity} coefficient controlling the diffusion strength between any pair at time $t$. The diffusivity is specified to be dependent on instances' states.
The gradient operator $\nabla$ measures the difference between source and target states, i.e., $(\nabla \mathbf Z(t))_{ij} = \mathbf z_j(t) - \mathbf z_i(t)$, and the divergence operator $\nabla^*$ sums up information flows at a point, i.e., $(\nabla^*)_i = \sum_{j=1}^N \mathbf S_{ij}(\mathbf Z(t), t) \left(\nabla \mathbf Z(t)\right)_{ij}$. Note that both operators are defined over a discrete space consisting of $N$ locations. The physical implication of Eq.~\ref{eqn-diffuse} is that the temporal change of heat at location $i$ equals to
the heat flux that spatially enters into the point. %Notice that we do not assume ground-truth dependence is available or the diffusivity functions reflect the correct data dependence, and instead, we expect the diffusion process can propagate useful information for yielding informative representation/prediction, which can be achieved by learning with labeled data in the dataset.
Eq.~\ref{eqn-diffuse} can be explicitly written as
\begin{equation}\label{eqn-diffuse2}
    \frac{\partial \mathbf z_i(t)}{\partial t} = \sum_{j=1}^N \mathbf S_{ij}(\mathbf Z(t), t) (\mathbf z_j(t) - \mathbf z_i(t)).
\end{equation} 
Such a diffusion process can serve as an inductive bias that guides the model to use other instances' information at every layer for learning informative instance representations. 
We can use numerical methods to solve the continuous dynamics in Eq.~\ref{eqn-diffuse2}, e.g., the explicit Euler scheme involving finite differences with step size $\tau$, which after some re-arranging gives:
\begin{equation}\label{eqn-diffuse-iter}
    \mathbf z_i^{(k+1)} = \left (1 - \tau \sum_{j=1}^N \mathbf S_{ij}^{(k)} \right ) \mathbf z_i^{(k)} + \tau \sum_{j=1}^N \mathbf S_{ij}^{(k)} \mathbf z_j^{(k)}.
\end{equation}
The numerical iteration can stably converge for $\tau\in(0, 1)$. We can adopt the state after a finite number $K$ of propagation steps  and use it for final predictions, i.e., $\hat y_i = \mbox{MLP}(\mathbf z_i^{(K)})$.



\looseness=-1\textbf{\emph{Remark.}} The diffusivity coefficient in Eq.~\ref{eqn-diffuse} is a measure of the rate at which heat can spread over the space~\citep{rosenberg1997laplacian}. Particularly in Eq.~\ref{eqn-diffuse2}, $\mathbf S(\mathbf Z(t), t)$ 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 $\mathbf S(\mathbf Z(t), 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 $\mathbf S(\mathbf Z(t), t)$ as the observed graph structure if available in some scenarios. In such a case, however, the information flows are restricted by neighboring nodes in a graph. An ideal case could be to allow $\mathbf S(\mathbf Z(t), t)$ 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. %\modify{This is exactly what we will adopt in this paper (see Section~\ref{sec:model}.3 and \ref{sec:inst} for details), contributing to superior flexibility and capacity of our diffusion model. We also compare the performance of using different $\mathbf S$'s in Section~\ref{sec:exp}.4.}










\vspace{-4pt}
\subsection{Diffusion Constrained by a Layer-wise Energy}
\vspace{-4pt}
As mentioned previously, the crux is how to define a proper diffusivity function to induce a 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 $\mathbf S^{(k)}$, we consider the diffusivity as a time-dependent latent variable and introduce an \emph{energy function} that measures the presumed quality of instance states at a given step $k$:
\begin{equation}\label{eqn-energy}
    E(\mathbf Z, k; \delta) = \|\mathbf Z - \mathbf Z^{(k)}\|_{\mathcal F}^2 + \lambda\sum_{i,j} \delta(\|\mathbf z_i - \mathbf z_j\|_2^2),
\end{equation}
where $\delta:\mathbb R^+ \rightarrow \mathbb R$ is defined as a function that is \emph{non-decreasing} and \emph{concave} on a particular interval of our interest, and promotes robustness against large differences~\citep{RWLS-icml21} among any pair of instances. Eq.~\ref{eqn-energy} assigns each state in $\mathbb R^d$ with an energy scalar which can be leveraged to regularize the updated states (towards lower energy desired). The weight $\lambda$ trades two effects: 1) for each instance $i$, the states not far from the current one $\mathbf z_i^{(k)}$ have low energy; 2) for all instances, the smaller differences their states have, the lower energy is produced. 

\textbf{\emph{Remark.}} Eq.~\ref{eqn-energy} can essentially be seen as a robust version of the energy introduced by \cite{globallocal-2003}, inheriting the spirit of regularizing the global and local consistency of representations. %The `robust' 
``Robust" here particularly implies that the $\delta$ adds uncertainty to each pair of the instances and could \emph{implicitly} filter the information of noisy links (potentially reflected by proximity in the latent space). %The second term of Eq.~\ref{eqn-energy} also serves as a generalized version of the graph-based regularization~\cite{graphkernel} which has no uncertainty on edges, i.e., $\delta(x) = x$, and only sums up the penalty for linked nodes, i.e., $(i,j)\in \mathcal E$. In contrast, our model does not rely on input graphs and could uncover new structure.  %across spurious potential links.\footnote{We use `potential' here to emphasize that the pairwise interaction is allowed for any node pair in our model beyond the observed topology.} 


\textbf{Energy Constrained Diffusion.} The diffusion process describes the \emph{microscopic} behavior of instance states through evolution, while the energy function provides a \emph{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 a steady point wherein the yielded instance representations have absorbed enough global information under a 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 a valid diffusivity function.
Formally, we aim to find a series of $\mathbf S^{(k)}$'s whose dynamics and constraints are given by 
\begin{equation}\label{eqn-diffuse-appx}
    \begin{split}
    &\mathbf z_i^{(k+1)} = \left (1 - \tau \sum_{j=1}^N \mathbf S_{ij}^{(k)} \right ) \mathbf z_i^{(k)} + \tau \sum_{j=1}^N \mathbf S_{ij}^{(k)} \mathbf z_j^{(k)} \\
    &\mbox{s. t.}~~\mathbf z_i^{(0)} = \mathbf x_i, \quad E(\mathbf Z^{(k+1)}, k; \delta) \leq E(\mathbf Z^{(k)}, k-1; \delta), \quad k\geq 1.
    \end{split}
\end{equation}
The formulation induces a new class of geometric flows on latent manifolds whose dynamics are \emph{implicitly} defined by optimizing
a time-varying energy function (see Fig.~\ref{fig:model} for an illustration). 

\vspace{-4pt}
\subsection{Tractability of Solving Diffusion Process with Energy Minimization}
\vspace{-4pt}

Unfortunately, Eq.~\ref{eqn-diffuse-appx} is hard to solve since we need to infer the value for a series of coupled $\mathbf S^{(k)}$'s 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 an explicit closed-form solution for $\mathbf S^{(k)}$ based on the current states $\mathbf Z^{(k)}$ that yields a rigorous decrease of the energy.
\begin{theorem}\label{thm-main}
    For any regularized energy defined by Eq.~\ref{eqn-energy} with a given $\lambda$, there exists $0<\tau<1$ such that the diffusion process of Eq.~\ref{eqn-diffuse-iter} with the diffusivity between pair $(i,j)$ at the $k$-th step given by
    \begin{equation}\label{eqn-optimal-diffuse}
        \mathbf{\hat S}_{ij}^{(k)} = \frac{\omega_{ij}^{(k)}}{\sum_{l=1}^N \omega_{il}^{(k)}}, \quad \omega_{ij}^{(k)} = \left. \frac{\partial \delta(z^2)}{\partial z^2} \right |_{z^2 = \|\mathbf z_i^{(k)} - \mathbf z_j^{(k)}\|_2^2},
    \end{equation}
    yields a descent step on the energy, i.e., $E(\mathbf Z^{(k+1)}, k; \delta) \leq E(\mathbf Z^{(k)}, k-1; \delta)$ for any $k\geq 1$.
\end{theorem}

Theorem~\ref{thm-main} suggests the existence for the optimal diffusivity in the form of a function over the $l_2$ distance between states at the current step, i.e., $\|\mathbf z_i^{(k)} - \mathbf z_j^{(k)}\|_2$. The result enables us to unfold the implicit process and compute $\mathbf S^{(k)}$ in a feed-forward way from the initial states. We thus arrive at a new family of neural model architectures with layer-wise computation specified by:
{\small
\vspace{-5pt}
\begin{center}
\fcolorbox{black}{gray!10}{\parbox{0.97\linewidth}{
\vspace{-5pt}
\begin{equation}\label{eqn-model}
    \begin{split}
        & \mbox{Diffusivity Inference:} \quad \mathbf {\hat S}_{ij}^{(k)} = \frac{f(\|\mathbf z_i^{(k)} - \mathbf z_j^{(k)}\|_2^2)}{\sum_{l=1}^N f(\|\mathbf z_i^{(k)} - \mathbf z_l^{(k)}\|_2^2)}, \quad 1 \leq i, j \leq N,\\
        & \mbox{State Updating:} \quad \mathbf z_i^{(k+1)} = \underbrace{\left (1 - \tau \sum_{j=1}^N \mathbf {\hat S}_{ij}^{(k)} \right ) \mathbf z_i^{(k)} }_{\mbox{state conservation}} + \underbrace{\tau \sum_{j=1}^N \mathbf {\hat S}_{ij}^{(k)} \mathbf z_j^{(k)} }_{\mbox{state propagation}}, \quad 1 \leq i \leq N.
    \end{split}
\end{equation}
\vspace{-10pt}
}
}
\end{center} }
\vspace{-5pt}
\textbf{\emph{Remark.}} The choice of function $f$ in above formulation is not arbitrary, but needs to be a non-negative and decreasing function of $z^2$, so that the associated $\delta$ in Eq.~\ref{eqn-energy} is guaranteed to be non-decreasing and concave w.r.t. $z^2$. Critically though, there remains much room for us to properly design the specific $f$, %(that is essentially connected to the regularized energy with particular $\delta$, given by the result of Theorem~\ref{thm-main}), 
so as to provide adequate capacity and scalability. Also, in our model presented by Eq.~\ref{eqn-model} we only have one hyper-parameter $\tau$ in practice, noting that the weight $\lambda$ in the regularized energy is implicitly determined through $\tau$ by Theorem \ref{thm-main}, which reduces the cost of hyper-parameter searching.


\vspace{-4pt}

\section{Instantiations of \model}\label{sec:inst}
\vspace{-4pt}


\subsection{Model Instantiations}
\vspace{-4pt}

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~\ref{appx-alg}. %For other possible forms potentially working well, we provide some non-exhaustive discussions in Appendix~\ref{} and leave more thorough studies for future works.
First, because $\|\mathbf z_i - \mathbf z_j\|_2^2 = \|\mathbf z_i\|_2^2 + \|\mathbf z_j\|_2^2 - 2\mathbf z_i^\top \mathbf z_j$, we can convert $f(\|\mathbf z_i - \mathbf z_j\|_2^2)$ into the form $g(\mathbf z_i^\top \mathbf z_j)$ using a change of variables on the condition that $\|\mathbf z_i\|_2$ remains constant. And we add layer normalization to each layer to loosely enforce such a property in practice.

\textbf{Simple Diffusivity Model.} A straightforward design is to adopt the linear function $g(x) = 1+x$:
\begin{equation}\label{eqn-S1}
    \omega_{ij}^{(k)} = f(\|\tilde{\mathbf z}_i^{(k)} - \tilde{\mathbf z}_j^{(k)}\|_2^2) = 1 + \left (\frac{\mathbf z_i^{(k)}}{\|\mathbf z_i^{(k)}\|_2} \right)^\top 
    \left (\frac{\mathbf z_j^{(k)}}{\|\mathbf z_j^{(k)}\|_2} \right ),
\end{equation}
Assuming $\tilde{\mathbf z}_i^{(k)} = \frac{\mathbf z_i^{(k)}}{\|\mathbf z_i^{(k)}\|_2}$, $\tilde{\mathbf z}_j^{(k)} = \frac{\mathbf z_j^{(k)}}{\|\mathbf z_j^{(k)}\|_2}$ and $z = \|\tilde{\mathbf z}_i^{(k)} - \tilde{\mathbf z}_j^{(k)}\|_2$, Eq.~\ref{eqn-S1} can be written as $f(z^2) = 2 - \frac{1}{2}z^2$, which yields a non-negative result and is decreasing on the interval $[0, 2]$ in which $z^2$ lies.
One scalability concern for the model Eq.~\ref{eqn-model} arises because of the need to compute pairwise diffusivity and propagation for each individual, inducing $\mathcal O(N^2)$ complexity. Remarkably, the simple diffusivity model allows a significant acceleration by noting that the state propagation can be re-arranged via %(assume $\overline{\mathbf z}_i = \tilde{\mathbf z}_i^{(l)} / \|\tilde{\mathbf z}_i^{(l)}\|_2$ for clean notation)
\begin{equation}
     \sum_{j=1}^N \mathbf S_{ij}^{(k)} \mathbf z_j^{(k)} = \sum_{j=1}^N \frac{1 + (\tilde{\mathbf z}_i^{(k)})^\top \tilde{\mathbf z}_j^{(k)}}{\sum_{l=1}^N\left (1 + (\tilde{\mathbf z}_i^{(k)})^\top \tilde{\mathbf z}_l^{(k)} \right ) } \mathbf z_j^{(k)} = \frac{\sum_{j=1}^N\mathbf z_j^{(k)} + \left (\sum_{j=1}^N \tilde{\mathbf z}_j^{(k)}\cdot  (\mathbf z_j^{(k)})^\top \right ) \cdot \tilde{\mathbf z}_i^{(k)} }{N + (\tilde{\mathbf z}_i^{(k)})^\top \sum_{l=1}^N \tilde{\mathbf z}_l^{(k)}}.
\end{equation}
The two summation terms above can be computed once and shared to every instance $i$, reducing the complexity in each iteration to $\mathcal O(N)$ (see Appendix~\ref{appx-alg} for how we achieve linear complexity w.r.t. $N$ in the matrix form for model implementation). We refer to this implementation as \model-s. %\modify{Similar strategies for complexity reduction is utilized by NodeFormer~\citep{wunodeformer}, a recently proposed scalable graph Transformer based on random feature map, while in contrast we resort to a more straightforward design that does not require any stochastic component or approximation.}

\textbf{Advanced Diffusivity Model.} The simple model facilitates efficiency/scalability, yet may sacrifice the  capacity for complex latent geometry. We thus propose an advanced version with $g(x) = \frac{1}{1+\exp(-x)}$:
\begin{equation}\label{eqn-S2}
    \omega_{ij}^{(k)} = f(\|\tilde{\mathbf z}_i^{(k)} - \tilde{\mathbf z}_j^{(k)}\|_2^2) = \frac{1}{1+\exp{\left(- (\mathbf z_i^{(k)} )^\top 
    (\mathbf z_j^{(k)} ) \right ) }},
\end{equation}
which corresponds with $f(z^2) = \frac{1}{1+e^{z^2/2-1}}$ guaranteeing monotonic decrease and non-negativity. We dub this version as \model-a. Appendix~\ref{appx-inst} further compares the two models (i.e., different $f$'s and $\delta$'s) through synthetic results. Real-world empirical comparisons are in Section~\ref{sec:exp}.

\begin{table}[t!]
\centering
\caption{A unified view for MLP, GCN and GAT from our energy-driven geometric diffusion framework regarding %. The table provides a head-to-head comparison with \model in terms of 
energy function forms, diffusivity specifications and algorithmic complexity. \label{tbl-existing}}
\vspace{-5pt}
\small
    \resizebox{0.99\textwidth}{!}{
\begin{tabular}{@{}c|c|c|c@{}}
\toprule
Models & Energy Function $E(\mathbf Z, k; \delta)$ & Diffusivity $\mathbf S^{(k)}$ & Complexity \\
\midrule
MLP & $\|\mathbf Z - \mathbf Z^{(k)}\|_2^2$ & $\mathbf S^{(k)}_{ij} = \left\{ 
    \begin{aligned}
         &1, \quad \mbox{if} \; i = j  \\
         &0, \quad otherwise
    \end{aligned}
    \right. $ & $\mathcal O(NKd^2)$ \\
\hline
GCN & $\sum_{(i,j)\in \mathcal E} \|\mathbf z_i - \mathbf z_j\|_2^2$ & $\mathbf S^{(k)}_{ij} = \left\{ 
    \begin{aligned}
         &\frac{1}{\sqrt{d_id_j}}, \quad \mbox{if} \; (i,j) \in \mathcal E  \\
         &0, \quad otherwise
    \end{aligned}
    \right. $ & $\mathcal O(|\mathcal E|Kd^2)$  \\
\hline
GAT & $\sum_{(i,j)\in \mathcal E} \delta(\|\mathbf z_i - \mathbf z_j\|_2^2)$ & $\mathbf S^{(k)}_{ij} = \left\{ 
    \begin{aligned}
         &\frac{f(\|\mathbf z_i^{(k)} - \mathbf z_j^{(k)}\|_2^2)}{\sum_{l: (i,l)\in \mathcal E} f(\|\mathbf z_i^{(k)} - \mathbf z_l^{(k)}\|_2^2)}, \quad \mbox{if} \; (i,j) \in \mathcal E  \\
         &0, \quad otherwise
    \end{aligned}
    \right.  $& $\mathcal O(|\mathcal E|Kd^2)$  \\
\hline
\model & $\|\mathbf Z - \mathbf Z^{(k)}\|_2^2 + \lambda \sum_{i,j} \delta(\|\mathbf z_i - \mathbf z_j\|_2^2)$ & $\mathbf S^{(k)}_{ij} = \frac{f(\|\mathbf z_i^{(k)} - \mathbf z_j^{(k)}\|_2^2)}{\sum_{l=1}^N f(\|\mathbf z_i^{(k)} - \mathbf z_l^{(k)}\|_2^2)}, \quad 1\leq i,j \leq N$ & $\begin{aligned} &\mbox{\model-s}:\quad \mathcal O(NKd^2) \\ &\mbox{\model-a}:\quad \mathcal O(N^2Kd^2) \end{aligned}$ \\
\bottomrule
\end{tabular}}
\vspace{-10pt}
\end{table}

\subsection{Model Extensions and Further Discussion}
\textbf{Incorporating Layer-wise Transformations.} Eq.~\ref{eqn-model} does not use feature transformations for each layer. To further improve the representation capacity, we can add such transformations after the updating, i.e., $\mathbf z_i^{(k)} \leftarrow h^{(k)}(\mathbf z_i^{(k)})$ where $h^{(k)}$ can be a fully-connected layer (see Appendix~\ref{appx-alg} for details). In this way, each iteration of the diffusion yields a descent of a particular energy $E(\mathbf Z, k;\delta, h^{(k)}) = \|\mathbf Z - h^{(k)}(\mathbf Z^{(k)})\|_2^2 + \sum_{i,j} \delta(\|\mathbf z_i - \mathbf z_j\|_2^2)$ dependent on $k$. %The diffusion process becomes a composition of a series of steps, each of which optimizes a locally adaptive energy. 
The trainable transformation $h^{(k)}$ can be optimized w.r.t. the supervised loss to 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 effects for datasets with larger sizes. Furthermore, one can consider non-linear activations in the layer-wise transformation $h^{(k)}$ though we empirically found that using a linear model already performs well. We also note that Theorem~\ref{thm-main} can be extended to hold even when incorporating such a non-linearity in each layer (see Appendix~\ref{appx-non} for detailed discussions).





\textbf{Incorporating Input Graphs.} For the model presented so far, we do \emph{not} assume an input graph for the model formulation. For situations with observed structures as available input, we have $\mathcal G = (\mathcal V, \mathcal E)$ that can be leveraged as a geometric prior. We can thus modify the updating rule as:
\begin{equation}\label{eqn-diffuse-graph}
    \mathbf z_i^{(k+1)} = \left (1 - \frac{\tau}{2} \sum_{j=1}^N \left(\mathbf {\hat S}_{ij}^{(k)} + \tilde{\mathbf A}_{ij} \right ) \right ) \mathbf z_i^{(k)} + \frac{\tau}{2} \sum_{j=1}^N \left (\mathbf {\hat S}_{ij}^{(k)} + \tilde{\mathbf A}_{ij}  \right )\mathbf z_j^{(k)},
\end{equation}
where $\tilde{\mathbf A}_{ij} = \frac{1}{\sqrt{d_id_j}}$ if $(i,j)\in \mathcal E$ and 0 otherwise, and $d_i$ is instance $i$'s degree in $\mathcal G$. The diffusion iteration of Eq.~\ref{eqn-diffuse-graph} is essentially a descent step on a new energy additionally incorporating a graph-based penalty~\citep{graphkernel}, i.e., $\sum_{(i,j)\in \mathcal{E}} \|\mathbf z_i^{(k)} - \mathbf z_j^{(k)}\|_2^2$ (see Appendix~\ref{appx-inst} for details). %We call the model versions using input graphs as \model-s-g or \model-a-g for the sake of distinction.

\modify{\textbf{Scaling to Large Datasets.} Another advantage of \model over GNNs is the flexibility for mini-batch training. For datasets with prohibitive instance numbers that make it hard for full-batch training on a single GPU, the common practice for GNNs resorts to subgraph sampling~\cite{graphsaint} or graph clustering~\cite{clustergcn-kdd19} 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 \model, we can naturally partition the dataset into random mini-batches and feed one mini-batch for one feed-forward and backward computation. The flexibility for mini-batch 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~\citep{wunodeformer}.}



\textbf{Connection with Existing Models.} Our theory in fact gives rise to a general diffusion framework that unifies some existing models as special cases of ours. As a non-exhaustive summary and high-level comparison, Table~\ref{tbl-existing} presents the relationships with MLP, GCN and GAT (see more elaboration in Appendix~\ref{appx-connection}). %As increasing evidence suggesting the comparable empirical competitiveness of MLPs and propagation-based models (e.g., GNNs) in some practical tasks, the theoretical understandings for different inductive bias still remain unclear. One concurrent work \cite{yang2023graph} points out the fundamental connections among MLPs and GNNs with regard to their generalization effects, while our result in this paper serves as another aspect for understanding the implicit regularization effects w.r.t. embeddings by different propagation layers.
\modify{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)~\citep{GCN-vallina}, it only regularizes the global consistency term constrained within 1-hop neighbors of the observed structure and simplifies $\delta$ as an identity function. The diffusivity of GCN induces information flows through observed edges. For graph attention networks (GAT)~\citep{GAT}, the energy function is similar to GCN's except that the non-linearity $\delta$ remains (as a certain specific form), and the diffusivity is computed by attention over edges.
In contrast with the above models, \model is derived from the general diffusion model that enables non-zero diffusivity values between arbitrary instance pairs.}




\begin{table}[t]
\centering
\caption{Mean and standard deviation of testing accuracy on node classification (with five different random initializations). All the models are split into groups with a comparison of non-linearity (whether the model requires activation for layer-wise transformations), PDE-solver (whether the model requires PDE-solver) and Input-G (whether the propagation purely relies on input graphs).\label{tbl-bench}}
  \vspace{-5pt}
\small
    \resizebox{\textwidth}{!}{
\begin{tabular}{@{}c|c|cccccc@{}}
\toprule
\textbf{Type} & \textbf{Model} & \textbf{Non-linearity} & \textbf{PDE-solver} & \textbf{Input-G}  & \textbf{Cora} & \textbf{Citeseer} & \textbf{Pubmed}   \\ 
\midrule
\multirow{3}{*}{Basic models} & MLP & R   & -  & -   & 56.1 $\pm$ 1.6 & 56.7 $\pm$ 1.7 & 69.8 $\pm$ 1.5  \\
& LP & -  & -  & R   & 68.2 & 42.8 & 65.8 \\
& ManiReg & R    & -  & R  & 60.4 $\pm$ 0.8 & 67.2 $\pm$ 1.6 & 71.3 $\pm$ 1.4 \\
\midrule
\multirow{8}{*}{Standard GNNs} 
& GCN & R   & -  & R   & 81.5 $\pm$ 1.3  & 71.9 $\pm$ 1.9 & 77.8 $\pm$ 2.9 \\
& GAT & R   & -  & R   & 83.0 $\pm$ 0.7 & 72.5 $\pm$ 0.7 & 79.0 $\pm$ 0.3 \\
& SGC & -  & -  & R   & 81.0 $\pm$ 0.0 & 71.9 $\pm$ 0.1 & 78.9 $\pm$ 0.0 \\
& GCN-$k$NN & R   & -  & -    & 72.2 $\pm$ 1.8 & 56.8 $\pm$ 3.2 & 74.5 $\pm$ 3.2 \\
& GAT-$k$NN & R   & -  & -    & 73.8 $\pm$ 1.7 & 56.4 $\pm$ 3.8 & 75.4 $\pm$ 1.3 \\
& Dense GAT & R   & -  & -    & 78.5 $\pm$ 2.5 & 66.4 $\pm$ 1.5 & 66.4 $\pm$ 1.5 \\
& LDS & R   & -  & -    & 83.9 $\pm$ 0.6 & \color{brown}\textbf{74.8 $\pm$ 0.3} & out-of-memory \\
& GLCN & R   & -  & -    & 83.1 $\pm$ 0.5 & 72.5 $\pm$ 0.9 & 78.4 $\pm$ 1.5 \\
\midrule
\multirow{6}{*}{Diffusion-based models} 
& GRAND-l & -  & R   & R   & 83.6 $\pm$ 1.0 & 73.4 $\pm$ 0.5 & 78.8 $\pm$ 1.7 \\
& GRAND & R   & R   & R    & 83.3 $\pm$ 1.3 & 74.1 $\pm$ 1.7 & 78.1 $\pm$ 2.1 \\
& GRAND++ & R   & R   & R    & 82.2 $\pm$ 1.1 & 73.3 $\pm$ 0.9 & 78.1 $\pm$ 0.9 \\
& GDC & R   & -  & R   & 83.6 $\pm$ 0.2 & 73.4 $\pm$ 0.3 & 78.7 $\pm$ 0.4 \\
& GraphHeat & R   & -  & R  & 83.7 & 72.5 & 80.5 \\
& DGC-Euler & -  & -  & R   & 83.3 $\pm$ 0.0 & 73.3 $\pm$ 0.1 & 80.3 $\pm$ 0.1 \\
\midrule
\multirow{3}{*}{Graph Transformers} &
NodeFormer & -  & -  & -  & 83.4 $\pm$ 0.2 & 73.0 $\pm$ 0.3 & \color{brown}\textbf{81.5 $\pm$ 0.4} \\
& {\model}-s & -  & -  & -    & \color{purple}\textbf{85.9 $\pm$ 0.4} & 73.5 $\pm$ 0.3 & \color{purple}\textbf{81.8 $\pm$ 0.3} \\
& {\model}-a & -   & -  & -    & \color{brown}\textbf{84.1 $\pm$ 0.6} & \color{purple}\textbf{75.7 $\pm$ 0.3} & 80.5 $\pm$ 1.2 \\
\bottomrule
\end{tabular}}
\vspace{-15pt}
\end{table}


\vspace{-4pt}
\section{Experiments}\label{sec:exp}
\vspace{-4pt}

We apply \model to various tasks for evaluation: 1) 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. In each case, we compare a different set of competing models closely associated with \model and speciﬁcally 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.~\ref{eqn-diffuse-graph}. Due to space limit, we defer details of datasets to Appendix~\ref{appx-dataset} and the implementations to Appendix~\ref{appx-implementation}. Also, we provide additional empirical results in Appendix~\ref{appx-results}.


\vspace{-4pt}
\subsection{Semi-supervised Node Classification Benchmarks}
\vspace{-4pt}

We test \model on three citation networks \texttt{Cora}, \texttt{Citeseer} and \texttt{Pubmed}. Table~\ref{tbl-bench} reports the testing accuracy. We compare with several sets of baselines linked with our model from different aspects. 1) Basic models: \emph{MLP} and two classical graph-based SSL models Label Propagation (\emph{LP})~\citep{lp-icml2003} and \emph{ManiReg}~\citep{belkin2006manireg}. 2) GNN models: \emph{SGC}~\citep{SGC-icml19}, \emph{GCN}~\citep{GCN-vallina}, \emph{GAT}~\citep{GAT}, their variants \emph{GCN-$k$NN}, \emph{GAT-$k$NN} (operating on $k$NN graphs constructed from input features) and \emph{Dense GAT} (with a densely connected graph replacing the input one), and two strong structure learning models LDS~\citep{LDS-icml19} and GLCN~\citep{jiang2019glcn}. 3) PDE graph models: the SOTA models \emph{GRAND}~\citep{grand} (with its linear variant GRAND-l) and \emph{GRAND++}~\citep{GRAND++}. %(additionally using a source term). 
4) Diffusion-inspired GNN models: \emph{GDC}~\citep{klicpera2019diffusion}, \emph{GraphHeat}~\citep{xu2020heat} and a recent work \emph{DGC-Euler}~\citep{wang2021dissecting}. Table~\ref{tbl-bench} shows that \model achieves the best results on three datasets with significant improvements. Also, we notice that the simple diffusivity model \model-s significantly exceeds the counterparts without non-linearity (SGC, GRAND-l and DGC-Euler) and even comes to the first on \texttt{Cora} and \texttt{Pubmed}. These results suggest that \model 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.



\begin{figure}[t]
\begin{minipage}{0.33\linewidth}
    \centering
	\captionof{table}{Testing ROC-AUC for \texttt{Proteins} and Accuracy for \texttt{Pokec} on large-scale node classification datasets. $*$ denotes using mini-batch training. \label{tbl-largebench}}
	\vspace{-6pt}
    \resizebox{0.98\textwidth}{!}{
\begin{tabular}{@{}c|cc@{}}
\toprule
\textbf{Models} & \textbf{Proteins} & \textbf{Pokec}   \\ 
\midrule
MLP & 72.41 $\pm$ 0.10 & 60.15 $\pm$ 0.03 \\
LP & 74.73 & 52.73\\
SGC &49.03 $\pm$ 0.93 & 52.03 $\pm$ 0.84 \\
GCN & 74.22 $\pm$ 0.49$^*$ & 62.31 $\pm$ 1.13$^*$ \\
GAT & 75.11 $\pm$ 1.45$^*$ & 65.57 ± 0.34$^*$  \\
NodeFormer & \color{brown}\textbf{77.45 $\pm$ 1.15}$^*$ & \color{brown}\textbf{68.32 $\pm$ 0.45}$^*$ \\
\midrule
{\model}-s & \color{purple}\textbf{79.49 $\pm$ 0.44}$^*$ & \color{purple}\textbf{69.24 $\pm$ 0.76}$^*$ \\
\bottomrule
\end{tabular}}
    \end{minipage}
    \hspace{5pt}
    \begin{minipage}{0.66\linewidth}
   \subfigure[]{
    \begin{minipage}[t]{0.48\linewidth}
    \centering
    \includegraphics[width=\linewidth]{figure/cifar-ablation.pdf}
    \label{fig-abl}
    \vspace{-10pt}
    \end{minipage}
    }
    \subfigure[]{
    \begin{minipage}[t]{0.48\linewidth}
    \centering
    \includegraphics[width=\linewidth]{figure/cora_dis_depth2.pdf}
    \label{fig-hyper}
    \vspace{-10pt}
    \end{minipage}
    }
    \vspace{-10pt}
    \caption{(a) Ablation studies w.r.t. different diffusivity function forms on \texttt{CIFAR}. (b) Impact of $K$ and $\tau$ on \texttt{Cora}.}
    \end{minipage}
    \vspace{-10pt}
\end{figure}

\begin{table}[tb!]
		\centering
		\small
		\caption{Testing accuracy for image (\texttt{CIFAR} and \texttt{STL}) and text (\texttt{20News}) classification.}
  \vspace{-5pt}
		\label{tab:image classification}
		\resizebox{0.95\textwidth}{!}{
		\begin{tabular}{c|c|ccccccc|cc}
		\midrule
		\multicolumn{2}{c|}{\textbf{Dataset}}  & MLP & LP & ManiReg  & GCN-$k$NN & GAT-$k$NN & DenseGAT & GLCN  & \model-s & \model-a  \\
		\midrule
		\multirow{3}{*}{\textbf{CIFAR}} 
		& 100 labels &65.9 ± 1.3&66.2&67.0 ± 1.9&66.7 ± 1.5&66.0 ± 2.1 & out-of-memory &66.6 ± 1.4&\color{brown}\textbf{69.1 ± 1.1}&\color{purple}\textbf{69.3 ± 1.4} \\
		& 500 labels &73.2 ± 0.4&70.6&72.6 ± 1.2&72.9 ± 0.4&72.4 ± 0.5 & out-of-memory &72.8 ± 0.5&\color{purple}\textbf{74.8 ± 0.5}&\color{brown}\textbf{74.0 ± 0.6} \\
		& 1000 labels &75.4 ± 0.6&71.9&74.3 ± 0.4&74.7 ± 0.5&74.1 ± 0.5 & out-of-memory &74.7 ± 0.3&\color{purple}\textbf{76.6 ± 0.3}&\color{brown}\textbf{75.9 ± 0.3} \\
		\midrule
		\multirow{3}{*}{\textbf{STL}} 
		& 100 labels &66.2 ± 1.4&65.2&66.5 ± 1.9&\color{brown}\textbf{66.9 ± 0.5}&66.5 ± 0.8 & out-of-memory &66.4 ± 0.8&\color{purple}\textbf{67.8 ± 1.1}&66.8 ± 1.1 \\
		& 500 labels &\color{brown}\textbf{73.0 ± 0.8}&71.8&72.5 ± 0.5&72.1 ± 0.8&72.0 ± 0.8 & out-of-memory &72.4 ± 1.3&\color{purple}\textbf{73.7 ± 0.6}&72.9 ± 0.7 \\
		& 1000 labels &75.0 ± 0.8&72.7&74.2 ± 0.5&73.7 ± 0.4&73.9 ± 0.6 & out-of-memory &74.3 ± 0.7&\color{purple}\textbf{76.4 ± 0.5}&\color{brown}\textbf{75.3 ± 0.6} \\
		\midrule
		\multirow{3}{*}{\textbf{20News}} 
		& 1000 labels &54.1 ± 0.9&55.9&56.3 ± 1.2&56.1 ± 0.6&55.2 ± 0.8 & 54.6 ± 0.2 &56.2 ± 0.8&\color{brown}\textbf{57.7 ± 0.3}&\color{purple}\textbf{57.9 ± 0.7} \\
		& 2000  labels &57.8 ± 0.9&57.6&60.0 ± 0.8&60.6 ± 1.3&59.1 ± 2.2 & 59.3 ± 1.4 &60.2 ± 0.7&\color{brown}\textbf{61.2 ± 0.6}&\color{purple}\textbf{61.3 ± 1.0}  \\
		& 4000 labels &62.4 ± 0.6&59.5&63.6 ± 0.7&64.3 ± 1.0&62.9 ± 0.7 & 62.4 ± 1.0 &64.1 ± 0.8&\color{purple}\textbf{65.9 ± 0.8}&\color{brown}\textbf{64.8 ± 1.0}  \\
		\midrule
	\end{tabular}				
		}
\vspace{-15pt}
	\end{table}
 
\vspace{-4pt}
\subsection{Large-Scale Node Classification Graphs}
\vspace{-4pt}

We also consider two large-scale graph datasets \texttt{ogbn-Proteins}, a multi-task protein-protein interaction network, and \texttt{Pokec}, a social network. %For \texttt{Proteins}, we follow the splitting and evaluation protocol by \cite{ogb-nips20}; for \texttt{Pokec} we randomly choose 10\%/10\%/80\% instances for training/validation/test and use Accuracy for evaluation following \cite{newbench}.
Table~\ref{tbl-largebench} presents the results. Due to the dataset size (0.13M/1.63M nodes for two graphs) and scalability issues that many of the competitors in Table~\ref{tbl-bench} as well as \model-a would potentially experience, we only compare \model-s with standard GNNs. In particular, we found GCN/GAT/\model-s are still hard for full-graph training on a single V100 GPU with 16GM memory. We thus consider mini-batch training with batch size 10K/100K for \texttt{Proteins}/\texttt{Pokec}. %\modify{Besides, we compare with a recently proposed Graph Transformer (GT) model~\cite{graphformer-neurips21} as a strong competitor that considers all-pair propagation.}
We found that \model outperforms common GNNs by a large margin, which suggests its desired efficacy on large datasets. As mentioned previously, we prioritize the efficacy of \model as a 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~\citep{adadiffusion-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-batch sizes for training and study its impact on testing performance in Appendix~\ref{appx-batch}. Furthermore, we compare the training time and memory costs in Appendix~\ref{appx-time} where we found that \model-s is about 6 times faster than GAT and 39 times faster than DenseGAT on \texttt{Pokec}, which suggests superior scalability and efficiency of \model-s on large graphs.






\vspace{-4pt}

\subsection{Image and Text Classification with Low Label Rates}
\vspace{-4pt}


We next conduct experiments on \texttt{CIFAR-10}, \texttt{STL-10} and \texttt{20News-Group} datasets to test \model for standard classification tasks with limited label rates.
For \texttt{20News} provided by \cite{pedregosa2011scikit}, we take 10 topics and use words with TF-IDF more than 5 as features. %For evaluation purpose, we randomly select 100/200/400 instances per class as training set, 2000 instances for validation and the remaining xxx instances for testing.
For \texttt{CIFAR} and \texttt{STL}, two public image datasets, we first use the self-supervised approach SimCLR~\citep{chen2020simple} (that does not use labels for training) to train a ResNet-18 for extracting the feature maps as input features of instances.
These datasets contain no graph structure, so we use $k$NN to construct a graph over input features for GNN competitors and do \emph{not} use input graphs for \model.

Table~\ref{tab:image classification} reports the testing accuracy of \model and competitors including MLP, ManiReg, GCN-$k$NN, GAT-$k$NN, DenseGAT and GLCN. 
Two \model models perform much better than MLP in nearly all cases, suggesting the effectiveness of learning the inter-dependence over instances. Besides, \model yields large improvements over GCN and GAT which are in some sense limited by the handcrafted graph that leads to sub-optimal propagation. Moreover, \model significantly outperforms GLCN, a strong baseline that learns new (static) graph structures, which demonstrates the superiority of our evolving diffusivity that can adapt to different layers.




\vspace{-4pt}
\subsection{Spatial-Temporal Dynamics Prediction}
\vspace{-4pt}

We consider three spatial-temporal datasets with details in Appendix~\ref{appx-dataset}. Each dataset consists of a series of graph snapshots where nodes are treated as instances and each of them has a integer label (e.g., reported cases for \texttt{Chickenpox} or \texttt{Covid}). The task is to predict the labels of one snapshot based on the previous ones. %For each dataset, to make it more challenging and close to the typical real-world low-data regime, we split the snapshots into training, validation and test sets by 2:2:6. Since the testing snapshots contain new instances (or unseen graph structures), such a scenario requires inductive learning.
Table~\ref{tab:spatial_temporal_results} compares testing MSE of four \model variants (here \model-s w/o g denotes the model \model-s without using input graphs) against baselines. We can see that two \model 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 to the competitors, our models rank the first with significant improvements. %In particular, for \texttt{Covid} \model-s achieves 8.1\% lower MSE than the best competitor GAT-$k$NN and for \texttt{WikiMath} \model-a decreases the MSE by 12.8\% than the runner-up Dense GAT, which demonstrates the superiority of \model for temporal dynamics prediction. 


\begin{table}[t!]
\centering
\caption{Mean and standard deviation of MSE on spatial-temporal prediction datasets.}
  \vspace{-5pt}
\label{tab:spatial_temporal_results}
\small
    \resizebox{1.0\textwidth}{!}{
\begin{tabular}{c|cccccc|cccc}
\toprule
\textbf{Dataset}& MLP&GCN&  GAT&Dense GAT&GAT-$k$NN&GCN-$k$NN&{\model}-s& {\model}-a& {\model}-s w/o g& {\model}-a w/o g \\ 
\midrule
\multirow{2}{*}{\textbf{Chickenpox}}& 0.924&0.923&0.924&0.935&0.926&0.936&\color{purple}\textbf{0.914} & \color{brown}\textbf{0.915} & 0.916 & 0.916
\\
& ($\pm$0.001)&($\pm$0.001)&($\pm$0.002)&($\pm$0.005)&($\pm$0.004)&($\pm$0.004)&\color{purple}\textbf{(0.006)} & \color{brown}\textbf{(0.008)} & (0.006) & (0.006)\\
\midrule
\multirow{2}{*}{\textbf{Covid}}& 0.956&1.080&1.052&1.524&0.861&1.475& 0.779 & \color{brown}\textbf{0.757} & 0.779 & \color{purple}\textbf{0.741}\\
& ($\pm$0.198)&($\pm$0.162)&($\pm$0.336)&($\pm$0.319)&($\pm$0.123)&($\pm$0.560)& (0.037) & \color{brown}\textbf{(0.048)} & (0.028) & \color{purple}\textbf{(0.052)}\\
\midrule
\multirow{2}{*}{\textbf{WikiMath}}& 1.073&1.292&1.339&0.826&0.882&1.023&0.731 & 0.763 & \color{brown}\textbf{0.727} & \color{purple}\textbf{0.716}\\
& ($\pm$0.042)&($\pm$0.125)&($\pm$0.073)&($\pm$0.070)&($\pm$0.015)&($\pm$0.058)&(0.007) & (0.020) & \color{brown}\textbf{(0.025)} & \color{purple}\textbf{(0.030)}\\
\bottomrule
\end{tabular}}
\vspace{-15pt}
\end{table}


\vspace{-4pt}

\subsection{Further Results and Discussions}
\vspace{-4pt}

\textbf{How do different diffusivity functions perform?} Figure~\ref{fig-abl} compares \model with four variants using other diffusivity functions that have no essential connection with energy minimization: 1) \emph{Identity} sets $\mathbf S^{(k)}$ as a fixed identity matrix; 2) \emph{Constant} fixes $\mathbf S^{(k)}$ as all-one constant matrix; 3) \emph{Full Attn} parameterizes $\mathbf S^{(k)}$ by attention networks~\citep{transformer}; 4) \emph{Kernel} adopts Gaussian kernel for computing $\mathbf S^{(k)}$. More results on other datasets are in Appendix~\ref{appx-abl} and they consistently show that our adopted diffusivity forms produce superior performance, which verifies the effectiveness of our diffusivity designs derived from minimization of a principled energy.

\textbf{How do model depth and step size impact the performance?} We discuss the influence of model depth $K$ and step size $\tau$ on \texttt{Cora} in Fig.~\ref{fig-hyper}. More results are deferred to Appendix~\ref{appx-hyper}. The curves indicate that the step size impacts the model performance as well as the variation trend w.r.t. model depth. When using larger step size (e.g., $\tau=0.5$), the model can yield superior performance with shallow layers; when using smaller step size (e.g., $\tau=0.1$), the model needs to stack deep layers for performing competitively. The reason could be that larger $\tau$ contributes to more concentration on global information from other instances in each diffusion iteration, which brings up more beneficial information with shallow layers yet could lead to instability for deep layers. And, using small $\tau$ can enhance the model's insensitivity to deep model depth, which often leads to the performance degradation of other baselines like GCN and GAT.



\vspace{-4pt}

\section{Conclusions}
\vspace{-4pt}

\looseness=-1 This paper proposes an 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 a principled energy as implicit regularization. We further design two practical implementations with enough scalability and capacity for learning complex interactions over the underlying data geometry. Extensive experiments demonstrate the effectiveness and superiority of the model.
