DIFFormer: Scalable (Graph) Transformers Induced by Energy Constrained Diffusion

• Real-world interactions violate IID assumptions, while unknown data-generation principles make uncovering dependencies practically difficult.Source: §1 · p. 1
• Observed graphs can be incomplete, noisy, or artificially constructed, creating potential inconsistency between graph structure and underlying data geometry.Source: §1 · p. 1
• The proposed encoder treats dataset instances as locations on latent Riemannian manifolds and uses diffusion to propagate information between arbitrary instance pairs.Source: §1 · p. 2
• A principled energy function regularizes evolving representations toward internal consistency and guides the diffusion process.Source: §1 · p. 2
• Theory establishes equivalence between diffusion iterations and regularized-energy minimization, yielding closed-form optimal diffusivity updates.Source: §1 · p. 2
• DIFFORMER offers an O(N) variant for all-pair interactions and a more expressive variant for complex latent structures, with strong results across diverse tasks.Source: §1 · p. 2

• Graph-based semi-supervised learning uses partially labeled nodes and observed relations, but those structures may be unavailable or unreliable.Source: §2 · p. 2
• The paper explores a Transformer-like encoder that discovers data geometry through interdependence among labeled and unlabeled instances.Source: §2 · p. 3
• Unlike prior diffusion-based approaches, the proposed dynamics are implicitly defined by optimizing a regularized energy.Source: §2 · p. 3
• Its theory links numerical diffusion iterations with energy optimization and supports encoders for uncovering latent structures among many instances.Source: §2 · p. 3
• Transformers model token dependencies through all-pair attention, whereas this work targets analogous interdependence among instances in a dataset.Source: §2 · p. 3

Geometric diffusion model

• An anisotropic diffusion process evolves all instance representations jointly, allowing information to flow between dataset instances.The diffusivity controls pairwise flow strength and may depend on instance states.Source: §3.1 · pp. 3–4
• The diffusivity can range from self-loops or observed graph edges to adaptive interactions among arbitrary instance pairs.Self-loop-only diffusion reduces to an MLP, while graph-based diffusion restricts propagation to neighboring nodes.Source: §3.1 · p. 4

Diffusion constrained by a layer-wise energy

• The layer-wise energy regularizes states toward both proximity to current representations and global consistency across instances.The concave non-decreasing function δ promotes robustness against large pairwise differences.Source: §3.2 · p. 4

Energy-constrained diffusion

• The resulting geometric flow is implicitly defined by minimizing a time-varying energy on latent manifolds.The framework links microscopic state evolution with macroscopic consistency measured by the energy.Source: §3.2 · p. 4

Tractability of solving diffusion process with energy minimization

• Theorem 1 provides a closed-form diffusivity based on current state distances that guarantees an energy-descent step.For an appropriate step size 0 < τ < 1, the update satisfies E(Z(k+1), k; δ) ≤ E(Z(k), k −1; δ).Source: §3.3 · p. 5
• The descent guarantee unfolds the implicit diffusion process into layer-wise feed-forward neural computations.The resulting architecture computes diffusivity from current states and combines state conservation with state propagation.Source: §3.3 · p. 5

Model instantiations

• The simple diffusivity model uses a linear function and reduces each iteration's complexity from pairwise O(N^2) computation to O(N).Shared summation terms enable the linear-complexity implementation DIFFORMER-s.Source: §4.1 · pp. 5–6

Connections with existing models

• The energy-driven diffusion framework presents MLP, GCN, and GAT as special cases differing in energy forms, diffusivity, and complexity.The comparison is summarized in Table 1.Source: Table 1 · p. 6, §4.2 · p. 7

Model instantiations

• The advanced diffusivity model uses a sigmoid-based function to provide greater capacity for complex latent geometry.DIFFORMER-a preserves monotonic decrease and non-negativity through its associated function.Source: §4.1 · p. 6

Model extensions and further discussion

• Layer-wise feature transformations can increase representation capacity, particularly on larger datasets, while linear transformations perform well empirically.The transformations are not necessary for small datasets according to the experiments described.Source: §4.2 · p. 6
• Observed input graphs can be incorporated as geometric priors by restricting diffusion with graph-based penalties.The resulting iteration is a descent step on an energy that additionally includes a graph-based penalty.Source: §4.2 · p. 6

Scaling to large datasets

• DIFFORMER supports random mini-batch training, avoiding the subgraph sampling or clustering strategies commonly used to scale GNNs.Mini-batch training also supports parallel acceleration and federated learning.Source: §4.2 · pp. 6–7

