Dear reviewer,

Thank you for your comments.

Weakness 1:

We should perhaps make this more explicit in the text: all our baselines are more powerful than those presented in Law and Lucas, where no neural networks are used. We have added a more detail explanation in Appendix D.5.

In our experiments, a feature encoder is employed across all cases, including those where the geometry of the space is fixed. Specifically, for the Minkowski and De Sitter metrics, we use a neural network to map the original node features to the geometry, even when the metric itself is not learned. This approach ensures a fair comparison between using a fixed, closed-form metric and learning the geometry directly with NSTs.

It is important to note that the best possible distortion is 1, which corresponds exactly to an isometric embedding. Indeed, if the distortion=1, then Definition 2 (ii) becomes an equality. We understand that this may not be clear to researchers not used to this kind of evaluation measure with metric embedding problems, and we will add a remark in our paper. We disagree that Minkowski embeddings achieve similar levels of final distortion in real-world datasets. For instance, they perform particularly poorly in In silico, E. coli, and S. cerevisiae (see Table 2), these are real-world gene regulatory networks. Additionally, in line with neural snowflakes, we would like to emphasize that NSTs are especially effective at preserving low distortion in lower dimensions: when the embedding dimension is reduced to 2, the performance gap between NSTs and the fixed metrics becomes notably more pronounced.

Weakness 2:

We have included downstream experiments in Appendix D.6 using NST features to augment the input features to downstream classifiers such as MLPs, GCNs, GATs, CPGNN-MLPs, and CPGNN-Chebys for node classification. We observe that NST features help downstream accuracy. Please refer to the appendix for further details. We also discuss this in the response to reviewer vRyg.

Weakness 1:

We discuss experiments extensively in Appendix D and provide additional details about the computational implementation of Neural Spacetimes in Appendix C. However, we agree that additional discussion on the computational implementation of learnable geometries (such as Neural Spacetimes) versus fixed geometries (such as Minkowski and de Sitter space) would be helpful for readers (we have added appendix D.5).

For all experiments, we use a feature encoder that maps the graph node features to the relevant manifold. We assume that the output of the feature encoder is in Euclidean space in all cases. We then proceed according to the specific geometry. This approach aligns with previous works such as neural snowflakes and other studies on product manifold embeddings cited in the main paper. Intuitively, one can understand the problem as a two-step process, which in practice is learned jointly via end-to-end gradient descent. In the first step, the encoder learns to position the graph nodes in space—that is, it learns to map the original node features to coordinates in the manifold. In the second step, the distance between points on the manifold is evaluated, either based on a given metric when the geometry is fixed or by learning the geometry (the metric itself) in a data-driven fashion. In the cases of Minkowski and de Sitter space, the geometry is not learned but given. More specifically, in Minkowski space, we use the feature outputs of the Euclidean feature encoder, apply a change in the metric tensor, and recognize that one coordinate has a different signature. In this particular case, we do not require an exponential map since the space is flat. For de Sitter space, the situation is similar: only the positioning of points in the manifold is optimized via gradient descent, while the geometry remains fixed. In terms of implementation and metric calculation, de Sitter space is a curved manifold, making computations more complicated than for Minkowski space. We use an embedding approach that avoids exponential maps: we map points into a de Sitter hyperboloid in a higher-dimensional Minkowski space and compute geodesic distances there. The geometric operations we utilize include normalizing spatial components, computing the de Sitter inner product, and handling both timelike and spacelike separations. Neural spacetimes are intrinsically different in that their geometry is not fixed but rather random at initialization and parametrized by a neural network that always outputs a pseudo-quasi-metric by construction. However, the specific pseudo-quasi-metric it generates is learned during optimization based on the input graph. Hence, in this case, both the position of the points on the manifold and the geometry of the manifold itself are learned jointly via gradient descent. For Neural Spacetimes, we don't need exponential maps either since the construction is based on warping a norm. During optimization, the inputs are the graph node features, and the targets are both the distances between nodes (induced by the graph geometry) and the causality structure (given by the edge directions). Since Neural Spacetimes are able to reshape the geometry of the embedding space as a function of the data, they are inherently more flexible and they are able to embed DAGs with less distortion (the best possible value for distortion is 1).

Weakness 2:

