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 sufﬁcient to do so.

•

Combining permutation symmetries (S

) with continuous symmetries of

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 ﬁrst

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) deﬁned 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

). As we will explore in this Chapter, this

is often far from the full story, from the geometric perspective.

The insufﬁciency 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

∈R

. Molecules were the initial motivating example

for the graph domain in Chapter 5, yet they are also a common ﬁrst 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-

eﬁt from explicitly understanding the spatial arrangement in such geometric

graphs. This will be the ﬁnal domain we consider on our geometric journey.

210 Chapter 8

H H

C C

Figure 8.1

Based on its molecular graph (left), cyclooctane (C

) 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 speciﬁc 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, s

∈R

, but also with spatial

coordinates, x

∈R

. 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 S

. The exact way in which we expand will depend

on the speciﬁc assumptions on which quantities are equivariant under which

transformations of the given spatial coordinates, x

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 ﬁelds,

that property should remain the same even if we rotate, translate or reﬂect all of

the molecule’s atoms. This collection of roto-translations and reﬂections in 3D

space is the Euclidean group

, E(3). For any 3D orthogonal matrix Q ∈R

3×3

and translation vector

t ∈R

, we deﬁne a corresponding group action g ∈E(3)

operating on the coordinate space:

ρ(g)x

= Qx

+ t (8.101)

Crucially, applying this spatial transformation to the coordinates typically

does not affect the abstract node features s

(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 deﬁne 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 ∈R

n×n

and coordinate matrix

X ∈R

n×3

. Then, to be E(3)-invariant, such a layer must satisfy:

F(S, A, XQ + t) = F(S, A, X) (8.102)

for all possible

(Q, t) ∈E(3) – that is, its output should be unaffected by

roto-translations and reﬂections 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

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

×E(d). We can combine these constraints into just one, speciﬁed for

every (P, (Q, t)) ∈S

×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/reﬂecting 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 reﬂection

of the

coordinates—this trivially guarantees E(d) invariance for any dimensionality

d. By our assumption, s

are such features, whereas we can feed in distance

information, ∥(x

– x

)∥, 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):

= ϕ

v∈N



, s

, ∥x

– x

∥



(8.104)

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 ∈N

with respect to the receiver, u, expressed via the dot

product (x

– x

)

⊤

– x

). Incorporating this information requires explicitly

considering pairs of senders at once, modifying the GNN equation to:

= ψ





, s

, ∥x

– x

∥,

w∈N

\{v}



, s

, ∥x

– x

∥, (x

– x

)

⊤

– x

)







(8.105)

= ϕ

v∈N

, (8.106)

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 insufﬁcient 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,

∥x

– x

∥, might not have been expressive enough, the actual raw values of x

and x

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 insufﬁcient for

distinguishing them.

molecule yields different ﬁnal predictions. Since we cannot change how we

update the latents, h

, we need to ﬁnd a way to update the representations of

our coordinates, x

, in a way that will not affect how the latents, h

, 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

∈R

, and updating the node coordinates to

′

∈R

. 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 reﬂections, yielding an E(d) equivariance constraint:

F(PS, PAP

⊤

, PXQ + t) = (PH, PX

′

Q + t) (8.108)

for all permutation matrices P ∈R

n×n

, orthogonal matrices Q ∈R

d×d

and

translation vectors t ∈R

(see Figure 8.3). The ﬁnal-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

′

directly as the output—this is typical of n-body problems.

214 Chapter 8

′

ρ(180

◦

)

ρ(180

◦

)

F F

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

∈R

to h

∈R

) and the node coordinates (from x

∈R

to x

′

∈R

), in a way that is equivariant to actions from E(2). In this case, a speciﬁc

180

◦

rotation of node coordinates is considered. As exempliﬁed by the commutative

diagram, an E(2)-equivariant model must yield the same answer regardless of whether

the GNN or the rotation is applied ﬁrst.

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 speciﬁc 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 ∈R

n×m

, and coordinates, X ∈R

n×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 s

∈R

into latents h

∈R

. Matters get signiﬁcantly more interesting

when extending this model to (equivariantly) update the coordinates.

To see how, ﬁrst recall that our aim is for the updated coordinates to roto-

translate and/or reﬂect 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 reﬂections) as well as translations. An operation that covers all

of the above is a linear combination

of displacement vectors, x

– x

, as in:

′

= x

v̸=u

(

– x

)

(8.109)

where β

∈R is a scalar coefﬁcient, which we assume to be ﬁxed upfront

for now. It’s a good idea to quickly reﬂect on why these speciﬁc opera-

tions satisfy the constraints in Equation 8.108. Firstly, linear combinations

of this kind will commute with matrix products (orthogonal or otherwise!):

Q(x

v̸=u

– x

)) = Qx

v̸=u

(Qx

– Qx

). And secondly, the

displacement operator will naturally cancel out any impacts of translation vec-

tors: (x

+ t) +

v̸=u

