divergences

What are Large Language Models? What are they not?

Sigrid Keydana — Tue, 20 Jun 2023 00:00:00 GMT

“At this writing, the only serious ELIZA scripts which exist are some which cause ELIZA to respond roughly as would certain psychotherapists (Rogerians). ELIZA performs best when its human correspondent is initially instructed to”talk” to it, via the typewriter of course, just as one would to a psychiatrist. This mode of conversation was chosen because the psychiatric interview is one of the few examples of categorized dyadic natural language communication in which one of the participating pair is free to assume the pose of knowing almost nothing of the real world. If, for example, one were to tell a psychiatrist “I went for a long boat ride” and he responded “Tell me about boats”, one would not assume that he knew nothing about boats, but that he had some purpose in so directing the subsequent conversation. It is important to note that this assumption is one made by the speaker. Whether it is realistic or not is an altogether separate question. In any case, it has a crucial psychological utility in that it serves the speaker to maintain his sense of being heard and understood. The speaker furher defends his impression (which even in real life may be illusory) by attributing to his conversational partner all sorts of background knowledge, insights and reasoning ability. But again, these are the speaker’s contribution to the conversation.”

Joseph Weizenbaum, creator of ELIZA (Weizenbaum 1966).

Weizenbaum, Joseph. 1966. “ELIZA - a Computer Program for the Study of Natural Language Communication Between Man and Machine.” Commun. ACM (New York, NY, USA) 9 (1): 36–45. https://doi.org/10.1145/365153.365168.

GPT, the ancestor all numbered GPTs, was released in June, 2018 – five years ago, as I write this. Five years: that’s a long time. It certainly is as measured on the time scale of deep learning, the thing that is, usually, behind when people talk of “AI”. One year later, GPT was followed by GPT-2; another year later, by GPT-3. At this point, public attention was still modest – as expected, really, for these kinds of technologies that require lots of specialist knowledge. (For GPT-2, what may have increased attention beyond the normal, a bit, was OpenAI ’s refusal to publish the complete training code and full model weights, supposedly due to the threat posed by the model’s capabilities – alternatively, as argued by others, as a marketing strategy, or yet alternatively, as a way to preserve one’s own competitive advantage just a tiny little bit longer.

As of 2023, with GPT-3.5 and GPT-4 having followed, everything looks different. (Almost) everyone seems to know GPT, at least when that acronym appears prefixed by a certain syllable. Depending on who you talk to, people don’t seem to stop talking about that fantastic [insert thing here] ChatGPT generated for them, about its enormous usefulness with respect to [insert goal here]… or about the flagrant mistakes it made, and the danger that legal regulation and political enforcement will never be able to catch up.

What made the difference? Obviously, it’s ChatGPT, or put differently, the fact that now, there is a means for people to make active use of such a tool, employing it for whatever their personal needs or interests are¹. In fact, I’d argue it’s more than that: ChatGPT is not some impersonal tool – it talks to you, picking up your clarifications, changes of topic, mood… It is someone rather than something, or at least that’s how it seems. I’ll come back to that point in It’s us, really: Anthropomorphism unleashed. Before, let’s take a look at the underlying technology.

¹ Evidently, this is not about singling out ChatGPT as opposed to other chatbots; rather, I’m adopting it as the prototypical such application, since it is the one omnipresent in the media these days.

Large Language Models: What they are

How is it even possible to build a machine that talks to you? One way is to have that machine listen a lot. And listen is what these machines do; they do it a lot. But listening alone would never be enough to attain results as impressive as those we see. Instead, LLMs practice some form of “maximally active listening”: Continuously, they try to predict the speaker’s next utterance. By “continuously”, I mean word-by-word: At each training step, the model is asked to produce the subsequent word in a text.

Maybe in my last sentence, you noted the term “train”. As per common sense, “training” implies some form of supervision. It also implies some form of method. Since learning material is scraped from the internet, the true continuation is always known. The precondition for supervision is thus always fulfilled: A supervisor can just compare model prediction with what really follows in the text. Remains the question of method. That’s where we need to talk about deep learning, and we’ll do that in Model training.

Overall architecture

Today’s LLMs are, in some way or the other, based on an architecture known as the Transformer. This architecture was originally introduced in a paper catchily titled “Attention is all you need” (Vaswani et al. 2017). Of course, this was not the first attempt at automating natural-language generation – not even in deep learning, the sub-type of machine learning whose defining characteristic are many-layered (“deep”) artificial neural networks. But there, in deep learning, it constituted some kind of paradigm change. Before, models designed to solve sequence-prediction tasks (time-series forecasting, text generation…) tended to be based on some form of recurrent architecture, introduced in the 1990’s (eternities ago, on the time scale of deep-learning) by (Hochreiter and Schmidhuber 1997). Basically, the concept of recurrence, with its associated threading of a latent state, was replaced by “attention”. That’s what the paper’s title was meant to communicate: The authors did not introduce “attention”²; instead, they fundamentally expanded its usage so as to render recurrence superfluous.

Vaswani, Ashish, Noam Shazeer, Niki Parmar, et al. 2017. Attention Is All You Need. https://arxiv.org/abs/1706.03762.

Hochreiter, Sepp, and Jürgen Schmidhuber. 1997. “Long Short-Term Memory.” Neural Computation 9 (December): 1735–80. https://doi.org/10.1162/neco.1997.9.8.1735.

² I’m using quotes to refer to how attention is operationalized in deep learning, as opposed to how it is conceptualized in cognitive science or psychology.

³ If you’re wondering how that is possible – shouldn’t there be a separate, top-level module for generation? – no, there need not be. That’s because training implies prediction.

How did that ancestral Transformer look? – One prototypical task in natural language processing is machine translation. In translation, be it done by a machine or by a human, there is an input (in one language) and an output (in another). That input, call it a code. Whoever wants to establish its counterpart in the target language first needs to decode it. Indeed, one of two top-level building blocks of the archetypal Transformer was a decoder, or rather, a stack of decoders applied in succession. At its end, out popped a phrase in the target language³. What, then, was the other high-level block? It was an encoder, something that takes text (or tokens, rather, i.e., something that has undergone tokenization) and converts it into a form the decoder can make sense of. (Obviously, there is no analogue to this in human translation.)

From this two-stack architecture, subsequent developments tended to keep just one. The GPT family, together with many others, just kept the decoder stack. Now, doesn’t the decoder need some kind of input – if not to translate to a different language, then to reply to, as in the chatbot scenario? Turns out that no, it doesn’t – and that’s why you can also have the bot initiate the conversation. Unbeknownst to you, there will, in fact, be an input to the model – some kind of token signifying “end of input”. In that case, the model will draw on its training experience to generate a word likely to start out a phrase. That one word will then become the new input to continue from, and so forth. Summing up so far, then, GPT-like LLMs are Transformer Decoders.

The question is, how does such a stack of decoders succeed in fulfilling the task?

GPT-type models up close

In opening the black box, we focus on its two interfaces – input and output – as well as on the internals, its core.

Input

For simplicity, let me speak of words, not tokens. Now imagine a machine that is to work with – more even: “understand”⁴ – words. For a computer to process non-numeric data, a conversion to numbers necessarily has to happen. The straightforward way to effectuate this is to decide on a fixed lexicon, and assign each word a number. And this works: The way deep neural networks are trained, they don’t need semantic relationships to exist between entities in the training data to memorize formal structure. Does this mean they will appear perfect while training, but fail in real-world prediction? – If the training data are representative of how we converse, all will be fine. In a world of perfect surveillance, machines could exist that have internalized our every spoken word. Before that happens, though, the training data will be imperfect.

⁴ Why the quotes? See Large Language Models: What they are not.

⁵ As a fascinating example from dynamical systems theory, take delay coordinate embeddings.

⁶ Suitably named embedding layer.

⁷ See, for example, (Caliskan et al. 2022).

Caliskan, Aylin, Pimparkar Parth Ajay, Tessa Charlesworth, Robert Wolfe, and Mahzarin R. Banaji. 2022. “Gender Bias in Word Embeddings.” Proceedings of the 2022 AAAI/ACM Conference on AI, Ethics, and Society, July. https://doi.org/10.1145/3514094.3534162.

A much more promising approach than to simply index words, then, is to represent them in a richer, higher-dimensional space, an embedding space. This idea, popular not just in deep learning but in natural language processing overall, really goes far beyond anything domain-specific – linguistic entities, say⁵. You may be able to fruitfully employ it in virtually any domain – provided you can devise a method to sensibly map the given data into that space. In deep learning, these embeddings are obtained in a clever way: as a by-product of sorts of the overall training workflow. Technically, this is achieved by means of a dedicated neural-network layer⁶ tasked with evolving these mappings. Note how, smart though this strategy may be, it implies that the overall setting – everything from training data via model architecture to optimization algorithms employed – necessarily affects the resulting embeddings. And since these may be extracted and made use of in down-stream tasks, this matters⁷.

As to the GPT family, such an embedding layer constitutes part of its input interface – one “half”, so to say. Technically, the second makes use of the same type of layer, but with a different purpose. To contrast the two, let me spell out clearly what, in the part we’ve talked about already, is getting mapped to what. The mapping is between a word index – a sequence 1, 2, …, – on the one hand and a set of continuous-valued vectors of some length – 100, say – on the other. (One of them could like this: ) Thus, we obtain an embedding for every word. But language is more than an unordered assembly of words. Rearranging words, if syntactically allowed, may result in drastically changed semantics. In the pre-transformer paradigma, threading a sequentially-updated hidden state took care of this. Put differently, in that type of model, information about input order never got lost throughout the layers. Transformer-type architectures, however, need to find a different way. Here, a variety of rivaling methods exists. Some assume an underlying periodicity in semanto-syntactic structure. Others – and the GPT family, as yet and insofar we know, has been part of them⁸ – approach the challenge in exactly the same way as for the lexical units: They make learning these so-called position embeddings a by-product of model training. Implementation-wise, the only difference is that now the input to the mapping looks like this: 1, 2, …, where “maximum position” reflects choice of maximal sequence length supported.

⁸ For GPT-4, even high-level model information has not been released.

Summing up, verbal input is thus encoded – embedded, enriched – twofold as it enters the machine. The two types of embedding are combined and passed on to the model core, the already-mentioned decoder stack.

Core Processing

The decoder stack is made up of some number of identical blocks (12, in the case of GPT-2). (By “identical” I mean that the architecture is the same; the weights – the place where a neural-network layer stores what it “knows” – are not. More on these “weights” soon.)

Inside each block, some sub-layers are pretty much “business as usual”. One is not: the attention module, the “magic” ingredient that enabled Transformer-based architectures to forego keeping a latent state. To explain how this works, let’s take translation as an example.

In the classical encoder-decoder setup, the one most intuitive for machine translation, imagine the very first decoder in the stack of decoders. It receives as input a length-seven cypher, the encoded version of an original length-seven phrase. Since, due to how the encoder blocks are built, input order is conserved, we have a faithful representation of source-language word order. In the target language, however, word order can be very different. A decoder module, in producing the translation, had rather not do this by translating each word as it appears. Instead, it would be desirable for it to know which among the already-seen tokens is most relevant right now, to generate the very next output token. Put differently, it had better know where to direct its attention.

Thus, figure out how to distribute focus is what attention modules do. How do they do it? They compute, for each available input-language token, how good a match, a fit, it is for their own current input. Remember that every token, at every processing stage, is encoded as a vector of continuous values. How good a match any of, say, three source-language vectors is is then computed by projecting one’s current input vector onto each of the three. The closer the vectors, the longer the projected vector. ⁹ Based on the projection onto each source-input token, that token is weighted, and the attention module passes on the aggregated assessments to the ensuing neural-network module.

⁹ Mathematically, this is achieved by a pretty standard and pervasively-used, in machine learning, operation, the dot product.

To explain what attention modules are for, I’ve made use of the machine-translation scenario, a scenario that should lend a certain intuitiveness to the operation. But for GPT-family models, we need to abstract this a bit. First, there is no encoder stack, so “attention” is computed among decoder-resident tokens only. And second – remember I said a stack was built up of identical modules? – this happens in every decoder block. That is, when intermediate results are bubbled up the stack, at each stage the input is weighted as appropriate at that stage. While this is harder to intuit than what happened in the translation scenario, I’d argue that in the abstract, it makes a lot of sense. For an analogy, consider some form of hierarchical categorization of entities. As higher-level categories are built from lower-level ones, at each stage the process needs to look at its input afresh, and decide on a sensible way of subsuming similar-in-some-way categories.

Output

Stack of decoders traversed, the multi-dimensional codes that pop out need to be converted into something that can be compared with the actual phrase continuation we see in the training corpus. Technically, this involves a projection operation as well a strategy for picking the output word – that word in target-language vocabulary that has the highest probability. How do you decide on a strategy? I’ll say more about that in the section Mechanics of text generation, where I assume a chatbot user’s perspective.

Model training

Before we get there, just a quick word about model training. LLMs are deep neural networks, and as such, they are trained like any network is. First, assuming you have access to the so-called “ground truth”, you can always compare model prediction with the true target. You then quantify the difference – by which algorithm will affect training results. Then, you communicate that difference – the loss – to the network. It, in turn, goes through its modules, from back/top to start/bottom, and updates its stored “knowledge” – matrices of continuous numbers called weights. Since information is passed from layer to layer, in a direction reverse to that followed in computing predictions, this technique is known as back-propagation.

And all that is not triggered once, but iteratively, for a certain number of so-called “epochs”, and modulated by a set of so-called “hyper-parameters”. In practice, a lot of experimentation goes into deciding on the best-working configuration of these settings.

Mechanics of text generation

We already know that during model training, predictions are generated word-by-word; at every step, the model’s knowledge about what has been said so far is augmented by one token: the word that really was following at that point. If, making use of a trained model, a bot is asked to reply to a question, its response must by necessity be generated in the same way. However, the actual “correct word” is not known. The only way, then, is to feed back to the model its own most recent prediction. (By necessity, this lends to text generation a very special character, where every decision the bot makes co-determines its future behavior.)

Why, though, talk about decisions? Doesn’t the bot just act on behalf of the core model, the LLM – thus passing on the final output? Not quite. At each prediction step, the model yields a vector, with values as many as there are entries in the vocabulary. As per model design and training rationale, these vectors are “scores” – ratings, sort of, how good a fit a word would be in this situation. Like in life, higher is better. But that doesn’t mean you’d just pick the word with the highest value. In any case, these scores are converted to probabilities, and a suitable probability distribution is used to non-deterministically pick a likely (or likely-ish) word. The probability distribution commonly used is the multinomial distribution, appropriate for discrete choice among more than two alternatives. But what about the conversion to probabilities? Here, there is room for experimentation.

Technically, the algorithm employed is known as the softmax function. It is a simplified version of the Boltzmann distribution, famous in statistical mechanics, used to obtain the probability of a system’s state given that state’s energy and the temperature of the system. But for temperature¹⁰, both formulae are, in fact, identical. In physical systems, temperature modulates probabilities in the following way: The hotter the system, the closer the states’ probabilities are to each other; the colder it gets, the more distinct those probabilities. In the extreme, at very low temperatures there will be a few clear “winners” and a silent majority of “losers”.

¹⁰ … and the Boltzmann constant – but that being a constant, we don’t consider it here.

In deep learning, a like effect is easy to achieve (by means of a scaling factor). That’s why you may have heard people talk about some weird thing called “temperature” that resulted in [insert adjective here] answers. If the application you use lets you vary that factor, you’ll see that a low temperature will result in deterministic-looking, repetitive, “boring” continuations, while a high one may make the machine appear as though it were on drugs.

That concludes our high-level overview of LLMs. Having seen the machine dissected in this way may already have left you with some sort of opinion of what these models are – not. This topic more than deserves a dedicated exposition – and papers are being written pointing to important aspects all the time – but in this text, I’d like to at least offer some input for thought.

Large Language Models: What they are not

In part one,describing LLMs technically, I’ve sometimes felt tempted to use terms like “understanding” or “knowledge” when applied to the machine. I may have ended up using them; in that case, I’ve tried to remember to always surround them with quotes. The latter, the adding quotes, stands in contrast to many texts, even ones published in an academic context (Bender and Koller 2020). The question is, though: Why did I even feel compelled to use these terms, given I do not think they apply, in their usual meaning? I can think of a simple – shockingly simple, maybe – answer: It’s because us, humans, we think, talk, share our thoughts in these terms. When I say understand, I surmise you will know what I mean.

Now, why do I think that these machines do not understand human language, in the sense we usually imply when using that word?

A few facts

I’ll start out briefly mentioning empirical results, conclusive thought experiments, and theoretical considerations. All aspects touched upon (and many more) are more than worthy of in-depth discussion, but such discussion is clearly out of scope for this synoptic-in-character text.

First, while it is hard to put a number on the quality of a chatbot’s answers, performance on standardized benchmarks is the “bread and butter” of machine learning – its reporting being an essential part of the prototypical deep-learning publication. (You could even call it the “cookie”, the driving incentive, since models usually are explicitly trained and fine-tuned for good results on these benchmarks.) And such benchmarks exist for most of the down-stream tasks the LLMs are used for: machine translation, generating summaries, text classification, and even rather ambitious-sounding setups associated with – quote/unquote – reasoning.

How do you assess such a capability? Here is an example from a benchmark named “Argument Reasoning Comprehension Task” (Habernal et al. 2018).

Habernal, Ivan, Henning Wachsmuth, Iryna Gurevych, and Benno Stein. 2018. “The Argument Reasoning Comprehension Task: Identification and Reconstruction of Implicit Warrants.” Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers) (New Orleans, Louisiana), June, 1930–40. https://doi.org/10.18653/v1/N18-1175.

