Mechanistic Neural Networks for Scientific Machine Learning

Authors: adeel  pervez
francesco  locatello
efstratios  gavves
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 03:27:00

Abstract: This paper presents *Mechanistic Neural Networks*, a neural network design for machine learning applications in the sciences. It incorporates a new *Mechanistic Block* in standard architectures to explicitly learn governing differential equations as representations, revealing the underlying dynamics of data and enhancing interpretability and efficiency in data modeling.
Central to our approach is a novel fast, parallel and scalable *Relaxed Linear Programming Solver* (NeuRLP) using a differentiable optimization approach for ODE learning and solving. 
Mechanistic Neural Networks demonstrate their versatility for scientific machine learning applications on tasks from equation discovery to dynamic systems modeling.

Main paper: 
 introduction: 
 In this paper, we introduce Mechanistic Neural Networks, a new neural network design that contains one or more Mechanistic Block that explicitly integrate governing equations as symbolic elements in the form of ODE representations. To efficiently train them, we revisit classical results on linear programs (Young, 1961; Rabinowitz, 1968 ) and develop a GPU-friendly solver. Together, they enable automating modeling and discovery of best-fitting mechanisms from data in an efficient, scalable, and interpretable way.Mechanistic networks are composed of two parts: a mechanistic encoder and a solver. The output of the mechanistic encoder is an explicit symbolic "ODE representation" U x of the general formEQUATIONwhereEQUATIONU x is a family of ordinary differential equations F (α, x) = 0, governed by learnable coefficients α that may be time-dependent. Coefficients α are obtained from the mechanistic encoder f θ , and parameters θ are trained to optimally model the evolution of data x = [x 1 , ..., x t ] over time.When training ODE representations we must simultaneously learn the precise form of multiple independent ODEs (or independent systems of ODEs) and solve them over several time steps. Sequential numerical solvers such as Runge-Kutta used in Neural ODEs (Chen et al., 2018) are simply too inefficient for solving large batches of independent ODEs as required for Mechanistic Neural Networks.With Mechanistic Neural Networks, we address both challenges directly in a native neural network context, as shown in Figure 1 . We propose a novel neural Relaxed Linear Programming Solver (NeuRLP) for ODEs based on differentiable optimization methods (Amos and Kolter, 2017 ; Young, A G z p K v W a l g 9 J p m V q O S U z + J 1 Q x X Q 8 9 c s F u V J G i + W X U s J 2 0 L l Q I q m W s 9 p J C R L Q p m Z n h + O M z M R w y S 2 s u P m 6 X h x 4 H e Z y W T X d 4 O x S l 7 X X M u P X g p k w 1 m g O l o x A / F s h b 7 P A j s z 7 U d v n G 1 / y I 6 O 1 b l W l i z L U h k f Z 8 X / y y N E E a F u 4 / K Q t A h J l 0 3 i M E v E 4 2 z y t i S m O w 3 q v n s x z M t z 4 y 1 L i 7 p s l / k q 0 3 k 1 S D 8 r P i 0 L l c x 2 B z G C p 5 + X Z s J z I 3 W Z + f 8 r 5 V Q r c r / 5 k J n c 7 Z o e C 7 q e M F Y 3 1 g x z z V j d N A o q C z p 3 h y I Z D s f 3 G w 9 6 A Q X o Z K s b 4 2 p k f z 6 y P 4 t c s t x g 4 H k X a s H 5 + 9 I g N 5 C 4 / v z D s o v 3 7 e d x 5 m + 3 v z W t + t + W B Y k 8 S O i g Z V O x Y q 7 6 z 0 x m q a B G 2 7 W T 3 s v T R S q u V Y 6 j q W u 2 x t l 0 4 z n q 3 Q 1 G h k L z C d 5 t V 6 I U q k a x R W E d k 3 l R j 0 2 2 W W k L h j 7 L y 9 W 0 z u Z m a j j q 4 j q Z s 6 a W V o V 0 Y J o V 0 Q b o q i I A q L j V F Q H a r c i d i H 2 K m I P o l s R D y D 2 K 2 I f 4 m N F f I T o V U Q P o l 8 R f Y h B R Q w g h h U x h B h V x A j i U 0 V 8 g i g r o o Q Y V 8 Q Y Y l I R E 4 i D i j i A m F b E F O J zL m G 5 4 A X G l 4 A X m p 4 C X i l 4 R X g t Y b X g G 0 N 2 4 A 3 G t 4 A 3 m p 4 C 3 i n 4 Z 3 u k D B S n B o p O C l O n R S s F K d W C l 6 K U y 8 F M 8 W p m Y K b 4 t R N w U 5 x a q f g p z j 1 U z B U n B o q O C p O H R U s F a e W C p 6 K U 0 8 F U 8 W p q Y K r 4 t R V w V Z x a q v g q z j 1 V T B W n B o r O C t O n R W s F a f W C t 6 K U 2 8 F c 8 W p u Y K 7 4 t R d w V 5 x a q / g r z j 1 V z B Y n B o s O C w + c V j l 7 7 s t e 5 y 1 O n n n m u f 9 i d D p V y S / U P x Q f Z 0 v 5 s v P s w 4 f t z 7 W V L L p H / L y r 9 W q Q 6 P L S l 5 W 0 G i l V V b Q a 6 V d V t B u p S g r a L j S K S t o u b J b V t B 1 Z a + s o P F K t 6 y g 9 8 p + W U H 7 l Y 9 l B R 1 Y e m X F m x h g 6 R t u b D B P C m a P D G Z M E 7 9 m 9 A d x Y t h h 8 N f E y P + E Y b B g Z E i l h L 6 R M u k y 7 o n Y 8 K d z B + W 5 0 b J l W F b Q t W V U V t C 4 5 V N Z Q e + W s q y g f c u 4 r K C D y 6 S s o I n L Q V l B H 5 d p W U E r l 8 9 l B d 1 c D s s K G r o c l R X 0 d P m x r B R t X d 2 U 8 2 m e h Z W E s + Y t U x H F r h c G q H 1 z S s i 6 k 5 5 I G I k o P h O d a B R z Y L I 3 U g u Y b I z U B i a 7 I h X A Z E u k D j D Z D 2 k X m G y G t A d M d k I 6 A C b b I O 0 D k z 2 Q P g K T D Z B 6 w B 7 B P r B P M E k 0 z X A I T I o 5 j Y B J J a d P w K S M U w l M a j i N g W N 6 U 4 G T x T m h p Z s C k 7 p N n 4 F J 0 a Z D Y F K x 6 Q i Y l G v 6 E X h a q 2 b x W 2 y m 6 i y U w u 8 E Y V L 8 4 v x e V b W S O 6 9 X N + Z / Z K 4 e X G + + 3 / j + / f d n f 1 / 9 c W v 6 A / Q X t T / X / l L 7 t r Z R + 6 H 2 Y 2 2 / 1 q p d1 a y a q P 2 z 9 q / a v 7 / 5 7 5 t X b 7 5 8 8 9 U k 9 N V n 0 z F / r J V e b / 7 0 P 0 b z I u g = < / l a t e x i t >x < l a t e x i t s h a 1 _ b a s e 6 4 = " D t T a e E n t V 9 V M w V J w a K j g q T h 0 V L B W n l g q e i l N P B V P F q a m C q + L U V c F W c W q r 4 K s 4 9 V U w V p w a K z g r T p 0 V r B W n 1 g r e i l N v B X P F q b m C u + L U X c F e c W q v 4 K 8 4 9 V c w W J w a L D g s P n F Y 5 e + 7 L X u c t T p 5 5 5 r n / Y n Q 6 S s e k I m J R r + h F 4 W q t m 8 V t s p u o s l M L v B G F S / O j 8 X l W 1 k j u v V j f m f 2 S u H l x v v t / 4 4 f 0 P Z 3 9 b / X F r + g P 0 l 7 U / 1 / 5 S e 1 v b q P 2 9 9 m N t v 9 a q X d W s W q / 2 z 9 q / a v / + 9 r + v v 3 j 9 1 e u v J 6 G f f z Y d 8 8 d a 6 f X 6 T / 8 D h / I j y A = = < / l a t e x i t > u(t)Q M m K / d S g i / 1 A g s Z v 4 = " > A A A f X X i c l Z n b c u P G E Y Z p J 0 4 c x X H W y U U u U u V F R b W V t U u 7 J a l S T i 5 N U W e J E n W k D t C y Z o A B C Q o n D U C I X J i v k N v k 1 X K V V 8 k A J O d v E K S cG i Q i s y Y m c g W c k o Z F f u W G 7 U l i J N 1 I H z J K u W q t h 9 Z h k V q L y U z 6 L 1 w 1 V Q M 9 f s 1 i U 5 2 m 8 W H Y t J W w L l Q M p m m o 9 p 5 G Q L A l l Z n p + O M 7 M R A y T 2 M q O m 6 f j x Y H f Z y a T X d 8 N x i p 5 X X M t P 3 o p k A 1 n g e p o x Q z E s x X 6 P g v s z L Q f v X G 2 / S E 7 O l b n W l m y L E t l f J w V / y + P E E W E u o 3 L Q 9 I i J F 0 2 i c M s E Y + zp g V O 3 W j e G q M h h N n V J 4 r C K U / m c m Y r u J a y G Q 2 n y n z T 0 b x D F p R l + O Y n E 3 S 5 C a T d + 5 k r M h R I f K K y C F a F d G C a F d E G 6 K o i A K i 4 1 R U B 2 q 3 I n Y h 9 i p i D 6 J b E Q 8 g 9 i t i H + J j R X y E 6 F V E D 6 J f E X 2 I Q U U M I I Y V M Y Q Y V c Q I 4 l N F f I I o K 6 K E G F f E G G J S E R O I g 4 o 4 g J h W x B T i c 0 V 8 h j i s i E O I o 4 o 4 g v i x I n 4 s H n T v 6 G v l j b E V e r a h 3 L q b M M / w R D L x 9 M U U 9 W x i e d T Y O i b e A t 0 C b Y A 2 Q L d B t 0 F 3 Q H d A d 0 F 3 Q f d A 9 0 D 3 Q f d B D 0 B J Y R + C H o I e g R 6 B H o M e g z Z B m 6 A n o C e g p 6 C n o C 3 Q F u g Z 6 B n o O e g 5 6 A X o B e g l 6 C X o F e g V 6 D X o N W g b t A 1 6 A 3 o D e g t 6 C 3 o H e l d U m N F g n t u V L O r l X x Q X V Z Z V n w y x m D d x Q A X c A t z S s A H Y 0 H A b c F v D H c A d D X c B d z X c A 9 z T c B 9 w X 8 M D w A M N D w E P N T w C P N L w G P B Y w y Z g U 8 M T w B M N T w F P N W w B t j Q 8 A z z T 8 B z w X M M L w A s N L w E v N b w C v N L w G v B a w z Z g W 8 M b w B s N b w F v N b w D v F v 0 h I p Z E B u x k K 5 j e O G z k L l V m A y N H T Y t v p g 2 7 d j h w J x g C 9 g i 2 A a 2 C R Y C X B CV c k v 1 D 8 U H 2 d L + b L z 7 M O H 7 c + 1 l S y 6 R / y 8 q / V q k O j y 0 p e V t B o p V V W 0 G u l X V b Q b q U o K 2 i 4 0 i k r a L m y W 1 b Q d W W v r K D x S r e s o P f K f l l B + 5 W P Z Q U d W H p l x Z s Y Y O k b b m w w T w p m j w x m T B O / Z v Q H c W L Y Y f D X x M j / h G G w Y G R I p Y S + k T L p M u 6 J 2 P C n c w f l u d G y Z V h W 0 L V l V F b Q u O V T W U H v l r K s o H 3 L u K y g g 8 u k r K C J y 0 F Z Q R + X a V l B K 5 f P Z Q X d X A 7 L C h q 6 H J U V 9 H T 5 s a w U b V 3 d l P N p n o W V h L P m L V M R< l a t e x i t s h a 1 _ b a s e 6 4 = " 6 0 0 a + z E B 1 2 1 0 p 6 Z B o C F X t m 0 X W x I = " > A A A f Y X i c l Z n Z c u P G F Y Z p Z 7 O V x W P n 0 l U e V F R T c V K a K U m V c n J p i t o l S t R K L d C w u o E G C Q q b G i B E C u Z L 5 D Z 5 s V z n R d I A y f 4 P C F L O s G q G 4 P e f b j Q O T u P 8 F H n k u X G y v v 6 f z z 7 / x S 9 / 9 e v f f P H l y m 9 / 9 / s / f P X m 6 2 + u 4 3 A g L X F l h V 4 o b z i L h e c G 4 i p x E 0 / c R F I w n 3 u i z R 8 b u d 5 O h Y z d M L h M R p F 4 8 F k 3 c B 3 X Y o l C N 0 7 H T H o i Y Z 0 3 q + s f 1 o u X U T 3 Y m B 6 s 1 q a v V u f r 7 x q m H V o D X w S J 5 b E 4 v t 9 Y j 5 K H j M n E t T w x X j E H s Y i Y 9 c i 6 4 l 4 d B s w X 8 U N W L H h s v F P E N p x Q q n 9 B Y h S U j s i Y H 8 c j n 6 t I n y W 9 e F 7 L 4 S L t f p A 4 / 3 j I 3 C A a J C K w J i d y B p 6 R h E Z + 9 Y b t S m E l 3 k g d M E u 6 a q 2 G 1 W O S W Y n K U f k s X j d U A T 1 / z W J R n q v x Y t m 1 l L A t V A 6 k a K r 1 n E Z C s i S Ut G Q W V B 5 + 5 Q J M P h + H 7 j Q S + g A J 1 s d W N c j e z P R / Z n k U u W G w w 8 7 0 I t O H 9 f G u Q G E t e f f 1 h 2 8 b 7 9 P M 7 8 7 f b 3 p l X / y 7 I g k Q c J H b R s K l b M V f + Z y S w V 1 G i 7 d t J 7 f b p I x b X K c T R 1 z d Y 4 m 2 4 8 R 7 2 7 w c h Q a D 7 B u + 1 K l E L V K L Y o r G M y L + q x y T Y r b c H Q Z 3 m 5 m t b Z 3 E w N R 1 1 c J 3 P W 1 J L H Y 2 N e l D N V 9 s I l E T T k 4 + b i o P z a 9 V k + 5 q m o R t B 5 j D y D p j W I 3 h e Z X D a i 1 X 5 t T J 6 d r D 6 f u n w g x h X L 0 Y E L T j F b 1 W Z x i t e W 3 y i e T 1 h O E T t 3 j x q n J z u V m 5 m P p o M / Z d y Q D v z J 5 M N O Y v 7 U K R L x C b P I B p v N U R T Q p 6 6 / 3 n 5 l 7 f W i f N + V h u S P a N R k c 2 5 C f + D k z 4 J B P u O y l K s Y S Y L 0 g u s L 4 i Y 5 q s b + P x c 3 G 5 / v y 5 + f o X 6 x Y 2 x V 5 r C F l 7 B i g u K o k 2 / y B R H 5 J p l F 5 N e 7 k j 9 t U 9 c W e L w H g Z D F E 1 6 1 H + E k 0 2 d 8 / o h f V Q X q d n v J 7 A E / d w J l O 5 Q d K 5 7 c k 6 H 3 q x u T A Q 9 F 3 R O z Y q d u F E + N 0 X D i j M p z B a H 0 J z M Z 0 1 V c C 5 n M 5 j N l / s k o n k E r 6 n I c k 7 N J m t x k 8 s 6 d j B U 5 K k R e E T l E q y J a E O 2 K a E M U F V F A d J y K 6 k D t V s Q u x F 5 F 7 E F 0 K + I B x H 5 F 7 E N 8 r I i P E L 2 K 6 E H 0 K 6 I P M a i I A c S w I o Y Q o 4 o Y Q X y q i E 8 Q Z U W U E O O K G E N M K m I C c V A R B x D T i p h C f K 6 I z x C H F X E I c V Q R R x B f K u J L 8 a B 7 T 1 8 r 7 4 y t 0 L M N 5 d b d h H m G J 5 K J p y + m q G c T y 6 P G 1 j H x F u g W a A O 0 A b o N u g 2 6 A 7 o D u g u 6 C 7 o H u g e 6 D 7 o P e g B K C v s Q 9 B D 0 C P Q I 9 B j 0 G L Q J 2 g Q 9 A T 0 B P Q U 9 B W 2 B t k D P Q M 9 A z 0 H P Q S 9 A L 0 A v Q S 9 B r 0 C v Q K 9 B r 0 H b o G 3 Q G 9 A b 0 F v Q W 9 A 7 0 L u i w o w G 8 9 y u Z F E v / 6 K 4 q L K s + m S I x b y J A y r g F u C W h g 3 A h o b b g N s a 7 g D u a L g L u K v h H u C e h v u A + x o e A B 5 o e A h 4 q O E R 4 J G G x 4 D H G j Y B mx q e A J 5 o e A p 4 q m E L s K X h G e C Z h u e A 5 x p e A F 5 o e A l 4 q e E V 4 J W G 1 4 D X G r Y B 2 x r e A N 5 o e A t 4 q + E d 4 N 2 i J 1 T M g t i I h X Q d w w u f h c y t w m R o 7 L B p 8 c W 0 a c c O B + Y E W 8 A W w T a w T b A Q 4 I J w 1 Z 4 1 d w j v A n c J 7 g H 3 C H a B X Y L 7 w H 2 C H 4 E f C f a A P Y J 9 Y J / g A D g g O A Q O C Y 6 A I 4 K f g J 8 I l s C S 4 B g 4 J j g B T g g e A A 8 I T o F T g p + B n w k e A g 8 J H g G P C H 4 B X t w f L 1 B 9 C 5 9 l s V P X M 9 Q x 7 5 a G W 4 A N D R u A 2 x p u A + 5 o u A O 4 q + E u 4 J 6 G e 4 D 7 G u 4 D H m h 4 A H i o 4 S H g k Y Z H g M c a H g M 2 N W w C n m h 4 A n i q 4 S l g S 8 M W 4 J m G Z 4 D n G p 4 D X m h 4 A X i p 4 S X g l Y Z X g N c a X g O 2 N W w D 3 m h 4 A 3 i r 4 S 3 g n Y Z 3 u k P C S H F q p O C k O H V S s F K c W i l 4 K U 6 9 F M w U p 2 Y K b o p T N w U 7 x a m d g p / i 1 E / B U H F q q O C o O H V U s F S c W i p 4 K k 4 9 F U w V p 6 Y K r o p T V w V b x a m t g q / i 1 F f B W H F q r O C s O H V W s F a c W i t 4 K 0 6 9 F c w V p + Y K 7 o p T d w V 7 x a m 9 g r / i 1 F / B Y H F q s O C w + M R h l b / v t u x x 1 u r k n W u e 9 y d C p 1 + R / E L x Q / V 1 v p g v P 8 8 6 f N z 6 W F P J p n / I y 7 9 W q w 6 N L i t 5 W U G j l V Z Z Q a + V d l l B u 5 W i r K D h S q e s o O X K b l l B 1 5 W 9 s o L G K 9 2 y g t 4 r + 2 U F 7 V c + l h V 0 Y O m V F W 9 i g K V v u L H B P C m Y P T K Y M U 3 8 m t E f x I l h h 8 G f E y P / E 4 b B g p E h l R L 6 R s q k y 7 g n Y s O f z h 2 U 5 0 b L l m F Z Q d e W U V l B 4 5 Z P Z Q W 9 W 8 q y g v Y t 4 7 K C D i 6 T s o I m L g d l B X 1 c p m U F r V w + l x V 0 c z k s K 2 j o c l R W 0 N P l S 1 k p 2 r q 6 K e f T P A s r C W f N W 6 Y i i l 0 v D F D 7 5 p S Q d R c / 9 p G I 4 j P R i U Y x B y Z 7 I 7 W A y c Z I b W C y K 1 I B T L Z E 6 g C T / Z B 2 g c l m S H v A Z C e k A 2 C y D d I + M N k D 6 S M w 2 Q C p B + w R 7 A P 7 B J N E 0 w y H w K S Y 0 w i Y V H L 6 B E z K O J X A p I b T G D i m N x U 4 W Z w T W r o p M K n b 9 B m Y F G 0 6 B C Y V m 4 6 A S b m m L 8 D T W j W L 3 2 I z V W e h F H 4 n C J P i h + c P q q q V 3 H m z u j H / I 3 P 1 4 H r z w 8 Y P H 3 4 4 + 9 v q j 1 v T H 6 C / q H 1 b + 1 P t + 9 p G 7 e + 1 H 2 v 7 t V b t q m b V v N o / a / + q / f u 7 / 7 7 9 8 u 2 b t 9 9 M Q j / / b D r m j 7 X S 6 + 2 3 / w M t M C X t < / l a t< l a t e x i t s h a 1 _ b a s e 6 4 = " 6 0 0 a + z E B 1 21 0 p 6 Z B o C F X t m 0 X W x I = " > A A A f Y X i c l Z n Z c u P G F Y Z p Z 7 O V x W P n 0 l U e V F R T c V K a K U m V c n J p i t o l S t R K L d C w u o E G C Q q b G i B E C u Z L 5 D Z 5 s V z n R d I A y f 4 P C F L O s G q G 4 P e f b j Q O T u P 8 F H n k u X G y v v 6 f z z 7 / x S 9 / 9 e v f f P H l y m 9 / 9 / s / f P X m 6 2 + u 4 3 A g L X F l h V 4 o b z i L h e c G 4 i p x E 0 / c R F I w n 3 u i z R 8 b u d 5 O h Y z d M L h M R p F 4 8 F k 3 c B 3 X Y o l C N 0 7 H T H o i Y Z 0 3 q + s f 1 o u X U T 3 Y m B 6 s 1 q a v V u f r 7 x q m H V o D X w S J 5 b E 4 v t 9 Y j 5 K H j M n E t T w x X j E H s Y i Y 9 c i 6 4 l 4 d B s w X 8 U N W L H h s v F P E N p x Q q n 9 B Y h S U j s i Y H 8 c j n 6 t I n y W 9 e F 7 L 4 S L t f p A 4 / 3 j I 3 C A a J C K w J i d y B p 6 R h E Z + 9 Y b t S m E l 3 k g d M E u 6 a q 2 G 1 W O S W Y n K U f k s X j d U A T 1 / z W J R n q v x Y t m 1 l L A t V A 6 k a K r 1 n E Z C s i S U m e n 5 4 T g z E z F M Y i s 7 b p 6 O F w f + N T O Z 7 P p u M F b J 6 5 p r + d F r g W w 4 C 1 R H K 2 Y g n q 3 Q 9 1 l g Z 6 b 9 6 I 2 z 7 Y / Z 0 b E 6 1 8 q S Z V k q 4 + O s + H 9 5 h C g i 1 G 1 c H p I W I e m y S R x m i X i c T d 6 W x H S n Q d 3 3 r 4 Z 5 e W 6 8 Z W l R l + 0 y X 2 U 6 r w b p Z 8 W n Z a G S 2 e 4 g R v D 0 8 9 J M e G 6 k L j P / f 6 W c a k X u N x 8 y k 7 t d 0 2 N B 1 x P G 6 s a a Y a 4 Z q 5 t G Q W V B 5 + 5 Q J M P h + H 7 j Q S + g A J 1 s d W N c j e z P R / Z n k U u W G w w 8 7 0 I t O H 9 f G u Q G E t e f f 1 h 2 8 b 7 9 P M 7 8 7 f b 3 p l X / y 7 I g k Q c J H b R s K l b M V f + Z y S w V 1 G i 7 d t J 7 f b p I x b X K c T R 1 z d Y 4 m 2 4 8 R 7 2 7 w c h Q a D 7 B u + 1 K l E L V K L Y o r G M y L + q x y T Y r b c H Q Z 3 m 5 m t b Z 3 E w N R 1 1 c J 3 P W 1 J L H Y 2 N e l D N V 9 s I l E T T k 4 + b i o P z a 9 V k + 5 q m o R t B 5 j D y D p j W I 3 h e Z X D a i 1 X 5 t T J 6 d r D 6 f u n w g x h X L 0 Y E L T j F b 1 W Z x i t e W 3 y i e T 1 h O E T t 3 j x q n J z u V m 5 m P p o M / Z d y Q D v z J 5 M N O Y v 7 U K R L x C b P I B p v N U R T Q p 6 6 / 3 n 5 l 7 f W i f N + V h u S P a N R k c 2 5 C f + D k z 4 J B P u O y l K s Y S Y L 0 g u s L 4 i Y 5 q s b + P x c 3 G 5 / v y 5 + f o X 6 x Y 2 x V 5 r C F l 7 B i g u K o k 2 / y B R H 5 J p l F 5 N e 7 k j 9 t U 9 c W e L w H g Z D F E 1 6 1 H + E k 0 2 d 8 / o h f V Q X q d n v J 7 A E / d w J l O 5 Q d K 5 7 c k 6 H 3 q x u T A Q 9 F 3 R O z Y q d u F E + N 0 X D i j M p z B a H 0 J z M Z 0 1 V c C 5 n M 5 j N l / s k o n k E r 6 n I c k 7 N J m t x k 8 s 6 d j B U 5 K k R e E T l E q y J a E O 2 K a E M U F V F A d J y K 6 k D t V s Q u x F 5 F 7 E F 0 K + I B x H 5 F 7 E N 8 r I i P E L 2 K 6 E H 0 K 6 I P M a i I A c S w I o Y Q o 4 o Y Q X y q i E 8 Q Z U W U E O O K G E N M K m I C c V A R B x D T i p h C f K 6 I z x C H F X E I c V Q R R x B f K u J L 8 a B 7 T 1 8 r 7 4 y t 0 L M N 5 d b d h H m G J 5 K J p y + m q G c T y 6 P G 1 j H x F u g W a A O 0 A b o N u g 2 6 A 7 o D u g u 6 C 7 o H u g e 6 D 7 o P e g B K C v s Q 9 B D 0 C P Q I 9 B j 0 G L Q J 2 g Q 9 A T 0 B P Q U 9 B W 2 B t k D P Q M 9 A z 0 H P Q S 9 A L 0 A v Q S 9 B r 0 C v Q K 9 B r 0 H b o G 3 Q G 9 A b 0 F v Q W 9 A 7 0 L u i w o w G 8 9 y u Z F E v / 6 K 4 q L K s + m S I x b y J A y r g F u C W h g 3 A h o b b g N s a 7 g D u a L g L u K v h H u C e h v u A + x o e A B 5 o e A h 4 q O E R 4 J G G x 4 D H G j Y B m x q e A J 5 o e A p 4 q m E L s K X h G e C Z h u e A 5 x p e A F 5 o e A l 4 q e E V 4 J W G 1 4 D X G r Y B 2 x r e A N 5 o e A t 4 q + E d 4 N 2 i J 1 T M g t i I h X Q d w w u f h c y t w m R o 7 L B p 8 c W 0 a c c O B + Y E W 8 A W w T a w T b A Q 4 I J w 1 Z 4 1 d w j v A n c J 7 g H 3 C H a B X Y L 7 w H 2 C H 4 E f C f a A P Y J 9 Y J / g A D g g O A Q O C Y 6 A I 4 K f g J 8 I l s C S 4 B g 4 J j g B T g g e A A 8 I T o F T g p + B n w k e A g 8 J H g G P C H 4 B X t w f L 1 B 9 C 5 9 l s V P X M 9 Q x 7 5 a G W 4 A N D R u A 2 x p u A + 5 o u A O 4 q + E u 4 J 6 G e 4 D 7 G u 4 D H m h 4 A H i o 4 S H g k Y Z H g M c a H g M 2 N W w C n m h 4 A n i q 4 S l g S 8 M W 4 J m G Z 4 D n G p 4 D X m h 4 A X i p 4 S X g l Y Z X g N c a X g O 2 N W w D 3 m h 4 A 3 i r 4 S 3 g n Y Z 3 u k P C S H F q p O C k O H V S s F K c W i l 4 K U 6 9 F M w U p 2 Y K b o p T N w U 7 x a m d g p / i 1 E / B U H F q q O C o O H V U s F S c W i p 4 K k 4 9 F U w V p 6 Y K r o p T V w V b x a m t g q / i 1 F f B W H F q r O C s O H V W s F a c W i t 4 K 0 6 9 F c w V p + Y K 7 o p T d w V 7 x a m 9 g r / i 1 F / B Y H F q s O C w + M R h l b / v t u x x 1 u r k n W u e 9 y d C p 1 + R / E L x Q / V 1 v p g v P 8 8 6 f N z 6 W F P J p n / I y 7 9 W q w 6 N L i t 5 W U G j l V Z Z Q a + V d l l B u 5 W i r K D h S q e s o O X K b l l B 1 5 W 9 s o L G K 9 2 y g t 4 r + 2 U F 7 V c + l h V 0 Y O m V F W 9 i g K V v u L H B P C m Y P T K Y M U 3 8 m t E f x I l h h 8 G f E y P / E 4 b B g p E h l R L 6 R s q k y 7 g n Y s O f z h 2 U 5 0 b L l m F Z Q d e W U V l B 4 5 Z P Z Q W 9 W 8 q y g v Y t 4 7 K C D i 6 T s o I m L g d l B X 1 c p m U F r V w + l x V 0 c z k s K 2 j o c l R W 0 N P l S 1 k p 2 r q 6 K e f T P A s r C W f N W 6 Y i i l 0 v D F D 7 5 p S Q d R c / 9 p G I 4 j P R i U Y x B y Z 7 I 7 W A y c Z I b W C y K 1 I B T L Z E 6 g C T / Z B 2 g c l m S H v A Z C e k A 2 C y D d I + M N k D 6 S M w 2 Q C p B + w R 7 A P 7 B J N E 0 w y H w K S Y 0 w i Y V H L 6 B E z K O J X A p I b T G D i m N x U 4 W Z w T W r o p M K n b 9 B m Y F G 0 6 B C Y V m 4 6 A S b m m L 8 D T W j W L 3 2 I z V W e h FH 4 n C J P i h + c P q q q V 3 H m z u j H / I 3 P 1 4 H r z w 8 Y P H 3 4 4 + 9 v q j 1 v T H 6 C / q H 1 b + 1 P t + 9 p G 7 e + 1 H 2 v 7 t V b t q m b V v N o / a / + q / f u 7 / 7 7 9 8 u 2 b t 9 9 M Q j / / b D r m j 7 X S 6 + 2 3 / w M t M C X t < / l a t e x i t > f✓ Decoder Basis parametersBasis < l a t e x i t s h a 1 _ b a s e 6 4 = " M 2 o O g N 5 N c 6 B C 1 h C L 3 U O c B / 9 a Y 9 E = " > A A A f e 3 i c l Z n Z b u P I F Y Y 1 k 2 3 i b D 3 J Z Y B p I k Y j P Q N 3 w z a S S R A g w M j y v s q L L C 9 0 C 1 V k U a L M z U W K l p q j 2 z x N b p N 3 y c M E S J G S 6 j 8 U J U 9 a Q L e o 7 z 9 1 V D y q q v P L 4 p H n x s n 6 + n 8 + + / x H P / 7 J T 3 / 2 x c 9 X f v H L X / 3 6 N 6 + + / O 1 1 H A 6 k J V p W 6 I X y h r N Y e G 4 g W o m b e O I m k o L 5 3 B N t / t j I 9 X Y q Z O y G w V U y i s S D z 7 q B 6 7 g W S x T q v D J M n y U 9 i 3 l Z a 9 w Z G n 8 z d t + a z I t 6 b M 0 Y f m 3 8 3 V j v v F p d f 7 9 e P I z q x c b 0 Y r U 2 f T Q 7 X 3 7 V M O 3 Q G v g i S C y P x f H 9 x n q U P G R M J q 7 l i f G K O Y h F x K x H 1 h X 3 6 j J g v o g f s u J W x s Y b R W z D C a X 6 F y R G Q e m I j P l x P P K 5 i s w n H s 9 r O V y k 3 Q 8 S 5 6 8 P m R t E g 0 Q E 1 u S N n I F n J K G R 1 8 W w X S m s x B u p C 2 Z J V 8 3 V s H p M M i t R 1 S u / i 9 c N V U D P X 7 N Y l F d x v F h 2 L S V s C 1 U D K U 7 U f M 4 i I V k S y s z 0 / H C c m Y k Y J r G V H Z + c j R c H f p O Z T H Z 9 N x i r 4 n X N t f z q p U A 2 n A W q q x U z E M 9 W 6 P s s s D P T f v T G 2 f a H 7 O h Y v d f K k m l Z q u L j r P h / e Y Q o I t T H u D w k L U L S Z U k c Z o l 4 n E 2 e l s R 0 p 0 H d d y + G e X l t v G V l U b f t M l 9 V O l 8 N 0 s + K V 8 t C J b P d Q Y z g 6 e u l l f D c S N 1 m / v 9 K u d S K 3 G 8 + Z C Z 3 u 6 b H g q 4 n j N W N N c N c M 1 Y 3 j Y L K g s 5 9 Q p E M h + P 7 j Q c 9 g Q J 0 s t W N c T W y P x / Z n 0 U u m W 4 w 8 L x L N e H 8 e W m Q G 0 j c f / 5 i 2 c 3 7 9 v M 4 8 7 f b b 0 2 r / v W y I J E H C R 2 0 L B U r c t V / I J m l g h p t 1 0 5 6 L 6 e L V F y z H E d L d 9 I c Z 9 O N 5 6 h n N x g Z C s 0 X e L d d i V K o G s U W h X U m x + d k m 5 W 2 Y O i z f L m a 1 v l c p o a j b q 6 T O W t q y u O x M S / K m S p 7 4 Z I I G v J h c 3 F Q f u / 6 X T 7 k p a h G 0 D x G X k H T G k T v i k o u G 9 F s v z Q m r 0 5 W n y 9 d P h D j i u n o w A V v M Z v V Z v E W L 0 2 / U Z x P m E 4 R O / c Z N c 5 O d y o f Z j 6 a D v 6 U c U M 6 8 H u T D z u J + X 2 n K M Q n Z J E N N s t R L K B P n X + 9 / c L c 6 8 X y f V M a k h / R W J M n c w n 9 g Z O f B Y M 8 4 7 K S q x h J g v S E 6 w v i J j W q x v 4 / N z c b n + / L H 8 5 Q v 9 w x t i o 5 b O E l r E h Q X H X y T b 4 g I t 8 k s 4 j 8 f l f y 0 z Z 1 b Y H j P Q i E L E 5 4 1 X 6 E k 0 z P + P y I X 1 U L 1 O 3 2 k t k B P / c G y n Y o o 1 a c 3 J O h 9 6 s b k w E P x b o n Z s V O 3 S i e G q P h x B m V c w W h 9 C e Z j O k s r o V M Z v l M m b 8 y i j N o R d 2 O Y 3 I 2 K Z O b T J 6 5 k 7 G i R o X I K y K H a F V E C 6 J d E W 2 I o i I K i I 5 T U R 2 o 3 Y r Y h d i r i D 2 I b k U 8 g N i v i H 2 I j x X x E a J X E T 2 I f k X 0 I Q Y V M Y A Y V s Q Q Y l Q R I 4 h P F f E J o q y I E m J c E W O I S U V M I A 4 q 4 g B i W h F T i M 8 V 8 R n i s C I O I Y 4 q 4 g j i x 4 r 4 s T j o 3 t H H y h t j K / R s Q 7 l 1 N 2 G e 4 Y l k 4 u m L F P V s Y n n U 2 D o S b 4 F u g T Z A G 6 D b o N u g O 6 A 7 o L u g u 6 B 7 o H u g + 6 D 7 o A e g Z G E f g h 6 C H o E e g R 6 D H o O e g J 6 A n o K e g p 6 B n o E 2 Q Z u g 5 6 D n o B e g F 6 C X o J e g V 6 B X o C 3 Q F u g 1 6 D V o G 7 Q N e g N 6 A 3 o L e g t 6 B 3 p X r D C j w T y 3 K 1 n U y 7 8 o L l p Z V j 3 T 3 6 3 r s 0 T W F u C W h g 3 A h o b b g N s a 7 g D u a L g L u K v h H u C e h v u A + x o e A B 5 o e A h 4 q O E R 4 J G G x 4 D H G p 4 A n m h 4 C n i q 4 R n g m Y Z N w K a G 5 4 D n G l 4 A X m h 4 C X i p 4 R X g l Y Y t w J a G 1 4 D X G r Y B 2 x r e A N 5 o e A t 4 q + E d 4 N 2 i E y p m Q W z E Q r q O 4 Y X P Q u Z W Y T I 0 d t h 0 8 c W 0 a c c O B + Y E W 8 A W w T a w T b A Q 4 I J w 1 Z 4 1 d w j v A n c J 7 g H 3 C H a B X Y L 7 w H 2 C H 4 E f C f a A P Y J 9 Y J / g A D g g O A Q O C Y 6 A I 4 K f g J 8 I l s C S 4 B g 4 J j g B T g g e A A 8 I T o F T g p + B n w k e A g 8 J H g G P C P 4 I v L g / X m L 1 L T z L Y q e u M 9 S R d 0 v D L c C G h g 3 A b Q 2 3 A X c 0 3 A H c 1 X A X c E / D P c B 9 D f c B D z Q 8 A D z U 8 B D w S M M j w G M N j w F P N D w B P N X w F P B M w z P A p o Z N w H M N z w E v N L w A v N T w E v B K w y v A l o Y t w G s N r w H b G r Y B b z S 8 A b z V 8 B b w T s M 7 3 S F h p D g 1 U n B S n D o p W C l O r R S 8 F K d e C m a K U z M F N 8 W p m 4 K d 4 t R O w U 9 x 6 q d g q D g 1 V H B U n D o q W C p O L R U 8 F a e e C q a K U 1 M F V 8 W p q 4 K t 4 t R W w V d x 6 q t g r D g 1 V n B W n D o r W C t O r R W 8 F a f e C u a K U 3 M F d 8 W p u 4 K 9 4 t R e w V 9 x 6 q 9 g s D g 1 W H B Y f O K w y t 9 3 m / Y 4 a 3 b y z j X P + x O h 0 6 9 I f q H 4 o f o 6 X + T L 3 2 c d P m 5 9 r K l k 0 z / k 5 V + r V Y d G l 5 W 8 r K D R S q u s o N d K u 6 y g 3 U p R V t B w p V N W 0 H J l t 6 y g 6 8 p e W U H j l W 5 Z Q e + V / b K C 9 i s f y w o 6 s P T K i j c x w N I 3 3 N h g n h T M H h n M m B Z + z e g P 4 s S w w + C P i Z H / C c N g w c i Q S g l 9 I 2 X S Z d w T s e F P c w f l 3 G j Z M i w r 6 N o y K i t o 3 P K p r K B 3 S 1 l W 0 L 5 l X F b Q w W V S V t D E 5 a C s o I / L t K y g l c v n s o J u L o d l B Q 1 d j s o K e r r 8 W F a K t q 4 + l I t p n Y W V h L P m L V M R x a 4 X B l j 7 5 p S Q e S c 9 k T A S U b w m O t E o 5 s B k b 6 Q W M N k Y q Q 1 M d k U q g M m W S B 1 g s h / S L j D Z D G k P m O y E d A B M t k H a B y Z 7 I H 0 E J h s g 9 Y A 9 g n 1 g n 2 B S a F r h E J g s 5 j Q C J i s 5 f Q I m y z i V w G Q N p z F w T D 9 U 4 G R x T e j S T Y H J u k 2 f g c m i T Y f A Z M W m I 2 C y X N O P w N O 1 a h a / x WZ q n Y V S + J 0 g T I q f p N + r V a 3 k z q v V j f k f m a s X 1 5 v v N 7 5 9 / + 3 5 n 1 a / 2 5 r + A P 1 F 7 f e 1 P 9 T e 1 j Z q f 6 l 9 V 9 u v N W u t m l X 7 R + 2 f t X / V / v 3 V f 1 + v v v 7 m 9 d o k 9 P P P p m N + V y s 9 X v / 5 f 7 q q L h U = < / l a t e x i t > U x : F (↵, x) = 0 < l a t e x i t s h a _ b a s e = " o D F u s Z / t T v x q / c Y W J Z r m P i j E Q = " > A A A f l X i c l Z n Z b u P I F Y Y k i b D J R S C x E S M R i a B u E b w S Q I M M D I m L l l d o V u o I o s S J W u U r T U H L L n i a y W e J k V K r v Q l D z T A m y R + q W D y s q v N L p H n x s n G x v + + / x H P / J T / x c X f v H L X / N + + / O N H A l J a t A v l L W e x N x A X C d u o n b S A r m c + + a C R + U y N g N g t k H I l H n U D E t l i j U e f N P k l R j x n f G A W J P X k / X R Z N w o b y Q a z M J Z j j R Z d R u t G E k e O / W N t v F C + j e r A O i r z V t z p d / b J h a A E S S W x + L Y X M j S h z J h P X s R k x R z G I m L W g H X F g z o M m C / i x y y Y n x V h H b c E K p / o L E K C h t k T E / j s c + V E + S r x v J b D R d r D M H H + Z i Q T R M R G B N L + Q M P S M J j T x j h u K Y S X e W B w S p q r I b V Y J Z i c p r + S p e N Q B P X / d Y l G e l i b W U s C N U D q R o q v G c R U K y J J S Z f n h J D M T M U p i K z t p n k W B / M n s + m w U c n r m u v W u B b P Q S q I W z E A W H v s D O T H v g T b K d D n x i b r W y p J h W S r j k z v z x C F B H q M S P S Y u Q d F k n D r N E P M m m b t i u r O g r t X w w N y t K j b d p m v M p P B u l n x d m y U M l s d x g j e H a + N B O e G n b z P + v l F O t y M P W Y Z y t t L O h w l j b V K t q V j b M g o q C z r h C I Z j i Y P m A A X o Z G u b k p k f z y / x K Z L j B P M u Y D z V B b i B x / / n J s p v e d J u + v z K t + l + W B Y k S O i g Z V x o q / R m q a B G W T u v d R S q u V Y j q W u J t l s T n q Q G h k L z C d r V I U q k a x R W G d V Y X W a l J R j L J + u p n U + P D U T f X y Z x N e T J x J g X Y s q e + G S C B r y Y W t x U H v + i o f l R U I g / R p B x p G p M L m v R a r / W J s O V p P X d Q Y r h M A F l g Z V Z x i d e G y j J w y n i J R o z K w x b a f m E G n l E n M b / r F I n h F k g U U y g T x / v f K O v F H b a p J v Z i T z b k O / a G T w X D v M d l K V c x k g T p A d c X x E z V I I T f j f l / f Q / y i u G E L L F F B V R J / k C y L y R f I S k d / v S r b p q t s L H g Z D F D q / K j C S R f b / F r a o K V y s s H P X U D Z D m X h i p v R h b X P a L G Y S s K k b x T N j N J o o J f Q S j a U / G b B Q Q i Y v / Z k y P z O K P W h F Y j c j Z N k t M m T s S J H h c g r I o d o V U Q L o l R b Y i i I g q I j l N R H a j d i t i F K u I P Y h u R T y E K + I f Y i D i j i A F V E D J f E X I Q U U M I I Y V M Y Q Y V c Q I l N F f I I o K K E G F f E G G J S E R O I w o h J h W x B T i c V h j i q i C O I o h v i x I n s N r p L X y t g O P d t Q b t N m G d I p l + q K L eV Z s R i t Q B F X A b c F v D B m B D w x A H Q A X c A P c A f c F / D A A D D Q B D z U A j z S B j w W M M T w B M N m B N D U B T z U A z z T s A X Y v A c F z D C A L D S B L z W A r z S B r w W s M b w B s N B t D W B b z W A z T B w f t E O F b M g N m I h X c f w w m c h c s w b R o b D b Y l q Y c D c I t Y I t g G g m W A h w Q b g q z o h H e B u w T g H s E u A u w X g P s E D A H B H r B H s A / s E x w A B w S H w C H B E X B E B P w E E S W B I c A c E J A J w U P g I c E p c E r w M / A z w S P g E c F j D H B H E X d L z L F e n s H U P d f S r e E Y E P D B u C O h j u A u x r u A u p u A e r + E + I G G B C H G h C H m l B H i s T H g i Y Y n g E N m C n G p C n m l B t j S s A V r u E I W G F C X G l C X m l B X i t T X g j Y Y g G N C G t C m l B i v b u k D B S n B o p O C l O n R S s F K d W C l K U y F M W p m Y K b t R N w U x a q f g p z j U z B U n B o q O C p O H R U s F a e W C p K U F U W p q Y K r t R V w V Z x a q v g q z j V T B W n B o r O C t O n R W s F a f W C t K U F c W p u Y K t R d w V x a q / g r z j V z B Y n B o s O C w + d V j l z s t e K O n n l m u f q d D p V y S / U P x Q f Z w v + s u v s w E f t z H R V L L Z F n x p V o V F l J S r K L T S K i u o t d I u K y i U p Q V F F z p l B W U X N k t K i s l d W U H i l W Z Q e W / r K D y k F Z Q Q W W X l n x p g Z Y + o Y b G y T g t l j g x m z x K b / W G c G H Y Y / D k x q w D B a M D a m U D d S J l G P R E b / q z v o N w S r Y M y w q q t o z K C g q f C o r q N S l h W U b x m X F V R w m Z Q V F H E L C u o z I t K y j l r m s o J r L U V l B Q Z f j s o K a L j + W l a K s q d y M c u z s J L w p X j L V E S x U B r I T c S U k j E Q U Q n G s U c m K y N A I m C y O g c m q S A U w W R K p A z W Q o F J o s h Q G T l Z A O g c k y S P v A Z A k A C y A F I P C P Y B / Y J J o m m G Q B y W R O I A y k M n Y D K N U w l M n A a A f o Q I n i N C p K T O Z t + g x M J m A i Y z N h D k + m a f g S e z V W z + C U / M s l M L v B G F S/ F j X s q J X f e r G O / h c P b j Z e r / f u v z / + u A f o L p / q P p l V t s / b r e g q r d l z a v + q / b v n p / V + / + s q z u r e N P T z z Z t f l c r v V b P / g s j g A < / l a t e x i t > 1962 ). NeuRLP has three critical advantages over traditional solvers, leading to more efficient learning over longer sequences than traditional sequential solvers. These are: (i) step parallelism, i.e., being able to solve for hundreds of ODE time steps in parallel, allowing for faster solving and efficient gradient flow; (ii) batch parallelism, where we can solve in parallel on GPU batches of independent systems of ODEs in a single forward pass; (iii) learned step sizes, where the step sizes are differentiable and learnable by a neural network. This makes the NeurRLP solver ideal for training efficiently neural networks that model complex dynamic systems.↵ = [c i,t,x , k,t,x , b t,x , ! x , s t ]Relevance for scientific applications. Machine Learning for dynamical systems has adopted specialized methodologies such Physics-Informed Neural Networks (Raissi et al., 2019) for solving PDEs, neural operators (Li et al., 2020c) for prediction, linear regression with basis functions for discovery (Brunton et al., 2016 ), Neural ODEs (Chen et al., 2018) for dynamical systems. Being able to weave in governing equations in neural representations and solve them efficiently, Mechanistic Neural Networks potentially offer a stepping stone for broad scientific application of machine learning (Figure 1 ). We demonstrate this with experiments on tasks from each of above settings. 
