Flow Equivariant Recurrent Neural Networks

Summary

We thank the reviewer for their time spent with our manuscript. We appreciate that they find our work to be novel in both goals and methods, agreeing with sentiments of the other reviewers.

The reviewer has suggested numerous areas for improvement of the clarity of our work, which we have taken to heart. Below we propose how we aim implement these clarifying modifications to our manuscript.

Clarity

We sincerely thank the reviewer for their suggestions about how to improve the clarity of our submission. We readily admit that in pursuit of technical precision, and to match the conventions of prior equivariance literature, the manuscript has become less accessible to those outside the field. We believe this complication is not strictly necessary, and therefore appreciate the Reviewer encouraging us strive for greater clarity.

To facilitate communication of the main ideas and findings of our work, we plan to change the following:

• We plan to add a visualization of the difference between static equivariance and flow equivariance to the main text.
• We plan to add an additional visualization of the representation of the flow action in the FERNN latent space for the simple bump example (visualizing the result of the main Theorem, 4.1).
• We plan to add visualizations of data flows to the main text (on KTH), for a more immediate conceptual picture of what a flow symmetry is.
• Finally, as reviewer vWda has suggested, our detailed review of G-equivariance can be shortened and moved to the appendix to reduce extended mathematical complication of the main text.

We additionally plan to release a blog post accompanying the paper which will have the following clarifying properties:

• a motivational introduction to the 'flow equivariance problem'.
• a direct visual and motivational comparison between static equivariance and flow equivariance.
• an extension of the 'bump' example which runs throughout the blog post, showing the difference between the G-RNN and FERNN, and the visualization of the flow action.

We propose to include a link to this blog post in the paper to assist those who stumble on the paper but prefer a gentler introduction.

Questions

• The reviewer is correct that $f$ is missing in line 84. It should be: $(g \cdot \Phi [f])(y) := \rho_Y(g) \Phi [f] (\pi_Y(g^{-1}) y)$. We appreciate them mentioning this error.
• We appreciate the reviewer bringing up where we missed defining indices $i$ and $k$ in Equation 3. These are the output and input channels for the convolution respectively, following the original group convolution notation of Cohen & Welling (2016). We have made this explicit in the camera ready draft, and have further made explicit the domains of the lifting and group convolutions, as the reviewer noted this may have been ill defined.
• With respect to output mappings, this choice is typically left up to the modeler, and is not part of traditional equivariance frameworks. Again we defer to the formulation of Cohen & Welling (2016), of which our Equation 3 is merely a restatement.
• In lines 130-131 the reviewer is correct that $f$ is again missing. It should be: $(\psi (\nu) \cdot \Phi [f] )_t(y) := \Phi_t [f] ( \psi_t ( \nu )^{-1} \cdot y)$. We greatly appreciate the reviewer catching these notational inconsistencies and we have amended them in the camera ready draft.

We warmly welcome discussion on any further points the reviewer may have in regards to how we may improve the clarity of the text.

Summary

We sincerely thank the Reviewer scMK for their time spent with our work, and for greatly helping us to improve the manuscript. We appreciate that they commend the originality, clarity and exposition of our work.

In our response, we directly address the Reviewer's primary concern about the practical applicability of our method. This concern appears to be due to the perception that the benefits of our proposed model are limited to spatially global flows, constant speed flows, and flows which exactly match the model's small discrete set of generators.

Through a suite of new results, we demonstrate that the practical benefits of FERNNs are not so limited, even if the theory does not yet extend this far. Explicitly, we present positive empirical results on:

1. Non-global flows,
2. Variable-velocity flows,
3. Data flows landing between the model's generator set.

In each of these cases, we demonstrate FERNNs still excel over non-flow-equivariant baselines.

We subsequently include a quantitative analysis of the computational overhead of flow-equivariance, both in terms of memory and compute time, and how this scales with the size of the generator set $V$ (equivalent to a notion of the 'amount of flow equivariance').