Claim: Google is not a harmful monopoly
Reason: People can choose not to use Google
Warrant: Other search engines don’t redirect to Google
Alternative: All other search engines redirect to Google

Here claim and reason together make up the argument. But what, exactly, is it that links them? At first look, this can even be confusing to a human. The missing link is what is called warrant here – add it in, and it all starts to make sense. The task, then, is to decide which of warrant or alternative supports the conclusion, and which one does not.

If you think about it, this is a surprisingly challenging task. Specifically, it seems to inescapingly require world knowledge. So if language models, as has been claimed, perform nearly as well as humans, it seems they must have such knowledge – no quotes added. However, in response to such claims, research has been performed to uncover the hidden mechanism that enables such seemingly-superior results. For that benchmark, it has been found (Niven and Kao 2019) that there were spurious statistical cues in the way the dataset was constructed – those removed, LLM performance was no better than random.

Niven, Timothy, and Hung-Yu Kao. 2019. “Probing Neural Network Comprehension of Natural Language Arguments.” CoRR abs/1907.07355. http://arxiv.org/abs/1907.07355.

¹¹ That choice of species is probably not a coincidence: see https://en.wikipedia.org/wiki/Cephalopod_intelligence.

World knowledge, in fact, is one of the main things an LLM lacks. Bender et al. (Bender and Koller 2020) convincingly demonstrate its essentiality by means of two thought experiments. One of them, situated on a lone island, imagines an octopus¹¹ inserting itself into some cable-mediated human communication, learning the chit-chat, and finally – having gotten bored – impersonating one of the humans. This works fine, until one day, its communication partner finds themselves in an emergency, and needs to build some rescue tool out of things given in the environment. They urgently ask for advice – and the octopus has no idea what to respond. It has no ideas what these words actually refer to.

The other argument comes directly from machine learning, and strikingly simple though it may be, it makes its point very well. Imagine an LLM trained as usual, including on lots of text involving plants. It has also been trained on a dataset of unlabeled photos, the actual task being unsubstantial – say it had to fill out masked areas. Now, we pull out a picture and ask: How many of that blackberry’s blossoms have already opened? The model has no chance to answer the question.

Now, please look back at the Joseph Weizenbaum quote I opened this article with. It is still true that language-generating machine have no knowledge of the world we live in.

Before moving on, I’d like to just quickly hint at a totally different type of consideration, brought up in a (2003!) paper by Spärck Jones (Spaerck 2004). Though written long before LLMs, and long before deep learning started its winning conquest, on an abstract level it is still very applicable to today’s situation. Today, LLMs are employed to “learn language”, i.e., for language acquisition. That skill is then built upon by specialized models, of task-dependent architecture. Popular real-world¹² down-stream tasks are translation, document retrieval, or text summarization. When the paper was written, there was no such two-stage pipeline. The author was questioning the fit between how language modeling was conceptualized – namely, as a form of recovery – and the character of these down-stream tasks. Was recovery – inferring a missing, for whatever reasons – piece of text a good model, of, say, condensing a long, detailed piece of text into a short, concise, factual one? If not, could the reason it still seemed to work just fine be of a very different nature – a technical, operational, coincidental one?

¹² As opposed to the aforementioned problems subsumed under “reasoning”, those having been constructed for research purposes.

[…] the crucial characterisation of the relationship between the input and the output is in fact offloaded in the LM approach onto the choice of training data. We can use LM for summarising because we know that some set of training data consists of full texts paired with their summaries.¹³

¹³ From (Spaerck 2004).

Spaerck, Karen. 2004. “Language Modelling’s Generative Model : Is It Rational?”

It seems to me that today’s two-stage process notwithstanding, this is still an aspect worth giving some thought.

It’s us: Language learning, shared goals, and a shared world

We’ve already talked about world knowledge. What else are LLMs missing out on?

In our world, you’ll hardly find anything that does not involve other people. This goes a lot deeper than the easily observable facts: our constantly communicating, reading and typing messages, documenting our lives on social networks… We don’t experience, explore, explain a world of our own. Instead, all these activities are inter-subjectively constructed. Feelings are¹⁴. Cognition is; meaning is. And it goes deeper yet. Implicit assumptions guide us to constantly look for meaning, be it in overheard fragments, mysterious symbols, or life events.

¹⁴ See https://lisafeldmanbarrett.com/books/how-emotions-are-made/.

How does this relate to LLMs? For one, they’re islands of their own. When you ask them for advice – to develop a research hypothesis and a matching operationalization, say, or whether a detainee should be released on parole – they have no stakes in the outcome, no motivation (be it intrinsic or extrinsic), no goals. If an innocent person is harmed, they don’t feel the remorse; if an experiment is successful but lacks explanatory power, they don’t sense the shallowness; if the world blows up, it won’t have been their world.

Secondly, it’s us who are not islands. In Bender et al.’s octopus scenario, the human on one side of the cable plays an active role not just when they speak. In making sense of what the octopus says, they contribute an essential ingredient: namely, what they think the octopus wants, thinks, feels, expects… Anticipating, they reflect on what the octopus anticipates.

As Bender et al. put it:

It is not that O’s utterances make sense, but rather, that A can make sense of them.

That article (Bender and Koller 2020) also brings impressive evidence from human language acquisition: Our predisposition towards language learning notwithstanding, infants don’t learn from the availability of input alone. A situation of joint attention is needed for them to learn. Psychologizing, one could hypothesize they need to get the impression that these sounds, these words, and the fact they’re linked together, actually matters.

Bender, Emily M., and Alexander Koller. 2020. “Climbing Towards NLU: On Meaning, Form, and Understanding in the Age of Data.” Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics (Online), July, 5185–98. https://doi.org/10.18653/v1/2020.acl-main.463.

Let me conclude, then, with my final “psychologization”.

It’s us, really: Anthropomorphism unleashed

Yes, it is amazing what these machines do. (And that makes them incredibly dangerous power instruments.) But this in no way affects the human-machine differences that have been existing throughout history, and continue to exist today. That we are inclined to think they understand, know, mean – that maybe even they’re conscious: that’s on us. We can experience deep emotions watching a movie; hope that if we just try enough, we can sense what a distant-in-evolutionary-genealogy creature is feeling; see a cloud encouragingly smiling at us; read a sign in an arrangement of pebbles.

Our inclination to anthropomorphize is a gift; but it can sometimes be harmful. And nothing of this is special to the twenty-first century.

Like I began with him, let me conclude with Weizenbaum.

Some subjects have been very hard to convince that ELIZA (with its present script) is not human.

Photo by Marjan Blan on Unsplash

Group-equivariant neural networks with escnn

Sigrid Keydana — Tue, 09 May 2023 00:00:00 GMT

Today, we resume our exploration of group equivariance. This is the third post in the series. The first was a high-level introduction: what this is all about; how equivariance is operationalized; and why it is of relevance to many deep-learning applications. The second sought to concretize the key ideas by developing a group-equivariant CNN from scratch. That being instructive, but too tedious for practical use, today we look at a carefully designed, highly-performant library that hides the technicalities and enables a convenient workflow.

First though, let me again set the context. In physics, an all-important concept is that of symmetry¹, a symmetry being present whenever some quantity is being conserved. But we don’t even need to look to science. Examples arise in daily life, and – otherwise why write about it - in the tasks we apply deep learning to.

¹ This is nicely explained in, for example, Jakob Schwichtenberg’s Physics from symmetry.

In daily life: Think about speech – me stating “it is cold”, for example. Formally, or denotation-wise, the sentence will have the same meaning now as in five hours. (Connotations, on the other hand, can and will probably be different!). This is a form of translation symmetry, translation in time.

In deep learning: Take image classification. For the usual convolutional neural network, a cat in the center of the image is just that, a cat; a cat on the bottom is, too. But one sleeping, comfortably curled like a half-moon “open to the right”, will not be “the same” as one in a mirrored position. Of course, we can train the network to treat both as equivalent by providing training images of cats in both positions, but that is not a scaleable approach. Instead, we’d like to make the network aware of these symmetries, so they are automatically preserved throughout the network architecture.

Purpose and scope of this post

Here, I introduce escnn, a PyTorch extension that implements forms of group equivariance for CNNs operating on the plane or in (3d) space. The library is used in various, amply illustrated research papers; it is appropriately documented; and it comes with introductory notebooks both relating the math and exercising the code. Why, then, not just refer to the first notebook, and immediately start using it for some experiment?

In fact, this post should – as quite a few texts I’ve written – be regarded as an introduction to an introduction. To me, this topic seems anything but easy, for various reasons. Of course, there’s the math. But as so often in machine learning, you don’t need to go to great depths to be able to apply an algorithm correctly. So if not the math itself, what generates the difficulty? For me, it’s two things.

First, to map my understanding of the mathematical concepts to the terminology used in the library, and from there, to correct use and application. Expressed schematically: We have a concept A, which figures (among other concepts) in technical term (or object class) B. What does my understanding of A tell me about how object class B is to be used correctly? More importantly: How do I use it to best attain my goal C? This first difficulty I’ll address in a very pragmatic way. I’ll neither dwell on mathematical details, nor try to establish the links between A, B, and C in detail. Instead, I’ll present the characters² in this story by asking what they’re good for.

² If you have some background on representations: No, not those characters …

³ Yes, that word exists, although I must admit I didn’t know before typing it into a search engine.

⁴ One paper particularly stood out to me: Coordinate Independent Convolutional Networks – Isometry and Gauge Equivariant Convolutions on Riemannian Manifolds.

⁵ Directly pertinent to today’s topic, thinking of the materials produced for University of Amsterdam’s course on group-equivariant deep learning.

Second – and this will be of relevance to just a subset of readers – the topic of group equivariance, particularly as applied to image processing, is one where visualizations can be of tremendous help. The quaternity³ of conceptual explanation, math, code, and visualization can, together, produce an understanding of emergent-seeming quality… if, and only if, all of these explanation modes “work” for you. (Or if, in an area, a mode that does not wouldn’t contribute that much anyway.) Here, it so happens that from what I saw, several papers have excellent visualizations⁴, and the same holds for some lecture slides and accompanying notebooks⁵. But for those among us with limited spatial-imagination capabilities – e.g., people with Aphantasia – these illustrations, intended to help, can be very hard to make sense of themselves. If you’re not one of these, I totally recommend checking out the resources linked in the above footnotes. This text, though, will try to make the best possible use of verbal explanation to introduce the concepts involved, the library, and how to use it.

That said, let’s start with the software.

Using escnn

Escnn depends on PyTorch. Yes, PyTorch, not torch; unfortunately, the library hasn’t been ported to R yet.⁶ For now, thus, we’ll employ reticulate⁷ to access the Python objects directly.

⁶ If, after reading this post, you feel that maybe you would be interested in porting it – I can definitely say I think that’s a great idea!

⁷ If you’re new to reticulate, please consult its excellent documentation on topics like calling Python from R, determining the Python version used, installing Python packages, as well as an outstanding introduction to Python for R users.

The way I’m doing this is install escnn in a virtual environment, with PyTorch version 1.13.1. As of this writing, Python 3.11 is not yet supported by one of escnn’s dependencies; the virtual environment thus builds on Python 3.10. As to the library itself, I am using the development version from GitHub, running pip install git+https://github.com/QUVA-Lab/escnn.

Once you’re ready, issue

library(reticulate)
# Verify correct environment is used.
# Different ways exist to ensure this; I've found most convenient to configure this on
# a per-project basis in RStudio's project file (.Rproj)
py_config()

# bind to required libraries and get handles to their namespaces
torch <- import("torch")
escnn <- import("escnn")

Escnn loaded, let me introduce its main objects and their roles in the play.

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

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

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

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.

Where to from here

This “introduction to an introduction” has been the attempt to draw a high-level map of the terrain, so you can decide if this is useful to you. If it’s not just useful, but interesting theory-wise as well, you’ll find lots of excellent materials linked from the README. The way I see it, though, this post already should enable you to actually experiment with different setups.

One such experiment, that would be of high interest to me, might investigate how well different types and degrees of equivariance actually work for a given task and dataset. Overall, a reasonable assumption is that, the higher “up” we go in the feature hierarchy, the less equivariance we require. For edges and corners, taken by themselves, full rotation equivariance seems desirable, as does equivariance to reflection; for higher-level features, we might want to successively restrict allowed operations, maybe ending up with equivariance to mirroring merely. Experiments could be designed to compare different ways, and levels, of restriction.⁸

⁸ See the documentation for escnn$nn$RestrictionModule for how to do this.

Thanks for reading!

Photo by Volodymyr Tokar on Unsplash

Deep Learning and Scientific Computing with R torch: the book

Sigrid Keydana — Wed, 05 Apr 2023 00:00:00 GMT

First things first: Where can you get it? As of today, you can download the e-book or order a print copy from the publisher, CRC Press; the free online edition is here. There is, to my knowledge, no problem to perusing the online version – besides one: It doesn’t have the squirrel that’s on the book cover.

So if you’re a lover of amazing creatures…

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:

Who’s it for?

In short, Deep Learning and Scientific Computing with R torch – being the only comprehensive text, as of this writing, on this topic – addresses a wide audience. The hope is that there’s something in it for everyone (well, most everyone).

If you’ve never used torch, nor any other deep-learning framework, starting right from the beginning is the thing to do. No prior knowledge of deep learning is expected. The assumption is that you know some basic R, and are familiar with machine-learning terms such as supervised vs. unsupervised learning, training-validation-test set, et cetera. Having worked through part one, you’ll find that parts two and three – independently – continue right from where you left off.

If, on the other hand, you do have basic experience with torch and/or other automatic-differentiation frameworks, and are mostly interested in applied deep learning, you may be inclined to skim part one, and go to part two, checking out the applications that interest you most (or just browse, looking for inspiration). The domain-dependent examples were chosen to be rather generic and straightforward, so as to have the code generalize to a whole range of similar applications.

Finally, if it was the “scientific computing” in the title that caught your attention, I certainly hope that part three has something for you! (As the book’s author, I may say that writing this part was an extremely satisfying, incredibly engaging experience.¹) Part three really is where it makes sense to talk of “browsing” – its topics hardly depend on each other, just look around for what appeals to you.

¹ though challenging, as well (though … or because?)

To wrap up, then:

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!

Implementing rotation equivariance: Group-equivariant CNN from scratch

Sigrid Keydana — Mon, 27 Mar 2023 00:00:00 GMT

Convolutional neural networks (CNNs) are great – they’re able to detect features in an image no matter where. Well, not exactly. They’re not indifferent to just any kind of movement. Shifting up or down, or left or right, is fine; rotating around an axis is not. That’s because of how convolution works: traverse by row, then traverse by column (or the other way round). If we want “more” (e.g., successful detection of an upside-down object), we need to extend convolution to an operation that is rotation-equivariant. An operation that is equivariant to some type of action will not only register the moved feature per se, but also, keep track of which concrete action made it appear where it is.

This is the second post in a series that introduces group-equivariant CNNs (GCNNs). The first was a high-level introduction to why we’d want them, and how they work. There, we introduced the key player, the symmetry group, which specifies what kinds of transformations are to be treated equivariantly. If you haven’t, please take a look at that post first, since here I’ll make use of terminology and concepts it introduced.