mechanistic neural networks: 
 Formally, a Mechanistic Network contains a mechanistic encoder in a mechanistic block that takes an input x and generates a differential equation U x as representation according to equations 2 and 1. The family of ordinary differential equationsEQUATIONrepresents a broad parameterization for the ODE representation, with an arbitrary number d of linear terms with derivatives u (i) and an arbitrary number r of nonlinear terms g k including derivatives u (k) , k = 1, ..., d. The coefficients c i (t; x), ϕ k (t; x) for the linear and nonlinear terms are functions that possibly depend on the time variable t (thus non-autonomous ODEs), and on input x.After computing the ODE representation U x , we solve it with our specially designed parallel solver NeuRLP for n time steps and get a numerical solution as output of the mechanistic block: z = solve_ode(U x , ω x , n). ω x includes initial or boundary conditions and variables controlling step sizes that may be specific to the input x. ω x can be learned by NeurRLP, unlike traditional solvers. Efficient Quadratic relaxation and extension. Solving ODEs using the LP method inside neural networks has three main obstacles: 1) the solutions to the LP are not continuously differentiable (Wilder et al., 2019) , 2) solving linear programs is generally done using specialized solvers that do not take advantage of GPU parallelization and are too inefficient for neural networks applications, and 3) The matrices A are highly sparse where dense methods (such as from Amos and Kolter (2017)) are infeasible for large problems. We solve these problems by reframing the problem as an equality constraint-only quadratic program with regularization to ensure boundedness and the use of sparse methods for large problems. We cover further details including error analysis, complexity and numerical validation in A.4, A.2, A.3. We also extend the LP method to nonlinear ODEs by combining learning with solving in A.6. The presentation above applies the quadratic approximation to the linear program 5. The same approximation can also be applied to the dual of 5 and in practice can lead to a more efficient implementation. 
