Title: value: Integrating GNN and Neural ODEs for Estimating Non-Reciprocal Two-Body Interactions in Mixed-Species Collective Motion
Authors: value: Masahito Uwamichi
Simon K. Schnyder
Tetsuya J. Kobayashi
Satoshi Sawai
Venue: NeurIPS 2024
Keywords: value: Deep Learning
Neural Differential Equations
Graph Neural Networks
System Identification
Active Matter
Collective Motion
Non-Reciprocal

Abstract: value: Analyzing the motion of multiple biological agents, be it cells or individual animals, is pivotal for the understanding of complex collective behaviors. 
With the advent of advanced microscopy, detailed images of complex tissue formations involving multiple cell types have become more accessible in recent years. However, deciphering the underlying rules that govern cell movements is far from trivial. Here, we present a novel deep learning framework for estimating the underlying equations of motion from observed trajectories, a pivotal step in decoding such complex dynamics. 
Our framework integrates graph neural networks with neural differential equations, enabling effective prediction of two-body interactions based on the states of the interacting entities. We demonstrate the efficacy of our approach through two numerical experiments. First, we used simulated data from a toy model to tune the hyperparameters. Based on the obtained hyperparameters, we then applied this approach to a more complex model with non-reciprocal forces that mimic the collective dynamics of the cells of slime molds. Our results show that the proposed method can accurately estimate the functional forms of two-body interactions -- even when they are nonreciprocal -- thereby precisely replicating both individual and collective behaviors within these systems.


