8 Geometric Graphs

• In many cases, graphs are endowed with spatial geometry (e.g. 3D point
clouds). It can be highly useful to leverage this geometry, yet the graph
domains we studied in the Graphs chapter were not sufficient to do so.
• Combining permutation symmetries (S ) with continuous symmetries of
n
physical space (e.g. E(3)) yields a rich class of geometric graph models.
• Geometric graphs also arise through studying manifold domains, which

generalise the homogeneous domains we studied in the Groups chapter.

• Expressing local geometry in a manifold and discretising it yields mesh

convolutional networks, which are strongly relatable to geometric GNNs.

• We can also use this framework to elegantly study AlphaFold 2, the first

equivariant neural network to ever win the Nobel prize (in 2024).
Throughout our initial exploration of graph representation learning in Chapter
5, we assumed our graphs were entirely abstract. We operated on a domain
Ω = G = (V , E ) defined simply by a set of nodes, V , and edges E , and the only
requirement of our neural networks was equivariant behaviour with respect to
permutations of the n nodes (G = S n). As we will explore in this Chapter, this
is often far from the full story, from the geometric perspective.
The insufficiency of “pure” graphs becomes evident in natural systems,
wherein a graph may often be endowed with some spatial geometry within,
e.g., its node features x u ∈ Rk. Molecules were the initial motivating example
for the graph domain in Chapter 5, yet they are also a common first motivation
for going beyond purely-combinatorial views of graphs. While a molecule can
be represented as an abstract chemical graph of atoms and bonds, in reality,
these atoms do not just form an abstract set; they occupy 3D space in a way
that minimises potential energy, known as a conformer (Figure 8.1). To accu-
rately predict a molecule’s chemical reactivity, or how it might bind to a protein
target—all of high interest to drug development—neural networks might ben-
efit from explicitly understanding the spatial arrangement in such geometric
graphs. This will be the final domain we consider on our geometric journey.
210 Chapter 8

H H

H
H H

H
H H
C C
H H
C H

H H
H
C C H C C C
H C C H
H H H H
C H
H H
C
H
H H C
C C H H
H
C C H C
H
C
H
H
H H
H
C C H C C C
H H
H H

H H
H H
H H

Figure 8.1
Based on its molecular graph (left), cyclooctane (C 8 H16 ) is a seemingly very simple
molecule; eight carbon atoms with identical local structure. However, the actual energy-
minimising conformers of cyclooctane tell a very different story (Pakes, Rounds, and
Strauss 1981). Based on the specific energy-minimising positioning of atoms, this
molecule may appear in positions such as the crown (middle) or the twist boat-chair
(right). Note all of the conformations have identical graph structure, but the relative
coordinates of the atoms (which are encoded in the node features) are very different.

8.1 Domain
To exemplify and help us formalise the notion of geometric graphs, we will
utilise a graph of points in 3D space. Namely, assume our nodes u ∈ V are
endowed not only with abstract (scalar) features, su ∈ Rm , but also with spatial
coordinates, x u ∈ R3 . This is the simplest example of a geometric graph.
If we are to apply the Geometric Deep Learning blueprint to this new domain
in a way that respects this spatial structure, we must expand our assumed group
of symmetries G beyond Sn . The exact way in which we expand will depend
on the specific assumptions on which quantities are equivariant under which
transformations of the given spatial coordinates, xu .
For example, if we’d like to predict some property of a given molecule (such
as its toxicity or solubility), assuming no background effects or force fields,
that property should remain the same even if we rotate, translate or reflect all of
the molecule’s atoms. This collection of roto-translations and reflections in 3D
space is the Euclidean group1, E(3). For any 3D orthogonal matrix Q ∈ R3×3
and translation vector2 t ∈ R3, we define a corresponding group action g ∈ E(3)
operating on the coordinate space:
ρ(g)x u = Qxu + t (8.101)
Crucially, applying this spatial transformation to the coordinates typically
does not affect the abstract node features su (e.g., atomic mass, charge, or atom
Geometric Graphs 211

type in the context of molecular modelling), as they are scalar quantities. We

would expect any graph neural network operating over these feature to respect
this constraint. This allows us to explicitly define what an E(3)-invariant GNN
should satisfy. Let F(S, A, X) be a layer of such a GNN, consuming as input
a feature matrix S ∈ R n×m , adjacency matrix A ∈ Rn×n and coordinate matrix
X ∈ Rn×3 . Then, to be E(3)-invariant, such a layer must satisfy:
F(S, A, XQ + t) = F(S, A, X) (8.102)
3
for all possible (Q, t) ∈ E(3) – that is, its output should be unaffected by
roto-translations and reflections of the input. There is nothing special about
the 3-dimensionality of the input coordinates here—one may express E(d )-
invariant models using exactly the same constraints—simply changing the
dimensionality of Q ∈ R d×d and t ∈ R d accordingly. Further, note that this
is an additional constraint on top of existing permutation symmetries that
would make F a valid GNN layer. In practice, this means that E(d )-invariant
GNNs are actually constrained with respect to the product symmetry group,
G = S n × E(d ). We can combine these constraints into just one, specified for
every (P, (Q, t)) ∈ S n × E(d ), where P ∈ R n×n is a permutation matrix:
F(PS, PAP⊤ , PXQ + t) = PF(S, A, X) (8.103)
where we are both permuting and roto-translating/reflecting the input—
constraining the output to be equivariant with respect to the permutation, and
invariant with respect to the spatial transformation. Note we have already dis-
cussed product groups at length when covering subgraph GNNs in Chapter
5 (Bevilacqua et al. 2021), although in that context both of the groups being
combined were permutation groups.

8.1.1 Invariant geometric GNNs: SchNet and DimeNet

A typical simple way to satisfy this kind of constraint is to only supply inputs to
the model that would not change under any roto-translation or reflection 4 of the
coordinates—this trivially guarantees E(d ) invariance for any dimensionality
d. By our assumption, s u are such features, whereas we can feed in distance
information, ∥(x u – x v)∥, as a simple coordinate-derived invariant edge feature.
Plugging it into our usual message-passing formulation from Chapter 5, we
obtain the following GNN layer, also known as SchNet (Schütt et al. 2018):
!
M  
h u = ϕ s u, ψ su , sv , ∥xu – xv∥ (8.104)
v∈Nu

It is a simple exercise to prove that SchNet is both permutation-equivariant

and E(d )-invariant. While it offers a rudimentary way to incorporate spatial
212 Chapter 8

information, SchNet is still a powerful representation learning method with

many recorded downstream uses in molecular modelling.
It is possible to extend this invariant modelling idea to incorporate more
diverse geometric features—for example, we can consider the angle between
pairs of senders v, w ∈ Nu with respect to the receiver, u, expressed via the dot
product (x v – x u)⊤ (xw – xu). Incorporating this information requires explicitly
considering pairs of senders at once, modifying the GNN equation to:
 
M  
muv = ψ su, s v, ∥x u – xv ∥, χ sv , sw , ∥xu – xv∥, (xv – xu)⊤ (xw – xu) 
w∈Nu \{v}
(8.105)
!
M
h u = ϕ su, muv , (8.106)
v∈Nu

where χ is a learnable neural network architecture, just like ψ and ϕ. This

architecture is commonly known as DimeNet (Gasteiger, Groß, and Gün-
nemann 2020) and it remains state-of-the-art on many important molecular
modelling tasks, striking a balance between simplicity and respecting struc-
tural constraints. The reader might wonder, then, are there meaningful ways to
build geometric GNNs beyond simply providing them with invariant features?

8.1.2 Towards E(d) equivariance

As pairwise distances and relative angles already offer a wealth of information
about a geometric graph’s spatial layout, it might not be immediately obvi-
ous that they are insufficient to distinguish all such graphs. However, one can
indeed construct counterexamples where not only the node set is identical, but
so are the distances and angles (Bartók, Kondor, and Csányi 2013; Pozdnyakov
and Ceriotti 2022).
One such pair of graphs is illustrated in Figure 8.2: as there are no distin-
guishing nodes, distances or angles across them, models such as SchNet and
DimeNet have no hope of assigning different representations to them. This fol-
lows a similar argument to the one that aligns GNNs to the 1-WL test (Morris et
al. 2019), which we discussed in Chapter 5. The main difference is that now the
set of hashable entities also contains distances and angles, yielding a geometric
WL test (Joshi et al. 2023). We need to introduce additional spatial quanti-
ties to meaningfully distinguish such inputs. While the distance magnitudes,
∥xu – xv ∥, might not have been expressive enough, the actual raw values of xu
and x v would convey such information usefully. However, computing scalar
representations cannot directly make advantage of such information without
breaking the E(d)-invariance constraint, which would mean roto-translating a
Geometric Graphs 213

Figure 8.2
An example of two geometric graphs which cannot be distinguished by invariant models
such as SchNet, due to Pozdnyakov and Ceriotti (2022). Note that the set of nodes, their
colours, as well as edge distances and relative angles between them, are identical across
the two graphs—therefore, the E(d)-invariant GNNs studied thus far are insufficient for
distinguishing them.

molecule yields different final predictions. Since we cannot change how we

update the latents, hu , we need to find a way to update the representations of
our coordinates, x u, in a way that will not affect how the latents, h u, evolve.
Hence, our layer now needs to compute a pair of outputs:
F(S, A, X) = (H, X′) (8.107)
computing node latent features h u ∈ Rk, and updating the node coordinates to
x′u ∈ Rd. The latents still must obey E(d )-invariance, as we assume them to
represent scalar quantities. The coordinates, however, can transform together
with roto-translations and reflections, yielding an E(d ) equivariance constraint:
F(PS, PAP⊤ , PXQ + t) = (PH, PX′Q + t) (8.108)
n×n d×d
for all permutation matrices P ∈ R , orthogonal matrices Q ∈ R and
d
translation vectors t ∈ R (see Figure 8.3). The final-layer representations and
coordinates can be used to make equivariant predictions over the input geomet-
ric graphs. For dynamics modelling tasks, one might prefer to use the predicted
coordinates x ′u directly as the output—this is typical of n-body problems.
214 Chapter 8

s1
x1
s4 x4 s3
◦
x3 s2 ρ(180 )
x2 x2
x3
s2 x4
x1
s3 s4
s1

F F

h2
h3 x′2 h 1 x′1
x′3 x′4
◦
ρ(180 ) h4

h4 x ′4 x ′3
x ′1 x′2
h1 h 3
h2

Figure 8.3
An illustration of an E(2)-equivariance constraint, inspired by Satorras, Hoogeboom,
and Welling (2021). An E(2)-equivariant GNN, F(S, A, X) = H, X′ , updates both the
node representations (from s u ∈ R 2 to hu ∈ R3 ) and the node coordinates (from xu ∈ R2
to x ′u ∈ R2), in a way that is equivariant to actions from E(2). In this case, a specific
180◦ rotation of node coordinates is considered. As exemplified by the commutative
diagram, an E(2)-equivariant model must yield the same answer regardless of whether
the GNN or the rotation is applied first.

8.2 Model
An important question remains: how do we actually build graph neural net-
works that respect these kinds of continuous spatial equivariance conditions?
While the space of possible solutions to this equivariance condition is certainly
vast, there are generally two primary strategies for resolving them. Namely, we
may either focus on supplying specific equivariant quantities for the particular
input shapes provided, or we can attempt solving the constraint for all possible
spatial input shapes (relying on irreducible representations of E(d )).

8.2.1 E(d)-equivariant GNNs

For now, we continue with the previously stated assumption of our geometric
graph inputs comprising scalars, S ∈ Rn×m , and coordinates, X ∈ Rn×d . Under
this claim, we can design an elegant and scalable equivariant GNN (Satorras,
Hoogeboom, and Welling 2021) which explicitly supplies invariants relevant
Geometric Graphs 215

to these particular input types—such as pairwise distances—and recombines

directional information in simple ways—such as linear combinations.
As the scalar properties are expected to be invariant to roto-translations,
we can re-use E(d)-invariant models such as SchNet for updating the scalar
features su ∈ Rm into latents h u ∈ Rk . Matters get significantly more interesting
when extending this model to (equivariantly) update the coordinates.
To see how, first recall that our aim is for the updated coordinates to roto-
translate and/or reflect together and identically with their input coordinates.
This means that we cannot compute arbitrary new coordinates, and instead
must derive them from already-known ones. Further, the way in which they
are derived must commute with orthogonal matrix products (in order to handle
rotations and reflections) as well as translations. An operation that covers all
of the above is a linear combination5 of displacement vectors, x u – xv, as in:
X
x′u = xu + βuv (x u – xv ) (8.109)
v̸=u

