A quick intro to Geometric Deep Learning(WIP)

Coauthored with Zikai Wang.

Most of the time we consider deep learning as an engineering problem – we are told what works, what doesn’t, and what the standard solution is to a specific problem. Yet it is often less clear why these architectures work. Geometric Deep Learning provides us with mathematical tools for understanding why architectures such as CNNs, LSTMs, and Transformers can exploit structure that a plain MLP does not explicitly encode.

To quickly give you a sense of how it works and why it is useful, we will skip most of the math in the front part of Original Geometric Deep Learning book and give you a walk-through derivation of convolution from linearity and shift equivariance.

Let’s start with a 1D image with 4 pixels.

You can imagine each value of $𝑥$ as the brightness of one pixel.

Since we want to discuss shift equivariance, we first need to define a shift operation. To make things simple, we assume images with no boundary, which means positions wrap around in a circle. We call this property periodic boundary condition. For example, if we shift the image by 1 pixel to the right, the new image shall be:

Let’s call a shift by $𝑎$ pixels ${𝑆}_{𝑎}$. In particular,

For an $𝑛$-pixel signal, the $𝑢$-th component of the shifted image is

Now suppose we have an image-processing layer $𝐹$ that takes an image as input and outputs a feature map with the same spatial size:

The layer is shift-equivariant if applying a shift to the input is equivalent to applying the same shift to the output:

In other words, if the input moves, the feature map moves with it by the same amount.

We call this kind of structural constraint Inductive Bias: an assumption built into a learning algorithm about which functions are likely to be useful for a given task.

For example, shift equivariance means that if a cat moves a few pixels in an image, the intermediate feature map should move by the same amount. Shift invariance, on the other hand, means that the final output, such as a classification score, remains unchanged under the same shift.

These properties can be built into the network architecture or encouraged through regularization. Architectural constraints can enforce exact equivariance under appropriate assumptions, whereas regularization alone does not generally guarantee it.

We will see exactly how that works in the next chapter.

On a discrete grid with periodic boundary conditions, the unit shift ${𝑆}_1$ corresponds to a simple permutation matrix:

Suppose our image-processing layer $𝐹\left(𝑥\right)=𝑊𝑥$, where $𝑊\in {ℝ}^{𝑛\times 𝑛}$ is the weight matrix to be determined. ($𝑊$ is simply the matrix representation of the same linear operator $𝐹$)

Using elementary linear algebra, let us examine the restrictions that this commutativity constraint $𝑊{𝑆}_1={𝑆}_1𝑊$ imposes on the entries ${𝑊}_{𝑢,𝑣}$ of matrix $𝑊$. Right-multiplication by the shift operator ${\left(𝑊{𝑆}_1\right)}_{𝑢,𝑣}$ shifts column $𝑣+1$ of $𝑊$ to column $𝑣$, which yields

.

Whereas the left-multiplication by the shift operator, ${\left({𝑆}_1𝑊\right)}_{𝑢,𝑣}$ shifts row $𝑢-1$ of $𝑊$ to row $𝑢$, which yields

From $𝑊{𝑆}_1={𝑆}_1𝑊$, we obtain the fundamental constraint equation:

letting ${𝑢}^{\prime }=𝑢-1$,

This relations reveals that a shift equivariant operation is represented by a matrix where elements along each diagonal are constant, and any diagonal that runs off one edge wraps around and continues with the same value on the opposite edge. This is also known as a circulant matrix. If we define an $𝑛$-dimensional parameter vector ${𝜃}_{𝑘}≔{𝑊}_{𝑘,0}$, where $𝑘=\left(𝑢-𝑣\right)mod𝑛$, then

This demonstrates that the entry ${𝑊}_{𝑢,𝑣}$ depends solely on the index difference $\left(𝑢-𝑣\right)$. Substituting this back into $𝑊$, its structure is necessarily

Now consider the $𝑢$-th component of the output feature map $𝑦=𝑊𝑥=𝐶\left(𝜃\right)𝑥$:

This is precisely the definition of discrete convolution. In other words, the mathematical framework naturally deduces the convolution formula from linearity and shift equivariance