example applications in scientific ml: 
 We benchmark Mechanistic Networks in five different settings including discovery of equations, PDE solving, , PDE solving, n-body prediction, physical parameters discovery and time series forecasting (with the last two in the appendix) from scientific ML applications. (1) the Lorenz system, and (2) ODEs with complex nonlinear function of the form d dt x(t) = F (Ax(t)), where A is a linear transformation and F is a nonlinear function such as tanh, and (3) ODEs with rational function derivatives, d dt x(t) = p(x) q(x) , where p and q are polynomials. ODEs with non-linear functions cannot be modeled by the approach employed by SINDy. Results are shown in Figure 2 and Figures 6 and 7 in the appendix. For the Lorenz system which can be described as linear basis combinations, both SINDy and variants, as well as Mechanistic NNs recover the exact equations. For complex nonlinear and rational function ODEs which requires We adapt Mechanistic NNs from ODEs to PDEs. For 1-d PDEs, we simply model spatial dimensions with independent ODEs. With a spatial dimension of 256 and prediction over 10 time steps, we learn 256 ODEs for 10 time steps each. For 2-d PDEs we use a neural operator model with stacked MNN layers. We provide further details of the model and training and visualizations in appendix C.5. Following Brandstetter et al. (2022) , we compare relative MSE loss using Lie-symmetry augmented ResNet, FNO and autoregressive FNO on the 1-d KdV equation (Figure 3 ) and with FNO on 2-d Darcy Flow (Li et al., 2020c) (Table 2 in the appendix).We use 50 second 1-d KdV equation data and predict for 100 steps. We use 10 time steps of history as opposed to the baselines which use 20 time steps. We also show visualizations for KdV prediction in Figure 3 on a 100 sec dataset. Mechanistic NNs are competitive with FNO and augmented models without using any specialized adaptions for stable rollout (Brandstetter et al., 2022) N-body prediction. The task is to predict future locations and velocities of all bodies in a system given past observations of locations and velocities. We use Neural ODEs (Chen et al., 2018; Norcliffe et al., 2020) as a gold standard baseline and an MNN with second-order ODEs for our model. We use planetary ephemerides data from the JPL Horizons database for solar system dynamics (Giorgini, 2015) . The data is positions and velocities for the 25 largest bodies in the solar system from 1980 to 2015 with a step size of 12 hours. We use the first 70% of the data for training and the rest for evaluation.At training the MNN model predicts the next 50 steps given 50 input steps. At testing we rollout predictions for 2000 steps given the starting 50 steps. See prediction rollouts for Earth and Mars in and evaluation losses in Figure 4 . Mechanistic NNs improve Neural ODEs significantly by at least a factor of 10. 