Finally, we conclude with a comparison of the limitations raised by the Reviewer and those present in the original Group Equivariant Convolutional Neural Networks (G-CNNs) of Cohen & Welling (2016). We highlight that G-CNNs have delivered remarkable practical value since their introduction (as noted by Reviewer vWda) despite very similar limitations, and thus suggest that our work may be interpreted analogously as the introduction of a new foundational framework for sequence-model equivariance, opening the door for future work to address the theoretical limitations.

Non-Global Flows

To demonstrate that the benefits of FERNNs extend to non-spatially global flows, we provide new results on a 2-Digit Flowing MNIST task. The task is identical to that in the main paper, except that 2 random digits are simultaneously placed on the canvas, and each digit is 'flowed' independently. Visually, this looks like two digits moving in two separate random directions (for translation), or two overlapping digits rotating at different angular velocities (for rotation). The full image sequence is thus not describable by a single global flow (in the formal mathematical sense).

As can be seen, FERNNs still achieve roughly an order of magnitude lower forward prediction error than their non-flow-equivariant counterparts. This strong performance is explained by the fact that while each $\nu$ channel flows globally, each pixel is simultaneously represented in the full lifted space of all $\nu \in V$. Therefore, the readout from the hidden state (in our case, max-pooling) can select an independent velocity channel to represent each spatial location.

Variable-Velocity Flows

The reviewer is correct to note that constant speed is inherent to the mathematical definition of a flow: flows exponentiate a single constant generator element of the Lie algebra. This property is in fact necessary to get the flow composition property described on line 125 ($\psi_s(\nu) \cdot \psi_t(\nu) = \psi_{s+t}(\nu)$), and this flow composition property is essential to the equivariance proof of the FERNN (see line 230, & Eqn. 83 of Supplement).

However, even though the theory only extends to constant velocity flows, in practice, we observe that sequences with variable velocity flows are still better modeled by FERNNs than non-flow-equivariant counterparts. We provide tables below demonstrating this for two non-constant-flow datasets:

'Bouncing MNIST', where 2-digits bounce off each other and the image boundary.

KTH with a random velocity switch halfway through the sequence.

We see that on Bouncing MNIST, the FERNN model nearly halves the error, and the addition of velocity mixing features further improves the performance. We refer the Reviewer to our reply to Reviewer vWda for an extended discussion $V$-mixing and the dataset details. We similarly see on the variable velocity KTH dataset, FERNNs again significantly outperform their non-flow-equivariant counterparts.

Handling Data Flows Outside the Model's Generator Set $V$

Finally, to demonstrate that FERNNs yield improved performance even when data 'lands between' the elements of the generator set, we generate a Translating MNIST dataset with x-velocity flows from -5 to 5 pix/step (0 y-velocity), and train models with equivariance to subsets of these data-generating flows.

In the table below, we report the test set MSE (mean over 3 seeds) computed for each x-velocity separately (Avg. at end). We see that any flow equivariance improves the average loss over non-flow-equivariant models, and further, when the data lands between the elements of the generator set (bolded), the FERNNs always achieve lower error than the G-RNN, sometimes by up to an order of magnitude, despite not being analytically equivariant to these flows. In the camera ready draft, we plan to include a figure of this data to improve readability.

Computational Overhead

Finally, we study the computational overhead of flow equivariance. Theoretically, both compute and memory should scale approximately linearly with the number of flow generators (denoted $\|V\|$). In the tables below, we validate this for FERNNs with different sized generator sets. Note the G-RNN is equivalent to a FERNN with $\|V\| = 1$, i.e. $V =${$0$}.

If we fit a linear model to these results, we obtain extremely good fits, and the following parameters:

We note that our current FERNN implementation is extremely computationally inefficient (looping over each flow generator), while a scan-based optimized implementation of this model is certainly possible and would dramatically reduce runtime and improve its scalability.

Conclusion

In conclusion, we thank the reviewer for encouraging these new experiments. We believe many readers are likely to share the same concerns, so we are grateful for the opportunity to demonstrate empirically that these limitations are likely not as severe as theory alone might suggest.

