Geometric Generative Modeling with Noise-Conditioned Graph Networks

We thank the reviewer for their time, effort, and constructive feedback. We address the reviewer’s concerns and questions below:

Learned scheduler

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.

While developing differentiable approximations for these operations represents an exciting future direction, we believe it warrants dedicated study beyond our current paper's scope.

Coarsening via voxel clustering or K-means could introduce artifacts; learnable pooling (e.g., DiffPool) might improve performance

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

E(n)-GNNs and Equivariant Diffusion models explicitly encode geometric symmetries [...]. Is there a way to integrate equivariance constraints into NCGNs?

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.

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

Correlation function in Theorem 3.2 is idealized; real-world spatial correlations may be more complex.

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

Broader baseline comparisons

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

Thank you again for your helpful feedback. We welcome any additional discussions.

We thank the reviewer for their time, effort, and constructive feedback. We address the reviewer’s concerns and questions below:

In section 3.1, this work cited no work on geometric generative models

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

[1] An autoregressive generative model of 3d meshes. ICML, 2020.

[2] Pointflow: 3d point cloud generation with continuous normalizing flows. ICCV, 2019.

[3] Highly accurate protein structure prediction with alphafold. Nature, 2021.

[4] Diffdock: Diffusion steps, twists, and turns for molecular docking. ICLR, 2023.

[5] Spatiotemporal modeling of molecular holograms. Cell, 2024.

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.

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

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

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.

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:

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

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

Will the change of architecture and graph during generation leads to significant computation overhead?

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

Thank you again for your helpful feedback. We welcome any additional discussions.

We thank the reviewer for their time, effort, and constructive feedback. We address the reviewer’s concerns and questions below:

Essential References Not Discussed

Thank you for pointing out these references. We will ensure these works are discussed in the related works of the camera-ready version:

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.

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.

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.

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.

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.

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

[1] Denoising Diffusion Probabilistic Models. NeurIPS, 2020.

[2] Flow Matching for Generative Modeling. ICLR, 2023.

[3] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. ICLR, 2023.

[4] Action Matching: Learning Stochastic Dynamics from Samples. ICML, 2023.

The choice of an exponential scheduling function is somewhat heuristic, and a deeper study of learned schedules could provide stronger justification.

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.

While developing differentiable approximations for these operations represents an exciting future direction, we believe it warrants dedicated study beyond our current paper's scope.

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.

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

Thank you again for your helpful feedback. We welcome any additional discussions.