Review discussion:
Two-body interactions are not my research field.
Thus, I don't have enough expertise to evaluate this work. I sent an email to AC to re-assign this paper when I received the assignment. Unfortunately, I didn't receive any feedback. My justification is just based on the writing without any methodology consideration.
see strengths
see strengths
see strengths
No ethics review needed.
- Using a GNN and Neural ODE to describe the dynamics of two-body interaction state changes is well-motivated and reasonable.
- Experiments from interaction models are implemented from simple to complex setups, showing the efficiency and capacity of the proposed method.
- The paper structure and writing is clear and easy to follow.
One of the main contribution claimed in this paper (as well as in its title) is the combination of GNN and Neural ODE, but this approach has been studied in multiple works back since 2020, the authors seem to overlook the important and related body of works [1,2,3,4]. [2] is a direct competing method (GDE) where the authors proposed a model integrating graph NN and Neural ODEs. GDE-related applications exist as well, for example, social network embedding [5], action recognition [6].
Given the missing discussion of related work, and my concern about the major technical contribution claim in this paper, I do not see much merits of this study to be novel and interesting to the community.
[1] L. Xhonneux et al.: Continuous Graph Neural Networks. PMLR 2020.
[2] M. Poli et al.: Graph Neural Ordinary Differential Equations. AAAI 2020.
[3] L. Chen et al.: Signed Graph Neural Ordinary Differential Equation for Modeling Continuous-Time Dynamics. AAAI 2024.
[4] A. Han et al.: From Continuous Dynamics to Graph Neural Networks: Neural Diffusion and Beyond. TMLR 2024.
[5] Y. Zhang et al.: Improving Social Network Embedding via New Second-Order Continuous Graph Neural Networks. KDD 2022.
[6] L. Pan et al.: Spatial-temporal graph neural ODE networks for skeleton-based action recognition. Scientific Reports 2024.
The author should discuss in their response regarding the missing related works that involve integrating GNN+Neural ODEs, and highlight their differences/contribution from previous work.
Limitations were discussed in the paper.
No ethics review needed.
Thank you for bringing to our attention these papers.  We will cite these papers in the final edition.
Re – Question and Weakness (Relationship to Graph ODE):
The pre-existing methods [1-5] solve and train ODEs on a graph with fixed edges and consider changes in the edge weight. These approaches require a fully or almost fully connected graph in solving ODE for an application to collective motion where the adjacency drastically changes in time, such as collective cell migration, flocks, and pedestrian dynamics. In such systems addressing the adjacency at each time is critical because the behavior of individuals is determined by sensible information, that is, the state of neighboring others. This is computationally expensive making it difficult to be extrapolated to situations with unknown graph structures.
In contrast, since edge structures are continuously being updated in our approach, the pruning effect substantially reduces memory requirement - in our case of 400-body simulations with the edge density of 2%,  the required memory of 200 GB for fully connected graph is reduced to 30GB making it feasible on an off-the-shelf GPU board.
As for reference [6], although the method introduces a dynamic edge structure, the time series of the graph (edge) structure needs to be estimated first using an entire time series and thus the trained model cannot be further extrapolated to situations outside of the particular graph structure. Our method resolves this issue by setting a rule to define edges at each time point instead of explicitly giving graph structure time series.
Thanks for the authors clarifying my concerns on the difference of the method from previous work, and highlighting their contributions on computational costs, I have increased my score to 5.
Strengths:
- Novel integration of GNNs and Neural ODEs for collective motion analysis
- Ability to handle mixed-species systems and estimate species-dependent interactions (and demonstrating its effectiveness across simple and complex models)
- Successful generalization to single-species scenarios after training on mixed-species data
- Potential applicability to a wide range of biological systems exhibiting collective motion
Weaknesses:
- High computational cost and long training times .
- Current limitation to deterministic motion equations and pairwise interactions.
- Lack of ablative studies on hyper parameter optimizations and method’s sensitivity to noise.
- Lack comparative analysis with existing methods for estimating collective motion dynamics.
Questions for the authors:
1. How does the performance of your method compare to existing approaches for estimating collective motion dynamics, such as SINDy or Bayesian optimization methods?
2. Have you explored the method’s sensitivity to noise in the input data? How robust is the estimation process to measurement errors or stochastic fluctuations in the trajectories?
3. I’m intrigued by the use of LAMB optimizer. What specific characteristics of your problem or empirical results led you to choose LAMB over AdamW?
4. What are the primary factors limiting the scalability of your approach to systems with stochastic motion equations and higher-order interactions?
yes
No ethics review needed.
Thank you for the feedback and comments.    We are encouraged by your overall evaluation and many of the technical suggestions which are very helpful in future extension of the work.
Re – Question 1 and Weakness 4 (Comparative analysis with SINDy and Bayesian optimization methods):
Thank you for pointing out the lack of comparative analysis.
Although both SINDy and Bayesian optimization are in principle applicable to our target system, they require drastic change of codes.
Thus we have not tested them.
Re – Question 2 and Weakness 3 (The method’s sensitivity to noise in the input data):
No, we have not systematically studied dependency on the noise parameter, nor introduced additional noise, and this is on-going and part of our future plans.  Due to the very fact that stochastic differential equations are employed to generate the dataset and to train ODEs, we expect the method itself to be robust against stochasticity. In the mixed-species case, linear relationship between the true vs estimation was robust (fig. 2F and 2J).  We do however see consistent under-estimation, albeit slight, of the interaction term of \phi (F^{(2)}_\phi) which presumably is caused by the stochasticity that is added to d\phi/dt.
Re – Question 3 (Advantage to use LAMB optimizer in comparison with AdamW):
Our module has a layer where the output of the fully connected network is multiplied by a trainable scalar.
When using Adam, the scalar behaved erratically compared with weights and biases in fully connected layers. The ability of LAMB to automatically adjust the learning rate between layers helped us stably train the module. We are aware that there is also an option to set a different learning rate for the scalar using AdamW, which we did not test.
Re – Question 4 and Weakness 2 (Primary factors which limits the scalability of our approach to systems with stochastic motion equations and higher-order interactions):
Computational load is the primary factor when extending the present approach to stochastic systems.  With the current stand-alone off-the shelf GPU in the market (NVIDIA A6000), it takes weeks to train ODEs.  Since the bottleneck seems to be at neural ODE calculation, performing bayesian inference to fit SDEs would take an unrealistically long time.  Some form of parallelization and use of super computers may help resolve this issue in the future.
As for higher-order interactions, the limitation is the lack of a good library to deal with hypergraphs, at least to the best of our knowledge. The other option is to use message passing twice, which we have not tried, because of the caveat of making the training and interpretation of neural networks more difficult.
Re – Weakness 1 (High computational cost and long training times):
Since our target of estimation is pairwise force which would work as a total force to drive the entities, we are essentially dealing with an inverse problem. The costs are necessary for stable and accurate estimation of each pairwise force against noise.  Given the nature of the problem, we believe costs of memory and time for estimation is a difficult issue to resolve.
Still, as we mentioned in the last section, we agree that it is an important avenue to explore in future research.
I thank the reviewers for their detailed response. I increase my score to 7.
1. The paper describes a very interesting problem which (to my knowledge) hasn’t been extensively explored through learning based approaches
2. It uses state-of-the-art message passing gnns and Neural ODEs as part of the
3. The model is accurately replicating simulation behaviour
4. Code is available
1. The paper is really hard to read. It took me a very long time to understand how it was implemented and although the code was generously submitted, this isn’t commented making it difficult to follow in the time available. Gnns and ODEs are in the title of the paper but are barely mentioned in the text.
2. The methodological contribution is limited
3. The model is trained in a supervised way meaning that it cannot learn new behaviours or extrapolate beyond the training regime. As this was stated as the main limitation of previous works it suggests this isn’t adding anything beyond these?
4. Results are based on validating relative to simple simulations - and run time is apparently anyway slow, so what is the value of the proposed approach
1. The main contribution seems to be to use a fully connected network to learn the force terms through a data term that forces the model to match the simulation results. Why not use a Physics informed neural network to embed the physics into the model, and thus hopefully extrapolate beyond the training regime?
2. Wouldn’t it make more sense to model this as a graph ODE? http://arxiv.org/pdf/1911.07532
3. It would help future readers if the authors explain in more detail how the whole network is trained - including where the gnn and neural ODE are called (it took me a long time to work this out as they are only briefly mentioned in the text)
4. This paper seems like it might be relevant and seems to be going beyond what is proposed here through learning models of interactions https://arxiv.org/abs/2303.09906 with PiNNs - can the authors explain the advs/disadvs of their work relative to an approach such as this
Limitations are not really discussed beyond runtimes
No ethics review needed.
Thank you for the feedback and comments. Your careful reading was very helpful in identifying places that we missed to explain in the draft.
Re – Question 1 and Weakness 2 (Why not use a Physics informed neural network):
This was briefly described in the Background section, however it was not fully conveyed in relation to Physics-Informed-Neural-Networks (PINN).
While there is no strict definition of PINN, we believe our model absolutely falls into that category, since we build our architecture on equations of motion for active matter (a class of coupled equations of motion important in non-equilibrium and biological systems) with forces being represented by neural networks. Most works on PINN centers on the estimation of relatively simple close to equilibrium systems such as diffusion and fluid dynamics where semi or almost non-supervised learning is partially possible based on first principles.
Instead, actively propelled particles as such studied here are far from equilibrium and Hamiltonian is not available. It is that we are exploring this difficult class of systems which is nevertheless fundamental in biological physics that makes the study novel and important.
Re – Question 2 & Weakness 2 (Relationship with Graph ODEs):
Our model is in fact a graph ODE except that the graph is time evolving due to change in particle positions. Valid edges were detected at every timepoint based on particle-to-particle distance. Existing graph ODEs instead have explicitly defined (i.e. fixed or semi-fixed) graphs. We realized that this wasn’t fully described and would be happy to clarify in the final edition.
Re – Question 3 & Weakness 1 (Clarify the details of training):
We recognize now that the section that describes this (p.4 L. 130-136) should have been clearer.
Apart from adding more comments in the codes, we propose to reorganize and extend the Method section, and add a flow chart in the final version. The reorganized section would make the following clear:
- Both for generating the training data with predetermined forces and the learning approach with neural networks, we solve a differential equation, eq. (1) in the original draft, over time. Depending on whether there is noise in the system, we employed either neural ODE or neural SDE. By neural ODE we simply mean that we solve an ODE in which the forces are given by neural networks.
In either case, the r.h.s. of eq. (1) has to be evaluated at least once per time step.
- To generate the training data, predetermined force functions were simply evaluated. During the learning process, the forces are represented by neural networks. Graph neural networks (GNN) techniques were employed to make this evaluation efficient. We note that the state of the system at any given time can be expressed by a graph, with each entity representing a vertex, and interactions between pairs of entities representing edges. As the positions evolve over time, the graph is time-dependent, too.
- Whenever the r.h.s. of eq. (1) needs to be evaluated by the solver, we pass the system state to the GNN wrapper. The GNN wrapper defines the graph structure and performs a message passing process, which evaluates the neural network (or predetermined function to generate training data) of force F^{(1)} for each vertex, evaluates the neural network for the force F^{(2)} for each edge, and finally adds the forces on the edges to each connecting vertex.
- Loss function: In order to infer the forces F^{(1)} and F^{(2)}, we split training data into pairs of states which are a time tau apart.
Then, for each pair of states, we integrate the equation of motion from the first state for a time tau using the neural network forces and then compare the resulting state to the second state. This exploits that the noise in the training data has the tendency to partially cancel out over time.
Re – Question 4 (advantages and disadvantages of our method in comparison with arxiv:2303.09906)
The preprint addresses learning the dynamics of an order parameter of alignment of agents instead of the underlying full trajectory of the interacting agents. The order parameter encodes some of the information about the collective dynamics of the agents, since it is an average over the agent’s states at a given time, but it does not come close to capturing their full dynamics. Their approach infers the dynamics of a one-dimensional quantity, whereas we infer the full dynamics of all the agents in the system. Their approach is useful in situations where one is only interested in understanding the dynamics of a mesoscopic or macroscopic quantity such as an order parameter and where having to keep track of the complete dynamics of the system is a hindrance, e.g.
in terms of data storage or computational cost. Their approach however cannot be applied to inferring the microscopic interactions between the agents (because they are intentionally averaged out), which is the aim of this study.
Re – Weaknesses 3 & 4 (Unable to extrapolate as other previous works due to supervised approach. The value of the proposed approach):
This is related to Question1. We should note that in complex systems where first principle is not easily available, an unsupervised approach is not feasible. To the best of our knowledge, the only studies of learning the interaction of actively propelled particles is on the Viscek model (cited refs [4, 24]) which still has an order parameter that can be exploited. The aim of our work instead is to go beyond such system where no clear order parameter is given. From application point of view, too, the interaction term in Viscek model leaves out crucial aspects of interaction between biological cells and thus we turned to the current model equation. Whether learning such models is even possible was not at all trivial and this work demonstrates its first realization and thus novel. We realized that this aspect was not sufficiently explained and will be amended in the corrected edition.
I've read your response and thank you for your attempt to qualify but this has raised more questions. For sure, although I have some basic understanding of PINNs its not my specialism and I do not know well the systems you are modelling -  so I have reduced my confidence in my review.
What I was wondering was whether the equations with which you use to define your simulations could not be used to generate a physics based loss for your ode - where my understanding is taken from this paper (which does not model interaction terms):
Linka, K. et al. Bayesian Physics-Informed Neural Networks for real-world nonlinear dynamical systems. arXiv [cs.LG] (2022).
I ask this because from what I understand from your paper section 5.2 (last paragraph) the model is overfitting on the interaction forces, and my limited understanding is that adding some form of physics loss can help regularise and allow models to extrapolate beyong the training regime?
As it stands the current results perform well for simulations - but for these the model is known, and a neural network is a universal approximator, so I assume it is not entirely surprising that it can learn the function given enough data?
In general, I'm not clear what the longterm projection of this work is? Can the authors comment on how this would be translated to real world data for which a model doesn't exist and what the potential advantage of that would be? How would they extrapolate beyond the training regime? How would they validate that the model works?
I'm still not clear why the graph neural ODE paper (Poli et al 2019) is not relevant here - I admit to not having studied this paper in depth but they show results on multi agent trajectory extrapolation and discuss a spatio-temporally evolving graph.
We appreciate the reviewer's comments and the thoughtful questions raised. Below, we provide a detailed response to address the concerns highlighted.
>whether the equations in your simulations could not be used to generate a physics based loss for your ode
The paper by Linka et al aimed at learning daily new COVID-19 cases (\hat{x}). Since it was evident that the data oscillates in time, the authors stabilized the prediction of \hat{x} by constraining a neural network (NN) model (which takes time t as input and outputs x) to exhibit a similar behavior to a harmonic oscillator model (where f\{x\} := \ddot{x} + (c/m)\dot{x} + (k/m)x = 0 for x).
Specifically, the constraint was applied by adding the term f\{x\}^2 to the loss function. Using this approach, the authors successfully predicted future time series which is not included in the training data.
In contrast, our work aims at studying a system where we can not use a term like f\{x\}^2. The very purpose of the learning process is to discover the model we used to create the training data. We used one of the models that has been proposed to explain the collective behavior of cells.
However, this model does not fully account for the characteristics observed in the experimental data we aim to analyze, such as phase separation-like behavior between different cell types.
Additionally, when considering other systems, both the collective behavior as a whole and the individual cell behavior differ significantly between systems, lacking any unifying pattern.
So far, no universal model nor principle governing these systems have been found. Given this background, we avoided assuming a specific model in the training process.
> what I understand from your paper section 5.2 (last paragraph) the model is overfitting on the interaction forces, … adding some form of physics loss can help regularise and allow models to extrapolate …
The last sentence of Section 5.2 may have caused confusion.  There was no overfitting but rather, it was underestimation described here. We will more clarify this point in the final version.
We should note that the extrapolation performed in the Linka et al paper is indeed achieved in our study. In Section 5, This is demonstrated by our model’s ability to reproduce the dynamics of collective motion from initial states not included in the training data.
>a neural network is a universal approximator, so … it is not entirely surprising
While the universal approximation theorem guarantees the existence of a solution where a NN can approximate any function, this does not imply that such a solution can be found, nor that the solution will generalize well to similar data. Our case is an inverse problem where each individual in a group is influenced by forces from numerous neighbors. The challenge is to estimate individual forces from the integrated trajectories observed, which are the result of summing and integrating those forces. It is not guaranteed that such forces can always be computed, and even if the computation converges to a solution that explains the trajectory, due to the nature of inverse problems, this does not ensure that the individual forces before summation are correct.
However, we have confirmed that our computations do converge, and that the individual forces are accurately estimated.
>what the longterm projection of this work is?
In biological physics, many including us are attempting to  address various types of collective motion exhibited by a group of groups of cells. Understanding such dynamics is fundamental to knowing how tissues form. Our present approach should be applicable in extracting individual cell behavioral rules from real data.
>How would they validate that the model works?
This is related to the above comments. Our plan is to test this approach by applying real experimental data and from there, make corrections and improve the assumed model.
>why is the graph neural ODE paper (Poli et al 2019) not relevant?
The method by Poli et al is relevant to the present work. We will clarify the difference between our method and theirs in the camera-ready version.
Poli et al addresses systems where the graph structure changes much slower than ours. Their extrapolation is limited to a few steps during which the edge structure remains almost unchanged. If the prediction extends beyond a few steps, the predicted node states deviate from the correct values, leading to inconsistencies with the edge structure. To perform long-term predictions with their method, one would need to compute the neural ODE over short intervals, updating the edge structure based on the resulting node states, and then repeat the neural ODE computation for another short interval. This approach is inconvenient and assumes that the edge structure remains unchanged during the neural ODE computation.
In contrast, our method does not require pre-defined edges, allowing one to make long-term predictions based solely on the initial state of nodes.
After rebuttal I have raised score to 5, and my confidence down. It would probably still be more convincing if the method was validated on experimental data. I'm still of the opinion that imposing a physics loss as regularisation doesn't mean that the model is forced to follow it exactly.
We thank reviewers for careful consideration of our work and are thankful for their astute criticism, which was very helpful in vastly improving the quality of the manuscript.  The manuscript was reviewed by 4 reviewers, with two of them (Shpk and Sm49) recommending acceptance. The other two, KAZi and T7k9, expressed disapproval out of concerns of novelty of the work which we believe are mainly due to the relatively short description provided in our original draft on the research background and the methodology.
We have provided a point-by-point response to the reviewers’ comments which we incorporate in an updated draft.  We feel strongly that we have been able to address the concerns of the referees.
In particular, we would like to briefly point out the following modifications:
- Referees KAZi and T7k9 both asked about the relationship to graph ODE.  We clarify the relationship of our method with pre-existing approaches to combine GNNs and neural ODE. Our new approach does not require any given graph structure, but defines it using a given rule at every time step of neural ODE calculation. This enabled us to extrapolate the trained model to unknown initial conditions, even in case the adjacency drastically changes in time.
- Describe Methods more clearly.
In particular, how Graph Neural Networks (GNNs)  are used to evaluate the equations of motion effectively. We added a flow chart which describes how our model evaluates the equations of motion (attached PDF file).
- Describe Physics Informed Neural Network and where the present work sits relative to related literatures.
- Referee KAZi expressed difficulty in readability while referee T7k9 found the writing to be rather strong.  Grammar and expression will be checked and corrected.

Meta-review:
role: PC
data: comment: value: The paper studies the problem of learning models for collective behavior from data. While the initial reviews were quite mixed, the rebuttal and discussion phase lead to some convergence in terms of understanding, clarity and the overall score (including a score increase that is clearly reflected in the comments, but not entered correctly). From my point of view, this is a quite interesting paper, and while there are certainly some edges to be ironed out in terms of presentation and clarity (see reviewer comments), I am recommending acceptance at this point. I am also strongly encouraging the authors to take all comments seriously and adjust the manuscript accordingly.

Decision: Accept (poster)