where βuv ∈ R is a scalar coefficient, which we assume to be fixed upfront

for now. It’s a good idea to quickly reflect on why these specific opera-
tions satisfy the constraints in Equation 8.108. Firstly, linear combinations
of this kind will commute with matrix products (orthogonal or otherwise!):
P P
Q(xu + v̸=u βuv(x u – xv)) = Qxu + v̸=u β uv(Qxu – Qx v). And secondly, the
displacement operator will naturally cancel out any impacts of translation vec-
P P
tors: (xu + t) + v̸=u βuv((x u + t) – (xv + t)) = (x u + v̸=u βuv (xu – x v)) + t. The
choice to sum over all nodes v ̸= u is an arbitrary one which works well in
the context of modelling physics. In practice, aggregation over any specific
neighbourhoods Nu will also suffice.
Keeping βuv fixed, however, is an underwhelming choice—we haven’t made
advantage of the possible wealth of information stored in our nodes. Accord-
ingly, we can define an attentional mechanism neural network which will
leverage the s u and x u vectors to dynamically compute the coefficients of
the combination. However, this mechanism cannot take arbitrary inputs—for
example, if raw coordinates, x u, were passed to it directly, this neural network
could easily start performing multiplicative interactions between coordinates
that would not commute with orthogonal Q matrices. As such, we require the
coefficients to be robust to any roto-translations or reflections of xu—meaning
it needs to be E(d )-invariant by definition! We definitely know how to build
such networks—we can simply reuse the SchNet blueprint, for example.
216 Chapter 8

90°

Figure 8.4
While scalar features (heatmap) do not change under rotations, vector features (arrows)
may change direction. The simple E(d ) equivariant GNN given in Equations 8.110–
8.111 does not take this into account, and requires an explicit modification (Equation
8.114).

Putting it all together, E(d)-equivariant GNNs take the following form:

!
M  
hu = ϕ s u, ψ su , sv , ∥xu – xv ∥ (8.110)
v∈Nu
X  
x′u = xu + β su, sv , ∥xu – x v ∥ (xu – xv ) (8.111)
v̸=u

Note that ψ and β both take exactly the same, invariant, inputs, to guarantee
desirable behaviours. However, their output spaces are very different: k-
dimensional latent messages in the case of computing ψ , and scalar coefficients
in the case of β .

8.2.2 Vector-structured features and “tensoring up”

As we have seen, in contrast to ordinary GNNs, the model discussed thus far
enables the correct treatment of ‘coordinates’ for each node in the graph. They
are now treated as a member of the E(d ) group, which means the network
outputs behave correctly under rotations, reflections and translations of the
input. The node features, su , however, are treated in a channel-wise manner
and still assumed to be scalars that do not change under these transformations.
Geometric Graphs 217

This limits the type of spatial information that can be captured within our
framework. For example, different elements of the node features may encode
vectors—e.g. velocities or forces—which should potentially change direction
under such transformations (Figure 8.4).
This issue may be partially alleviated by explicitly introducing the concept
of vector node features, v u ∈ R d. Note that vectors transform together with any
rotation or reflection of coordinates but remain invariant to translations. This
justifies tracking them separately, making our GNN take the following form:
F(S, V, A, X) = (H, V ′, X′ ) (8.112)
with the explicit addition of the input vector feature matrix, V ∈ Rn×d , and
the output vector features v ′u ∈ Rd . The relevant equivariance condition also
changes appropriately:
F(PS, PVQ, PAP ⊤, PXQ + t) = (PH, PV ′ Q, PX′ Q + t) (8.113)
Much like in the case of coordinate updates, we will once again rely on lin-
ear combinations of vector features, vu, with invariant coefficients, to ensure
equivariance to orthogonal matrix products. A useful addition is that, as the
displacements, x u – xv , are themselves vector features, we may also include
them in the update rule. All together, we recover the following equation:
X   X  
v ′u = vu + α su , sv , ∥xu – xv ∥ v v + β s u, s v, ∥xu – xv ∥ (xu – xv)
v̸=u v̸=u
(8.114)
where α and β are neural networks outputting scalar values. The specific way
in which this equation is implemented varies across proposals, with notable
early examples including the already-discussed E(d )-equivariant GNNs of
Satorras, Hoogeboom, and Welling (2021), the concurrently published PaiNN
(Schütt, Unke, and Gastegger 2021), and GVP-GNNs (Jing et al. 2021).
While the vector extension remains elegant, it also reveals an issue with
building equivariant graph neural networks in this manner: every time a new
input tensor order is added to our tasks, a novel architecture extension needs
to be prepared, taking into account a new kind of equivariance condition. For
example, our inputs might include node matrix features6, Mu ∈ Rd×d , which
may arise as outer products M u = a ⊗ b of two vectors a, b ∈ R d. Such inputs
transform under E(d ) as QMu Q⊤ for every orthogonal matrix, Q, inducing
appropriate equivariance constraints 7. Keeping track of these constraints may
quickly get out of hand as we “tensor up” our inputs further.
Hence, while the architectures discussed so far already present capable
equivariant models for most practical geometric graph inputs, in some cases
it may be desirable to explore a broader collection of functions that satisfy the
218 Chapter 8

= ⊕ ⊕

Matrix trace (type-0) antisymmetric (type-1) traceless symmetric (type-2)

ℓ=0

ℓ=1

ℓ=2

Figure 8.5
Illustration of how an arbitrary matrix decomposes into its constituent type-ℓ vectors of
2ℓ + 1 dimensions.

E(d ) equivariance property in a principled manner that generalises to arbitrary

input tensor orders. We study a popular version of such an approach next.

8.2.3 *Learning on irreducible representations: Tensor Field Networks

and SE(3)-Transformers
For the content that follows, we make a deliberate choice to restrict our sym-
metry group from E(d) to SE(3): the group of 3D roto-translations. This
means that, for our group actions, we replace the previously-considered d-
dimensional orthogonal matrices, Q ∈ Rd×d , with 3D rotation matrices, R ∈
R 3×3 , and still including 3D translations, t ∈ R3 . This restriction is done
merely for pedagogical reasons of ease of tracking mathematical state, and
is not strictly required for building tractable architectures.

[Link] Type-ℓ vectors Since we would like to be able to reason about how
to update any possible spatial geometric graph input (irrespective of tensor
shape), it will be useful to formalise what all possible tensor shapes might look
like. To do this, we leverage the concept of type-ℓ vectors, for ℓ ∈ {0, 1, 2, . . . }.
Taken together, these span all possible inputs that can transform under SE(3),
with type-ℓ vectors containing 2ℓ + 1 elements. Accordingly, type-0 vectors
are single-element scalars s ∈ R, type-1 vectors are 3D vectors v ∈ R3 , type-
2 vectors are matrices M ∈ R3×3 that are symmetric (M = M⊤ ) and traceless
Geometric Graphs 219

(tr(M) = 0), meaning they can be written using only five parameters w ∈ R 5 :
 
w 1 w2 w3
M = w 2 w4 w5  (8.115)
w 3 w5 –(w1 + w4)
Similarly, type-3 vectors are symmetric traceless 3D tensors T ∈ R 3×3×3 ,
which can be represented using seven parameters, and so on.
Note that, while the type-2 vectors in this case do not span all possible input
matrices on their own, any input matrix N ∈ R 3×3 can be represented using a
combination of type-ℓ vectors. To see why, first note that we can use a type-0
vector, s ∈ R, to encode 13 tr(N). Then, we can use a type-1 vector, v ∈ R3, to
encode the antisymmetric components—the non-zero elements of 12 (N – N⊤ ):
 1   
0 (n – n 21) 12 (n13 – n31 )
2 12
0 v1 v 2
1
(N – N ⊤) =  12 (n21 – n12 ) 0 1
(n – n32 ) =  –v1 0 v3 
2 23
2 1 1
2 (n31 – n13 ) 2 (n 32 – n 23) 0 –v2 –v3 0
(8.116)
Once we subtract the trace from N’s diagonal, and the antisymmetric part
1 ⊤ e 1 ⊤
2 (N – N ) from N, we obtain N = N – sI – 2 (N – N ):
   1 1 
n11 – s n12 – v1 n13 – v 2 n11 – s 2 (n12 + n 21) 2 (n13 + n 31 )
Ne = n 21 + v 1 n 22 – s n23 – v 3 =  1 (n12 + n 21 ) n22 – s 1
(n + n 32 ) 
2 2 23
1 1
n 31 + v 2 n32 + v3 n33 – s (n
2 13
+ n 31 ) (n
2 23
+ n 32 ) –(n 11 + n 22 – 2s)
(8.117)
for which we indeed only need a type-2 vector, as the resulting matrix is sym-
metric and traceless. We have hence demonstrated that any 3 × 3 matrix can
be expressed by combining parameters from type-0, type-1 and type-2 vectors
(see Figure 8.5). Similar results hold for other input shapes relevant to SE(3)
symmetry.
Since all possible inputs to an SE(3)-equivariant model will be representable
as type-ℓ vectors, we will assume that our equivariant model is provided with
input feature vectors of the kind f ℓu ∈ R2ℓ+1 . Note that here we will be ignoring
the channel axis (e.g. we will assume there is exactly one type-0 scalar feature,
fu0 ∈ R, exactly one type-1 features, fu1 ∈ R3, etc.) for simplification purposes,
but everything we discuss naturally distributes across multiple channels. Our
equivariant network is then tasked to map this input representation into a suit-
able space of latent vectors: {f uℓ} u∈V ,ℓ≥0 → {h ℓu} u∈V ,ℓ≥0. We will once again,
for simplicity, assume that the latent space hℓu ∈ R 2ℓ+1 has no channel axis.

[Link] Irreducible representations of rotations The strategy by which

we will construct this generic set of equivariant architectures is by leveraging
220 Chapter 8

the irreducible representations (irreps) of the chosen symmetry group. We pre-

viously discussed irreps in Chapter 3; their key property is that representations
ρ(g) of all g ∈ G can be expressed as some combination of the irreps of G (up
to a change of basis). Irreps can be pre-computed for many groups of interest,
and they can massively reduce the number of cases we need to consider when
building the set of all possible G -equivariant models.
Fortunately, the 3D rotation group SO(3), a key part of the roto-translations
we care about, has a well-known, computable, set of irreps—the (real)
Wigner D-matrices 8, Dℓ (g) ∈ R(2ℓ+1)×(2ℓ+1) (Wigner 1927). Specifically, for
any representation, ρ, and change-of-basis orthogonal matrix, Q:
" #
M
⊤
ρ(g) = Q Dℓ (g) Q (8.118)
ℓ∈L
L (a+b)×(a+b)
where : Ra×a × R b×b → R is the direct sum, and L ⊆ {0, 1, . . . } is
the set of vector orders that we are acting on (i.e. the specific type-ℓ vectors
we have access to for our task). Note that the shape of Q ∈ R N ×N is then spec-
P
ified by N = ℓ∈L 2ℓ + 1. A useful special case of this are “pure” type-ℓ vector
inputs (for which L = {ℓ}); they transform only according to Dℓ (g)—in fact,
type-ℓ vectors are rigorously defined exactly by using this property.
As the Dℓ (g) matrices can be efficiently computed (for example, by libraries
such as e3nn (Geiger and Smidt 2022b)), and diving deeper into their con-
struction is not necessary for understanding their implications on equivariant
architectures, we will not discuss them further here.
Instead, we immediately divert our attention to how we can support mes-
sage passing between our node features fu and fv . Note this was much simpler
to do in the previous equivariant GNNs we studied, as they only exchanged
type-0 and type-1 features, with carefully constructed invariant and equivari-
ant transformations guaranteed to compose well. To truly span the entire space
of possible SE(3)-equivariant GNNs, we need to carefully take into account all
of the type-ℓ vector orders and how they might interact.

[Link] Constructing equivariant filters via spherical harmonics The

first stage of any message passing layer is, for each edge (u, v) ∈ E , to convert
sender features, fv, into messages, m uv . Note we will assume that receivers do
not participate in the message computation, for reasons of simplicity. As the
messages are expected to transform together with SE(3), they too will need to
be decomposable into type-ℓ vectors.
In order to ensure that these vectors respect the equivariance, the way we
compute them must depend on the relative positions of the two nodes. As we
already know, node coordinates, x u∈ R 3, are 3D vectors and hence part of the
Geometric Graphs 221

