Source-linked AI summary
Flow Equivariant Recurrent Neural Networks
T. Anderson Keller
TL;DR
Equivariant sequence models have largely been studied for static transformations, leaving time-parameterized symmetries less addressed. This paper defines flow equivariance and constructs recurrent models that satisfy it. The authors report generalization to unseen flows and improved length generalization and performance on datasets with strong flow symmetries.
Problem
Equivariance has largely been limited to static transformations, leaving time-parameterized symmetries less addressed in sequence models.
Source: §1 · p. 1Method
The paper introduces flow equivariance and constructs recurrent networks by lifting hidden states across flow generators with shared weights.
Source: §1 · pp. 1–2Results
The authors report zero-shot generalization to unseen flows, improved length generalization, and improved performance on datasets with strong flow symmetries.
Source: §1 · p. 2Takeaways & Limitations
The work provides an initial framework for constructing sequence models equivariant to time-parameterized group transformations.
Source: §6 · p. 9FERNNs are formally restricted to constant-velocity transformations, and their flow-indexed activations, compute, and memory grow linearly with the size of the flow set; truncation also induces errors.
Source: §6 · p. 10
Abstract
from arXiv · showhide
Data arrives at our senses as a continuous stream, smoothly transforming from one instant to the next. These smooth transformations can be viewed as continuous symmetries of the environment that we inhabit, defining equivalence relations between stimuli over time. In machine learning, neural network architectures that respect symmetries of their data are called equivariant and have provable benefits in terms of generalization ability and sample efficiency. To date, however, equivariance has been considered only for static transformations and feed-forward networks, limiting its applicability to sequence models, such as recurrent neural networks (RNNs), and corresponding time-parameterized sequence transformations. In this work, we extend equivariant network theory to this regime of 'flows' -- one-parameter Lie subgroups capturing natural transformations over time, such as visual motion. We begin by showing that standard RNNs are generally not flow equivariant: their hidden states fail to transform in a geometrically structured manner for moving stimuli. We then show how flow equivariance can be introduced, and demonstrate that these models significantly outperform their non-equivariant counterparts in terms of training speed, length generalization, and velocity generalization, on both next step prediction and sequence classification. We present this work as a first step towards building sequence models that respect the time-parameterized symmetries which govern the world around us.
1 Introduction
Equivariance has been used to encode data symmetries, but sequence models have largely addressed static transformations rather than continuous motion. The paper defines flow equivariance and introduces recurrent models designed to respect time-parameterized symmetries.
1 Introduction
- Flow equivariance extends equivariance to sequence models under time-parameterized transformations, a regime largely missed by work focused on temporally static transformations.The paper considers flows, formally one-parameter Lie subgroups.Source: §1 · p. 1
- The authors introduce recurrent neural networks explicitly equivariant to time-parameterized symmetries by lifting hidden states across flows indexed by generating vector fields.Weight sharing across the flow-indexed bank makes it a single group-equivariant dynamical system.Source: §1 · p. 2
- The paper reports zero-shot generalization to unseen flows, improved length generalization, and improved performance on datasets with strong flow symmetries.The paper also previews next-step prediction and sequence classification experiments.Source: §1 · p. 2
2 Equivariant Neural Networks
Equivariance requires input transformations to correspond to structured output transformations. Group convolutions enable this for static group actions, and group-equivariant recurrent networks preserve that property for static transformations of an input sequence.
2 Equivariant Neural Networks
- Equivariance means that transforming an input corresponds to a structured transformation of the output, with group structure supporting transformations of network activations.Group closure and associativity underpin the account of equivariant activations as living on the group.Source: §2 · p. 2
- Group convolution lifts signals to functions on the group, where convolution commutes with the group action.The paper describes an initial lifting convolution followed by group convolutions.Source: §2 · p. 3
- A G-RNN replaces its usual input and recurrent maps with group convolutions, with an invariant hidden-state initialization and an equivariant pointwise nonlinearity.The paper uses discrete time, while noting that continuous-time sequence domains require only minor analytical changes.Source: §2 · p. 3
- Group-equivariant recurrent and input maps preserve equivariance to a static transformation of the entire input sequence, but this does not establish equivariance to motion.Static transformations can be pulled through both maps and the nonlinearity to the hidden state.Source: §2 · p. 3
3 Flow Equivariance
Flow equivariance requires a sequence model’s output to transform consistently when its input undergoes a time-parameterized symmetry. Frame-wise equivariant processing satisfies this trivially, but recurrent G-RNNs generally do not.
3 Flow Equivariance
- A flow is a time-parameterized symmetry generated by a Lie algebra element, and flow equivariance requires the sequence model’s output transformation to commute with transforming its input.The definition applies to flows generated by a set of generators.Source: §3 · pp. 3–4
- Frame-wise equivariant maps are trivially flow equivariant, but this construction achieves no sequence modeling because it removes recurrence.A G-RNN with zero recurrent weights reduces to this frame-wise case.Source: §3 · p. 4
- A G-RNN with nonzero recurrent weights is not flow equivariant except in the degenerate flow-invariant case.The counterexample uses additive recurrence: a static bump becomes a growing bump, while a moving input produces a lagging train of bumps.Source: §3 · p. 4
- The counterexample is stated to apply to most other existing sequence models, including state-space models, gated recurrent networks, neural ODEs, and transformers.The authors say the proposed RNN construction can likely carry over to other sequence models.Source: §3 · p. 4
4 Flow Equivariant Recurrent Neural Networks
FERNNs lift hidden states over flow generators and perform recurrent updates in dynamically transformed reference frames. Under stated initialization conditions, this construction is flow equivariant, with practical implementations using a truncated flow group.
4 Flow Equivariant Recurrent Neural Networks
- FERNNs augment the hidden state with a flow-generator dimension and apply the instantaneous flow element during each recurrent update.This lets each flow-indexed state evolve according to its corresponding vector field.Source: §4 · pp. 4–5
- The input is lifted to the product space by copying the group-convolution output across flow indices, while the hidden state is defined over the product of flow and symmetry groups.The paper uses a trivial input lift in its experiments and proofs, while also describing a non-trivial alternative.Source: §4 · p. 5
- With hidden-state initialization invariant to the group action and constant across flow indices, FERNNs satisfy the paper’s flow-equivariance condition.The transformed-input hidden state corresponds to a shift along the flow index and an action along the group dimension.Source: §4 · p. 5
- The flow-indexed bank lets an input’s velocity match a hidden-state channel, shifting the moving coordinate reference frame to that flow.The model’s hidden-state action then includes corresponding relative shifts across the other flow channels.Source: §4 · p. 6
- Exact equivariance requires a flow group closed under the output action, but practical models truncate the group and incur boundary errors.Most experiments do not mix flow channels, to avoid propagating truncation-induced errors; flow mixing can model inter-velocity interactions.Source: §4 · p. 6
5 Experiments
Experiments compare flow-equivariant recurrent networks with non-equivariant baselines on motion-based prediction and classification tasks. FERNNs show stronger training speed, length and velocity generalization, and performance on moving action-recognition data.
Moving KTH action recognition
- Moving KTH experiments add translation flows to simulate camera motion; unlike MNIST, velocities are not resampled for each example..Source: §5 · pp. 6, 9
Next-step prediction
- FERNNs reached nearly an order of magnitude lower error than G-RNNs on both Flowing MNIST datasets and did so in significantly fewer iterations; partial equivariance also helped.The validation-loss plot tracks training steps, and Table 1 reports test MSE.Source: Table 1 · p. 7, Figure 3 · p. 7, §5 · p. 7
Length generalization
- On Translating MNIST, FERNN predictions remain consistent beyond the 10-step training horizon while G-RNNs degrade; gains are weaker for Rotating MNIST.The evaluation extends forward prediction to 70 steps, beyond the 10 future steps used in training.Source: §5 · pp. 7–8, Figure 4 · p. 8
Velocity generalization
- FERNNs predict nearly perfectly at unseen translation and rotation velocities within V, and interpolate outside V with slightly increased error and improved performance across a wider velocity range.Figure 5 compares predictions across test velocities; Figure 6 reports test MSE for varying V-set coverage.Source: Figure 5 · p. 8, §5 · p. 8, Figure 6 · p. 8
Moving KTH action recognition
- On Moving KTH, FERNNs outperform non-flow-equivariant models on actions under out-of-distribution flows and significantly outperform them on strongly moving data without a train-test disparity.Table 2 reports test accuracy for models trained and tested with V T 2.Source: Table 2 · p. 9, Figure 8 · p. 9, §5 · p. 9
Action-recognition models
- The G-RNN+ ablation keeps FERNN's lifted hidden state and max-pooling but replaces its flow action with a learned convolution, controlling for those architectural additions.The authors use this comparison to argue that FERNN's gains are not simply due to increased hidden-state size or an extra max-pooling nonlinearity.Source: §5 · p. 9
6 Discussion
The discussion frames flow equivariance as a foundation for sequence models, while the appendix establishes component-level equivariance results. It situates the approach alongside prior work and notes its limitations and connections to neural dynamics.
6 Discussion
- Flow equivariance is presented as a new framework for time-parameterized symmetries in sequence models, with a causal, model-general scope distinct from prior spacetime-convolution approaches.The discussion contrasts it with prior recurrent work limited to static coordinate transformations and with hand-crafted equivariant filters lacking a familiar training framework.Source: §6 · pp. 9–10
- The formal construction is limited to constant-velocity transformations, and activations, compute, and memory grow linearly with the size of V.The authors note that inter-V flow-convolution kernels can model variable velocities in practice, despite the formal restriction.Source: §6 · p. 10
- The discussion connects flow equivariance to traveling-wave-like neural dynamics and proposes further empirical investigation of this possible relationship.It argues that accurate, stable equivariant representation of a flowing stimulus requires hidden-state dynamics that autonomously realize a homomorphic representation of that flow.Source: §6 · p. 10
A.1 Equivariance of Group-Lifting Convolution & Group-Convolution
- Group-lifting and group-convolution layers satisfy equivariance: transforming the input corresponds to the matching transformation of the output.The appendix re-derives both convolution identities in the paper's notation.Source: §A.1 · p. 17
A.2 Equivariance of Group Equivariant RNN
- A group-equivariant RNN preserves equivariance across timesteps when its latent group action is linear, its nonlinearity is equivariant, and its initialization is invariant.The proof uses induction and assumes zero input before time zero.Source: §A.2 · p. 18
A.3 Frame-wise Flow Equivariance
- A framewise map equivariant to each flow-generated group element yields a flow-equivariant sequence model when applied independently at every timestep.The sequence model concatenates the per-timestep outputs of the map.Source: §A.3 · p. 18
A.4 Group Equivariant RNNs are not Flow Equivariant
The appendix shows that ordinary group-equivariant RNNs generally fail flow equivariance, then verifies flow equivariance for the lifting and flow-convolution components.
A.4 Group Equivariant RNNs are not Flow Equivariant
- A group-equivariant RNN with nonzero recurrent weights is not generally flow equivariant; the degenerate constant-kernel case is flow invariant.The proof shows that equivariance for nonconstant kernels would require the recurrent convolution output to be constant along all flows, contradicting the nonconstant-kernel assumption.Source: §A.4 · pp. 19–20
A.5 Equivariance of Lifting Flow Convolution
- The flow-lifting convolution is verified to be equivariant under the action of individual flow elements.The appendix treats it as a trivial lift of the group-lifting convolution with an additional duplicated index.Source: §A.5 · p. 21
A.6 Equivariance of Flow Convolution
- The appendix proves that the flow convolution is equivariant to the action of individual flow elements.This is the stated result of the flow-convolution proof.Source: §A.6 · p. 21
A.7 Equivariance of FERNN
The appendix proves FERNN flow equivariance under specified hidden-state initialization conditions and illustrates how its flow-indexed channels respond to moving inputs. It also describes the Flowing MNIST data and model setups.
A.7 Equivariance of FERNN
- FERNNs are flow equivariant when the hidden-state initialization is invariant to the group action and constant across the flow dimension.The induced output action shifts the flow index by the input flow and applies the corresponding group action; the proof proceeds by induction.Source: §A.7 · pp. 21–22
- For a matching input motion, the corresponding hidden-state channel processes the input in its own reference frame, and the flow action permutes the channels.The visualization shows separate V channels flowing at their respective ν velocities.Source: §A.7 · p. 22, Figure 9 · p. 23
B.1 Flowing MNIST: Dataset Creation
- Flowing MNIST sequences are generated by applying randomly selected admissible flows to MNIST samples; training and test sequences are typically 20 steps, with longer tests for length generalization.The set of flows defining FERNN equivariance need not match the dataset's training-flow set.Source: §B.1 · p. 24
B.2 Flowing MNIST: Models
- The Flowing MNIST models use matched trainable parameter counts, while FERNNs add a flow-indexed hidden-state dimension and pool across it before decoding.The translation FERNN uses an otherwise matching RNN and decoder architecture.Source: §B.2 · pp. 25–26
B.3 Flowing MNIST: Next-Step Prediction Training & Evaluation
In the velocity-generalization evaluation, FERNNs perform significantly better on test flows when made equivariant to those velocities. The appendix describes the sequence lengths, training procedure, and evaluation splits used in these experiments.
B.3 Flowing MNIST: Next-Step Prediction Training & Evaluation
- For next-step prediction, models receive the first 10 steps of 20-step training sequences and are evaluated on reconstruction loss over the remaining 10 steps.Models are trained for 50 epochs; reported test performance is from the validation-best model, averaged across five random initializations.Source: §B.3 · p. 26
- Length generalization is evaluated by testing on 70-step sequences after training and validation on 20-step sequences, using only the first 10 test steps as input.Loss is reported for each of the remaining 60 test steps.Source: §B.3 · p. 26
- FERNNs perform significantly better on unseen test flows when made equivariant to those velocities.The velocity-generalization setup uses training flows that are a subset of the test flows.Source: §B.3 · pp. 26–27
B.4 Moving KTH: Datasets
The appendix describes the Moving KTH dataset variants, model architectures, and training and evaluation procedure. It compares recurrent and convolutional baselines with flow-equivariant models, including an ablation of the FERNN recurrence.
B.4 Moving KTH: Datasets
- Moving KTH uses 16-frame clips from 32-frame KTH segments, downsampled to 32 × 32 and split by participant into training, validation, and test sets.The original stationary-camera dataset is supplemented by two variants with translation flows to simulate a moving viewpoint.Source: §B.4 · p. 27
B.5 Moving KTH: Models
- The five compared models are a 3D-CNN, a G-RNN, two FERNN variants, and G-RNN+, an ablation that replaces the flow action with learned velocity-specific 5 × 5 convolutions.The FERNN variants lift velocity channels and implement the flow action by rolling spatial dimensions; their flow-convolution weights are shared across velocities.Source: §B.5 · pp. 27–28
B.6 Moving KTH: Action Recognition Training & Evaluation
- Models are trained for 500 epochs with three random seeds per learning-rate setting, and reported test loss is averaged across the three selected runs.The selected learning rate is based on mean best validation loss across seeds; lower rates generally suit RNNs, while higher rates suit CNNs.Source: §B.6 · p. 28
B.8 Compute
These appendix sections report FERNN compute and scaling characteristics, describe a non-trivial lift and related approaches, and provide additional prediction visualizations. The visualizations show strong length generalization on rotation and translation despite weaker rotation results than on translation.
B.8 Compute
- Roughly 67 hours for 50 epochs: rotation FERNN-VR4 training was slowed by inefficient rotation and recurrence implementations.The recurrence used a naive PyTorch implementation with for-loops; the authors suggest scan-based JAX or a custom CUDA kernel could improve runtime.Source: §B.8 · pp. 28–29
B.9 Runtime and Memory Usage vs. |V |
- Runtime and memory scale linearly with generator-set size |V| in the Translating MNIST experiment restricted to x-direction translation.The authors identify the G-RNN as the |V| = 1 case.Source: Figure 11 · p. 29, §B.9 · p. 29
C.1 FERNN with Non-Trivial Lift
- With a non-trivial lift, input features flow along the velocity index, so the recurrence removes the additional ψ shift and the output action permutes velocity channels without shifting spatial position.The authors note that this construction makes kernel indexing explicitly time-dependent, unlike the trivial lift.Source: §C.1 · p. 30
D.1 Flow Equivariance without a Hidden State; D Related Work
- Frame-wise equivariant models and space-time group convolutions lack recurrent hidden states and have finite temporal context, whereas recurrent models can theoretically maintain unbounded context.The appendix distinguishes these approaches from flow-equivariant RNNs, which extend equivariant models to online recurrence.Source: §D.1 · p. 31
D.3 Neuroscience and Biologically Inspired Neural Networks; D.6 Equivariant Dynamical Systems
- FERNNs offer a theoretical lens for studying convolutional recurrent models in neuroscience, and differ from control-observer work by learning dynamics rather than relying on handwritten models.The contrast with control-observer methods is that FERNNs lift hidden states to a group-indexed field and enforce equivariance through shared group-convolution weights.Source: §D.3 · p. 32, §D.6 · p. 33
E.1 In-Distribution Next-Step Prediction (Table 1 & Figure 3)
- In-distribution predictions on Translating and Rotating MNIST appear visually strong for all models.The appendix evaluates models on the same flows used for training.Source: Figure 12 · p. 34, §E.1 · p. 34
E.2 Length Generalization MSE vs. Forward Prediction Plot (Rotation); E.3 Length Generalization Visualizations (Translation)
- On rotation, FERNN-VR4 strongly generalizes up to 30 steps and outperforms the non-equivariant model; on translation, FERNN forecasts remain consistent beyond training length while G-RNN predictions degrade.The rotation length-generalization gap is less pronounced than for translation, which the authors suspect reflects accumulated interpolation errors from small-angle rotations on a discrete grid.Source: Figure 13 · p. 35, §E.2 · p. 35, Figure 14 · p. 36
E.4 Velocity Generalization Visualizations
FERNNs generalize visually to higher test velocities than those seen in training, while their formal equivariance theory assumes constant velocity. In a Moving KTH setting with an unpredictable velocity switch, FERNNs still outperform non-flow-equivariant baselines.
E.4 Velocity Generalization Visualizations
- FERNNs trained on lower velocities generalize to higher unseen velocities in translation and rotation, while non-flow-equivariant RNNs fail on the shown examples.The examples test Rotating MNIST at ±40 degrees per step and Translating MNIST at ν = (±2, ±2).Source: Figure 15 · p. 37, §E.4 · p. 37
E.5 Flow Equivariant RNNs and Variable Velocity Flows; E.5.1 Variable Velocity Moving KTH Action Recognition
- The formal equivariance theory requires constant-velocity flows, but FERNNs empirically model variable-velocity sequences better than non-equivariant counterparts.The appendix tests Moving KTH sequences with a random, unpredictable velocity switch halfway through each clip; FERNNs still outperform the baselines by a significant margin.Source: §E.5 · p. 38, §E.5.1 · p. 38
E.5.2 Mixing Across the V Dimension: Bouncing MNIST
The bouncing-MNIST experiment tests whether FERNNs can model elastic collisions that change digit velocities. It reports lower error than the G-RNN baseline and further improvement from mixing across velocity channels.
E.5.2 Mixing Across the V Dimension: Bouncing MNIST
- The dataset uses elastic collisions with boundaries and between digits, making velocity changes depend on the objects’ velocities and requiring models to represent interactions between velocities.The experiment enlarges the canvas and marks digit bounding boxes and image boundaries to make collisions easier to identify.Source: §E.5.2 · p. 38
- On bouncing MNIST, FERNN nearly halves G-RNN error, and adding V-mixing further improves performance despite lacking an equivariance guarantee for these complex flows.The authors associate V-mixing’s benefit with interactions between velocity features, such as transferring activation between opposing velocity channels after collisions.Source: §E.5.2 · p. 38, Table 4 · p. 39
- For the main constant-velocity Translating MNIST task, adding V-mixing did not improve results because inter-velocity features were not needed.Source: §E.5.2 · p. 39
E.6 Moving KTH with Data Augmentation
The appendix tests whether velocity resampling on Moving KTH and wider non-equivariant baselines change FERNNs’ advantage. FERNNs remain best with data augmentation, while larger G-RNNs do not close the performance gap.
E.6 Moving KTH with Data Augmentation
- With Moving KTH velocities resampled each training iteration, FERNNs still significantly outperformed comparable non-flow-equivariant G-RNN baselines.All models improved with resampling, and the 3D-CNN improved the most.Source: Table 5 · p. 39, §E.6 · p. 39
E.7 Larger Baselines
- On KTH, a 25×-channel G-RNN reached 54.6 ± 1% test accuracy versus 66.5% for the original G-RNN; on Translating MNIST, 1024-channel G-RNN MSE was 2.0e−3 versus FERNN-VT²’s 1.5e−4.Increasing baseline channels yielded only marginal gains at best on MNIST and significantly worse performance in other cases.Source: §E.7 · p. 39