Today, we code a simple GCNN from scratch. Code and presentation tightly follow a notebook provided as part of University of Amsterdam’s 2022 Deep Learning Course. They can’t be thanked enough for making available such excellent learning materials.

In what follows, my intent is to explain the general thinking, and how the resulting architecture is built up from smaller modules, each of which is assigned a clear purpose. For that reason, I won’t reproduce all the code here; instead, I’ll make use of the package gcnn. Its methods are heavily annotated; so to see some details, don’t hesitate to look at the code.

As of today, gcnn implements one symmetry group: , the one that serves as a running example throughout post one. It is straightforwardly extensible, though, making use of class hierarchies throughout.

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!

Photo by Marjan Blan | @marjanblan on Unsplash

Upside down, a cat’s still a cat: Evolving image recognition with Geometric Deep Learning

Sigrid Keydana — Thu, 09 Mar 2023 00:00:00 GMT

This is the first in a series of posts on group-equivariant convolutional neural networks (GCNNs). Today, we keep it short, high-level, and conceptual; examples and implementations will follow. In looking at GCNNs, we are resuming a topic we first wrote about in 2021: Geometric Deep Learning, a principled, math-driven approach to network design that, since then, has only risen in scope and impact.

From alchemy to science: Geometric Deep Learning in two minutes

In a nutshell, Geometric Deep Learning is all about deriving network structure from two things: the domain, and the task. The posts will go into a lot of detail, but let me give a quick preview here:

By domain, I’m referring to the underlying physical space, and the way it is represented in the input data. For example, images are usually coded as a two-dimensional grid, with values indicating pixel intensities.¹
The task is what we’re training the network to do: classification, say, or segmentation. Tasks may be different at different stages in the architecture. At each stage, the task in question will have its word to say about how layer design should look.

¹ That’s when the image has a single channel. Otherwise, there are a number of grids, each mapped to a separate channel.

For instance, take MNIST. The dataset consists of images of ten digits, 0 to 10, all gray-scale. The task – unsurprisingly – is to assign each image the digit represented.

First, consider the domain. A is a wherever it appears on the grid. We thus need an operation that is translation-equivariant: It flexibly adapts to shifts (translations) in its input. More concretely, in our context, equivariant operations are able to detect some object’s properties even if that object has been moved, vertically and/or horizontally, to another location. Convolution, ubiquitous not just in deep learning, is just such a shift-equivariant operation.

Let me call special attention to the fact that, in equivariance, the essential thing is that “flexible adaptation.” Translation-equivariant operations do care about an object’s new position; they record a feature not abstractly, but at the object’s new position. To see why this is important, consider the network as a whole. When we compose convolutions, we build a hierarchy of feature detectors. That hierarchy should be functional no matter where in the image. In addition, it has to be consistent: Location information needs to be preserved between layers.

Terminology-wise, thus, it is important to distinguish equivariance from invariance. An invariant operation, in our context, would still be able to spot a feature wherever it occurs; however, it would happily forget where that feature happened to be. Clearly, then, to build up a hierarchy of features, translation-invariance is not enough.

What we’ve done right now is derive a requirement from the domain, the input grid. What about the task? If, finally, all we’re supposed to do is name the digit, now suddenly location does not matter anymore. In other words, once the hierarchy exists, invariance is enough. In neural networks, pooling is an operation that forgets about (spatial) detail. It only cares about the mean, say, or the maximum value itself. This is what makes it suited to “summing up” information about a region, or a complete image, if at the end we only care about returning a class label.

In a nutshell, we were able to formulate a design wishlist based on (1) what we’re given and (2) what we’re tasked with.

After this high-level sketch of Geometric Deep Learning, we zoom in on this series of posts’ designated topic: group-equivariant convolutional neural networks.

The why of “equivariant” should not, by now, pose too much of a riddle. What about that “group” prefix, though?

The “group” in group-equivariance

As you may have guessed from the introduction, talking of “principled” and “math-driven”, this really is about groups in the “math sense.” Depending on your background, the last time you heard about groups was in school, and with not even a hint at why they matter. I’m certainly not qualified to summarize the whole richness of what they’re good for, but I hope that by the end of this post, their importance in deep learning will make intuitive sense.

Groups from symmetries

Here is a square.

Now close your eyes.

Now look again. Did something happen to the square?

You can’t tell. Maybe it was rotated; maybe it was not. On the other hand, what if the vertices were numbered?

Now you’d know.

Without the numbering, could I have rotated the square in any way I wanted? Evidently not. This would not go through unnoticed:

There are exactly four ways I could have rotated the square without raising suspicion. Those ways can be referred to in different ways; one simple way is by degree of rotation: 90, 180, or 270 degrees. Why not more? Any further addition of 90 degrees would result in a configuration we’ve already seen.

The above picture shows three squares, but I’ve listed three possible rotations. What about the situation on the left, the one I’ve taken as an initial state? It could be reached by rotating 360 degrees (or twice that, or thrice, or …) But the way this is handled, in math, is by treating it as some sort of “null rotation”, analogously to how acts in addition, in multiplication, or the identity matrix in linear algebra.

Altogether, we thus have four actions that could be performed on the square (an un-numbered square!) that would leave it as-is, or invariant. These are called the symmetries of the square. A symmetry, in math/physics, is a quantity that remains the same no matter what happens as time evolves. And this is where groups come in. Groups – concretely, their elements – effectuate actions like rotation.

Before I spell out how, let me give another example. Take this sphere.

How many symmetries does a sphere have? Infinitely many. This implies that whatever group is chosen to act on the square, it won’t be much good to represent the symmetries of the sphere.

Viewing groups through the action lens

Following these examples, let me generalize. Here is typical definition. ²

² From Wolfram Alpha (beginning of article).

A group is a finite or infinite set of elements together with a binary operation (called the group operation) that together satisfy the four fundamental properties of closure, associativity, the identity property, and the inverse property. The operation with respect to which a group is defined is often called the “group operation,” and a set is said to be a group “under” this operation. Elements , , , … with binary operation between and denoted form a group if

Closure: If and are two elements in , then the product is also in .

Associativity: The defined multiplication is associative, i.e., for all ,, in , .

Identity: There is an identity element (a.k.a. , , or ) such that for every element in .

Inverse: There must be an inverse (a.k.a. reciprocal) of each element. Therefore, for each element of , the set contains an element such that .

In action-speak, group elements specify allowable actions; or more precisely, ones that are distinguishable from each other. Two actions can be composed; that’s the “binary operation”. The requirements now make intuitive sense:

A combination of two actions – two rotations, say – is still an action of the same type (a rotation).
If we have three such actions, it doesn’t matter how we group them. (Their order of application has to remain the same, though.)
One possible action is always the “null action”. (Just like in life.) As to “doing nothing”, it doesn’t make a difference if that happens before or after a “something”; that “something” is always the final result.
Every action needs to have an “undo button”. In the squares example, if I rotate by 180 degrees, and then, by 180 degrees again, I am back in the original state. It is if I had done nothing.

Resuming a more “birds-eye view”, what we’ve seen right now is the definition of a group by how its elements act on each other. But if groups are to matter “in the real world”, they need to act on something outside (neural network components, for example). How this works is the topic of the following posts, but I’ll briefly outline the intuition here.

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!

Photo by Ihor OINUA on Unsplash

AO, NAO, ENSO: A wavelet analysis example

Sigrid Keydana — Thu, 19 Jan 2023 00:00:00 GMT

Recently, we showed how to use torch for wavelet analysis. A member of the family of spectral analysis methods, wavelet analysis bears some similarity to the Fourier Transform, and specifically, to its popular two-dimensional application, the spectrogram.

As explained in that book excerpt, though, there are significant differences. For the purposes of the current post, it suffices to know that frequency-domain patterns are discovered by having a little “wave” (that, really, can be of any shape) “slide” over the data, computing degree of match (or mismatch) in the neighborhood of every sample.

With this post, then, my goal is two-fold.

First, to introduce torchwavelets, a tiny, yet useful package that automates all of the essential steps involved. Compared to the Fourier Transform and its applications, the topic of wavelets is rather “chaotic” – meaning, it enjoys much less shared terminology, and much less shared practice. Consequently, it makes sense for implementations to follow established, community-embraced approaches, whenever such are available and well documented. With torchwavelets, we provide an implementation of Torrence and Compo’s 1998 “Practical Guide to Wavelet Analysis” (Torrence and Compo (1998)), an oft-cited paper that proved influential across a wide range of application domains. Code-wise, our package is mostly a port of Tom Runia’s PyTorch implementation, itself based on a prior implementation by Aaron O’Leary.

Torrence, C., and G. P. Compo. 1998. “A Practical Guide to Wavelet Analysis.” Bulletin of the American Meteorological Society 79 (1): 61–78.

Second, to show an attractive use case of wavelet analysis in an area of great scientific interest and tremendous social importance (meteorology/climatology). Being by no means an expert myself, I’d hope this could be inspiring to people working in these fields, as well as to scientists and analysts in other areas where temporal data arise.

Concretely, what we’ll do is take three different atmospheric phenomena – El Niño–Southern Oscillation (ENSO), North Atlantic Oscillation (NAO), and Arctic Oscillation (AO) – and investigate them using wavelet analysis. In each case, we also look at the overall frequency spectrum, given by the Discrete Fourier Transform (DFT), as well as a classic time-series decomposition into trend, seasonal components, and remainder.

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.

Analysis: ENSO

We begin with the best-known of these phenomena: ENSO. Data are available from 1854 onwards; however, for comparability with AO, we discard all records prior to January, 1950. For analysis, we pick NINO34_MEAN, the monthly average sea surface temperature in the Niño 3.4 region (i.e., the area between 5° South, 5° North, 190° East, and 240° East). Finally, we convert to a tsibble, the format expected by feasts::STL().

library(tidyverse)
library(tsibble)

download.file(
  "https://bmcnoldy.rsmas.miami.edu/tropics/oni/ONI_NINO34_1854-2022.txt",
  destfile = "ONI_NINO34_1854-2022.txt"
)

enso <- read_table("ONI_NINO34_1854-2022.txt", skip = 9) %>%
  mutate(x = yearmonth(as.Date(paste0(YEAR, "-", `MON/MMM`, "-01")))) %>%
  select(x, enso = NINO34_MEAN) %>%
  filter(x >= yearmonth("1950-01"), x <= yearmonth("2022-09")) %>%
  as_tsibble(index = x)

enso

# A tsibble: 873 x 2 [1M]
          x  enso
       
 1 1950 Jan  24.6
 2 1950 Feb  25.1
 3 1950 Mar  25.9
 4 1950 Apr  26.3
 5 1950 May  26.2
 6 1950 Jun  26.5
 7 1950 Jul  26.3
 8 1950 Aug  25.9
 9 1950 Sep  25.7
10 1950 Oct  25.7
# … with 863 more rows

As already announced, we want to look at seasonal decomposition, as well. In terms of seasonal periodicity, what do we expect? Unless told otherwise, feasts::STL() will happily pick a window size for us. However, there’ll likely be several important frequencies in the data. (Not wanting to ruin the suspense, but for AO and NAO, this will definitely be the case!). Besides, we want to compute the Fourier Transform anyway, so why not do that first?

Here is the power spectrum:

library(torch)
fft <- torch_fft_fft(as.numeric(scale(enso$enso)))

In the below plot, the x axis corresponds to frequencies, expressed as “number of times per year”. We only display frequencies up to and including the Nyquist frequency, i.e., half the sampling rate, which in our case is 12 (per year).

num_samples <- nrow(enso)
nyquist_cutoff <- ceiling(num_samples / 2) # highest discernible frequency
bins_below_nyquist <- 0:nyquist_cutoff

sampling_rate <- 12 # per year
frequencies_per_bin <- sampling_rate / num_samples
frequencies <- frequencies_per_bin * bins_below_nyquist

df <- data.frame(f = frequencies, y = as.numeric(fft[1:(nyquist_cutoff + 1)]$abs()))
df %>% ggplot(aes(f, y)) +
  geom_line() +
  xlab("frequency (per year)") +
  ylab("magnitude") +
  ggtitle("Spectrum of Niño 3.4 data")

Frequency spectrum of monthly average sea surface temperature in the Niño 3.4 region, 1950 to present.

There is one dominant frequency, corresponding to about once a year. From this component alone, we’d expect one El Niño event – or equivalently, one La Niña – per year. But let’s locate important frequencies more precisely. With not many other periodicities standing out, we may as well restrict ourselves to three:

strongest <- torch_topk(fft[1:(nyquist_cutoff/2)]$abs(), 3)
strongest

[[1]]
torch_tensor
233.9855
172.2784
142.3784
[ CPUFloatType{3} ]

[[2]]
torch_tensor
74
21
7
[ CPULongType{3} ]

What we have here are the magnitudes of the dominant components, as well as their respective bins in the spectrum. Let’s see which actual frequencies these correspond to:

important_freqs <- frequencies[as.numeric(strongest[[2]])]
important_freqs

[1] 1.00343643 0.27491409 0.08247423

That’s once per year, once per quarter, and once every twelve years, approximately. Or, expressed as periodicity, in terms of months (i.e., how many months are there in a period):

num_observations_in_season <- 12/important_freqs  
num_observations_in_season

[1] 11.95890  43.65000 145.50000

We now pass these to feasts::STL(), to obtain a five-fold decomposition into trend, seasonal components, and remainder.

library(feasts)
enso %>%
  model(STL(enso ~ season(period = 12) + season(period = 44) +
              season(period = 145))) %>%
  components() %>%
  autoplot()

Decomposition of ENSO data into trend, seasonal components, and remainder by feasts::STL().

According to Loess decomposition, there still is significant noise in the data – the remainder remaining high despite our hinting at important seasonalities. In fact, there is no big surprise in that: Looking back at the DFT output, not only are there many, close to one another, low- and lowish-frequency components, but in addition, high-frequency components just won’t cease to contribute. And really, as of today, ENSO forecasting – tremendously important in terms of human impact – is focused on predicting oscillation state just a year in advance. This will be interesting to keep in mind for when we proceed to the other series – as you’ll see, it’ll only get worse.

By now, we’re well informed about how dominant temporal rhythms determine, or fail to determine, what actually happens in atmosphere and ocean. But we don’t know anything about whether, and how, those rhythms may have varied in strength over the time span considered. This is where wavelet analysis comes in.

In torchwavelets, the central operation is a call to wavelet_transform(), to instantiate an object that takes care of all required operations. One argument is required: signal_length, the number of data points in the series. And one of the defaults we need to override: dt, the time between samples, expressed in the unit we’re working with. In our case, that’s year, and, having monthly samples, we need to pass a value of 1/12. With all other defaults untouched, analysis will be done using the Morlet wavelet (available alternatives are Mexican Hat and Paul), and the transform will be computed in the Fourier domain (the fastest way, unless you have a GPU).

library(torchwavelets)
enso_idx <- enso$enso %>% as.numeric() %>% torch_tensor()
dt <- 1/12
wtf <- wavelet_transform(length(enso_idx), dt = dt)

A call to power() will then compute the wavelet transform:

power_spectrum <- wtf$power(enso_idx)
power_spectrum$shape

[1]  71 873

The result is two-dimensional. The second dimension holds measurement times, i.e., the months between January, 1950 and September, 2022. The first dimension warrants some more explanation.

Namely, we have here the set of scales the transform has been computed for. If you’re familiar with the Fourier Transform and its analogue, the spectrogram, you’ll probably think in terms of time versus frequency. With wavelets, there is an additional parameter, the scale, that determines the spread of the analysis pattern.

The aforementioned book excerpt discusses this in detail.

Some wavelets have both a scale and a frequency, in which case these can interact in complex ways. Others are defined such that no separate frequency appears. In the latter case, you immediately end up with the time vs. scale layout we see in wavelet diagrams (scaleograms). In the former, most software hides the complexity by merging scale and frequency into one, leaving just scale as a user-visible parameter. In torchwavelets, too, the wavelet frequency (if existent) has been “streamlined away”. Consequently, we’ll end up plotting time versus scale, as well. I’ll say more when we actually see such a scaleogram.

For visualization, we transpose the data and put it into a ggplot-friendly format:

times <- lubridate::year(enso$x) + lubridate::month(enso$x) / 12
scales <- as.numeric(wtf$scales)