Semi-supervised node classification

• DIFFormer achieves the best results on Cora, Citeseer, and Pubmed, with DIFFormer-s leading on Cora and Pubmed.The comparison includes classical, GNN, PDE-based, and diffusion-inspired baselines.Source: §5.1 · p. 8

Large-scale node classification graphs

• On ogbn-Proteins and Pokec, DIFFormer-s outperforms common GNNs by a large margin under mini-batch training.The datasets contain 0.13M and 1.63M nodes, respectively.Source: §5.2 · p. 8, Table 3 · p. 8

Image and text classification with low label rates

• On CIFAR-10, STL-10, and 20News-Group, DIFFormer substantially outperforms MLP and graph-based competitors under limited-label classification.GNN competitors use kNN graphs, whereas DIFFormer does not use input graphs.Source: Table 4 · p. 8, §5.3 · p. 9

Spatial-temporal dynamics prediction

• DIFFormer variants rank first with significant improvements on spatial-temporal prediction, and graph-free variants outperform graph-using counterparts in four of six cases.The task predicts one snapshot from previous snapshots using testing MSE.Source: §5.4 · p. 9, Table 5 · p. 9

Further results and discussions

• Energy-derived diffusivity functions outperform identity, constant, attention, and Gaussian-kernel alternatives across reported datasets.The ablation supports diffusivity designs derived from energy minimization.Source: Figure 2 · p. 8, §5.5 · p. 9
• Larger step sizes such as τ = 0.5 perform better with shallow models, whereas smaller τ = 0.1 requires deeper models for competitive performance.The analysis is conducted on Cora and examines model depth K and step size τ.Source: Figure 2 · p. 8, §5.5 · p. 9

Contributions

• The proposed model encodes all instances jointly through geometric diffusion aimed at minimizing a principled energy.The energy acts as implicit regularization for evolving instance representations.Source: §6 · p. 9
• DIFFormer provides a scalable O(N) implementation and a more expressive implementation for learning complex latent interactions.The two variants target prohibitive instance counts and richer latent structures, respectively.Source: §6 · p. 9

Scope and positioning

• The approach addresses settings where observed graph structures are unavailable, incomplete, noisy, or unreliable by learning adaptive instance dependencies.This positions DIFFormer alongside graph structure learning and Transformer-like all-pair dependency modeling.Source: §A.4 · p. 15

Energy minimization

• The variational upper bound converts the original energy minimization into a tractable surrogate optimization problem.Fenchel-duality-based decomposition supplies the upper bound and its equality condition.Source: B PROOF FOR THEOREM 1 · p. 15
• Explicit-Euler diffusion updates are equivalent to one-step gradient descent on the surrogate and therefore decrease the original energy.The construction yields E(Z(k+1), k; δ) ≤ E(Z(k), k; δ).Source: B PROOF FOR THEOREM 1 · pp. 15–16
• The energy remains non-increasing across successive diffusion steps because each update decreases the current energy and the optimized state minimizes the relevant intermediate objective.The proof combines E(Z(k+1), k; δ) ≤ E(Z(k), k; δ) with E(Z(k), k; δ) ≤ E(Z(k), k −1; δ).Source: B PROOF FOR THEOREM 1 · p. 16

Diffusivity design

• Valid diffusivity functions must be non-negative and decreasing in z2, ensuring a non-decreasing, concave penalty δ(z2).The simple and advanced models instantiate these conditions with different diffusivity and penalty forms.Source: D DIFFERENT ENERGY FORMS · p. 17

DIFFORMER implementation

• The encoder can incorporate input-graph propagation, produce logits for classification or regression, and use identity or nonlinear layer updates.Graph-based propagation is added to all-pair propagation before output prediction.Source: D DIFFERENT ENERGY FORMS · p. 17, §E.1 · p. 18
• DIFFORMER combines diffusion-induced propagation with either linear-complexity all-pair aggregation or more expressive similarity computation.The simple variant uses linear complexity in N, while explicit all-pair computation requires O(Nd2 + N^2d).Source: §E.1 · p. 18

Dataset overview

• The section provides detailed information about experimental datasets, preprocessing, and evaluation protocols used in the experiments.Node-classification dataset statistics are summarized in Table 6.Source: G DATASET INFORMATION · p. 21, Table 6 · p. 21

Citation networks

• Cora, Citeseer, and Pubmed are citation networks used to classify document topics from bag-of-words features and citation links.These networks contain approximately 2K–20K nodes and use 20 labeled instances per class for training.Source: §G.1 · p. 21