conclusion: 
 Mechanistic Neural Networks (MNNs) are an approach for modeling complex dynamical and physical systems in terms of explicit governing mechanism. MNNs represent the evolution of complex dynamics in terms of families of differential equations making them flexible and able to model the dynamics of complex systems combined with a specialized solver. We demonstrate the effectiveness of the method with experiments in diverse settings.A FURTHER DETAILS FOR SECTION 3A.1 LINEAR PROGRAMS A linear program in the primal form is specified by a linear objective and a set of linear constraints.minimize c t x subject to Ax = b, x ≥ 0 (9)where A ∈ R m×n , c ∈ R n , b ∈ R m the following specifies a linear program. Matrix A and vector b define the equality constraints that the solution for x must comply with. c t x is a cost that the solution x must minimize. The linear program can also be written in dual form, minimize b t λ subject to A t λ = c.(10) 
 
 Appendix: We consider the case of a second order linear ODE with an N -step grid. For simplicity we consider a fixed step size h, i.e.,EQUATIONEQUATIONas Taylor approximations and ũ′′ t is obtained by plugging the approximate values in the ODE 11. We consider the following Taylor constraints (expressions 7 8) for the function and its first derivative. We use the absolute-value error inequalities for conciseness, the case for equalities is similar.EQUATION)Step t = 1. From Taylor's theorem we have that for the first step, t = 1,u(h) = ũ1 + O(h 3 ) (16) u ′ (h) = ũ′ 1 + O(h 2 ) (17) From 14, 15 u 1 = ũ1 + O(ϵ + h 3 ) (18) u ′ 1 = ũ′ 1 + O(ϵ/h + h 2 ) (19)This implies a local error at each step of O(ϵ + h 3 ) in u t .Step t = 2. To estimate the error at step 2 we need to estimate the error in u ′′ 1 at step 1. For u ′′ 1 we get the error by multiplying the error in u 1 by c0 c2 and that of u ′ 1 by c1 c2 and adding.EQUATIONNotice that u ′′ 1 always appears with a coefficient ofEQUATION)Each of the terms u 1 , hu ′ 1 , h 2 u ′′ 1 contribute an error of O(ϵ + h 3 ) to u 2 plus an additional error of O(h 3 ) arising from the Taylor approximation and an error of ϵ arising from the inequalities 14, 15.N Steps. Proceeding similarly, after N steps we get a cumulative error of O(N (ϵ + h 3 )).For ϵ ≈ h 3 and N ≈ 1/h, we get an error of O(h 2 ) for N -steps under the assumption that c0 c2 is O( 1h 2 ) and c1 c2 is O( 1 h ). The analysis implies that the cumulative error can become large for equations where c0 c2 , c1 c2 are large. Or, in little omega notation, c0 c2 is ω( 1 h 2 ) and c1 c2 is ω( 1 h ) .A.3 EFFICIENT QUADRATIC RELAXATION Solving ODEs using the LP method inside neural networks has three main obstacles: 1) the solutions to the LP are not continuously differentiable (Wilder et al., 2019) with respect to the variables A, b, c that interest us and 2) solving linear programs is generally done using specialized solvers that do not take advantage of GPU parallelization and are too inefficient for neural networks applications, and 3) The matrices A are highly sparse where dense methods for solving and computing gradients (such as from Amos and Kolter ( 2017)) are infeasible for large problems.We can avoid the non-differentiability of linear programs by including a diagonal convex quadratic term (Wilder et al., 2019) as a regularization term, converting inequalities into equalities by slack variables and removing non-negativity constraints (Pervez et al., 2023) to obtain an equality-constrained quadratic program,EQUATIONwhere G = γI, γ ∈ R is a multiple of the identity for a relaxation parameter γ and ξ are slack variables. Importantly, equality-constrained quadratic programs can be directly and very efficiently solved in parallel on GPU (Pervez et al., 2023) . This is why we rewrite inequalities as equalities using slack variables. Although with equalities only we lose the ability to explicitly encode non-negativity constraints, we mitigate this by regularization making sure that solutions remain bounded.Forward propagation and solving the quadratic program. We can solve the quadratic program directly with well-known techniques (Wright and Nocedal, 1999) , namely by simplifying and solving the following KKT system for some λ ∈ R m ,EQUATIONFor smaller problems, we can solve this system efficiently using a dense Cholesky factorization. For larger problems, we use an indirect conjugate gradient method to solve the KKT system using only sparse matrix computations. Both methods are performed batch parallel on GPU.Backward propagation and gradients computation. In the backward pass, we need to update the ODE coefficients in the constraint matrix A and β. We obtain the gradient relative to constraint matrix A by computing∇ A ℓ(A) = ∇ A z∇ z ℓ(z), where ℓ(.) is our loss function and z is a solution of the quadratic program.We can compute the individual gradients using already established techniques for differentiable optimization Amos and Kolter (2017) with the addition of computing sparse gradients only for the constraint matrix A. Briefly, computing the gradient requires solving the system equation 23 for with a right-hand side containing the incoming gradient g:EQUATIONThe gradient ∇ A ℓ(A) can then be computed by first solving for d z , d λ and then computing d λ z ⊤ + λd ⊤ z (Amos and Kolter, 2017) . In general, this would produce a dense gradient matrix, which is very memory inefficient for sparse A. We avoid this by computing gradients only for the non-zero terms of A by computing sparse outer products.The method proposed in Young (1961) does not add a smoothness constraint for the highest order derivative term, since derivative information is not available. In cases where a more accurate highest order term is required, we also add a central difference constraint as a smoothness condition on the highest order term.A.6 NONLINEAR ODES The standalone solver described above works for linear ODEs. When combined with neural networks, we can extend the approach to nonlinear ODEs by combining solving with learning.For each non-linear ODEs term g k , we add an extra variable ν k,t with coefficients ϕ k,t,x corresponding to the non-linear term to our linear program.We rewrite our nonlinear ODE in equation 4 asEQUATIONEQUATIONEQUATIONthat is, for each nonlinear term k and for every time step t we also add in z an auxiliary variable ν k,t . Additionally, we include derivative variables ν i k,t that are part of the Taylor approximations to ensure smoothness. We then solve the linear part of the above ODE, that is equation 25 subject to 27 with the linear programs we described in the previous subsections. Further, we convert the nonlinear part in equation 26 to a loss term ν k,tg k (t, u t , u ′ t , . . .) 2 , which is added to the loss function of the neural network. With the extra losses, we learn the parameters ϕ such that ν k is close to the required non-linear function of the solution.Figure 14 shows solving and fitting of a non-linear ODE.Nonlinear ODEs for Discovery. When building MNN models for governing equation discovery, we incorporate nonlinear ODEs using a set of predefined basis functions {θ i }, such as the polynomial basis functions (Brunton et al., 2016) , to build an equation of the formEQUATIONThe input u ′ x (t) to the basis functions are generated by a neural network with input x as u ′ x = NN(x), where u ′x (and possibly x) depends on time t. To ensure that this is a proper nonlinear ODE we add a consistency term to the loss function to minimize the squared loss (u(t; x)u ′ x (t)) 2 . This ensures that the basis input and ODE solution are close.The computational and memory complexity of MNNs is determined by the size of the time grid n, and the order d of the ODEs to be generated. The last layer of f outputs n × (d + 2) ODE parameters. This means that the memory required to store the coefficients can be large depending on the grid size and dimension. The main computational effort in solving the system equation 23 for a batch of ODEs, which we do by a Cholesky factorization for small problems or sparse conjugate gradient for large ones. Cholesky factorization has complexity cubic in nd while conjugate gradient has quadratic complexity.Benchmarking against RK4 from scipy and torchdiffeq. We compare with traditional ODE solvers on second-and third-order linear ODEs with constant coefficients from the scipy package. For a learning comparison, we also compare with the RK4 solver with the adjoint method from the torchdiffeq package on the benchmark task of fitting long and noisy sinusoidal functions of varying lengths. The quantitative and qualitative results in Appendix D show that NeuRLP is comparable to standard solvers on the linear ODE-solving task. On the fitting task NeurLP significantly improves upon the baseline and is about 200x faster with a lower MSE loss than the torchdiffeq baseline for 1000 steps.NeuRLP can learn time steps. Unlike traditional solvers, NeuRLP can learn the discretization grid for learning and solving ODEs, becoming adaptively finer in regions where the fit is poor. We validate this on fitting a damped sinusoid, see results in Figure 13 , where we begin with a uniform grid and with steps becoming denser in regions with bad fit.Problem. Often, the problem is not to discover the governing equations in a system but the most fitting physical parameters explaining the observations. Applications include inverse problems in dynamical systems (Wenk et al., 2020) .Gold standard. We use second order Neural ODEs (Norcliffe et al., 2020) to fit ODE models of corresponding to Newton's second Law, matching corresponding derivative coefficients to infer the physical parameters. Mechanistic NNs for discovering physical parameters.We use a second order ODE MNN with a time invariant 2nd order coefficient to match Newton's second Law. The force is learned by a neural networks as a function of position.Experiment. We design an experiment with two bodies with masses m 1 = 10, m 2 = 20, distance d and initial velocities v 1 , v 2 , moving under the influence of Newtonian laws, and gravitational force, F = G m1m2 r 2 r, r being the unit vector of direction of force, G = 2 the gravitational constant. We generate a single random train trajectory for the two bodies for 40k steps. The physical parameters we infer are mass ratio m1 m2 and distribution of force values F = [F x , F y ], by combining Newton's second and third law. We show quantitative and qualitative results in Figure 5 . Since forces are only determined up to a constant, to compare forces we normalize by dividing by the force at the first step. Neural ODE and Mechanistic NNs estimate the mass ratio while MNNs perform significantly better at estimating the force distribution and Neural ODE forces often have the incorrect direction as shown by the negative cosine similarity averaged over the entire trajectory.Problem. Time-series modelling and future forecasting is a classical statistical and learning problem, usually with low-dimensional signals, like financial data or complex dynamical phenomena from sciences.Gold standard. We compare with Neural ODE and variants including second order and augmented Neural ODEs.Mechanistic NNs for time series. We use a basic Mechanistic NN second-order ODE for this experiment.Experiment. We validate on the benchmark of modeling the accelerations a 2 produced over time by a shaker under a wing in aircrafts (Norcliffe et al., 2020) . The model sees 1,000 past time accelerations a 2 and predicts the next 4000. We show quantitative results in and qualitative results in Figure 10 and C .3 in the appendix. The distribution of predicted a 2 at test time are very close to the true ones. Mechanistic NNs are on-par with second-order ODE, converge significantly faster, and achieve two times lower training error showing they can model complex phenomena.We give further details regarding the ODE discovery setup. This method follows the SINDy Brunton et al. ( 2016) approach for discovering sparse differential equations using a library of basis functions. Unlike SINDy, which resorts to linear regression, the MNN method uses deep neural networks and builds a non-linear model which allows modeling of a greater class of ODEs.The method requires a set of basis functions such as the polynomial basis functions up to some maximum degree. Over two variables x, y this is the set of functions {0, x, y, x 2 , xy, y 2 , xy 2 , . . . , y d } for some maximum degree d. Let k denote the total number of basis functions.Next we are given some observations X = [(x 0 , y 0 ), (x 1 , y 1 ), . . . , (x n-1 , y n-1 )] for n steps. We first transform the sequence by applying an MLP to the flattened observations producing another sequence of the same shape.X = [(x 0 , ỹ0 ), . . . , (x n-1 , ỹn-1 )] = MLP(X)We apply the basis functions to X to build the basis matrix Θ ∈ R n×k .EQUATIONLet ξ ∈ R n×2 be a set of parameters, with each column specifying the active basis functions for the corresponding variable in [ ẋ, ẏ] .The ODE to be discovered is then modeled asEQUATION)where f is some arbitrary differentiable function. Note that for SINDy X = X and f is the identity function and the problem is reduced to a form of linear regression adapted to promote sparsity in ξ. SINDy estimates the derivatives using finite differences with some smoothing methods.With MNN the ODE 30 is solved using the quadratic programming ODE solver to obtain the solution xt , ȳt for t ∈ {0, . . . , n -1}. The loss is then computed as the MSE loss between xt , ỹt , xt , ȳt and the data x t , y t .loss = 1 N t (x t -x t ) 2 + (ỹ t -y t ) 2 + (x t -x t ) 2 + (ȳ t -y t ) 2We show two examples of cases where F (u) is a rational function (a ratio of polynomials) and when F (u) is a nonlinear function of Θξ. Moreover, unlike SINDy, MNN can learn a single governing equation from multiple trajectories each with a different initial state making MNN more flexible. In many situations a single trajectory sample is not enough represent to the entire state space while multiple trajectories allow discovery of a more representative solution.Planar and Lorenz System. We first examine the ability of MNN equation discovery for systems where the true ODE can be exactly represented as a linear combination of polynomial basis functions. We use a two variable planar system and the chaotic Lorenz system as examples. Both MNN and SINDy are able to recover the planar system. Simulation of the learned Lorenz ODE are shown in Figure 6 for MNN and SINDy.Next we consider ODE systems where the derivative cannot be written as a linear combination of polynomial (or other) basis function. Nonlinear Function of Basis. First, we consider systems where the derivative is given by a nonlinear function of a polynomial. For simplicity we assume that the nonlinear function is known. As an example we solve the system from Figure 2 with the tanh nonlinear function.Vector fields for the learned systems are shown in Figure 2 for MNN and SINDy. We see that although SINDy fits the training example, the directions diverge further away. With MNN we see that the learned vector field is consistent with the ground truth far from the training example even though we use only a single trajectory.Rational Function Derivatives. Second, we consider the case where the derivative is given by a rational function, i.e., F (u) = p(u)/q(u), where p and q are polynomials. Such functions cannot be represented by the linear combination of polynomials considered by standard SINDy, however such functions can be represented by MNNs by taking p and q to be two separate combinations of basis polynomials and dividing. An example is shown in Figure 7 in the appendix for the system where we see again MNNs learning much better equations compared to SINDy with a second-order polynomial basis tha overfits. Further, by including more trajectories in the training, results improve further, see x ′ = -10.0003x + 10.0003y y ′ = 27.9760x + -0.9934y -0.9996xz z ′ = -2.6660z + 0.9995xyx ′ = -10.000x + 10.000y y ′ = 27.998x + -1.000y + -1.000xz z ′ = -2.667z + 1.000xy MNN Non-linear x ′ = tanh(-0.7314x + 0.5545y+ -1.2524x 2 + -0.1511xy + 0.2134y 2 ) y ′ = tanh(0.9879x + 1.0005y + 0.1742x 2 ) SINDy Non-linear x ′ = -1.968x + 0.985y + -0.054x 2 y ′ = 1.x ′ = -0.9287x + 0.4386y + -1.1681x 2 + 0.3545y 2 0.4871 + 0.8123x + 0.0984x 2 + 0.3700xy + 0.3081x 2 y ′ = 0.6360x + 0.5971y + 0.3267x 2 0.6090 + 0.7507x 2 + 0.5694y 2x ′ = -1.705x + 0.899y + -0.318x 2 y ′ = -0.795 + 3.072y + 4.777x 2 + 6.892xy + -4. We test MNN on the nested spheres dataset (Dupont et al., 2019) , where we must classify each particle as one of two classes. This task is not possible for unaugmented Neural ODEs since they are limited to differomorphisms (Dupont et al., 2019) . We show the results in Figure 9 , including comparisons with Neural ODE (Chen et al., 2018), Augmented Neural ODE and second-order Neural ODE (Norcliffe et al., 2020) . MNNs can comfortably classify the dataset without augmentation and can also derive a governing equation.We use a second order ODE with coefficients computed with a single layer and the right hand side is set to 0. We use a step size of 0.1 and length 30. However, as we note, 5 time steps are enough for accurate classification. The loss function is the cross entropy loss.MNNs obtain an explicit linear ODE per datapoint that governs the evolution of the point. The example we give is for one of the ODEs for one point and for a 5-time step evolution. This computed equation is sufficient for perfect classification.MNNs can learn complex dynamical phenomena significantly faster than Neural ODE and second order Neural ODE. We reproduce an experiment with a real-world aircraft benchmark dataset (Noël and Schoukens, 2017; Norcliffe et al., 2020) . In this dataset the effect of a shaker producing acceleration under a wing gives rise to acceleration a 2 on another point For this part of the experiment we use an MNN with a restricted ODE to match Newton's second law. In the MNN model for this experiment, we use the same coefficient for the second derivative term for all time steps with the remaining coefficients fixed to 0, that is c 2 (t) = c and c 1 (t) = 0, c 0 (t) = 0. b(t) = F t corresponds to the force term which is computed by a neural network from the initial position and velocity with two hidden layers of 1024 units and Newton's second law F 21 = -F 12 . We use a step size of 0.01 and run for 50 epochs.The baseline is an SONODE designed to correspond to Newton's second and third law with an MLP for force as above.C.5 PDE SOLVING 1d Model. For 1d problems we use a simplest possible model of modeling the spatial dimension by independent ODEs. We use a history of 10 time steps and predict for 9 time steps in one iteration, using the last time step as initial condition for the ODE. During evaluation we predict and evaluate for 100 steps. We use 3rd and 4th order ODEs. The coefficients for ODEs, step sizes the right hand side (b) are computed by 1d ResNets with 10 blocks. We use the L1 loss which we find improves rollout performance.2d Model. In Figure 11 we show the MNN architecture we used to solve PDEs. We use the 2d Darcy Flow dataset used by Li et al. (2020c) scaled to 85x85. The ODE is solved for 30 steps and the entire soluton trajectory is then upsampled and combined with the input features map. The network is built by stacking three such modules together plus an input MLP layer and an output layer. We perform an simple experiment illustrate how the ODE learning method can fit ODEs to noisy data. We generate a sine wave with dynamic Gaussian noise added during each training step. We train two models: the first a homogeneous second order ODE with arbitrary coefficients and the second a homogeneous second order ODE with constant coefficients. We also train a model without noise. The results are shown in Figure 15 . The figures show that the method can learn an ODE in the presence of noise giving a smooth solution. The model with constant coefficients learns the following ODE. 0.92023u ′′ -0.00016u ′ + 0.228u = 0, with (learned) initial conditions u(0) = -0.031799 and u ′ (0) = 2.3657. We compare NeuRLP with the RK4 solver from torchdiffeq on a task of fitting noisy sinusoidal waves of varying lengths. We compare MSE and time in Table 3 and Figure 16 . We show learned trajectories for a 2-body prediction problem with an MNN on synthetic data in Figure 18 . The objects are generated using the gravitation force law for 4000 steps and the first half are used for training and we predict the second half. Source code is available at https://github.com/alpz/mech-nn 