df <- as_tibble(as.matrix(power_spectrum$t()), .name_repair = "universal") %>%
  mutate(time = times) %>%
  pivot_longer(!time, names_to = "scale", values_to = "power") %>%
  mutate(scale = scales[scale %>%
    str_remove("[\\.]{3}") %>%
    as.numeric()])
df %>% glimpse()

Rows: 61,983
Columns: 3
$ time   1950.083, 1950.083, 1950.083, 1950.083, 195…
$ scale  0.1613356, 0.1759377, 0.1918614, 0.2092263,…
$ power  0.03617507, 0.05985500, 0.07948010, 0.09819…

There is one additional piece of information to be incorporated, still: the so-called “cone of influence” (COI). Visually, this is a shading that tells us which part of the plot reflects incomplete, and thus, unreliable and to-be-disregarded, data. Namely, the bigger the scale, the more spread-out the analysis wavelet, and the more incomplete the overlap at the borders of the series when the wavelet slides over the data. You’ll see what I mean in a second.

The COI gets its own data frame:

coi <- wtf$coi(times[1], times[length(enso_idx)])
coi_df <- data.frame(x = as.numeric(coi[[1]]), y = as.numeric(coi[[2]]))

And now we’re ready to create the scaleogram:

labeled_scales <- c(0.25, 0.5, 1, 2, 4, 8, 16, 32, 64)
labeled_frequencies <- round(as.numeric(wtf$fourier_period(labeled_scales)), 1)

ggplot(df) +
  scale_y_continuous(
    trans = scales::compose_trans(scales::log2_trans(), scales::reverse_trans()),
    breaks = c(0.25, 0.5, 1, 2, 4, 8, 16, 32, 64),
    limits = c(max(scales), min(scales)),
    expand = c(0, 0),
    sec.axis = dup_axis(
      labels = scales::label_number(labeled_frequencies),
      name = "Fourier period (years)"
    )
  ) +
  ylab("scale (years)") +
  scale_x_continuous(breaks = seq(1950, 2020, by = 5), expand = c(0, 0)) +
  xlab("year") +
  geom_contour_filled(aes(time, scale, z = power), show.legend = FALSE) +
  scale_fill_viridis_d(option = "turbo") +
  geom_ribbon(data = coi_df, aes(x = x, ymin = y, ymax = max(scales)),
              fill = "black", alpha = 0.6) +
  theme(legend.position = "none")

Scaleogram of ENSO data.

What we see here is how, in ENSO, different rhythms have prevailed over time. Instead of “rhythms”, I could have said “scales”, or “frequencies”, or “periods” – all those translate into one another. Since, to us humans, wavelet scales don’t mean that much, the period (in years) is displayed on an additional y axis on the right.

So, we see that in the eighties, an (approximately) four-year period had exceptional influence. Thereafter, yet longer periodicities gained in dominance. And, in accordance with what we expect from prior analysis, there is a basso continuo of annual similarity.

Also, note how, at first sight, there seems to have been a decade where a six-year period stood out: right at the beginning of where (for us) measurement starts, in the fifties. However, the dark shading – the COI – tells us that, in this region, the data is not to be trusted.

Summing up, the two-dimensional analysis nicely complements the more compressed characterization we got from the DFT. Before we move on to the next series, however, let me just quickly address one question, in case you were wondering (if not, just read on, since I won’t be going into details anyway): How is this different from a spectrogram?

In a nutshell, the spectrogram splits the data into several “windows”, and computes the DFT independently on all of them. To compute the scaleogram, on the other hand, the analysis wavelet slides continuously over the data, resulting in a spectrum-equivalent for the neighborhood of each sample in the series. With the spectrogram, a fixed window size means that not all frequencies are resolved equally well: The higher frequencies appear more frequently in the interval than the lower ones, and thus, will allow for better resolution. Wavelet analysis, in contrast, is done on a set of scales deliberately arranged so as to capture a broad range of frequencies theoretically visible in a series of given length.

Take a look at the aforementioned book excerpt to see spectrogram and scaleogram compared in a nice application.

Analysis: NAO

The data file for NAO is in fixed-table format. After conversion to a tsibble, we have:

download.file(
 "https://crudata.uea.ac.uk/cru/data//nao/nao.dat",
 destfile = "nao.dat"
)

# needed for AO, as well
use_months <- seq.Date(
  from = as.Date("1950-01-01"),
  to = as.Date("2022-09-01"),
  by = "months"
)

nao <-
  read_table(
    "nao.dat",
    col_names = FALSE,
    na = "-99.99",
    skip = 3
  ) %>%
  select(-X1, -X14) %>%
  as.matrix() %>%
  t() %>%
  as.vector() %>%
  .[1:length(use_months)] %>%
  tibble(
    x = use_months,
    nao = .
  ) %>%
  mutate(x = yearmonth(x)) %>%
  fill(nao) %>%
  as_tsibble(index = x)

nao

# A tsibble: 873 x 2 [1M]
          x   nao
       
 1 1950 Jan -0.16
 2 1950 Feb  0.25
 3 1950 Mar -1.44
 4 1950 Apr  1.46
 5 1950 May  1.34
 6 1950 Jun -3.94
 7 1950 Jul -2.75
 8 1950 Aug -0.08
 9 1950 Sep  0.19
10 1950 Oct  0.19
# … with 863 more rows

Like before, we start with the spectrum:

fft <- torch_fft_fft(as.numeric(scale(nao$nao)))

num_samples <- nrow(nao)
nyquist_cutoff <- ceiling(num_samples / 2)
bins_below_nyquist <- 0:nyquist_cutoff

sampling_rate <- 12 
frequencies_per_bin <- sampling_rate / num_samples
frequencies <- frequencies_per_bin * bins_below_nyquist

df <- data.frame(f = frequencies, y = as.numeric(fft[1:(nyquist_cutoff + 1)]$abs()))
df %>% ggplot(aes(f, y)) +
  geom_line() +
  xlab("frequency (per year)") +
  ylab("magnitude") +
  ggtitle("Spectrum of NAO data")

Spectrum of NAO data, 1950 to present.

Have you been wondering for a tiny moment whether this was time-domain data – not spectral? It does look a lot more noisy than the ENSO spectrum for sure. And really, with NAO, predictability is much worse - forecast lead time usually amounts to just one or two weeks.

Proceeding as before, we pick dominant seasonalities (at least this still is possible!) to pass to feasts::STL().

strongest <- torch_topk(fft[1:(nyquist_cutoff/2)]$abs(), 6)
strongest

[[1]]
torch_tensor
102.7191
80.5129
76.1179
75.9949
72.9086
60.8281
[ CPUFloatType{6} ]

[[2]]
torch_tensor
147
99
146
59
33
78
[ CPULongType{6} ]

important_freqs <- frequencies[as.numeric(strongest[[2]])]
important_freqs

[1] 2.0068729 1.3470790 1.9931271 0.7972509 0.4398625 1.0584192

num_observations_in_season <- 12/important_freqs  
num_observations_in_season

[1]  5.979452  8.908163  6.020690 15.051724 27.281250 11.337662

Important seasonal periods are of length six, nine, eleven, fifteen, and twenty-seven months, approximately - pretty close together indeed! No wonder that, in STL decomposition, the remainder is even more significant than with ENSO:

nao %>%
  model(STL(nao ~ season(period = 6) + season(period = 9) +
              season(period = 15) + season(period = 27) +
              season(period = 12))) %>%
  components() %>%
  autoplot()

Decomposition of NAO data into trend, seasonal components, and remainder by feasts::STL().

Now, what will we see in terms of temporal evolution? Much of the code that follows is the same as for ENSO, repeated here for the reader’s convenience:

nao_idx <- nao$nao %>% as.numeric() %>% torch_tensor()
dt <- 1/12 # same interval as for ENSO
wtf <- wavelet_transform(length(nao_idx), dt = dt)
power_spectrum <- wtf$power(nao_idx)

times <- lubridate::year(nao$x) + lubridate::month(nao$x)/12 # also same
scales <- as.numeric(wtf$scales) # will be same because both series have same length

df <- as_tibble(as.matrix(power_spectrum$t()), .name_repair = "universal") %>%
  mutate(time = times) %>%
  pivot_longer(!time, names_to = "scale", values_to = "power") %>%
  mutate(scale = scales[scale %>%
    str_remove("[\\.]{3}") %>%
    as.numeric()])

coi <- wtf$coi(times[1], times[length(nao_idx)])
coi_df <- data.frame(x = as.numeric(coi[[1]]), y = as.numeric(coi[[2]]))

labeled_scales <- c(0.25, 0.5, 1, 2, 4, 8, 16, 32, 64) # same since scales are same 
labeled_frequencies <- round(as.numeric(wtf$fourier_period(labeled_scales)), 1)

ggplot(df) +
  scale_y_continuous(
    trans = scales::compose_trans(scales::log2_trans(), scales::reverse_trans()),
    breaks = c(0.25, 0.5, 1, 2, 4, 8, 16, 32, 64),
    limits = c(max(scales), min(scales)),
    expand = c(0, 0),
    sec.axis = dup_axis(
      labels = scales::label_number(labeled_frequencies),
      name = "Fourier period (years)"
    )
  ) +
  ylab("scale (years)") +
  scale_x_continuous(breaks = seq(1950, 2020, by = 5), expand = c(0, 0)) +
  xlab("year") +
  geom_contour_filled(aes(time, scale, z = power), show.legend = FALSE) +
  scale_fill_viridis_d(option = "turbo") +
  geom_ribbon(data = coi_df, aes(x = x, ymin = y, ymax = max(scales)),
              fill = "black", alpha = 0.6) +
  theme(legend.position = "none")

Scaleogram of NAO data.

That, literally, is a much more colorful picture than with ENSO! High frequencies are present, and continually dominant, over the whole time period.

Interestingly, though, we see similarities to ENSO, as well: In both, there is an important pattern, of periodicity four or slightly more years, that exerces influence during the eighties, nineties, and early two-thousands – only with ENSO, it shows peak impact during the nineties, while with NAO, its dominance is most visible in the first decade of this century. Also, both phenomena exhibit a strongly visible peak, of period two years, around 1970. So, is there a close(-ish) connection between both oscillations? This question, of course, is for the domain experts to answer. At least I found a recent study (Scaife et al. (2014)) that not only suggests there is, but uses one (ENSO, the more predictable one) to inform forecasts of the other:

Scaife, A. A., Alberto Arribas Herranz, E. Blockley, et al. 2014. “Skillful Long-Range Prediction of European and North American Winters.” Geophysical Research Letters 41 (7): 2514–19. https://www.microsoft.com/en-us/research/publication/skillful-long-range-prediction-of-european-and-north-american-winters/.

Previous studies have shown that the El Niño–Southern Oscillation can drive interannual variations in the NAO [Brönnimann et al., 2007] and hence Atlantic and European winter climate via the stratosphere [Bell et al., 2009]. […] this teleconnection to the tropical Paciﬁc is active in our experiments, with forecasts initialized in El Niño/La Niña conditions in November tending to be followed by negative/positive NAO conditions in winter.

Will we see a similar relationship for AO, our third series under investigation? We might expect so, since AO and NAO are closely related (or even, two sides of the same coin).

Analysis: AO

First, the data:

download.file(
 "https://www.cpc.ncep.noaa.gov/products/precip/CWlink/daily_ao_index/monthly.ao.index.b50.current.ascii.table",
 destfile = "ao.dat"
)

ao <-
  read_table(
    "ao.dat",
    col_names = FALSE,
    skip = 1
  ) %>%
  select(-X1) %>%
  as.matrix() %>% 
  t() %>%
  as.vector() %>%
  .[1:length(use_months)] %>%
  tibble(x = use_months,
         ao = .) %>%
  mutate(x = yearmonth(x)) %>%
  fill(ao) %>%
  as_tsibble(index = x) 

ao

# A tsibble: 873 x 2 [1M]
          x     ao
        
 1 1950 Jan -0.06 
 2 1950 Feb  0.627
 3 1950 Mar -0.008
 4 1950 Apr  0.555
 5 1950 May  0.072
 6 1950 Jun  0.539
 7 1950 Jul -0.802
 8 1950 Aug -0.851
 9 1950 Sep  0.358
10 1950 Oct -0.379
# … with 863 more rows

And the spectrum:

fft <- torch_fft_fft(as.numeric(scale(ao$ao)))

num_samples <- nrow(ao)
nyquist_cutoff <- ceiling(num_samples / 2)
bins_below_nyquist <- 0:nyquist_cutoff

sampling_rate <- 12 # per year
frequencies_per_bin <- sampling_rate / num_samples
frequencies <- frequencies_per_bin * bins_below_nyquist

df <- data.frame(f = frequencies, y = as.numeric(fft[1:(nyquist_cutoff + 1)]$abs()))
df %>% ggplot(aes(f, y)) +
  geom_line() +
  xlab("frequency (per year)") +
  ylab("magnitude") +
  ggtitle("Spectrum of AO data")

Spectrum of AO data, 1950 to present.

Well, this spectrum looks even more random than NAO’s, in that not even a single frequency stands out. For completeness, here is the STL decomposition:

strongest <- torch_topk(fft[1:(nyquist_cutoff/2)]$abs(), 5)

important_freqs <- frequencies[as.numeric(strongest[[2]])]
important_freqs
# [1] 0.01374570 0.35738832 1.77319588 1.27835052 0.06872852

num_observations_in_season <- 12/important_freqs  
num_observations_in_season
# [1] 873.000000  33.576923   6.767442   9.387097 174.600000 

ao %>%
  model(STL(ao ~ season(period = 33) + season(period = 7) +
              season(period = 9) + season(period = 174))) %>%
  components() %>%
  autoplot()

Decomposition of AO data into trend, seasonal components, and remainder by feasts::STL().

Finally, what can the scaleogram tell us about dominant patterns?

ao_idx <- ao$ao %>% as.numeric() %>% torch_tensor()
dt <- 1/12 # same interval as for ENSO and NAO
wtf <- wavelet_transform(length(ao_idx), dt = dt)
power_spectrum <- wtf$power(ao_idx)

times <- lubridate::year(ao$x) + lubridate::month(ao$x)/12 # also same
scales <- as.numeric(wtf$scales) # will be same because all series have same length

df <- as_tibble(as.matrix(power_spectrum$t()), .name_repair = "universal") %>%
  mutate(time = times) %>%
  pivot_longer(!time, names_to = "scale", values_to = "power") %>%
  mutate(scale = scales[scale %>%
    str_remove("[\\.]{3}") %>%
    as.numeric()])

coi <- wtf$coi(times[1], times[length(ao_idx)])
coi_df <- data.frame(x = as.numeric(coi[[1]]), y = as.numeric(coi[[2]]))

labeled_scales <- c(0.25, 0.5, 1, 2, 4, 8, 16, 32, 64) # same since scales are same 
labeled_frequencies <- round(as.numeric(wtf$fourier_period(labeled_scales)), 1)

ggplot(df) +
  scale_y_continuous(
    trans = scales::compose_trans(scales::log2_trans(), scales::reverse_trans()),
    breaks = c(0.25, 0.5, 1, 2, 4, 8, 16, 32, 64),
    limits = c(max(scales), min(scales)),
    expand = c(0, 0),
    sec.axis = dup_axis(
      labels = scales::label_number(labeled_frequencies),
      name = "Fourier period (years)"
    )
  ) +
  ylab("scale (years)") +
  scale_x_continuous(breaks = seq(1950, 2020, by = 5), expand = c(0, 0)) +
  xlab("year") +
  geom_contour_filled(aes(time, scale, z = power), show.legend = FALSE) +
  scale_fill_viridis_d(option = "turbo") +
  geom_ribbon(data = coi_df, aes(x = x, ymin = y, ymax = max(scales)),
              fill = "black", alpha = 0.6) +
  theme(legend.position = "none")

Scaleogram of AO data.

Having seen the overall spectrum, the lack of strongly dominant patterns in the scaleogram does not come as a big surprise. It is tempting – for me, at least – to see a reflection of ENSO around 1970, all the more since by transitivity, AO and ENSO should be related in some way. But here, qualified judgment really is reserved to the experts.

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!

Photo by ActionVance on Unsplash

Getting started with deep learning on graphs

Sigrid Keydana — Tue, 22 Nov 2022 00:00:00 GMT