Interestingly, the limitations brought up by the reviewer are nearly identical to those of the original Group Equivariant Neural Networks of Cohen & Welling (2016) -- namely, scalability, discrete group equivariance and the restriction to global transformations. With time, however, researchers have developed methods to ameliorate these limitations: such as steerable equivariance (Weiler et al. 2019), addressing the discrete approximation limitation, and simultaneously scalability; and gauge equivariance (Cohen et al. 2019), addressing the local vs. global transformation limitation. In doing so, as Reviewer vWda has pointed out, group-equivariance has become incredibly practically relevant, including being used in highly popularized models such as AlphaFold 2.

We therefore suggest that our paper may be seen as a the first step in the introduction of a new theoretical foundation -- namely a foundation for dynamic equivariance in sequence models that opens the door for future work. We welcome further discussion.

Summary

We sincerely thank Reviewer ngxM for their strong support of our work. We appreciate that they acknowledge the timeliness and importance of introducing equivariance for sequence models at this point in time. We are very grateful that they find our analysis and exposition exceptional, and that they find our claims to be well supported by our experiments.

In our response, we will address the two weaknesses listed by the reviewer: namely, simplicity of the modeled data flows, and the computational complexity of the FERNN model.

Specifically, we first highlight a number of more complex flow prediction problems that we have tested FERNNs on since our submission of the original manuscript. Secondly, we clarify the computational requirements of FERNNs both theoretically and empirically.

More Complex Flow Prediction

Below we present results for three new sequence modeling tasks with more complex flow symmetries as part of the data generating process. We see in each case that the flow equivariant models maintain their advantage over non-flow-equivariant counterparts.

2-Digit Independently Flowing MNIST

The first dataset we explore is a 2-digit variant of the flowing MNIST task. The task is identical to that in the main paper, except that 2 random digits are simultaneously placed on the canvas, and each digit is 'flowed' according to its own independent randomly sampled flow generator. Visually, this looks like two digits moving in two separate random directions (for translation), or two overlapping digits rotating at different angular velocities (for rotation). Below we report the forward prediction test set MSE in the same setting as in the main paper (10 steps of input, forward predict the next 10 steps).

Although it is not apparent from this table, we have also explored the length generalization and velocity generalization of the FERNN on this 2-digit MNIST task and we find that it performs equally as well as for the 1-digit task, demonstrating that all of the benefits of flow equivariance appear to function as expected in this more complex flow setting.

While still a toy dataset, this modification represents a significant new ability of the model -- namely, the ability to model multiple input flows simultaneously. One can imagine this as a highly simplified version of the visual transformations that would occur when one simultaneously combines camera motion with the independent motion of the video subject: while the camera motion may induce a global flow in one direction, the subject is free to move in an entirely independent direction (e.g. a camera pans to the left while a basketball player runs up on the screen). We are pleased to show that FERNNs can handle this type of scenario.

KTH with a random velocity switch halfway through the sequence

The second more complex flow type we explore is non-constant-velocity flows. Specifically we explore the ability of FERNNs to model sequences where the flow velocity switches to a new random velocity mid-way through the sequence. We test this on the KTH action recognition dataset by modifying our simulated 'camera motion' with such a switch. We again see that FERNNs outperform their non-flow-equivariant counterparts significantly.

'Bouncing MNIST', where 2-digits bounce off each other and the image boundary

The third and perhaps most complex dataset we experiment with is a 'bouncing' variant of the flowing MNIST dataset. Specifically, we construct this dataset identically to the Translating MNIST dataset from the manuscript, but with 2 digits which now have elastic collisions with the image boundary and with each-other. This causes digits to change velocity at multiple points during the sequence as a function of their own velocity, and the velocity of the other objects that they collide with. Thus to forward predict these sequences, models must be good at modeling interactions between velocities.

In order to have space for both digits and the boundary, we increase the canvas size to 28x48. We additionally draw tight bounding boxes around the digits (in white, value 1.0) and draw a one-pixel image boundary on the canvas (in gray, value 0.5). In the table below we report the results of the FERNN and the G-RNN baseline on this dataset. We additionally report the results of a FERNN with convolutional kernels that are non-zero for the cross-$V$ terms (denoted $V$-Mixing). We refer the Reviewer to our discussion with Reviewer vWda for a more in-depth discussion of the role of $V$-mixing, and why it helps here.

