Sapolsky

“Determined” is wonderfully written, an easier read than “Behave”. If you’ve read “Behave”, you won’t find too much new content here. But the focus is different; Sapolsky dedicates the major portion of this book striving to answer the, undoubtedly, hardest question: If there is no free will, how should we organize our living together? If no-one is responsible for their actions, no sensible meaning assignable to “self control” and “willpower”, if there is no “crime” and no “punishment”, how can we make our societies safe to live in? (As an aside, this is not the only societal question ensuing. Sapolsky doesn’t discuss this at length – it being more urgent to address the “crime question” – but: If, just like responsibility, merit dissolves into nothing, how then do we build a just, a fair, an ethical society?)

Why is there no free will? Because at all times, what you think, what you feel, what you consider doing depends on how your life looked before. Here Sapolsky gives a condensation of what 800-page “Behave” is dedicated to: hormonal, neurological, environmental, genetic, cultural, evolutionary factors that have shaped, every moment of your life and even when still in the womb, how your brain, your body, your ways of feeling and thinking look now.

Following that summarization, he addresses popular arguments, as brought up by various flavors of free-will proponents. Some are philosophical, meaning they’re profferred by professional philosophers. For example, Sapolsky cites Daniel Dennett arguing that, sure, there’s luck, or bad luck, but in the long run, it will all average out : “[…] a good runner who starts at the back of the pack, if he is really good enough to DESERVE winning, will probably have had plenty of opportunity to overcome the initial disadvantage”. (Capitalization by Sapolsky, who I imagine to have felt some complex emotion of revulsion and “cruel satisfaction” in citing this passage.) – But does luck average out? Not in the marathon, where the person starting from the back is, in some endurance sports, likely to have to wait in pile-ups, causing them to lose yet more time, and to arrive at refreshment stations already emptied from all the good stuff. And not in life either – in that regard, Sapolsky’s book is full to the brim of shocking, disillusioning, depressing examples.

Supplementing this “averaging out” postulate is the view, popular among philosophers as well as “normal people”, that sure, you might sense whatever strong temptation, but if so, you just have to resist – to exercise self-control, a.k.a. discipline, a.k.a. “will power”. If you’re driven to binge-eat, throw the food out; if you can’t motivate yourself to study, give yourself a kick in the ass; if you’re tempted to give up, then just don’t give up. This certainly makes for narratives more than compliant with Western, meritocracy-shaped societies. But where does that two-step process – “me no. 1”, the subject that’s driven, coerced to do (or let be) something, and “me no. 2”, the boss who orders and forbids – come from? If your past determines what you feel and think now, then certainly it, too, decides over access to willpower now.

Unsurprisingly, Sapolsky does not seem to be much impressed by arguments like this. What about a second line of arguments, based on results from the natural sciences (first and foremost, from physics, the science that keeps us all in awe and subjection)? There are three fields frequently advertised as exhibiting indeterminism, and thus – it is said – to enable free will, and to each of those Sapolsky dedicates two chapters: one giving a readable, lucid introduction, and one, discussing the claim that from the quantum (right…) of indeterminacy they introduce, free will does emerge.

I won’t lie and admit upfront that in all three cases, either I find his respective arguments very convincing, or I’d want to put things even more sharply, or both.

First, there is the fascinating subject of chaos theory. But magical though its results may look to us humans, all it implies is unpredictability – not randomness, however, not indeterminism. Chaos theory is all about how, under certain circumstances, deterministic processes can produce completely different results. All that’s needed are incredibly small differences in initial conditions, i.e., the starting state. In practice this means – seeing how we never can determine a starting state to 100% precision - that no reliable forecast is possible. But that’s very different from randomness.

Second, there’s complexity science, with its impressive demonstrations that out of very simple building blocks, complex, organized behavior can emerge. But here again, we struggle to demonstrate a sensible connection to free will. Not only are the building blocks – neurons – anything but simple. Also, the emergent behavior – ant foraging, say, or network layout – has nothing to do with free will. Concretely, the individua – single ant, network node - aren’t acquiring free will (nor do they gain in agency, to use the term figuring prominently in Mitchell’s text).

Finally, and most importantly, there’s quantum physics. Most importantly, because this is the field whose theories and findings have inspired people, physicists and non-physicists alike, to the wildest fantasies. (I’ll talk a little bit about this when we get to Mitchell, in a minute.) The fundamental problem is very down-to-earth, though. Whatever tiny structure is claimed to possess quantum-like indeterminacy, we need to show how from that layer, bigger structures arise that, in the end, produce what we call free will. Sapolsky calls it the bubbling-up, and of this, it looks like there does not exist any even half-way plausible account.

Mitchell

A convincing story of that bubbling up, that’s exactly what Mitchell is trying to present. He presents a detailed account of how, with the very first appearance of life, something like agency came into being, and how from there, it continued to evolve. It’s a fascinating read, especially as regards the very early stages. To me, at least, it is; I’ve never thought about very simple organisms – single-celled organisms! – in that way.

Mitchell explains how life itself is intimately linked to agency: No life is possible without continuous work being done to keep an “inner” separate from an “outer”.

In consequence, to not go extinct, organisms have to evolve in the direction of – call it however you want: agency (in the sense of actively shaping one’s fate), intelligent decision-making, creative adaptation… I find this super convincing - it’s just that, employing words like “intelligent” or “creative”, I’m using the language I’m accustomed to think in. Just like I use them for describing my own behavior, I can use them referring to any living organism’s actions. (Yes, I think that there’s no more “truth” in using these terms for humans than for other creatures. To me, it’s totally fine to use them, though, as long as I’m conscious of my just utilizing some form of agreed-upon linguistic code here.)

End of “irresistible aside”, back to Mitchell. I like the agency story, but I fail to see how from – gradual, as he himself demonstrates so nicely! – increase in complexity might arise what we call free will. In fact, Mitchell dedicates several chapters to elaborating this link. The main role is played by quantum physics. Mitchell cites prominent contemporary physicists who have built whole high-level (and quite audacious, I’d say, sometimes) world views on quantum-scale indeterminacy. For example (citing Mitchell):

“Gisin and del Santo argue that, beyond some decimal point, the numbers describing the physical parameters of classical systems in the future simply become indefinite. It’s not just that we don’t know what they are—the universe also doesn’t know what they are, not yet. The numbers are not given with infinite precision all at once; instead, they evolve as processes through time. Again, it is through interaction that these parameters take on a definite value, paralleling the “collapse” of quantum possibilities. It’s not that these parameters can take any value; they’re not completely random. They are constrained, often tightly, by the physics of the system in question —they are, however, not completely constrained.”

This, to me, reads, in a way, as “indefinite” as the future parameters are posited to be. What does Mitchell make of it?

“Indeed, the behavior of everything in the universe would have been predetermined from the dawn of time, including me writing this sentence and you reading it. In my view, this claim is absurd on its face. That is not an argument for whether it’s true or not but just an observation that this claim jars so bracingly with our actual experience of the world that we should be strongly suspicious of it. […] There is only one future open, with no possibilities, no deciding or acting, no purposiveness, no mattering, no trying, no goals, or functions—it would be literally a meaningless universe.”

I’m sorry, but all I can think reading this is “wishful thinking”. We don’t want to live in a meaningless universe. But meaning is something we construct, or fail to construct, in our lives, not something outside.

More on the “hard facts side”, even the big bang jumps in to help:

“Indeed, the observed organization of matter in the universe depended on the presence of quantum fluctuations in a phase of “cosmic inflation” that preceded the “hot” Big Bang. As the universe was rapidly expanding, it was initially a homogeneous field of energy. Without random quantum fluctuations that broke up this symmetrical field, matter and energy would have been so evenly distributed that pretty much nothing would have happened. Instead, these tiny blips introduced enough inhomogeneity into gravitational fields that galaxies and stars and planets could form.“

I can say absolutely nothing about the facts, but again, I have problems following Mitchell’s conclusions:

“No, the really crucial point is that the introduction of chance undercuts necessity’s monopoly on causation. The low-level physical details and forces are not causally comprehensive: they are not sufficient to determine how a system will evolve from state to state. This opens the door for higher-level features to have some causal influence in determining which way the physical system will evolve. This influence is exerted by establishing contextual constraints: in other words, the way the system is organized can also do some causal work. In the brain, that organization embodies knowledge, beliefs, goals, and motivations—our reasons for doing things. This means some things are driven neither by necessity nor by chance; instead, they are up to us.”

“This opens the door…” – does it? How exactly do higher-level features come to exert causal influence?

Wrapping up,- - when Mitchell cites Bergson, his concept of durations, I am, emotionally, all on his side: I like that concept a lot, I find it totally emotionally attractive. But still, I think it’s a description of how our reality feels to us. If you don’t want to make statements about an external reality (and I’m lucky there; I don’t want to) that’s fine, but otherwise you somehow have to, in Sapolsky’s words, make that “bubbling up” plausible.

Speaking of Sapolsky – I think that as a proponent of free will, you also have to explain of how bubbling up, if it’s successful, circumvents his account of how the past – your past, immediate, recent and less recent; your environment, starting from when in the womb; and finally, evolution, both cultural and biological – make up what you are now.

Spaces, groups, and representations: escnn$gspaces

We start by peeking into gspaces, one of the two sub-modules we are going to make direct use of.

gspaces <- escnn$gspaces
py_list_attributes(gspaces) |> (\(vec) grep("On", vec, value = TRUE))() |> sort()

[1] "conicalOnR3" "cylindricalOnR3" "dihedralOnR3" "flip2dOnR2" "flipRot2dOnR2" "flipRot3dOnR3"
[7] "fullCylindricalOnR3" "fullIcoOnR3" "fullOctaOnR3" "icoOnR3" "invOnR3" "mirOnR3 "octaOnR3"
[14] "rot2dOnR2" "rot2dOnR3" "rot3dOnR3" "trivialOnR2" "trivialOnR3"    

The methods I’ve listed instantiate a gspace. If you look closely, you see that they’re all composed of two strings, joined by “On”. In all instances, the second part is either R2 or R3. These two are the available base spaces – and – an input signal can live in. Signals can, thus, be images, made up of pixels, or three-dimensional volumes, composed of voxels. The first part refers to the group you’d like to use. Choosing a group means choosing the symmetries to be respected. For example, rot2dOnR2() implies equivariance as to rotations, flip2dOnR2() guarantees the same for mirroring actions, and flipRot2dOnR2() subsumes both.