((x

+ t) – (x

+ t)) = (x

v̸=u

– x

)) + 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 speciﬁc

neighbourhoods N

will also sufﬁce.

Keeping β

ﬁxed, however, is an underwhelming choice—we haven’t made

advantage of the possible wealth of information stored in our nodes. Accord-

ingly, we can deﬁne an attentional mechanism neural network which will

leverage the s

and x

vectors to dynamically compute the coefﬁcients of

the combination. However, this mechanism cannot take arbitrary inputs—for

example, if raw coordinates, x

, 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

coefﬁcients to be robust to any roto-translations or reﬂections of x

—meaning

it needs to be E(d)-invariant by deﬁnition! We deﬁnitely 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 modiﬁcation (Equation

8.114).

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

= ϕ

v∈N



, s

, ∥x

– x

∥



(8.110)

′

= x

v̸=u



, s

, ∥x

– x

∥



(

– x

)

(8.111)

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 coefﬁcients

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, reﬂections and translations of the

input. The node features, s

, 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

∈R

. Note that vectors transform together with any

rotation or reﬂection of coordinates but remain invariant to translations. This

justiﬁes 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 ∈R

n×d

, and

the output vector features v

′

∈R

. 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, v

, with invariant coefﬁcients, to ensure

equivariance to orthogonal matrix products. A useful addition is that, as the

displacements, x

– x

, are themselves vector features, we may also include

them in the update rule. All together, we recover the following equation:

′

= v

v̸=u



, s

, ∥x

– x

∥



v̸=u



, s

, ∥x

– x

∥



(

– x

)

(8.114)

where α and β are neural networks outputting scalar values. The speciﬁc 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 features

, M

∈R

d×d

, which

may arise as outer products M

= a ⊗b of two vectors a, b ∈R

. Such inputs

transform under E(d) as QM

⊤

for every orthogonal matrix, Q, inducing

appropriate equivariance constraints

. 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 ∈R

d×d

, with 3D rotation matrices, R ∈

3×3

, and still including 3D translations, t ∈R

. This restriction is done

merely for pedagogical reasons of ease of tracking mathematical state, and

is not strictly required for building tractable architectures.

8.2.3.1 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 ∈R

, type-

2 vectors are matrices M ∈R

3×3

that are symmetric (M = M

⊤

) and traceless

Geometric Graphs 219

(tr(M) = 0), meaning they can be written using only ﬁve parameters w ∈R

M =





–(w

+ w

)





(8.115)

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, ﬁrst note that we can use a type-0

vector, s ∈R, to encode

tr(N). Then, we can use a type-1 vector, v ∈R

, to

encode the antisymmetric components—the non-zero elements of

(N – N

⊤

(N – N

⊤

) =





– n

)

– n

)

– n

) 0

– n

)

– n

)

– n

) 0









0 v

–v

0 v

–v





(8.116)

Once we subtract the trace from N’s diagonal, and the antisymmetric part

(N – N

⊤

) from N, we obtain

N = N – sI –

(N – N

⊤

N =





– s n

– v

+ v

– s n

– v

+ v

– s









– s

+ n

)

+ n

)

+ n

) n

– s

+ n

)

+ n

)

+ n

) –(n

+ n

– 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

ℓ

∈R

2ℓ+1

. Note that here we will be ignoring

the channel axis (e.g. we will assume there is exactly one type-0 scalar feature,

∈R, exactly one type-1 features, f

∈R

, etc.) for simpliﬁcation 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∈V,ℓ≥0

→{h

ℓ

}

u∈V,ℓ≥0

. We will once again,

for simplicity, assume that the latent space h

ℓ

∈R

2ℓ+1

has no channel axis.

8.2.3.2 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

, D

ℓ

(g) ∈R

(2ℓ+1)×(2ℓ+1)

(Wigner 1927). Speciﬁcally, for

any representation, ρ, and change-of-basis orthogonal matrix, Q:

ρ(g) = Q

⊤

ℓ∈L

ℓ

(g)

Q (8.118)

where

: R

a×a

×R

b×b

→R

(a+b)×(a+b)

is the direct sum, and L⊆{0, 1, . . . } is

the set of vector orders that we are acting on (i.e. the speciﬁc type-ℓ vectors

we have access to for our task). Note that the shape of Q ∈R

N×N

is then spec-

iﬁed 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 deﬁned exactly by using this property.

As the D

ℓ

(g) matrices can be efﬁciently 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 f

and f

. 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.

8.2.3.3 Constructing equivariant ﬁlters via spherical harmonics The

ﬁrst stage of any message passing layer is, for each edge (u, v) ∈E, to convert

sender features, f

, into messages, m

. 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

∈R

, are 3D vectors and hence part of the

Geometric Graphs 221

type-1 node features, f

. To make the computations explicit, we will treat x

as a completely separate vector from f

, but please note that it is completely

possible to track and update coordinates within f

; 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 r

= x

– x

Keeping in mind that the type-ℓ message, m

ℓ

, may itself depend on fea-

tures of several types, f

, we need to construct appropriate ﬁlters W

ℓk

: R

→

(2ℓ+1)×(2k+1)

. that will allow us to combine our messages as

ℓ

k≥0

ℓk

(8.119)

The ﬁlter, W

ℓk

), 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

