LLM + RoboticsUnit 4 · Specialised architectures · Day 21

Neural Message Passing for Quantum Chemistry

Most graph neural networks are one algorithm in different clothes: each node collects messages from its neighbours, updates its state, repeats, and then the graph is pooled into an answer.

AuthorsJustin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, George E. Dahl (Google Brain, Google, DeepMind)

Year / venue2017 · ICML

Core ideaMessage, update, readout: one framework for graph nets

Study time3–4 h

Original: arXiv:1704.01212 · Code, notebook and exercises: 30u30 Day 21

Start here (no maths needed)

Many computer programs for networks of things work the same way: each thing swaps messages with its neighbours, updates what it knows, and repeats.

Think of it like this

Think of news spreading through a group of friends. In each round, everyone tells their direct friends what they know. After one round, only your friends have heard. After two rounds, friends of friends have heard too. News travels one step per round. If there's a group chat everyone is in, it reaches everyone much faster. This paper's programs pass messages the same way, between atoms in a molecule.

How it works

Draw the things as dots and their connections as lines. For a molecule, dots are atoms and lines are bonds.
Each dot sends a message to the dots it is connected to.
Each dot reads its messages and updates what it knows.
Repeat for a few rounds. Each round, information travels one more step.
Finally, combine what every dot knows into one answer, such as how the molecule behaves.

Why robots care

A team of delivery robots can only talk to the robots close by. If each one passes on what it hears, a warning like "the hallway is blocked" spreads through the whole team, one robot-hop per round.

Try it

With friends in a line, whisper a fact to the first person. Each round, every person who knows it tells the person next to them. Count how many rounds it takes to reach the end.
In the demo further down, click a dot to make it the starting point. Move the Rounds T slider to see how far the news spreads. Then tick Add master node, a hub connected to every dot, and see how much faster it reaches everyone.

Words to know

Graph: A set of dots (nodes) joined by lines (edges).
Node: One dot in the graph, like an atom or a robot.
Message: The information one node sends to a neighbour.
Round: One turn where every node sends and receives messages.
Master node: An extra dot connected to all the others, like a group chat.

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Summary

TL;DR

By 2017 there were many neural networks for graphs, each with its own notation. This paper shows at least eight of them fit one template, the Message Passing Neural Network (MPNN), defined by three learned functions: message, update and readout. The authors then searched inside that template for the best model for predicting 13 quantum-chemical properties of small molecules (the QM9 dataset). These normally come from DFT simulations that take about an hour per molecule. Their best MPNN set the state of the art on all 13 targets and reached the chemists' "chemical accuracy" bar on 11 of them.

Template: T rounds of message passing along edges, then a readout that doesn't depend on node order.
Best variant: an edge network message function (edge features become a matrix), a GRU update, a set2set readout, and hydrogens as explicit nodes.
Numbers: average error ratio 0.68 across targets for one model and 0.52 for a 5-model ensemble, vs 1.35 for the best earlier method (1.0 = chemical accuracy).
Long range matters: without 3D coordinates, adding a global "master node" or a set2set readout cut the average error ratio from 3.47 to about 2.6.
Speed: inference is roughly 300,000× faster than the DFT calculation it approximates.

2. The core idea

Each node v has a hidden state h_v, starting from its features (atom type, charge and so on). Each edge has features e_vw (bond type, distance). For t = 0 … T−1:

m_v^{t+1} = Σ_{w ∈ N(v)} M_t(h_v^t, h_w^t, e_vw)     # gather messages
h_v^{t+1} = U_t(h_v^t, m_v^{t+1})                    # update state
ŷ = R({ h_v^T : v ∈ G })                             # read out the graph

N(v): the neighbours of v. Messages only travel along edges.
M_t: the message function. What w tells v, given both states and the edge between them.
Σ: summing over neighbours, so the order of neighbours doesn't matter and any degree works.
U_t: the update function, here a GRU. Weights are tied across rounds.
R: the readout. It must give the same answer for any ordering of the nodes.

The paper's best message function is the edge network:

M(h_v, h_w, e_vw) = A(e_vw) · h_w