Decision: Authors propose Mechanistic Neural Networks (MNNs) with a differentiable ODE solver NeuRLP for tasks like equation discovery and PDE solving. Some aspects of MNNs need clarification like the model architecture details, comparisons to state-of-the-art methods, and providing more examples of discovered equations. For the camera-ready version, I encourage authors to clarify the MNN components and information flow, provide more implementation specifics on experiments, and further highlight where NeuRLP offers advances over prior differential equation solvers.

Reviews: review: title: title: Proposed NeuRLP solver is novel and important addition for SciML; Experiments with proposed architecture need more clarifications
main_review: main_review: Paper proposes NeuRLP solver that uses differentiable quadratic programming to solve ODEs as an alternative to traditional solvers. The solver is used within the proposed Mechanistic NN block that consists of an ODE representation whose parameters are learned with an neural encoder. Output of the MNN block is obtained by solving the corresponding ODE. 


Strengths:
1. NeuRLP uses differentiable quadratic programming with equality constraints to solve the ODE. Method can be extensively parallelized, even across time steps. Table 2 and Figure 18 show that proposed method is significantly faster than the traditional solvers like RK4, especially on longer sequences. This component can be beneficial in many SciML methods that incorporate ODE within the architecture.
2. NeuRLP is able to solve ODEs with nonlinear terms using auxiliary variables in the quadratic programming problem. 
3. Experiments using NeurLP and MNN show improvements in forecasting tasks in ODEs and PDEs. 