type-1 node features, fu1. To make the computations explicit, we will treat x u
as a completely separate vector from f u1 , but please note that it is completely
possible to track and update coordinates within fu1 ; we’re mainly splitting them
to keep the operations easily interpretable. Assuming we’re sending a message
from v to u, the relevant displacement vector is ruv = xv – x u.
Keeping in mind that the type-ℓ message, mℓuv , may itself depend on fea-
tures of several types, fvk, we need to construct appropriate filters Wℓk : R3 →
R (2ℓ+1)×(2k+1) . that will allow us to combine our messages as
X
mℓuv = Wℓk (ruv )fvk (8.119)
k≥0

The filter, Wℓk (r uv), transforms a neighbour input of type k into a partial mes-
sage of type ℓ. While assembling such messages by summation, as is done here,
appears arbitrary, it is actually provably universal (Dym and Maron 2021). The
filter is a function of relative displacement: as we discussed above, if this would
not be done, there would be no way to ensure equivariance holds.
To allow us to account for how the filter should behave under roto-translation
in a principled manner, we separate the displacement into its magnitude and
direction components: ruv = ∥ruv ∥r̂ uv, where r̂uv = r uv /∥ruv ∥ is the normalised
displacement vector.
The magnitude ∥ruv ∥ is an invariant quantity with respect to SE(3), and
hence we can act on it directly with an arbitrary neural network without dis-
rupting equivariance. We call such a model a radial neural network, φℓ : R + →
 
R. Note that φℓ ∥r uv∥ must be a scalar (type-0 vector) in order to preserve
equivariance properties (cf. the need for a scalar output from β in Equation
8.111). The inclusion of ℓ in φ ℓ implies that we can have separate radial net-
works for different vector types. As for the normalised displacement r̂uv , first
of all, note that it is now defined on the sphere, S 2. Since we need to produce a
filter of type ℓ, and type-ℓ vectors have 2ℓ + 1 entries, we require a function of
the form Y ℓ : S 2 → R 2ℓ+1 which is equivariant to rotations in SO(3). Fortunately,
it is well-known how to construct an orthonormal basis of such functions: we
may leverage the spherical harmonics (Figure 8.6) as our choice of Yℓ . Cru-
cially, if the input coordinate is rotated by g ∈ SO(3), the resulting type-ℓ vector
is guaranteed to transform according to the corresponding Wigner D-matrix:
Y ℓ (Rgr̂ uv ) = Dℓ (g)Yℓ (r̂uv ), (8.120)
hence guaranteeing rotation equivariance. Note that spherical harmonics are
fundamental basis functions on the sphere (with significant implications
to chemistry, as they represent various quantised configurations of atomic
orbitals), and are readily pre-computable. Therefore, as was the case for
222 Chapter 8

Figure 8.6
Visualisation of the first four orders of spherical harmonics, Y ℓ (r̂) ∈ R 2ℓ+1 , for ℓ ∈
{0, 1, 2, 3} (one per row). Note how increased orders increase subdivisions of the
sphere, not unlike the various-order orbital arrangements of electrons within an atom.
Lastly, the rows show only the first half of the harmonic dimensions (including the 0th
mode)—the other half are identical in shape to what is plotted here, but appropriately
rotated.

Wigner D-matrices, we will not further comment on the specific equations

used to compute the spherical harmonics, Yℓ.
Recapping our progress so far, we’ve successfully combined a (learnable)
radial neural network with a (pre-computed) spherical harmonic, to obtain
 
a type-ℓ filter vector, φ ℓ ∥r uv∥ Yℓ (r̂uv ) ∈ R2ℓ+1 . The only part that remains
to specify our final filter, Wℓk (ruv ), is the describe how this filter vector
needs to interact with the sender’s type-k feature vector, fvk, while preserving
equivariance. We study how to principally construct such interactions next.

[Link] Feature-filtering interactions via Clebsch-Gordan coefficients

Thus far, by leveraging spherical harmonics, we have constructed an equivari-
ant, type-ℓ filter vector, which we will denote by ϕℓ ∈ R2ℓ+1 . Now, we would
like to explicitly interact this filter vector with a given neighbour feature vector
Geometric Graphs 223

of type k, f vk ∈ R2k+1 . More generally, for any given pair of type-ℓ filter vec-
tor and type-k input vector, we require a recipe allowing us to combine their
features in a way that still transforms according to SO(3).
Expressive feature interactions typically come from cross-multiplying them.
In the case of our two feature vectors, ϕ ℓ and fvk , we obtain all possible pairwise
multiplications through their tensor product, ϕℓ ⊗ fvk ∈ R(2ℓ+1)×(2k+1), specif-
 
ically, ϕℓ ⊗ fvk ab = ϕℓa(fvk ) b. The resulting vector then transforms according
to the product representation of relevant Wigner D-matrices, Dℓ ⊗ Dk . Hence
the tensor product is no longer a “pure” type-ℓ vector, and we cannot use its
values directly in our neural network, as we assumed it would rely solely on
maintaining and manipulating pure vectors. This means that we need to dis-
entangle the various pure-type components of ϕℓ ⊗ fvk and appropriately route
them. Fortunately, we know how to reduce tensor products of type-ℓ vectors
and type-k vectors into a collection of type-j vectors (of various orders, j).
Interestingly, due to the underlying group theory, the pure orders j contained
within the tensor product must satisfy the triangle inequality w.r.t. ℓ and k:
|ℓ – k| ≤ j ≤ ℓ + k .
We can explicitly express this decomposition9 as follows:
2ℓ+1 X
X 2k+1 2ℓ+1 X
X 2k+1
 ℓ   k
(ϕℓ ⊗ fvk) jm = C jm k
ℓ,m1 ,k,m2 ϕ ⊗ f v m ,m = Cjm ℓ
ℓ,m1 ,k,m 2 ϕm 1 fv m
1 2 2
m 1 =1 m2 =1 m1=1 m 2 =1
(8.121)
through the Clebsch-Gordan (CG) coefficients, Cjm ℓ,m 1,k,m2 ∈ R. One way to
read this notation is: “the coefficient sending the result of cross-multiplying
the m 2-th component of a type-k vector with the m1 -th component of a type-ℓ
filter, to the m-th component of a type-j vector result”. CG coefficients offer a
recipe for a change-of-basis operation that can be precomputed and kept fixed.
As such, similarly to Wigner D-matrices and spherical harmonics, many estab-
lished equivariant learning libraries support computing these coefficients out
of the box, and we will not dwell further on how CG coefficients are con-
structed. The CG coefficients are the final piece needed to build our filter
matrices W ℓk (ruv). Combining all we’ve seen so far, with slightly modified
notation:
Xℓ+k X  ruv 
2j+1
ℓk
W (ruv) = φ j(∥ruv ∥) Yj Cℓ (8.122)
∥ruv ∥ m jmk
j=|ℓ–k| m=1

In this equation, j ranges over all type-j filters that will yield a message of
type ℓ when applied to type-k vector features. The radial neural network φ j
scales the intensity of the interaction, the spherical harmonic Y j provides the
directional spatial awareness, and the matrix Cℓjmk ∈ R(2ℓ+1)×(2k+1) comprises
224 Chapter 8

f 0v f 1v f 2v

C 00,:,0 C 01,:,1 C02,:,2 C11,:,0 C10,:,1 C 11,:,1 C12,:,1 C11,:,2 C12,:,2 C22,:,0 C 21,:,1 C22,:,1 C20,:,2 C21,:,2 C 22,:,2

Σ Σ Σ

m0uv m1uv m2uv

Figure 8.7
A schematic of the “Clebsch-Gordan switchboard” for a Tensor Field Network handling
vectors up to and including type 2. This is meant to illustrate which neighbour feature
orders, fvℓ , contribute to which message orders, mℓuv , as well as which Clebsch-Gordan
coefficients, Cℓj,:,k handle each part of the communication. Note that this diagram, for
simplicity, does not include any specifics about the filtering operation through the
spherical harmonic. Rather, the tensor product with the appropriate order-j spherical
harmonic, Yj (r̂ uv), is implicitly computed within each coefficient block.

all the relevant Clebsch-Gordan coefficients, C ℓ,:j,m,k,: . This matrix hardwires the
rotationally equivariant pathways mapping the type-k neighbour feature vector
to the type-ℓ message, via the type-j filtering. This “switchboard” is illustrated
in more detail in Figure 8.7. As before, one might wonder whether this kind
of composition is universal: several concurrent works have proved this fact
(Weiler et al. 2018; Thomas et al. 2018; Kondor 2018). That is, the space of
possible W ℓk(r uv) filters fully spans all SE(3)-equivariant layers.

[Link] Tensor Field Networks Throughout prior sections, we formalised:

• The key information channels that allow us to seamlessly characterise roto-

translation equivariant computation—the type-ℓ vectors;

• Constructing filtering operations via their orthonormal basis—the spherical

harmonics—weighted by radial neural networks; and

• Properly disentangling the outcome of filtering operations into “pure” type-ℓ

vectors, via the fixed exchange parameters—the Clebsch-Gordan coefficients.

Geometric Graphs 225

We are now ready to put all of the above together into tensor field networks
(Thomas et al. 2018, TFNs), a popular model for SE(3)-equivariant learning
over 3D point clouds. Note that, while most of our discussions so far focussed
on SO(3), the fact we are solely relying on relative displacements (which do
not vary with translations) implies that the corresponding model will have
translation equivariance as well.
Keeping in mind Equation 8.119’s form of computing the message mℓuv , we
simply combine the results across all neighbours:
X
h ℓu = mℓuv (8.123)
v∈Nu
XX
= Wℓk (xv – xu )fvk (8.124)
v∈Nu k≥0

ℓ+k
XXX 2
Xj+1  
xv – xu
= φj (∥x v – xu∥) Yj Cℓjmkf kv (8.125)
∥xv – xu ∥ m
v∈Nu k≥0 j=|ℓ–k| m=1

where, as before, we leverage summations throughout as, per Dym and Maron
(2021), this is sufficient to express all SE(3)-equivariant operators. Please note
a specific edge case of node u’s interactions with itself (v = u). While the case
might appear problematic due to the division by zero in the spherical harmonic,
by Schur’s lemma (as discussed in Chapter 3), the only allowed interactions are
via a single self-interaction scalar: w ℓℓ fuℓ . We accordingly separate this case:
X XX ℓ+k X  xv – x u 
2j+1
ℓ ℓℓ ℓ
h u = w fu + φj (∥x v – x u∥) Yj Cℓ f k
∥xv – x u ∥ m jmk v
v∈Nu \{u} k≥0 j=|ℓ–k| m=1
(8.126)
There is a very worthwhile point to be made about Equation 8.126, as it will
recur for the remainder of this section. Namely, while it is certainly alge-
braically complicated, this Equation has very few learnable parameters—in
fact, the only controllable components are interaction strengths wℓℓ and φj —
and those only directly depend on the invariant features of the problem. All
other computations, tricky as they may be, derive strictly from our knowledge
and understanding of the rotations in SO(3), and can be fully pre-computed.
Looked at in this way, the Equation is no more than yet another instance of
the familiar message passing operator, with complicated scaffolding depending
on the spatial arrangement of the input nodes, and a small parametric compo-
nent controlling how all the principal directions of interactions are recombined.
In fact, this perspective can readily allow us to transform other popular
graph neural network instances into their irrep-projected SE(3)-equivariant
variants—we will now demonstrate this for the Transformer.
226 Chapter 8

[Link] SE(3)-Transformers Recall from Chapter 5 that a Transformer

(Vaswani et al. 2017) relies on a dot-product self-attention mechanism (Equa-
tion 5.69); allowing receiver nodes to dynamically alter the interaction strength
with every incoming message, through a coefficient α uv ∈ R. These coeffi-
cients are derived from normalised dot-products of query and key vectors in
the receiver and sender node, correspondingly; and the messages themselves
are derived as value vectors in every sender. Interpreting Equation 8.126 as a
message-passing equation, we find that only simple alterations are needed to
make it attentional:
X X X ℓ+k 2j+1
X  
V x v – xu
h ℓu = wℓℓf
V u
ℓ
+ α uv φ j (∥x v – xu ∥) Y j Cℓ f k
∥x v – xu ∥ m jmk v
v∈N u\{u} k≥0 j=|ℓ–k| m=1
(8.127)
Specifically, we explicitly use the V identifier on our parametric components
V
(wℓℓ
V and φj ) to denote that these parameters pertain to the computation of
value vectors. Further, the computed value vectors are now scaled by the atten-
tion coefficient, αuv . In order for the coefficient’s multiplication to not disturb
the message’s equivariance properties, it needs to be SE(3)-invariant (just like
in the case of φVj ).
Fortunately, to ensure this invariance property, all we need is for our query
and key vectors to be SE(3)-equivariant. To see why in the case of rotations,
consider two vectors q(x), k(y) and an attention coefficient α = q(x)⊤ k(y).
Further, assume the corresponding equivariance conditions on both q and k,
i.e., q(Rx (θ )x) = Rqk (θ )q(x) and k(Ry (θ )y) = Rqk(θ )k(y) for every rotation
angle θ . Then, when rotating our frame by θ, the dot product becomes:
 ⊤
