1 of 42

Equivariant Neural Networks�for Dynamics

​

​

​

Robin Walters

Khoury College of Computer Sciences

Roux Institute, Northeastern University

Mathematics in Imaging, Data, and Optimization,

Rensselaer Polytechnic Institute

February 2, 2022

2 of 42

Dynamics and Symmetry

trajectory prediction

biomedical engineering

molecular dynamics

mechanical engineering

fluid dynamics

climate science

robotics

ecology

economics

disease modeling

3 of 42

Predicting Dynamics

Numerical Methods for

Differential Equations

Is there a way we can combine the advantages of DL with those of traditional mathematical modeling?

​

Deep Learning

  • First-Principles
  • Conservation Laws
  • Extremely flexible
  • Requires little data
  • Black Box
  • No Guarantees
  • Poor generalizability
  • High Data Requirements
  • Must know exact dynamics
  • Computationally demanding
  • Requires fully observed state, boundary conditions, and forces
  • Dynamics can be unknown
  • Orders of magnitude faster
  • Able to work with partially observed or noisy data

Solution: Equivariant Neural Networks

4 of 42

Equivariant Neural Networks

​

​

● Equivariant Neural Networks explicitly incorporate symmetry.

​

● Pipeline:

​

​

​

​

  • Advantages:
    • Sample efficiency (less data needed)
    • Parameter efficiency (smaller model)
    • Consistency (provable behavior)
    • More physically realistic (energy conservation)
    • Improved generalization under distributional shift

​

​

5 of 42

Symmetry Groups

​

 

 

6 of 42

Same Group, Different Representations

​

 

 

 

 

1

1

2

3

4

r

7 of 42

Equivariant Functions

Model functions which are equivariant with respect to symmetry group.

 

8 of 42

Equivariant Neural Networks

​

 

 

1

1

2

3

4

weight sharing

9 of 42

Warm-up: Equivariance by Weight Sharing

Q: Which M are equivariant?

10 of 42

Warm-up: Equivariance by Weight Sharing

Q: Which M are equivariant?

11 of 42

Warm-up: Equivariance by Weight Sharing

12 of 42

Turbulent Flow Prediction

​

​

​

​

Robin Walters*, Rui Wang*,, and Rose Yu. "Incorporating Symmetry into Deep Dynamics Models for Improved Generalization." International Conference on Learning Representations (ICLR), 2021.

Applications

  • Climate Science
  • Mechanical Engineering
  • Medical Devices (blood)

13 of 42

Fluid Flow and Symmetry

​

Has symmetries coming Galilean invariance and scaling laws.

The forward prediction function ��

is equivariant wrt these sym.

​

Goal: Incorporate symmetry of DNS into deep model of fΘ .

Navier Stokes Equation (Nonlinear PDE)

14 of 42

Symmetry: Galilean Invariance

​

independence of physics from arbitrary choice of coordinate system 🠆 Galilean invariance

Translation

Rotation

Uniform Motion

To a NN, these data all look very different

15 of 42

Symmetry: Scaling

​

t

t

16 of 42

Scaling: Truncated G-Convolution

​

Standard Convolution shares weights across the input by translating a kernel across the input.

​

​

​

​

For scale-equivariant convolution, we must translate and scale a kernel across the input

Related to Worrall & Welling, 2019

17 of 42

Scaling: Truncated G-Convolution

​

18 of 42

Results: Improved Generalization on Rayleigh-Bénard Convection

  • Rows are models with different equivariance; Columns are test sets with transformed data.
  • denotes ResNet/Unet model with equivariance to symmetry group G.

19 of 42

Results: Equivariance vs. Data Augmentation

20 of 42

Results: Improved Generalization on Rayleigh-Bénard Convection

  • EquNets Robust to Distributional Transformation.

UM

Rot

Scale

Target

Non-Equ

Equ

21 of 42

  • Reanalysis ocean current data from NEMO engine.
  • Approx 20 deg lat/long area in Indian, Pacific, Atlantic.
  • Lower RMSE. Much lower ESE.

Results: Real-world Ocean Currents Data

22 of 42

Trajectory Prediction

​

​

​

​

Walters, Robin*, Jinxi Li*, and Rose Yu. "Trajectory Prediction using Equivariant Continuous Convolution." International Conference on Learning Representations (ICLR), 2021.

​

23 of 42

Overview

  • Trajectory prediction in multi-agent systems
    • Necessary to avoid collisions in autonomous vehicles
    • Need consistency and physical fidelity

​

  • Encode physical dynamics, interactions, human behavior
    • make predictions for all agents
    • represent features in each agents perspective simultaneously
  • Equivariant continuous convolutions to reflect local symmetry
  • Our models achieve improvement on trajectory prediction
    • Vehicle and pedestrian data
    • Better generalization and consistency
    • Sample efficiency
    • Fewer parameters, memory
    • Faster runtime

Vehicles: Argoverse

Pedestrians: TrajNet++

24 of 42

Local Rotational Equivariance

25 of 42

Our Model - ECCO

Equivariant Continuous Convolution

26 of 42

ECCO Architecture Overview

Encode Past

27 of 42

ECCO Architecture Overview

Decode Future

28 of 42

Input, Output, Hidden features in SO(2)-rep

  • Hidden features are functions on circle

​

​

​

  • SO(2) action

​

​

29 of 42

x

x

=

=

x

x

Torus Convolution

Features are Circles. Kernels are Tori.

​

30 of 42

x

x

=

=

x

x

Features are Circles. Kernels are Tori.

​

31 of 42

Equivariant Continuous Convolution

Grid Conv2D

polar coordinates

32 of 42

SO(2)-Equivariant Continuous Convolution

Thm (Weiler & Cesa ‘19, Cohen, Geiger & Weiler ‘18)

K Equivariant ⇔

33 of 42

SO(2)-Equivariant Continuous Convolution

Thm (Weiler & Cesa ‘19, Cohen, Geiger & Weiler ‘18)

K Equivariant ⇔

34 of 42

SO(2)-Equivariant Continuous Convolution

weights shared along orbits

Thm (Weiler & Cesa ‘19, Cohen, Geiger & Weiler ‘18)

K Equivariant ⇔

35 of 42

SO(2)-Equivariant Continuous Convolution

weights shared along orbits

weights constrained (and shared) by stabilizer

Thm (Weiler & Cesa ‘19, Cohen, Geiger & Weiler ‘18)

K Equivariant ⇔

36 of 42

Equivariance Error

​

37 of 42

Results: Prediction Accuracy/Parameter Efficient

38 of 42

Results: Consistency

Truth Non-Equ Equ

39 of 42

Results: Consistency

Truth Non-Equ Equ

Right ❌

Left

40 of 42

Results: Consistency

Truth Non-Equ Equ

Left

Right ❌

Left

Left

41 of 42

Conclusion

  • Incorporate various symmetries by designing equivariant convolutional networks.
  • Variety of data types: gridded and point clouds.
  • Improved generalization, consistency, data and parameter efficiency.
  • Predictions of real-world data: ocean currents, vehicle and pedestrian trajectories.
  • Simple and effective weight sharing algorithm.
  • Novel equivariant neural network using regular representations for an infinite group.

​

​

42 of 42

Acknowledgment

Thank you for your attention.