Ultimately, we see that FERNNs again outperform the non-flow-equivariant counterparts on this dataset, despite having no guarantee of equivariance with respect to these highly complex flow transformations. We find the fact that $V$-mixing improves modeling of object interactions to be very exciting, and a direction certainly worth exploring further in future work.

Computational Requirements of FERNN layers

Finally, we address the computational complexity of the FERNN layer. The flow convolution in Equation 9 performs a convolution over standard group elements ($g \in G$, which can be thought of as spatial positions for the simple translation case), and also over flow generators ($\nu \in V$, which can be thought of as movement velocities). This can be implemented efficiently as a 3D convolution for the case of 2D images, and a 1D generator group -- (as we do in the above model denoted $V$-mixing). For sets with higher dimensional generators, one can use N-D convolutions, which are also implemented efficiently in frameworks like Jax.

The computational complexity of the model scales linearly with the size of the set of generators $V$ (which we denote $|V|$). Below we report the memory requirements and runtime per epoch as a function of the size of this set. Note the G-RNN is equivalent to a FERNN with $|V| = 1$, i.e. $V =$ {$0$}.

If we fit a linear model to these results, we obtain extremely good fits, and the following parameters:

We note that our current FERNN implementation is extremely computationally inefficient (looping over each flow generator), while a scan-based optimized implementation of this model is certainly possible and would dramatically reduce runtime and improve its scalability. In our limitations and future work section, we describe how something like this may be possible within the State Space Model (SSM) framework.

Conclusion

We thank the reviewer for their engagement with our work and strong support. We welcome further discussion of any of the above points.

Summary

We thank Reviewer vWda for their comments, suggestions, and critical analysis of our work. We appreciate that they find our work original, technically sound, and that it addresses an under explored topic. We also appreciate that they acknowledge the connection of our work to Group Convolution [Cohen & Welling, 2016], supporting our belief that FERNNs lay the theoretical groundwork for exploiting a new type of time-dependent symmetric structure in data, just as Group Convolution did for 'static' symmetric structure.

Below, we reply to a few of the reviewer's comments and include results that directly address their suggested improvements regarding:

• Mixing across velocity ($V$) channels
• Data augmentation baselines
• Realistic non-constant camera motion

As a high level summary of our response:

1. We provide new results with mixing across $V$ channels for both the original MNIST data (where we see no improvement) and a new non-constant velocity 'bouncing MNIST' dataset (where we see improvement), and explain these results with a mechanistic hypothesis for the functional role of $V$-mixing.
2. We clarify that our baselines on MNIST can already be understood as 'data augmented' from our dataset construction procedure. We provide new results on the KTH action recognition dataset with proper data augmentation and still find the FERNN performs best.
3. We provide new results on both MNIST and KTH with non-constant velocity flows. We see that the FERNN still outperforms non-flow-equivariant baselines, despite the mismatch with theory, supporting the practical relevance of our approach.

Mixing Across Velocity ($V$) Channels

We thank the reviewer for highlighting the importance of mixing across $V$ channels, and encouraging us to add these experiments; we find the results quite illuminating.

In the table below we show the performance (Test MSE $\pm$ std. over 3 seeds) of the G-RNN, original FERNN, and FERNN with $V$-mixing (as described in Eqn. 9, with a convolutional kernel of size 3 in the $V$ dimension). We find that on the original Translating MNIST dataset (first column), the results are not significantly different with $V$-Mixing added, and the model actually performs measurably worse. As mentioned in the main text (lines 248-255), the reason for not including $V$-mixing in our initial results is not due to computational limitations, but rather a conscious design choice to limit the propagation of equivariance error arising from the truncation of the infinite velocity group to a finite range. This is similar to what was found in prior work on scale-equivariance [Worall and Welling, 2019, Sosnovik et al., 2020b], and is likely the reason for the drop in performance.

However, we believe we can also understand why $V$-mixing does not improve performance significantly here -- namely, we believe $V$-mixing is most beneficial when there are dynamical computations that require the interaction of multiple different velocity features simultaneously to predict the next frame. The constant velocity global flows in Translating MNIST require no such features.