Protein-interaction network

• OGBN-Proteins evaluates molecule-property prediction on a multi-task protein-protein interaction network using its original data split.The dataset is described as a protein-interaction network with molecule instances as prediction targets.Source: §G.1 · p. 21

Social network

• Pokec is a large-scale social network whose profile features support user-gender prediction with a 10%/10%/80% train/validation/test split.Profile information includes geographical region, registration time, and age.Source: §G.1 · p. 21

Image classification

• STL-10 and CIFAR-10 provide image-classification tasks with ten classes and limited labeled training instances per class.Training uses 10, 50, or 100 instances per class, with 1,000 validation instances.Source: §G.2 · p. 21
• Image inputs are represented by feature maps extracted from a SimCLR-trained ResNet-18 without using labels during SimCLR training.The extracted feature maps serve as instance input features for the classifier.Source: §G.2 · p. 21

Text classification

• 20Newsgroup supplies a text-classification task with 9,607 instances, using ten classes and TFIDF features selected by frequency.The setup follows prior work and retains words with frequency greater than 5%.Source: §G.2 · p. 21

Dataset setup

• Snapshots are split into training, validation, and test sets at a 2:2:6 ratio for low-data spatial-temporal prediction.Node features evolve across snapshots for every dataset.Source: §G.3 · p. 22

Datasets

• Chickenpox models weekly county-level cases in Hungary using neighborhood edges and four lagged weekly-count features.The target is the following week's signed case count.Source: §G.3 · p. 22
• Covid models daily regional mobility in England using recent case counts to predict each region's cases one day ahead.Edges encode daily movement between NUTS3 regions.Source: §G.3 · p. 22
• WikiMath models daily Wikipedia visits over a directed, weighted page-link graph.Node features are historical daily visits, and edge weights count links from source to target pages.Source: §G.3 · p. 22

Dataset statistics

• Table 7 summarizes graph dynamics, node-feature meanings and dimensions, snapshot counts, node counts, and edge semantics for the spatial-temporal datasets.These properties distinguish whether graph structure is dynamic or static.Source: Table 7 · p. 22

Experimental setup

• Image and text datasets use learned diffusion structures because they do not provide input graphs.The experiments use feature transformation for layer-wise updating and set τ = 0.5.Source: §H.2 · p. 22

Energy-function ablation

• Figure 4 compares alternative energy-function forms used by DIFFormer on image and text datasets.The accompanying discussion analyzes these energy-function ablations.Source: Figure 4 · p. 23, §I.1 · p. 23

Evaluation

• The study averages results over five runs and reports standard deviations while selecting test accuracy at the best validation epoch.Training runs for at most 600 epochs.Source: §H.2 · p. 23

Supplementary analyses

• Supplementary experiments add ablations, hyper-parameter studies, and visualizations on datasets omitted from the main section.These results address energy functions, step size, model depth, and learned representations.Source: I MORE EXPERIMENT RESULTS · p. 23

Hyper-parameter study

• Figure 5 plots testing performance against model size K for several baselines and DIFFormer variants under different step sizes τ.The comparison is intended to assess how step size interacts with increasing model depth.Source: §I.2 · p. 23

Representation visualization

• On 20News and STL, learned diffusivity connects instances across classes, while embeddings show small intra-class and large inter-class distances.The resulting separation makes the representations easier for a classifier to distinguish.Source: §I.3 · pp. 23–24

Diffusion visualization

• On Chickenpox, large diffusion strengths usually connect nodes with similar ground-truth labels.DIFFormer-s concentrates large weights, whereas DIFFormer-a spreads them more broadly to model more complex structures.Source: §I.3 · p. 24

Mini-batch impact

• Small mini-batches sacrifice performance, whereas large batches produce decent, low-variance results with acceptable memory costs.The study specifically examines this trade-off on Pokec.Source: §I.4 · p. 24, Table 8 · p. 24

Efficiency and scalability

• DIFFormer-s has comparable training time to GAT on smaller datasets and is much faster on Pokec.It also reduces memory costs several times relative to DenseGAT.Source: §I.5 · p. 24, Table 9 · p. 25

Pseudo-label compatibility

• DIFFormer is technically orthogonal to pseudo-label methods because it provides an encoder backbone trained with standard supervised loss.Meta Pseudo Labels can therefore be used as a plug-in with DIFFormer and competing backbones.Source: §I.6 · p. 25, Table 10 · p. 25