ﬁlter 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 ﬁlter should behave under roto-translation

in a principled manner, we separate the displacement into its magnitude and

direction components: r

= ∥r

∥

, where

= r

/∥r

∥ is the normalised

displacement vector.

The magnitude ∥r

∥ 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

∥



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

, ﬁrst

of all, note that it is now deﬁned on the sphere, S

. Since we need to produce a

ﬁlter of type ℓ, and type-ℓ vectors have 2ℓ + 1 entries, we require a function of

the form Y

ℓ

: S

→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:

ℓ

) = D

ℓ

(g)Y

ℓ

(

), (8.120)

hence guaranteeing rotation equivariance. Note that spherical harmonics are

fundamental basis functions on the sphere (with signiﬁcant implications

to chemistry, as they represent various quantised conﬁgurations of atomic

orbitals), and are readily pre-computable. Therefore, as was the case for

222 Chapter 8

Figure 8.6

Visualisation of the ﬁrst 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 ﬁrst half of the harmonic dimensions (including the 0

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 speciﬁc 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-ℓ ﬁlter vector, φ

ℓ



∥r

∥



ℓ

(

)

∈R

2ℓ+1

. The only part that remains

to specify our ﬁnal ﬁlter, W

ℓk

), is the describe how this ﬁlter vector

needs to interact with the sender’s type-k feature vector, f

, while preserving

equivariance. We study how to principally construct such interactions next.

8.2.3.4 Feature-ﬁltering interactions via Clebsch-Gordan coefﬁcients

Thus far, by leveraging spherical harmonics, we have constructed an equivari-

ant, type-ℓ ﬁlter vector, which we will denote by ϕ

ℓ

∈R

2ℓ+1

. Now, we would

like to explicitly interact this ﬁlter vector with a given neighbour feature vector

Geometric Graphs 223

of type k, f

∈R

2k+1

. More generally, for any given pair of type-ℓ ﬁlter 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 f

, we obtain all possible pairwise

multiplications through their tensor product, ϕ

ℓ

⊗f

∈R

(2ℓ+1)×(2k+1)

, specif-

ically,



ℓ

⊗f



= ϕ

ℓ

)

. The resulting vector then transforms according

to the product representation of relevant Wigner D-matrices, D

ℓ

⊗D

. 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 ϕ

ℓ

⊗f

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 decomposition

as follows:

(ϕ

ℓ

⊗f

)

2ℓ+1

2k+1

ℓ,m

,k,m



ℓ

⊗f



2ℓ+1

2k+1

ℓ,m

,k,m

ℓ





(8.121)

through the Clebsch-Gordan (CG) coefﬁcients, C

ℓ,m

,k,m

∈R. One way to

read this notation is: “the coefﬁcient sending the result of cross-multiplying

the m

-th component of a type-k vector with the m

-th component of a type-ℓ

ﬁlter, to the m-th component of a type-j vector result”. CG coefﬁcients offer a

recipe for a change-of-basis operation that can be precomputed and kept ﬁxed.

As such, similarly to Wigner D-matrices and spherical harmonics, many estab-

lished equivariant learning libraries support computing these coefﬁcients out

of the box, and we will not dwell further on how CG coefﬁcients are con-

structed. The CG coefﬁcients are the ﬁnal piece needed to build our ﬁlter

matrices W

ℓk

). Combining all we’ve seen so far, with slightly modiﬁed

notation:

ℓk

) =

ℓ+k

j=|ℓ–k|

(∥r

∥)

2j+1

m=1



∥r

∥



ℓ

jmk

(8.122)

In this equation, j ranges over all type-j ﬁlters that will yield a message of

type ℓ when applied to type-k vector features. The radial neural network φ

scales the intensity of the interaction, the spherical harmonic Y

provides the

directional spatial awareness, and the matrix C

ℓ

jmk

∈R

(2ℓ+1)×(2k+1)

comprises

224 Chapter 8

0,:,0

1,:,1

2,:,2

1,:,0

0,:,1

1,:,1

2,:,1

1,:,2

2,:,2

2,:,0

1,:,1

2,:,1

0,:,2

1,:,2

2,:,2

Σ Σ Σ

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, f

ℓ

, contribute to which message orders, m

ℓ

, as well as which Clebsch-Gordan

coefﬁcients, C

ℓ

j,:,k

handle each part of the communication. Note that this diagram, for

simplicity, does not include any speciﬁcs about the ﬁltering operation through the

spherical harmonic. Rather, the tensor product with the appropriate order-j spherical

harmonic, Y

(

), is implicitly computed within each coefﬁcient block.