q(Rx(θ )x)⊤ k(R y(θ )y) = Rqk (θ )q(x) Rqk (θ )k(y)
= q(x)⊤Rqk (θ ) ⊤Rqk (θ )k(y)
= q(x)⊤k(y) = α
where the final cancellation is due to the orthogonality of the rotation matrix
R qk(θ ). Conveniently, we already have a recipe for building general-purpose
equivariant vectors; therefore, we can construct our q and k vectors by re-using
the entirety of the theory that led us to TFNs!
To seamlessly do this, we will express both query and key vectors as a
collection of type-ℓ vectors; qℓu , kℓuv ∈ R2ℓ+1 . Note the difference in indexing:
queries inherently depend on a particular receiver, whereas keys will actually
depend on both the sender and receiver, allowing us to explicitly take into
account their displacement10 , xv – xu . Further, this implies that query vectors
will not depend on any spherical harmonic or Clebsch-Gordan components,
Geometric Graphs 227

while key vectors will be derived similarly to value vectors:

XX ℓ+k 2j+1
X  
K x v – xu
ℓ ℓℓ ℓ
qu = w Q fu ℓ
k uv = φj (∥xv – xu ∥) Yj Cℓ fk
∥x v – x u ∥ m jmk v
k≥0 j=|ℓ–k| m=1
(8.128)
while keeping in mind that the equation above assumes u = ̸ v , and we need to
fall back to scalar multiplication kℓuu = wℓℓ
K fu for the special case.
ℓ

Once both queries and keys are prepared, we can simply combine all of their
dot products together or, equivalently, stack them in a single vector and then
perform the dot product. Then we use the softmax, as usual, to normalise:
 ⊤ L 
L
exp q
ℓ≥0 u
ℓ k
ℓ≥0 uv
ℓ

α uv = X   L  (8.129)
L ⊤
exp ℓ
ℓ≥0 qu
ℓ
ℓ≥0 k uw
w∈N u

Plugging these coefficients back into Equation 8.127, we recover the SE(3)-
Transformer model of Fuchs et al. (2020). While, as per our previous
discussion, such a layer does not amplify the expressive power of TFNs
(Thomas et al. 2018), it does provide learning flexibility to the model, and
often yields stronger predictions in practice. Simply put, it might take a
TFN substantially more resources to exactly emulate an attention operator; in
SE(3)-Transformers, that functionality—which we know works quite well for
representation learning—is baked in from the start.

8.3 Case Study: Manifold And Mesh CNNs

Throughout this Chapter, we have observed a geometric graph’s nodes primar-
ily as a point cloud : a set of points floating freely in 3D space, perhaps with
some observed abstract connections between them. However, very often these
points might span a more interesting surface structure, which we refer to as
a mesh. In fields like computer graphics and biology, 3D objects are almost
always modelled as triangulated meshes—a collection of vertices, edges, and
triangular faces. Mathematically, meshes typically arise as discretisations of
a manifold —a continuous, smooth surface (Figure 8.8). We’ve already seen
examples of how to build neural networks over specific manifolds, when defin-
ing the spherical CNN (T. S. Cohen et al. 2018) in Chapter 7. That said, spheres
are unusually simple manifolds in the sense that every point locally “looks
exactly the same” (recall we referred to such spaces as homogeneous). Most
manifolds do not satisfy this property, and hence, general-purpose manifold
neural networks necessitate greater care over each point’s local shape. In this
section, we will observe some common strategies for building general-purpose
228 Chapter 8

Figure 8.8
Meshes represent a triangulated geometric surface; they can be represented at different
levels of resolution, as visualised here. As the resolution becomes finer and finer, the
mesh approaches the original manifold it was derived from.

convolutional neural networks over such inputs—concluding that they are not
dissimilar from geometric GNNs after all.
Understanding several of the mathematical objects we will discuss here
relies on a working knowledge of differential geometry. As this field is per-
haps less familiar to the machine learning audience, we will aim to intuitively
introduce the key concepts required for our discussion, and refer the reader to
Penrose (2005) for their detailed exposition.

8.3.1 Riemannian manifolds

Since the formal definition of a manifold is somewhat involved, we prefer to
provide an intuitive picture at the expense of some precision. To start, imagine
you are standing on the surface of the Earth. Globally, the Earth is a curved sur-
face (specifically, an ellipsoid ). But, locally, right under your feet, the ground
feels perfectly flat. This property is essential for defining a manifold, Ω.
Specifically, one can think of a (differentiable or smooth) manifold as a
smooth, curved surface that is locally Euclidean, in the sense that any of its
small regions around any point can be deformed to appear like a patch of R s;
in this case, the manifold is said to be s-dimensional (in the Earth example,
s = 2, as around any point it appears like we are walking on a flat 2D plane).
Having identified this patch of Rs around a point, u ∈ Ω, we can extend it
infinitely in all possible directions, as a local approximation of the manifold,
called the tangent space, Tu Ω. In the case of spheres, S2, the tangent space,
T uS2 = R2 , becomes the unique plane which is tangent to any given point u; see
Geometric Graphs 229

Figure 8.9. Once such a space is defined, we can use it to express displacements
of various points in u’s local neighbourhood, by way of elements X ∈ Tu Ω,
which we call tangent vectors. In the sphere example, X ∈ R2 gives us the
displacement vector in the tangent plane, which we can then project onto the
sphere to find a particular neighbouring point.
Since the tangent space is only an approximation of the underlying mani-
fold, we cannot directly use standard Euclidean metrics on them as a substitute
for the points they represent. For example, the Euclidean norm of X ∈ Tu Ω
will not generally tell us how far away the corresponding point actually is on
the manifold. In order to properly measure such quantities, we need to equip
our tangent space with additional structure, expressed as a positive-definite
bilinear function gu : TuΩ × TuΩ → R, depending smoothly on u. Such a func-
tion is called a Riemannian metric11 , and it can be thought of as an inner
product on the tangent space, ⟨X , Y ⟩ u = g u (X , Y ), which evaluates the angle
between any two tangent vectors X , Y ∈ T u Ω. Using it, we can also induce a
norm, ∥X ∥u = g1/2 u (X , X ), allowing to locally measure lengths of vectors. For
the sphere example, the Riemannian metric is just the usual dot product of
vectors in R 2: g u (x, y) = x⊤y.
We must stress that tangent vectors are abstract geometric entities that exists
in their own right and are coordinate-free. If we are to express a tangent vector
X ∈ T uΩ numerically as an array of numbers, we can only represent it as a list
of coordinates x = (x1 , . . . , x s) relative to some local basis {X1, . . . X s} ⊆ Tu Ω.
Similarly, the metric can be expressed as an s × s matrix G with elements
gij = gu (Xi , Xj ) in that basis (often known as the Gram matrix).
A manifold equipped with a metric is called a Riemannian manifold, and
properties that can be expressed entirely in terms of the metric are said to
be intrinsic. This is a crucial notion for our discussion, as according to our
template, we will be seeking to construct functions acting on signals defined
on Ω that are invariant to metric-preserving transformations called isometries,
Iso(Ω), that deform the manifold without affecting its local structure.

8.3.2 Geodesics
Going back to our Earth manifold example: if you want to describe the location
of a nearby city to where you’re standing, you would not draw a straight 3D
line through the crust of the Earth. Instead, you’d trace the shortest path along
its surface. This “surface-hugging” shortest path is called a geodesic.
To see how to formalise geodesics, consider a smooth curve γ : [0, T ] → Ω
on the manifold, which traces a path between endpoints u = γ (0) and v = γ (T ).
We can express the direction this curve has in each of its points, t , via its
derivative. In this context, the derivative, γ ′ (t), is a tangent vector in Tγ(t) , and
230 Chapter 8

Figure 8.9
Basic notions of Riemannian geometry illustrated on the example of the two-
dimensional sphere S 2 = {u ∈ R3 : ∥u∥ = 1}, realised a subset (sub-manifold) of R 3. The
tangent space to the sphere is given as T uS2 = {x ∈ R3 : x ⊤u = 0} and is a 2D plane –
hence this is a 2-dimensional manifold. The Riemannian metric is simply the Euclidean
inner product restricted to the tangent plane, ⟨x, y⟩u = x ⊤y for any x, x ∈ Tu S2 . The
exponential map is given by expu (x) = cos(∥x∥)u + sin(∥∥xx∥∥) x, for x ∈ T uS2 . Geodesics
are great arcs of length d(u, v ) = arccos(u⊤ v).

we can take the norm of this vector to estimate its local velocity. Among all
the curves connecting points u and v, we are interested in those of minimum
length, i.e., we are seeking γ minimising the cumulative velocity:
Z T Z T
ℓ(γ ) = ∥γ ′(t)∥ γ(t) dt = g1/2 ′ ′
γ(t) (γ (t), γ (t)) dt .
0 0
Such curves are called geodesics (from the Greek γϵoδαισια, literally ‘divi-
sion of Earth’) and they play a very important role in differential geometry.
Note that the way we defined geodesics is intrinsic, as they depend solely on
the Riemannian metric (through the length functional). In the example case of
the sphere, geodesics are great arcs between two points.
Using the concept of a geodesic, we can define several other important
intrinsic concepts: parallel transport, exponential maps and geodesic distance.
Parallel transport Since the tangent spaces around two distinct points,
u, v ∈ Ω, are themselves distinct, it’s not trivial to compare any two vectors
X ∈ Tu Ω, Y ∈ T vΩ across them. Making use of geodesics, we can construct a
mechanism to “transport” a tangent vector in one space to a tangent vector in
another (Figure 8.10).
To see how, let γ be a geodesic connecting u = γ (0) with v = γ (T ), and let X ∈
T uΩ be a tangent vector at u. We can now define a collection of tangent vectors
along the geodesic, X(t ) ∈ Tγ(t) Ω, that preserves the angle to the derivative and
Geometric Graphs 231

Figure 8.10
Assume we want to transport a vector in A’s tangent plane to C ’s tangent plane. Trans-
porting it using Euclidean translation, without changing its direction (shown in red),
results in vectors that are not part of C ’s trangent plane. Parallel transport, in contrast,
constantly reorients the vector along the geodesic (shown in blue). This also means that
it is dependent on the path taken—AC and ABC paths result in different final vector
directions in C’s tangent plane.

norm, for all points t ∈ [0, T ]:

gγ(t) (X (t), γ′ (t)) = gu(X , γ ′ (0)) ∥X (t)∥γ(t) = ∥X ∥ (8.130)
and this uniquely defines X (T ) ∈ T vΩ. We can consider the corresponding map,
Γu→v : Tu Ω → TvΩ, with Γ u→v(X ) = X (T ), as the parallel transport of tangent
vector X into TvΩ, allowing us to compare X and Y directly. Given that we
preserve lengths and angles, parallel transport amounts to only rotation of the
tangent vector, so it can be associated with an element of SO(s).
Exponential maps We can interpret tangent vectors, X ∈ T uΩ, as ways to
“take steps” of a certain direction and magnitude away from a given point
u ∈ Ω. Therefore, following X should land us at another point, v ∈ Ω; it would
be very useful if we could rigorously identify this point. Using geodesics, we
can indeed define such a quantity, the exponential map, expu : T uΩ → Ω, so
long as the geodesics on Ω can be sufficiently extended in any direction around
u (if we can infinitely extend them, Ω is said to be geodesically complete).
To see how to build the exponential map, assume we have a point u ∈ Ω and
a tangent vector X ∈ Tu Ω. Then, we define a X-directed geodesic, γX , such that:
γ X(0) = u γX′ (0) = X (8.131)
232 Chapter 8