Weaknesses:
1. Section 2 can be clarified to better distinguish the "ODE discovery" and "Prediction/forecasting" cases of the MNN models. Currently, these cases seem disconnected and require a back-and-forth with Appendix. 
2. Nonlinear ODE discovery experiments: I am not convinced if the proposed MNN method is better than SINDy for discovering ODEs.
	1. In Appx B.3, MLP(X) is used to obtained a nonlinear representation $\tilde{X}$ of X and the ODE is learnt over this representation. But the loss function seems to enforce that $\tilde{X} = X$ which negates the use of this MLP. This can also be seen from the learnt equations in C.1 that do not contain the effect of MLP. 
	2. When the ground truth dynamics has a nonlinear function applied at **last**, e.g., tanh or of the form p/q, MNN is given the exact form of this function whereas SINDy is not. I believe it should be easy to adapt SINDy to the case when this final function is known and nonlinear.  
3. Since the paper claims interpretability, it would be good to show the equations learnt within MNN for all experiments (e.g., n-body experiments). 
4. Please provide the following details in the paper: 
    1. What was the exact family of ODEs $U_x$ used in experiments (e.g., what were $g_k$)?
    2. What was the maximum number of ODE terms used in experiments? Do the proposed methods work even with large number of ODE terms (during both discovery and prediction)?
   