If, in deep-learning world, the first half of the last decade has been the age of images, and the second, that of language, one could say that now, we’re living in the age of graphs. At least, that’s what commonly cited research metrics suggest. But as we’re all aware, deep-learning research is anything but an ivory tower. To see real-world implications, it suffices to reflect on how many things can be modeled as graphs. Some things quite naturally “are” graphs, in the sense of having nodes and edges: neurons, underground stations, social networks. Other things can fruitfully be modeled as graphs: molecules, for example; or language, concepts, three-dimensional shapes … If deep learning on graphs is desirable, what are the challenges, and what do we get for free?

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.

A minimal graph

Now, a basic Data object is created from three tensors. The first holds the node features: two features each for five nodes. (Both features are identical on purpose, for “cognitive ease” – on our, not the algorithm’s, part.)

import torch

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

The second specifies existing connections. For undirected graphs (like ours), each edge appears twice. The tensor you see here is specified in one-edge-per-line form for convenience reasons; to the Data() constructor we’ll pass its transpose instead.

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)

The third tensor holds the node labels. (The task will be one of node – not edge, not graph – classification.)

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

Constructing and inspecting the resulting graph, we have:

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

For our upcoming experiments, it’s more helpful, though, to visualize the graph:

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)

Although our experiments won’t be about training performance (how could they be, with just five nodes), let me remark in passing that this graph is small, but not boring: The middle node is equally connected to both “sides”, yet feature-wise, it would pretty clearly appear to belong on just one of them. (Which is true, given the provided class labels). Such a constellation is interesting because, in the majority of networks, edges indicate similarity.

A minimal GNN

Now, we code and run the minimal GNN. We’re not interested in class labels (yet); we just want to see each node’s embeddings after a single pass.

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

Evidently, we just had to start the process – but what process, exactly? From what we know about the three stages of message passing, an essential question is what nodes do with the information that flows over the edges. Our first experiment, then, is to inspect the incoming messages.

Poking into `message()`

In message(), we have access to a structure named x_j. This tensor holds, for each node , the embeddings of all nodes connected to it via incoming edges. We’ll print them, and then, just return them, unchanged.

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

Let me spell this out. In data.edge_index, repeated here for convenience:

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

the first pair denotes the edge from node 0 (that had features (1, 1)) to node 1. This information is found in x_j’s first row. Then the second row holds the information flowing in the opposite direction, namely, the features associated with node 1. And so on.

Interestingly, since we’re passing through this module just once, we can see the messages that will be sent without even running it.

Namely, since data.edge_index[0] designates the source nodes for each edge:

data.edge_index[0]

we can index into data.x to pick up what will be the incoming features for each connection.

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

Now, what does this tell us? Node 0, for example, received messages from nodes 1 and 2: (2, 2) and (3, 3), respectively. We know that the default aggregation mode is add; and so, would expect an outcome of (5, 5). Indeed, this is the new embedding for node 0.

In a nutshell, thus, the minimal GNN updates every node’s embedding so as to prototypically reflect the node’s neighborhood. Take care though: Nodes represent their neighborhoods, but themselves, they count for nothing. We will change that now.

Adding self loops

All we need to do is modify the adjacency matrix to include edges going from each node back to itself.

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

As expected, the neighborhood summary at each node now includes a contribution from each node itself.

Now we know how to access the messages, we’d like to aggregate them in a non-standard way.

Customizing `aggregate()`

Instead of message(), we now override aggregate(). If we wanted to use another of the “standard” aggregation modes (mean, mul, min, or max), we could just override __init__(), like so:

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

To implement custom summaries, however, we make use of torch_scatter (one of PyG’s installation prerequisites) for optimal performance. Let me show this by means of a simple example.

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

In aggregate(), we have two types of tensors to work with. One, inputs, holds what was returned from message(). In our case, this is identical to x_j, since we didn’t make any modifications to the default behavior. The second, index, holds the recipe for where in the aggregation those features should go. Here, the very first tuple, (1, 1), will contribute to the summary for node 1; the second, (2, 2), to that for node 0 – and so on. By the way, just like x_j (in a single-layer, single-pass setup) is “just” data.x[data.edge_index[0]], that index is “just” data.edge_index[1]. Meaning, this is the list of target nodes connected to the edges in question.

At this point, all kinds of manipulations could be done on either inputs or index; however, we content ourselves with just passing them through to torch_scatter.scatter(), and returning the negated sums. We’ve successfully built a network of contrarians.

By now, we’ve played with message() as well as aggregate(). What about update()?

Adding memory to `update()`

There’s one thing really strange in what we’re doing. It doesn’t jump to the eye, since we’re not simulating a real training phase; we’ve been calling the layer just once. If we hadn’t, we’d have noticed that at every call, the nodes happily forget who they were before, dutifully assuming the new identities assigned. In reality, we probably want them to evolve in a more consistent way.

For example:

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

In update(), we have access to both the final message aggregate (inputs) and the nodes’ prior states (x). Here, I’m just averaging those two.

At this point, we’ve successfully acquainted ourselves with the three stages of message passing: acting on individual messages, aggregating them, and self-updating based on past state and new information. But none of our models so far could be called a neural network, since there was no learning involved.

Adding parameters

If we look back at the generic message passing formulation:

we see two places where neural network modules can act on the computation: before message aggregation, and as part of the node update process. First, we illustrate the former option. For example, we can apply a MLP in forward(), before the call to aggregate():

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

Finally, we can apply network modules in both places, as exemplified next.

General message passing

We keep the MLP from the previous class, and add a second in update():

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

At this point, I hope you’ll feel comfortable to play around, subclassing the MessagePassing base class. Also, if now you consult the above-mentioned documentation page (Creating message passing networks), you’ll be able to map the example implementations (dedicated to popular GNN layer types) to where they “hook into” the message passing process.

Experimentation with MessagePassing was the point of this post. However, you may be wondering: How do I actually use this for node classification? Didn’t the graph have a class defined for each node? (It did: data.y.)

So let me conclude with a (minimal) end-to-end example that uses one of the above modules.

A minimal workflow

To that purpose, we compose that module with a linear one that performs node classification:

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)

We can then train the model like any other:


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

And that’s it for this time. Stay tuned for examples of how graph models are applied in the sciences, as well as illustrations of bleeding-edge developments in Geometric Deep Learning, the principles-based, heuristics-transcending approach to neural networks.

Thanks for reading!

Photo by Alina Grubnyak on Unsplash

Forecasting El Niño-Southern Oscillation (ENSO)

Sigrid Keydana — Wed, 02 Feb 2022 00:00:00 GMT

Today, we use the convLSTM introduced in a previous post to predict El Niño-Southern Oscillation (ENSO).

El Niño, la Niña

ENSO refers to a changing pattern of sea surface temperatures and sea-level pressures occurring in the equatorial Pacific. From its three overall states, probably the best-known is El Niño. 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. The opposite conditions are termed La Niña. Everything in-between is classified as normal.

ENSO has great impact on the weather worldwide, and routinely harms ecosystems and societies through storms, droughts and flooding, possibly resulting in famines and economic crises. The best societies can do is try to adapt and mitigate severe consequences. Such efforts are aided by accurate forecasts, the further ahead the better.

Here, deep learning (DL) can potentially help: Variables like sea surface temperatures and pressures are given on a spatial grid – that of the earth – and as we know, DL is good at extracting spatial (e.g., image) features. For ENSO prediction, architectures like convolutional neural networks (Ham et al. (2019)) or convolutional-recurrent hybrids¹ are habitually used. One such hybrid is just our convLSTM; it operates on sequences of features given on a spatial grid. Today, thus, we’ll be training a model for ENSO forecasting. This model will have a convLSTM for its central ingredient.

¹ E.g., Forecasting El Niño with Convolutional and Recurrent Neural Networks.

² E.g., CMIP5, CNRM-CM5, or HADGEM2-ES.

Before we start, a note. While our model fits well with architectures described in the relevant papers, the same cannot be said for amount of training data used. For reasons of practicality, we use actual observations only; consequently, we end up with a small (relative to the task) dataset. In contrast, research papers tend to make use of climate simulations², resulting in significantly more data to work with.

From the outset, then, we don’t expect stellar performance. Nevertheless, this should make for an interesting case study, and a useful code template for our readers to apply to their own data.

Data

We will attempt to predict monthly average sea surface temperature in the Niño 3.4 region³, as represented by the Niño 3.4 Index, plus categorization as one of El Niño, La Niña or neutral⁴. Predictions will be based on prior monthly sea surface temperatures spanning a large portion of the globe.

³ This region extends over latitudes from 5° south to 5° north and longitudes from 120° west to 170° west.

⁴ That classification is based on ONI (Oceanic Niño Index), a measure representing 3-month average anomalies in the Niño 3.4 Index.

On the input side, public and ready-to-use data may be downloaded from Tokyo Climate Center; as to prediction targets, we obtain index and classification here.

Input and target data both are provided monthly. They intersect in the time period ranging from 1891-01-01 to 2020-08-01; so this is the range of dates we’ll be zooming in on.

Input: Sea Surface Temperatures

Monthly sea surface temperatures are provided in a latitude-longitude grid of resolution 1°. Details of how the data were processed are available here.

Data files are available in GRIB format; each file contains averages computed for a single month. We can either download individual files or generate a text file of URLs for download. Once you’ve saved these URLs to a file, you can have R get the files for you like so:

purrr::walk(
   readLines("files"),
   function(f) download.file(url = f, destfile = basename(f))
)

From R, we can read GRIB files using stars. For example:

# let's just quickly load all libraries we require to start with

library(torch)
library(tidyverse)
library(stars)
library(viridis)
library(ggthemes)

torch_manual_seed(777)

read_stars(file.path(grb_dir, "sst189101.grb"))

stars object with 2 dimensions and 1 attribute
attribute(s):
 sst189101.grb   
 Min.   :-274.9  
 1st Qu.:-272.8  
 Median :-259.1  
 Mean   :-260.0  
 3rd Qu.:-248.4  
 Max.   :-242.8  
 NA's   :21001   
dimension(s):
  from  to offset delta                       refsys point values    
x    1 360      0     1 Coordinate System importe...    NA   NULL [x]
y    1 180     90    -1 Coordinate System importe...    NA   NULL [y]

So in this GRIB file, we have one attribute - which we know to be sea surface temperature – on a two-dimensional grid. As to the latter, we can complement what stars tells us with additional info found in the documentation:

The east-west grid points run eastward from 0.5ºE to 0.5ºW, while the north-south grid points run northward from 89.5ºS to 89.5ºN.

We note a few things we’ll want to do with this data. For one, the temperatures seem to be given in Kelvin, but with minus signs.⁵ We’ll remove the minus signs and convert to degrees Celsius for convenience. We’ll also have to think about what to do with the NAs that appear for all non-maritime coordinates.

⁵ This may also be an artifact produced by the software stack involved in reading the file.

Before we get there though, we need to combine data from all files into a single data frame. This adds an additional dimension, time, ranging from 1891/01/01 to 2020/01/12:

grb <- read_stars(
  file.path(grb_dir, map(readLines("files", warn = FALSE), basename)), along = "time") %>%
  st_set_dimensions(3,
                    values = seq(as.Date("1891-01-01"), as.Date("2020-12-01"), by = "months"),
                    names = "time"
                    )

grb

stars object with 3 dimensions and 1 attribute
attribute(s), summary of first 1e+05 cells:
 sst189101.grb   
 Min.   :-274.9  
 1st Qu.:-273.3  
 Median :-258.8  
 Mean   :-260.0  
 3rd Qu.:-247.8  
 Max.   :-242.8  
 NA's   :33724   
dimension(s):
     from   to offset delta                       refsys point                    values    
x       1  360      0     1 Coordinate System importe...    NA                      NULL [x]
y       1  180     90    -1 Coordinate System importe...    NA                      NULL [y]
time    1 1560     NA    NA                         Date    NA 1891-01-01,...,2020-12-01

Let’s visually inspect the spatial distribution of monthly temperatures for one year, 2020:

ggplot() +
  geom_stars(data = grb %>% filter(between(time, as.Date("2020-01-01"), as.Date("2020-12-01"))), alpha = 0.8) +
  facet_wrap("time") +
  scale_fill_viridis() +
  coord_equal() +
  theme_map() +
  theme(legend.position = "none")

Target: Niño 3.4 Index

For the Niño 3.4 Index, we download the monthly data and, among the provided features, zoom in on two: the index itself (column NINO34_MEAN) and PHASE, which can be E (El Niño), L (La Niño) or N (neutral).

nino <- read_table2("ONI_NINO34_1854-2020.txt", skip = 9) %>%
  mutate(month = as.Date(paste0(YEAR, "-", `MON/MMM`, "-01"))) %>%
  select(month, NINO34_MEAN, PHASE) %>%
  filter(between(month, as.Date("1891-01-01"), as.Date("2020-08-01"))) %>%
  mutate(phase_code = as.numeric(as.factor(PHASE)))

nrow(nino)

Next, we look at how to get the data into a format convenient for training and prediction.

Preprocessing

Input

First, we remove all input data for points in time where ground truth data are still missing.

sst <- grb %>% filter(time <= as.Date("2020-08-01"))

Next, as is done by e.g. Ham et al. (2019), we only use grid points between 55° south and 60° north. This has the additional advantage of reducing memory requirements.

Ham, Yoo-Geun, Jeong-Hwan Kim, and Jing-Jia Luo. 2019. “Deep learning for multi-year ENSO forecasts.” 573 (7775): 568–72. https://doi.org/10.1038/s41586-019-1559-7.

sst <- grb %>% filter(between(y,-55, 60))

dim(sst)

360, 115, 1560

As already alluded to, with the little data we have we can’t expect much in terms of generalization. Still, we set aside a small portion of the data for validation, since we’d like for this post to serve as a useful template to be used with bigger datasets.

sst_train <- sst %>% filter(time < as.Date("1990-01-01"))
sst_valid <- sst %>% filter(time >= as.Date("1990-01-01"))

From here on, we work with R arrays.

sst_train <- as.tbl_cube.stars(sst_train)$mets[[1]]
sst_valid <- as.tbl_cube.stars(sst_valid)$mets[[1]]

Conversion to degrees Celsius is not strictly necessary, as initial experiments showed a slight performance increase due to normalizing the input, and we’re going to do that anyway. Still, it reads nicer to humans than Kelvin.

sst_train <- sst_train + 273.15
quantile(sst_train, na.rm = TRUE)

     0%     25%     50%     75%    100% 
-1.8000 12.9975 21.8775 26.8200 34.3700

Not at all surprisingly, global warming is evident from inspecting temperature distribution on the validation set (which was chosen to span the last thirty-one years).

sst_valid <- sst_valid + 273.15
quantile(sst_valid, na.rm = TRUE)

    0%    25%    50%    75%   100% 
-1.800 13.425 22.335 27.240 34.870

The next-to-last step normalizes both sets according to training mean and variance.

train_mean <- mean(sst_train, na.rm = TRUE)
train_sd <- sd(sst_train, na.rm = TRUE)

sst_train <- (sst_train - train_mean) / train_sd

sst_valid <- (sst_valid - train_mean) / train_sd

Finally, what should we do about the NA entries? We set them to zero, the (training set) mean. That may not be enough of an action though: It means we’re feeding the network roughly 30% misleading data. This is something we’re not done with yet.

sst_train[is.na(sst_train)] <- 0
sst_valid[is.na(sst_valid)] <- 0

Target

The target data are split analogously. Let’s check though: Are phases (categorizations) distributedly similarly in both sets?

nino_train <- nino %>% filter(month < as.Date("1990-01-01"))
nino_valid <- nino %>% filter(month >= as.Date("1990-01-01"))

nino_train %>% group_by(phase_code, PHASE) %>% summarise(count = n(), avg = mean(NINO34_MEAN))

# A tibble: 3 x 4
# Groups:   phase_code [3]
  phase_code PHASE count   avg
          
1          1 E       301  27.7
2          2 L       333  25.6
3          3 N       554  26.7

nino_valid %>% group_by(phase_code, PHASE) %>% summarise(count = n(), avg = mean(NINO34_MEAN))

# A tibble: 3 x 4
# Groups:   phase_code [3]
  phase_code PHASE count   avg
          
1          1 E        93  28.1
2          2 L        93  25.9
3          3 N       182  27.2

This doesn’t look too bad. Of course, we again see the overall rise in temperature, irrespective of phase.

Lastly, we normalize the index, same as we did for the input data.

train_mean_nino <- mean(nino_train$NINO34_MEAN)
train_sd_nino <- sd(nino_train$NINO34_MEAN)

