PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Qitian Wu, Chenxiao Yang, Wentao Zhao, Yixuan He, David Wipf, Junchi Yan

arXiv:2301.09474v4cs.LGcs.AI

TL;DR

Real-world data violate IID assumptions and make latent instance dependencies difficult to uncover, especially when observed graphs are incomplete or unreliable. The paper introduces energy-constrained diffusion encoders that propagate information across instances while minimizing a regularized energy, with scalable and expressive DIFFORMER variants. Experiments report strong performance across node classification, image/text classification, and spatial-temporal prediction tasks.

  • Problem

    Real-world data are non-IID, while unknown generation processes and unreliable observed graphs make underlying instance dependencies difficult to uncover.

  • Method

    The model encodes all dataset instances through diffusion dynamics constrained by a principled energy, with variants supporting O(N) interactions or complex latent structures.

  • Results

    DIFFORMER outperforms state-of-the-art approaches on semi-supervised node classification and performs competitively on large-scale graphs, while also supporting image/text classification and spatial-temporal prediction.

  • Takeaways & Limitations

    DIFFORMER provides a general-purpose encoder backbone for learning inter-instance interactions across diverse datasets and prediction tasks.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper addresses the difficulty of learning useful instance dependencies when real-world data are non-IID and observed graphs may be incomplete or noisy. It proposes an energy-constrained diffusion encoder with scalable and expressive variants, supported by theory and broad experiments.

  • Real-world interactions violate IID assumptions, while unknown data-generation principles make uncovering dependencies practically difficult.
  • Observed graphs can be incomplete, noisy, or artificially constructed, creating potential inconsistency between graph structure and underlying data geometry.
  • The proposed encoder treats dataset instances as locations on latent Riemannian manifolds and uses diffusion to propagate information between arbitrary instance pairs.
  • A principled energy function regularizes evolving representations toward internal consistency and guides the diffusion process.
  • Theory establishes equivalence between diffusion iterations and regularized-energy minimization, yielding closed-form optimal diffusivity updates.
  • 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.

2 RELATED WORK

The paper builds on graph-based learning, neural diffusion, and Transformers while targeting latent inter-instance dependencies without relying solely on observed structures. Its distinctive contribution is an energy-optimized diffusion framework with theory supporting new scalable encoders.

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

3 ENERGY CONSTRAINED GEOMETRIC DIFFUSION TRANSFORMERS

The model treats dataset instances as states in an anisotropic diffusion process, with pairwise information flow guided by an energy measuring representation consistency. A theorem supplies diffusivities that make each update decrease this energy and enables feed-forward neural encoding.

  • 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.
  • Geometric diffusion model: 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.
  • 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.
  • 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.
  • 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; δ).
  • Tractability of solving diffusion process with energy minimization: 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.

4 INSTANTIATIONS OF DIFFORMER

DIFFORMER is instantiated through simple and advanced diffusivity functions, trading scalability against capacity for complex latent geometry. The framework also supports transformations, input graphs, mini-batch training, and connections to established neural architectures.

  • 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.
  • 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.
  • 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.
  • 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.
  • Model extensions and further discussion: 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.
  • 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.

5 EXPERIMENTS

Experiments evaluate DIFFormer across node classification, image/text classification, and spatial-temporal prediction. Results show strong performance across scales and settings, while ablations examine diffusivity choices, depth, and step size.

  • 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.
  • 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.
  • 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.
  • 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.
  • 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.
  • Further results and discussions: 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 τ.

6 CONCLUSIONS

The paper presents DIFFormer as an energy-driven diffusion encoder for learning interactions among dataset instances. Its theory motivates scalable implementations, while experiments support its applicability across graph, image, text, and spatial-temporal tasks.

  • 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.
  • Contributions: 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.
  • 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.

B PROOF FOR THEOREM 1

The proof links diffusion iterations to energy minimization, establishing monotonic energy descent and motivating diffusivity choices and DIFFORMER implementations.

  • 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.
  • Energy minimization: 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; δ).
  • Energy minimization: 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; δ).
  • 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.
  • 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.
  • DIFFORMER implementation: 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).

G DATASET INFORMATION

The dataset information section introduces the evaluation datasets and points to a table summarizing node-classification dataset statistics.

  • 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.

G.1 NODE CLASSIFICATION DATASETS

Node-classification evaluation covers citation, protein-interaction, and social-network datasets with task-specific features, splits, and labels.

  • 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.
  • 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.
  • 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.

G.2 IMAGE AND TEXT CLASSIFICATION DATASETS

Image and text classification experiments use STL-10, CIFAR-10, and 20Newsgroup with sparse supervision and feature representations derived from images or text.

  • 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.
  • Image classification: 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.
  • 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%.

G.3 SPATIAL-TEMPORAL DATASETS

The spatial-temporal experiments use evolving node features across graph snapshots and evaluate forecasting under challenging train/validation/test splits. The datasets cover epidemiological, mobility, and Wikipedia-traffic dynamics with distinct graph structures and prediction targets.

  • 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.
  • 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.
  • Datasets: 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.
  • Datasets: 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.
  • 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.

H.2 IMAGE AND TEXT CLASSIFICATION EXPERIMENT

The image and text experiments apply layer-wise feature transformation and learn diffusion structures without input graphs. Supplementary analyses examine energy-function ablations and the effects of step size and model depth.

  • 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.
  • 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.
  • 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.
  • 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.
  • 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.

I.3 VISUALIZATION

Visualization and scalability analyses show how DIFFormer’s learned diffusion strengths shape representations and how mini-batch size affects large-graph training. The results connect diffusion patterns to class structure while highlighting efficiency trade-offs.

  • 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.
  • 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.
  • 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.
  • 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.
  • 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.
Loading 2301.09474v4…