Overall, NeuRLP proposed in the paper is novel and significant for the SciML community. Some clarifications are required for its proposed use in the MNN architecture (especially for ODE discovery).
title: title: Reasonable and exhaustive experimental results with several unclear aspects
main_review: main_review: ### Summary:
Mechanistic Neural Networks (MNN) are introduced for equation discovery, PDE solving, and forecasting. The method implements a mechanistic encoder that translates the input into an ODE and, at the same time, predicts the step size to solve the equation. Subsequently, a neural Relaxed Linear Programming Solver integrates the inputs to generate a prediction. Results underline MNN's ability to discover equations and accurately forecast dynamic systems, while not qyite reaching the accuracy of state-of-the-art methods in PDE solving.

### Strengths:
- Exhaustive experimentation with additional details provided in the appendix.
- Wide applicability of the method to different tasks.
- Accelerating equation discovery by providing an efficient GPU implementation for parallel computing.

### Weaknesses and Concerns
- Unclear what the exact inputs and outputs of the model components are. E.g., what does $f_{\theta}$ receive and concretely produce as output and how are these outputs then composed as an ODE?
- MNN is compared against a fairly old version of SINDy. More recent variants, such as SINDy-PI [[1]](https://royalsocietypublishing.org/doi/full/10.1098/rspa.2020.0279) in fact are able to recover the coefficients of nonlinear systems.
- While MNNs can be applied to various tasks, it does not seem to reach FNO's performance in PDE solving.
- In Figure 2, it is unclear to me how MNN can technically learn the entire vector field accurately (including the region of the test sample) by only seeing the train sample. The SINDy solution looks realistic to me, given the method only receives the plotted training trajectory in the left end of the vector field. I understand that the two trajectories plotted in Figure 8 are more informative about the shape of the vector field, allowing MNN to better capture the overall vector field in that case, i.e., when receiving two trajectories. Can you explain why Figure 2 only contains one trajectory, verify whether MNN is only receiving that single trajectory, and explain how MMN technically can get such an accurate notion of the complete vector field then? Also, how does SINDy (or SINDy-PI) perform when receiving two trajectories during training?

### Questions:
1. On page 1, what are large batches and do you seek to learn many ODEs simultaneously from different families, or do you vary some coefficients in an ODE and call this many ODEs?
2. If the step size is a learnable parameter, how do you ensure that it does not converge to zero in order to generate the most accurate results?
3. Can you show some examples of what $\omega_x$ converges to? Would be interesting to see what step size it induces under what circumstances.
4. In Equation (2), can you specify what $b$ is and why the right-hand-side is not zero as in traditional ODE representations? That is, do you allow for a bias depending on $t$ and $x$ and how do you account for it and model it?
5. in "Efficient Quadratic relaxation and extension", is the LP method standing for linear programming? Maybe add as explanation.

### Minor comments:
- On page 1, space missing in "When training ODE representationswe..."
- In the \textbf{} commands, when opening new paragraphs, you sometimes do and do not use a dot, which might be made consistent.
- At the beginning of Section 4, PDE solving printed two times in "... discovery of equations, PDE solving, , PDE solving ..."
- Typo in **N-body prediciton** on page 4.

### Assessment
In balance I am undecided about this paper, which, on the one hand, offers exhaustive and reasonable results on a relevant topic, but, on the other hand, is vague about several aspects that need clarification.

### References:
[1] https://royalsocietypublishing.org/doi/full/10.1098/rspa.2020.0279