Let’s define such a gspace. Here we ask for rotation equivariance on the Euclidean plane, making use of the same cyclic group – – we developed in our from-scratch implementation:

r2_act <- gspaces$rot2dOnR2(N = 4L)
r2_act$fibergroup

In this post, I’ll stay with that setup, but we could as well pick another rotation angle – N = 8, say, resulting in eight equivariant positions separated by forty-five degrees. Alternatively, we might want any rotated position to be accounted for. The group to request then would be SO(2), called the special orthogonal group, of continuous, distance- and orientation-preserving transformations on the Euclidean plane:

(gspaces$rot2dOnR2(N = -1L))$fibergroup

SO(2)

Going back to , let’s investigate its representations:

r2_act$representations

$irrep_0
C4|[irrep_0]:1

$irrep_1
C4|[irrep_1]:2

$irrep_2
C4|[irrep_2]:1

$regular
C4|[regular]:4

A representation, in our current context and very roughly speaking, is a way to encode a group action as a matrix, meeting certain conditions. In escnn, representations are central, and we’ll see how in the next section.

First, let’s inspect the above output. Four representations are available, three of which share an important property: they’re all irreducible. On , any non-irreducible representation can be decomposed into into irreducible ones. These irreducible representations are what escnn works with internally. Of those three, the most interesting one is the second. To see its action, we need to choose a group element. How about counterclockwise rotation by ninety degrees:

elem_1 <- r2_act$fibergroup$element(1L)
elem_1

1[2pi/4]

Associated to this group element is the following matrix:

r2_act$representations[[2]](elem_1)

             [,1]          [,2]
[1,] 6.123234e-17 -1.000000e+00
[2,] 1.000000e+00  6.123234e-17

This is the so-called standard representation,

