Geometric Generative Modeling with Noise-Conditioned Graph Networks
Authors: Peter Pao-Huang
Mitchell Black
Xiaojie Qiu
Venue: ICML 2025 poster (2025)
Keywords: graph neural networks
generative models
diffusion models
flow-matching
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: Noise-Conditioned Graph Networks is a new class of graph neural networks for more efficient and expressive generative modeling.

Abstract: Generative modeling of graphs with spatial structure is essential across many applications from computer graphics to spatial genomics. Recent flow-based generative models have achieved impressive results by gradually adding and then learning to remove noise from these graphs. Existing models, however, use graph neural network architectures that are independent of the noise level, limiting their expressiveness. To address this issue, we introduce *Noise-Conditioned Graph Networks* (NCGNs), a class of graph neural networks that dynamically modify their architecture according to the noise level during generation. Our theoretical and empirical analysis reveals that as noise increases, (1) graphs require information from increasingly distant neighbors and (2) graphs can be effectively represented at lower resolutions. Based on these insights, we develop Dynamic Message Passing (DMP), a specific instantiation of NCGNs that adapts both the range and resolution of message passing to the noise level. DMP consistently outperforms noise independent architectures on a variety of domains including $3$D point clouds, spatiotemporal transcriptomics, and images.

Geometric Generative Modeling with Noise-Conditioned Graph Networks
Peter Pao-Huang' Mitchell Black? Xiaojie Qiu
Abstract
biology, atoms amino acids serve nodes with 3D co- ordinates while chemical properties constitute the features, accelerating the design of novel drugs Ko6mez-Bombarelll et al., 2018; Bengio et al., 2021) and predicting the structure of proteins (Watson et al., 2023; Jumper et al., 2021). In emerging fields like spatial single-cell genomics, geomet- ric graphs capture both the spatial distribution of cells and their gene expression profiles, providing insights into tissue organization and development (Qiu al., 2024).
Generative modeling of graphs with spatial struc- ture is essential across many applications from computer graphics to spatial genomics. Recent flow-based generative models have achieved im- pressive results by gradually adding and then learning to remove noise from these graphs. Ex- isting models, however, use graph neural net- work architectures that are independent of the noise level, limiting their expressiveness. To ad- dress this issue, we introduce Noise- Conditioned Graph Networks (NCGNs), a class of graph neu- ral networks that dynamically modify their archi- tecture according to the noise level during gen- eration. Our theoretical and empirical analysis reveals that as noise increases, (1) graphs require information from increasingly distant neighbors and (2) graphs can be effectively represented lower resolutions. Based on these insights, we develop Dynamic Message Passing (DMP), spe- cific instantiation of NCGNs that adapts both the range and resolution of message passing to the noise level. DMP consistently outperforms noise- independent architectures on a variety of domains including 3D point clouds, spatiotemporal tran- scriptomies, and images. Code is available
To model the distribution of these geometric graphs, recent advances in generative modeling, particularly flow-based approaches such as diffusion (Ho et al., 2020; Song & Er- mon, 2019) and flow-matching (Lipman et al., 2022: Tong et al., 2023) models, have shown remarkable success in many scientific domains (Jing al., 2023; Pao-Huang et al., 2023; Luo et al., 2022). These methods work by learning an iterative transformation between simple prior distribution (typically Gaussian) and the target data distribution. Their effectiveness can be partially attributed to the surprisingly simple training objective: transform or interpolate the data distribution towards a Gaussian and then learn to reverse this transformation. Colloquially, the interpolation towards Gaussian can be called a noising processing.
While existing graph neural networks for these flow-based generative modeling have demonstrated promising results, they rely on a afixed (aka noise- sniie-independent) graph neural net- work architecture across the noising process. For instance, previous works (Corso et al., 2022; Luo al., 2021; Jing et al., 2022) construct message-passie operations using constant radius k-nearest neighbor graph independent of the noise level. This static approach limits the ability to fully capture how the underlying graph signal evolves under different noise levels in the generative process.
1. Introduction
Geometric graphs, defined by both features and positional information, are powerful representation in many scien- tific domains. In computer vision and graphics, these graphs represent 3D point clouds and shapes, enabling generative applications like shape synthesis (Nash et al., 2020) and point cloud generation (Yang et al., 2019). In molecular
Contributions. This limitation motivates our investigation into the behavior of geometric graphs under varying noise levels. Through an information-theoretiic analysis, we prove that noise increases, recovering the underlying graph signal requires aggregating information from increasingly distant neighbors. This theoretical finding has two important implications. First, pesssage-passing architectures should increase their connectivity radius with the noise level since local structure becomes less informative with increasing noise. We show that this holds in practice by analyzing the
Department of Computer Science, Stanford University 2 Departmeet of Computer Science, University of California San Diego 'Department of Geneties, Stanford University. Correspon- dence Peter Pao- Huang <peterph @csbhhhcss.t edu.,, Xiaojie Qiu <xiaojie @ stanford.edu>..
Proceedings of the 4dnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Nolse-Conditooned Graph Networks
Range
Resolution
Position
Feature
Increasing
Increasing
Nase Candilo
Resolution
Noise Conditonal Resolution
Figure 1. Comparing the range and resolution between noise-independent graph neural networks and noise-conditioned graph neural networks (GNNs). (Top Row) Noise is introduced in the positions of the star while features are fixed. (Bottom Row) Noise is introduced in the features of the star while positions fixed. (Left Column) As noise increases, the range of noise-independent GNNs remains the same while noise-conditional GNNs increases. (Right Column) As noise increases, the resolution of noise-independent GNNs remains fixed (shown circular nodes) while noise-conditioned GNNs decreases (coarse-grained nodes shown nodes).
weights of a trained graph attention network used within flow-based generative model. Second, at higher noise levels, coarse-grasd graniiigg through pooling operations can better preserve the graph signal while reducing computational cost. also show this empirically by measuring discrepancy between dataset of geometric graphs and their noised, coarse-grained counterparts.
2. Background
2.1. Geometric Generative Modeling
In this work, we focus on modeling geometric graphs, which extend traditional graphs by incorporating spatial informa- tion. Formally, we define geometric graph tuple G = (X, R, E) where X {x E R are node fea- tures, R = (7)()) E RdJN Ji=1 are spatial positions (in this paper, d 3)), and & ((x))
Leveraging these theoretical and empirical insights, pro- pose
Graph Networks (NCGNs), new class of graph neural networks that dynamically adapts the graph structure based on the noise level of the generative model. We then develop specific instantiation NCGNs called Dynamic Message Passing (DMP). DMP interpo- lates the range and resolution of the message passing graph from fully connected low-resolution graph high noise to sparsely connected high-resolution graph at low noise. By simultaneously adjusting both connectivity and node resolution according to the noise level, DMP achieves bet- ter expressivity than inseeenneennn architectures while maintaining linear-time message passing complexity.
the edges, where N is the number of nodes. In this paper, the edges & are function of the positions R; for example, they might be the edges of the k-nearest neighbor graph. We additionally constrain our definition of geometric generative modeling to learning either the positional distribution conditioned on features pl X) the feature distribution conditioned on positions (X))
This general learning problem has applications in many fields. For instance, quantum systems can be represented by treating particles as nodes with quantum numbers as features and spatial coordinates as positions. Climate sci- ence applications represent weather stations as nodes with meteorological measurements features. Even traditional computer vision tasks be reformulated under this ge- ometric paradigm by treating pixels as nodes on regular lattice with RGB features.
Through extensive experiments, we demonstrate the effec- tiveness of our approach across multiple domains. On the ModelNer40 3D shape generation task, DMP achieves a 16.15% average improvement in Wasserstein distance com- pared to baselines. We further validate our method on a simulated spatiotemporal transcriptomics dataset, where DMP often outperforms existing approaches across multiple biologically inspired generation tasks. Finally, show that state-of-the-an art image generative models can be easily mod- ified to incorporate DMP through minimal code changes, leading to significant improvements in FID with the same computational cost.
2.2. Flow-Based Generative Models
A common paradigm to model such distributions is through the unifying lens of flow-based generative models which en- compasses diffusion models, flow-matching, and other vari- ants (Liu et al., 2022: Neklyudov et al., 2023; Albergo et al., 2023). Suppose Z (X, R} and G. (X,, R, E) when
2

---