Please refer to Appendix C.2 for algorithmic descriptions of the components of the NST model. The feature encoder is an MLP, and both the neural partial order and the neural quasi-metric are based on linear projections and pointwise non-linearities. Consequently, the computational complexity is comparable to that of MLPs. We do not use any computationally expensive operations, such as attention.

For an MLP with $L$ layers, the total time complexity sum across layers (total number of layers L) is $O\left(n \sum_{i=1}^{L} d_{\text{in}}^{(i)} \times d_{\text{out}}^{(i)}\right)$ where $d_{\text{in}}^{(i)}$ and $d_{\text{out}}^{(i)}$ are the input and output dimensions of the linear layer $i$, and $n$ is the batch size. The same would apply in our case.

Question 1:

Thank you, this was a typo and also pointed out by Reviewer Crpy. The correct formula (which we have used in its correct form throughout the manuscript of course) should read: $c\,\rho(f(x_1),f(x_2)) \le d(x_1,x_2) \le D \,c\rho(f(x_1),f(x_2)).$ We have made the correction; thank you for noticing this typo.

Dear reviewer,

Thank you for your comments, see responses below.

Weakness 1: Further experiments

We have included the experiments addressing your concern regarding different distributions of graphs, see Appendix D.2.1 in which we vary the graph connectivity. We find that given the same number of nodes, a less connected graph is easier to embed in space and harder to embed in time according to our experiments.

Weakness 2.1: Language in lines 420-422.

Thank you for pointing out the inconsistency in notation. We have fixed it for the final version of the manuscript. We have rewritten everything in terms of $\hat{x}_u,\hat{x}_v,\hat{x}_w$.

Weakness 2.2: T and W variables

Thank you for pointing out a possible confusion for readers. We chose to use W in that theorem statement, as we are emphasizing that for a graph with width W one may take T=W; i.e.\ W is given before T is determined to be T:=W. Nevertheless, if the paper would be clearer by adding some details on that then we would be more than happy to.

Weakness 2.3: Typo in line 215

Indeed, you are right; thanks for noticing that typo. We have corrected it.

Question: Space and Time dimensions in the latent space embedding.

Both the features passed to $\mathcal{D}$ and $\mathcal{T}$ are functions of the original node features. The original node features lie in $\mathbb{R}^N$ and are mapped to $\mathbb{R}^D$ and $\mathbb{R}^T$. Manually partitioning these features in $\mathbb{R}^N$ into spatial and temporal components is impractical, as the original graph lacks a defined notion of spatial and temporal features. Thus, we use a neural network encoder to map the features, learning an optimal transformation of the data via gradient descent in an end-to-end manner. Alternatively, the feature encoder network can be seen as creating two parallel projections of the original input features: one for space and one for time. In summary, the NST model learns the optimal projection of the original input features for space and time by propagating gradients through $\mathcal{D}$, $\mathcal{T}$, and $\mathcal{E}$ simultaneously. The choices of $D$ and $T$ are left as hyperparameters, which may be determined experimentally. Alternatively, domain knowledge may guide these choices; for example, if the number of antichains in the underlying DAG is known a priori, this information can help set the time dimension accordingly. In Appendix C.5, we also discuss potential classical algorithms from the literature to assist in tuning the $D$ and $T$ hyperparameters.

Dear reviewer,

Thank you for the suggestions. See the responses below.

Weakness 1:

To be able to fully capture the causality structure we must have the number of time dimensions be equal to the number of antichains. Please see Appendix C.5 for additional suggestions on this matter.

Weakness 2:

The NST algorithm is not limited to optimizing for local geometry alone. However, from a computational standpoint, the problem appears intractable. We discuss the global optimization setup in more detail in Appendix C.4. The main challenge lies in generating data to train the algorithm—specifically, generating ground-truth global distances and inferring all global causal relationships within a graph.

Nevertheless, in Appendix D.1, we conduct global embedding experiments on unweighted, undirected trees where all local distances are set to 1. In this particular setup, the problem is simplified, allowing us to focus on evaluating the global spatial embedding capabilities of NST using only the learnable quasi-metric and the feature encoder.

Question 1:

Following the suggestion of reviewer vRyg we have run experiments for ogbn-arxiv, see Appendix D.3 'Ogbn-arxiv Large Citation Graph Embedding'. We observe that in this case NSTs are also superior to the baseline geometries.

Question 2:

This is an interesting question, but we do not have a straightforward answer at the moment, this would be an interesting extension of our work.

Question 3:

Please refer to Section 3.2: ‘Local geometry optimization and global geometry implications’. The model is limited by optimizing the local geometry instead of global geodesic distances over the DAG. Although preserving global distances is not guaranteed, the local transitivity of causal connectivity will implicitly hold globally, even when performing local optimization only.

Question 4:

Similarly to Question 3, the hierarchical causal structure will be preserved globally thanks to the transitivity property even when optimizing locally. The main concern is guaranteeing lack of causality between anti-chains: this will naturally happen with high probability as the number of time dimensions are increased.

Dear reviewer,

Thank you for your comments. We have updated the manuscript to address your concerns. Note that changes are highlighted in red.

Downstream experiments

We understand the concern regarding downstream tasks. First, we would like to emphasize that the primary objective of this paper is focused on embeddings—specifically, the ability to faithfully represent discrete graph structures within a manifold representation. This focus forms the foundation of both our theoretical guarantees and computational framework. Accordingly, the main experiments in the text concentrate on this objective, encompassing both synthetic DAG embedding experiments and real-world applications such as hyperlink and gene regulatory networks.

Nevertheless, we agree that exploring the downstream performance of features derived from these embeddings is an interesting direction for future research. In response to comments raised in Weakness Points 1, 2, and 3, we have included Appendix D.6 at the end of the paper. This appendix proposes a method for leveraging NST-based features in downstream tasks and evaluates them using MLPs, GCNs, GATs, CPGNN-MLPs, and CPGNN-Chebys on the Cornell, Texas, and Wisconsin datasets. We hope this addition addresses:

• Node classification experiments (Point 1),
• Benchmarking with CPGNN and heterophilic graphs (Point 2, as Cornell, Texas, and Wisconsin are heterophilic), and
• Benchmarking with GCN and GAT on real-world learning tasks (Point 3). We have reproduced all the baselines, including the full CPGNN framework (pretraining, initialization, and regularization).

Hyperparameters and results are extensively discussed in the appendix. We also reiterate that this was not the original focus of the paper and invite the community to further explore downstream applications of NSTs.

Finally, we would like to clarify that our framework does not address spatiotemporal graphs. We only consider static graphs with directed edges, not dynamic graphs that evolve over time. In our case, the "time" manifold captures causality (i.e., directionality within the static graph). Spatiotemporal graph learning and DAG learning are distinct problems, and our work focuses on the latter.

Large Graph Embeddings

We also address point 4 in the weaknesses section by including an additional subsection called 'Ogbn-arxiv Large Citation Graph Embedding' inside Appendix D.3. In which we embed the arxiv dataset from OGBN as suggested by the reviewer, which is two orders of magnitude bigger than the other graph datasets we considered before. We consistently observe the superiority of NSTs as compared to fixed geometries.

Questions

• Question 1: The reviewer asks how do NSTs perform in real-world tasks using well-established metrics such as accuracy. Please refer to the newly included Appendix D.6. We can see that including NST based features boost the performance in terms of accuracy for most model configurations and ablations. We would also like to highlight that our distortion metric used to evaluate embedding is a standard 'well-established metric' in the graph embedding literature. Please refer for reference such as 'Cuts, Trees and `1-Embeddings of Graphs' and also ' Embedding Tree Metrics into Low-Dimensional Euclidean Spaces' by Gupta. For the 'directionality' embedding we do use accuracy to quantify the proportion of edges the NST represents correctly.

• Question 2: We have addressed with in Appendix D.6 too. Rather than how NSTs compare to GNNs, we use NST features to enhance downstream classifiers, which may include GNNs. Interestingly, we actually find that MLPs enhanced with NSTs work best at least for Texas, Cornell, and Wisconsin.

• Question 3: See Appendix D.3. We still show superiority of NSTs over Minkowski and De Sitter space. As expected, increasing the graph makes the task harder for all models.

• Question 4: In appendix D.6 we focus on heterophilic graphs and use the baselines suggested by the reviewer. NST features do help in general for this task according to our initial experimentation.