, evaluated at . (It is called the standard representation because it directly comes from how the group is defined (namely, a rotation by in the plane).

The other interesting representation to point out is the fourth: the only one that’s not irreducible.

r2_act$representations[[4]](elem_1)

[1,]  5.551115e-17 -5.551115e-17 -8.326673e-17  1.000000e+00
[2,]  1.000000e+00  5.551115e-17 -5.551115e-17 -8.326673e-17
[3,]  5.551115e-17  1.000000e+00  5.551115e-17 -5.551115e-17
[4,] -5.551115e-17  5.551115e-17  1.000000e+00  5.551115e-17

This is the so-called regular representation. The regular representation acts via permutation of group elements, or, to be more precise, of the basis vectors that make up the matrix. Obviously, this is only possible for finite groups like , since otherwise there’d be an infinite amount of basis vectors to permute.

To better see the action encoded in the above matrix, we clean up a bit:

round(r2_act$representations[[4]](elem_1))

    [,1] [,2] [,3] [,4]
[1,]    0    0    0    1
[2,]    1    0    0    0
[3,]    0    1    0    0
[4,]    0    0    1    0

This is a step-one shift to the right of the identity matrix. The identity matrix, mapped to element 0, is the non-action; this matrix instead maps the zeroth action to the first, the first to the second, the second to the third, and the third to the first.

We’ll see the regular representation used in a neural network soon. Internally – but that need not concern the user – escnn works with its decomposition into irreducible matrices. Here, that’s just the bunch of irreducible representations we saw above, numbered from one to three.

Having looked at how groups and representations figure in escnn, it is time we approach the task of building a network.

Representations, for real: escnn$nn$FieldType

So far, we’ve characterized the input space (), and specified the group action. But once we enter the network, we’re not in the plane anymore, but in a space that has been extended by the group action. Rephrasing, the group action produces feature vector fields that assign a feature vector to each spatial position in the image.

Now we have those feature vectors, we need to specify how they transform under the group action. This is encoded in an escnn$nn$FieldType . Informally, we could say that a field type is the data type of a feature space. In defining it, we indicate two things: the base space, a gspace, and the representation type(s) to be used.

In an equivariant neural network, field types play a role similar to that of channels in a convnet. Each layer has an input and an output field type. Assuming we’re working with grey-scale images, we can specify the input type for the first layer like this:

nn <- escnn$nn
feat_type_in <- nn$FieldType(r2_act, list(r2_act$trivial_repr))

The trivial representation is used to indicate that, while the image as a whole will be rotated, the pixel values themselves should be left alone. If this were an RGB image, instead of r2_act$trivial_repr we’d pass a list of three such objects.

So we’ve characterized the input. At any later stage, though, the situation will have changed. We will have performed convolution once for every group element. Moving on to the next layer, these feature fields will have to transform equivariantly, as well. This can be achieved by requesting the regular representation for an output field type:

feat_type_out <- nn$FieldType(r2_act, list(r2_act$regular_repr))

Then, a convolutional layer may be defined like so:

conv <- nn$R2Conv(feat_type_in, feat_type_out, kernel_size = 3L)

Group-equivariant convolution

What does such a convolution do to its input? Just like, in a usual convnet, capacity can be increased by having more channels, an equivariant convolution can pass on several feature vector fields, possibly of different type (assuming that makes sense). In the code snippet below, we request a list of three, all behaving according to the regular representation.

feat_type_in <- nn$FieldType(r2_act, list(r2_act$trivial_repr))
feat_type_out <- nn$FieldType(
  r2_act,
  list(r2_act$regular_repr, r2_act$regular_repr, r2_act$regular_repr)
)

conv <- nn$R2Conv(feat_type_in, feat_type_out, kernel_size = 3L)

We then perform convolution on a batch of images, made aware of their “data type” by wrapping them in feat_type_in:

x <- torch$rand(2L, 1L, 32L, 32L)
x <- feat_type_in(x)
y <- conv(x)
y$shape |> unlist()

[1]  2  12 30 30

The output has twelve “channels”, this being the product of group cardinality – four distinguished positions – and number of feature vector fields (three).

If we choose the simplest possible, approximately, test case, we can verify that such a convolution is equivariant by direct inspection. Here’s my setup:

feat_type_in <- nn$FieldType(r2_act, list(r2_act$trivial_repr))
feat_type_out <- nn$FieldType(r2_act, list(r2_act$regular_repr))
conv <- nn$R2Conv(feat_type_in, feat_type_out, kernel_size = 3L)

torch$nn$init$constant_(conv$weights, 1.)
x <- torch$vander(torch$arange(0,4))$view(tuple(1L, 1L, 4L, 4L)) |> feat_type_in()
x

g_tensor([[[[ 0.,  0.,  0.,  1.],
            [ 1.,  1.,  1.,  1.],
            [ 8.,  4.,  2.,  1.],
            [27.,  9.,  3.,  1.]]]], [C4_on_R2[(None, 4)]: {irrep_0 (x1)}(1)])

Inspection could be performed using any group element. I’ll pick rotation by :

all <- iterate(r2_act$testing_elements)
g1 <- all[[2]]
g1

Just for fun, let’s see how we can – literally – come whole circle by letting this element act on the input tensor four times:

all <- iterate(r2_act$testing_elements)
g1 <- all[[2]]

x1 <- x$transform(g1)
x1$tensor
x2 <- x1$transform(g1)
x2$tensor
x3 <- x2$transform(g1)
x3$tensor
x4 <- x3$transform(g1)
x4$tensor

tensor([[[[ 1.,  1.,  1.,  1.],
          [ 0.,  1.,  2.,  3.],
          [ 0.,  1.,  4.,  9.],
          [ 0.,  1.,  8., 27.]]]])
          
tensor([[[[ 1.,  3.,  9., 27.],
          [ 1.,  2.,  4.,  8.],
          [ 1.,  1.,  1.,  1.],
          [ 1.,  0.,  0.,  0.]]]])
          
tensor([[[[27.,  8.,  1.,  0.],
          [ 9.,  4.,  1.,  0.],
          [ 3.,  2.,  1.,  0.],
          [ 1.,  1.,  1.,  1.]]]])
          
tensor([[[[ 0.,  0.,  0.,  1.],
          [ 1.,  1.,  1.,  1.],
          [ 8.,  4.,  2.,  1.],
          [27.,  9.,  3.,  1.]]]])

You see that at the end, we are back at the original “image”.

Now, for equivariance. We could first apply a rotation, then convolve.

Rotate:

x_rot <- x$transform(g1)
x_rot$tensor

This is the first in the above list of four tensors.

Convolve:

y <- conv(x_rot)
y$tensor

tensor([[[[ 1.1955,  1.7110],
          [-0.5166,  1.0665]],

         [[-0.0905,  2.6568],
          [-0.3743,  2.8144]],

         [[ 5.0640, 11.7395],
          [ 8.6488, 31.7169]],

         [[ 2.3499,  1.7937],
          [ 4.5065,  5.9689]]]], grad_fn=<ConvolutionBackward0>)

Alternatively, we can do the convolution first, then rotate its output.

Convolve:

y_conv <- conv(x)
y_conv$tensor

tensor([[[[-0.3743, -0.0905],
          [ 2.8144,  2.6568]],

         [[ 8.6488,  5.0640],
          [31.7169, 11.7395]],

         [[ 4.5065,  2.3499],
          [ 5.9689,  1.7937]],

         [[-0.5166,  1.1955],
          [ 1.0665,  1.7110]]]], grad_fn=<ConvolutionBackward0>)

Rotate:

y <- y_conv$transform(g1)
y$tensor

tensor([[[[ 1.1955,  1.7110],
          [-0.5166,  1.0665]],

         [[-0.0905,  2.6568],
          [-0.3743,  2.8144]],

         [[ 5.0640, 11.7395],
          [ 8.6488, 31.7169]],

         [[ 2.3499,  1.7937],
          [ 4.5065,  5.9689]]]])

Indeed, final results are the same.

At this point, we know how to employ group-equivariant convolutions. The final step is to compose the network.

A group-equivariant neural network

Basically, we have two questions to answer. The first concerns the non-linearities; the second is how to get from extended space to the data type of the target.

First, about the non-linearities. This is a potentially intricate topic, but as long as we stay with point-wise operations (such as that performed by ReLU) equivariance is given intrinsically.

In consequence, we can already assemble a model:

feat_type_in <- nn$FieldType(r2_act, list(r2_act$trivial_repr))
feat_type_hid <- nn$FieldType(
  r2_act,
  list(r2_act$regular_repr, r2_act$regular_repr, r2_act$regular_repr, r2_act$regular_repr)
  )
feat_type_out <- nn$FieldType(r2_act, list(r2_act$regular_repr))

model <- nn$SequentialModule(
  nn$R2Conv(feat_type_in, feat_type_hid, kernel_size = 3L),
  nn$InnerBatchNorm(feat_type_hid),
  nn$ReLU(feat_type_hid),
  nn$R2Conv(feat_type_hid, feat_type_hid, kernel_size = 3L),
  nn$InnerBatchNorm(feat_type_hid),
  nn$ReLU(feat_type_hid),
  nn$R2Conv(feat_type_hid, feat_type_out, kernel_size = 3L)
)$eval()

model

SequentialModule(
  (0): R2Conv([C4_on_R2[(None, 4)]:
       {irrep_0 (x1)}(1)], [C4_on_R2[(None, 4)]: {regular (x4)}(16)], kernel_size=3, stride=1)
  (1): InnerBatchNorm([C4_on_R2[(None, 4)]:
       {regular (x4)}(16)], eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU(inplace=False, type=[C4_on_R2[(None, 4)]: {regular (x4)}(16)])
  (3): R2Conv([C4_on_R2[(None, 4)]:
       {regular (x4)}(16)], [C4_on_R2[(None, 4)]: {regular (x4)}(16)], kernel_size=3, stride=1)
  (4): InnerBatchNorm([C4_on_R2[(None, 4)]:
       {regular (x4)}(16)], eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (5): ReLU(inplace=False, type=[C4_on_R2[(None, 4)]: {regular (x4)}(16)])
  (6): R2Conv([C4_on_R2[(None, 4)]:
       {regular (x4)}(16)], [C4_on_R2[(None, 4)]: {regular (x1)}(4)], kernel_size=3, stride=1)
)

Calling this model on some input image, we get:

x <- torch$randn(1L, 1L, 17L, 17L)
x <- feat_type_in(x)
model(x)$shape |> unlist()

[1]  1  4 11 11

What we do now depends on the task. Since we didn’t preserve the original resolution anyway – as would have been required for, say, segmentation – we probably want one feature vector per image. That we can achieve by spatial pooling:

avgpool <- nn$PointwiseAvgPool(feat_type_out, 11L)
y <- avgpool(model(x))
y$shape |> unlist()

[1] 1 4 1 1

We still have four “channels”, corresponding to four group elements. This feature vector is (approximately) translation-invariant, but rotation-equivariant, in the sense expressed by the choice of group. Often, the final output will be expected to be group-invariant as well as translation-invariant (as in image classification). If that’s the case, we pool over group elements, as well:

invariant_map <- nn$GroupPooling(feat_type_out)
y <- invariant_map(avgpool(model(x)))
y$tensor

tensor([[[[-0.0293]]]], grad_fn=<CopySlices>)

We end up with an architecture that, from the outside, will look like a standard convnet, while on the inside, all convolutions have been performed in a rotation-equivariant way. Training and evaluation then are no different from the usual procedure.

What’s in the book?

Deep Learning and Scientific Computing with R torch has three parts.

The first covers the indispensible basics: tensors, and how to manipulate them; automatic differentiation, the sine qua non of deep learning; optimization, the strategy that drives most of what we call artificial intelligence; and neural-network modules, torch's way of encapsulating algorithmic flow. The focus is on understanding the concepts, on how things “work” – that’s why we do things like code a neural network from scratch, something you’ll probably never do in later use.

Foundations laid, part two – considerably more sizeable – dives into deep-learning applications. It is here that the ecosystem surrounding core torch enters the spotlight. First, we see how luz automates and considerably simplifies many programming tasks related to network training, performance evaluation, and prediction. Making use of the wrappers and instrumentation facilities it provides, we next learn about two aspects of deep learning no real-world application can afford to neglect: How to make models generalize to unseen data, and how to accelerate training. Techniques we introduce keep re-appearing throughout the use cases we then look at: image classification and segmentation, regression on tabular data, time-series forecasting, and classifying speech utterances. It’s in working with images and sound that essential ecosystem libraries, namely, torchvision and torchaudio, make their appearance, to be used for domain-dependent functionality.

In part three, we move beyond deep learning, and explore how torch can figure in general mathematical or scientific applications. Prominent topics are regression using matrix decompositions, the Discrete Fourier Transform, and the Wavelet Transform. The primary goal here is to understand the underlying ideas, and why they are so important. That’s why, here just like in part one, we code algorithms from scratch, before introducing the speed-optimized torch equivalents.

Now that you know about the book’s content, you may be asking:

What do I get?

Content-wise, I think I can consider this question answered. If there were other books on torch with R, I’d probably stress two things: First, the already-referred-to focus on concepts and understanding. Second, the usefulness of the code examples. By using off-the-shelf datasets, and performing the usual types of tasks, we write code fit to serve as a start in your own applications – providing templates ready to copy-paste and adapt to a purpose.

Thanks for reading, and I hope you enjoy the book!

Step 1: The symmetry group

In coding a GCNN, the first thing we need to provide is an implementation of the symmetry group we’d like to use. Here, it is , the four-element group that rotates by 90 degrees.

We can ask gcnn to create one for us, and inspect its elements.

# remotes::install_github("skeydan/gcnn")
library(gcnn)
library(torch)

C_4 <- CyclicGroup(order = 4)
elems <- C_4$elements()
elems

torch_tensor
 0.0000
 1.5708
 3.1416
 4.7124
[ CPUFloatType{4} ]

Elements are represented by their respective rotation angles: , , , and .

Groups are aware of the identity, and know how to construct an element’s inverse:

C_4$identity

g1 <- elems[2]
C_4$inverse(g1)

torch_tensor
 0
[ CPUFloatType{1} ]

torch_tensor
4.71239
[ CPUFloatType{} ]

Here, what we care about most is the group elements’ action. Implementation-wise, we need to distinguish between them acting on each other, and their action on the vector space , where our input images live. The former part is the easy one: It may simply be implemented by adding angles. In fact, this is what gcnn does when we ask it to let g1 act on g2:

g2 <- elems[3]

# in C_4$left_action_on_H(), H stands for the symmetry group
C_4$left_action_on_H(torch_tensor(g1)$unsqueeze(1), torch_tensor(g2)$unsqueeze(1))

torch_tensor
 4.7124
[ CPUFloatType{1,1} ]

What’s with the unsqueeze()s? Since ’s ultimate raison d’être is to be part of a neural network, left_action_on_H() works with batches of elements, not scalar tensors.

Things are a bit less straightforward where the group action on is concerned. Here, we need the concept of a group representation. This is an involved topic, which we won’t go into here. In our current context, it works about like this: We have an input signal, a tensor we’d like to operate on in some way. (That “some way” will be convolution, as we’ll see soon.) To render that operation group-equivariant, we first have the representation apply the inverse group action to the input. That accomplished, we go on with the operation as though nothing had happened.

To give a concrete example, let’s say the operation is a measurement. Imagine a runner, standing at the foot of some mountain trail, ready to run up the climb. We’d like to record their height. One option we have is to take the measurement, then let them run up. Our measurement will be as valid up the mountain as it was down here. Alternatively, we might be polite and not make them wait. Once they’re up there, we ask them to come down, and when they’re back, we measure their height. The result is the same: Body height is equivariant (more than that: invariant, even) to the action of running up or down. (Of course, height is a pretty dull measure. But something more interesting, such as heart rate, would not have worked so well in this example.)

Returning to the implementation, it turns out that group actions are encoded as matrices. There is one matrix for each group element. For , the so-called standard representation is a rotation matrix:

In gcnn, the function applying that matrix is left_action_on_R2(). Like its sibling, it is designed to work with batches (of group elements as well as vectors). Technically, what it does is rotate the grid the image is defined on, and then, re-sample the image. To make this more concrete, that method’s code looks about as follows.

Here is a goat.

img_path <- system.file("imgs", "z.jpg", package = "gcnn")
img <- torchvision::base_loader(img_path) |> torchvision::transform_to_tensor()
img$permute(c(2, 3, 1)) |> as.array() |> as.raster() |> plot()

First, we call C_4$left_action_on_R2() to rotate the grid.

# Grid shape is [2, 1024, 1024], for a 2d, 1024 x 1024 image.
img_grid_R2 <- torch::torch_stack(torch::torch_meshgrid(
    list(
      torch::torch_linspace(-1, 1, dim(img)[2]),
      torch::torch_linspace(-1, 1, dim(img)[3])
    )
))

# Transform the image grid with the matrix representation of some group element.
transformed_grid <- C_4$left_action_on_R2(C_4$inverse(g1)$unsqueeze(1), img_grid_R2)

Second, we re-sample the image on the transformed grid. The goat now looks up to the sky.

transformed_img <- torch::nnf_grid_sample(
  img$unsqueeze(1), transformed_grid,
  align_corners = TRUE, mode = "bilinear", padding_mode = "zeros"
)

transformed_img[1,..]$permute(c(2, 3, 1)) |> as.array() |> as.raster() |> plot()

Step 2: The lifting convolution

We want to make use of existing, efficient torch functionality as much as possible. Concretely, we want to use nn_conv2d(). What we need, though, is a convolution kernel that’s equivariant not just to translation, but also to the action of . This can be achieved by having one kernel for each possible rotation.

Implementing that idea is exactly what LiftingConvolution does. The principle is the same as before: First, the grid is rotated, and then, the kernel (weight matrix) is re-sampled to the transformed grid.

Why, though, call this a lifting convolution? The usual convolution kernel operates on ; while our extended version operates on combinations of and . In math speak, it has been lifted to the semi-direct product .

lifting_conv <- LiftingConvolution(
    group = CyclicGroup(order = 4),
    kernel_size = 5,
    in_channels = 3,
    out_channels = 8
  )

x <- torch::torch_randn(c(2, 3, 32, 32))
y <- lifting_conv(x)
y$shape

[1]  2  8  4 28 28

Since, internally, LiftingConvolution uses an additional dimension to realize the product of translations and rotations, the output is not four-, but five-dimensional.

Step 3: Group convolutions

Now that we’re in “group-extended space”, we can chain a number of layers where both input and output are group convolution layers. For example:

group_conv <- GroupConvolution(
  group = CyclicGroup(order = 4),
    kernel_size = 5,
    in_channels = 8,
    out_channels = 16
)

z <- group_conv(y)
z$shape

[1]  2 16  4 24 24

All that remains to be done is package this up. That’s what gcnn::GroupEquivariantCNN() does.

Step 4: Group-equivariant CNN

We can call GroupEquivariantCNN() like so.

cnn <- GroupEquivariantCNN(
    group = CyclicGroup(order = 4),
    kernel_size = 5,
    in_channels = 1,
    out_channels = 1,
    num_hidden = 2, # number of group convolutions
    hidden_channels = 16 # number of channels per group conv layer
)

img <- torch::torch_randn(c(4, 1, 32, 32))
cnn(img)$shape

[1] 4 1

At casual glance, this GroupEquivariantCNN looks like any old CNN … weren’t it for the group argument.

Now, when we inspect its output, we see that the additional dimension is gone. That’s because after a sequence of group-to-group convolution layers, the module projects down to a representation that, for each batch item, retains channels only. It thus averages not just over locations – as we normally do – but over the group dimension as well. A final linear layer will then provide the requested classifier output (of dimension out_channels).

And there we have the complete architecture. It is time for a real-world(ish) test.

Rotated digits!

The idea is to train two convnets, a “normal” CNN and a group-equivariant one, on the usual MNIST training set. Then, both are evaluated on an augmented test set where each image is randomly rotated by a continuous rotation between 0 and 360 degrees. We don’t expect GroupEquivariantCNN to be “perfect” – not if we equip with as a symmetry group. Strictly, with , equivariance extends over four positions only. But we do hope it will perform significantly better than the shift-equivariant-only standard architecture.

First, we prepare the data; in particular, the augmented test set.

dir <- "/tmp/mnist"

train_ds <- torchvision::mnist_dataset(
  dir,
  download = TRUE,
  transform = torchvision::transform_to_tensor
)

test_ds <- torchvision::mnist_dataset(
  dir,
  train = FALSE,
  transform = function(x) {
    x |>
      torchvision::transform_to_tensor() |>
      torchvision::transform_random_rotation(
        degrees = c(0, 360),
        resample = 2,
        fill = 0
      )
  }
)

train_dl <- dataloader(train_ds, batch_size = 128, shuffle = TRUE)
test_dl <- dataloader(test_ds, batch_size = 128)

How does it look?

test_images <- coro::collect(
  test_dl, 1
)[[1]]$x[1:32, 1, , ] |> as.array()

par(mfrow = c(4, 8), mar = rep(0, 4), mai = rep(0, 4))
test_images |>
  purrr::array_tree(1) |>
  purrr::map(as.raster) |>
  purrr::iwalk(~ {
    plot(.x)
  })

We first define and train a conventional CNN. It is as similar to GroupEquivariantCNN(), architecture-wise, as possible, and is given twice the number of hidden channels, so as to have comparable capacity overall.

 default_cnn <- nn_module(
   "default_cnn",
   initialize = function(kernel_size, in_channels, out_channels, num_hidden, hidden_channels) {
     self$conv1 <- torch::nn_conv2d(in_channels, hidden_channels, kernel_size)
     self$convs <- torch::nn_module_list()
     for (i in 1:num_hidden) {
       self$convs$append(torch::nn_conv2d(hidden_channels, hidden_channels, kernel_size))
     }
     self$avg_pool <- torch::nn_adaptive_avg_pool2d(1)
     self$final_linear <- torch::nn_linear(hidden_channels, out_channels)
   },
   forward = function(x) {
     x <- x |>
       self$conv1() |>
       (\(.) torch::nnf_layer_norm(., .$shape[2:4]))() |>
       torch::nnf_relu()
     for (i in 1:(length(self$convs))) {
       x <- x |>
         self$convs[[i]]() |>
         (\(.) torch::nnf_layer_norm(., .$shape[2:4]))() |>
         torch::nnf_relu()
     }
     x <- x |>
       self$avg_pool() |>
       torch::torch_squeeze() |>
       self$final_linear()
     x
   }
 )

fitted <- default_cnn |>
    luz::setup(
      loss = torch::nn_cross_entropy_loss(),
      optimizer = torch::optim_adam,
      metrics = list(
        luz::luz_metric_accuracy()
      )
    ) |>
    luz::set_hparams(
      kernel_size = 5,
      in_channels = 1,
      out_channels = 10,
      num_hidden = 4,
      hidden_channels = 32
    ) %>%
    luz::set_opt_hparams(lr = 1e-2, weight_decay = 1e-4) |>
    luz::fit(train_dl, epochs = 10, valid_data = test_dl) 

Train metrics: Loss: 0.0498 - Acc: 0.9843
Valid metrics: Loss: 3.2445 - Acc: 0.4479

Unsurprisingly, accuracy on the test set is not that great.

Next, we train the group-equivariant version.

fitted <- GroupEquivariantCNN |>
  luz::setup(
    loss = torch::nn_cross_entropy_loss(),
    optimizer = torch::optim_adam,
    metrics = list(
      luz::luz_metric_accuracy()
    )
  ) |>
  luz::set_hparams(
    group = CyclicGroup(order = 4),
    kernel_size = 5,
    in_channels = 1,
    out_channels = 10,
    num_hidden = 4,
    hidden_channels = 16
  ) |>
  luz::set_opt_hparams(lr = 1e-2, weight_decay = 1e-4) |>
  luz::fit(train_dl, epochs = 10, valid_data = test_dl)

Train metrics: Loss: 0.1102 - Acc: 0.9667
Valid metrics: Loss: 0.4969 - Acc: 0.8549

For the group-equivariant CNN, accuracies on test and training sets are a lot closer. That is a nice result! Let’s wrap up today’s exploit resuming a thought from the first, more high-level post.

A challenge

Going back to the augmented test set, or rather, the samples of digits displayed, we notice a problem. In row two, column four, there is a digit that “under normal circumstances”, should be a 9, but, most probably, is an upside-down 6. (To a human, what suggests this is the squiggle-like thing that seems to be found more often with sixes than with nines.) However, you could ask: does this have to be a problem? Maybe the network just needs to learn the subtleties, the kinds of things a human would spot?

The way I view it, it all depends on the context: What really should be accomplished, and how an application is going to be used. With digits on a letter, I’d see no reason why a single digit should appear upside-down; accordingly, full rotation equivariance would be counter-productive. In a nutshell, we arrive at the same canonical imperative advocates of fair, just machine learning keep reminding us of:

Always think of the way an application is going to be used!

In our case, though, there is another aspect to this, a technical one. gcnn::GroupEquivariantCNN() is a simple wrapper, in that its layers all make use of the same symmetry group. In principle, there is no need to do this. With more coding effort, different groups can be used depending on a layer’s position in the feature-detection hierarchy.

Here, let me finally tell you why I chose the goat picture. The goat is seen through a red-and-white fence, a pattern – slightly rotated, due to the viewing angle – made up of squares (or edges, if you like). Now, for such a fence, types of rotation equivariance such as that encoded by make a lot of sense. The goat itself, though, we’d rather not have look up to the sky, the way I illustrated action before. Thus, what we’d do in a real-world image-classification task is use rather flexible layers at the bottom, and increasingly restrained layers at the top of the hierarchy.

Thanks for reading!

Outlook: Group-equivariant CNN

Above, we noted that, in image classification, a translation-invariant operation (like convolution) is needed: A is a whether moved horizontally, vertically, both ways, or not at all. What about rotations, though? Standing on its head, a digit is still what it is. Conventional convolution does not support this type of action.

We can add to our architectural wishlist by specifying a symmetry group. What group? If we wanted to detect squares aligned to the axes, a suitable group would be , the cyclic group of order four. (Above, we saw that we needed four elements, and that we could cycle through the group.) If, on the other hand, we don’t care about alignment, we’d want any position to count. In principle, we should end up in the same situation as we did with the sphere. However, images live on discrete grids; there won’t be an unlimited number of rotations in practice.

With more realistic applications, we need to think more carefully. Take digits. When is a number “the same”? For one, it depends on the context. Were it about a hand-written address on an envelope, would we accept a as such had it been rotated by 90 degrees? Maybe. (Although we might wonder what would make someone change ball-pen position for just a single digit.) What about a standing on its head? On top of similar psychological considerations, we should be seriously unsure about the intended message, and, at least, down-weight the data point were it part of our training set.

Importantly, it also depends on the digit itself. A , upside-down, is a .

Zooming in on neural networks, there is room for yet more complexity. We know that CNNs build up a hierarchy of features, starting from simple ones, like edges and corners. Even if, for later layers, we may not want rotation equivariance, we would still like to have it in the initial set of layers. (The output layer – we’ve hinted at that already – is to be considered separately in any case, since its requirements result from the specifics of what we’re tasked with.)

That’s it for today. Hopefully, I’ve managed to illuminate a bit of why we would want to have group-equivariant neural networks. The question remains: How do we get them? This is what the subsequent posts in the series will be about.

Till then, and thanks for reading!

Three oscillations

By far the best-known – the most infamous, I should say – among the three is El Niño–Southern Oscillation (ENSO), a.k.a. El Niño/La Niña. The term refers to a changing pattern of sea surface temperatures and sea-level pressures occurring in the equatorial Pacific. Both El Niño and La Niña can and do have catastrophic impact on people’s lives, most notably, for people in developing countries west and east of the Pacific.

El Niño occurs when surface water temperatures in the eastern Pacific are higher than normal, and the strong winds that normally blow from east to west are unusually weak. From April to October, this leads to hot, extremely wet weather conditions along the coasts of northern Peru and Ecuador, continually resulting in major floods. La Niña, on the other hand, causes a drop in sea surface temperatures over Southeast Asia as well as heavy rains over Malaysia, the Philippines, and Indonesia. While these are the areas most gravely impacted, changes in ENSO reverberate across the globe.

Less well known than ENSO, but highly influential as well, is the North Atlantic Oscillation (NAO). It strongly affects winter weather in Europe, Greenland, and North America. Its two states relate to the size of the pressure difference between the Icelandic High and the Azores Low. When the pressure difference is high, the jet stream – those strong westerly winds that blow between North America and Northern Europe – is yet stronger than normal, leading to warm, wet European winters and calmer-than-normal conditions in Eastern North America. With a lower-than-normal pressure difference, however, the American East tends to incur more heavy storms and cold-air outbreaks, while winters in Northern Europe are colder and more dry.

Finally, the Arctic Oscillation (AO) is a ring-like pattern of sea-level pressure anomalies centered at the North Pole. (Its Southern-hemisphere equivalent is the Antarctic Oscillation.) AO’s influence extends beyond the Arctic Circle, however; it is indicative of whether and how much Arctic air flows down into the middle latitudes. AO and NAO are strongly related, and might designate the same physical phenomenon at a fundamental level.

Now, let’s make these characterizations more concrete by looking at actual data.

Conclusion

Like I said in the beginning, this post would be about inspiration, not technical detail or reportable results. And I hope that inspirational it has been, at least a little bit. If you’re experimenting with wavelets yourself, or plan to – or if you work in the atmospheric sciences, and would like to provide some insight on the above data/phenomena – we’d love to hear from you!

As always, thanks for reading!

What’s so special about deep learning on graphs?

Graphs are different from images, language, as well as tabular data in that node numbering does not matter. In other words, graphs are permutation-invariant. Already this means that architectures established in other domains cannot be transferred verbatim. (The ideas underlying them can be transferred though. Thus, in the graph neural network (henceforth: GNN) model zoo you’ll see lots of allusions to “convolutional”, “attention”, and other established terms.) Put very simply, and in concordance with common sense, whatever algorithm is used, it will fundamentally be based on how nodes are connected: the edges, that is.

When relationships are modeled as graphs, both nodes and edges can have features. This, too, adds complexity. But not everything is harder with graphs. Think of how cumbersome it can be to obtain labeled data for supervised learning. With graphs, often an astonishingly small amount of labeled data is needed. More surprisingly still, a graph can be constructed when not a single edge is present. Put differently, learning on sets can morph into learning on graphs.

At this point, let me switch gears and move on to the practical part: the raison d’être of this post.

Matching concepts and code: PyTorch Geometric

In this (and future) posts, we’ll make use of PyTorch Geometric (from hereon: PyG), the most popular, at this time, and fastest-growing in terms of functionality as well as user base, library dedicated to graph DL.

Deep learning on graphs, in its most general form, is usually characterized by the term message passing. Messages are passed between nodes that are linked by an edge: If node has three neighbors, it will receive three messages. Those messages have to be summarized in some meaningful way. Finally – GNNs consisting of consecutive layers – the node will have to decide how to modify its previous-layer features (a.k.a. embeddings) based on that summary.

Together, these make up a three-step sequence: collect messages; aggegate; update. What about the “learning” in deep learning, though? There are two places where learning can happen: Firstly, in message collection: Incoming messages could be transformed by a MLP, for example. Secondly, as part of the update step. All in all, this yields mathematical formulae like this, given in the PyG documentation:

Scary though this looks, once we read it from the right, we see that it nicely fits the conceptual description. The are the three types of incoming messages a node can receive: its own state at the previous layer, the states of its neighbors (the nodes ) at the previous layer, and features/embeddings associated to the edge in question. (I’m leaving out edge features in this discussion completely, so as to not further enhance complexity.) These messages are (optionally) transformed by the neural network , and whatever comes out is summarized by the aggregator function . Finally, a node will update itself based on that summary as well as its own previous-layer state, possibly by means of applying neural network .

Now that we have this conceptual/mathematical representation, how does it map to code we see, or would like to write? PyG has excellent, extensive documentation, including at the beginner level. But here, I’d like to spell things out in detail – pedantically, if you like, but in a way that tells us a lot about how GNNs work.

Let’s start by the information given in one of the key documentation pages, Creating message passing networks:

PyG provides the MessagePassing base class, which helps in creating such kinds of message passing graph neural networks by automatically taking care of message propagation. The user only has to define the functions , i.e. message(), and , i.e. update(), as well as the aggregation scheme to use, i.e. aggr="add", aggr="mean" or aggr="max".

Scrolling down that page and looking at the two example implementations, however, we see that an implementation of update() does not have to be provided; and from inspecting the source code, it is clear that, technically, the same holds for message(). (And unless we want a form of aggregation different from the default add, we do not even need to specify that, either.)

Thus, the question becomes: What happens if we code the minimal PyG GNN? To find out, we first need to create a minimal graph, one minimal enough for us to track what is going on.

import torch

x = torch.tensor([[1, 1], [2, 2], [3, 3], [11, 11], [12, 12]], dtype=torch.float)

edge_index = torch.tensor([
  [0, 1],
  [1, 0],
  [0, 2],
  [2, 0],
  [1, 2],
  [2, 1],
  [2, 3],
  [3, 2],
  [2, 4],
  [4, 2],
  [3, 4],
  [4, 3]
], dtype=torch.long)

y = torch.tensor([[0], [0], [0], [1], [1]], dtype=torch.float)

from torch_geometric.data import Data

data = Data(x = x, edge_index = edge_index.t().contiguous(), y = y)
data.x
data.edge_index
data.y

tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 3.,  3.],
        [11., 11.],
        [12., 12.]])
        