all the relevant Clebsch-Gordan coefﬁcients, 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 ﬁltering. 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

) ﬁlters fully spans all SE(3)-equivariant layers.

8.2.3.5 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 ﬁltering operations via their orthonormal basis—the spherical

harmonics—weighted by radial neural networks; and

•

Properly disentangling the outcome of ﬁltering operations into “pure” type-ℓ

vectors, via the ﬁxed exchange parameters—the Clebsch-Gordan coefﬁcients.

Geometric Graphs 225

We are now ready to put all of the above together into tensor ﬁeld 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

ℓ

, we

simply combine the results across all neighbours:

ℓ

v∈N

ℓ

(8.123)

v∈N

k≥0

ℓk

– x

(8.124)

v∈N

k≥0

ℓ+k

j=|ℓ–k|

(∥x

– x

∥)

2j+1

m=1



– x

∥x

– x

∥



ℓ

jmk

(8.125)

where, as before, we leverage summations throughout as, per Dym and Maron

(2021), this is sufﬁcient to express all SE(3)-equivariant operators. Please note

a speciﬁc 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

ℓℓ

ℓ

. We accordingly separate this case:

ℓ

= w

ℓℓ

ℓ

v∈N

\{u}

k≥0

ℓ+k

j=|ℓ–k|

(∥x

– x

∥)

2j+1

m=1



– x

∥x

– x

∥



ℓ

jmk

(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 φ

—

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

8.2.3.6 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 coefﬁcient α

∈R. These coefﬁ-

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 ﬁnd that only simple alterations are needed to

make it attentional:

ℓ

= w

ℓℓ

ℓ

v∈N

\{u}

k≥0

ℓ+k

j=|ℓ–k|

(∥x

– x

∥)

2j+1

m=1



– x

∥x

– x

∥



ℓ

jmk

(8.127)

Speciﬁcally, we explicitly use the V identiﬁer on our parametric components

ℓℓ

and φ

) to denote that these parameters pertain to the computation of

value vectors. Further, the computed value vectors are now scaled by the atten-

tion coefﬁcient, α

. In order for the coefﬁcient’s multiplication to not disturb

the message’s equivariance properties, it needs to be SE(3)-invariant (just like

in the case of φ

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 coefﬁcient α = q(x)

⊤

k(y).

Further, assume the corresponding equivariance conditions on both q and k,

i.e., q(R

(θ)x) = R

(θ)q(x) and k(R

(θ)y) = R

(θ)k(y) for every rotation

angle θ. Then, when rotating our frame by θ, the dot product becomes:

q(R

(θ)x)

⊤

k(R

(θ)y) =



(θ)q(x)



⊤

(θ)k(y)

= q(x)

⊤

(θ)

⊤

(θ)k(y)

= q(x)

⊤

k(y) = α

where the ﬁnal cancellation is due to the orthogonality of the rotation matrix

(θ). 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

ℓ

, k

ℓ

∈R

2ℓ+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 displacement

, x

– x

. 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:

ℓ

= w

ℓℓ

ℓ

k≥0

ℓ+k

j=|ℓ–k|

(∥x

– x

∥)

2j+1

m=1



– x

∥x

– x

∥



ℓ

jmk

(8.128)

while keeping in mind that the equation above assumes u ̸= v, and we need to

fall back to scalar multiplication k

ℓ

= w

ℓℓ

ℓ

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:

exp





ℓ≥0

ℓ



⊤



ℓ≥0

ℓ





w∈N

exp





ℓ≥0

ℓ



⊤



ℓ≥0

ℓ





(8.129)

Plugging these coefﬁcients 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 ﬂexibility 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 ﬂoating 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 ﬁelds 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 speciﬁc manifolds, when deﬁn-

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 ﬁner and ﬁner, 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 ﬁeld 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 deﬁnition 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 (speciﬁcally, an ellipsoid). But, locally, right under your feet, the ground

feels perfectly ﬂat. This property is essential for deﬁning a manifold, Ω.

Speciﬁcally, 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

;

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 ﬂat 2D plane).

Having identiﬁed this patch of R

around a point, u ∈Ω, we can extend it

inﬁnitely in all possible directions, as a local approximation of the manifold,

called the tangent space, T

Ω. In the case of spheres, S

, the tangent space,

= R

, becomes the unique plane which is tangent to any given point u; see

Geometric Graphs 229

Figure 8.9. Once such a space is deﬁned, we can use it to express displacements

of various points in u’s local neighbourhood, by way of elements X ∈T

Ω,

which we call tangent vectors. In the sphere example, X ∈R

gives us the

displacement vector in the tangent plane, which we can then project onto the

sphere to ﬁnd 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 ∈T

Ω

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-deﬁnite

bilinear function g

: T

Ω ×T

Ω →R, depending smoothly on u. Such a func-

tion is called a Riemannian metric

, and it can be thought of as an inner

product on the tangent space, ⟨X, Y⟩

= g

(X, Y), which evaluates the angle

between any two tangent vectors X, Y ∈T

