google-deepmind / google-deepmind/deepmind-research

MeshGraphNet: Residual Connection Mismatch - Node Update Uses Pre-Residual Edge Features Instead of Updated Ones

Open
#609 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
15.2k
Forks
2.9k
PR merge metrics
No merged PRs in 30d

Description

According to the paper, a block of the processor should perform the following updates:

$e'_{ij} \gets f(e_{ij}, v_i, v_j)$
$v'_i \gets v_i + f(v_i, \sum_j e'_{ij}) $

where \( f \) is an MLP with a residual connection (and I have omitted the distinction between world and mesh edges for simplicity).

This means that the model should update as follows:

$e'_{ij} \gets e_{ij} + \text{MLP}(e_{ij}, v_i, v_j)$
$v'_i \gets v_i + \text{MLP}(v_i, \sum_j e'_{ij})$

However, the current implementation instead does:

$e^{\text{update}}_{ij} = \text{MLP}(e_{ij}, v_i, v_j)$
$v'_i \gets v_i + \text{MLP}(v_i, \sum_j e^{\text{update}}_{ij})$
$e'_{ij} \gets e_{ij} + e^{\text{update}}_{ij}$

That is, the node processor MLP is fed the **difference between the updated and old edge representation** rather than the **updated edge representation** itself, which is what is described in the paper.

### Proposed Fix

The issue can be resolved by modifying the `_build` method in `core_model.py` as follows:

```python
def _build(self, graph):
"""Applies GraphNetBlock and returns updated MultiGraph per MeshGraphNets paper."""

# Apply edge functions with immediate residual connections
new_edge_sets = []
for edge_set in graph.edge_sets:
# Compute edge update
edge_update = self._update_edge_features(graph.node_features, edge_set)
# Apply residual connection immediately
updated_features = edge_set.features + edge_update
new_edge_sets.append(edge_set._replace(features=updated_features))

# Apply node function with residual connection
node_update = self._update_node_features(graph.node_features, new_edge_sets)
# Apply residual connection to node features
new_node_features = graph.node_features + node_update

return MultiGraph(new_node_features, new_edge_sets)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.