tensor([[0, 1, 0, 2, 1, 2, 2, 3, 2, 4, 3, 4],
        [1, 0, 2, 0, 2, 1, 3, 2, 4, 2, 4, 3]])
        
tensor([[0.],
        [0.],
        [0.],
        [1.],
        [1.]])

import networkx as nx
import matplotlib.pyplot as plt
from torch_geometric.utils import to_networkx

def visualize_graph(G, color, labels):
    plt.figure(figsize=(7,7))
    plt.axis('off')
    nx.draw_networkx(
      G,
      pos = nx.spring_layout(G, seed = 777),
      labels = labels,
      node_color = color,
      cmap = "Set3"
      )
    plt.show()

G = to_networkx(data, to_undirected = True, node_attrs = ["x"])
labels = nx.get_node_attributes(G, "x")
visualize_graph(G, color = data.y, labels = labels)

from torch_geometric.nn import MessagePassing

class IAmLazy(MessagePassing):
    def forward(self, x, edge_index):
        out = self.propagate(edge_index, x = x)
        return out
      
module = IAmLazy()
out = module(data.x, data.edge_index)
out

tensor([[ 5.,  5.],
        [ 4.,  4.],
        [26., 26.],
        [15., 15.],
        [14., 14.]])

class IAmMyOthers(MessagePassing):
    def forward(self, x, edge_index):
        out = self.propagate(edge_index, x = x)
        return out
    def message(self, x_j):
        print("in message, x_j is")
        print(x_j)
        return x_j
      