A is a small neural network that turns the edge feature vector into a d × d matrix. So a short double bond and a long single bond transform the neighbour's state differently, and the edge features can be continuous (like distance), not just a few categories.

3. Deep dive

One template, many models

The paper rewrites earlier models as choices of M, U and R. Gated Graph Neural Networks use a learned matrix per edge type and a GRU update. The Kipf and Welling graph convolution uses a degree-normalised neighbour sum and a ReLU update. Interaction Networks, built for physics, use an MLP on the pair of states plus the edge. Once you see the template, a new graph paper is mostly a question of which three functions it picked. The Relation Network from Day 19 fits too: think of it as one round on a fully connected graph with a summed readout.

Information moves one hop per round

After T rounds, a node only knows about nodes up to T edges away. In a long molecule, atoms at opposite ends may never hear from each other. The paper tested two fixes. Virtual edges connect non-bonded pairs with a special edge type. A master node is connected to every node, and acts as shared scratch space that all nodes read and write each round. With it, any two nodes are at most two hops apart. The demo shows the difference. In training, T was searched between 3 and 8.

Interactive: how far does information spread?

Click a node to make it the source. Filled nodes have received its information after T rounds. Tick the master node to add a global hub.

Rounds T: Add master node

Nodes reached

–

Rounds to reach all

–

The readout and the set2set trick

A plain sum of final node states works, but it can wash out detail. The best model used set2set, an attention-based pooling over sets from Day 16. It reads the node states several times with a small LSTM and attention, and stays order-invariant. Without spatial inputs, set2set alone brought the average error ratio from 3.47 to 2.57.

Making it cheaper: towers

One round on a dense graph costs O(n²d²) for n nodes with d-dimensional states. The towers variant splits each state into k chunks of size d/k, runs message passing on each chunk separately, then mixes them. That cuts cost to roughly O(n²d²/k). With k = 8, n = 9, d = 200, inference was about 2× faster, and accuracy improved slightly too.

4. Why it matters for robotics

Learned physics simulators. Interaction Networks, one of the models this paper unifies, were built to predict how objects push and pull each other. Later learned simulators for particles, cloth and rigid bodies use exactly this message-update loop, with nodes as particles or mesh points and edges as contacts.
Multi-robot coordination. Put robots on nodes and radio links on edges. One round of message passing is one exchange with neighbours, so a policy with T rounds uses information from T hops away. That matches how decentralised teams actually communicate.
Bodies and scenes are graphs. A robot's links and joints, or the objects in a scene, form natural graphs. Graph policies handle a different number of legs or objects without retraining the input layer.
Fast surrogates for slow computation. Here an MPNN replaced an hour of DFT with a fast prediction. In robotics the same move replaces a slow simulator or contact model with a learned one, but only inside the range it was trained on. The authors themselves flag generalising to larger graphs as unsolved.

5. How to study it

Read section 2 (the framework) and section 5 (the variants) of the paper, about 1.5 h. For each model in section 2, write down its M, U and R.
In the 30u30 Day 21 folder, read implementation.py: find the message, update and readout pieces and match them to the equations above.
Run train_minimal.py on the synthetic molecular data. Change the number of message rounds and watch the per-property error.
Do the exercises: message functions, graph construction and batching, readouts, edge networks, then multi-target property prediction.
Stretch: add a master node to the implementation and compare on a long chain-shaped graph.

6. Self-check

What are the three functions that define an MPNN?

The message function M, the node update function U, and the graph-level readout R.

Why must the readout be invariant to node order?

Node numbering is arbitrary. The same molecule written with atoms in a different order must give the same prediction.

After 3 rounds, what can a node know about?

Only nodes within 3 edges of it. Information travels one hop per round.

What does the edge network add over one matrix per bond type?

It maps continuous edge features, like distance, to a transformation matrix. So messages can depend smoothly on geometry, not just on a small set of bond categories.

Why does a master node help without 3D coordinates?

It connects every node through one hub, so any two nodes are at most two hops apart. Long-range effects that the bond graph alone can't carry in a few rounds get through.

7. Further reading

The original paper, with the supplementary material deriving each earlier model as an MPNN
30u30 Day 21: implementation, notebook, cheatsheet, exercises
Next: Day 22, Deep Speech 2