Ω. Using it, we can also induce a

norm, ∥X∥

= g

1/2

(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

: g

(x, y) = x

⊤

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

Ω numerically as an array of numbers, we can only represent it as a list

of coordinates x = (x

, . . . , x

) relative to some local basis {X

, . . . X

} ⊆T

Ω.

Similarly, the metric can be expressed as an s ×s matrix G with elements

= g

, X

) 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 deﬁned

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

= {u ∈R

: ∥u∥= 1}, realised a subset (sub-manifold) of R

. The

tangent space to the sphere is given as T

= {x ∈R

: 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⟩

= x

⊤

y for any x, x ∈T

. The

exponential map is given by exp

(x) = cos(∥x∥)u +

sin(∥x∥)

∥x∥

x, for x ∈T

. 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:

ℓ(γ) =

∥γ

′

(t)∥

γ(t)

dt =

1/2

γ(t)

(γ

′

(t), γ

′

(t)) dt.

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 deﬁned 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 deﬁne 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 ∈T

Ω, Y ∈T

Ω 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 ∈

Ω be a tangent vector at u. We can now deﬁne 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 ﬁnal vector

directions in C’s tangent plane.

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

γ(t)

(X(t), γ

′

(t)) = g

(X, γ

′

(0)) ∥X(t)∥

γ(t)

= ∥X∥ (8.130)

and this uniquely deﬁnes X(T) ∈T

Ω. We can consider the corresponding map,

u→v

: T

Ω →T

Ω, with Γ

u→v

(X) = X(T), as the parallel transport of tangent

vector X into T

Ω, 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

Ω, 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 deﬁne such a quantity, the exponential map, exp

: T

Ω →Ω, so

long as the geodesics on Ω can be sufﬁciently extended in any direction around

u (if we can inﬁnitely 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 ∈T

Ω. Then, we deﬁne a X-directed geodesic, γ

, such that:

(0) = u γ

′

(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”: exp

(X) = γ

(1). In the case of spheres, we rely on

projecting tangent vectors, x ∈R

, back onto the surface, using trigonometry:

exp

(x) = cos(∥x∥)u +

sin(∥x∥)

∥x∥

x (8.132)

Geodesic distance Last but not least, geodesics and their length function-

als can be used to properly deﬁne distances between pairs of points u, v ∈Ω,

assuming geodesic completeness. Speciﬁcally, 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 theorem

(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 deﬁne what it might look like to con-

volve a signal on a manifold, of the form x : Ω →R, with an appropriate ﬁlter,

θ. Analogously to the group convolutions covered in Chapter 7:

(x ⋆ θ)(u) =

Ω

x(exp

Y)θ

(Y) dY, (8.133)

Note that this convolution is expressed within the tangent space, T

Ω, 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

Ω, we look up the

ﬁlter in the corresponding point (θ

(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 ﬁlter θ now needs to be deﬁned

separately for each point’s tangent space, heavily crippling any weight sharing

prospects. Secondly, we generally need local ﬁlters, as we cannot integrate

over the entire tangent space T

Ω, 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 deﬁned as ω

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).

→T

Ω, and once chosen, they allow us to represent Y as an s-dimensional

array of coordinates x = ω

–1

(Y).

Assuming we have chosen an appropriate gauge, we can rewrite the

convolution in Equation 8.133 as

(x ⋆ θ)(u) =

[0,1]

x(exp

(ω

y))θ(y) dy, (8.134)

with the ﬁlter now deﬁned 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 ﬁltering 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, ω

, which we survey next.

Geodesic patches Assuming our input is a two-dimensional manifold within

3D space, we can deﬁne geodesic patches around a given point; these will be

excellent anchors for our ﬁlters. Let us choose a local polar gauge ˜ω

: [0, R] ×

[0, 2π) →T

Ω, using which we deﬁne x(u, r, ϑ) = x(exp

˜ω

(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. Reﬁning Equation

234 Chapter 8

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

(x ⋆ θ)(u) =

2π

x(u, r, ϑ)θ(r, ϑ) drdϑ. (8.135)

Figure 8.11 (Left) illustrates several geodesic patches resulting from this kind

of operation, over a reference manifold.

Gauge-invariant ﬁltering A key ambiguity in deﬁning 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 ﬁlter on the angle altogether, yielding a simpliﬁed ﬁlter θ(r):

(x ⋆ θ)(u) =

2π

x(u, r, ϑ)θ(r) drdϑ.

While elegant and useful, such isotropic ﬁlters 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 ﬁlter:

(x ⋆ θ)(u) = max

∈[0,2π)

2π

x(u, r, ϑ)θ(r, ϑ + ϑ

) drdϑ.

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 efﬁcient 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 ﬁrst intrinsic CNN-like architectures on meshes.

To make these models practical, it is not sufﬁcient 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 N

bins and our angle

space into N

bins. This means we ﬁre N

geodesic rays out of u that equally

partition the tangent space, and record the ﬁlter (and signal) values along those

geodesic rays, in increments of R/N

. 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 coefﬁ-

cients shown in red.

discrete angular max-pooling equation:

–1

max

∆r=0

–1

∆ϑ=0



u, R

∆r

, 2π

∆ϑ





∆r

, 2π

(∆ϑ + ϑ

) mod N



(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, exp

˜ω



∆r

, 2π

∆ϑ



∈N

). 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 ﬁrst 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 speciﬁc cases of atten-

tional GNNs, wherein attention coefﬁcients are derived from the coordinates

assigned to each point (Figure 8.12).

Another relevant way of connecting geodesic CNNs to geometric GNNs is

by observing that their ﬁlters 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

= 120

◦

Gauge transformation

g : V → SO(2)

Figure 8.13

An example of a gauge symmetry on the sphere surface. Gauges in each node’s tangent

space, T

M, are transformed according to their own assigned 2D rotation, g

∈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

ﬁne-grained notion of gauge equivariance

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, x

∈R

. As before,

we will assign a tangent space, T

M, at each point; all of the nodes in N

will be observed through this tangent space

. Then, we make arbitrary gauge

choices, ω

: R

→T

M, in each node’s tangent space: for example, we can

pick an arbitrary neighbour, v ∈N

, and assume that it encodes the base direc-

tion from u; that is, that the angle ϑ

= 0. From this choice of gauge, exact

angles, ϑ

, and tangent plane coordinates, c

∈R

, can be computed for all

nodes w ∈N

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 ﬁxed origin – a common set of symmetry transformations between such

gauges are 2D rotations, and possibly reﬂections, 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 reﬂections) 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, f

, 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 modiﬁcations 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

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, g

∈SO(2). This rotation

modiﬁes the gauge by ω

′

(ρ(g

)c), where ρ(g

) ∈R

2×2

is the corre-

sponding rotation matrix. It must transform the individual node coordinates

by c

′

= ρ(g

)

–1

, to ensure the underlying tangent vectors remain the same:

′

) = ω

(ρ(g

′

) = ω

(ρ(g

)ρ(g

)

–1

) = ω

) (8.137)

For similar reasons, this gauge transformation also must transform the input

features by f

′

= ρ

)

–1

, for appropriately chosen input representation ρ

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 : R

→R

, writing f (f

) = h

, not at

all unlike the set neural networks we discussed in Chapter 5. Further, assume

that our network’s output space transforms under ρ

out

) in each node u.

What happens to f under a gauge transformation g ∈G? f (f

) now becomes

′

(f) = ρ

out

)

–1

f (ρ

)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

′

= f

′

) = ρ

out

)