module = IAmMyOthers()
out = module(data.x, data.edge_index)
print("result is:")
out

in message, x_j is
tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 1.,  1.],
        [ 3.,  3.],
        [ 2.,  2.],
        [ 3.,  3.],
        [ 3.,  3.],
        [11., 11.],
        [ 3.,  3.],
        [12., 12.],
        [11., 11.],
        [12., 12.]])
        
result is:
tensor([[ 5.,  5.],
        [ 4.,  4.],
        [26., 26.],
        [15., 15.],
        [14., 14.]])

tensor([[0, 1, 0, 2, 1, 2, 2, 3, 2, 4, 3, 4],
        [1, 0, 2, 0, 2, 1, 3, 2, 4, 2, 4, 3]])

data.edge_index[0]

data.x[data.edge_index[0]]

tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 1.,  1.],
        [ 3.,  3.],
        [ 2.,  2.],
        [ 3.,  3.],
        [ 3.,  3.],
        [11., 11.],
        [ 3.,  3.],
        [12., 12.],
        [11., 11.],
        [12., 12.]])

from torch_geometric.utils import add_self_loops

class IAmMyOthersAndMyselfAsWell(MessagePassing):
    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes = x.size(0))
        print("in forward, augmented edge index now has shape")
        print(edge_index.shape)
        out = self.propagate(edge_index, x = x)
        return out
    def message(self, x_j):
        return x_j

module = IAmMyOthersAndMyselfAsWell()
out = module(data.x, data.edge_index)
print("result is:")
out

in forward, augmented edge index now has shape:
torch.Size([2, 17])

result is:
tensor([[ 6.,  6.],
        [ 6.,  6.],
        [29., 29.],
        [26., 26.],
        [26., 26.]])

def __init__(self):
        super().__init__(aggr = "mean")

from torch_scatter import scatter

class IAmJustTheOppositeReally(MessagePassing):
    def forward(self, x, edge_index):
        out = self.propagate(edge_index, x = x)
        return out
    def aggregate(self, inputs, index):
        print("in aggregate, inputs is")
        # same as x_j (incoming node features)
        print(inputs)
        print("in aggregate, index is")
        # this is data.edge_index[1]
        print(index)
        # see https://pytorch-scatter.readthedocs.io/en/1.3.0/index.html
        # for other aggregation modes
        # default dim is -1
        return - scatter(inputs, index, dim = 0, reduce = "add") 
      
module = IAmJustTheOppositeReally()
out = module(data.x, data.edge_index)
print("result is:")
out

in aggregate, inputs is
tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 1.,  1.],
        [ 3.,  3.],
        [ 2.,  2.],
        [ 3.,  3.],
        [ 3.,  3.],
        [11., 11.],
        [ 3.,  3.],
        [12., 12.],
        [11., 11.],
        [12., 12.]])
        
in aggregate, index is
tensor([1, 0, 2, 0, 2, 1, 3, 2, 4, 2, 4, 3])

result is:
tensor([[ -5.,  -5.],
        [ -4.,  -4.],
        [-26., -26.],
        [-15., -15.],
        [-14., -14.]])