Note that such a geodesic always exists and is unique, assuming geodesic com-
pleteness. Then, to project the tangent vector back to Ω, we simply query this
geodesic at “unit length”: expu (X ) = γX (1). In the case of spheres, we rely on
projecting tangent vectors, x ∈ R2, back onto the surface, using trigonometry:
sin(∥x∥)
expu (x) = cos(∥x∥)u + x (8.132)
∥x∥
Geodesic distance Last but not least, geodesics and their length function-
als can be used to properly define distances between pairs of points u, v ∈ Ω,
assuming geodesic completeness. Specifically, we can form a geodesic, γ u→v :
[0, T ] → Ω, such that γ u→v(0) = u and γu→v (T ) = v. Then, the geodesic distance
between u and v is simply the cumulative velocity of the geodesic between
them: d (u, v) = ℓ(γu→v ). Note that this quantity is guaranteed to exist, owing to
the Hopf-Rinow theorem12 (Hopf and Rinow 1931). On spheres, the geodesic
distance is just the familiar arc-cosine of the corresponding angle between the
two points: d(u, v ) = arccos(u⊤ v).

8.3.3 Spatial manifold convolutions

Using all of the above concepts, we can define what it might look like to con-
volve a signal on a manifold, of the form x : Ω → R, with an appropriate filter,
θ. Analogously to the group convolutions covered in Chapter 7:
Z
(x ⋆ θ )(u) = x(expu Y)θu (Y ) dY , (8.133)
TuΩ
Note that this convolution is expressed within the tangent space, Tu Ω, to make
advantage of its associated metrics and inner products (which might not be
easy to express on Ω itself). For each tangent vector Y ∈ T uΩ, we look up the
filter in the corresponding point (θu(Y )), apply it to the corresponding point on
the signal, which we retrieve using the exponential map of Y , and we integrate
across the entire space.
While such an equation is tempting to use, it does have several issues when
the space is non-homogeneous. Firstly, the filter θ now needs to be defined
separately for each point’s tangent space, heavily crippling any weight sharing
prospects. Secondly, we generally need local filters, as we cannot integrate
over the entire tangent space T uΩ, as past a certain point geodesics start to
intersect (e.g., in spheres, eventually any geodesic direction will “wrap around”
the sphere and no longer be shortest).
Finally, we cannot work with θ(Y ), as Y is an abstract geometric object:
in order for it to be used for computations, we must represent its coordinates,
relative to some local basis (also known as a gauge). Gauges are defined as ωu :
Geometric Graphs 233

Figure 8.11
Left: Various geodesic patches illustrated on a human pose manifold. Right: A method
for constructing discrete geodesic rays along a mesh surface, which can then be used as
base directions for angular max pooling in geodesic CNNs (Masci et al. 2015).

R s → T uΩ, and once chosen, they allow us to represent Y as an s-dimensional

array of coordinates x = ω–1u (Y ).
Assuming we have chosen an appropriate gauge, we can rewrite the
convolution in Equation 8.133 as
Z
(x ⋆ θ )(u) = x(exp u (ωuy))θ (y) dy , (8.134)
[0,1] s

with the filter now defined locally—in this case, on the s-dimensional unit
hypercube. Since the exponential map is intrinsic, the resulting convolution is
invariant to isometries in Iso(Ω).

8.3.4 Geodesic CNNs

The convolutional filtering of the form in Equation 8.134 is the most common
way of realising architectures for deep learning on meshes—we simply have
to discretise the corresponding exponential map. However, there are several
intricacies arising from the choice of gauge, ωu , which we survey next.
Geodesic patches Assuming our input is a two-dimensional manifold within
3D space, we can define geodesic patches around a given point; these will be
excellent anchors for our filters. Let us choose a local polar gauge ω̃u : [0, R] ×
[0, 2π ) → Tu Ω, using which we define x(u, r , ϑ) = x(exp u ω̃u (r, ϑ)). This allows
us to look up our signal x : Ω → R at a point which is at geodesic distance r
from u ∈ Ω, with the geodesic’s derivative at u forming an angle of ϑ with
respect to a given reference direction. Note that this direction is chosen (arbi-
trarily) by the gauge, to represent an angle of zero radians. Refining Equation
234 Chapter 8

8.134 to use such a coordinate system, we obtain the following form:

Z R Z 2π
(x ⋆ θ )(u) = x(u, r , ϑ)θ (r, ϑ) dr dϑ. (8.135)
0 0
Figure 8.11 (Left) illustrates several geodesic patches resulting from this kind
of operation, over a reference manifold.
Gauge-invariant filtering A key ambiguity in defining geodesic patches is
the arbitrary choice of the gauge’s principal direction (for which ϑ = 0). Indeed,
any rotation of the gauge would alter the ϑ coordinates—and, as many times
before in this book, it would be useful to construct models which do not depend
on such arbitrary decisions.
There exist straightforward ways to make our computation invariant to such
choices of gauge; the arguably simplest of which is to remove the dependence
of the filter on the angle altogether, yielding a simplified filter θ(r ):
Z R Z 2π
(x ⋆ θ )(u) = x(u, r , ϑ)θ (r) dr dϑ.
0 0
While elegant and useful, such isotropic filters have no way of representing
meaningful directional information on the manifold. Masci et al. (2015) extend
their expressive power by instead performing angular max-pooling: attempt-
ing all possible gauges, and keeping only the outcome which maximises the
response of our filter:
Z R Z 2π
(x ⋆ θ )(u) = max x(u, r , ϑ)θ (r, ϑ + ϑ 0) drdϑ.
ϑ 0∈[0,2π ) 0 0

Back to meshes! To actually practically implement such convolutions, we

need to discretise the manifold into a mesh. Upon meshing, geodesics become
poly-lines that traverse the mesh’s triangular faces. Traditionally, geodesics
have been computed using the Fast Marching algorithm (Kimmel and Sethian
1998), an efficient numerical approximation of a nonlinear PDE called the
eikonal equation encountered in physical models of wave propagation in a
medium. This was adapted by Kokkinos et al. (2012) for the computation of
local geodesic patches, and later reused by Masci et al. (2015) for constructing
geodesic CNNs, the first intrinsic CNN-like architectures on meshes.
To make these models practical, it is not sufficient to discretise the
manifold—we also need to discretise the set of considered coordinates within
the gauge. That is, we partition our radial space into Nr bins and our angle
space into Nϑ bins. This means we fire N ϑ geodesic rays out of u that equally
partition the tangent space, and record the filter (and signal) values along those
geodesic rays, in increments of R/Nr . See Figure 8.11 (Right) for an illustration
of how these geodesic rays are computed. All together, this yields the following
Geometric Graphs 235

Figure 8.12
Left-to-right: examples of patch operators used in geodesic CNNs (Masci et al. 2015),
anisotropic CNNs (Boscaini, Masci, Rodolà, Bronstein, and Cremers 2016) and MoNet
(Monti et al. 2017), with the level sets of the corresponding pairwise interaction coeffi-
cients shown in red.

discrete angular max-pooling equation:

N r N ϑ –1    
Nϑ –1 X X ∆r ∆ϑ ∆r (∆ϑ + ϑ0) mod Nϑ
hu = max x u, R , 2π θ R , 2π
ϑ0=0 Nr Nϑ Nr Nϑ
∆r=0 ∆ϑ=0
(8.136)
While it might not be immediately obvious, this equation is a form of geomet-
ric message passing! Namely, each of the recorded points  along thegeodesics
∆ϑ
can be treated as a neighbour of node u (that is, expu ω̃ R ∆r
N r , 2π N ϑ ∈ Nu ). If
our input is already given in discretised mesh form and we do not have direct
access to the underlying manifold, then features will generally not be available
in these points. However, they can be approximated via feature interpolation
of the corresponding triangle endpoints. A few years after geodesic CNNs
were first published, the MoNet framework of Monti et al. (2017) has showed
how they, along with similar proposals such as anisotropic CNNs (Boscaini,
Masci, Rodolà, and Bronstein 2016) may be observed as specific cases of atten-
tional GNNs, wherein attention coefficients are derived from the coordinates13
assigned to each point (Figure 8.12).
Another relevant way of connecting geodesic CNNs to geometric GNNs is
by observing that their filters fundamentally depend on invariant spatial prop-
erties, such as distances and angles. This makes geodesic CNNs effectively
the mesh equivalent of point-cloud GNN models like SchNet and DimeNet. An
interesting point of difference is how invariance in angles is achieved: DimeNet
relies on angles between triplets of nodes, whereas geodesic CNNs compute
angles with respect to a choice of “base direction” (a gauge) and then perform
angular max-pooling to remove the explicit dependence on this choice.
236 Chapter 8

Gauge transformation

gu = 120 ◦
g : V → SO(2)

Figure 8.13
An example of a gauge symmetry on the sphere surface. Gauges in each node’s tangent
space, Tu M, are transformed according to their own assigned 2D rotation, gu ∈ SO(2).
In this case, the highlighted node rotates its own gauge by 120◦ ; note other gauges
rotating by differing amounts.

8.3.5 From gauge invariance to gauge equivariance

The gauge invariant models described so far offer a tried-and-tested way to
remove dependence on what is fundamentally an arbitrary geometric choice;
that said, as we already previously discussed in this Chapter, invariance is also
needlessly destructive of relative spatial information, particularly early on in
a mesh CNN’s computations. To preserve this wealth of information while
also making sure our model does not depend on arbitrary decisions, a more
fine-grained notion of gauge equivariance14 is needed.
As a starting point, we will assume a general 3D mesh, M = (V , E , F ) with
nodes u ∈ V , edges (u, v) ∈ E and triangular faces (u, v, w) ∈ F . Furthermore,
each node has a 3D position, expressed by its coordinate, xu ∈ R3 . As before,
we will assign a tangent space, TuM, at each point; all of the nodes in Nu
will be observed through this tangent space15. Then, we make arbitrary gauge
choices, ωu : R 2 → Tu M, in each node’s tangent space: for example, we can
pick an arbitrary neighbour, v ∈ Nu, and assume that it encodes the base direc-
tion from u; that is, that the angle ϑuv = 0. From this choice of gauge, exact
angles, ϑuw , and tangent plane coordinates, cw ∈ R2 , can be computed for all
nodes w ∈ Nu.
Gauge symmetries Since tangent spaces in this context are 2D planes, and
gauges represent a choice of “base direction” – effectively a 2D unit vector
with a fixed origin – a common set of symmetry transformations between such
gauges are 2D rotations, and possibly reflections, depending on whether we
Geometric Graphs 237

assume the mesh to be oriented. These transformations can be neatly organ-

ised by either O(2) (2D rotations and reflections) or SO(2) (2D rotations only).
Since the choice of gauge is arbitrary, the outputs of our neural networks
over the gauge-derived node features, fu , should transform together with any
transformations applied to the gauges.
Within the context of gauges in 2D tangent spaces transforming under
SO(2), we can perform arbitrary gauge modifications by rotating the angle-
0 direction in any particular node. Importantly, this is different from previous
symmetry groups we study that act consistently across an input—now, we can
apply a different change of gauge in different mesh nodes’ tangent spaces
Tu M, all simultaneously. Generally, we describe an SO(2) gauge symmetry
group, G , as containing maps, g : V → SO(2) (see Figure 8.13). Every node on
the mesh, u ∈ V , picks its own gauge transformation, gu ∈ SO(2). This rotation
modifies the gauge by ωu′ (c) = ωu(ρ(g u)c), where ρ(g u ) ∈ R2×2 is the corre-
sponding rotation matrix. It must transform the individual node coordinates
by c′v = ρ(gu )–1 cv , to ensure the underlying tangent vectors remain the same:
ω′u (c′v ) = ωu (ρ(gu)c ′v) = ωu(ρ(g u)ρ(gu )–1 cv ) = ωu(c v ) (8.137)
For similar reasons, this gauge transformation also must transform the input
features by fu′ = ρin (gu )–1 fu , for appropriately chosen input representation ρin.
Pointwise transformations An important component of any neural network
over a collection of nodes is a pointwise function which is applied in a shared
manner to each node; it is also an excellent way to warm up to designing gauge-
equivariant models, in preparation for interactive network parts down the line.
As such, assume we have a function f : Rk → Rm , writing f (fu) = hu , not at
all unlike the set neural networks we discussed in Chapter 5. Further, assume
that our network’s output space transforms under ρout (gu) in each node u.
What happens to f under a gauge transformation g ∈ G ? f (f u) now becomes
f ′ (f ) = ρout (gu )–1f (ρ in(g u)f ). That is, we gauge-transform the features of the
input, apply f to the new features, and then make an appropriate transformation
in the output space, going against the applied gauge transform. We can prove
that this ensures that output features will gauge-transform as expected via ρ out :
h ′u = f ′(f u′) = ρ out(g u)–1 f (ρin(gu )fu′ ) (8.138)
–1 –1
= ρout(g u) f (ρin(gu )ρin (gu ) f u) (8.139)
–1
= ρout(g u) f (fu ) (8.140)
What would it mean for f to be gauge equivariant? Recall that, generally, equiv-
ariance implies that transformations may be applied at either the input stage or
238 Chapter 8