nino_train <- nino_train %>% mutate(NINO34_MEAN = scale(NINO34_MEAN, center = train_mean_nino, scale = train_sd_nino))
nino_valid <- nino_valid %>% mutate(NINO34_MEAN = scale(NINO34_MEAN, center = train_mean_nino, scale = train_sd_nino))

On to the torch dataset.

`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:  #182,080 parameters
● conv1:  #801 parameters
● conv2:  #25,632 parameters
● conv3:  #25,632 parameters
● linear:  #2,138,176 parameters
● b1:  #128 parameters
● cont:  #8,320 parameters
● cat:  #8,320 parameters
● cont_output:  #129 parameters
● cat_output:  #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.

Predictions

Now losses tend to be abstract; let’s see what actually gets predicted. We obtain predictions for index values and phases like so …

net$eval()

pred_index <- c()
pred_phase <- c()

coro::loop(for (b in valid_dl) {

  output <- net(b$x$to(device = device))

  pred_index <- c(pred_index, output[[2]]$to(device = "cpu"))
  pred_phase <- rbind(pred_phase, as.matrix(output[[3]]$to(device = "cpu")))

})

… and combine these with the ground truth, stripping off the first six rows (six was the number of timesteps used as predictors):

valid_perf <- data.frame(
  actual_temp = nino_valid$NINO34_MEAN[(batch_size + 1):nrow(nino_valid)] * train_sd_nino + train_mean_nino,
  actual_nino = factor(nino_valid$phase_code[(batch_size + 1):nrow(nino_valid)]),
  pred_temp = pred_index * train_sd_nino + train_mean_nino,
  pred_nino = factor(pred_phase %>% apply(1, which.max))
)

For the phase, we can generate a confusion matrix:

yardstick::conf_mat(valid_perf, actual_nino, pred_nino)

          Truth
Prediction   1   2   3
         1  70   0  43
         2   0  47  10
         3  23  46 123

This looks better than expected (based on the losses). Phases 1 and 2 correspond to El Niño and La Niña, respectively, and these get sharply separated.

What about the Niño 3.4 Index? Let’s plot predictions versus ground truth:

valid_perf <- valid_perf %>% 
  select(actual = actual_temp, predicted = pred_temp) %>% 
  add_column(month = seq(as.Date("1990-07-01"), as.Date("2020-08-01"), by = "months")) %>%
  pivot_longer(-month, names_to = "Index", values_to = "temperature")

ggplot(valid_perf, aes(x = month, y = temperature, color = Index)) +
  geom_line() +
  scale_color_manual(values = c("#006D6F", "#B2FFFF")) +
  theme_classic()

Nino 3.4 Index: Ground truth vs. predictions (validation set).

This does not look bad either. However, we need to keep in mind that we’re predicting just a single time step ahead. We probably should not overestimate the results. Which leads directly to the discussion.

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

Beyond alchemy: A first look at geometric deep learning

Sigrid Keydana — Thu, 26 Aug 2021 00:00:00 GMT

To the practitioner, it may often seem that with deep learning, there is a lot of magic involved. Magic in how hyper-parameter choices affect performance, for example. More fundamentally yet, magic in the impacts of architectural decisions. Magic, sometimes, in that it even works (or not). Sure, papers abound that strive to mathematically prove why, for specific solutions, in specific contexts, this or that technique will yield better results. But theory and practice are strangely dissociated: If a technique does turn out to be helpful in practice, doubts may still arise to whether that is, in fact, due to the purported mechanism. Moreover, level of generality often is low.

In this situation, one may feel grateful for approaches that aim to elucidate, complement, or replace some of the magic. By “complement or replace”, I’m alluding to attempts to incorporate domain-specific knowledge into the training process. Interesting examples exist in several sciences, and I certainly hope to be able to showcase a few of these, on this blog at a later time. As for the “elucidate”, this characterization is meant to lead on to the topic of this post: the program of geometric deep learning.

Geometric deep learning: An attempt at unification

Geometric deep learning (henceforth: GDL) is what a group of researchers, including Michael Bronstein, Joan Bruna, Taco Cohen, and Petar Velicković, call their attempt to build a framework that places deep learning (DL) on a solid mathematical basis.

Prima facie, this is a scientific endeavor: They take existing architectures and practices and show where these fit into the “DL blueprint”. DL research being all but confined to the ivory tower, though, it’s fair to assume that this is not all: From those mathematical foundations, it should be possible to derive new architectures, new techniques to fit a given task. Who, then, should be interested in this? Researchers, for sure; to them, the framework may well prove highly inspirational. Secondly, everyone interested in the mathematical constructions themselves — this probably goes without saying. Finally, the rest of us, as well: Even understood at a purely conceptual level, the framework offers an exciting, inspiring view on DL architectures that – I think – is worth getting to know about as an end in itself. The goal of this post is to provide a high-level introduction .

Before we get started though, let me mention the primary source for this text: Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges (Bronstein et al. (2021)).

Bronstein, Michael M., Joan Bruna, Taco Cohen, and Petar Velickovic. 2021. “Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges.” CoRR abs/2104.13478. https://arxiv.org/abs/2104.13478.

Geometric priors

A prior, in the context of machine learning, is a constraint imposed on the learning task. A generic prior could come about in different ways; a geometric prior, as defined by the GDL group, arises, originally, from the underlying domain of the task. Take image classification, for example. The domain is a two-dimensional grid. Or graphs: The domain consists of collections of nodes and edges.

In the GDL framework, two all-important geometric priors are symmetry and scale separation.

Symmetry

A symmetry, in physics and mathematics, is a transformation that leaves some property of an object unchanged. The appropriate meaning of “unchanged” depends on what sort of property we’re talking about. Say the property is some “essence”, or identity — what object something is. If I move a few steps to the left, I’m still myself: The essence of being “myself” is shift-invariant. (Or: translation-invariant.) But say the property is location. If I move to the left, my location moves to the left. Location is shift-equivariant. (Translation-equivariant.)

So here we have two forms of symmetry: invariance and equivariance. One means that when we transform an object, the thing we’re interested in stays the same. The other means that we have to transform that thing as well.

The next question then is: What are possible transformations? Translation we already mentioned; on images, rotation or flipping are others. Transformations are composable; I can rotate the digit 3 by thirty degrees, then move it to the left by five units; I could also do things the other way around. (In this case, though not necessarily in general, the results are the same.) Transformations can be undone: If first I rotate, in some direction, by five degrees, I can then rotate in the opposite one, also by five degrees, and end up in the original position. We’ll see why this matters when we cross the bridge from the domain (grids, sets, etc.) to the learning algorithm.

Scale separation

After symmetry, another important geometric prior is scale separation. Scale separation means that even if something is very “big” (extends a long way in, say, one or two dimensions), we can still start from small patches and “work our way up”. For example, take a cuckoo clock. To discern the hands, you don’t need to pay attention to the pendulum. And vice versa. And once you’ve taken inventory of hands and pendulum, you don’t have to care about their texture or exact position anymore.

In a nutshell, given scale separation, the top-level structure can be determined through successive steps of coarse-graining. We’ll see this prior nicely reflected in some neural-network algorithms.

From domain priors to algorithmic ones

So far, all we’ve really talked about is the domain, using the word in the colloquial sense of “on what structure”, or “in terms of what structure”, something is given. In mathematical language, though, domain is used in a more narrow way, namely, for the “input space” of a function. And a function, or rather, two of them, is what we need to get from priors on the (physical) domain to priors on neural networks.

The first function maps from the physical domain to signal space. If, for images, the domain was the two-dimensional grid, the signal space now consists of images the way they are represented in a computer, and will be worked with by a learning algorithm. For example, in the case of RGB images, that representation is three-dimensional, with a color dimension on top of the inherited spatial structure. What matters is that by this function, the priors are preserved. If something is translation-invariant before “real-to-virtual” conversion, it will still be translation-invariant thereafter.

Next, we have another function: the algorithm, or neural network, acting on signal space. Ideally, this function, again, would preserve the priors. Below, we’ll see how basic neural-network architectures typically preserve some important symmetries, but not necessarily all of them. We’ll also see how, at this point, the actual task makes a difference. Depending on what we’re trying to achieve, we may want to maintain some symmetry, but not care about another. The task here is analogous to the property in physical space. Just like in physical space, a movement to the left does not alter identity, a classifier, presented with that same shift, won’t care at all. But a segmentation algorithm will – mirroring the real-world shift in position.

Now that we’ve made our way to algorithm space, the above requirement, formulated on physical space – that transformations be composable – makes sense in another light: Composing functions is exactly what neural networks do; we want these compositions to work just as deterministically as those of real-world transformations.

In sum, the geometric priors and the way they impose constraints, or desiderates, rather, on the learning algorithm lead to what the GDL group call their deep learning “blueprint”. Namely, a network should be composed of the following types of modules:

Linear group-equivariant layers. (Here group is the group of transformations whose symmetries we’re interested to preserve.)
Nonlinearities. (This really does not follow from geometric arguments, but from the observation, often stated in introductions to DL, that without nonlinearities, there is no hierarchical composition of features, since all operations can be implemented in a single matrix multiplication.)
Local pooling layers. (These achieve the effect of coarse-graining, as enabled by the scale separation prior.)
A group-invariant layer (global pooling). (Not every task will require such a layer to be present.)

Having talked so much about the concepts, which are highly fascinating, this list may seem a bit underwhelming. That’s what we’ve been doing anyway, right? Maybe; but once you look at a few domains and associated network architectures, the picture gets colorful again. So colorful, in fact, that we can only present a very sparse selection of highlights.

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.

Images and CNNs

Vanilla CNNs are easily mapped to the four types of layers that make up the blueprint. Skipping over the nonlinearities, which, in this context, are of least interest, we next have two kinds of pooling.

First, a local one, corresponding to max- or average-pooling layers with small strides (2 or 3, say). This reflects the idea of successive coarse-graining, where, once we’ve made use of some fine-grained information, all we need to proceed is a summary.

Second, a global one, used to effectively remove the spatial dimensions. In practice, this would usually be global average pooling. Here, there’s an interesting detail worth mentioning. A common practice, in image classification, is to replace global pooling by a combination of flattening and one or more feedforward layers. Since with feedforward layers, position in the input matters, this will do away with translation invariance.

Having covered three of the four layer types, we come to the most interesting one. In CNNs, the local, group-equivariant layers are the convolutional ones. What kinds of symmetries does convolution preserve? Think about how a kernel slides over an image, computing a dot product at every location. Say that, through training, it has developed an inclination toward singling out penguin bills. It will detect, and mark, one everywhere in an image — be it shifted left, right, top or bottom in the image. What about rotational motion, though? Since kernels move vertically and horizontally, but not in a circle, a rotated bill will be missed. Convolution is shift-equivariant, not rotation-invariant.

There is something that can be done about this, though, while fully staying within the framework of GDL. Convolution, in a more generic sense, does not have to imply constraining filter movement to horizontal and vertical translation. When reflecting a general group convolution, that motion is determined by whatever transformations constitute the group action. If, for example, that action included translation by sixty degrees, we could rotate the filter to all valid positions, then take these filters and have them slide over the image. In effect, we’d just wind up with more channels in the subsequent layer – the intended base number of filters times the number of attainable positions.

This, it must be said, it just one way to do it. A more elegant one is to apply the filter in the Fourier domain, where convolution maps to multiplication. The Fourier domain, however, is as fascinating as it is out of scope for this post.

The same goes for extensions of convolution from the Euclidean grid to manifolds, where distances are no longer measured by a straight line as we know it. Often on manifolds, we’re interested in invariances beyond translation or rotation: Namely, algorithms may have to support various types of deformation. (Imagine, for example, a moving rabbit, with its muscles stretching and contracting as it hobbles.) If you’re interested in these kinds of problems, the GDL book goes into those in great detail.

For group convolution on grids – in fact, we may want to say “on things that can be arranged in a grid” – the authors give two illustrative examples. (One thing I like about these examples is something that extends to the whole book: Many applications are from the world of natural sciences, encouraging some optimism as to the role of deep learning (“AI”) in society.)

One example is from medical volumetric imaging (MRI or CT, say), where signals are represented on a three-dimensional grid. Here the task calls not just for translation in all directions, but also, rotations, of some sensible degree, about all three spatial axes. The other is from DNA sequencing, and it brings into play a new kind of invariance we haven’t mentioned yet: reverse-complement symmetry. This is because once we’ve decoded one strand of the double helix, we already know the other one.

Finally, before we wrap up the topic of CNNs, let’s mention how through creativity, one can achieve – or put cautiously, try to achieve – certain invariances by means other than network architecture. A great example, originally associated mostly with images, is data augmentation. Through data augmentation, we may hope to make training invariant to things like slight changes in color, illumination, perspective, and the like.

Graphs and GNNs

Another type of domain, underlying many scientific and non-scientific applications, are graphs. Here, we are going to be a lot more brief. One reason is that so far, we have not had many posts on deep learning on graphs, so to the readers of this blog, the topic may seem fairly abstract. The other reason is complementary: That state of affairs is exactly something we’d like to see changing. Once we write more about graph DL, occasions to talk about respective concepts will be plenty.

In a nutshell, though, the dominant type of invariance in graph DL is permutation equivariance. Permutation, because when you stack a node and its features in a matrix, it doesn’t matter whether node one is in row three or row fifteen. Equivariance, because once you do permute the nodes, you also have to permute the adjacency matrix, the matrix that captures which node is linked to what other nodes. This is very different from what holds for images: We can’t just randomly permute the pixels.

Sequences and RNNs

With RNNs, we are going be very brief as well, although for a different reason. My impression is that so far, this area of research – meaning, GDL as it relates to sequences – has not received too much attention yet, and (maybe) for that reason, seems of lesser impact on real-world applications.

In a nutshell, the authors refer two types of symmetry: First, translation-invariance, as long as a sequence is left-padded for a sufficient number of steps. (This is due to the hidden units having to be initialized somehow.) This holds for RNNs in general.

Second, time warping: If a network can be trained that correctly works on a sequence measured on some time scale, there is another network, of the same architecture but likely with different weights, that will work equivalently on re-scaled time. This invariance only applies to gated RNNs, such as the LSTM.

What’s next?

At this point, we conclude this conceptual introduction. If you want to learn more, and are not too scared by the math, definitely check out the book. (I’d also say it lends itself well to incremental understanding, as in, iteratively going back to some details once one has acquired more background.)

Something else to wish for certainly is practice. There is an intimate connection between GDL and deep learning on graphs; which is one reason we’re hoping to be able to feature the latter more frequently in the future. The other is the wealth of interesting applications that take graphs as their input. Until then, thanks for reading!

Photo by NASA on Unsplash

Starting to think about AI Fairness

Sigrid Keydana — Thu, 15 Jul 2021 00:00:00 GMT

If you use deep learning for unsupervised part-of-speech tagging of Sanskrit ¹, or knowledge discovery in physics ², you probably don’t need to worry about model fairness. If you’re a data scientist working at a place where decisions are made about people, however, or an academic researching models that will be used to such ends, chances are that you’ve already been thinking about this topic. — Or feeling that you should. And thinking about this is hard.

¹ Srivastava et al. (2018)

Srivastava, Prakhar, Kushal Chauhan, Deepanshu Aggarwal, Anupam Shukla, Joydip Dhar, and Vrashabh Prasad Jain. 2018. “Deep Learning Based Unsupervised POS Tagging for Sanskrit.” Proceedings of the 2018 International Conference on Algorithms, Computing and Artificial Intelligence (New York, NY, USA), ACAI 2018. https://doi.org/10.1145/3302425.3302487.

² Cranmer et al. (2020)

Cranmer, Miles D., Alvaro Sanchez-Gonzalez, Peter W. Battaglia, et al. 2020. “Discovering Symbolic Models from Deep Learning with Inductive Biases.” CoRR abs/2006.11287. https://arxiv.org/abs/2006.11287.

It is hard for several reasons. In this text, I will go into just one.

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.

Situating fairness criteria

Think back to the example of a hiring algorithm. What does it mean for this algorithm to be fair? We approach this question under two – incompatible, mostly – assumptions:

The algorithm is fair if it behaves the same way independent of which demographic group it is applied to. Here demographic group could be defined by ethnicity, gender, abledness, or in fact any categorization suggested by the context.
The algorithm is fair if it does not discriminate against any demographic group.

I’ll call these the technical and societal views, respectively.

Fairness, viewed the technical way

What does it mean for an algorithm to “behave the same way” regardless of which group it is applied to?

In a classification setting, we can view the relationship between prediction () and target () as a doubly directed path. In one direction: Given true target , how accurate is prediction ? In the other: Given , how well does it predict the true class ?

Based on the direction they operate in, metrics popular in machine learning overall can be split into two categories. In the first, starting from the true target, we have recall, together with “the rates”: true positive, true negative, false positive, false negative. In the second, we have precision, together with positive (negative, resp.) predictive value.

If now we demand that these metrics be the same across groups, we arrive at corresponding fairness criteria: equal false positive rate, equal positive predictive value, etc. In the inter-group setting, the two types of metrics may be arranged under headings “equality of opportunity” and “predictive parity”. You’ll encounter these as actual headers in the summary table at the end of this text. (Said table organizes concepts from different areas into a three-category format. The overall narrative builds up towards that “map” in a bottom-up way – meaning, most entries will not make sense at this point.}

While overall, the terminology around metrics can be confusing (to me it is), these headings have some mnemonic value. Equality of opportunity suggests that people similar in real life () get classified similarly (). Predictive parity suggests that people classified similarly () are, in fact, similar ().

The two criteria can concisely be characterized using the language of statistical independence. Following Barocas et al. (2019), these are:

Separation: Given true target , prediction is independent of group membership ().
Sufficiency: Given prediction , target is independent of group membership ().

Given those two fairness criteria – and two sets of corresponding metrics – the natural question arises: Can we satisfy both? Above, I was mentioning precision and recall on purpose: to maybe “prime” you to think in the direction of “precision-recall trade-off”. And really, these two categories reflect different preferences; usually, it is impossible to optimize for both. The most famous, probably, result is due to Chouldechova (2016) : It says that predictive parity (testing for sufficiency) is incompatible with error rate balance (separation) when prevalence differs across groups. This is a theorem (yes, we’re in the realm of theorems and proofs here) that may not be surprising, in light of Bayes’ theorem, but is of great practical importance nonetheless: Unequal prevalence usually is the norm, not the exception.

Chouldechova, Alexandra. 2016. “Fair prediction with disparate impact: A study of bias in recidivism prediction instruments.” arXiv e-Prints, October, arXiv:1610.07524. https://arxiv.org/abs/1610.07524.

Yeom, Samuel, and Michael Carl Tschantz. 2018. “Discriminative but Not Discriminatory: A Comparison of Fairness Definitions Under Different Worldviews.” CoRR abs/1808.08619. http://arxiv.org/abs/1808.08619.

This necessarily means we have to make a choice. And this is where the theorems and proofs do matter. For example, Yeom and Tschantz (2018) show that in this framework – the strictly technical approach to fairness – separation should be preferred over sufficiency, because the latter allows for arbitrary disparity amplification. Thus, in this framework, we may have to work through the theorems.

What is the alternative?

Fairness, viewed as a social construct

Starting with what I just wrote: No one will likely challenge fairness being a social construct. But what does that entail?

Let me start with a biographical reminiscence. In undergraduate psychology (a long time ago), probably the most hammered-in distinction relevant to experiment planning was that between a hypothesis and its operationalization. The hypothesis is what you want to substantiate, conceptually; the operationalization is what you measure. There necessarily can’t be a one-to-one correspondence; we’re just striving to implement the best operationalization possible.

In the world of datasets and algorithms, all we have are measurements. And often, these are treated as though they were the concepts. This will get more concrete with an example, and we’ll stay with the hiring software scenario.

Assume the dataset used for training, assembled from scoring previous employees, contains a set of predictors (among which, high-school grades) and a target variable, say an indicator whether an employee did “survive” probation. There is a concept-measurement mismatch on both sides.

For one, say the grades are intended to reflect ability to learn, and motivation to learn. But depending on the circumstances, there are influence factors of much higher impact: socioeconomic status, constantly having to struggle with prejudice, overt discrimination, and more.

And then, the target variable. If the thing it’s supposed to measure is “was hired for seemed like a good fit, and was retained since was a good fit”, then all is good. But normally, HR departments are aiming for more than just a strategy of “keep doing what we’ve always been doing”.

Unfortunately, that concept-measurement mismatch is even more fatal, and even less talked about, when it’s about the target and not the predictors. (Not accidentally, we also call the target the “ground truth”.) An infamous example is recidivism prediction, where what we really want to measure – whether someone did, in fact, commit a crime – is replaced, for measurability reasons, by whether they were convicted. These are not the same: Conviction depends on more then what someone has done – for instance, if they’ve been under intense scrutiny from the outset.

Fortunately, though, the mismatch is clearly pronounced in the AI fairness literature. Friedler et al. (2016) distinguish between the construct and observed spaces; depending on whether a near-perfect mapping is assumed between these, they talk about two “worldviews”: “We’re all equal” (WAE) vs. “What you see is what you get” (WYSIWIG). If we’re all equal, membership in a societally disadvantaged group should not – in fact, may not – affect classification. In the hiring scenario, any algorithm employed thus has to result in the same proportion of applicants being hired, regardless of which demographic group they belong to. If “What you see is what you get”, we don’t question that the “ground truth” is the truth.

This talk of worldviews may seem unnecessary philosophical, but the authors go on and clarify: All that matters, in the end, is whether the data is seen as reflecting reality in a naïve, take-at-face-value way.

For example, we might be ready to concede that there could be small, albeit uninteresting effect-size-wise, statistical differences between men and women as to spatial vs. linguistic abilities, respectively. We know for sure, though, that there are much greater effects of socialization, starting in the core family and reinforced, progressively, as adolescents go through the education system. We therefore apply WAE, trying to (partly) compensate for historical injustice. This way, we’re effectively applying affirmative action, defined as

A set of procedures designed to eliminate unlawful discrimination among applicants, remedy the results of such prior discrimination, and prevent such discrimination in the future.

In the already-mentioned summary table, you’ll find the WYSIWIG principle mapped to both equal opportunity and predictive parity metrics. WAE maps to the third category, one we haven’t dwelled upon yet: demographic parity, also known as statistical parity. In line with what was said before, the requirement here is for each group to be present in the positive-outcome class in proportion to its representation in the input sample. For example, if thirty percent of applicants are Black, then at least thirty percent of people selected should be Black, as well. A term commonly used for cases where this does not happen is disparate impact: The algorithm affects different groups in different ways.

Similar in spirit to demographic parity, but possibly leading to different outcomes in practice, is conditional demographic parity ³. Here we additionally take into account other predictors in the dataset; to be precise: all other predictors. The desiderate now is that for any choice of attributes, outcome proportions should be equal, given the protected attribute and the other attributes in question. I’ll come back to why this may sound better in theory than work in practice in the next section.

³ Wachter et al. (2020b)

Summing up, we’ve seen commonly used fairness metrics organized into three groups, two of which share a common assumption: that the data used for training can be taken at face value. The other starts from the outside, contemplating what historical events, and what political and societal factors have made the given data look as they do.

Before we conclude, I’d like to try a quick glance at other disciplines, beyond machine learning and computer science, domains where fairness figures among the central topics. This section is necessarily limited in every respect; it should be seen as a flashlight, an invitation to read and reflect rather than an orderly exposition. The short section will end with a word of caution: Since drawing analogies can feel highly enlightening (and is intellectually satisfying, for sure), it is easy to abstract away practical realities. But I’m getting ahead of myself.

A quick glance at neighboring fields: law and political philosophy

In jurisprudence, fairness and discrimination constitute an important subject. A recent paper that caught my attention is Wachter et al. (2020a) . From a machine learning perspective, the interesting point is the classification of metrics into bias-preserving and bias-transforming. The terms speak for themselves: Metrics in the first group reflect biases in the dataset used for training; ones in the second do not. In that way, the distinction parallels Friedler et al. (2016) ’s confrontation of two “worldviews”. But the exact words used also hint at how guidance by metrics feeds back into society: Seen as strategies, one preserves existing biases; the other, to consequences unknown a priori, changes the world.

To the ML practitioner, this framing is of great help in evaluating what criteria to apply in a project. Helpful, too, is the systematic mapping provided of metrics to the two groups; it is here that, as alluded to above, we encounter conditional demographic parity among the bias-transforming ones. I agree that in spirit, this metric can be seen as bias-transforming; if we take two sets of people who, per all available criteria, are equally qualified for a job, and then find the whites favored over the Blacks, fairness is clearly violated. But the problem here is “available”: per all available criteria. What if we have reason to assume that, in a dataset, all predictors are biased? Then it will be very hard to prove that discrimination has occurred.

A similar problem, I think, surfaces when we look at the field of political philosophy, and consult theories on distributive justice for guidance. Heidari et al. (2018) have written a paper comparing the three criteria – demographic parity, equality of opportunity, and predictive parity – to egalitarianism, equality of opportunity (EOP) in the Rawlsian sense, and EOP seen through the glass of luck egalitarianism, respectively. While the analogy is fascinating, it too assumes that we may take what is in the data at face value. In their likening predictive parity to luck egalitarianism, they have to go to especially great lengths, in assuming that the predicted class reflects effort exerted. In the below table, I therefore take the liberty to disagree, and map a libertarian view of distributive justice to both equality of opportunity and predictive parity metrics.

Heidari, Hoda, Michele Loi, Krishna P. Gummadi, and Andreas Krause. 2018. “A Moral Framework for Understanding of Fair ML Through Economic Models of Equality of Opportunity.” CoRR abs/1809.03400. http://arxiv.org/abs/1809.03400.

In summary, we end up with two highly controversial categories of fairness criteria, one bias-preserving, “what you see is what you get”-assuming, and libertarian, the other bias-transforming, “we’re all equal”-thinking, and egalitarian. Here, then, is that often-announced table.

	Demographic parity	Equality of opportunity	Predictive parity
A.K.A. / subsumes / related concepts	statistical parity, group fairness, disparate impact, conditional demographic parity ⁴	equalized odds, equal false positive / negative rates	equal positive / negative predictive values, calibration by group
Statistical independence criterion ⁵	independence	separation	sufficiency
Individual / group	group	group (most) or individual (fairness through awareness)	group
Distributive Justice	egalitarian	libertarian (contra Heidari et al., see above)	libertarian (contra Heidari et al., see above)
Effect on bias ⁶	transforming	preserving	preserving
Policy / “worldview” ⁷	We’re all equal (WAE)	What you see is what you get (WYSIWIG)	What you see is what you get (WYSIWIG)

⁴ Wachter et al. (2020b)

Wachter, Sandra, Brent D. Mittelstadt, and Chris Russell. 2020b. “Why Fairness Cannot Be Automated: Bridging the Gap Between EU Non-Discrimination Law and AI.” CoRR abs/2005.05906. https://arxiv.org/abs/2005.05906.

⁵ Barocas et al. (2019)

Barocas, Solon, Moritz Hardt, and Arvind Narayanan. 2019. Fairness and Machine Learning. Fairmlbook.org.

⁶ Wachter et al. (2020a)

Wachter, Sandra, Brent D. Mittelstadt, and Chris Russell. 2020a. “Bias Preservation in Machine Learning: The Legality of Fairness Metrics Under EU Non-Discrimination Law.” West Virginia Law Review, Forthcoming abs/2005.05906. https://ssrn.com/abstract=3792772.

⁷ Friedler et al. (2016)

Friedler, Sorelle A., Carlos Scheidegger, and Suresh Venkatasubramanian. 2016. “On the (Im)possibility of Fairness.” CoRR abs/1609.07236. http://arxiv.org/abs/1609.07236.

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

Photo by Anders Jildén on Unsplash

Convolutional LSTM for spatial forecasting

Sigrid Keydana — Thu, 17 Dec 2020 00:00:00 GMT

This post is the first in a loose series exploring forecasting of spatially-determined data over time. By spatially-determined I mean that whatever the quantities we’re trying to predict – be they univariate or multivariate time series, of spatial dimensionality or not – the input data are given on a spatial grid.

For example, the input could be atmospheric measurements, such as sea surface temperature or pressure, given at some set of latitudes and longitudes. The target to be predicted could then span that same (or another) grid. Alternatively, it could be a univariate time series, like a meteorological index.

But wait a second, you may be thinking. For time-series prediction, we have that time-honored set of recurrent architectures (e.g., LSTM, GRU), right? Right. We do; but, once we feed spatial data to an RNN, treating different locations as different input features, we lose an essential structural relationship. Importantly, we need to operate in both space and time. We want both: recurrence relations and convolutional filters. Enter convolutional RNNs.

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.

A `torch` convLSTM

The code discussed here may be found on GitHub. (Depending on when you’re reading this, the code in that repository may have evolved though.)

My starting point was one of the PyTorch implementations found on the net, namely, this one. If you search for “PyTorch convGRU” or “PyTorch convLSTM”, you will find stunning discrepancies in how these are realized – discrepancies not just in syntax and/or engineering ambition, but on the semantic level, right at the center of what the architectures may be expected to do. As they say, let the buyer beware. (Regarding the implementation I ended up porting, I am confident that while numerous optimizations will be possible, the basic mechanism matches my expectations.)

What do I expect? Let’s approach this task in a top-down way.

Input and output

The convLSTM’s input will be a time series of spatial data, each observation being of size (time steps, channels, height, width).

Compare this with the usual RNN input format, be it in torch or Keras. In both frameworks, RNNs expect tensors of size (timesteps, input_dim)¹. input_dim is for univariate time series and greater than for multivariate ones. Conceptually, we may match this to convLSTM’s channels dimension: There could be a single channel, for temperature, say – or there could be several, such as for pressure, temperature, and humidity. The two additional dimensions found in convLSTM, height and width, are spatial indexes into the data.

¹ Leaving aside the batch dimension in this discussion.

In sum, we want to be able to pass data that:

consist of one or more features,
evolve in time, and
are indexed in two spatial dimensions.

How about the output? We want to be able to return forecasts for as many time steps as we have in the input sequence. This is something that torch RNNs do by default, while Keras equivalents do not. (You have to pass return_sequences = TRUE to obtain that effect.) If we’re interested in predictions for just a single point in time, we can always pick the last time step in the output tensor.

However, with RNNs, it is not all about outputs. RNN architectures also carry through hidden states.

What are hidden states? I carefully phrased that sentence to be as general as possible – deliberately circling around the confusion that, in my view, often arises at this point. We’ll attempt to clear up some of that confusion in a second, but let’s first finish our high-level requirements specification.

We want our convLSTM to be usable in different contexts and applications. Various architectures exist that make use of hidden states, most prominently perhaps, encoder-decoder architectures. Thus, we want our convLSTM to return those as well. Again, this is something a torch LSTM does by default, while in Keras it is achieved using return_state = TRUE.

Now though, it really is time for that interlude. We’ll sort out the ways things are called by both torch and Keras, and inspect what you get back from their respective GRUs and LSTMs.

Interlude: Outputs, states, hidden values … what’s what?

For this to remain an interlude, I summarize findings on a high level. The code snippets in the appendix show how to arrive at these results. Heavily commented, they probe return values from both Keras and torch GRUs and LSTMs. Running these will make the upcoming summaries seem a lot less abstract.

First, let’s look at the ways you create an LSTM in both frameworks. (I will generally use LSTM as the “prototypical RNN example”, and just mention GRUs when there are differences significant in the context in question.)

In Keras, to create an LSTM you may write something like this:

lstm <- layer_lstm(units = 1)

The torch equivalent would be:

lstm <- nn_lstm(
  input_size = 2, # number of input features
  hidden_size = 1 # number of hidden (and output!) features
)

Don’t focus on torch‘s input_size parameter for this discussion. (It’s the number of features in the input tensor.) The parallel occurs between Keras’ units and torch’s hidden_size. If you’ve been using Keras, you’re probably thinking of units as the thing that determines output size (equivalently, the number of features in the output). So when torch lets us arrive at the same result using hidden_size, what does that mean? It means that somehow we’re specifying the same thing, using different terminology. And it does make sense, since at every time step current input and previous hidden state are added²:

² In theory, it would be possible for them to be of different sizes if the respective weight matrices transformed their operands to the same output size.

Now, about those hidden states.

When a Keras LSTM is defined with return_state = TRUE, its return value is a structure of three entities called output, memory state, and carry state. In torch, the same entities are referred to as output, hidden state, and cell state. (In torch, we always get all of them.)

So are we dealing with three different types of entities? We are not.

The cell, or carry state is that special thing that sets apart LSTMs from GRUs deemed responsible for the “long” in “long short-term memory”. Technically, it could be reported to the user at all points in time; as we’ll see shortly though, it is not.

What about outputs and hidden, or memory states? Confusingly, these really are the same thing. Recall that for each input item in the input sequence, we’re combining it with the previous state, resulting in a new state, to be made used of in the next step³:

³ Yes, this is the same formula as above.

Now, say that we’re interested in looking at just the final time step – that is, the default output of a Keras LSTM. From that point of view, we can consider those intermediate computations as “hidden”. Seen like that, output and hidden states feel different.

However, we can also request to see the outputs for every time step. If we do so, there is no difference – the outputs (plural) equal the hidden states. This can be verified using the code in the appendix.

Thus, of the three things returned by an LSTM, two are really the same. How about the GRU, then? As there is no “cell state”, we really have just one type of thing left over – call it outputs or hidden states.

Let’s summarize this in a table.

Table 1: RNN terminology. Comparing torch-speak and Keras-speak. In row 1, the terms are parameter names. In rows 2 and 3, they are pulled from current documentation.
Referring to this entity:	`torch` says:	Keras says:
Number of features in the output This determines both how many output features there are and the dimensionality of the hidden states.	`hidden_size`	`units`
Per-time-step output; latent state; intermediate state … This could be named “public state” in the sense that we, the users, are able to obtain all values.	hidden state	memory state
Cell state; inner state … (LSTM only) This could be named “private state” in that we are able to obtain a value only for the last time step. More on that in a second.	cell state	carry state

Now, about that public vs. private distinction. In both frameworks, we can obtain outputs (hidden states) for every time step. The cell state, however, we can access only for the very last time step. This is purely an implementation decision. As we’ll see when building our own recurrent module, there are no obstacles inherent in keeping track of cell states and passing them back to the user.

If you dislike the pragmatism of this distinction, you can always go with the math. When a new cell state has been computed (based on prior cell state, input, forget, and cell gates – the specifics of which we are not going to get into here), it is transformed to the hidden (a.k.a. output) state making use of yet another, namely, the output gate:

Definitely, then, hidden state (output, resp.) builds on cell state, adding additional modeling power.

Now it is time to get back to our original goal and build that convLSTM. First though, let’s summarize the return values obtainable from torch and Keras.

Table 2: Contrasting ways of obtaining various return values in `torch` vs. Keras. Cf. the appendix for complete examples.
To achieve this goal:	in `torch` do:	in Keras do:
access all intermediate outputs ( = per-time-step outputs)	`ret[[1]]`	`return_sequences = TRUE`
access both “hidden state” (output) and “cell state” from final time step (only!)	`ret[[2]]`	`return_state = TRUE`
access all intermediate outputs and the final “cell state”	both of the above	`return_sequences = TRUE, return_state = TRUE`
access all intermediate outputs and “cell states” from all time steps	no way	no way

`convLSTM`, the plan

In both torch and Keras RNN architectures, single time steps are processed by corresponding Cell classes: There is an LSTM Cell matching the LSTM, a GRU Cell matching the GRU, and so on. We do the same for ConvLSTM. In convlstm_cell(), we first define what should happen to a single observation; then in convlstm(), we build up the recurrence logic.

Once we’re done, we create a dummy dataset, as reduced-to-the-essentials as can be. With more complex datasets, even artificial ones, chances are that if we don’t see any training progress, there are hundreds of possible explanations. We want a sanity check that, if failed, leaves no excuses. Realistic applications are left to future posts.

A single step: `convlstm_cell`

Our convlstm_cell’s constructor takes arguments input_dim , hidden_dim, and bias, just like a torch LSTM Cell.

But we’re processing two-dimensional input data. Instead of the usual affine combination of new input and previous state, we use a convolution of kernel size kernel_size. Inside convlstm_cell, it is self$conv that takes care of this.

Note how the channels dimension, which in the original input data would correspond to different variables, is creatively used to consolidate four convolutions into one: Each channel output will be passed to just one of the four cell gates. Once in possession of the convolution output, forward() applies the gate logic, resulting in the two types of states it needs to send back to the caller.

library(torch)
library(zeallot)

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

    c(h_prev, c_prev) %<-% prev_states
    
    combined <- torch_cat(list(x, h_prev), dim = 2)  # concatenate along channel axis
    combined_conv <- self$conv(combined)
    c(cc_i, cc_f, cc_o, cc_g) %<-% torch_split(combined_conv, self$hidden_dim, dim = 2)
    
    # 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))
  }
)

Now convlstm_cell has to be called for every time step. This is done by convlstm.

Iteration over time steps: `convlstm`

A convlstm may consist of several layers, just like a torch LSTM. For each layer, we are able to specify hidden and kernel sizes individually.

During initialization, each layer gets its own convlstm_cell. On call, convlstm executes two loops. The outer one iterates over layers. At the end of each iteration, we store the final pair (hidden state, cell state) for later reporting. The inner loop runs over input sequences, calling convlstm_cell at each time step.

We also keep track of intermediate outputs, so we’ll be able to return the complete list of hidden_states seen during the process. Unlike a torch LSTM, we do this for every layer.

convlstm <- nn_module(
  
  # hidden_dims and kernel_sizes are vectors, with one element for each layer in n_layers
  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) {
    
    c(batch_size, seq_len, num_channels, height, width) %<-% x$size()
   
    # 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)
      c(h, c) %<-% hidden_states[[i]]
      # 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 time steps
      for (t in 1:seq_len) {
        c(h, c) %<-% self$cell_list[[i]](cur_layer_input[ , t, , , ], list(h, c))
        # keep track of output (h) for every time step
        # h has dim (batch_size, hidden_size, height, width)
        output_sequence[[t]] <- h
      }

      # stack hs for all time steps 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)
  }
    
)

Calling the `convlstm`

Let’s see the input format expected by convlstm, and how to access its different outputs.

Here is a suitable input tensor.

# batch_size, seq_len, channels, height, width
x <- torch_rand(c(2, 4, 3, 16, 16))

First we make use of a single layer.

model <- convlstm(input_dim = 3, hidden_dims = 5, kernel_sizes = 3, n_layers = 1)

c(layer_outputs, layer_last_states) %<-% model(x)

We get back a list of length two, which we immediately split up into the two types of output returned: intermediate outputs from all layers, and final states (of both types) for the last layer.

With just a single layer, layer_outputs[[1]]holds all of the layer’s intermediate outputs, stacked on dimension two.

dim(layer_outputs[[1]])
# [1]  2  4  5 16 16

layer_last_states[[1]]is a list of tensors, the first of which holds the single layer’s final hidden state, and the second, its final cell state.

dim(layer_last_states[[1]][[1]])
# [1]  2  5 16 16
dim(layer_last_states[[1]][[2]])
# [1]  2  5 16 16

For comparison, this is how return values look for a multi-layer architecture.

model <- convlstm(input_dim = 3, hidden_dims = c(5, 5, 1), kernel_sizes = rep(3, 3), n_layers = 3)
c(layer_outputs, layer_last_states) %<-% model(x)

# for each layer, tensor of size (batch_size, seq_len, hidden_size, height, width)
dim(layer_outputs[[1]])
# 2  4  5 16 16
dim(layer_outputs[[3]])
# 2  4  1 16 16

# list of 2 tensors for each layer
str(layer_last_states)
# List of 3
#  $ :List of 2
#   ..$ :Float [1:2, 1:5, 1:16, 1:16]
#   ..$ :Float [1:2, 1:5, 1:16, 1:16]
#  $ :List of 2
#   ..$ :Float [1:2, 1:5, 1:16, 1:16]
#   ..$ :Float [1:2, 1:5, 1:16, 1:16]
#  $ :List of 2
#   ..$ :Float [1:2, 1:1, 1:16, 1:16]
#   ..$ :Float [1:2, 1:1, 1:16, 1:16]

# h, of size (batch_size, hidden_size, height, width)
dim(layer_last_states[[3]][[1]])
# 2  1 16 16

# c, of size (batch_size, hidden_size, height, width)
dim(layer_last_states[[3]][[2]])
# 2  1 16 16

Now we want to sanity-check this module with the simplest-possible dummy data.

Sanity-checking the `convlstm`

We generate black-and-white “movies” of diagonal beams successively translated in space.

Each sequence consists of six time steps, and each beam of six pixels. Just a single sequence is created manually. To create that one sequence, we start from a single beam:

library(torchvision)

beams <- vector(mode = "list", length = 6)
beam <- torch_eye(6) %>% nnf_pad(c(6, 12, 12, 6)) # left, right, top, bottom
beams[[1]] <- beam

Using torch_roll() , we create a pattern where this beam moves up diagonally, and stack the individual tensors along the timesteps dimension.

for (i in 2:6) {
  beams[[i]] <- torch_roll(beam, c(-(i-1),i-1), c(1, 2))
}

init_sequence <- torch_stack(beams, dim = 1)

That’s a single sequence. Thanks to torchvision::transform_random_affine(), we almost effortlessly produce a dataset of a hundred sequences. Moving beams start at random points in the spatial frame, but they all share that upward-diagonal motion.

sequences <- vector(mode = "list", length = 100)
sequences[[1]] <- init_sequence

for (i in 2:100) {
  sequences[[i]] <- transform_random_affine(init_sequence, degrees = 0, translate = c(0.5, 0.5))
}

input <- torch_stack(sequences, dim = 1)

# add channels dimension
input <- input$unsqueeze(3)
dim(input)
# [1] 100   6  1  24  24

That’s it for the raw data. Now we still need a dataset and a dataloader. Of the six time steps, we use the first five as input and try to predict the last one.

dummy_ds <- dataset(
  
  initialize = function(data) {
    self$data <- data
  },
  
  .getitem = function(i) {
    list(x = self$data[i, 1:5, ..], y = self$data[i, 6, ..])
  },
  
  .length = function() {
    nrow(self$data)
  }
)

ds <- dummy_ds(input)
dl <- dataloader(ds, batch_size = 100)

Here is a tiny-ish convLSTM, trained for motion prediction:

model <- convlstm(input_dim = 1, hidden_dims = c(64, 1), kernel_sizes = c(3, 3), n_layers = 2)

optimizer <- optim_adam(model$parameters)

num_epochs <- 100

for (epoch in 1:num_epochs) {
  
  model$train()
  batch_losses <- c()
  
  for (b in enumerate(dl)) {
    
    optimizer$zero_grad()
    
    # last-time-step output from last layer
    preds <- model(b$x)[[2]][[2]][[1]]
  
    loss <- nnf_mse_loss(preds, b$y)
    batch_losses <- c(batch_losses, loss$item())
    
    loss$backward()
    optimizer$step()
  }
  
  if (epoch %% 10 == 0)
    cat(sprintf("\nEpoch %d, training loss:%3f\n", epoch, mean(batch_losses)))
}

Epoch 10, training loss:0.008522

Epoch 20, training loss:0.008079

Epoch 30, training loss:0.006187

Epoch 40, training loss:0.003828

Epoch 50, training loss:0.002322

Epoch 60, training loss:0.001594

Epoch 70, training loss:0.001376

Epoch 80, training loss:0.001258

Epoch 90, training loss:0.001218

Epoch 100, training loss:0.001171

Loss decreases, but that in itself is not a guarantee the model has learned anything. Has it? Let’s inspect its forecast for the very first sequence and see.

For printing, I’m zooming in on the relevant region in the 24x24-pixel frame. Here is the ground truth for time step six:

b$y[1, 1, 6:15, 10:19]

0  0  0  0  0  0  0  0  0  0
0  0  0  0  0  0  0  0  0  0
0  0  1  0  0  0  0  0  0  0
0  0  0  1  0  0  0  0  0  0
0  0  0  0  1  0  0  0  0  0
0  0  0  0  0  1  0  0  0  0
0  0  0  0  0  0  1  0  0  0
0  0  0  0  0  0  0  1  0  0
0  0  0  0  0  0  0  0  0  0
0  0  0  0  0  0  0  0  0  0

And here is the forecast. This does not look bad at all, given there was neither experimentation nor tuning involved.

round(as.matrix(preds[1, 1, 6:15, 10:19]), 2)

       [,1]  [,2]  [,3]  [,4]  [,5]  [,6]  [,7]  [,8]  [,9] [,10]
 [1,]  0.00  0.00  0.00  0.00  0.00  0.00  0.00  0.00  0.00     0
 [2,] -0.02  0.36  0.01  0.06  0.00  0.00  0.00  0.00  0.00     0
 [3,]  0.00 -0.01  0.71  0.01  0.06  0.00  0.00  0.00  0.00     0
 [4,] -0.01  0.04  0.00  0.75  0.01  0.06  0.00  0.00  0.00     0
 [5,]  0.00 -0.01 -0.01 -0.01  0.75  0.01  0.06  0.00  0.00     0
 [6,]  0.00  0.01  0.00 -0.07 -0.01  0.75  0.01  0.06  0.00     0
 [7,]  0.00  0.01 -0.01 -0.01 -0.07 -0.01  0.75  0.01  0.06     0
 [8,]  0.00  0.00  0.01  0.00  0.00 -0.01  0.00  0.71  0.00     0
 [9,]  0.00  0.00  0.00  0.01  0.01  0.00  0.03 -0.01  0.37     0
[10,]  0.00  0.00  0.00  0.00  0.00  0.00 -0.01 -0.01 -0.01     0

This should suffice for a sanity check. If you made it till the end, thanks for your patience! In the best case, you’ll be able to apply this architecture (or a similar one) to your own data – but even if not, I hope you’ve enjoyed learning about torch model coding and/or RNN weirdness ;-)

I, for one, am certainly looking forward to exploring convLSTMs on real-world problems in the near future. Thanks for reading!

Appendix

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

Keras

LSTM

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)

GRU

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

`torch`

LSTM (non-stacked architecture)

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)

GRU (non-stacked architecture)

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

divergences

What are Large Language Models? What are they not?

Large Language Models: What they are

Overall architecture

GPT-type models up close

Input

Core Processing

Output

Model training

Mechanics of text generation

Large Language Models: What they are not

A few facts

It’s us: Language learning, shared goals, and a shared world

It’s us, really: Anthropomorphism unleashed

Group-equivariant neural networks with escnn

Purpose and scope of this post

Using escnn

Spaces, groups, and representations: escnn$gspaces

Representations, for real: escnn$nn$FieldType

Group-equivariant convolution

A group-equivariant neural network

Where to from here

Deep Learning and Scientific Computing with R torch: the book

What’s in the book?

Who’s it for?

What do I get?

Implementing rotation equivariance: Group-equivariant CNN from scratch

Step 1: The symmetry group

Step 2: The lifting convolution

Step 3: Group convolutions

Step 4: Group-equivariant CNN

Rotated digits!

A challenge

Upside down, a cat’s still a cat: Evolving image recognition with Geometric Deep Learning

From alchemy to science: Geometric Deep Learning in two minutes

The “group” in group-equivariance

Groups from symmetries

Viewing groups through the action lens

Outlook: Group-equivariant CNN

AO, NAO, ENSO: A wavelet analysis example

Three oscillations

Analysis: ENSO

Analysis: NAO

Analysis: AO

Conclusion

Getting started with deep learning on graphs

What’s so special about deep learning on graphs?

Matching concepts and code: PyTorch Geometric

A minimal graph

A minimal GNN

Poking into message()

Adding self loops

Customizing aggregate()

Adding memory to update()

Adding parameters

General message passing

A minimal workflow

Forecasting El Niño-Southern Oscillation (ENSO)

El Niño, la Niña

Data

Input: Sea Surface Temperatures

Target: Niño 3.4 Index

Preprocessing

Input

Target

Torch dataset

Dataloaders

Model

Training

Predictions

Discussion

Appendix: convlstm code

Beyond alchemy: A first look at geometric deep learning

Geometric deep learning: An attempt at unification

Geometric priors

Symmetry

Scale separation

From domain priors to algorithmic ones

Domains, priors, architectures

Images and CNNs

Spaces, groups, and representations: `escnn$gspaces`

Representations, for real: `escnn$nn$FieldType`

Poking into `message()`

Customizing `aggregate()`

Adding memory to `update()`

`Torch` dataset

Appendix: `convlstm` code

A `torch` convLSTM

`convLSTM`, the plan

A single step: `convlstm_cell`

Iteration over time steps: `convlstm`

Calling the `convlstm`

Sanity-checking the `convlstm`

`torch`