class IDoEvolveOverTime(MessagePassing):
    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        out = self.propagate(edge_index, x = x)
        return out
    def update(self, inputs, x):
        print("in update, inputs is")
        print(inputs)
        print("in update, x is")
        print(x)
        return (inputs + x)/2

module = IDoEvolveOverTime()
out = module(data.x, data.edge_index)
print("result is:")
out

in update, inputs is
tensor([[ 6.,  6.],
        [ 6.,  6.],
        [29., 29.],
        [26., 26.],
        [26., 26.]])
in update, x is
tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 3.,  3.],
        [11., 11.],
        [12., 12.]])
result is:
tensor([[ 3.5000,  3.5000],
        [ 4.0000,  4.0000],
        [16.0000, 16.0000],
        [18.5000, 18.5000],
        [19.0000, 19.0000]])

from torch.nn import Sequential as Seq, Linear, ReLU

class ILearnAndEvolve(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr = "sum")
        self.mlp = Seq(Linear(in_channels, out_channels),
                       ReLU(),
                       Linear(out_channels, out_channels))
    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        x = self.mlp(x)
        out = self.propagate(edge_index = edge_index, x = x)
        return out
    def update(self, inputs, x):
        return (inputs + x)/2

module = ILearnAndEvolve(2, 2)
out = module(data.x, data.edge_index)
print("result is:")
out

result is:
tensor([[-0.8724, -0.4407],
        [-0.9056, -0.4623],
        [-2.0229, -1.1240],
        [-1.8691, -1.0867],
        [-1.9024, -1.1082]], grad_fn=<DivBackward0>)

class ILearnAndEvolveDoubly(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr = "sum")
        self.mlp_msg = Seq(Linear(in_channels, out_channels),
                       ReLU(),
                       Linear(out_channels, out_channels))
        self.mlp_upd = Seq(Linear(out_channels, out_channels),
                       ReLU(),
                       Linear(out_channels, out_channels))
    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        x = self.mlp_msg(x)
        out = self.propagate(edge_index = edge_index, x = x)
        return out
    def update(self, inputs, x):
        return self.mlp_upd((inputs + x)/2)

module = ILearnAndEvolveDoubly(2, 2)
out = module(data.x, data.edge_index)
print("result is:")
out

result is:
tensor([[ 0.0573, -0.6988],
        [ 0.0358, -0.6894],
        [-0.1730, -0.6450],
        [-0.5855, -0.4171],
        [-0.5890, -0.4141]], grad_fn=<AddmmBackward0>)

class Network(torch.nn.Module):
    def __init__(self, in_channels, out_channels, num_classes):
        super().__init__()
        self.conv = ILearnAndEvolveDoubly(in_channels, out_channels)
        self.classifier = Linear(out_channels, num_classes)
    def forward(self, x, edge_index):
        x, edge_index = data.x, data.edge_index
        x = self.conv(x, edge_index)
        return self.classifier(x)

model = Network(2, 2, 1) 

import torch.nn.functional as F

optimizer = torch.optim.Adam(model.parameters(), lr = 0.01)
model.train()

for epoch in range(5):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.binary_cross_entropy_with_logits(out, data.y)
    loss.backward()
    optimizer.step()

preds = torch.sigmoid(out)
preds

tensor([[0.6502],
        [0.6532],
        [0.7027],
        [0.7145],
        [0.7165]], grad_fn=<SigmoidBackward0>)

Torch dataset

The dataset is responsible for correctly matching up inputs and targets.

Our goal is to take six months of global sea surface temperatures and predict the Niño 3.4 Index for the following month. Input-wise, the model will expect the following format semantics:

batch_size * timesteps * width * height * channels, where

• batch_size is the number of observations worked on in one round of computations,

• timesteps chains consecutive observations from adjacent months,

• width and height together constitute the spatial grid, and

• channels corresponds to available visual channels in the “image”.

In .getitem(), we select the consecutive observations, starting at a given index, and stack them in dimension one. (One, not two, as batches will only start to exist once the dataloader comes into play.)

Now, what about the target? Our ultimate goal was – is – predicting the Niño 3.4 Index. However, as you see we define three targets: One is the index, as expected; an additional one holds the spatially-gridded sea surface temperatures for the prediction month. Why? Our main instrument, the most prominent constituent of the model, will be a convLSTM, an architecture designed for spatial prediction. Thus, to train it efficiently, we want to give it the opportunity to predict values on a spatial grid. So far so good; but there’s one more target, the phase/category. This was added for experimentation purposes: Maybe predicting both index and phase helps in training?

Finally, here is the code for the dataset. In our experiments, we based predictions on inputs from the preceding six months (n_timesteps <- 6). This is a parameter you might want to play with, though.

n_timesteps <- 6

enso_dataset <- dataset(
  name = "enso_dataset",
  
  initialize = function(sst, nino, n_timesteps) {
    self$sst <- sst
    self$nino <- nino
    self$n_timesteps <- n_timesteps
  },
  
  .getitem = function(i) {
    x <- torch_tensor(self$sst[, , i:(n_timesteps + i - 1)]) # (360, 115, n_timesteps)
    x <- x$permute(c(3,1,2))$unsqueeze(2) # (n_timesteps, 1, 360, 115))
    
    y1 <- torch_tensor(self$sst[, , n_timesteps + i])$unsqueeze(1) # (1, 360, 115)
    y2 <- torch_tensor(self$nino$NINO34_MEAN[n_timesteps + i])
    y3 <- torch_tensor(self$nino$phase_code[n_timesteps + i])$squeeze()$to(torch_long())
    list(x = x, y1 = y1, y2 = y2, y3 = y3)
  },
  
  .length = function() {
    nrow(self$nino) - n_timesteps
  }
  
)

valid_ds <- enso_dataset(sst_valid, nino_valid, n_timesteps)

Dataloaders

After the custom dataset, we create the – pretty typical – dataloaders, making use of a batch size of 4.

batch_size <- 4

train_dl <- train_ds %>% dataloader(batch_size = batch_size, shuffle = TRUE)

valid_dl <- valid_ds %>% dataloader(batch_size = batch_size)

Next, we proceed to model creation.

Model

The model’s main ingredient is the convLSTM introduced in a prior post. For convenience, we reproduce the code in the appendix.

Besides the convLSTM, the model makes use of three convolutional layers, a batchnorm layer and five linear layers. The logic is the following.

First, the convLSTM job is to predict the next month’s sea surface temperatures on the spatial grid. For that, we almost just return its final state, - almost: We use self$conv1 to reduce the number channels to one.

For predicting index and phase, we then need to flatten the grid, as we require a single value each. This is where the additional conv layers come in. We do hope they’ll aid in learning, but we also want to reduce the number of parameters a bit, downsizing the grid (strides = 2 and strides = 3, resp.) a bit before the upcoming torch_flatten().

Once we have a flat structure, learning is shared between the tasks of index and phase prediction (self$linear), until finally their paths split (self$cont and self$cat, resp.), and they return their separate outputs.

(The batchnorm? I’ll comment on that in the Discussion.)

model <- nn_module(
  
  initialize = function(channels_in,
                        convlstm_hidden,
                        convlstm_kernel,
                        convlstm_layers) {
    
    self$n_layers <- convlstm_layers
    
    self$convlstm <- convlstm(
      input_dim = channels_in,
      hidden_dims = convlstm_hidden,
      kernel_sizes = convlstm_kernel,
      n_layers = convlstm_layers
    )
    
    self$conv1 <-
      nn_conv2d(
        in_channels = 32,
        out_channels = 1,
        kernel_size = 5,
        padding = 2
      )
    
    self$conv2 <-
      nn_conv2d(
        in_channels = 32,
        out_channels = 32,
        kernel_size = 5,
        stride = 2
      )
    
    self$conv3 <-
      nn_conv2d(
        in_channels = 32,
        out_channels = 32,
        kernel_size = 5,
        stride = 3
      )
    
    self$linear <- nn_linear(33408, 64)
    
    self$b1 <- nn_batch_norm1d(num_features = 64)
        
    self$cont <- nn_linear(64, 128)
    self$cat <- nn_linear(64, 128)
    
    self$cont_output <- nn_linear(128, 1)
    self$cat_output <- nn_linear(128, 3)
    
  },
  
  forward = function(x) {
    
    ret <- self$convlstm(x)
    layer_last_states <- ret[[2]]
    last_hidden <- layer_last_states[[self$n_layers]][[1]]
    
    next_sst <- last_hidden %>% self$conv1() 
    
    c2 <- last_hidden %>% self$conv2() 
    c3 <- c2 %>% self$conv3() 
    
    flat <- torch_flatten(c3, start_dim = 2)
    common <- self$linear(flat) %>% self$b3() %>% nnf_relu()

    next_temp <- common %>% self$cont() %>% nnf_relu() %>% self$cont_output()
    next_nino <- common %>% self$cat() %>% nnf_relu() %>% self$cat_output()
    
    list(next_sst, next_temp, next_nino)
    
  }
  
)

Next, we instantiate a pretty small-ish model. You’re more than welcome to experiment with larger models, but training time as well as GPU memory requirements will increase.

net <- model(
  channels_in = 1,
  convlstm_hidden = c(16, 16, 32),
  convlstm_kernel = c(3, 3, 5),
  convlstm_layers = 3
)

device <- torch_device(if (cuda_is_available()) "cuda" else "cpu")

net <- net$to(device = device)
net

An `nn_module` containing 2,389,605 parameters.

── Modules ───────────────────────────────────────────────────────────────────────────────
● convlstm: <nn_module> #182,080 parameters
● conv1: <nn_conv2d> #801 parameters
● conv2: <nn_conv2d> #25,632 parameters
● conv3: <nn_conv2d> #25,632 parameters
● linear: <nn_linear> #2,138,176 parameters
● b1: <nn_batch_norm1d> #128 parameters
● cont: <nn_linear> #8,320 parameters
● cat: <nn_linear> #8,320 parameters
● cont_output: <nn_linear> #129 parameters
● cat_output: <nn_linear> #387 parameters

Training

We have three model outputs. How should we combine the losses?

Given that the main goal is predicting the index, and the other two outputs are essentially means to an end, I found the following combination rather effective:

# weight for sea surface temperature prediction
lw_sst <- 0.2