Gauge transformation
g : V → SO(2)

   
0.7 −1
vu = vu =
0.7 0

f f

Gauge transformation
g : V → SO(2)

  
0 −0.7
vu = vu =
1 0.7

Figure 8.14
Illustrating the gauge equivariance constraint on the particular case of a pointwise func-
tion f , acting on a vector input vu ∈ R2 (shown in red), whose coordinates are computed
relative to node u’s gauge (depicted in blue). Note the significance of the provided com-
mutative diagram: when f is gauge equivariant, we can either first apply f to our vector,
then rotate our gauges by an amount – or first rotate the gauges, then apply f ; the end
result is a vector with the same coordinates.

at the output stage of a function, without causing any difference in outcome:

f (ρin (gu )fu) = ρ out(gu )f (fu) (8.141)
for all gauge transformations g : V → SO(2) (Figure 8.14). But, if we combine
this condition with the transformed function, f ′, we recover a significant result:
f ′(f u) = ρout (gu )–1 f (ρin (gu )f u) = ρout (gu ) –1ρ out(g u)f (fu) = f (fu) (8.142)
Intuitively, if we know a pointwise function f is gauge equivariant, its
implementation does not change when applying gauge transforms!
How do we ensure a pointwise f is gauge-equivariant? Since, in the point-
wise case, both sides of Equation 8.141 concern the same local transformation,
gu ∈ SO(2), we conclude f simply needs to be equivariant to SO(2). In our pre-
vious sections, we demonstrated several mechanisms to build such functions in
Geometric Graphs 239

three-dimensional spaces; ranging from manual construction to projecting onto

irreducible representations. In both cases, the SO(2) group is significantly eas-
ier to work with than its 3D counterparts, and in practice this part is relatively
simple to construct given the previous context discussed in this Chapter.
Things will, of course, get more interesting when we allow our function to
process multiple nodes’ features at once.
Interactive transformations To ease us into analysing the case where mul-
tiple nodes’ features can interact, we can first consider the case where all node
input features, fu , are “pure scalars” (i.e. type-0 vectors).
Interestingly, since scalars do not have any spatial orientation, their features
are unaffected by change of gauge. As such, gauge-equivariant models admit
any linear map between scalar-valued node inputs and outputs, including fully
weight-unshared options such as the following:
X
hu = θuvfv (8.143)
v∈V

Note that such an equation uses a different filter θu for different receiver nodes
u ∈ V, yielding a model that is all but impractically general.
To actually extract value out of gauge equivariance conditions in the interac-
tive setting, we need to “tensor up” to higher-type vectors in SO(2) (note that
these are not the same vector types as in the 3D version we previously studied,
as they relate to the irreducible representations of SO(2), not SO(3)!).
Let us assume that we have a vector feature in each node’s tangent space,
interpreted as fu ∈ R 2 relative to a choice of gauge. If we attempt to naïvely
adapt our parametric linear function to this case, it would be of the form:
X
hu = Θuv fv (8.144)
v∈V

where Θ uv ∈ R2×2is a weight matrix, designed to convert vector features in v

into vector messages in u. However, here we hit the opposite kind of road-
block: since the gauges used to measure fv and hu are allowed to change
independently and arbitrarily, we obtain an unrealistic equivariance constraint:
Θ uv = ρ(gu )–1 Θuvρ(gv ) (8.145)
and more importantly, it must hold for all possible choices of g u, gv ∈ SO(2).
This constraint states that, under any change of gauge either in its input (ρ(g v))
or output (ρ(g u)) space, this filter’s actions need to remain consistent. Due to
the wielding of two gauges at once, it should come as no surprise that the only
possible choice of Θuv in this case is an all-zero matrix. Once again, the con-
straint proves impractical in this case. The only way forward is to reduce the
number of gauges we’re tracking from two to one—that is, we first transport
240 Chapter 8

xv xv
v v x naive
v→u
Parallel transport gv →u
u v→u u x v →u

Figure 8.15
We want to parallel-transport a vector, xv, in v’s tangent space, Tv M, into a neigh-
bour’s tangent space, Tu M. Specifically, we transport x v onto u along the geodesic
specified by the mesh, possibly with additional gauge alignment, giving a transported
vector xv→u . It is cumbersome to do this procedure for every possible vector on-the-fly,
therefore what we do instead is compute the 2D angle, gv→u , that xv→u makes with the
coordinates of x v interpreted within the gauge of v, which we call xnaive
v→u . The result-
ing angle between xv→u and xnaivev→u is then exactly the transformation, gv→u , needed. In
pracitce, the transporters and their corresponding matrices are pre-computed for a given
mesh before model training or inference are executed.

the vectors from v to u’s tangent space, and only then do we invoke the filter:
X
hu = Θuvρ(g v→u)f v (8.146)
v∈V

Here, gv→u denotes the parallel transport from v to u along the geodesic con-
necting them (see Figure 8.15). In general, this transport will lead to a rotation
of the vector (aligning it with u’s gauge), which we encode using the represen-
tation ρ(g v→u) ∈ R2×2 . To avoid ambiguity, the geodesic from v to u must be
unique, which means the filters must necessarily have local support only (i.e.
around specific neighbourhoods Nu ). This ensures that our transport matrix,
ρ(gv→u ), remains nonparametric, leaving all learnable parameters with Θuv.
If we perform a gauge transformation g ∈ G in this setting, the neigh-
bour’s features transform as ρ(gv) –1fv, and the transport matrix transforms as
ρ(gu ) –1ρ(gv→u )ρ(gv). And since our filter, Θuv , now operates only within u’s
tangent space, its equivariance constraint only requires commuting with gu:
Θuv ρ(gu) = ρ(g u)Θ uv (8.147)
Geometric Graphs 241

Provided this constraint is satisfied, we can show that the overall equation is
gauge-equivariant:
X
h ′u = Θuv ρ(g′v→u )fv′
v∈N u
X
= Θuv ρ(gu)–1 ρ(gv→u)ρ(g v)ρ(g v )–1 fv
v∈N u
X
= ρ(gu )–1 Θuv ρ(gv→u)f v
v∈N u
X
= ρ(gu )–1 Θuv ρ(g v→u )fv = ρ(gu )–1 hu
v∈Nu

8.3.6 Gauge equivariant mesh CNNs

The just-discussed blueprint turns out to offer a fruitful, practical method
for constructing mesh neural networks with gauge equivariance. These gauge
equivariant mesh CNNs (GEM-CNNs) have been initially proposed by T.
Cohen et al. (2019) (over the icosahedron, offering a fast approximation to
spherical CNNs16 ) and de Haan et al. (2020) (over arbitrary 2D meshes).
Let fu ∈ Rk be a feature vector at node u ∈ V . GEM-CNNs compute a latent
feature h u ∈ Rm , through a localised form of Equation 8.146, while also noting
that tracking Θuv across nodes is cumbersome, and instead it is easier to make
the weight matrix depend on a geometric pairwise property of u and v such
as their angle, ϑuv , relative to node u’s local choice of gauge. The model is
realised as follows:
X
h u = Θ self fu + Θ(ϑuv)ρ(g v→u)f v (8.148)
v∈Nu

Here, ρ(g v→u) ∈ Rk×k is the previously described transporter matrix, which
parallel transports the neighbour’s input features into the receiver’s tangent
space. Note that the specific instantiation of the transporter will depend on
the specifics of the input feature space; for example, which type-ℓ vectors are
present in the input (explicitly or implicitly). Further, our parameters Θself ∈
R m×k and Θ : [0, 2π ) → Rm×k are chosen to appropriately commute with gauge
transformations. As discussed before, this means they must satisfy:
Θ selfρ in(ϑ) = ρout (ϑ)Θ self Θ(ϑ uv – ϑ)ρin (ϑ) = ρ out(ϑ)Θ(ϑuv) (8.149)
where ρin (ϑ) ∈ Rk×k and ρout (ϑ) ∈ Rk×k are the representations of 2D rotations
over the input space and its representations over the output space, respectively.
Since these constraints are linear, the space of matrices Θself and matrix-
valued functions Θ satisfying these constraints is a linear subspace, and so we
242 Chapter 8

can parameterise them as a linear combination of basis kernels with learnable

P P
coefficients: Θself = i αi Θ(self
i)
and Θ(ϑ) = i βi Θ(i) (ϑ).
The Geometric Bridge While it took us a lot of mathematical juggling to
derive it, the final GEM-CNN (Equation 8.148) looks strikingly similar to the
usual message passing model between the node features on the mesh. The core
two differences are that (1) we rely on our understanding of the mesh’s geome-
try to precompute transporter matrices ρ(g v→u), and (2) we restrict our weight
space to a learned combination of basis matrices conditioned on some invari-
ant property (such as the angle, ϑuv ). Once the transporter and basis matrices
are pre-computed, they act simply as multiplicative factors to the incoming
message, with no additional corrections needed.
Beyond its connection to GNNs, we can make a more specific connection to
geometric GNNs such as the previously discussed Tensor Field Networks. The
same way TFNs used spherical harmonics as their transformation basis to learn
an angle-dependent linear combination from, and Clebsch-Gordan coefficients
to perfectly preserve global 3D rotation symmetries across type-ℓ vectors dur-
ing message passing, the exact same point can be made about the GEM-CNNs’
basis matrices Θ(i) (ϑ) and the transporter ρ(gv→u ). The main distinction is that
we are now preserving a set of local SO(2) symmetries (relative to each node’s
gauge) rather than a single SE(3) symmetry that is globally applied.
Ultimately, whether one models the atoms of a freely rotating molecule in
3D space, or aerodynamic effects across a given manifold or mesh, the geo-
metric GNN blueprint remains exactly the same: define the relevant (global or
local) symmetries, design filters that respect them, and use those filters to pass
messages along the edges of the underlying graph.

8.4 Case Study: AlphaFold 2

One of the greatest achievements of deep learning is definitely its use in solv-
ing17 the protein folding problem: predicting the 3D spatial configuration of
proteins, given their amino acid sequences. This success was achieved through
the AlphaFold 2 architecture (Jumper et al. 2021), which was then promptly
used to approximately fold, at massive scale, many other protein structures for
which structures were not known, starting with the human proteome (Tunya-
suvunakool et al. 2021). For its contributions to computational biology, one
half18 of the 2024 Nobel Prize in Chemistry was awareded to John Jumper and
Demis Hassabis for this result.
Of particular interest to this Chapter is the fact that AlphaFold 2 heavily built
on top of its predecessor, which was a rather standard deep CNN architecture
Geometric Graphs 243

applied directly to the amino acid sequence (Senior et al. 2020), to more heav-
ily lean into the geometry, physics and bioinformatics of the protein folding
task. As such, and given the architecture’s thorough contribution to science,
we found it to be a fitting way to close out the geometric graphs Chapter.
AlphaFold 2 is, in its entirety, a rather complex system19 and, given our spe-
cific interests, we will refrain from studying it in its entirety—rather, we will
focus on two key subcomponents of its structure module, which is responsible
for predicting the individual 3D coordinates of amino acids in the protein, and
incorporates an interesting kind of equivariance, not unlike some of the archi-
tectures we previously discussed here. We will also discuss the specific impact
these components had, both within AlphaFold 2 and in relation to subsequent
iterations of the model.

8.4.1 Problem setup and architecture overview

We can interpret the input to AlphaFold 2 as a sequence of amino acid tokens
(from a vocabulary of twenty), along with additional information coming
from its multiple-sequence alignment (MSA) against a database of known
sequences. The MSA information allows us to introduce context on the most
similar known sequences to the given one, which proves valuable for enriching
the input sequence’s representation with key evolutionary information. Further,
if there are any similar sequences for which the folded structure is already
known, this information is extracted in the form of template features, and is
also extensively leveraged within the architecture.
While the input is a 1D sequence of amino acids (not at all unlike the 1D
grids studied in Chapter 6), eventually these amino acids will arrange them-
selves in a complicated way within 3D space, with long-range interactions
commonly occurring, e.g. whenever pairs of amino acids are close enough to
each other in 3D space in spite of being distant in the input sequence (Figure
8.16). As such, we treat the individual amino acids, u ∈ V , as nodes in a graph.
Eventually, our target is to predict the relative 3D positions of the atoms inside
a
each amino acid, x au ∈ RNu ×3 , where N au ∈ N is the number of atoms in the
amino acid u. To support such predictions, we will be endowing each of the
nodes with relevant spatial features, giving rise to a geometric graph output.
AlphaFold 2’s dataflow can be roughly partitioned into two phases. The first
phase is the embedding module, which is tasked with computing meaningful
scalar representations of individual amino acids, s u ∈ Rk , as well as pair rep-
resentations, zuv ∈ Rm across pairs of amino acids u, v ∈ V . This information is
then provided to the structure module, which is tasked with inferring 3D spa-
a
tial positions of all atoms, xau ∈ RN u×3 , as well as an intrinsic estimate of the
model’s confidence in each amino acid’s predicted position.
244 Chapter 8

