Learning Physics with Graph Neural Networks: Part 2

In Part 1

[embed]Learning Physics with Graph Neural Networks: Part 1 I am fascinated by a different kind of machine learning architecture: Graph Neural Networks, or simply GNNs. As a…medium.com

we built a mental model of Graph Neural Networks: graphs, message passing, and the three-step loop, compute messages, aggregate, update. We even wrote a minimal implementation in PyTorch.

Now we put it to work. Can a GNN learn to simulate a physical system, specifically, heat diffusing through a metal rod?

We’ll build a complete pipeline: a classical solver for ground truth, a GNN architecture, a training loop, and a 3500-step rollout test. The result is both encouraging and sobering. The GNN nails single-step predictions, but when we let it run freely, it collapses. Understanding why is the point of this article.

Before we build anything, it’s worth asking: why use a GNN for a physics problem? Why not a standard feedforward network, or a CNN?

The answer is that GNNs encode three inductive biases that happen to mirror how physical systems actually work:

1. Pairwise interactions. In a GNN, the state of node i evolves based on its interactions with individual neighbors j. This is exactly how physics works, forces act between pairs of particles. Gravity between two masses, heat flow between two adjacent points, stress between connected elements in a mesh.

2. Shared rules. The same message and update functions are applied at every node and every edge. This mirrors the universality of physical laws the heat equation doesn’t change depending on where you are on the rod. Gravity doesn’t have different formulas at different locations.

3. Locality. Each node is directly influenced only by its neighbors. Just as a point on a metal rod is heated or cooled by the points immediately next to it (not by a point two meters away), a GNN node aggregates messages from its local neighborhood.

These aren’t just nice analogies. Consider Newton’s second law applied to a network of interacting particles and the GNN message passing:

Physics vs GNN Message passing

The rate of change of particle i state is the sum of pairwise interactions with its neighbors. Compare this to the GNN message-passing update. The structure is identical. A GNN already speaks the language of physics.

So let’s give it a physics problem to solve. We’ll start with one of the simplest and most beautiful equations in physics: the heat equation.

The Heat Equation

The heat equation describes how temperature diffuses through a material. Imagine a thin metal rod: you heat it in the middle, clamp both ends to ice (zero temperature), and watch the heat spread and dissipate.

Mathematically:

Why this problem? Because it’s one of the easiest PDE (Partial Differential Equation) to learn:

Pure diffusion, no turbulence, no shocks

• The solution is always smooth
• We have an exact numerical solver to compare against
• If a GNN fails here, the failure is in the learning pipeline, not the physics complexity

Ground Truth: Crank-Nicolson

To train our GNN, we need ground truth data. We generate it using the Crank-Nicolson scheme, a classical finite difference method that is unconditionally stable and second-order accurate in both space and time.

The idea is simple: discretize the rod into N = 50 interior points, and advance the solution one time step at a time. At each step, we solve a small linear system that averages the explicit and implicit Euler methods.

We won’t derive the scheme here (it’s standard numerical methods), but the key point is: Crank-Nicolson gives us near-exact solutions that we trust completely. It’s our ground truth and our benchmark.

# Full code: https://github.com/blockskode/medium/tree/main/gnn/heat_equation
from heat_equation.crank_nicolson import CrankNicolson1D

solver = CrankNicolson1D(N=50, L=1.0, alpha=0.01, dt=0.001)
u = solver.initial_condition(gaussian_bump)  # starting temperature profile
for step in range(3500):
    u = solver.step(u)  # advance one time step

We now have our ground truth solver and a clear picture of what the solution looks like. The next question is: how do we feed this problem to a GNN? The heat equation lives on a continuous rod, but a GNN operates on a graph. We need to bridge the two.

From Partial Differential Equation to Graph

We need to represent our 1D rod as a graph.

The 1D Chain

The discretized rod has 50 interior points. Each point becomes a node. Each pair of adjacent points gets a bidirectional edge, node i connects to node i-1 and node i+1, just like the spatial stencil of the heat equation.

def build_1d_graph(N):
    """Build edge_index for a 1D chain of N interior nodes."""
    src, dst = [], []
    for i in range(N - 1):
        src.append(i);     dst.append(i + 1)   # i -> i+1
        src.append(i + 1); dst.append(i)        # i+1 -> i
    return torch.tensor([src, dst], dtype=torch.long)

Handling the Boundaries

How should we handle the boundaries? The Dirichlet conditions fix both endpoints at u = 0 for all time and there’s nothing to predict there. But the GNN still needs to know the boundaries exist, because they drive the physics: heat flows out through them. There are several options:

Option A: Keep boundary nodes, freeze their values. Include the boundary nodes in the graph with u = 0, let them send messages to their neighbors at every step, but never update them. The GNN directly “sees” the cold wall through message passing , node 1 receives a message from a neighbor stuck at zero. The boundary influence is explicit.

Option B: Remove boundary nodes, add positional features. Only include the 50 interior nodes. Each node receives its normalized x-coordinate as an extra input, so the GNN can learn that nodes near pos = 0 or pos = 1 are close to cold walls. The boundary influence is implicit, the GNN must learn what position means.

Option C: Encode boundary proximity as edge features. Mark edges near boundaries with a special flag or distance value.

We go with Option B. It’s the simplest to implement: we just remove the boundaries and give each node its position as an extra input.

The GNN then has to figure out on its own that nodes with position close to 0 or 1 are near cold walls. This might sound like a lot to ask, but it works. The network picks this up from the training data without any trouble.

The GNN Architecture

Our GNN follows the same message-passing pattern from Part 1, with three stages:

1. Encoder: A small Multi-Layer Perceptron (2 linear layers with ReLU activation) that maps the 2D input (temperature + position) into a 64-dimensional hidden representation 2. Message passing: 4 message-passing layers, each containing two Multi-Layer Perceptrons: one for computing messages, one for updating node states, with residual connection 3. Decoder: A small Multi-Layer Perceptron (2 linear layers with ReLU activation) that maps the hidden representation back to a single scalar, the predicted temperature change

#  GitHub: https://github.com/blockskode/medium/tree/main/gnn/heat_equation
class HeatGNN(nn.Module):
    def __init__(self, hidden_dim=64, n_layers=4, input_dim=2):
        super().__init__()
        # Encoder: 2D input -> hidden space
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )
        # 4 message-passing layers
        self.mp_layers = nn.ModuleList([
            MessagePassingLayer(hidden_dim, hidden_dim)
            for _ in range(n_layers)
        ])
        # Decoder: hidden space -> scalar delta
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
        )

Why residual prediction?

The model predicts the change in temperature, not the next temperature directly:

The temperature changes per time step are small (on the order of 10^-4 to 10^-3). Predicting the full next-state value would require the network to first learn an almost-perfect identity mapping, then add a tiny correction. Predicting just the correction directly is much easier.

Residual connections in message passing

Each MP ( Massage Passing) layer adds a correction rather than computing from scratch:

def forward(self, u, edge_index):
        x = torch.cat([u, self.pos_feat], dim=-1)  # (N, 2)
        h = self.encoder(x)                          # (N, 64)

        for mp_layer in self.mp_layers:
            h = h + mp_layer(h, edge_index)          # residual

        delta = self.decoder(h).squeeze(-1)          # (N,)
        u_next = u.squeeze(-1) + delta               # residual prediction
        return u_next

Training and Generating Data

The GNN learns entirely from data produced by the Crank-Nicolson solver. We generate 100 trajectories, each running for 50 time steps, from diverse random initial conditions:

• Gaussian bumps: a bell curve with random center, width, and amplitude
• Sine modes: a random sum of 1–3 Fourier modes
• Smoothed steps: a sharp transition smoothed by a hyperbolic tangent

This diversity forces the GNN to learn the general dynamics of heat diffusion, not just memorize one shape. From each trajectory, we extract consecutive pairs (ut, u{t+1}) as training samples. 100 trajectories x 50 steps = 5000 training pairs.

Training Setup

Training Curve

Training converges quickly and the loss drops by two orders of magnitude in the first 2 epochs and plateaus by epoch 3. Final training loss: ~2.7 x 10^-5. The validation loss tracks closely, indicating no overfitting at this scale.

This looks great. The GNN has learned the one-step dynamics accurately. So let’s test it.

Results: The Catastrophic Failure

We test on a Gaussian bump that the GNN has never seen during training, and roll it out autoregressively for 3500 steps (t = 3.5 seconds). At each step, the GNN’s own prediction becomes the input for the next step, just like a real simulator would work.

GNN vs Crank-Nicolson:

The first few frames look perfect. By step 500, the GNN starts drifting. By step 1000, the prediction is qualitatively wrong. By step 3500, the GNN output has no resemblance to the true solution.

The Crank-Nicolson solver shows the heat smoothly dissipating toward zero, exactly what physics demands. The GNN instead produces a bloated, non-physical temperature profile that grows rather than shrinks.

Error Growth

The error growth is exponential before saturating. The MSE shoots from ~10^-15 at step 1 to ~10^-1 by step 1000. That’s twelve orders of magnitude of error accumulation.

Energy Drift

Here’s the most damning metric. The heat equation with Dirichlet boundary conditions must lose energy over time, heat flows out through the cold boundaries. It is physically impossible for the total energy to increase.

The classical solver correctly shows energy decreasing from 10.23 to 8.78. The GNN’s energy quadruples to 41.43. The model is literally creating heat out of nothing, which is a physically impossible result.

Why It Fails

Three observations:

1. Single-step prediction is excellent. At step 1, the MSE is ~10^-7. The GNN has learned the one-step dynamics to high accuracy. This is not a capacity or training problem.

2. Errors compound catastrophically during rollout. This is the distribution shift problem. During training, the GNN sees ground truth inputs. During inference, it sees its own predictions. Each small prediction error shifts the input distribution slightly, leading to a slightly larger error at the next step, which shifts the distribution further, and so on. Over 3500 steps, this unchecked compounding destroys the solution entirely.

3. The GNN violates energy dissipation. The heat equation can only **lose** energy through the boundaries and never gain it. But the GNN has no mechanism to enforce this. Nothing in the architecture prevents it from predicting temperature increases that violate conservation laws. And once it starts creating energy, the distribution shift accelerates: the inputs become increasingly out-of-distribution, and the errors compound faster.

What’s Next

These results are not a failure of the GNN concept, they’re a failure of the naive approach. The GNN learned the one-step dynamics well. What it lacks are the physical guarantees that prevent errors from compounding: energy must decrease, solutions must stay bounded, trajectories must remain stable over thousands of steps

The standard response would be: train longer, add more data, tune hyperparameters. But hope is not a design strategy (Did I just said that?). No amount of training will guarantee that energy decreases, or that errors don’t compound, or that the solution stays smooth.

In the next article, we’ll take a different approach. Instead of hoping the GNN learns the right physics, we’ll build the physics into the architecture, so that violations become structurally impossible, not just unlikely.

References

• Crank, J. & Nicolson, P. (1947). A practical method for numerical evaluation of solutions of partial differential equations of the heat-conduction type

• Sanchez-Gonzalez et al. (2020). Learning to Simulate Complex Physics with Graph Networks. arXiv:2002.09405

• Brandstetter et al. (2022). Message Passing Neural PDE Solvers. arXiv:2202.03376

• Battaglia et al. (2018). Relational inductive biases, deep learning, and graph neural networks. arXiv:1806.01261