# weight for prediction of El Nino 3.4 Index
lw_temp <- 0.4

# weight for phase prediction
lw_nino <- 0.4

The training process follows the pattern seen in all torch posts so far: For each epoch, loop over the training set, backpropagate, check performance on validation set.

But, when we did the pre-processing, we were aware of an imminent problem: the missing temperatures for continental areas, which we set to zero. As a sole measure, this approach is clearly insufficient. What if we had chosen to use latitude-dependent averages? Or interpolation? Both may be better than a global average, but both have their problems as well. Let’s at least alleviate negative consequences by not using the respective pixels for spatial loss calculation. This is taken care of by the following line below:

sst_loss <- nnf_mse_loss(sst_output[sst_target != 0], sst_target[sst_target != 0])

Here, then, is the complete training code.

optimizer <- optim_adam(net$parameters, lr = 0.001)

num_epochs <- 50

train_batch <- function(b) {
  
  optimizer$zero_grad()
  output <- net(b$x$to(device = device))
  
  sst_output <- output[[1]]
  sst_target <- b$y1$to(device = device)
  
  sst_loss <- nnf_mse_loss(sst_output[sst_target != 0], sst_target[sst_target != 0])
  temp_loss <- nnf_mse_loss(output[[2]], b$y2$to(device = device))
  nino_loss <- nnf_cross_entropy(output[[3]], b$y3$to(device = device))
  
  loss <- lw_sst * sst_loss + lw_temp * temp_loss + lw_nino * nino_loss
  loss$backward()
  optimizer$step()

  list(sst_loss$item(), temp_loss$item(), nino_loss$item(), loss$item())
  
}

valid_batch <- function(b) {
  
  output <- net(b$x$to(device = device))
  
  sst_output <- output[[1]]
  sst_target <- b$y1$to(device = device)
  
  sst_loss <- nnf_mse_loss(sst_output[sst_target != 0], sst_target[sst_target != 0])
  temp_loss <- nnf_mse_loss(output[[2]], b$y2$to(device = device))
  nino_loss <- nnf_cross_entropy(output[[3]], b$y3$to(device = device))
  
  loss <-
    lw_sst * sst_loss + lw_temp * temp_loss + lw_nino * nino_loss

  list(sst_loss$item(),
       temp_loss$item(),
       nino_loss$item(),
       loss$item())
}

for (epoch in 1:num_epochs) {
  
  net$train()
  
  train_loss_sst <- c()
  train_loss_temp <- c()
  train_loss_nino <- c()
  train_loss <- c()

  coro::loop(for (b in train_dl) {
    losses <- train_batch(b)
    train_loss_sst <- c(train_loss_sst, losses[[1]])
    train_loss_temp <- c(train_loss_temp, losses[[2]])
    train_loss_nino <- c(train_loss_nino, losses[[3]])
    train_loss <- c(train_loss, losses[[4]])
  })
  
  cat(
    sprintf(
      "\nEpoch %d, training: loss: %3.3f sst: %3.3f temp: %3.3f nino: %3.3f \n",
      epoch, mean(train_loss), mean(train_loss_sst), mean(train_loss_temp), mean(train_loss_nino)
    )
  )
  
  net$eval()
  
  valid_loss_sst <- c()
  valid_loss_temp <- c()
  valid_loss_nino <- c()
  valid_loss <- c()

  coro::loop(for (b in valid_dl) {
    losses <- valid_batch(b)
    valid_loss_sst <- c(valid_loss_sst, losses[[1]])
    valid_loss_temp <- c(valid_loss_temp, losses[[2]])
    valid_loss_nino <- c(valid_loss_nino, losses[[3]])
    valid_loss <- c(valid_loss, losses[[4]])
    
  })
  
  cat(
    sprintf(
      "\nEpoch %d, validation: loss: %3.3f sst: %3.3f temp: %3.3f nino: %3.3f \n",
      epoch, mean(valid_loss), mean(valid_loss_sst), mean(valid_loss_temp), mean(valid_loss_nino)
    )
  )
  
  torch_save(net, paste0(
    "model_", epoch, "_", round(mean(train_loss), 3), "_", round(mean(valid_loss), 3), ".pt"
  ))
  
}

When I ran this, performance on the training set decreased in a not-too-fast, but continuous way, while validation set performance kept fluctuating. For reference, total (composite) losses looked like this:

Epoch     Training    Validation
   
   10        0.336         0.633
   20        0.233         0.295
   30        0.135         0.461
   40        0.099         0.903
   50        0.061         0.727

Thinking of the size of the validation set - thirty-one years, or equivalently, 372 data points – those fluctuations may not be all too surprising.

Discussion

When working with small amounts of data, a lot can be learned by quick-ish experimentation. However, when at the same time, the task is complex, one should be cautious extrapolating.

For example, well-established regularizers such as batchnorm and dropout, while intended to improve generalization to the validation set, may turn out to severely impede training itself. This is the story behind the one batchnorm layer I kept (I did try having more), and it is also why there is no dropout.

One lesson to learn from this experience then is: Make sure the amount of data matches the complexity of the task. This is what we see in the ENSO prediction papers published on arxiv.

If we should treat the results with caution, why even publish the post?

For one, it shows an application of convLSTM to real-world data, employing a reasonably complex architecture and illustrating techniques like custom losses and loss masking. Similar architectures and strategies should be applicable to a wide range of real-world tasks – basically, whenever predictors in a time-series problem are given on a spatial grid.

Secondly, the application itself – forecasting an atmospheric phenomenon that greatly affects ecosystems as well as human well-being – seems like an excellent use of deep learning. Applications like these stand out as all the more worthwhile as the same cannot be said of everything deep learning is – and will be, barring effective regulation – used for.

Thanks for reading!

Appendix: convlstm code

For an in-depth explanation of convlstm, see the blog post.

library(torch)
library(torchvision)

convlstm_cell <- nn_module(
  
  initialize = function(input_dim, hidden_dim, kernel_size, bias) {
    
    self$hidden_dim <- hidden_dim
    
    padding <- kernel_size %/% 2
    
    self$conv <- nn_conv2d(
      in_channels = input_dim + self$hidden_dim,
      # for each of input, forget, output, and cell gates
      out_channels = 4 * self$hidden_dim,
      kernel_size = kernel_size,
      padding = padding,
      bias = bias
    )
  },
  
  forward = function(x, prev_states) {

    h_prev <- prev_states[[1]]
    c_prev <- prev_states[[2]]
    
    combined <- torch_cat(list(x, h_prev), dim = 2)  # concatenate along channel axis
    combined_conv <- self$conv(combined)
    gate_convs <- torch_split(combined_conv, self$hidden_dim, dim = 2)
    cc_i <- gate_convs[[1]]
    cc_f <- gate_convs[[2]]
    cc_o <- gate_convs[[3]]
    cc_g <- gate_convs[[4]]
    
    # input, forget, output, and cell gates (corresponding to torch's LSTM)
    i <- torch_sigmoid(cc_i)
    f <- torch_sigmoid(cc_f)
    o <- torch_sigmoid(cc_o)
    g <- torch_tanh(cc_g)
    
    # cell state
    c_next <- f * c_prev + i * g
    # hidden state
    h_next <- o * torch_tanh(c_next)
    
    list(h_next, c_next)
    
  },
  
  init_hidden = function(batch_size, height, width) {
    list(torch_zeros(batch_size, self$hidden_dim, height, width, device = self$conv$weight$device),
         torch_zeros(batch_size, self$hidden_dim, height, width, device = self$conv$weight$device))
  }
)

convlstm <- nn_module(
  
  initialize = function(input_dim, hidden_dims, kernel_sizes, n_layers, bias = TRUE) {
 
    self$n_layers <- n_layers
    
    self$cell_list <- nn_module_list()
    
    for (i in 1:n_layers) {
      cur_input_dim <- if (i == 1) input_dim else hidden_dims[i - 1]
      self$cell_list$append(convlstm_cell(cur_input_dim, hidden_dims[i], kernel_sizes[i], bias))
    }
  },
  
  # we always assume batch-first
  forward = function(x) {
    
    batch_size <- x$size()[1]
    seq_len <- x$size()[2]
    height <- x$size()[4]
    width <- x$size()[5]
   
    # initialize hidden states
    init_hidden <- vector(mode = "list", length = self$n_layers)
    for (i in 1:self$n_layers) {
      init_hidden[[i]] <- self$cell_list[[i]]$init_hidden(batch_size, height, width)
    }
    
    # list containing the outputs, of length seq_len, for each layer
    # this is the same as h, at each step in the sequence
    layer_output_list <- vector(mode = "list", length = self$n_layers)
    
    # list containing the last states (h, c) for each layer
    layer_state_list <- vector(mode = "list", length = self$n_layers)

    cur_layer_input <- x
    hidden_states <- init_hidden
    
    # loop over layers
    for (i in 1:self$n_layers) {
      
      # every layer's hidden state starts from 0 (non-stateful)
      h_c <- hidden_states[[i]]
      h <- h_c[[1]]
      c <- h_c[[2]]
      # outputs, of length seq_len, for this layer
      # equivalently, list of h states for each time step
      output_sequence <- vector(mode = "list", length = seq_len)
      
      # loop over timesteps
      for (t in 1:seq_len) {
        h_c <- self$cell_list[[i]](cur_layer_input[ , t, , , ], list(h, c))
        h <- h_c[[1]]
        c <- h_c[[2]]
        # keep track of output (h) for every timestep
        # h has dim (batch_size, hidden_size, height, width)
        output_sequence[[t]] <- h
      }

      # stack hs for all timesteps over seq_len dimension
      # stacked_outputs has dim (batch_size, seq_len, hidden_size, height, width)
      # same as input to forward (x)
      stacked_outputs <- torch_stack(output_sequence, dim = 2)
      
      # pass the list of outputs (hs) to next layer
      cur_layer_input <- stacked_outputs
      
      # keep track of list of outputs or this layer
      layer_output_list[[i]] <- stacked_outputs
      # keep track of last state for this layer
      layer_state_list[[i]] <- list(h, c)
    }
 
    list(layer_output_list, layer_state_list)
      
  }
    
)

Domains, priors, architectures