Tu T Tw
v
Tl

Tk

Tm

Figure 8.16
While proteins are formed out of polypeptide sequences of amino acids (i.e., a 1D
sequence), its structure folds in 3D space, allowing distant atoms in the sequence to
more directly interact. AlphaFold represents each amino acid as a separate node in a
geometric graph, and iteratively refines a local coordinate frame, Tu = (R u, t u), for each
of them. Once the frames are computed along with any relevant torsion angles, they can
be used to directly estimate positions of individual atoms within the molecule.

Importantly, once the 3D atom positions are first predicted, they are not
immediately presented as the final result. Instead, the model undergoes several
rounds of recycling, wherein the two-phase architecture is iterated for multiple
steps, and at each step the predicted embeddings and coordinates20 from the
previous step are directly re-used as part of the input features. Initially, these
previous-step embeddings are initialised to zero-vectors, and previous-step
coordinates are all initialised to the origin (3D zero-vector).
Note that, similarly to our discussion throughout this Chapter, whenever the
recycled coordinate information is used within the embedding module, it is
always converted to pairwise-invariant properties first (such as distances, duv =
∥xu – xv ∥), to preserve the equivariance properties that will be guaranteed by
the structure module.

8.4.2 The structure module, at a glance

As described, the structure module is tasked with mapping individual amino
acid embeddings s u and pair embeddings zuv into atom coordinate predictions
xau and their associated confidence estimates. To maintain a focussed discus-
sion, here we will direct our attention only on parts of the structure module
which are responsible for estimating the coordinates.
While the structure module largely operates on amino acid level, it eventu-
ally will need to make predictions for individual atoms within each amino acid.
Geometric Graphs 245

To allow for this transition to proceed seamlessly, the structure module mod-
els each amino acid through its coordinate frame, T u = (R u, t u), which details
how to translate local coordinates for each amino acid to global (shared) ones
across the entire protein. The frame relies on a translation vector tu ∈ R3 to
denote the (negated) global coordinate of the alpha-carbon atom, x Cα u = tu ,
which is present robustly in every amino acid and hence serves as an excellent
anchor point around which other atoms can be organised. Additionally, the
rotation matrix Ru ∈ R 3×3 determines how the other atoms are oriented in 3D
space relative to the alpha-carbon. Please see Figure 8.16 for an illustration.
Cautious readers have likely already noticed that there is a very interesting
similarity between the coordinate frames within AlphaFold 2 and the node-
centered gauges studied in geodesic and mesh CNNs: in both cases, the model
keeps track (implicitly or explicitly) of local frames of reference. A key differ-
ence is that, for a 3D folded protein, we generally aren’t allowed to arbitrarily
modify each reference frame in isolation, as we did while studying gauge sym-
metries previously, given that this will likely produce chemically implausible
atom positions. Rather, once optimal frames are determined, we can only act
on them globally, applying exactly the same roto-translation to all of them.
Similar to the previously mentioned recycling idea, Tu is not predicted in a
single shot, but it is estimated iteratively, starting with all frames locked onto
the same orientation, and placed at the origin: Tu = (I, 0). Then, at each step:
1. Individual amino acid representations, su , are updated using the currently
estimated frames, T u , in a way that is unaffected by any global roto-translations
of all the frames—via a module known as the invariant point attention (IPA).
2. Updated amino acid representations are then used to compute an update to
each frame, Tu . Given that su is guaranteed to be invariant to roto-translations
by the IPA, effectively acting as scalar features—as derived throughout this
Chapter—it is guaranteed that the frame update is roto-translation equivariant.
Finally, the model also predicts seven different torsion angles21 , αu ∈ R7×2 ,
which specify the remaining rotational degrees of freedom along individual
chemical bonds between the amino acids’ atoms. Note that the architecture
explicitly maintains a 2D representation of the angle (via its sine-cosine pair)
to make it easier to predict by a neural network in a way that plays nicely
with gradient-based optimisation. With T u and αu known, all atom coordinates
a
xau ∈ RNu ×3 within amino acid u can be computed exactly.
What remains to complete our overview of the coordinate-prediction archi-
tecture is an account of the invariant point attention module, and the subsequent
coordinate frame updates. We now study these two components in turn.
246 Chapter 8

8.4.3 Invariant point attention

As already hinted, the invariant point attention (IPA) module takes in amino
acid embeddings, su ∈ Rk , pair embeddings, zuv ∈ Rm , and current frame esti-
mates, Tu = (R u , tu). Its aim is to produce updated amino acid embeddings,
s′u ∈ Rk , in a manner invariant to globally roto-translating the frames (SE(3)).
In order to meaningfully interact with the frame features (which, effectively,
represent a choice of gauge in each amino acid), we need to produce 3D vectors
that will transform within them. However, the features we have at our disposal
so far (s u, zuv) are all scalars and will not interact in interesting ways with T u.
As such, the IPA comprises two distinct attention components – one that acts
purely on scalar features, and one that acts on (derived) vector features. We
prepare relevant query, key and value embeddings for both cases as follows:
q(0)
u = WQ 0 su k(0)
u = WK0 s u v(0)
u = WV 0 su WQ0 , WK0 , WV 0 ∈ Rk×l
q(1)
u = WQ 1 su k(1)
u = WK1 s u v(1)
u = WV 1 su W Q1 , WK 1 , WV1 ∈ Rk×3
Note the important distinction – the scalar (type-0) queries/keys/values all have
l-dimensional features, whereas their (type-1) vector versions are strictly three-
dimensional, to allow for clearly specified interactions with their gauges.
The pair embeddings are used mainly to derive an attentional bias, not
unlike concurrent proposals such as the Graphormer (C. Ying et al. 2021):
buv = w ⊤ m
B zuv, where wB ∈ R . Once all of these quantities are known, we
combine them into an (unnormalised) attention coefficient accordingly:
⊤ (0)  
(1)2
euv = C1 q (0) k + b uv +C2 ρ(T u)q(1)
u – ρ(Tv )k v (8.150)
| u {z v } |{z} | {z }
scalar attention pairwise bias vector attention
for appropriately chosen constants C1 , C2 ∈ R. ρ(T u) is the usual representa-
tion of the roto-translation from Tu’s frame into global coordinates: ρ(Tu)x =
R ux + tu . Note this equation can distribute across multiple attention heads (as
typically done in Transformers) as well as multiple query points (in which
case q (1) (1) (1) 3
ui , k ui , vui ∈R are separately computed for 1 ≤ i ≤ P, and recombined
PP  (1)

(1) 2
additively: C2 i=1 ρ(Tu )qu – ρ(T v )kv ).
Once the attention coefficients are computed, they can be normalised across
each receiver node via the softmax function:
αuv = softmax v (e uv) (8.151)
and then used to appropriately recombine the value scalars, value vectors, and
pair embeddings:
X X X
h (0)
u = αuv v(0)
v h(1)
u = ρ(Tu )
–1
αuv ρ(Tv )v(1)
v hpair
u = αuv zuv
v∈V v∈V v∈V
(8.152)
Geometric Graphs 247

Noting that we need to appropriately transform the vector values to follow

their corresponding frames, as we did with gauge-equivariant models. Here, we
denote by ρ(T u) –1x = R –1u (x – tu), converting from the global coordinates back
into Tu ’s frame. Finally, a simple affine combination of all of these features
(with residual) yields the updated node-wise representation of amino acids:
  (1) 
s′u = s u + W out h(0) (1) pair   + bout (8.153)
u ⊕ hu ⊕ h u ⊕ hu

where ⊕ denotes concatenation, Wout ∈ R k×(l+3+m+1) is a learnable weight

matrix, and b out ∈ R l a learnable bias vector. If we distribute this across H
heads and P value points, we would concatenate all the embeddings into a
vector of H (l + 3P + m + P) dimensions (also extending W out accordingly).
Given that this system is careful in how it transports vectors to and from
local frames, and we generally provide appropriate treatment of type-1 vec-
tors (as discussed several times throughout this Chapter), perhaps it comes
as no surprise that the updated features s ′u ∈ Rl are invariant to any global
roto-translation. We provide brief proofs that all vector- and frame-dependent
operations are indeed invariant to any global roto-translation, Tglob . First, we
show this for the vector attention component:
  
(1) 2
∥ρ(Tglob )ρ(Tu )q(1) (1) 2  (1)
u – ρ(T glob)ρ(T v )kv ∥ = ρ(Tglob ) ρ(T u)q u – ρ(Tv )kv

= ∥ρ(T u)q (1) (1) 2

u – ρ(T v)kv ∥

where we can cancel out the application of ρ(Tglob ) because (a) rotations do not
change norm, and (b) global translations get cancelled out by the subtraction.
Next, we prove that Tglob does not affect the updated type-1 vectors, h(1)u :
X
ρ(Tglob )h (1) –1
u = ρ(Tu) ρ(Tglob)
–1
αuv ρ(Tglob )ρ(Tv)vv(1)
v∈V
X
–1 –1
= ρ(Tu) ρ(Tglob) ρ(Tglob ) αuv ρ(Tv)vv(1)
v∈V
X
= ρ(Tu) –1 αuv ρ(Tv)vv(1) = h(1)
u
v∈V

Given these two invariances, we can conclude that s′u will be unaffected by any
global transformations of the local Tu frames. Armed with such an update, we
can now equivariantly update the frames in response to s′u.

8.4.4 Frame updates

In order to use the new amino acid embeddings to relocate their frames,
one needs to compute a new roto-translation ∆Tu = (R̃u ,t̃ u), indicating the
additional roto-translation derived from s′u. AlphaFold 2 does not predict
rotation angles directly from this vector—instead, it computes a quaternion,
248 Chapter 8

p
(1, iu , ju, ku )/ 1 + i 2u + j2u + k2u , and the translation vector, t̃ u, directly from u
(e.g. via a standard linear layer). We can convert a quaternion (au, bu , cu, d u)
to a 3D rotation matrix like so:
 2 2 2 2 
au + bu – cu – d u 2bu cu – 2au du 2b udu + 2auc u
R̃ u =  2bu cu + 2au du a2u – b2u + c2u – d2u 2cudu – 2a ub u  (8.154)
2b udu – 2au c u 2cu du + 2a ubu a2u – b 2u – c 2u + du2
Note that by keeping the first (real) quaternion component at 1, we ensure a
valid rotation is applied. Once ∆Tu is known, we simply update Tu as:
T u′ = (Ru R̃u, t u + R ut̃ u) (8.155)
Since the new quaternion entries and translation vector are derived from an
invariant quantity, they are themselves invariant to global roto-translations. We
can use this property and the associativity of SE(3) to prove that Tu′ transforms
equivariantly. If a global transformation Tglob is applied to the system, the
base frame representation becomes ρ(Tglob)ρ(Tu), while the predicted update
representation, ρ(∆Tu ), remains unchanged. The resulting updated frame is:
(ρ(Tglob)ρ(Tu))ρ(∆T u ) = ρ(Tglob)(ρ(Tu )ρ(∆Tu )) = ρ(Tglob )ρ(T ′u) (8.156)
This mathematically guarantees that applying a global transformation to the
input precisely applies the same global transformation to the updated output
frames, strictly preserving equivariance.

8.4.5 Perspectives
We have successfully demonstrated how AlphaFold 2, a truly seminal architec-
ture for protein folding, makes use of equivariant computation to meaningfully
maintain local coordinate frames for every amino acid within the input protein.
While this is in many ways a satisfying derivation, and one that very neatly
concludes the flow of this Chapter, it is also important to contextualise it: from
the perspective of AlphaFold 2’s overall implementation, it is a relatively small
contribution. As such, and to inspire further discussions, we conclude this case
study with our own perspective on how important this modification truly was,
and whether it continues to be relevant past the Nobel prize-winning moment.
Ablations One of the key scientific traits of the AlphaFold 2 paper is its
extensive set of ablation studies: evaluating the model’s predictive power if
some of its components are meaningfully held-out; for example, in the case
of holding out equivariance, the IPA-based frame update simply gets replaced
with directly predicting the rotation quaternions from non-invariant vectors.
In this case, the story definitely does not spell a clear-cut big win for equiv-
ariance; rather, it is a matter of “many small bits adding up”: the use of IPA,
Geometric Graphs 249