Noise- Coise-Conditioned Graph Networks
Z = X and G. = (X, R2,& E) when Z R. These models aim to learn the time-varying vector field dZz vo t)dt where E [0, which transports a prior distribution Pu(Z) (e.g. a Gaussian distribution) to the data distribu- tion P1 (Z) such that it satisfies the continuity equation
scribed in Section 2 to recover the unnoised vectors =
(8)
given the noised vectors 2/"
3.1. Theoretical Analysis
Up
de = (p)v)(t)).
To enable tractable analysis, introduce several technical assumptions (detailed in Appendix A). However, the results we prove continue to hold empirically even when these assumptions are relaxed, as we demonstrate in Section 3.2. For brevity, we defer all proofs to Appendix A.
We learn vo(Ge, t) with a graph neural network by re- gressing against the ground truth vector field ut (Ze) where Z ~ (Z, / (Zu 11), o.) with Zo (Z( drawn from the prior and Zi (Z) drawn from the data distribu- tion. Intuitively, u, represents the vector field that guides the noised sample Z, toward its unnoised target Z He and of are the mean and variance at noise level t. Different flow-based models specify different forms for μt, of, and Ue; we refer readers to Table of Tong et al. (2023) that details these different forms. In the rest of this work, we will refer to noise level and time l interchangeably (where close to 1 is low noise and t close to 0 is high noise).
1.1.1. NOISING OVER FEATURES
To recover the target (i.e. unnoised) features from noised features, we ideally want to maximize the mutual informa-
(3) and the accessible tion between the original feature z information to the graph neural network, which the aggre-
gation of noisy features. Here, we denote the aggregation of
noisy features as y((,)), which is a function of the radius of aggregation r around node at noise level
2.3. Fixed Graph Structure
To learn vi (G., 1) with graph neural network, previous works (Xu et al., 2022: Corso et al., 2022; Luo et al., 2021) construct a fixed (aka noise- poise-independent) graph structure from the positions. For instance, the message passing edges in Torsional Diffusion (Jing et al., 2022) are based on a constant-radius graph around node positions. In the case
0.-
Optimal Radlus
8.5
of positional generation, the message passing edges are redrawn for each t based on the noised positions; however,
the radius itself or h in K-NN are still a constant value across t. While one might consider using fully connected architectures like transformers as an alternative (Abramson et al., 2024), these approaches suffer from over-smoothing (Wu et al., 2024) and quadratic computational complexity.
62
10-1
100
32
322
10a
105
Signal-to-Nolse Ratio
Figure 2. Optimal Radius for Different Noise Levels. Using the correlation ((r)") (n(()) 1 (11() ,7(0) plot the radius that maximizes mutual information for each hignal-to-noise ratio using the formula from Lemma 3.1.
3. Gaussianized Geometric Graphs
Through theoretical and empirical analysis, we aim to show that the graph neural network structure should adapt as we change the amount of noise a geometric graph. To do so, we investigate the behavior of geometric graphs under varying levels of Gaussianization (aka the progressive move- ment from the data distribution graph's 's positions or features. Specifically, we look at two important components in graph neural networks under noise: the range of connectivity and resolution of nodes.
For our first theoretical result, we assume that the positions of the nodes are fixed and we only noise the features.
Lemma 3.1. Let the signal-to-noise ratio be SNR(() = of
Gaussian) either in the
where of is the initial variance and of is the variance at t, and assume SNR(t) > 0 for all t. Let pla (,(()) be the correlation between the features of node and j based their positions. Further, let D,ss) be the closed ball of radius around node based on its position (8), The mutual information [(x y((r)) between clean node feature xi and the aggregation y((tir) of noised features is
Preliminaries. Under the framework of flow-based gen- erative models, the node positions or features of a graph can be approximated an isotropic multivariate Gaussian (Ze, di))( We denote of as the data distribution's 's variance. Because the covariance matrix is diagonal, we can treat each node as independent univariate Gaussians p( p(=(")) (x(
2r
A
[(x] ((x))y((0n y((r))
SNR(T)
log
+ A B2
SNIT(E)
(0) of of). The target objective as de-
3

---

Noise-Conditioned Graph Networks
where
Proposition 3.3. Let 1111 - 11 = <<><<>>>- Given
variance-preserving noise and correlation structure between positions as p. the expected squared distance between node and j with respect to noise level Lis
A
p(n),), ,1 ydjdk,
J,KED .10)
B
Bi(mk (0) - - (1 = +)00)) ty?.
This lemma provides a formula for calculating the mutual information between the original signal of and the aggre- gated noisy signal y-(6.r) based only on the correlation p.
While this proposition offers more elementary treatment than the analysis of noising over features, it reinforces the notion that the radius must increase with the noise level.
If we assume () n7(5)) = 1 (1)0 Figure 2 visu- alizes the mutual information (calculated using Lemma 3.1). Clearly, increasing the signal-to-noise ratio causes the ra- dius with maximum mutual information to decrease, and decreasing the signal-to-noise ratio causes the radius that achieves maximum mutual information to increase. This relationship formalized with the following theorem:
3.2. Empirical Analysis
To validate that the theoretical analysis holds in practice without the simplifying assumptions, we conduct an empiri- cal demonstration on a dataset of 3D geometric objects (Fey & Lenssen, 2019) (spheres, cubes, triangular prisms, etc).
Increasing Noise - Increasing Connectivity. We train two separate Up networks two separate tasks: (1) generate the positions of the geometric objects and (2) generate the features of the geometric objects. Since this dataset does not come with node features, we construct features z; for each node i with position 17/ by computing the displacement vector from the node to the shape' 's center of mass p.: Ig ") / where μ = Σj= Mij
Theorem 3.2. Assume that
. where 0 < < 1. Let ri be the radius that maximizes mutual information for given signal-to-noie noise ratio ct- If we decrease the nigise ratio 02 < C1 then there exists 12 >> 11 such that
I(x1)) y(ma)) ()())))
The network Uy is fully connected single-layer graph atten- tion network. The goal of this task is to show what portions (and more importantly, what ranges) of the geometric object are relevant to the graph attention network at different levels of noise.
Since aggregation is simplified form of message passing, Theorem 3.2 suggests that message passing architectures should adapt their connectivity pattern with noise level. At high noise, nodes can increase information by communicat- ing with a broader neighborhood, while at low noise, local connectivity is optimal for maximizing information.
Figure shows the average (over the dataset) attention weight distribution across noise levels. At low noise, both position and feature generation models focus primarily on nearby nodes. As noise increases, both models shift atten- tion to more distant nodes, but in distinctive ways: position generation becomes sharply focused on distant nodes while ignoring nearby ones, whereas feature generation maintains a more uniform attention distribution across all distances.
Importantly, aggregation can also be viewed as form of pooling or coarse graining Both operations aggregate in- formation over a local neighborhood. Theorem 3.2 there- fore also provides theoretical justification for why reduc- ing graph resolution becomes advantageous at higher noise levels. By pooling nodes together, we perform a form of aggregation that helps denoise the graph while preserving essential structural information.
These empirical attention patterns support Theorem 3.2 and Proposition 3.3. As noise increases, the radius of effective communication must increase. However, this does not neces- sitate a fully connected graph, which would incur quadratic complexity and problems of oversmoothing. Instead, our theory suggests we can achieve optimal performance by dy- namically structuring connectivity based on the noise level, connecting only nodes within the relevant radius.
3.1.2. NOISING OVER POSITIONS
Another common case is learning to reverse the noise over positions. However, we cannot follow the same proof struc- ture as the previous case, as the aggregation of nodes is now also dependent on the noise level rather than just the radius of aggregation; as data becomes more Gaussian, the num- ber of nodes in some radius of node i is dependent the positional center of node i. Instead, Proposition 3.3 proves that as we increase the noise level, the expected distance between any close pair of nodes increases. Therefore, the radius of message passing must increase, so that originally nearby nodes can communicate.
Increasing Noise - Decreasing Resolution. To under- stand how noise affects the optimal graph resolution, we conduct systematic analysis of coarse-graining effects. For each geometric shape in our dataset, we (1) generate ver- sions with different noise levels, (2) create coarse-graing representations using K- K-means clustering with varying num-

---

Noise-Condition Graph Networks
Cenerative Nodeling Positions
Generative Hodeinng Fentures
Pooling Positions
Paciigg
Highar
Optimal Number of Clusters
Nederata
Pooling Over Features
foolin
Optimal Number Clusters
Halue
Clovatt
Intermediate
Partheat
Closest
Intermerd collree
Porthest
Figure 3. Attention Distribution over Node Distances. Average attention weight distribution between pairs of nodes function of their spatial distance, shown for different noise levels in a graph attention network.
x
(1
bers of clusters, and then (3) compare each noised and coarse-grained version its original, unnoised shape.
Figure 4. Optimal Graph Resolution Different Noise Lev- els. For each noise level, we plot the number of clusters that minimizes
distance hetween the coarse- grained noisy graph and the original graph. Results plotted separately for max and mean pooling operations.
To quantify the similarity between shapes, adopt the Gromov-Wasserster (GW) distance from object registration literature (Mémoli. 201 1). Given two point clouds, the GW distance measures how well their geometric structures match while being invariant to rigid transformations (i.e. a lower GW distance signifies better match).
Our key insight from Section 3 is that GNN should adapt to the noise level t. This leads us to propose Noise- Conditioned Graph Networks (NCGNs):
Figure 4 shows how the optimal level of goarse-graining varies with noise level, where optimality is defined min- imizing the Gromov-Wasserstei Wasserstein distance to the original shape. It is clear that as noise increases, the optimal num- ber of clusters decreases. This pattern indicates that higher levels of coarse- nouse-graining (fewer clusters) better preserve the underlying geometric structure when noise is high. This relationship holds across different pooling operations (max and mean) and across noising the node features positions.
(G,, t) NCGN,(G.).
The core difference here is that the previously static network GNN now conditioned on the noise level as NCGN, For example, NCGN, may different values of \ or may perform message passing over the hi nearest neighbors, where k, is function of I. While NCGN, is intentionally general, in the following section, we develop specific instantiation of NCGNs that dynamically changes its message passing connectivity and resolution according the noise level of the generative process.
different architecture for
4. loist-Conditionee Graph Networks
Let to G x 10 > Z be the graph neural network used to predict the vector field dZ = u (GG dt. Previous work uses fixed network independent of noise level as described in Section 2.3:
4.1. Dynamic Range Resolution Message Passing
We introduce an instance of NCGN called Dynamic Mes- sage Passing (DMP). DMP specifically modifies the mes- sage passing connectivity (i.e. the edges in the message passing graph) and the message passing resolution (how information is propagated through intermediary nodes). For the latter, DMP introduces coarse-graine' "supernodes" that serve only to relay messages between different parts of the graph during the message passing phase. Importantly, the network output no remains defined over all original nodes;
GG t) GNN(G.)
Here, GNN is an abstract message passing network with flexible design choices like structural elements (edge con- nectivity) and algorithmic components (message passing functions, aggregation schemes, and node update rules).
5

---

Coise-Conditioned Graph Networks
the coarse-graining only affects how messages are propa-
Algorithm Forward Pass of DMP
gated internally.
1: Input: G2, f, To, T1, S0, $1
Let ry be the range (e.g. kNN connectivity or radius graph)
Output: Vig
(pred)
and S, the number of supernodes used for message passing
3: rt, 8t ro,
at noise level t. Further, suppose we are given boundary
4: X, R lift((X, RR) Ri
conditions (ro. 5o) and (ri, $1) that specify the connectivity and NCGN, solely based on rt and 8t. To determine rt and St, use adaptive scheduler,
5: for k = 1o K GNN layers do
and resolution
respectively. DMP defines
6:
X. F.R.C coarsen X. R, E)
7:
Et concat(E, constructEdges/ R, r))
X, ft +- 3MP+ (X, R X, uncoarseng (X, R, X,
Tty St [(,, To, 80,
10: end for (pred)
11: 1/8
proj X, R)
Building on Section 3, we require the scheduler / to satisfy the following constraints:
(Monotonicity) For < (where represents higher noise), we must have r' 2 and s' S where r',s' [('', Fo: r1, 50, 81). This follows from Section 3, which shows that increasing noise requires greater radius for information aggregation and reduced graph resolution through boarser-grained nodes.
Uncoarsening Step. Finally, since vp operates on the full graph resolution, DMP maps the processed information back to the original nodes:
X, R uncoarsen X, R, X,
This operation can be the original nodes aggregating from their associated coarse-grainee nodes. uncoarsen could also be another learnable message passing layer where the coarse- grained node message passes to the original nodes.
(Boundary Condition Satisfaction) T1, 81 [(1, ro, T1, 50, 81) and ro, 50 = f(O, ro, F1, 80, 80,8)); 81).
Since deriving an optimal scheduler f is unclear, we rely on a predefined f that satisfies the above conditions. Some examples of f are shown in Figure 5. Section 5 provides guidance on which schedule performs best in practice.
Minimal Implementation Effort. Converting an existing GNN architecture DMP requires minimal modification. The primary changes involve implementing the coarsening and uncoarsening operations, while the core GNN architec- ture remains unchanged, operating on the coarse-grained representation. This modification also preserves all other aspects of the flow-based generative model. We show this in practice by modifying an existing state-of-the-ard generative model to DMP in Section 5.3.
Graph Construction. With specific r; and & given by f. DMP implements NCGN, with the following steps:
Graph Coarsening. DMP first maps the noisy graph G2 to coarse representation based on 8, as
8, Ri, & coarsen(s E,
Linear Time Complexity. DMP can achieve linear time complexity while maintaining expressiveness through the choice of boundary conditions and the constructEdges oper- ation. Consider a setup where at no noise (t = 1) we have full resolution (s1 = nodes) with sparse connectivity (r1 neighbors per node). while high noise (t = we have a coarse resolution (sg vr N nodes) with full connectiv- ity. When f interpolates between these boundaries such that TtSt r N (where Tr is the number of neighbors and st is the number of coarse- granine nodes at time t), message passing maintains linear time complexity throughout the generation process.
where [R] | = For instance, coarsen can be deterministic mean pooling or a learnable message passing layer from original to coarse-graines nodes. This coarse rep- resentation serves exclusively as an intermediate structure for efficient message passing, while uy maintains its output mapping the complete set of original nodes.
Edge Structure. Given coarse-grained edges E, DMP de- fines additional message passing edges between coarse- grained nodes according to the connectivity range r;:
غ = Co (nnstruuttddges((((r))
where constructEdges implements a proximity- -ased con- nectivity pattern such as a kNN or radius graph.
5. Experiments
We evaluate the DMP instantiation of NCGNs across three settings: generating geometric graphs over positions, gener- ating geometric graphs over features, and adapting existing generative models.
Message Passing Step. Rather than propagating messages between original nodes, DMP performs message passing over the gaaseegrained representation:
X MP E)
Setup. For the first two experiments, we implement DMP
6

---

Nolse-Conditoneed Graph Networks
NODEL
W2 (x10
M0DELNET40
RANDOM
KNN
LONG-SHOR'S RANGE
DMP
LINEAR
4.495
EXPONENTIAL
4.263
GCN
8.624
5.882
4.315
4.215
GAT
8.624
5.598
4.741
4.263
LOGARITHM
4.860
RELU
4.485
Table Wasserstein distance (x10-3 )) between the predicted distribution and the ground distribution of the ModelNer40 test set (the lower the better).
Table 2. Comparing different adaptive schedules for DMP. Each result is reported on flow- matchinthh models with GAT as Us.
in the same way. We have 3 message passing layers where we evaluate two popular message passing architectures, con- volutions (GCNs) (Kipf & Welling, 2016) and attention (GAT') (Veličković et al., 2017), to show that in either case, our method can still perform well. The coarsening operation a mean over the nodes in the cluster where the clusters are determined by voxel clustering. Our boundary conditions and edge construction match that of "Linear Time Complex- ity" in Section 4. As such, the DMP implementations in these two experiments have linear time complexity.
It is particularly important to test against long-short range because we aim to show that the performance benefit of DMP does not come from having both long and short range message passing, but rather that at each noise level, it is beneficial to only look at a specific neighborhood range. Additionally, DMP uses the exponential function as the adaptive schedule for this experiment, However, we also investigate the performance of other adaptive schedules in Section 5.1.1
Results. We evaluate DMP and the baselines using the Wasserstein (W2) distance between the test set distribution and the generated distribution. The results are shown in Table 1. Across both GCN and GAT message passing ar- chitectures, we see a consistent decrease of 16.15% in W2 distance on average by simply having noise-conditioned graph network rather than the fixed form of baselines.
We use DMP within two classes of flow-based generation models: diffusion models (Ho et al., 2020) in the first exper- iment and flow-matching models (Tong et al., 2023) in the second experiment. For diffusion models, we use 1000 num- ber of function evaluations (NFEs), and for flow-matching models, we use 200 NFEs. All samples (both positions and features) are normalized to [-0.5, 0.5). The rest of our hyperparameters are recorded in Appendix B.2 and fixed between both experiments; further experiment details in- cluding the DMP architecture can be found in Appendix B.
5.1.1. COMPARING ADAPTIVE SCHEDULES
From Section 4.1, it is clear that many functional forms sat- isfy the constraints for the adaptive schedule. To investigate the impact of different adaptive schedules, we compare the performance between the schedules shown in Figure on ModelNet40 using GAT the message passing architecture and diffusion as the generation model. Table 2 displays the results for each schedule.
For the third experiment, we demonstrate how exist- ing flow- generative model can be easily adapted to incorporate noise-conditioned graph networks by modify- ing state-of-the-art image generation model: Diffusion Transformer (DiT (Peebles& & Xie, 2023)). We denote the resulting modified architecture as DiT-DMP.
As shown in the table, while there is variability in perfor- mance between schedules, all schedules except the loga- rithm schedule still perform better than the baselines. Based on this experiment, since the exponential schedule per- formed best, will continue using only the exponential schedule for all following experiments.
5.1. Generating 3D Objects
Dataset. DMP is evaluated on ModelNet40 (Wu et al., 2015), a dataset of 3D CAD models for everyday objects (e.g. airplanes, beds, benches, etc) with 40 classes in total. From each object, we sample 400 points and treat the points as a point cloud. The training and testing set consists of 9843 and 2232 objects, respectively. On this dataset, we aim to learn the distribution over positions of the point clouds.
5.2. Generating Spatiotemporal Transcriptomies
Dataset. We evaluate spatiotemporal transcriptomies dataset simulating gene expression patterns across both space and time in developing limbs. This is an important task with applications in fundamental biological discov- ery (Tong et al., 2020; Qiu et al., 2022) and therapeutics (Bunne et al., 2023). Each node in the graph represents a cell with features corresponding to the expression levels of three genes (Bmp, Sox9, and Wnt), while node positions encode both temporal progression and spatial location in the
Baselines. We evaluate DMP against two baselines: (1) a full resolution (i.e. no coarse-graine (ouss--raiinig) point cloud with message passing edges constructed based on a kNN neigh- borhood and (2) full resolution point cloud with message passing edges constructed with both short and long-range edges (denoted as "long-short range"). For each node in the long-short range baseline, we uniformly sample & edges.
7

---

Noise- Coise-Conditioned Graph Networks
W2
RANDOM
KNN
FULLY CONNECTED
DMP
UNCONDITIONAL GENERATION
2.452
2.030/0.840
0.880/1.094
0.812/0.834
TEMPORAL TRAJECTORY PREDICTION
2.681
2.290/1.022
1.060/1.271
0.960/1 0.02
TEMPORAL INTERPOLATION
3.982
3.390/2.946
2.923/3.026
2.880/2.912
GENE IMPUTATION
2.810
2.943/2.527
2.507/2.534
2.529/2.575
SPACE IMPUTATION
3.344
3.222/2.928
2.923/2.946
2.907/2.913
GENE KNOCKOUT
3.202
2.890/2.390
2.424/2.473
2.392/2.397
Table Wasserstein distance between predicted distribution and ground distribution the spatial transcriptomic dataset (GCN performance/ / GAT performance).
MMAGNET 256 256
DIT
DIT-DMP
representations of dimensions 16 x 16, 8 x 8, or 4 x 4 respec- tively. transformer then learns to reverse the noise process. We modify DiT's architecture to implement DMP through two simple changes: (1) we adopt Flexi (Beyer et al., 2023) dynamic patch sizing to increase patch sizes higher noise levels, and (2) we replace the transformer's quadratic attention mechanism with neighborhood attention (Hassani et al., 2023) where neighborhood size scales with noise level. Detailed implementation specifications are provided in Appendix C.
FID SFID
84.051
63.983
18.787
12.823
1St PRECISIONT
16.735
24.681
0.296
0.446
RECALL?
0.394
0.426
Table 4. Performance between DiT and its modified DMP version.
tissue. Following Raspopovic et al. (2014), we simulate the data using a saaction-diffusion system that captures the key gene regulatory dynamics in limb formation. The training set contains 10000 samples with 100 nodes each, while the test set uses a different discretization with 96 nodes to eval- uate model robustness to varying resolutions. Detailed data generation procedures and biological context are provided in Appendix B.3.
Dataset. Evaluation conducted on the ImageNet 256 256 dataset (Deng al., 2009), composed of 1281167 train- ing images and 100000 test images across 1000 object classes.
Baselines. We evaluate baseline DiT using patch size p 4 and the Small model configuration (Peebles Xie, 2023) with 12 layers, hidden size of 384, and attention heads. Our modified DiT-DMP has the same computational com- plexity as the baseline.
Baselines. We compare DMP to the following baselines; (1) a KNN neighborhood graph and (2) a fully connected graph on important biologically inspired tasks. Note that unlike Section 5.1, we only have 100 nodes for each graph, so it is computationally feasible to baseline against fully connected graph
Results. As shown in Table 4, the modified DIT-DMP sig- nificantly outperforms the state-of-the-art DiT model across various metrics on unconditional generation. Notably, these gains are achieved while maintaining identical computa- tional complexity the baseline and requiring only two straightforward code changes: replacing the fixed patch size with FlexiViT's dynamic patch sizing and swapping the transformer's 's full attention for neighborhood attention.
Because we found from the previous experiment that the exponential adaptive schedule works best, we proceed with only using the exponential schedule. To showcase that DMP compatible with other flow-based generative models out- side of diffusion models, we will use the more recent formu- lation of flow-matching in this experiment.
6. Conclusion
Results. Summarized in Table 3, we see that DMP has lower W22 distance on GCN and comparable performance on GAT against baselines. Especially for specific tasks like gene imputation or gene knockout, there is not signifi- cant difference in performance between our method and baselines.
We introduced Noise-Conditioned Graph Networks, demon- strating that adapting graph neural network architectures to noise levels improve performance for geometric gen- erative modeling. We showed that optimal information ag- gregation requires increasing radius with noise level, while graph resolution can be reduced without losing much infor- mation. However, DMP is currently limited to predefined range and resolution schedule. Future works could intro- duce learnable range and resolution parameters. Further, several other components in graph neural networks could be adapted outside of range and resolution including the number of layers, message passing type, or width of the
5.3. Adopting an Existing Generative Model
For a given noised image in the generative process, DiT first encodes the image 32 x 32 latent representation using an autoencoder. Then, DiT partitions this representation into patches of size p {2, 4),, yielding coarse bours--grained
8

---

Noise-Conditioned Graph Networks
layers.
Corso, G., Stärk, H., Jing, B., Barzilay, R., and Jaakkola, T. Diffdock: Diffusion steps, twists, and turns for molecular docking, arXiv preprint arXiv:2210.01776 2022.
Impact Statement
Deng. J., Dong, W., Socher, R., Li, L.-J.,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248 258, Icee, 2009.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must specifically highlighted here
Fey, and Lenssen, J. F. Fast graph representation learning with pytorch geometric. arXiv preprint arXiv:1903. 19230022,, 2019.
Acknowledgements
We thank Francesco Di Giovanni, Wenzhuo Tang, Yifan Lu, and anonymous ICML reviewers for helpful feedback and discussions. PPH is supported by the National Sci- ence Foundation GRFP and acknowledges the support by the National Institutes of Health DP2 New Innovator Award (IDP2HG014282-001 and MorPhiC consortia grant (5U0LHG013116-002, MB is supported by NSF grants CCF- 2217058 and CCF-2310411.
Gómez-Bonne Bombarelli R., Wei, J. N., Duvenaud, D., Hernández-Lobato, M., Sánchez-Lengelee berla, D., eguilerr-Iparraguinn
B., She- J., Hirzel, T. D., Adams, P. and Aspuru-Guzik, A. Automatic chemical de- sign using a data-driven continuous representation of molecules. ACS central science, 4(2):268-2) 276, 2018.
Hassani, A., Walton, S., Li, Li, S., and Shi, H. Neigh- borhood attention transformer. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 6185-6194, 2023.
References
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T.,
Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J., et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-663 6851, 2020.
Jing. B., Corso, G., Chang. J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35: 2022.
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.0879) 2023.
Alon, U. and Yahav, E. On the bottleneck of graph neural networks and its practical implications. In In- ternational Conference on Learning Representations, 2021. URL
Jing, B., Erives, E., Pao-Huang P., Corso, G., Berger, B., and Jaakkola, Eigenfold: Generative protein struc- ture prediction with diffusion models. arXiv preprint arXiv:2304.02 /98, 2023.
wwt.arevvee.nne//eeeeee net/forun?
id=i 800hhCCV
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneherger, O., Tunyasuvunakool, K., Bates, R., 7idek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873) 2021.
Bengio, E., Jain, M., Korablyov. M.. Precup, D., and Ben- gio, Y. Flow network based generative models for non- iterative diverse candidate generation. Advances in Neu- ral Information Processing Systems, 34:27381-27394, 2021.
Kipf, T. N. and Welling, M. Semi-supervised classifica- tion with graph convolutional networks. arXiv preprint arXiv: 1e0%iv1609.000 wet.io.09900000 2016.
Beyer, I.,, Izmailov, P., Kolesnikov, A., Caron, M., Korn- blith, S., Zhai, X., Minderer, M., Tschannen, M., Alab- dulmohsin, I., and Pavetic, Flexivit: One model for all patch sizes. In Proceedings of the IEEE/CVF Confer- ence Computer Vision and Pattern Recognition, pp. 14496-14506, 2023.
Li, Z., Kovachki, N., Azizzzaeneshelli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03456 0038000000355... 2020.
Bunne, C., Stark, S. G., Gut, G., Del Castillo, J. S., Levesque, M. Lehmann, K.-V., Pelkmans, L., Krause, A., and Rätsch, G. Learning single-cell perturbation sponses using neural optimal transport. Nature methods, 2023.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint 27.00722 2022.
9

---

Noise-Conditinee Graph Networks
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint ar arXiv:2209.03000 2022.
Raspopovic, J., Marcon, L.,, Russo, L.., and Sharpe. J. Digit patterning is controlled by bmp-sox9-wnt turing net- work modulated by morphogen gradients. Science, 345 (6 (6196):566-) 570, 2014.
Luo, S., Shi, C., Xu, M., and Tang. J. Predicting molecular conformation via dynamic graph score matching. Ad- vances in Neural Information Processing Systems, 34: 19784-19795, 2021.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Luo, S., Su, Y., Peng. X., Wang. S., Peng, J., and Ma, Antigen-specific antibody design and optimization with diffusion-based generative models for protein structures. Advances in Neural Information Processing Systems, 35: 9754-9761 2022.
Tong, A., Huang. J., Wolf, G., Van Dijk, D., and Krish- naswamy, S. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference machine learning, pp. 9526-9536. PMLR, 2020.
Mémoli, Gromov-wassersiee distances and the metric approach to object matching. Foundations of computa- tional mathematics, 11:417-453 2011.
Tong. A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with mini- batch optimal transport. arXiv preprint --Xiv:2332.00445 2023.
Nash, C., Ganin, Y., Eslami, A., and Battaglia, P. Polygen: An autoregressive generative model of 3d meshes. In International conference on machine learning, pp. 7220- 7229. PMLR, 2020.
Veličković, Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y. Graph attention networks. arXiv preprint arXiv: 10903, 2017.
Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. Action matching: Learning stochastic dynamics from samples. In International conference on machine learn- ing. pp. 25858 35858-25889. PMLR, 2023.
Watson, J. L., Juergens, D., Bennett, N. R., Trippe, B. L., Yim, J., Eisenach, H. E., Ahern, W., Borst, A. J., Ragotte, R. J., Milles, L. F., et al. De novo design of protein struc- ture and function with rfdiffusion. Nature, 620(7976); 1089-1100, 2023.
Oono, K. and Suzuki, T. Graph neural networks exponen- tially lose expressive power for node classification. In International Conference on Learning Representations, 2020. URL https: http://opeeeeeeeeeeeeeee forum? id=S1ld02EFPr.
Wu, X., Ajorlou, A., Wu, Z., and Jadbabaie, A. Demys- tifying oversmoothing in attention-basel graph neural networks. Advances in Neural Information Processing Systems, 36, 2024.
Pao-Huang, P., Jing, B., and Berger, Scalable multimer structure prediction using diffusion models. In NeurIPS 2023 Al for Science Workshop, 2023.
Wu, Z., Song, S., Khosla, A., Yu, F,, Zhang, L., Tang, X., and Xiao, J. 3d shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1912- 1920, 2015.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision, pp. 4195-4205, 2023.
Xu, M., Yu, L., Song, Y., Shi, C., Ermon, S., and Tang, J. Geodiff: A geometric diffusion model for molecular con- formation generation. arXiv preprint axXiv:2003.09923, 2022.
Peng, Q., Zhou, P., and Li, T. stver: Reconstructing spatio- temporal dynamics of cell development using optimal transport. bioRxiv, pp. 2024-06, 2024.
Yang, G., Huang, X., Hao, Z., Liu, M.-Y., Belongie, S., and Hariharan, B. Pointflow: 3d point cloud generation with continuous normalizing flows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4541-4550, 2019.
Qiu, X., Zhang, Y., Murtin-Rufino, J. D., Weng, C., Hos- seinzadeh, S., Yang, D., Pogson, A. N., Hein, M. Y. Min, K. H. J., Wang, L., et al. Mapping transcriptomic vector fields of single cells. Cell, 185(4))690-711. 2022.
Qiu, X., Zhu, D. Y., Lu, Y., Yao, J.. Jing, Z., Min, K. H., Cheng, M., Pan, H., Zuo, L.,, King, S., et al. Spatiotem- poral modeling molecular holograms. Cell, 2024.
10

---

Noise-Conditioned Graph Networks
A. Proofs
A.1. Noising over Features
Assumptions. We make the following assumptions throughout the analysis noising over features:
The noising process defined as 2, = 21 + ex where er NN of)
iii) Node positions and features are supported on the real line: x ,7) E R.
17())
(iii) We treat the graph in continuous manner rather than discrete object, meaning that we have corresponding node (x()), () for all i E I, with the position v7 = i.
(iv)
The message passing scheme is summation over all node features in some closed ball D,60 of radius around node i based on its position 17 This aggregated sum of features y-(,.) SJED, .(1) the information accessible to node at noise level L.
it (5) dj
SJED 20) x4 + dj represents
Lemma 3.1. Let the signal-to-noise ratio be SNR(()= = नें where of is the initial variance and of is the variance at t, and assume SNR(t) > 0 for all t. Let pl m((0) be the correlation between the features of node and j based on their positions. Further, let D, be the closed ball of radius r around node i based on its position 17 The mutual information (x) (8) y((.)) between clean node feature ( and the aggregation Y, (i.r) of noised features is
[( y((r))
1
SNR(C) SNR(()
2r
+ A
log
2
2r
+ A B2
where
A=
p(n)1), (7 (djdkk
JikED .10
B
10 (1
JED
Proof. We aim to calculate
1( , ((())))))))))))))) =((x(")))))))))) + H(x((t))) - Hx()), yr y((r))).
The entropy of Gaussians are known in closed form as
(x) =
2 log
2TTeO²
x1
((x)( = log 2TTeO p((nm)
2
where
oi 2ye0 = var(y var
=g'dj]
-
/
cov(x 12 20 Jajak.
,.NED.
Following variance-explee variance-exploding noise,
f J.MED, .0
cor(z + Et 21 + + )ajak
(1) (1) (8)
where e, = N( of)
= 1
opp(((),
+ S of
-2ro² + +
10 11 )i)k
j.kED,,))

---

Coise-Conditioned Graph Networks
and
2 of
Further, we have
H(x y((r)) = log (22Te)
x1 (0
299 p(a))))))
2
with
(1()
2² 1 y(hi)) = cov
( + dj dj =
cov(2] 2 = of
,
41
20D,00
JED,00
jED,(0
Consequently,
[<x]) (6) Y, y((r))) log
1
of p((
2
"1 y(u
o
y(v.r) 2
of (2ro? + of J, .n) p(n()), )djk)
1
log
2
of + of J,,ED,,co .0)) p())()) 1y [djck) - (oî SJED, (1)11), ,()))))
2ro + of JIAED,I, .1) p(y(()), (t)
1
log
2
2ro of JI,ED, pn) n djjkk of dg p(nn(),))) nnf)
Factoring out of.
2r of
SJIAED, pn 17 )djkk
= log
2
ct + 2r
Jdjck - USED,,. (11 nD)))jj)² JIKEDD p(1(6), (djck
Define the ignal-to-noiss ratio SNR(t)= di
Then,
2r SNR())
+ JAKED p(n(1), (1) ddjkk
10 y((rr))=
1
log
JIKEDD,O ,m p(n(1), 1) djjck UJED,,o p((((),, sdf)2
SNR(C)
Lemma A.1. Assume that the correlation follows the following form where 0 <y<rsı,
(0 (7)()) = 1 -0 1/92
Then,
T(x11) (1) (0,(),,(()))) y((ir))
1 log( log((s(r,c))
2+ and SNR(t).
where k(r,c) c)
2+
Proof.
First, solve for
A
[ J.MED,
p) (k) [djkk J.MEDD,co
1- () - (k()) 4r2
8r
and
B = J,DD
/
1- (n) = 2r
3,3
3
12

---

Since des(nee) is continuous over [e, 1]. switch of signs between the endpoints shows that there exists an interior point dr where de(()es)) Cnise-oonditooee Graph Networks
Based Lemma 3.1, have
2r SNI(C)
+
A
[(x] (0)
y((r))
log
+ A
B
SNR(E
+ 4n
8
log
3
4p
0
9
2
+
8
log
3
4
9
where SNR(f).
Lemma A.2. Let k(r, c) and be defined as in Lemma A. Assume we restrict the support of to be (c, 1] where is sufficiently small. For a given signal- toniis ratio e at time t (denoted as cch. there exists a r* such that
max K(T,, c) and dx(r" .c) esrsl
dr*
Proof. By the extreme value theorem, since the domain is compact [c. 1)) and continuous, I must achieve maximum within this interval. Further, need show that the maximum is achieved interior point r (c, 1) to prove that dx(r".e.) O. dr
Taking the derivative of from Lemma A.I,
dx
18c (4cr
10cr 6r 18r + (2cr + 9)2
dr
du(r.ee) dr
For where e is sufficiently small,
0 because the numerator is positive and the denominator is positive:
lim 4cr 8 10cr 6r 18r2 + 9 9 0.
c++++
However, when r = 1. the numerator becomes negative and as such dx(r.c+) dr
< 0:
18c (4c 10c - 6 18 + 9) 18c (- 6c 15) 0.
dr 0.
Lemma A.3. Let rı and o₁ be the values guaranteed to exist by Lemma A.2 such that dx(rr.e)) 0. For 02 C1. have dri
dK(r1, 02) > 0.
dri
Proof. Since 7 is the radius that maximizes mutual information (which we know exist based on Lemma A.1), know that x(ri,ct) x(r,ei) A < < 1. Further, with
dx 18c (4er 10cr 6, - 18, + 9)
dr
(2cr6 9)
where de(r),01) dri
0, we focus on only the numerator:
4crri 10c 6 Gr,
2 + 18r,
4crr 10crr1 = 6r 5 + 18r 9
13

---

Noise-Conditioned Graph Networks
Suppose that 72 ri + ò. Since we know that de(r1-01) is maximum, we need show that dx(r) e2) Since the
dri
dr
denominator of dis is strictly positive, we only need to show when the numerator is always positive dr1
= 1802 (402r1 1002r1 - 6r1 - -1811 +
Substituting our equality from earlier,
= 18cy (4c211 - 1002r1 4cir1 &- - 1000r11
18c2r((02 (1) -00).
Since 02 O, need when 4r, 10 < 0 to maintain that the entire numerator is positive. This equality is true when
which is always true since 0 < As such, we have established that de(r).ca) 0.
dri
Theorem 3.2. Assume that ( = 1- - (1700) - 122)) where 0 < <rs 1. Let vı be the radius that maximizes mutual information for a given signal-to-noice ratio 61. If we decrease the signal-to-noise ratio 02 C1, then there exists 12 > Ti such that
[[ ( y(rr)) [(x1 y(m)).
Proof. Because log is strictly increasing, any maximizer for is also maximizer for I.
As such, with Lemma A.3, we know that du(r),0)) > 0. Therefore, r2 = + would increase mutual information, proving that we need to increase r beyond ri to reach the maximum mutual information at SNR level c2.
drj
A.2. Noising over Positions
Proposition 3.3. Let [|m]" - 1|2 7 << 1. Given variance-preseerinng noise and correlation structure between positions p. the expected squared distance between node i and j with respect to noise level t is
Bi(
The (1) - 2 t
Proof. Let v7(h 0 and n =y without loss of generality. By definition of variance- preseevvn noise, we have
T) (
N00 - t), Th ~N
We then calculate for var(m)
[[ (
1/: (1)))1] var(n.") ()) E(m! m) 1012
where
) var(m) (0) + var(n)) + 2 cov (m ())))) = 2(1 - 1) t)p(r),
which results
E [[ (t)
( (
2(1 t) 2(1 2(1-)))(7) + +2
B. Experiment & Details
B.1. Architecture
B.1.1. DYNAMIC MESSAGE PASSING
The model architecture used in the two experiment follows the same operations listed in Algorithm 1. The sole difference between the implementations of the models in the two experiments is the input and output node values. The remaining operations are identical between both experiments. We describe each operation in detail below.
14

---

Coise-Conditioned Graph Networks
KNN Connectivity Number of Clusters
Logarithm Schedule
Exponential Schedule
Linear Schedule
ReLU Schedule
2000
800
800
800
800
.000
600
600
600
200
200
200
200
0.0
0.
04
0.6
0.8
1.0
0.0
0.
04
0.6
0.8 1.0
02
0.4
0.5
0.0
1.0
0.0
0.2
0.4
0.0 1.0
Noise Level (1
Noise Level
Noise Level
Noise Level (1
Figure 5. Adaptive Schedules. The schedule determines the level of coarse-graining and connectivity of the graph neural network corresponding to a noise level. (Top) The number of clusters (aka fewer clusters means higher coarse-grate sause-graining). The higher the noise, the fewer number of clusters. (Bottom) The level connectivity (i.e. range) determined by the in kNN. The higher the noise, the greater the k.
Node Input Construction. We define two sets of nodes: (I) the original nodes and (2) the roarse-grained nodes. Let Hm = (1 E ]= be the original node features to the model, where N is the number of nodes and d is the dimension of the input data. Each coarse-grained node feature and position is defined as the average of the original nodes belonging to
(1) )i=1 and ft {i)() be the coarse mnars--griined node features and positions, respectively, where is the number of clusters determined by adaptive schedule (Figure 5) at the current noise level. We determine which original node belongs to which coarse coarse-grained grained node using voxel clustering. For each dimension j of the node positions, we uniformly split the interval (with range [min( ma)(R;)) into v/s partitions, which determines the boundaries each cluster the specific dimension j. the position of an original node is within the boundaries of cluster (for all dimensions), belongs to that cluster.
the specific coarse-grained node. Specifically, let H = (
For ModelNet40, the initial node input im = () = 0, 17 .) where R3 and t E [0, 1]. Since ModelNet40 only has positions and no features of each node, set to be no-op; hence, we have dimensions of d for node feature n!
The spatiotemporal transcriptomic dataset follows similar form, except now each node (i.e. cell) has both positional and feature information (where features are the gene expressions). Thus, the initial input (as described in Section 5.2) is ( co) (0()) f) where x E R since are tracking 3 genes, 1) R2 since concatenate the 1-dimensional temporal axis (progression of cell through time) and the 1-dimension spatial axis (position of cell across a line), and E [0, 1] is the noising time of the generative model (not to be confused with physical time of cells). Therefore, we have dimensions of d for node feature (1)
Edge Construction. Given the current noise level t, we leverage the adaptive schedule (Figure 5) to provide the ANN connectivity between the coarse-grained nodes. Specifically, we define our edges
E = KNN-Graph((it)
Lifting Operation. Before message passing, we first lift or embed each initial node input (both original and coarse- Searse-grained nodes) with
HR= MLP [d, hdim, hdim, hdim)]
where hdim is the hidden dimensions and the original and coarse-graines nodes have separately parameterized MLPs. Note that k = 0 refers to the node features before the first message passing layer and is not related to the ANN connectivity.
Coarsening. For each layer in our model (the layer index is denoted as k), we define the message-passingg scheme from the original nodes to the gaasee-rained nodes as the following:
Σ
MLPj [iinears linearx .100) ())))) (i)ear
-)))), linearz(eis)) [3hdim, hdim, hdim)))
JENamm (2)
15