Given cues like “local” and “pooling”, what better architecture is there to start with than CNNs, the (still) paradigmatic deep learning architecture? Probably, it’s also the one a prototypic practitioner would be most familiar with.

The forest for the trees

Nowadays, it is hard to find a modeling framework that does not include functionality to assess fairness. (Or is at least planning to.) And the terminology sounds so familiar, as well: “calibration”, “predictive parity”, “equal true [false] positive rate”… It almost seems as though we could just take the metrics we make use of anyway (recall or precision, say), test for equality across groups, and that’s it. Let’s assume, for a second, it really was that simple. Then the question still is: Which metrics, exactly, do we choose?

In reality things are not simple. And it gets worse. For very good reasons, there is a close connection in the ML fairness literature to concepts that are primarily treated in other disciplines, such as the legal sciences: discrimination and disparate impact (both not being far from yet another statistical concept, statistical parity). Statistical parity means that if we have a classifier, say to decide whom to hire, it should result in as many applicants from the disadvantaged group (e.g., Black people) being hired as from the advantaged one(s). But that is quite a different requirement from, say, equal true/false positive rates!

So despite all that abundance of software, guides, and decision trees, even: This is not a simple, technical decision. It is, in fact, a technical decision only to a small degree.

Common sense, not math

Let me start this section with a disclaimer: Most of the sources referenced in this text appear, or are implied on the “Guidance” page of IBM’s framework AI Fairness 360. If you read that page, and everything that’s said and not said there appears clear from the outset, then you may not need this more verbose exposition. If not, I invite you to read on.

Papers on fairness in machine learning, as is common in fields like computer science, abound with formulae. Even the papers referenced here, though selected not for their theorems and proofs but for the ideas they harbor, are no exception. But to start thinking about fairness as it might apply to an ML process at hand, common language – and common sense – will do just fine. If, after analyzing your use case, you judge that the more technical results are relevant to the process in question, you will find that their verbal characterizations will often suffice. It is only when you doubt their correctness that you will need to work through the proofs.

At this point, you may be wondering what it is I am contrasting those “more technical results” with. This is the topic of the next section, where I’ll try to give a birds-eye characterization of fairness criteria and what they imply.

(A) Conclusion

In line with its original goal – to provide some help in starting to think about AI fairness metrics – this article does not end with recommendations. It does, however, end with an observation. As the last section has shown, amidst all theorems and theories, all proofs and memes, it makes sense to not lose sight of the concrete: the data trained on, and the ML process as a whole. Fairness is not something to be evaluated post hoc; the feasibility of fairness is to be reflected on right from the beginning.

In that regard, assessing impact on fairness is not that different from that essential, but often toilsome and non-beloved, stage of modeling that precedes the modeling itself: exploratory data analysis.

Thanks for reading!

What to expect from this post

Today, we won’t jump into real-world applications just yet. Instead, we’ll take our time to build a convolutional LSTM (henceforth: convLSTM) in torch. For one, we have to – there is no official PyTorch implementation.

(Keras, on the other hand, has one. If you’re interested in quickly playing around with a Keras convLSTM, check out this nice example.)

What’s more, this post can serve as an introduction to building your own modules. This is something you may be familiar with from Keras or not – depending on whether you’ve used custom models or rather, preferred the declarative define -> compile -> fit style. (Yes, I’m implying there’s some transfer going on if one comes to torch from Keras custom training. Syntactic and semantic details may be different, but both share the object-oriented style that allows for great flexibility and control.)

Last but not least, we’ll also use this as a hands-on experience with RNN architectures (the LSTM, specifically). While the general concept of recurrence may be easy to grasp, it is not necessarily self-evident how those architectures should, or could, be coded. Personally, I find that independent of the framework used, RNN-related documentation leaves me confused. What exactly is being returned from calling an LSTM, or a GRU? (In Keras this depends on how you’ve defined the layer in question.) I suspect that once we’ve decided what we want to return, the actual code won’t be that complicated. Consequently, we’ll take a detour clarifying what it is that torch and Keras are giving us. Implementing our convLSTM will be a lot more straightforward thereafter.

Appendix

This appendix contains the code used to create tables 1 and 2 above.

library(keras)

# batch of 3, with 4 time steps each and a single feature
input <- k_random_normal(shape = c(3L, 4L, 1L))
input

# default args
# return shape = (batch_size, units)
lstm <- layer_lstm(
  units = 1,
  kernel_initializer = initializer_constant(value = 1),
  recurrent_initializer = initializer_constant(value = 1)
)
lstm(input)

# return_sequences = TRUE
# return shape = (batch_size, time steps, units)
#
# note how for each item in the batch, the value for time step 4 equals that obtained above
lstm <- layer_lstm(
  units = 1,
  return_sequences = TRUE,
  kernel_initializer = initializer_constant(value = 1),
  recurrent_initializer = initializer_constant(value = 1)
  # bias is by default initialized to 0
)
lstm(input)

# return_state = TRUE
# return shape = list of:
#                - outputs, of shape: (batch_size, units)
#                - "memory states" for the last time step, of shape: (batch_size, units)
#                - "carry states" for the last time step, of shape: (batch_size, units)
#
# note how the first and second list items are identical!
lstm <- layer_lstm(
  units = 1,
  return_state = TRUE,
  kernel_initializer = initializer_constant(value = 1),
  recurrent_initializer = initializer_constant(value = 1)
)
lstm(input)

# return_state = TRUE, return_sequences = TRUE
# return shape = list of:
#                - outputs, of shape: (batch_size, time steps, units)
#                - "memory" states for the last time step, of shape: (batch_size, units)
#                - "carry states" for the last time step, of shape: (batch_size, units)
#
# note how again, the "memory" state found in list item 2 matches the final-time step outputs reported in item 1
lstm <- layer_lstm(
  units = 1,
  return_sequences = TRUE,
  return_state = TRUE,
  kernel_initializer = initializer_constant(value = 1),
  recurrent_initializer = initializer_constant(value = 1)
)
lstm(input)

# default args
# return shape = (batch_size, units)
gru <- layer_gru(
  units = 1,
  kernel_initializer = initializer_constant(value = 1),
  recurrent_initializer = initializer_constant(value = 1)
)
gru(input)

# return_sequences = TRUE
# return shape = (batch_size, time steps, units)
#
# note how for each item in the batch, the value for time step 4 equals that obtained above
gru <- layer_gru(
  units = 1,
  return_sequences = TRUE,
  kernel_initializer = initializer_constant(value = 1),
  recurrent_initializer = initializer_constant(value = 1)
)
gru(input)

# return_state = TRUE
# return shape = list of:
#    - outputs, of shape: (batch_size, units)
#    - "memory" states for the last time step, of shape: (batch_size, units)
#
# note how the list items are identical!
gru <- layer_gru(
  units = 1,
  return_state = TRUE,
  kernel_initializer = initializer_constant(value = 1),
  recurrent_initializer = initializer_constant(value = 1)
)
gru(input)

# return_state = TRUE, return_sequences = TRUE
# return shape = list of:
#    - outputs, of shape: (batch_size, time steps, units)
#    - "memory states" for the last time step, of shape: (batch_size, units)
#
# note how again, the "memory state" found in list item 2 matches the final-time-step outputs reported in item 1
gru <- layer_gru(
  units = 1,
  return_sequences = TRUE,
  return_state = TRUE,
  kernel_initializer = initializer_constant(value = 1),
  recurrent_initializer = initializer_constant(value = 1)
)
gru(input)

library(torch)

# batch of 3, with 4 time steps each and a single feature
# we will specify batch_first = TRUE when creating the LSTM
input <- torch_randn(c(3, 4, 1))
input

# default args
# return shape = (batch_size, units)
#
# note: there is an additional argument num_layers that we could use to specify a stacked LSTM - effectively composing two LSTM modules
# default for num_layers is 1 though 
lstm <- nn_lstm(
  input_size = 1, # number of input features
  hidden_size = 1, # number of hidden (and output!) features
  batch_first = TRUE # for easy comparability with Keras
)

nn_init_constant_(lstm$weight_ih_l1, 1)
nn_init_constant_(lstm$weight_hh_l1, 1)
nn_init_constant_(lstm$bias_ih_l1, 0)
nn_init_constant_(lstm$bias_hh_l1, 0)

# returns a list of length 2, namely
#   - outputs, of shape (batch_size, time steps, hidden_size) - given we specified batch_first
#       Note 1: If this is a stacked LSTM, these are the outputs from the last layer only.
#               For our current purpose, this is irrelevant, as we're restricting ourselves to single-layer LSTMs.
#       Note 2: hidden_size here is equivalent to units in Keras - both specify number of features
#  - list of:
#    - hidden state for the last time step, of shape (num_layers, batch_size, hidden_size)
#    - cell state for the last time step, of shape (num_layers, batch_size, hidden_size)
#      Note 3: For a single-layer LSTM, the hidden states are already provided in the first list item.

lstm(input)

# default args
# return shape = (batch_size, units)
#
# note: there is an additional argument num_layers that we could use to specify a stacked GRU - effectively composing two GRU modules
# default for num_layers is 1 though 
gru <- nn_gru(
  input_size = 1, # number of input features
  hidden_size = 1, # number of hidden (and output!) features
  batch_first = TRUE # for easy comparability with Keras
)

nn_init_constant_(gru$weight_ih_l1, 1)
nn_init_constant_(gru$weight_hh_l1, 1)
nn_init_constant_(gru$bias_ih_l1, 0)
nn_init_constant_(gru$bias_hh_l1, 0)

# returns a list of length 2, namely
#   - outputs, of shape (batch_size, time steps, hidden_size) - given we specified batch_first
#       Note 1: If this is a stacked GRU, these are the outputs from the last layer only.
#               For our current purpose, this is irrelevant, as we're restricting ourselves to single-layer GRUs.
#       Note 2: hidden_size here is equivalent to units in Keras - both specify number of features
#  - list of:
#    - hidden state for the last time step, of shape (num_layers, batch_size, hidden_size)
#    - cell state for the last time step, of shape (num_layers, batch_size, hidden_size)
#       Note 3: For a single-layer GRU, these values are already provided in the first list item.
gru(input)