Accelerating system simulation, identification and design with ML

Alex Beatson

Princeton University Department of Computer Science

�1

SIMULATION AND NUMERICAL METHODS

Numerical methods enable model-based reasoning about the world.

!2

SIMULATION AND NUMERICAL METHODS

Numerical methods enable model-based reasoning about the world.

e.g. ODEs (solved with Euler, Runge-Kutta, …)

!3

SIMULATION AND NUMERICAL METHODS

!4

Numerical methods enable model-based reasoning about the world.

e.g. PDEs (solved with Finite Element Analysis, Finite Volume Method, …)

SIMULATION AND NUMERICAL METHODS

!5

Numerical methods enable model-based reasoning about the world.

e.g. “inner loop” of optimization or iterative methods within variational methods, meta-learning, hyperparameter optimization

MODEL-BASED OPTIMIZATION

!6

Numerical methods enable model-based reasoning about the world.

Optimizing the output of a numerical method with respect to system parameters allows system identification, design, and control.

BOTTLENECKS

• Model design and parametrization

• Numerical method implementation and tuning

• Computation

!7

BOTTLENECKS

• Model design and parametrization

• Numerical method implementation and tuning

• Computation

!8

COMMON STRUCTURE

Discretization through time, space, or iterations

Increasing accuracy with increasing computation

!9

GOAL

Develop ML tools which exploit this structure to accelerate system simulation, identification, and design.

!10

Efficient Optimization of Loops and Limits with Randomized Telescoping Sums

�11

with Ryan P. Adams

ICML 2019

MOTIVATION

• Optimization with inner loops‣ Meta learning, hyperparameter optimization‣ Recurrent models

• Optimization with limits‣ Discretized numerical methods: PDEs, ODEs, …‣ Iterative methods: linear systems, inverses, eigenvalues, ‣ Integration with Monte Carlo or quadrature

• In both cases..‣ cheap truncations/approximations cause harmful bias‣ accurate approximations are computationally expensive

!12

RANDOMIZED TELESCOPES: UNBIASED ESTIMATION OF LIMITS

Consider:

Then:

Consider an estimator:

This is unbiased iff:

!13

DEMONSTRATION

General form

!14

Ground truth

“Russian roulette”

DEMONSTRATION

General form

!15

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!16

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!17

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!18

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!19

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!20

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!21

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!22

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!23

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!24

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!25

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!26

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!27

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!28

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!29

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!30

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!31

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!32

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!33

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!34

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!35

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!36

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!37

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!38

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!39

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!40

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!41

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!42

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!43

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!44

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!45

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!46

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!47

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!48

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!49

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!50

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!51

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!52

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!53

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!54

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!55

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!56

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!57

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!58

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!59

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!60

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!61

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!62

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!63

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!64

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!65

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!66

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!67

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!68

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!69

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!70

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!71

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!72

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!73

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!74

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!75

Ground truth

“Single sample”

“Russian roulette”

DEMONSTRATION

General form

!76

Ground truth

“Single sample”

RANDOMIZED TELESCOPES FOR OPTIMIZATION

• Randomized telescoping estimators for optimization = SGD for limits

• Can achieve finite variance and compute for any geometrically converging sequence, or sufficiently fast polynomially converging sequences (p > 3/2).

• This means we can optimize the limit itself rather than an approximation, with provable convergence rates and anytime guarantees!

!77

PRACTICAL RECIPE

• Pick a (large) maximal truncation

• Estimate the convergence properties of online

• Adapt sampling probabilities and learning rate: balance computation and variance to maximize a lower bound on increase in optimization efficiency relative to maximal truncation

!78

EXPERIMENTS

!79

LIMITATIONS AND EXTENSIONS

• Does not accelerate very high-dimensional problems (e.g., optimizing an RNN): our model of suffers in high dimensions and adaptive sampling falls back on the maximal truncation

• Theory only worked out for SGD: possible extensions to Adam, quasi-Newton, …

• Use series acceleration or predictive models to accelerate convergence to the limit

!80

SUMO: UNBIASED ESTIMATION OF LOG MARGINAL PROBABILITY FOR LATENT VARIABLE MODELS

• In latent variable modeling we often use a “decoder” and a “recognition model”

• This leads to a lower bound on the log-likelihood

• Debiasing this estimator improves image modeling and allows latent variable models to be used when lower bounds aren’t appropriate.

!81

Submittedwith Yucen Luo, Mohammad Norouzi, Jun Zhu, David Duvenaud, Ryan P. Adams, Ricky T. Q. Chen

PDE Model-Order Reduction with Neural Potential Energies

�82

with Tianju Xie, Jordan Ash, Geoffrey Roeder, Ryan P. Adams

In progress

CELLULAR MECHANICAL META MATERIALS

!83

Overvelde & Bertoldi, 2014

NONLINEAR ELASTICITY PDE

!84

BILEVEL DECOMPOSITION

!85

REDUCED-ORDER MODELING WITH NEURAL POTENTIAL ENERGIES

!86

• Collect data using Fenics to solve the PDE with random imposed boundary deformations (CPU)

• Train a neural network using Pytorch to model energy in cell as function of boundary deformation (GPU)

• Goal: “tile” these surrogate energy models to efficiently solve large systems

• Tricks: train on energy derivatives (stress-strain relations) as well as values; exploit physical invariances

• Assume boundary deformations are smooth: represent with a small number of cubic spline control points

• Spline interpolation from control points (Pytorch vector) to target locations (Fenics vector of Finite Element coefficients) is linear: ‣ Assemble the interpolation matrix once‣ Map deformations from Pytorch to Fenics, or energy gradients

from Fenics to Pytorch, by GPU matrix multiplication

SPLINE BOUNDARY REPRESENTATION

!87

INVARIANCES

!88

• Translation: subtract mean of displacements along each dimension

• Rotation: rotate coordinates to align with a reference (“Procrustes analysis”, closed form)

• Flips: use polar coordinates and one-dimensional convolution

DISTRIBUTED TRAINING WITH RAY+AWS

!89

CPU worker: data collection

CPU worker: deployment /

evaluation

GPU driver: model training

THANKS!