google-deepmind / google-deepmind/deepmind-research
MeshGraphNet: Residual Connection Mismatch - Node Update Uses Pre-Residual Edge Features Instead of Updated Ones
- 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
Assessment
This issue has not been assessed yet.