as well as the use of recyling, and many other components, yields significant-
but-marginal improvements. This does not imply that IPA is not useful, but
its impact might be interpreted as one in a series of incremental gains. Fur-
ther, there is certain evidence that the IPA synergistically combines with other
innovations in AlphaFold 2: for example, ablating out both IPA and recycling
yields significant, large performance reductions, even though removing both
of them in isolation was relatively marginal. Of course, this result comes with
important caveats, one of which is that we do not know the relative contribu-
tion of the two components to this effect; that said, it does at least indicate that
explicit equivariance can be a useful supplement to other methods.
The next generation Lastly, another important piece of the story behind why
the overall impact of IPA might be called into question is the design of this
protein folding system’s next generation, AlphaFold 3 (Abramson et al. 2024).
A key domain extension offered by AlphaFold 3 is the ability to handle systems
of molecules, including protein-ligand combinations. As ligands, for example,
do not need to be proteins, we no longer have it as a given that every node of our
geometric graph will be an amino acid. This yields the rigidity we established
with AlphaFold 2’s local frames, and requires transitioning into “pure” atom-
based systems in general. In these settings, keeping track of explicit frames
and maintaining their equivariance is not feasible, and AlphaFold 3 indeed has
no roto-translation equivariant components.
We offer three possible ways in which this result can be interpreted, while
stressing that determining the exact interplay between equivariant systems and
generating molecular coordinates is an actively ongoing discussion, and, in our
opinion, insufficient time has passed to make a clear conclusion:
• Given sufficiently broad training data, non-equivariant models can
learn a sufficiently-approximated equivariance to be useful. We believe
that this interpretation is valid: AlphaFold 3 is a clear existence proof. How-
ever, there might be significant nuance about just how much training data is
needed 22, and also what levels of symmetry deviations are acceptable 23 in the
model’s outputs.
• AlphaFold 3 still leverages geometry—just in a less rigid way. Specifi-

cally, roto-translations from SE(3) are not completely absent from the model:
they are heavily used in the model’s data augmentation pipeline, making sure
that the model sees a lot of bespoke examples illustrating the very kind of
equivariance we are holding out. That said, this is no guarantee that such tricks
will remain necessary in the future.
• AlphaFold 2 learned to “crawl” so AlphaFold 3 could learn to “run”.

Admittedly, saying that a Nobel prize-winning architecture is only “crawling”

250 Chapter 8

feels very off-putting, but we do believe there is a good analogy to be made

here. A very common development template of machine learning models is
to start them off with lots of human knowledge baked in—AlphaGo (Silver
et al. 2016) first learned to play Go from human games only, and only then
refined itself through self-play, after which it was able to completely yield the
human data altogether (Silver et al. 2017). Perhaps the first major breakthrough
in automated protein folding required baking in several important biases from
humans’ understanding of the underlying physics, but subsequent iterations
will not need that?
We hope these points give interesting food for thought to our readers, as much
as they’ve given to us. If nothing else, this is a clear indicator that there might
be significant nuance in some applications of geometry in deep learning; and
another interpretation is that we are simply “holding it wrong” if we restrict
ourselves to architectural design based on symmetry groups. This proposition
will be further broadened in our forthcoming, final Chapter, where we will
study what lies beyond geometric deep learning, both in terms of more general
structures, and moving beyond pure architectural constraints.

Suggested Further Reading

Introducing roto-translational symmetries into graph-structured inputs clearly
vastly expands the variability of architectural design choices. Since this kind
of symmetry acts directly on the data rather than the domain, this variability
mainly arises from the various types of data we may need to act symmetrically
towards. We already looked at some meaningful axes from which this arises:
different neural network architectures for different tensor orders (the ℓ in type-
ℓ vectors), as well as the number of nodes taken into account at once (two
for distances, three for angles, etc.). We invite the reader to consult Duval et
al. (2023), which carefully outlines a taxonomy of this kind—this paper was,
in fact, extremely helpful to us while settling on a narrative for this Chapter.
This variety only becomes more rich if we also want to account for models
that meaningfully understand manifold geometry, or correctly handle gauge
transformations on meshes. Our aim in this chapter was mainly to “rebase”
the key findings in this spaceonto the graphs, and we cannot claim to have
provided a truly thorough treatment of manifold neural networks or gauge
equivariance. Fortunately, Maurice Weiler et al. (2023) have written a fantas-
tic book that carefully focusses on exactly this topic, and we warmly direct
interested readers towards it.
Lastly, the fact we provided a rudimentary coverage of geometric graphs
allows a more natural point to introduce a bit of GDL history: Bronstein et
 Geometric Graphs 251

al. (2017)’s original geometric deep learning-branded position paper. In this

writeup, the focus was largely on moving beyond simple Euclidean geometries,
while still staying within the confines of spatial geometry. The original GDL
paper does not dwell as much on general-purpose symmetry constraints (in
fact, the word “equivariance” is mentioned only four times there!) as it does on
spatial (and also spectral!) data understanding. In the very least, this paper can
offer useful early insight into how our current material originally came to be.

Exercises

1. Throughout this Chapter, we discussed how roto-translations g ∈ SE(3) act on the

data (ρ(T u) = Rx + t), but recall that the real representation ρ should output a matrix.
What would this matrix look like?

2. A promising quantity for building invariant geometric GNNs is the cross prod-
uct of displacements, c = (x v – xu ) × (x w – xu ). What are possible issues that might
arise with using them in geometric architectures? When x u ∈ R 3 , are they E(3)
equivariant?

3. The Chapter demonstrates how a general 3 × 3 matrix can be decomposed into type-
0, type-1, and type-2 vectors. Consider the outer product of two 3D vectors a, b ∈
R 3 , giving M = ab⊤ . Decompose M into its type-0 (trace), type-1 (antisymmetric),
and type-2 (symmetric traceless) components, writing them out explicitly in terms
of the components of a and b.

4. Given that Tensor Field Networks can express the totality of point cloud geometric
architectures, one should be able to map to-and-from that formalism. Express the
vector-update E(3)-equivariant GNN of Satorras, Hoogeboom, and Welling (2021)
using the TFN equation, and show which simplifications might be needed to the
TFN equation to recover models like SchNet.

5. We discussed how Clebsch-Gordan coefficients are necessary to preserve SE(3)

equivariance when performing products of type-ℓ spherical harmonic filters with
type-k vector features. Provide a simple counterexample set of node inputs, for
which a direct product (without using coefficients in between) of this kind would
break the symmetry.

6. In order to implement the message passing equation used by gauge-equivariant mesh

CNNs, it is necessary to know the relative angles, ϑuv , between all neighbours v ∈ Nu
in node u’s tangent plane, T uM, as well as the corresponding transporter matrices,
ρ(gv→u ). Let nu ∈ R 3 be a unit-norm vector perpendicular to Tu M and pointing
outward from the mesh.

(i) Considering the two nodes’ absolute 3D coordinates, x u, x v ∈ R 3, in order to com-

pute these angles, we need to first project node v into Tu M (a function known as the
logarithmic map). Provide an equation of this function, logu (v).
 252 Chapter 8

(ii) For any other neighbour, w ∈ Nu, show how to compute ϑ uw under the assumption that
node u’s gauge assigns ϑuv = 0.
(iii) Given that transporter matrices account for mapping one node’s gauge within another
node’s tangent space, we know that ρ(gv→u ) must be a rotation matrix by some 2D angle
φ. Express φ under the assumption that nu = n v, that is, that the tangent spaces of u and
v are perfectly parallel.
(iv) What would the transporter look like if nu ̸= n v?

7. AlphaFold 2 updates frames by predicting a quaternion q = (1, iu , j u, k u ), where

i u , ju, k u ∈ R are all derived from a neural network. If the neural network pre-
dicts (iu , j u , k u) = (0, 0, 0), what is the resulting rotation matrix R̃u ? What is the
significance of this behaviour for the model’s learning dynamics?

Notes
1
Depending on the use case, we might restrict ourselves to the special
Euclidean group, SE(3), which excludes reflections, or the special orthogonal
group SO(3), which excludes both reflections and translations—we discussed
this group at length when constructing spherical CNNs in Chapter 7.
2
The 3D translation vectors studied here form a very different kind of sym-
metry compared to the (discrete or continuous) domain translations studied in
Chapters 6 and 7. This is because the translations within E(3) operate directly
on the data channels, C , of signals, X (Ω, C )—they inherently modify the data
contained within the signal. In contrast, domain translations operate directly
on the domain, Ω, and there are no inherent changes to the data within the
signal—the data is merely reordered according to the given domain translation.
3
We broadcast the translation vector t across the node axis in X.
4
Such properties are known as isometries.
5
It is interesting to note that, with appropriate care, cross products can also
produce vectors that transform equivariantly with respect to orthogonal trans-
formations and translations—however, they are only well-defined in specific
dimensionalities (d = 3 or d = 7), whereas the layer described in this section
will work for any d.
6
Strictly speaking, at this point, our feature is a node tensor, M ∈ Rn×d×d .
7
Note that, as in the case of vector features, the translations are assumed to
have an invariant effect. This will remain true for higher-order tensors as well.
8
Here the “D” stands for Darstellung, meaning “representation” in German.
Note that by default, these matrices are complex-valued, though the real-valued
version is common in equivariant deep learning frameworks such as e3nn
(Geiger and Smidt 2022b).
Geometric Graphs 253

9
In the given equation, we number the indices (m 1, m2 ) of type-ℓ vectors’
features from 1 till 2ℓ + 1, as one would typically expect; however, when per-
forming computations using spherical harmonics, it is much more common to
use a symmetric index range: from –ℓ to ℓ.
10
One might be tempted to also include the displacement in the queries, but
since we can only really control the scalar multipliers involved, it is possible
to show that this will not have a tangible change on the result of the query-key
dot product.
11
In honour of Bernhard Riemann, who introduced the concept in 1854.
12
The Hopf-Rinow theorem therefore establishes the equivalence between
geodesic completeness and metric completeness, the latter meaning every
Cauchy sequence converges in the geodesic distance metric.
13
In fact, the framework of MoNet predates graph attention networks
(Veli čković et al. 2018). The main innovation of GATs was to also include
node features as an input to the attention function.
14
One might be tempted to propose a “canonical gauge collection”
which smoothly changes everywhere along the mesh, yielding a consis-
tent set of coordinates. However, many manifolds (including spheres) are
non-parallelisable, meaning that one cannot define a smooth, non-vanishing
tangent vector field on them. The result on spheres is a special case of the
Poincaré-Hopf Theorem, colloquially referred to as the ‘hairy ball theorem’.
Put directly, ‘one cannot comb a hairy ball without creating a cowlick’.
15
Typically, this projection from neighbouring nodes back into the tangent
space is done via an opposite process to the exponential map, expu – known as
the logarithmic map.
16
The idea of discretising points along the surface of the sphere and
performing message passing along the resulting mesh is common in many
state-of-the-art architectures; a great example is the GraphCast model for
weather forecasting (Lam et al. 2023), which relies on multi-scale message
pasing along the Earth surface mesh.
17
AlphaFold 2 was cited as the “solution” to the protein folding problem, after
its overwhelming victory at the CASP14 competition, with very low error rates
compared to the experimentally observed structures at the competition.
18
The second half of the 2024 Chemistry Nobel Prize was awarded to
David Baker, whose group also pioneered several deep learning architectures
for protein analysis and design relevant to the scope of this book, such as
RoseTTaFold (Baek et al. 2021) and ProteinMPNN (Dauparas et al. 2022).
19
The authors warmly recommend reading the Supplementary Methods file
of the original paper by Jumper et al. (2021); this contains sixty-two pages’
254 Chapter 8

worth of information about the model architecture, inductive biases and physi-
cal priors, as well as countless ablations; a deep learning project writeup of the
highest quality.
20
In reality, only the coordinate of the amino acid’s beta carbon atom,
xCβ 3
u ∈ R (or pseudo-beta carbon for glycine), is recycled. This is because we
need one input per amino acid, and the beta carbon preserves all the necessary
information about the rotational orientations of individual amino acids.
21
As torsion angles are invariant quantities, they can be directly predicted
from su with a standard multilayer perceptron.
22
This might also relate to the size of the relevant symmetry group; what’s a
“large enough” training dataset to grok SE(3) is likely insufficient to grok Sn .
23
For example, if we need to feed the model’s outputs into a simulator
that depends on coordinates, small deviations in the coordinates might lead
to significant deviations after simulating.