–1

f (ρ

′

) (8.138)

= ρ

out

)

–1

f (ρ

)ρ

)

–1

) (8.139)

= ρ

out

)

–1

f (f

) (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

f f

Gauge transformation

g : V → SO(2)

Gauge transformation

g : V → SO(2)







−0.7

0.7





0.7





−1



Figure 8.14

Illustrating the gauge equivariance constraint on the particular case of a pointwise func-

tion f , acting on a vector input v

∈R

(shown in red), whose coordinates are computed

relative to node u’s gauge (depicted in blue). Note the signiﬁcance of the provided com-

mutative diagram: when f is gauge equivariant, we can either ﬁrst apply f to our vector,

then rotate our gauges by an amount – or ﬁrst 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 (ρ

) = ρ

out

)f (f

) (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 signiﬁcant result:

′

) = ρ

out

)

–1

f (ρ

) = ρ

out

)

–1

out

)f (f

) = f (f

) (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,

∈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 signiﬁcantly 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 ﬁrst consider the case where all node

input features, f

, 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:

v∈V

(8.143)

Note that such an equation uses a different ﬁlter θ

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 f

∈R

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:

v∈V

(8.144)

where Θ

∈R

2×2

is 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 f

and h

are allowed to change

independently and arbitrarily, we obtain an unrealistic equivariance constraint:

= ρ(g

)

–1

ρ(g

) (8.145)

and more importantly, it must hold for all possible choices of g

, g

∈SO(2).

This constraint states that, under any change of gauge either in its input (ρ(g

))

or output (ρ(g

)) space, this ﬁlter’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 Θ

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 ﬁrst transport

240 Chapter 8

naive

v→u

Parallel transport

v → u

Figure 8.15

We want to parallel-transport a vector, x

, in v’s tangent space, T

M, into a neigh-

bour’s tangent space, T

M. Speciﬁcally, we transport x

onto u along the geodesic

speciﬁed by the mesh, possibly with additional gauge alignment, giving a transported

vector x

v→u

. It is cumbersome to do this procedure for every possible vector on-the-ﬂy,

therefore what we do instead is compute the 2D angle, g

v→u

, that x

v→u

makes with the

coordinates of x

interpreted within the gauge of v, which we call x

naive

v→u

. The result-

ing angle between x

v→u

and x

naive

v→u

is then exactly the transformation, g

v→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 ﬁlter:

v∈V

ρ(g

v→u

(8.146)

Here, g

v→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

) ∈R

2×2

. To avoid ambiguity, the geodesic from v to u must be

unique, which means the ﬁlters must necessarily have local support only (i.e.

around speciﬁc neighbourhoods N

). This ensures that our transport matrix,

ρ(g

v→u

), remains nonparametric, leaving all learnable parameters with Θ

If we perform a gauge transformation g ∈G in this setting, the neigh-

bour’s features transform as ρ(g

)

–1

, and the transport matrix transforms as

ρ(g

)

–1

ρ(g

v→u

)ρ(g

). And since our ﬁlter, Θ

, now operates only within u’s

tangent space, its equivariance constraint only requires commuting with g

ρ(g

) = ρ(g

)Θ

(8.147)

Geometric Graphs 241

Provided this constraint is satisﬁed, we can show that the overall equation is

gauge-equivariant:

′

v∈N

ρ(g

′

v→u

′

v∈N

ρ(g

)

–1

ρ(g

v→u

)ρ(g

)

–1

v∈N

ρ(g

)

–1

ρ(g

v→u

= ρ(g

)

–1

v∈N

ρ(g

v→u

= ρ(g

)

–1

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 CNNs

) and de Haan et al. (2020) (over arbitrary 2D meshes).

Let f

∈R

be a feature vector at node u ∈V. GEM-CNNs compute a latent

feature h

∈R

, through a localised form of Equation 8.146, while also noting

that tracking Θ

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, ϑ

, relative to node u’s local choice of gauge. The model is

realised as follows:

= Θ

self

v∈N

Θ(ϑ

)ρ(g

v→u

(8.148)

Here, ρ(g

v→u

) ∈R

k×k

is the previously described transporter matrix, which

parallel transports the neighbour’s input features into the receiver’s tangent

space. Note that the speciﬁc instantiation of the transporter will depend on

the speciﬁcs of the input feature space; for example, which type-ℓ vectors are

present in the input (explicitly or implicitly). Further, our parameters Θ

self

∈

m×k

and Θ : [0, 2π) →R

m×k

are chosen to appropriately commute with gauge

transformations. As discussed before, this means they must satisfy:

self

(ϑ) = ρ

out

(ϑ)Θ

self

Θ(ϑ

– ϑ)ρ

(ϑ) = ρ

out

(ϑ)Θ(ϑ

) (8.149)

where ρ

(ϑ) ∈R

k×k

and ρ

out

(ϑ) ∈R

k×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

coefﬁcients: Θ

self

(i)

self

and Θ(ϑ) =

(i)

(ϑ).

The Geometric Bridge While it took us a lot of mathematical juggling to

derive it, the ﬁnal 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, ϑ

). 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 speciﬁc 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 coefﬁcients

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 ρ(g

v→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: deﬁne the relevant (global or

local) symmetries, design ﬁlters that respect them, and use those ﬁlters 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 deﬁnitely its use in solv-

ing

the protein folding problem: predicting the 3D spatial conﬁguration 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

half

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 ﬁtting way to close out the geometric graphs Chapter.

AlphaFold 2 is, in its entirety, a rather complex system

and, given our spe-

ciﬁc 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 speciﬁc 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

each amino acid, x

∈R

×3

, where N

∈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 dataﬂow can be roughly partitioned into two phases. The ﬁrst

phase is the embedding module, which is tasked with computing meaningful

scalar representations of individual amino acids, s

∈R

, as well as pair rep-

resentations, z

∈R

across pairs of amino acids u, v ∈V. This information is

then provided to the structure module, which is tasked with inferring 3D spa-

tial positions of all atoms, x

∈R

×3

, as well as an intrinsic estimate of the

model’s conﬁdence in each amino acid’s predicted position.

244 Chapter 8

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 reﬁnes a local coordinate frame, T

= (R

, t

), 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 ﬁrst predicted, they are not

immediately presented as the ﬁnal 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 coordinates

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 ﬁrst (such as distances, d

∥x

– x

∥), 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

and pair embeddings z

into atom coordinate predictions

and their associated conﬁdence 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

= (R

, t

), 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 t

∈R

denote the (negated) global coordinate of the alpha-carbon atom, x

Cα

= t

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 R

∈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, T

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: T

= (I, 0). Then, at each step:

1. Individual amino acid representations, s

, are updated using the currently

estimated frames, T

, 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, T

. Given that s

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 angles

, α

∈R

7×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

and α

known, all atom coordinates

∈R

×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, s

∈R

, pair embeddings, z

∈R

, and current frame esti-

mates, T

= (R

, t

). Its aim is to produce updated amino acid embeddings,

′

∈R

, 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

, z

) are all scalars and will not interact in interesting ways with T

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:

(0)

= W

(0)

= W

(0)

= W

, W

∈R

k×l

(1)

= W

(1)

= W

(1)

= W

, W

∈R

k×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 speciﬁed 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):

= w

⊤

, where w

∈R

. Once all of these quantities are known, we

combine them into an (unnormalised) attention coefﬁcient accordingly:

= C

(0)

⊤

(0)

| {z }

scalar attention

+ b

|{z}

pairwise bias



ρ(T

(1)

– ρ(T

(1)



| {z }

vector attention

(8.150)

for appropriately chosen constants C

, C

∈R. ρ(T

) is the usual representa-

tion of the roto-translation from T

’s frame into global coordinates: ρ(T

)x =

x + t

. 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)

, k

(1)

, v

(1)

∈R

are separately computed for 1 ≤i ≤P, and recombined

additively: C

i=1



ρ(T

(1)

– ρ(T

(1)



Once the attention coefﬁcients are computed, they can be normalised across

each receiver node via the softmax function:

= softmax

) (8.151)

and then used to appropriately recombine the value scalars, value vectors, and

pair embeddings:

(0)

v∈V

(0)

(1)

= ρ(T

)

–1

v∈V

ρ(T

(1)

pair

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

)

–1

x = R

–1

(x – t

), converting from the global coordinates back

into T

’s frame. Finally, a simple afﬁne combination of all of these features

(with residual) yields the updated node-wise representation of amino acids:

′

= s

+ W

out



(0)

⊕h

(1)

⊕h

pair

⊕



(1)





+ b

out

(8.153)

where ⊕ denotes concatenation, W

out

∈R

k×(l+3+m+1)

is a learnable weight

matrix, and b

out

∈R

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

′

∈R

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, T

glob

. First, we

show this for the vector attention component:

∥ρ(T

glob

)ρ(T

(1)

– ρ(T

glob

)ρ(T

(1)

∥



ρ(T

glob

)



ρ(T

(1)

– ρ(T

(1)





= ∥ρ(T

(1)

– ρ(T

(1)

∥

where we can cancel out the application of ρ(T

glob

) because (a) rotations do not

change norm, and (b) global translations get cancelled out by the subtraction.

Next, we prove that T

glob

does not affect the updated type-1 vectors, h

(1)

ρ(T

glob

(1)

= ρ(T

)

–1

ρ(T

glob

)

–1

v∈V

ρ(T

glob

)ρ(T

(1)

= ρ(T

)

–1

ρ(T

glob

)

–1

ρ(T

glob

)

v∈V

ρ(T

(1)

= ρ(T

)

–1

v∈V

ρ(T

(1)

= h

(1)

Given these two invariances, we can conclude that s

′

will be unaffected by any

global transformations of the local T

frames. Armed with such an update, we

can now equivariantly update the frames in response to s

′

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 ∆T

= (

), indicating the

additional roto-translation derived from s

′

. AlphaFold 2 does not predict

rotation angles directly from this vector—instead, it computes a quaternion,

248 Chapter 8

(1, i

, j

, k

1 + i

+ j

+ k

, and the translation vector,

, directly from u

(e.g. via a standard linear layer). We can convert a quaternion (a

, b

, c

, d

)

to a 3D rotation matrix like so:





+ b

– c

– d

– 2a

+ 2a

– b

+ c

– d

– 2a

+ 2a

– b

– c

+ d





(8.154)

Note that by keeping the ﬁrst (real) quaternion component at 1, we ensure a

valid rotation is applied. Once ∆T

is known, we simply update T

as:

′

= (R

, t

+ R

) (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 T

′

transforms

equivariantly. If a global transformation T

glob

is applied to the system, the

base frame representation becomes ρ(T

glob

)ρ(T

), while the predicted update

representation, ρ(∆T

), remains unchanged. The resulting updated frame is:

(ρ(T

glob

)ρ(T

))ρ(∆T

) = ρ(T

glob

)(ρ(T

)ρ(∆T

)) = ρ(T

glob

)ρ(T

′

) (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 ﬂow 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 modiﬁcation truly was,

and whether it continues to be relevant past the Nobel prize-winning moment.

Ablations One of the key scientiﬁc 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 deﬁnitely 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 signiﬁcant-

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 signiﬁcant, 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, insufﬁcient time has passed to make a clear conclusion:

•

Given sufﬁciently broad training data, non-equivariant models can

learn a sufﬁciently-approximated equivariance to be useful. We believe

that this interpretation is valid: AlphaFold 3 is a clear existence proof. How-

ever, there might be signiﬁcant nuance about just how much training data is

needed

, and also what levels of symmetry deviations are acceptable

in the

model’s outputs.

•

AlphaFold 3 still leverages geometry—just in a less rigid way. Speciﬁ-

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) ﬁrst learned to play Go from human games only, and only then

reﬁned itself through self-play, after which it was able to completely yield the

human data altogether (Silver et al. 2017). Perhaps the ﬁrst 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 signiﬁcant 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, ﬁnal Chapter, where we will

study what lies beyond geometric deep learning, both in terms of more general

structures, and moving beyond pure architectural constraints.