To test this mechanistic hypothesis for the role of $V$-mixing, we provide new results on a dataset we call 'bouncing MNIST', identical to translating MNIST, but where MNIST digits now have elastic collisions with the image boundary and with each-other (we have now added two digits per image, and increased the image size to 28x48 correspondingly). These collisions require exactly the types of muli-velocity computations we believe $V$-mixing would help with. Experimentally, we validate this hypothesis in the second column of the table below, where the FERNN with $V$-Mixing achieves the lowest test error on this bouncing MNIST dataset.

Data Augmentation

With respect to data augmentation, we appreciate the reviewer's suggestion; we include a general clarification of our data generation procedure below, followed by new results for the KTH data (both of which we will add to the main text).

Clarification: Since the datasets we use in our paper have no natural motion, they are already 'augmented' in some way in order to generate the training sequences. The main design choice is then whether the velocities ('motions') are sampled once for each image and fixed throughout training, or if velocities are re-sampled at each epoch (each time the image is presented). This second option, we call 're-sampling', is more akin to standard data augmentation.

MNIST results are already 'data augmented' according to the 're-sampling' definition: We highlight that the MNIST results in our paper are already 'data augmented' according to the 'resampling' definition above. This is because the velocities are randomly sampled for each training example in the data loader by default, meaning each digit is observed with a different velocity at each epoch, equivalent to data augmentation. This is described briefly on lines 270-272, and more thoroughly in Supplementary Material Section B.2. The only setting where this augmentation is not applied is on the velocity generalization tasks, where it would not make sense by definition of zero-shot generalization.

KTH results are not 're-sampling' data augmented, new results below: The KTH (action recognition) results in Table 2 are not with 're-sampled' augmentations for each epoch, instead they are sampled once (with a fixed seed for all models) at the beginning of training. We appreciate the reviewer bringing attention to this distinction. In the table below we show the new results for Table 2 replicated with 're-sampled' data augmentation for each model. We see that the performance of all models improves, and the performance gap between models reduces, although the FERNN still outperforms the baselines by a narrow margin.

In the camera ready version, we propose to include 're-sampled' data augmentation baselines in Figure 5 as well. We also propose to include non-augmented baselines on MNIST for completeness.

'Realistic' Non-Constant Camera Motion

The reviewer comments that natural camera motion is not a constant velocity, and suggests that this may make the method less applicable to real-world data. We agree, and further believe that the study of the performance of FERNNs on non-constant velocity flows is certainly of interest to us as well. As a preliminary step in this direction, we first appeal to the results above on 'Bouncing MNIST' where we see that the FERNN still achieves a significant performance improvement over non-flow-equivariant counterparts, despite the variable velocity flows.

Secondly, we provide new results on the KTH dataset where we simulate camera motion that changes its velocity vector randomly half-way through the sequence. We see in the table below that the FERNN still significantly outperforms the baselines.

In future work, we believe that it may be possible to train a 'flow predictor' to estimate the current flow between any two time-steps, and apply only this flow to this hidden state. This would avoid the 'exponential growth' problem that the reviewer mentions, while still maintaining equivariance. We leave this as a potential future direction which is motivated and supported by the theory of our proposed work.

Response to minor points:

• The reviewer requested pictures of the action recognition dataset with camera motion applied. We bring to the reviewer's attention that these are indeed already in Section B.7, Figure 6, of the supplementary material. We will move these to the main paper for the camera ready.
• In our camera ready draft we have split Figure 4 into two separate figures and removed empty space, significantly improving readability, and addressing the Reviewer's minor formatting concerns.
• In our camera ready we have expanded on the limitations of our model with respect to variable velocity flows, data augmentation, and computational scaling.
• We sincerely appreciate the authors suggestion to help simplify our exposition by reducing background on G-Equivariance. We agree that the paper could be simplified to improve the exposition, and reducing background mathematical clutter is a good option to do so. We do however believe we need to include a minimal example of G-equivariance in the text since it is explicitly required as a subset of a flow-equivariant network, i.e. it is not possible to build a flow-equivariant model which is not also G-equivariant.

We thank the reviewer again for their suggestions and we welcome further discussion if we may have misunderstood any of their comments.