In the example above, representing an image as a vector works perfectly well. However, this is not the most general way to represent a signal. For more complex domains, such as spheres or graphs, a vector representation can obscure the underlying structure of the domain. In GDL, a more general and natural approach is to represent a signal as a function, sometimes also called a mapping.

The standard form is

where $\mathrm{\Omega }$ is the domain. Given a position $𝑢\in \mathrm{\Omega }$, the signal returns its value $𝑥\left(𝑢\right)$. For a 2D image, $\mathrm{\Omega }$ is a grid. For a signal on a sphere, $\mathrm{\Omega }$ could be the sphere itself and $𝑥\left(𝑢\right)$ might be a temperature measured at point $𝑢$. For a graph, $\mathrm{\Omega }$ can be the set of nodes and $𝑥\left(𝑢\right)$ a feature attached to node $𝑢$.

For example, suppose we have a 2-by-2 image

Our previous vector representation flattens it as

In the function representation, we instead define $\mathrm{\Omega }=\left\{\left(0,0\right),\left(0,1\right),\left(1,0\right),\left(1,1\right)\right\}$ and the function as

The values are the same; the advantage is that the geometry lives explicitly in the domain $\mathrm{\Omega }$ instead of being hidden by an arbitrary flattening order. However, to keep things simple, we are still going to use a 1D example.

Now, we are going to represent a 1D image as the linear combination of impulse functions. But what’s even an impulse function? We all know that every vector can be reconstructed from its coordinates and the standard basis vectors:

When we are presenting an image as a function, we also want an analogous set of “basis functions”, which picks out each position of the domain. This is exactly what the impulse function does:

By itself, it means one bright pixel at the origin…

... 0 0 1 0 0...

        ^

        0

... 0 0 1 0 0...

        ^

        0

... 0 0 1 0 0...

        ^

        0

... 0 0 1 0 0...

        ^

        0

…, which is not very useful. But we can combine it with the shift operation we defined earlier to represent any image. To make things easier to read, we define ${𝛿}_{𝑣}={𝑆}^{𝑣}\left({𝛿}_0\right)$. Then any signal $𝑥$ can be assembled out of these one-pixel signals:

Pointwise, this says

For simplicity’s sake $𝑥$ is also written as $𝑥={\sum }_{𝑣\in \mathbb{Z}}𝑎\left(𝑣\right){𝛿}_{𝑣}$. Keep in mind that ${𝛿}_{𝑣}$ is a function.

For example, an image like this:

position:      -1   0   1

x:              2   5  -1

position:      -1   0   1

x:              2   5  -1

position:      -1   0   1

x:              2   5  -1

position:      -1   0   1

x:              2   5  -1

can be represented as $𝑥=2{𝛿}_{-1}+5{𝛿}_0-1{𝛿}_1$.

Let’s also start from a linear layer $𝐹$, which takes a signal $𝑥:\mathrm{\Omega }\to ℝ$ and outputs a signal $𝑦:\mathrm{\Omega }\to ℝ$.

(So when you see $\left[𝐹\left(𝑥\right)\right]\left(𝑢\right)$ below, it means the u-th component of the output signal $𝑦$.) This is the same kind of $𝐹$ as in the matrix section; once we choose the impulse basis, its coordinate matrix is the $𝑊$ used above.

Plug in the representation of $𝑥$ above to our layer $𝐹$, we have:

By linearity, we can move it inside the summation:

Because $𝐹$ is shift equivariant, we can replace $𝐹\left({𝑆}^{𝑣}{𝛿}_0\right)$ with ${𝑆}^{𝑣}𝐹\left({𝛿}_0\right)$.

The key insight is that shift equivariance allows us to determine the response to an impulse at any position from the response to a single impulse at the origin. $𝐹\left({𝛿}_0\right)$ is a constant “vector” to all the input signals. When we process different positions, we simply shift this “constant vector”. We call it a kernel in convolution term. Let’s denote $𝜃=𝐹\left({𝛿}_0\right)$ ($𝜃:\mathbb{Z}\to ℝ$), then

According to the definition of the shift operation, $\left({𝑆}^{𝑣}𝑔\right)\left(𝑢\right)=𝑔\left(𝑢-𝑣\right)$. We now get:

which is exactly the definition of discrete convolution.