---

Noise-Conditioned Graph Networks
where esy ||/ n/0 2 We conduct this operation for all original nodes j in the cluster i (denoted as Neluster All linear layers map Rhddi..
Message Passing Layer. After the coarsening operation, message passing is conducted between the coarse bears--rrained nodes. For each message passing layer MPk, we have the following
H MP MPp(H
where MPs E (GCNConv,GATCom)
For every experiment, we evaluate both GCNConv and GATConv as the message passing schemes demonstrate that the performance benefit of DMP generalizes more than just one type of message passing. We define GCNConv and GATConv following (Fey & Lenssen, 2019) for each node i as
GCNConv(h ( g0)) )-
= linear
Σ
GATConv(h, (i) ,,()) ((()) = as, le,,linear, (h + Σ
aq.linearr (b())
S
VEEO
JEE:
where g() give the nodes incident to node i. Further, 0i,j is calculated as
exp(LeakyReL.U(a] linear(h.)) a. (
ai,) Eregto exp(LeakyRel.U(a) linear (1))) + a linear (2))))
where and a, are learnable source and target attention vectors.
Uncoarsening. Similar to the coarsening operation, we now message pass from the coarse- sairneegrained nodes back to the original nodes since we ultimately want to predict the velocity field with respect to the original nodes:
spread MLLP (( (linear, (concat(h? (0) ()))). linear, (17 n)0)), linear, (e:g) }: [3hdim, hdim, hdim))
where Cij = ||>] n ll and is the coarse-grained node that the original node belongs to.
Then, we use a gating mechanism to control the flow of information from the coarse-grained nodes to the original nodes:
H6+1 MAH& + (1 A)spread;; [hdim, hdim, hdim))), where linear(concat( H (/read)))
Project Operation. Finally, after conducting the past operations over k K layers, the node features Hx+11 are projected back to the ambient space:
(pred)
Hout = MLP( Hı; [hdim, hdim, hdim, odim])
where odim is the dimension of the output. Depending the task (i.e. whether features), odim can either be the dimensions of the positions or the dimensions of the features. Specifically, for ModelNet40, since we are generating positions, odim = 3 (zyz coordinates), and for the spatial transcriptomicss dataset, since we are generating features, odim (number of genes).
generating positions generating
1.2. BASELINES
To ensure a fair comparison between baseline methods and DMP, we use the same architecture from Appendix B.I.I, except the message passing edges are fixed (kNN, fully connected, long-short range) and the graph always at full resolution. Specifically, the baselines consider the original nodes the coarse-grained nodes ("full-resolution") and conduct the coarsening and uncoarsening operations towards these "full-resolution" nodes (i.e. one-to-one mapping rather than DMP's many-to-one and one-to-many mappings). This consistency between baselines and DMP is important since it shows that the performance improvement purely comes from the adaptive range and resolution of the graph rather than architectural improvements.
B.2. Hyperparameters
We use the hyperparameter values listed Table 5 for both experiments. Many of these values are borrowed from the default values in Tong et al. (2023). Additionally, it was shown in previous work (Tong et al., 2023) that flow motching requires fewer NFEs than diffusion- based models, which is why we use more NFEs for diffusion compared matching.
16

---

Coise-Conditioned Graph Networks
Parameter
Value
Description
Activation
GELU
Neural network activation function
Normalization
Batch Norm
Layer normalization type
Message Passing Layers
3
Number of message passing layers
Hidden Dimensions
64/32
Dimension of hidden features (Exp. 1/Exp. 2)
Training Epochs
300
Total number of training epochs
Learning Rate
10001/0.000
Initial learning rate (Exp. 1/Exp. 2)
Scheduler
LambdaLR
Learning rate scheduler type Number of warmup epochs
Scheduler Warmup
10
EMA Decay
0.95
Exponential moving average decay rate Number of samples per batch k-nearest neighbors at time t=1 k-nearest neighbors at time t=0
Batch Size
128
KNN (t=1)
Fully connected
KNN (t=0)
c= [VN]
Clusters (t=1)
[veN]
Number of coarse-graine nodes at t=1 Number of coarse- -rasiieg nodes at t=0 Number of function evaluations for diffusion
Clusters (t=0)
N
NFEs (Diffusion)
1000
NFEs (Flow-Matchhigg
200
Number of function evaluations for now-matching
Table 5. Hyperparameter values and descriptions used for experiments & 2.
Time
Time 35
Cancentration
60
Space
Space
Figure 6. Example of the spatiotemporal data generated by the reaction reaction-diffusion equation.
B.3. Spatiotemporal Transcriptomics
Dataset Construction. Spatiotemporal transcriptomic measures the gene expression of many single-cells distributed spatially (c.g. tissue slice) over multiple timepoints. A graph in this context can be viewed as having N nodes (each node a cell) where the features represent the expression levels of three genes (Bmp, Sox9, and Wnt). The position of each node combines both temporal and spatial information; specifically, a one-dimensional temporal coordinate indicating the progression of the cell through time, and one-dimensional spatial coordinate representing the cell's position along a line across the developing hand tissue. These two coordinates together form the complete positional information for each cell.
This experiment aims to generate the distribution over features while holding the positions fixed. We simulate the trajectory of spatial transcriptomic following Raspopovic et al. (2014). Each gene expression value over time and space is determined by the following reaftion-diffluion equation:
i sox9
Ouaxx!
k2 bmp
Kz wnt 3
i bmp
a'bmp K4 sox9 ki bmp + +d 2 bmp
i wnt
k7 sox9 ki wnt + 2 wnt
Öt Ownt
where a R! represents the initial concentration for each gene, represents the regulation coefficient of one gene on
17

---

Noise-Conditonned Graph Networks
another gene, and sox9, bmp, wnt R represents the gene concentration the current time point. I is the number of spatial values (the higher the value, the larger the space) where each spatial value is scalar. The other variables and their specific values are further detailed in Table 6.
For our training set, we generate a dataset of 10000 samples by simulating the reaction-diffusion diffusion equation under different ini- tial conditions. For each simulation, we randomly initialize the concentration of each gene by sampling from 21(-0.01, 0.01). An example of the generated data is shown in Figure 6. We then divide the spatial dimension into 10 equally spaced points and record the system's 's evolution at 10 different timepoints, resulting in graphs with 100 nodes (10 spatial points 10 timepoints).
Since graph neural networks are discretization invariant (Li et al., 2020), we generate test set of 2000 samples using different spatial and temporal resolution. Specifically, we use 8 spatial points and 12 timepoints, resulting in graphs with 96 nodes.
We detail the values for variables used to simulate the from Raspopovic et al. (2014). Further, we can visualize the reaction-diffusiin equation with a Turing system displayed in Figure 7.
equation in Table 6, which are directly taken
Biologically-Inppin Tasks. We also visualize the conditional tasks that we evaluate DMP and baselines (Section 5.2) on in Figure 8. The details and biological relevance of each task is as follows:
1. (Temporal Trajectory Prediction). Given the gene expression values over space for the first timepoint, predict the future time trajectory of the gene expression values over space. This is often the default task of current spatiotemporal transcriptomics models.
2. (Temporal Interpolation). Predict the trajectories that begin at the given starting spatial gene expressions and terminate at the given ending spatial gene expressions. For example, we may want to find the path skin cells (starting timepoint) take to reprogram into stem cells (ending timepoint).
3. (Gene Imputation). Given missing gene over space and timepoints, predict the likely expression of this gene conditioned on the other genes. For example, some instruments can only measure a subset of genes and we may want to impute the expressions of the other genes not in the subset.
4. (Spatial Imputation). With missing spatial portion, impute the expressions of this missing spatial portion. For example, many spatial transcriptomics measurements may have portions of the tissue missing (due to experimental error). We may want to impute this missing portion of the tissue.
5. (Gene Knockout). If we knockout a gene, predict what the other gene expressions will look like. For example, we may be interested in how a gene contributes to disease; by knocking out this gene, what are the resulting spatial and temporal gene expressions?
Importance. Having the ability to model these types of tasks and answer biologically relevant questions that previous models (Peng et al., 2024; Qiu et al., 2024) could not-is incredibly significant. By developing an expressive yet efficient flow-based generative model, we take a first step towards foundation model for modeling spatiotemporal data, particularly for single-cell transcriptomics.
C. Experiment 3 Details
Here, we elaborate on the ImageNet experiment and the details of modifying the Diffusion Transformer (DiT) to DMP.
C.1. Modifying a Baseline SOTA Model to DMP
Baseline. For a given 256 x 256 pixel image from ImageNet, the baseline DiT model first encodes the image to 32 x 32 using the encoder of pretrained autoencoder. Since it is computationally infeasible to employ full (quadratic) attention over all 32 x 32 = 1024 elements, the model patchifies the 32 x 32 sample into coarser versions based on a chosen patch size p E (2, 8]. The larger the patch size, the coarser the representation of the image. For example, p = 2 results in a 16 x 16 representation while p 8 results in 4 x representation. The model then uses a transformer to fully attend these
18

---

Noise- Cise-Conditioned Graph Networks
Sox9
k3
k4
k7
k2
Wnt
Bmp
kg
kg
Inhibition
Activation
Figure 7. Turing Network for Bmp-Sox9-Wnt (Raspopovic al., 2014).
Parameter
Value
Description
Olgux9
a (0000, 0.01)
Base production rate of Sox9 protein
Olbmp
U(-0001, 0.01)
Base production rate of BMP
Owwt
((00.00,000))
Base production rate of WNT
l02
1
Positive regulation coefficient of BMP on Sox9
ks
-1
Negative regulation coefficient of WNT on Sox9
ka
1.27
Negative regulation coefficient Sox9 on BMP Self-regulation coefficient BMP
ks
-0.1
lot
1.59
Negative regulation coefficient of Sox9 on WNT Self-regulation coefficient of WNT
kg 1
-0.1
do
Diffusion coefficient for BMP
duy 2.5
Diffusion coefficient for WNT
Table 6. Parameter values and descriptions for the reaction-diffusion system that generated the spatiotemporal transcriptomics dataset in Section 5.2.
patches to learn how reverse the noise. We refer readers Peebles & Xie (2023) for details on the image transformations (moving image to latent space, patching, etc) and the transformer architecture used within this model.
We set the patch size for the baseline DiT to because enables fair comparison with DiT-DMP while allowing DiT-DMP to use multiple patch size configurations-wiihh p = 4, DiT-DMP can adaptively utilize patch sizes of 1, 2, and 4, whereas choosing smaller baseline patch size would constrain DiT-DMP's range patch sizes that can use. Further, we leverage on the Small model (Peebles & Xie, 2023) because while larger models like Large X-Large could potentially achieve better absolute performance, architecture to improve relative performance, which can be shown equally well with the Small model.
primary goal is demonstrate that DiT-DMP effectively adapt state-of-the-ant
DIT-DMP. To modify DiT to DiT-DMP, we condition (1) on the noise (i.e. pt) rather than being fixed value (i.c. E 42 8}) and (2) the range of attention on the noise rather than being fixed to full or quadratic attention. Like the two previous experiments, we leverage the exponential adaptive schedule to determine the level of coarse- searse-graining (aka patching) and the range of attention. However, because the patch size must be divisor of 32, we must round the coarse-grain -aase-ggaannin value given by the exponential schedule to the nearest divisor value.
Respecting the notation established in Section 4.1, let Sce be the length of side of the patchified representation (for instance, = 8 means that our patchified represention is 8 x 8 and our patch size is 4) at time t. We choose resolution endpoints so = 8 (resolution at complete noise) and 8 = 32 (resolution at no noise). Correspondingly, we choose range endpoints rn and r --ethe kernel size of the attention at = 0 and = 1-to be such that the computational complexity matches the baseline DiT some given patch size with quadratic complexity. Intermediate s, and ra values are the exponential interpolation between the endpoints rounded to the nearest divisor. Specifically, for some t, we have (re, 8t) ((8,8) 8), (4, (2, 32)
19

---

Coise-Conditione Graph Networks
Ground Truth
Temporal Trajectory
Temporal Interpolation
6.
?
?
?
86
?
?
?
8.0
?
?
?
?
?
?
?
15.0
5.0
Imputation
?
7,5
?
066
?
13.0
60
30-
?
12.5
?
?
?
04
Time
Time
Time
Figure Spatiotemporal Transcriptomic Tasks. In addition to unconditional generation, these are several other useful conditional tasks for single cell transcriptomics. For each task, we delineate what the model is given conditioning information and what the model needs generate with "?"
Notice that unlike DiT, which is limited to p E (2,4, 83 due to the computational complexity of quadratic attention, DiT-DMP can have p = (aka full resolution) since we adapt the attention to be of smaller range than full attention.
C.2. Hyperparameters
Parameter
Value
Layers
12
Hidden size
384
Heads
6
Batch size
64
Number of iterations
800K
NFEs
250
Global Seed
0
Table 7. Hyperparameter Values for DiT (and the DMP variant).
We list all hyperparameters used for this experiment in Table 7. All hyperparameters used in this experiment are the same as Peebles & Xie (2023) except batch size and number of training iterations. Rather than batch size of 256 and 400K training iterations, we have a batch size of 64 with 800K training iterations due compute limitations.
20

---

Woise-Conditiones Graph Networks
C.3. Examples of Generated Images
Figure 9. Select predictions from (Top Row) DiT and (Bottom Row) DiT-DMP.
D. Additional Experiments & Ablations
D.1. Empirical Analysis: Resolution
Using the same dataset from Section 3.2, we also visualize the average Maamov-Wasserstein (GW) distance between the geometric graphs and progressively noised and goarse-grained versions of the geometric graphs. Specifically, for each geometric graph, we generate increasingly noisy versions of the graph (isotropic noise added to the positions). Then, for each noised version of the graph, we coarse grain the graph by either max or mean pool the node positions. then calculate the GW distance between the original graph and the coarse-grained noisy version.
21

---

Noise- CiiseCoonditooed Graph Networks
Max Pooling Over Features
Mean Features
015.00
Gramov- Seoo--Wasserstee Wasserstein Distance
0:0.75
08-
24-
02
Max Pooling Over Positions
Mean Pooling Over Positions
08-
Gromov- ww--Wasssrteinn Distance
06-
04-
0.-
022
25
00
128
Number of Clusters
Number of Clusters
Figure 10. -a distance between the original unnoised geometric graph and coarse- graiine versions of the geometric graph under different levels of noise (a t, meaning higher a - higher noise). Two types of coarse-graining (max and mean pooling) are conducted either positions features.
Results. Figure 10 summarizes the results of this analysis. The key pattern to note is that as we increase the level of coarse-graining (i.e. lower the number of clusters), low-noise versions of the graph increase in GW distance drastically while high-noise versions decrease in GW distance we coarse-grain. This decrease in GW distance for high-noise versions makes sense:
increase noise, the coarse-graining effectively smooths out the noise, better preserving the original signal. This finding is also supported by Theorem 3.2, which shows that increasing the radius of aggregation (e.g. level of coarse-graining where the aggregation scheme is the mean or max) necessary to maximize mutual information for increasing noise levels.
Figure is derived from this experiment where we instead visualize the number of clusters that minimize the GW distance for each noise level.
D.2. Performance with Increasing Number of Layers
There are three popular issues related the number of layers in a graph neural network and performance: (1) oversmooth- ing (Oono & Suzuki, 2020), where node representations become too similar as layers increase; (2) underreaching (Alon & Yahav, 2021), where the network has too few layers capture long-range dependencies, preventing information from propagating between distant nodes that are relevant to the learning task; and (3) over-squashing (Alon & Yahav, 2021), where exponentially growing neighborhoods force the network to compress too much information into fixed-size vectors, leading to information bottlenecks in the message passing process.
Commonly, fully-connected graphs face issues of oversmoothing (since each node feature is the aggregation all the node features in the graph) while fixed range graphs face issues of underreaching and squashiua.n However, since DMP at high noise levels is fully connected while at low noise levels are sparsely connected (we effectively interpolate based on the noise level between these extremes), we investigate whether DMP can resolve these common GNN-relate issues by leveraging the benefits of both kinds of connectivity (i.e. both dense and sparse connectivity).
22

---

Ciise-oonditooned Graph Networks
KNN Fully Connected
Wasserstein Distance
1.4
1.2
10
0.8
2
16
64
Number of Layers
Figure 11. Performance of models on the spatiotemporal transcriptomies dataset across an incrcasing number of message passing layers. We report the Wasserstein distance for the unconditional generation task for DMP, kNN graph, and fully connected graph.
We train our method and baselines (kNN and fully connected) using the same architecture as described in Appendix B.1 on the spatiotemporal transcriptomics dataset, except with a hidden dimension of 32 due to computation limitations. Critically, vary the number of message passing layers K and evaluate the performance of each method the unconditional generation task with different number of layers K. We use GCNConv the message passing architecture.
Results. Figure 11 plots the results of this experiment. At the few number of layers regime (2 - 4 layers), DMP dominates in performance while kNN performs significantly worse compared to fully connected. This finding is expected because, without many layers, kNN faces the underreaching problem since far away nodes cannot exchange information; the fully connected baseline does not have this issue. However, at the higher number of layers regime (8+ layers), we see that the Wasserstein distance of both DMP and fully connected rises steadily while kNN decreases its Wasserstein distance between 8 16 layers before increasing again. Eventually, at the highest number of layers (64 layers), all methods trend towards effectively random predictions (since random predictions result in wasserstein distance of 2.449 for the unconditional generation task as provided in Table 3).
Surprisingly, DMP faces similar issues as fully connected where increasing the number of layers directly decreases its performance. Such a result demonstrates that DMP still suffers from the issues of oversmoothing as fully connected baselines. Ultimately, the lowest Wasserstein distance is still achieved by DMP with fewer number of layers.
23


Reviews:
[{"review_id": "6LMnUYoMtp", "reviewer": "Reviewer_5xLj", "summary": "This work propose to change the architecture of the backbone model according to the noise level in flow match models. It shows that the reception field should be expanded and the resolution should be coarsen in high noise level. Based on this insight, this work proposes  DMP, which consistently outperforms noise-independent architectures on 3D point clouds, spatio-temporal data and images.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We thank the reviewer for their time, effort, and constructive feedback. We address the reviewer’s concerns and questions below:\n\n> In section 3.1, this work cited no work on geometric generative models\n\nOur paper references key geometric generative modeling works in the introduction (e.g. [1-5]), but we agree that these references should be more explicitly discussed in Section 2.1 as well. We will ensure these important works are properly cited and discussed in Section 2.1 of the camera-ready version. Please let us know if there are relevant references we missed.\n\n[1] An autoregressive generative model of 3d meshes. ICML, 2020.\n\n[2] Pointflow: 3d point cloud generation with continuous normalizing flows. ICCV, 2019.\n\n[3] Highly accurate protein structure prediction with alphafold. Nature, 2021. \n\n[4] Diffdock: Diffusion steps, twists, and turns for molecular docking. ICLR, 2023.\n\n[5] Spatiotemporal modeling of molecular holograms. Cell, 2024. \n\n> Authors somehow overclaims their contribution. In abstract, this works claims to change GNN architecture according to noise level, however, only resolution and reception fields are changed, while other architecture designs are not discussed.\n\nOur original claim is for **the general framework of NCGN** since it is the first to formalize conditioning the GNN architecture on the noise-level of the generative process. **DMP is a specific implementation** of NCGNs that provides a practical model that can be used out-of-the-box with current flow-based generative models, as shown in Section 5.3. \n\nAnd while it is true that many other architectural components of GNNs could be conditioned on the noise level, Section 3 of our paper theoretically and empirically suggests that resolution and reception are two critical changes that impact expressivity, which is why we focus on these in our paper. Thus, our work serves as a first step in the potentially many implementations of NCGNs to improve performance of geometric generative models. As noted in the conclusion, NCGNs could adapt other architectural components (layer count, width, message passing type) in future work. We will clarify this distinction more explicitly in the camera-ready version to prevent any perception of overclaiming.\n\n> Baseline only includes the most vanilla models in the field, and the SOTA methods are not compared. Moreover, resolution/reception field in the center point in this work. However, graph transformers and GNNs with virtual node are known to capture global information.\n\n**DMP complements rather than replaces SOTA methods**. Our primary contribution is the noise-conditioning framework that can enhance existing architectures. In fact, DMP already incorporates concepts similar to what the reviewer suggests:\n\n1. **DMP utilizes virtual nodes by design** - Our coarse-grained \"supernodes\" in Section 4.1 function similarly to virtual nodes, but with a critical difference: they dynamically adapt their number and connectivity based on noise level. At high noise, we use fewer supernodes with wider receptive fields; at low noise, we use more supernodes with localized connectivity.\n\n2. **Our baselines effectively generalize many SOTA approaches** - Our fully-connected GAT baseline is effectively a graph transformer, the long-short range baselines capture similar connectivity patterns as architectures with static virtual nodes, and Section 5.3 demonstrates that incorporate noise-conditioning to a SOTA transformer-based model results in significant performance gains.\n\n> Will the change of architecture and graph during generation leads to significant computation overhead?\n\nDetailed in Section 4.1 under “Linear Time Complexity,” DMP maintains **linear-time message passing** throughout the generative process by ensuring that the product of connectivity ($r_t$) and resolution ($s_t$) remains constant: $r_ts_t = r_1N$. This design enables DMP to take \"the best of both worlds\" of having the expressivity benefits of more complex (e.g. fully connected) architectures while having the time complexity of sparsely connected methods.\n\n---\nThank you again for your helpful feedback. We welcome any additional discussions.", "claims_and_evidence": "Authors somehow overclaims their contribution. In abstract, this works claims to change GNN architecture according to noise level, however, only resolution and reception fields are changed, while other architecture designs are not discussed.", "methods_and_evaluation_criteria": "This work includes various tasks. However, the baseline only includes the most vanilla models in the field, and the SOTA methods are not compared. \nMoreover, resolution/reception field in the center point in this work. However, graph transformers and GNNs with virtual node are known to capture global information. This work includes DiT in image task, but these global GNN baselines are still important for other tasks.", "theoretical_claims": "Yes, I checked Section 3.", "experimental_designs_or_analyses": "Yes, I checked section 5.1.", "relation_to_broader_scientific_literature": "This work improves flow-based generative models. But its idea is novel to me.", "essential_references_not_discussed": "In section 3.1, this work cited no work on geometric generative models.", "other_strengths_and_weaknesses": "Clear theory and straightforward intuition on the relation between resolution and noise level.", "other_comments_or_suggestions": "Table's caption should be placed above the table.", "questions_for_authors": "Will the change of architecture and graph during generation leads to significant computation overhead?", "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": "", "extra_scores": {}}, {"review_id": "aYNe3q56HH", "reviewer": "Reviewer_qfiH", "summary": "This paper introduces Noise-Conditioned Graph Networks (NCGNs), a generative modeling approach for geometric graphs that dynamically adjusts graph structure based on noise levels rather than keeping it fixed throughout the process. The authors propose a method to adapt how information flows through the graph depending on the noise intensity.\n\nAt the core of this approach is Dynamic Message Passing (DMP), which modifies both the connectivity range (how far messages travel) and the resolution (level of detail in the graph representation) as noise changes. When noise is high, the model broadens connections and simplifies the graph structure; as noise decreases, it shifts to a more refined representation. The authors support this design with theoretical insights, showing that stronger noise necessitates information aggregation from more distant nodes and that coarser graph structures can reduce complexity while preserving essential information.\n\nTo validate the approach, the authors evaluate DMP on 3D shape generation, spatial transcriptomics, and image generation, demonstrating that it consistently outperforms existing graph-based generative models. Importantly, the method maintains linear-time complexity, making it scalable for large datasets.\n\nThis work highlights the importance of structured noise in generative models, arguing that their graph representations should evolve accordingly. The results suggest that noise-adaptive architectures lead to more expressive and efficient generative models across different domains.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "We thank the reviewer for their time, effort, and constructive feedback. We address the reviewer’s concerns and questions below:\n\n> Essential References Not Discussed\n\nThank you for pointing out these references. We will ensure these works are discussed in the related works of the camera-ready version:\n\n**Adaptive GNNs (Li et al., 2018; Shirzad et al., 2022)**: While these works and our paper both adapt the graph structure, they address fundamentally different objectives—discriminative tasks or evaluation metrics rather than generative processes. Their adaptation mechanisms also differ from our noise-conditioned approach, as they either learn task-specific static structures or improve evaluation methods without dynamically modifying architectures based on noise levels. \n\n**Noise-aware graph learning (e.g. Luo et al., 2021; Zhang et al., 2020)**: While Luo et al. (2021) redraws edges at different noise levels, they use a *static* receptive field and resolution throughout the process (same fixed radius at all noise levels, as shown in Figure 3 of their paper), contrasting with our *dynamic* approach that systematically varies both factors. Zhang et al. (2020) focus on subsampling with full attention rather than noise-conditioning.\n\n**Hierarchical graph pooling (e.g. Ying et al., 2018; Defferrard et al., 2019)**: There are similarities between these approaches and our work as they all aim to have a multiscale view of the graph. However, these methods employ fixed hierarchical structures for classification/regression tasks, while our approach continuously adapts the hierarchical representation based on the noise level of the generative model. Further, the pooling techniques in these works could complement NCGNs as coarse-graining procedures.\n\n**Multi-resolution diffusion models (e.g. Zhao et al., 2023)**: We share a common insight with Zhao et al. (2023) and other works like cascading diffusion models: it is beneficial to have coarser representations at high noise regimes and fine-grain representations at low noise regimes. However, these works operate with discrete, manually-determined resolution stages requiring separate models for each resolution, while NCGNs provide a single continuous model with automatically adapted resolution *and connectivity*.\n\n> The theoretical analysis assumes isotropic Gaussian noise and continuous graphs, which may not hold in all real-world settings. More validation on structured or irregular noise would strengthen this claim. \n\nWe acknowledge the limitations in our theoretical assumptions made in Section 3.1, but would like to clarify two key points: (1) **Isotropic Gaussian noise is the standard choice** in modern flow-based generative models, including diffusion models [1], flow-matching [2], and other variants [3,4]. This makes the isotropic Gaussian noise model assumption directly applicable to most real-world generative modeling settings. (2) The continuous graph assumption is indeed a simplification but is necessary to make the theoretical analysis tractable; to address the concerns, we specifically designed our empirical analysis in Section 3.2 to validate that our theoretical insights hold even when these assumptions are relaxed. \n\n[1] Denoising Diffusion Probabilistic Models. NeurIPS, 2020.\n\n[2] Flow Matching for Generative Modeling. ICLR, 2023.\n\n[3] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. ICLR, 2023.\n\n[4] Action Matching: Learning Stochastic Dynamics from Samples. ICML, 2023.\n\n> The choice of an exponential scheduling function is somewhat heuristic, and a deeper study of learned schedules could provide stronger justification.\n\n**Implementing learnable schedules is difficult** because the schedule is used for kNN graph construction and coarse-graining procedures, which involve non-differentiable operations (sorting, clustering, discrete assignments) that are not amenable to standard backpropagation. Instead of relying purely on heuristics, we conducted ablation studies in Table 2 and Section 5.1.1 comparing four scheduling functions (linear, exponential, logarithmic, ReLU), finding that all significantly outperform fixed baselines, with exponential consistently performing best. \n\nWhile developing differentiable approximations for these operations represents an exciting future direction, we believe it warrants dedicated study beyond our current paper's scope.\n\n> certain sections, particularly the theoretical derivations, could be made more accessible. The notation is sometimes dense, which may make it difficult for readers unfamiliar with flow-based generative models to follow.\n\nWe agree that some theoretical sections and notation could be made more accessible. Are there any specific sections, derivations, or notations you think could be improved? This would help us prioritize changes for the camera-ready version. \n\n---\nThank you again for your helpful feedback. We welcome any additional discussions.", "claims_and_evidence": "Most claims are well-supported by theory and experiments, but some areas need further validation. The theoretical analysis and empirical results strongly support the idea that noise-adaptive message passing improves generative modeling. Experiments across 3D shape generation, spatial transcriptomics, and image generation show DMP consistently outperforms fixed architectures, with clear improvements in Wasserstein distance. The paper also convincingly explains how its scheduling strategy ensures linear-time complexity.\n\nHowever, some claims require more justification. The theoretical analysis assumes isotropic Gaussian noise and continuous graphs, which may not hold in all real-world settings. More validation on structured or irregular noise would strengthen this claim. The choice of an exponential scheduling function is somewhat heuristic, and a deeper study of learned schedules could provide stronger justification. While the method performs well on tested datasets, its applicability to other domains like molecular modeling or social networks remains uncertain. A discussion of these limitations would improve clarity. This I'll get back to later.\n\nOverall, the paper presents strong evidence for its approach, but addressing theoretical assumptions, adaptive scheduling choices, and broader applicability would further strengthen the claims.", "methods_and_evaluation_criteria": "The methods and evaluation criteria align well with the problem. The authors use relevant benchmark datasets—ModelNet40, spatial transcriptomics, and ImageNet—covering diverse structured data types. Metrics like Wasserstein distance and FID are appropriate for measuring distribution alignment and sample quality. While comparisons to baseline graph-based models are useful, evaluating against other adaptive GNN architectures could further strengthen the analysis.", "theoretical_claims": "As far as I can tell, the proofs are done really well and are correct.", "experimental_designs_or_analyses": "Nothing to add here.", "relation_to_broader_scientific_literature": "The paper builds on and extends several key areas in generative modeling, geometric deep learning, and adaptive graph structures:\n\n- Flow-Based Generative Models: \nThe work extends diffusion models (Ho et al., 2020) and flow matching (Lipman et al., 2022) by introducing a noise-adaptive graph structure, rather than using a static neural network across the generative process. This aligns with prior work in score-based generative modeling (Song & Ermon, 2019) but introduces a dynamic graph representation to improve expressivity.\n\n- Graph-Based Generative Modeling: \nPrior works on graph generative models (Corso et al., 2022; Xu et al., 2022) typically use fixed message-passing radii (e.g., k-NN graphs) during training. This paper challenges that assumption, showing that adjusting graph connectivity and resolution based on noise level leads to better representation learning.\n\n- Multi-Scale Graph Representations: \nThe use of coarse-graining at high noise levels connects to ideas in hierarchical graph representations (Li et al., 2020) and graph signal processing (Oono & Suzuki, 2020). The paper strengthens the case that adaptive resolution improves generative modeling, a principle also explored in neural operators for PDEs.", "essential_references_not_discussed": "- The paper should discuss adaptive GNNs (e.g., Li et al., 2018; Shirzad et al., 2022) as related approaches for dynamically adjusting graph connectivity.\n- Noise-aware graph learning (e.g., Luo et al., 2021; Zhang et al., 2020) is conceptually close to NCGNs and could be cited.\n- Hierarchical graph pooling methods (e.g., Ying et al., 2018; Defferrard et al., 2019) share similarities with DMP's coarse-graining approach.\n- Multi-resolution diffusion models (e.g., Zhao et al., 2023) explore resolution adaptation in generative modeling and could strengthen the discussion around DiT-DMP.", "other_strengths_and_weaknesses": "The paper presents a compelling and well-motivated approach to generative modeling of geometric graphs by introducing Noise-Conditioned Graph Networks (NCGNs). The core idea—that graph neural network architectures should adapt dynamically to the noise level—is both intuitive and well-supported by theoretical and empirical analyses. I found the motivation strong, particularly in demonstrating that static graph structures are suboptimal for generative processes where noise plays a fundamental role. The experiments are thorough, covering diverse domains such as 3D point clouds, spatiotemporal transcriptomics, and image generation, which strengthens the paper’s claims regarding the generality of the approach.\n\nOne of the key strengths of the paper is its blend of theoretical insight and practical implementation. The information-theoretic analysis linking noise levels to optimal message-passing radius is a valuable contribution, providing a principled foundation for the proposed method. Additionally, the empirical analysis supports these theoretical findings in a convincing way, particularly through attention weight visualizations and coarse-graining experiments. The modification of an existing state-of-the-art model (DiT) to incorporate NCGNs with minimal changes is another highlight, showcasing the practical viability of the approach.\n\nHowever, I do have some concerns. The theoretical results, while interesting, rely on certain simplifying assumptions which I am not sure are in conflict with real-world data. Moreover, though the empirical validation mitigates this to some extent, it would be useful to discuss potential limitations or failure cases in scenarios with highly structured noise or irregular graph connectivity patterns. Additionally, while the authors compare different adaptive scheduling strategies for DMP, the choice of an exponential schedule is somewhat heuristic. A more systematic exploration, perhaps with learned schedules or adaptive mechanisms, could further strengthen the claims.\n\nFrom a clarity perspective, the paper is generally well-written, but certain sections, particularly the theoretical derivations, could be made more accessible. The notation is sometimes dense, which may make it difficult for readers unfamiliar with flow-based generative models to follow. More intuitive explanations or visualizations of key mathematical insights could improve readability.\n\nOverall, this is a strong and original contribution with clear significance for the field of geometric generative modeling. While there are some areas for refinement, particularly in the theoretical assumptions and clarity of exposition, the paper convincingly demonstrates that noise-adaptive graph structures offer a meaningful improvement over existing approaches. \n\nI actually really like the idea. I think if you add some more ablations and explain the limitations of the work better, it is a very solid piece of work.", "other_comments_or_suggestions": "No", "questions_for_authors": "See sections above.", "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": "", "extra_scores": {}}, {"review_id": "YmccfHFFRt", "reviewer": "Reviewer_GxaG", "summary": "The paper introduces ​Noise-Conditioned Graph Networks (NCGNs), a class of graph neural networks that dynamically adapt their architecture based on the noise level during flow-based generative modeling of geometric graphs. The key innovation is ​Dynamic Message Passing (DMP), which adjusts both the connectivity range (e.g., k-nearest neighbors) and graph resolution (via coarse-graining) as a function of noise.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We thank the reviewer for their time, effort, and constructive feedback. We address the reviewer’s concerns and questions below:\n\n> Learned scheduler\n\n**Implementing learnable schedules is difficult** because the schedule is used for kNN graph construction and coarse-graining procedures, which involve non-differentiable operations (sorting, clustering, discrete assignments) that are not amenable to standard backpropagation. Instead of relying purely on heuristics, we conducted ablation studies in Table 2 and Section 5.1.1 comparing four scheduling functions (linear, exponential, logarithmic, ReLU), finding that all significantly outperform fixed baselines, with exponential consistently performing best. \n\nWhile developing differentiable approximations for these operations represents an exciting future direction, we believe it warrants dedicated study beyond our current paper's scope.\n\n> Coarsening via voxel clustering or K-means could introduce artifacts; learnable pooling (e.g., DiffPool) might improve performance\n\nWe agree that learnable pooling methods like DiffPool could potentially enhance performance. We emphasize that NCGN is a general framework where specific implementation details (including coarsening operations) can be customized. We intentionally selected simple, deterministic pooling methods to clearly demonstrate our core contribution—the benefit of noise-conditioned architectures—without confounding factors from complex learnable components. Our experimental results show that even with these basic coarsening operations, NCGNs deliver significant performance improvements across multiple domains. Future work could certainly explore more advanced coarsening operations within the NCGN framework.\n\n> E(n)-GNNs and Equivariant Diffusion models explicitly encode geometric symmetries [...]. Is there a way to integrate equivariance constraints into NCGNs?\n\n**DMP complements these equivariant methods**. As a general framework, NCGNs can incorporate any type of GNN layer, including equivariant ones. For instance, DMP could be implemented with equivariant layers like those in E(n)-GNNs, SchNet, or DimeNet that incorporate pairwise distances between nodes. Our theoretical analysis in Section 3 could also inform best practices for equivariant architectures—for example, many equivariant models operate on radius graphs, and our work provides principled guidance on how the radius should vary with noise level. \n\nWhile further investigation is needed, we believe the performance gains from NCGNs could potentially translate to equivariant models as well, making the application of NCGNs in structural biology and other geometric domains a promising direction for future research.\n\n> Correlation function in Theorem 3.2 is idealized; real-world spatial correlations may be more complex.\n\nWe agree that the correlation function used in Theorem 3.2 is a simplified model of spatial correlation. This simplification was necessary to make theoretical analysis tractable while still capturing the essential property that correlation between nodes decreases with distance. Importantly, we designed our work to not rely solely on theoretical guarantees from idealized assumptions. Section 3.2 provides empirical analysis on real 3D geometric objects, demonstrating that our key insights about optimal reception fields and resolution scaling with noise also hold true in practice.\n\n> Broader baseline comparisons\n\nOur primary contribution is the noise-conditioning framework of NCGNs that can enhance existing architectures including many SOTA approaches. We implement NCGN on top of widely used architectures like GCNs and GATs and compare NCGN against general connectivity patterns including sparse connections (kNN), long and short connections, and full (quadratic) connections. In fact, **these effectively generalize many SOTA approaches**: the fully-connected GAT baseline is effectively a graph transformer and the long-short range baselines capture similar connectivity patterns as architectures with virtual nodes. Furthermore, Section 5.3 demonstrates that incorporating noise-conditioning to a SOTA transformer-based model results in significant performance gains.\n\n---\nThank you again for your helpful feedback. We welcome any additional discussions.", "claims_and_evidence": "**Strengths:**\n​\n- It is the first work to formalize noise-dependent graph architecture adaptation for geometric generative modeling.\nIntegrates theoretical insights (mutual information analysis) with practical algorithmic design (DMP).\n\n- Comprehensive Experiments: Validated on 3D shapes, biological data, and images, demonstrating broad applicability.\n​Theoretical-Experimental Alignment: Attention weight analysis (Fig. 3) and Gromov-Wasserstein distance measurements (Fig. 4) empirically confirm theoretical predictions.\n\n​**Weaknesses:**\n\n- The ​scheduler for adjusting resolution/connectivity (Fig. 5) is predefined (e.g., exponential), not learned. This may limit optimality across diverse datasets.\n- Coarsening via voxel clustering or K-means could introduce artifacts; learnable pooling (e.g., DiffPool) might improve performance.\n​Baseline Comparisons:", "methods_and_evaluation_criteria": "The method and Evaluation Criteria in this paper seems to be sound.", "theoretical_claims": "Correlation function in Theorem 3.2 is idealized; real-world spatial correlations may be more complex.", "experimental_designs_or_analyses": "The experimental design seems to be valid.", "relation_to_broader_scientific_literature": "-", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "-", "other_comments_or_suggestions": "The paper presents a ​theoretically grounded and empirically validated method for noise-adaptive graph generation. While some assumptions simplify real-world complexity, the core claims are well-supported, and the approach demonstrates clear practical value. Future work could explore learned schedulers and broader baseline comparisons.", "questions_for_authors": "**Q1:** The paper uses predefined schedulers (e.g., exponential) to interpolate between connectivity and resolution boundaries. Have you explored learned schedulers (e.g., via meta-learning or reinforcement learning) to dynamically optimize $r_t$ and $s_t$ during training? If not, do you foresee such approaches improving performance, particularly in domains where noise dynamics are non-monotonic or dataset-specific?\n\n**Q2:** Recent works like E(n)-GNNs and Equivariant Diffusion models explicitly encode geometric symmetries (e.g., rotation/translation invariance) into graph networks. How does DMP complement or contrast with these methods?  Is there a way to integrate equivariance constraints into NCGNs, is it possible that this would enhance performance on tasks like molecular generation, where symmetry preservation is critical?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes Noise-Conditioned Graph Networks (NCGNs), which adapt graph architecture based on noise level during generative modeling. The main instantiation, Dynamic Message Passing (DMP), adjusts receptive field and resolution and is theoretically motivated and empirically validated across several domains.

Reviewers found the paper novel, with clear theoretical grounding and strong empirical results. Concerns included missing comparisons to some SOTA methods, overclaiming in the abstract, and the use of heuristic (non-learned) schedulers. The authors addressed these points in detail, and reviewers were generally satisfied.
Decision: Accept (poster)